"""Step merging: at chosen steps, run the step a second time on the same picture with
different noise and blend the two answers.

Step N is answered twice. Picture A is the run's own answer. Picture B is the answer to
the same step after A's picture is re-noised to this step's noise level with a second
seed (and so is the sound, when the model has it). The two are then mixed:

    merged = base + strength * (other - base)

`Merge order` picks the base. "B into A" keeps A's picture and lets B's variation in;
"A into B" starts from B and lets A in. Two independent answers averaged lose contrast
(the independent detail partly cancels), so the mix is rescaled to the base's own spread.

"Second run" keeps a whole second video going beside yours: its first input is the second seed's noise and every
later input is its own last answer re-noised, the way a few-step sampler chains its steps. At the named steps its
answer is blended into yours the same way. It costs one extra pass at every step up to the last one named.

"Simultaneous actions": with `[A|B]` in the prompt (or `A|B` between two shortcuts) and Alternating prompts >
"Read A and B separately" on, the named steps are answered
once for each option -- the same picture and noise, a prompt that holds only that option -- and the answers are
merged the same way, instead of the second-seed pass. The prompt outside those steps holds both ("A while B").

Costs one extra model pass per merged step. The second pass is marked as a probe, so no
learning feature records it.
"""

import random
import time

import math

import torch
from comfy.model_management import InterruptProcessingException
from comfy.patcher_extension import WrappersMP

from ..._core import dit_hooks, log, streams

ID = "noise_merge"
TITLE = "Step merging"
MOUNT = "generation.sampling"
STAGE = "sampling"
CATEGORY = "sampling"
STATUS = "experimental"
REQUIRES = ["temporal_latent"]

SETTINGS = {
    "enabled": {
        "type": "bool", "default": False,
        "label": "Step merging",
        "hint": "At the steps below, answers the step twice with different noise and blends the two. Costs one extra pass per merged step.",
    },
    "what": {
        "type": "enum", "default": "answers", "options": [
            {"value": "answers", "label": "Step answers"},
            {"value": "parallel", "label": "Second run"},
            {"value": "noise", "label": "Starting noise"},
        ],
        "label": "Merge", "ui": "segmented",
        "hint": "Blends in a second seed. Step answers: at chosen steps. Second run: a whole second video (slower). Starting noise: free.",
        "when": {"enabled": True},
    },
    "steps": {
        "type": "text", "default": "1",
        "label": "Steps",
        "hint": "Step numbers from 1, e.g. 1 or 2-3 or 1,4. Steps past the end of the schedule are ignored.",
        "when": {"enabled": True, "what": ["answers", "parallel"]},
    },
    "order": {
        "type": "enum", "default": "b_into_a", "options": [
            {"value": "b_into_a", "label": "B into A"},
            {"value": "a_into_b", "label": "A into B"},
        ],
        "label": "Merge order", "ui": "segmented",
        "hint": "A is your seed's picture, B the same step on a second seed. The first one named is blended in; the other is the base.",
        "when": {"enabled": True, "what": ["answers", "parallel"]},
    },
    "strength": {
        "type": "float", "default": 0.5, "min": 0.0, "max": 1.0, "step": 0.05,
        "label": "Merge strength", "ui": "slider",
        "hint": "0 keeps your own result untouched, 1 replaces it with the second seed's. For starting noise, 0.5 is an even mix.",
        "when": {"enabled": True},
    },
    "fresh": {
        "type": "float", "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.05,
        "label": "Second pass new noise", "ui": "slider",
        "hint": "How much of the step's noise the second pass swaps for the second seed's. 1: all of it. 0: no change.",
        "when": {"enabled": True, "what": "answers"},
    },
    "sound": {
        "type": "bool", "default": True,
        "label": "Merge sound too",
        "hint": "Blends the sound the same way as the picture. Off keeps only the base's own sound.",
        "when": {"enabled": True, "what": ["answers", "parallel"]},
    },
    "seed": {
        "type": "int", "default": -1, "min": -1, "max": 2147483647,
        "label": "Second seed",
        "hint": "The noise for picture B (or the second noise). -1 picks a new one every run; a number gives the same B each time.",
        "when": {"enabled": True},
    },
}


def _say(message, level=log.ALERT):
    log.once(f"{ID}:{message}", level, "FunUI Step merging", message)


def merge(base, other, strength):
    """base + strength * (other - base), rescaled so each clip keeps the base's spread
    (std around its own mean) -- a plain mix of two independent answers is flatter."""
    mixed = base + strength * (other - base)
    dims = tuple(range(1, base.dim()))
    base_std = base.float().std(dim=dims, keepdim=True)
    mixed_std = mixed.float().std(dim=dims, keepdim=True).clamp(min=1e-8)
    mean = mixed.float().mean(dim=dims, keepdim=True)
    return ((mixed.float() - mean) * (base_std / mixed_std) + mean).to(base.dtype)


