"""The three things that make `Plaguekind/Minimax-H3` a *workflow* rather than just MiniMax-H3.

`Plaguekind/Minimax-H3` ships no weights: it is a ComfyUI graph
(`PlagueKind-MinimaxH3-V1.5.json`) over `Comfy-Org/MiniMax-H3`, and everything it contributes is in the
sampling and the post chain. Read off the graph, that is:

| ComfyUI node | widget | here |
|---|---|---|
| `KSamplerSelect` | `euler` | MiniMax-H3's only sampler; the checkpoint is CFG-distilled, one forward per step |
| `BasicScheduler` | `linear_quadratic`, 15 steps, denoise 1.0 | `linear_quadratic_sigmas` |
| `ImageSharpenKJ` | `rcas`, 0.3 | `rcas` |
| `FrameInterpolate` + `FrameInterpolationModelLoader` | `film_net_fp16.safetensors`, multiplier 2 | `interpolate` |
| `CreateVideo` | fps `24 * 2` | 48 fps out |
| `RTXVideoSuperResolution` | 2x, `ULTRA` | **not portable** — NVIDIA NGX, Windows/RTX driver only |

The sigma schedule is the one that changes the pixels most, and the one that is easy to get subtly wrong.
"""

from __future__ import annotations

import torch


# ----------------------------------------------------------------------------------------------------------------
# BasicScheduler(linear_quadratic)
# ----------------------------------------------------------------------------------------------------------------
# MiniMax-H3 carries two rectified-flow schedules per request, `shift = 12` for the video rows and `shift = 3` for
# the audio rows. diffusers builds both from one `linspace(1, 0, steps)` base grid; ComfyUI instead samples the
# *video* schedule and derives the audio one from it in closed form
# (`comfy/ldm/minimax/model.py::time_shift_sigma`). The two agree, because the shift is a bijection of the base
# grid — which is what lets a schedule chosen in ComfyUI's video-sigma space be transplanted here exactly.
#
# `linear_quadratic` is Mochi's schedule (`comfy/samplers.py::linear_quadratic_schedule`) and it does **not** go
# through the model's shift at all: it is `sigma_max = 1.0` scaled, so the grid PlagueKind's 15 steps actually run
# is this one verbatim, in the video stream, with the audio stream shifted off it.
VIDEO_SHIFT = 12.0
AUDIO_SHIFT = 3.0


def linear_quadratic_sigmas(
    steps: int, threshold_noise: float = 0.025, linear_steps: int | None = None
) -> torch.Tensor:
    """ComfyUI's `linear_quadratic` sigma grid, in MiniMax-H3's video-sigma space.

    Ported from `comfy/samplers.py::linear_quadratic_schedule` (itself from Mochi), with
    `model_sampling.sigma_max == 1.0`, which is what a rectified-flow model has. Returns `steps + 1` strictly
    decreasing sigmas from exactly 1.0 to exactly 0.0, so it drives `steps` forwards — ComfyUI's step count, not
    diffusers' (where the terminal zero is one of the `num_inference_steps`).

    Half the steps crawl through the first 2.5% of the trajectory and the rest sprint the remaining 97.5%: it is a
    front-loaded schedule, which is why 15 steps of it hold up against ~28 of the native grid.
    """
    steps = int(steps)
    if steps < 2:
        return torch.tensor([1.0, 0.0], dtype=torch.float32)
    if linear_steps is None:
        linear_steps = steps // 2

    linear = [i * threshold_noise / linear_steps for i in range(linear_steps)]
    threshold_noise_step_diff = linear_steps - threshold_noise * steps
    quadratic_steps = steps - linear_steps
    quadratic_coef = threshold_noise_step_diff / (linear_steps * quadratic_steps**2)
    linear_coef = threshold_noise / linear_steps - 2 * threshold_noise_step_diff / (quadratic_steps**2)
    const = quadratic_coef * (linear_steps**2)
    quadratic = [quadratic_coef * (i**2) + linear_coef * i + const for i in range(linear_steps, steps)]

    schedule = linear + quadratic + [1.0]
    return torch.tensor([1.0 - value for value in schedule], dtype=torch.float32)