def second_input(x, a, sigma, seed, step, fresh):
    """The input for picture B. Rectified flow: x = (1 - sigma) picture + sigma noise. A's answer is the picture;
    the noise is the step's own noise (what is left of x once A's picture is taken out) turned `fresh` of the
    way toward the second seed's noise, at unit variance. fresh 1 is all second-seed noise, 0 hands back x."""
    if sigma < 1e-4:
        return x
    own = (x.float() - (1.0 - sigma) * a.float()) / sigma
    gen = torch.Generator().manual_seed(int(seed) + 7919 * int(step))
    other = torch.randn(x.shape, generator=gen).to(device=x.device)
    theta = fresh * torch.pi / 2
    noise = float(torch.cos(torch.tensor(theta))) * own + float(torch.sin(torch.tensor(theta))) * other
    return ((1.0 - sigma) * a.float() + sigma * noise).to(x.dtype)


SIMUL = "funui_simultaneous"        # what Prompt markup leaves in transformer_options: {"options", "cond_len"}


def _mix(out_a, out_b, named, strength, b_first, with_sound):
    """Two full answers of one step -> (merged answer, how far apart they were), or None when either has no picture."""
    split_a, split_b = streams.video_of(out_a, named), streams.video_of(out_b, named)
    if split_a is None or split_b is None:
        return None
    (video_a, rebuild_a), (video_b, rebuild_b) = split_a, split_b
    tail_a, tail_b = _sound(out_a, video_a), _sound(out_b, video_b)
    diff = float((video_a - video_b).norm() / video_a.norm().clamp(min=1e-8))
    if b_first:
        base, other, base_tail, other_tail, rebuild = video_a, video_b, tail_a, tail_b, rebuild_a
    else:
        base, other, base_tail, other_tail, rebuild = video_b, video_a, tail_b, tail_a, rebuild_b
    merged = rebuild(merge(base, other, strength).to(base.dtype))
    if with_sound and base_tail is not None and other_tail is not None:
        n = merged.shape[-1] - base_tail.shape[-1]
        merged = torch.cat([merged[..., :n], merge(base_tail, other_tail, strength).to(merged.dtype)], dim=-1)
    return merged, diff


def _second_x(x, out, named, sigma, seed, step, fresh, with_sound):
    """The input for a second pass over `x` after `out` answered it: the step's noise turned `fresh` of the way toward
    the second seed's noise around out's picture (and sound), on x's own device and dtype. x itself when out has no picture."""
    split_x, split_a = streams.video_of(x, named), streams.video_of(out, named)
    if split_x is None or split_a is None:
        return x
    video_a, tail_a = split_a[0].detach(), _sound(out, split_a[0])
    x_b = split_x[1](second_input(split_x[0], video_a, sigma, seed, step, fresh))
    if tail_a is not None and with_sound:
        n = x_b.shape[-1] - tail_a.shape[-1]
        x_b = torch.cat([x_b[..., :n], second_input(x[..., n:], tail_a.detach(), sigma, seed + 1, step, fresh)], dim=-1)
    return x_b


def _sound(packed, video):
    """The part of a packed [B, 1, N] result that follows the picture (H3's sound), or None."""
    if getattr(packed, "dim", None) is None or packed.dim() != 3:
        return None
    n = video[0].numel()
    return packed[..., n:] if packed.shape[-1] > n else None


def _resized(video, shape):
    """`video` [B, C, T, H, W] brought to `shape` (only the picture's height and width may differ), or None."""
    if video.dim() != 5 or len(shape) != 5 or tuple(video.shape[:3]) != tuple(shape[:3]):
        return None
    return torch.nn.functional.interpolate(video.float(), size=tuple(shape[2:]), mode="trilinear",
                                           align_corners=False).to(device=video.device, dtype=video.dtype)


def mix_noise(noise, other, strength):
    """cos(t) noise + sin(t) other, t = strength * 90 degrees: two independent unit-variance noises
    stay unit variance (a plain average would shrink to 71% at an even mix, and the model reads that as a calmer, blurrier start)."""
    theta = strength * torch.pi / 2
    return (math.cos(theta) * noise + math.sin(theta) * other.to(device=noise.device, dtype=noise.dtype)).to(noise.dtype)       # plain numbers: no tensor to put on a device


def install_noise(patcher, values, key):
    strength = float(values.get("strength", 0.5))
    if strength <= 0.0:
        _say("off: strength is 0")
        return None
    chosen = int(values.get("seed", -1))
    run_seed = {}

    def sampler_sample(executor, model_wrap, sigmas, extra_args, callback, noise,
                       latent_image=None, denoise_mask=None, disable_pbar=False):
        if not isinstance(noise, torch.Tensor):
            _say("off this run: the starting noise is not a plain tensor")
            return executor(model_wrap, sigmas, extra_args, callback, noise, latent_image, denoise_mask, disable_pbar)
        second_stage = dit_hooks.stage_of(model_wrap) == 2      # Two stages: both stages mix with the SAME second seed
        if not second_stage or "seed" not in run_seed:
            run_seed["seed"] = chosen if chosen >= 0 else random.randrange(2 ** 31)
        seed = run_seed["seed"]
        other = torch.randn(noise.shape, generator=torch.Generator().manual_seed(seed)).to(noise)
        if not second_stage:
            log.info("FunUI Step merging", f"on: starting noise mixed with the second seed's ({seed}) at "
                     f"strength {strength:.2f}, no extra pass")
        return executor(model_wrap, sigmas, extra_args, callback, mix_noise(noise, other, strength),
                        latent_image, denoise_mask, disable_pbar)

    patcher.add_wrapper_with_key(WrappersMP.SAMPLER_SAMPLE, key, sampler_sample)
    return f"starting noise, strength {strength:.2f}, second seed {'random' if chosen < 0 else chosen}"