def time_shift_sigma(sigma: torch.Tensor, from_shift: float, to_shift: float) -> torch.Tensor:
    """Move a sigma between two exponential shifts of the same base grid.

    `comfy/ldm/minimax/model.py::time_shift_sigma`: invert `sigma = s*b / (1 + (s-1)*b)` back to the base grid `b`,
    then re-apply the other shift. Monotonic, and it fixes both 0.0 and 1.0, so a strictly decreasing schedule that
    ends at zero stays one.
    """
    if from_shift == to_shift:
        return sigma
    base = sigma / (from_shift + sigma * (1.0 - from_shift))
    return to_shift * base / (1.0 + (to_shift - 1.0) * base)


class use_linear_quadratic:
    """Force MiniMax-H3's two schedulers onto the `linear_quadratic` grid for one pipeline call.

    A context manager rather than a pipeline-block subclass on purpose: `MiniMaxH3Scheduler.set_timesteps` already
    takes a fully-formed `sigmas=` schedule as public API, so nothing here reaches into the modular blocks, and the
    override lives and dies inside one request.
    """

    def __init__(self, pipe, steps: int, threshold_noise: float = 0.025, enabled: bool = True):
        self.schedulers = [pipe.scheduler, pipe.audio_scheduler] if enabled else []
        self.steps = int(steps)
        self.threshold_noise = float(threshold_noise)

    def __enter__(self):
        video_sigmas = linear_quadratic_sigmas(self.steps, self.threshold_noise)
        for scheduler in self.schedulers:
            sigmas = time_shift_sigma(video_sigmas, VIDEO_SHIFT, float(scheduler.shift))
            unbound = type(scheduler).set_timesteps

            def forced(num_inference_steps=None, device=None, sigmas=None, _s=scheduler, _grid=sigmas, _f=unbound):
                return _f(_s, None, device, _grid)

            scheduler.set_timesteps = forced
        return self

    def __exit__(self, *_):
        for scheduler in self.schedulers:
            scheduler.__dict__.pop("set_timesteps", None)
        return False


# ----------------------------------------------------------------------------------------------------------------
# ImageSharpenKJ(rcas, 0.3)
# ----------------------------------------------------------------------------------------------------------------
def rcas(video: torch.Tensor, strength: float, chunk: int = 16) -> torch.Tensor:
    """AMD FidelityFX **RCAS** — Robust Contrast Adaptive Sharpening — on `(frames, 3, H, W)` in `[0, 1]`.

    The FidelityFX kernel, which is what `ImageSharpenKJ`'s `rcas` mode is: a 5-tap cross, a sharpening lobe whose
    strength is limited per pixel so the ring it would create cannot leave `[0, 1]`, and a renormalised blend.

        lobe = clamp(attenuation * min over channels of max(-min / 4*max, -(1 - max) / 4*(1 - min)), -0.1875, 0)
        out  = (center + lobe * (n + s + e + w)) / (1 + 4 * lobe)

    `lobe` is negative, so the neighbours are subtracted: a high-pass with a headroom-aware gain, which is why it
    sharpens MiniMax-H3's slightly soft VAE output without haloing it. PlagueKind's 0.3 is the strength; the note in
    the workflow calls it "very natural" and that matches — the lobe clamp caps it well below a visible ring.

    Batched over `chunk` frames at a time rather than ComfyUI's one, and written back in place: the clip is already
    resident on the card, but this runs immediately after the denoise loop's allocation peak, and a whole-clip pass at
    the full 1344x768x124 would ask the allocator for ~8 GB of intermediates at exactly the wrong moment.
    """
    if strength <= 0:
        return video

    frames, _, height, width = video.shape
    strength = float(strength)
    for start in range(0, frames, chunk):
        center = video[start : start + chunk]
        padded = torch.nn.functional.pad(center, (1, 1, 1, 1), mode="reflect")
        north = padded[:, :, 0:height, 1 : width + 1]
        south = padded[:, :, 2 : height + 2, 1 : width + 1]
        west = padded[:, :, 1 : height + 1, 0:width]
        east = padded[:, :, 1 : height + 1, 2 : width + 2]

        low = torch.minimum(torch.minimum(torch.minimum(torch.minimum(north, south), west), east), center)
        high = torch.maximum(torch.maximum(torch.maximum(torch.maximum(north, south), west), east), center)

        hit_min = -low / (high * 4.0 + 1e-6)
        hit_max = -(1.0 - high) / ((1.0 - low) * 4.0 + 1e-6)
        lobe = torch.maximum(hit_min, hit_max).amin(dim=1, keepdim=True)
        lobe = (lobe * strength).clamp_(-0.1875, 0.0)
        del low, high, hit_min, hit_max

        neighbours = north + south + east + west
        center.copy_(((center + lobe * neighbours) / (1.0 + 4.0 * lobe)).clamp_(0.0, 1.0))
    return video