def install(patcher, values, key):
    if not values.get("enabled"):
        return None
    if values.get("what") == "noise":
        return install_noise(patcher, values, key)
    wanted, problems = dit_hooks.parse_blocks(values.get("steps"), 100000)
    wanted = {i - 1 for i in wanted if i >= 1}
    for p in problems:
        _say(f"ignored step entry: {p}")
    strength = float(values.get("strength", 0.5))
    if not wanted:
        _say("off: no valid steps named")
        return None
    b_first = values.get("order", "b_into_a") == "b_into_a"
    with_sound = bool(values.get("sound", True))
    fresh = float(values.get("fresh", 1.0))
    parallel = values.get("what") == "parallel"
    last = max(wanted)
    chosen = int(values.get("seed", -1))
    live = {"seed": max(chosen, 0), "merged": 0, "prev": None, "options": {}}

    def start(executor, *args, **kwargs):
        # Two stages is two sampler calls of one run: merges counted in the first still count, and the verdict comes after the last.
        stage = dit_hooks.stage_of(getattr(executor, "class_obj", None))
        live["seed"] = chosen if chosen >= 0 else random.randrange(2 ** 31)
        live["prev"], live["options"] = None, {}
        if stage != 2:
            live["merged"] = 0
        log.info("FunUI Step merging", f"on: {'second run, ' if parallel else ''}steps {values.get('steps')}, strength {strength:.2f}, "
                 f"second seed {live['seed']}")
        try:
            return executor(*args, **kwargs)
        finally:
            live["options"].clear()
            if stage != 1 and not live["merged"]:
                _say("Inactive | no step of this run matched the steps you named "
                     "(or the model was not called as a plain prompt step)")

    patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, key, start)

    def apply_model(executor, x, t, *args, **kwargs):
        named = streams.model_args(args, kwargs)
        to = named.get("transformer_options") or {}
        where = dit_hooks.current_step(to)
        simul = to.get(SIMUL) if not parallel else None
        context = named.get("c_crossattn")
        if (simul and where is not None and where[0] in wanted and not dit_hooks.probing(to)
                and all(int(i) == 0 for i in (to.get("cond_or_uncond") or [0]))):
            if (isinstance(context, torch.Tensor) and context.dim() == 3
                    and int(context.shape[1]) == simul["cond_len"] and len(simul["options"]) > 1):
                try:
                    # Strength 0 keeps the base's answer and strength 1 (with fully new noise, so the second run does
                    # not start from the first one's picture) is the other's alone: that one is the only pass made.
                    base_k, other_k = (0, 1) if b_first else (1, 0)
                    only = base_k if strength <= 0.0 else other_k if strength >= 1.0 and fresh >= 1.0 else None
                    outs, x_k = {}, x
                    for k in ((only,) if only is not None else (0, 1)):
                        args_k, kwargs_k = streams.with_context(args, kwargs, streams.option_context(
                            getattr(executor, "class_obj", None), simul["options"][k], context, live["options"], k))
                        args_k, kwargs_k = streams.with_options(args_k, kwargs_k, {**to, streams.OPTION_PINNED: True})
                        outs[k] = executor(x_k, t, *args_k, **kwargs_k)
                        if k == 0 and fresh > 0.0 and only is None:
                            # The second option starts from partly new noise: the step's own noise, turned `fresh`
                            # of the way toward the second seed's, around the first option's picture.
                            x_k = _second_x(x, outs[0], named, float(t.max()), live["seed"], where[0], fresh, with_sound)
                    if only is not None:
                        live["merged"] += 1
                        log.info("FunUI Step merging", f"step {where[0] + 1}: only {'AB'[only]} read "
                                 f"(strength {strength:.2f}, new noise {fresh:.2f}), nothing to merge")
                        return outs[only]
                    mixed = _mix(outs[0], outs[1], named, strength, b_first, with_sound)
                    if mixed is not None:
                        live["merged"] += 1
                        log.info("FunUI Step merging", f"step {where[0] + 1}: both actions answered separately, the "
                                 f"answers differ by {mixed[1]:.1%}, merged at {strength:.2f}")
                        return mixed[0]
                except InterruptProcessingException:
                    raise
                except Exception as exc:            # noqa: BLE001
                    log.broke("FunUI Step merging", exc, "merging the simultaneous actions")
                    return executor(x, t, *args, **kwargs)
            else:
                _say("Simultaneous actions skipped: another feature changed the prompt's conditioning")
        out = executor(x, t, *args, **kwargs)
        if strength <= 0.0 or where is None or dit_hooks.probing(to) or where[0] > last or (where[0] not in wanted and not parallel):
            return out                              # strength 0 keeps your own answer: no second pass to blend in
        if any(int(i) != 0 for i in (to.get("cond_or_uncond") or [0])):
            return out                              # a negative-prompt call: guidance is for the prompt's picture
        try:
            split_x, split_a = streams.video_of(x, named), streams.video_of(out, named)
            if split_x is None or split_a is None:
                _say("off this run: could not find the picture in this model's latent")
                return out
            sigma = float(t.max())
            video_a, rebuild_a = split_a
            tail_a = _sound(out, video_a)
            gen_step = where[0]
            if parallel:
                # The second run: each step's input is its own last answer (zero at the start) re-noised with the
                # second seed's noise, the way a few-step sampler chains its steps.
                prev = live["prev"] or (torch.zeros_like(video_a), torch.zeros_like(tail_a) if tail_a is not None else None)
                if prev[0].shape != video_a.shape:
                    prev = (_resized(prev[0], video_a.shape), prev[1])      # Two stages: the second run moves to the full size too
                    if prev[0] is None:
                        _say("Inactive | the latent changed size during the run, so the second run was dropped")
                        return out
                base_v, base_t, mix = prev[0], prev[1], 1.0
            else:
                base_v, base_t, mix = video_a.detach(), tail_a, fresh
            x_b = split_x[1](second_input(split_x[0], base_v, sigma, live["seed"], gen_step, mix))
            if tail_a is not None and with_sound:
                n = x_b.shape[-1] - tail_a.shape[-1]
                x_b = torch.cat([x_b[..., :n], second_input(x[..., n:], base_t.detach(), sigma, live["seed"] + 1,
                                                            gen_step, mix)], dim=-1)
            args_b, kwargs_b = streams.with_options(args, kwargs, {**to, dit_hooks.PROBE: True})
            began = time.perf_counter()
            out_b = executor(x_b, t, *args_b, **kwargs_b)
            took = time.perf_counter() - began
            split_b = streams.video_of(out_b, named)
            if split_b is None:
                return out
            video_b, rebuild_b = split_b
            tail_b = _sound(out_b, video_b)
            if parallel:
                live["prev"] = (video_b.detach(), None if tail_b is None else tail_b.detach())
                if gen_step not in wanted:
                    return out
            live["merged"] += 1
            diff = float((video_a - video_b).norm() / video_a.norm().clamp(min=1e-8))
            log.info("FunUI Step merging",
                     f"step {where[0] + 1}: second answer differs from the first by {diff:.1%}, "
                     f"merged at {strength:.2f}{' (picture and sound)' if tail_a is not None and with_sound else ''}; "
                     f"the second pass took {took:.1f}s")
            mixed = _mix(out, out_b, named, strength, b_first, with_sound)
            return out if mixed is None else mixed[0]
        except InterruptProcessingException:
            raise                                   # Cancel stops the run, it is not a failed merge
        except Exception as exc:                    # noqa: BLE001
            log.broke("FunUI Step merging", exc, "merging the second answer")
            return out

    patcher.add_wrapper_with_key(WrappersMP.APPLY_MODEL, key, apply_model)
    return (f"steps {values.get('steps')}, {'B into A' if b_first else 'A into B'}, "
            f"strength {strength:.2f}, second seed {'random' if chosen < 0 else chosen}")


PROVIDES = {"modifier": install}