# ----------------------------------------------------------------------------------------------------------------
# FrameInterpolate(film_net_fp16, multiplier=2)
# ----------------------------------------------------------------------------------------------------------------
FILM_REPO = "Comfy-Org/frame_interpolation"
FILM_FILE = "frame_interpolation/film_net_fp16.safetensors"


def load_film():
    """FILM, off the same checkpoint the workflow names. CPU work; `None` on any failure, and the caller skips."""
    from huggingface_hub import hf_hub_download
    from safetensors.torch import load_file

    from film_net import FILMNet

    path = hf_hub_download(FILM_REPO, FILM_FILE)
    model = FILMNet()
    model.load_state_dict(load_file(path))
    return model.eval().to(torch.float16)


@torch.no_grad()
def interpolate(model, video: torch.Tensor, multiplier: int = 2) -> torch.Tensor:
    """`multiplier`x frame interpolation of `(frames, 3, H, W)` in `[0, 1]`, FILM, on the card.

    Mirrors ComfyUI's `FrameInterpolate`: one pass per adjacent pair, the flow computed once per pair and reused for
    every intermediate timestep (`forward_multi_timestep`), and the feature pyramid of frame `i + 1` carried over as
    frame `i` of the next pair — which halves the feature extractions. Output length is
    `(frames - 1) * multiplier + 1`, i.e. 24 fps in, `24 * multiplier` fps out.
    """
    frames = video.shape[0]
    if model is None or frames < 2 or multiplier < 2:
        return video

    dtype = torch.float16
    timesteps = [t / multiplier for t in range(1, multiplier)]
    # float16, not the input's float32: the buffer is the largest allocation of the whole post chain (a 2x pass over
    # 124 frames at 1344x768 is 247 of them) and it happens right after the denoise loop's peak.
    out = torch.empty(((frames - 1) * multiplier + 1, *video.shape[1:]), dtype=dtype, device=video.device)
    out[0] = video[0]
    cursor = 1

    cache: dict = {}
    for index in range(frames - 1):
        first = video[index : index + 1].to(dtype)
        second = video[index + 1 : index + 2].to(dtype)
        cache["img0"] = cache.pop("next") if "next" in cache else model.extract_features(first)
        cache["img1"] = model.extract_features(second)
        cache["next"] = cache["img1"]

        middles = model.forward_multi_timestep(first, second, timesteps, cache=cache)
        out[cursor : cursor + len(timesteps)] = middles.to(video.dtype).clamp_(0.0, 1.0)
        cursor += len(timesteps)
        out[cursor] = video[index + 1]
        cursor += 1

    return out