# Copyright 2026 Viggle AI. Licensed under the Apache License, Version 2.0 (see LICENSE-CODE).
# SPDX-License-Identifier: Apache-2.0
"""The MiniMax-H3 side: VAE encode, the packed `ref2va` sequence, and the canvas ladders.

The render is a *second video reference*: `<Video 1>` is the source clip, `<Video 2>` the point-cloud
render, both at the 480 class, packed as condition rows by the pipeline's own
`build_ref2va_packed_sequence` after the frozen text rows, then the noised target at the 768 class.
Tokens, not channels -- nothing about the transformer is touched, so both the re-camera teacher and the
DMD student are ordinary PEFT LoRAs over a plain, unmodified `transformer/` directory.
"""

import math

import torch
from diffusers import MiniMaxH3Scheduler
from diffusers.modular_pipelines.minimax_h3.before_denoise import (
    MiniMaxH3Ref2VAPrepareLayoutStep, MiniMaxH3SetTimestepsStep, patchify_video_latents,
)

# Canvases, one per aspect bucket; entry `i` is the same aspect in every class (to within 1.3 %), which is what
# lets `<Video 1>`/`<Video 2>` sit at the 480 class under a 768-class target: the rotary grid is normalised by
# sqrt(area), so the condition block spans the same rotary rectangle as the target it conditions.
LADDERS = {
    480: [(416, 960), (448, 896), (480, 832), (544, 736), (640, 640), (736, 544), (832, 480), (896, 448), (960, 416)],
    768: [(672, 1536), (704, 1408), (768, 1344), (864, 1184), (1024, 1024), (1184, 864), (1344, 768), (1408, 704), (1536, 672)],
}
TARGET_CLASS, COND_CLASS = 768, 480

PATCH_SIZE = (1, 2, 2)
AUDIO_CHANNELS = 2
VIDEO_TAG, TEXT_TAG, AUDIO_TAG = 0, 1, 2
KEYFRAME_NOISE_AUG = 0.999  # the pipeline's noise level on condition rows
VIDEO_SHIFT = 12.0  # scheduler/scheduler_config.json of MiniMax-H3
AUDIO_SHIFT = 3.0  # audio_scheduler/scheduler_config.json
PIXEL_MEAN = (0.485, 0.456, 0.406)
PIXEL_STD = (0.229, 0.224, 0.225)
ENCODE_SEED = 42
FPS = 24.0  # the frame rate the frozen prompt declares; the model has no other notion of time


def bucket(w, h):
    """`(target canvas, condition canvas)` for a `w x h` source: the ladder entry nearest in log-aspect."""
    i = min(range(len(LADDERS[TARGET_CLASS])),
            key=lambda i: abs(math.log(LADDERS[TARGET_CLASS][i][0] / LADDERS[TARGET_CLASS][i][1] * h / w)))
    return LADDERS[TARGET_CLASS][i], LADDERS[COND_CLASS][i]


class _Reference:
    """All `build_ref2va_packed_sequence` reads off a reference is its modality."""

    def __init__(self, kind):
        self.kind, self.has_audio = kind, False


def encode_video(vae, video):
    """[F,H,W,3] uint8 -> `(anchor, target)` normalised latents `[1,C,T,h,w]` from one `vae.encode`: the
    anchor is the pipeline's seed-42 fp16-rounded draw (what a condition block gets), the target a fresh draw."""
    latents_mean = torch.tensor(vae.config.latents_mean).view(1, -1, 1, 1, 1)
    latents_std = torch.tensor(vae.config.latents_std).view(1, -1, 1, 1, 1)
    mean = torch.tensor(PIXEL_MEAN, device=video.device).view(1, -1, 1, 1, 1)
    std = torch.tensor(PIXEL_STD, device=video.device).view(1, -1, 1, 1, 1)
    pixels = video.permute(3, 0, 1, 2)[None].to(torch.float32).div(255.0)
    posterior = vae.encode((pixels - mean) / std, return_dict=False)[0]
    anchor = posterior.sample(generator=torch.Generator().manual_seed(ENCODE_SEED)).to(torch.float16).float().cpu()
    return (anchor - latents_mean) / latents_std, (posterior.sample().float().cpu() - latents_mean) / latents_std


def decode_video(vae, rows, shape):
    """Packed target rows -> [F,H,W,3] uint8. `shape` is the target latent's `(C, T, h, w)`."""
    C, T, H, W = shape
    pt, ph, pw = PATCH_SIZE
    dev = rows.device
    x = rows.reshape(1, T // pt, H // ph, W // pw, C, pt, ph, pw).permute(0, 4, 1, 5, 2, 6, 3, 7).reshape(1, C, T, H, W)
    x = x * torch.tensor(vae.config.latents_std, device=dev).view(1, -1, 1, 1, 1) \
        + torch.tensor(vae.config.latents_mean, device=dev).view(1, -1, 1, 1, 1)
    with torch.autocast("cuda", dtype=torch.float16):
        video = vae.decode(x, return_dict=False)[0]
    video = video.float() * torch.tensor(PIXEL_STD, device=dev).view(1, -1, 1, 1, 1) \
        + torch.tensor(PIXEL_MEAN, device=dev).view(1, -1, 1, 1, 1)
    return (video.clamp(0, 1)[0].permute(1, 2, 3, 0) * 255).round().to(torch.uint8).cpu()


def pack(cond, render, target, prompt_embeds, text_token_tags, audio_x0):
    """`[text | <Video 1> = cond | <Video 2> = render | target]` as the transformer's packed inputs.

    The layout, rotary geometry and modality tags come from the pipeline's own builder. Condition rows carry
    the pipeline's `KEYFRAME_NOISE_AUG` noise augmentation; the target rows (and the audio rows) start as
    pure noise. Returns the dict of tensors the denoise loop needs plus `num_condition_video_rows` /
    `num_condition_audio_rows`.
    """
    refs = [cond.float(), render.float()]
    position_ids, token_tags, video_indices, audio_indices, text_indices, n_cond_video, n_cond_audio = \
        MiniMaxH3Ref2VAPrepareLayoutStep.build_ref2va_packed_sequence(
            text_token_tags, [_Reference("video"), _Reference("video")], refs, [],
            target.shape[2], target.shape[3], target.shape[4],
            audio_x0.shape[0] // AUDIO_CHANNELS,
            PATCH_SIZE, AUDIO_CHANNELS, AUDIO_TAG, VIDEO_TAG,
        )
    cond_rows = torch.cat([
        patchify_video_latents(KEYFRAME_NOISE_AUG * x + (1.0 - KEYFRAME_NOISE_AUG) * torch.randn_like(x), PATCH_SIZE)
        for x in refs
    ])
    target_rows = torch.randn_like(patchify_video_latents(target.float(), PATCH_SIZE))
    return {
        "hidden_states": torch.cat([cond_rows, target_rows]),
        "audio_hidden_states": torch.randn_like(audio_x0),
        "encoder_hidden_states": prompt_embeds,
        "token_tags": token_tags,
        "position_ids": position_ids,
        "video_indices": video_indices,
        "audio_indices": audio_indices,
        "text_indices": text_indices,
        "num_condition_video_rows": n_cond_video,
        "num_condition_audio_rows": n_cond_audio,
    }


def denoise(transformer, batch, steps, flow_shift, device, on_step=None):
    """Walk the scheduler over the target rows in place. `--steps 4 --flow-shift 3` is the DMD student's grid
    (3 forwards); the teacher wants `--steps 50 --flow-shift 12`."""
    batch = {k: v.to(device) if torch.is_tensor(v) else v for k, v in batch.items()}
    n_cond, n_cond_audio = batch["num_condition_video_rows"], batch["num_condition_audio_rows"]
    latents, audio = batch["hidden_states"], batch["audio_hidden_states"]
    sch, asch = MiniMaxH3Scheduler(shift=flow_shift), MiniMaxH3Scheduler(shift=AUDIO_SHIFT)
    sch.set_timesteps(steps, device)
    asch.set_timesteps(steps, device)
    layout = {k: batch[k] for k in ("token_tags", "position_ids", "video_indices", "audio_indices", "text_indices")}
    prompt = batch["encoder_hidden_states"].to(torch.bfloat16)[None]
    for i, (t, at) in enumerate(zip(sch.timesteps, asch.timesteps)):
        if on_step:
            on_step(i, len(sch.timesteps))
        timestep, ti = MiniMaxH3SetTimestepsStep.build_row_timesteps(
            batch["video_indices"], batch["audio_indices"], n_cond, n_cond_audio, batch["text_indices"].numel(),
            float(t), float(at), max(float(t), KEYFRAME_NOISE_AUG), 1.0)
        v, av = transformer(hidden_states=latents[None], audio_hidden_states=audio[None], encoder_hidden_states=prompt,
                            timestep=timestep.to(device), timestep_indices=ti.to(device), **layout, return_dict=False)
        latents[n_cond:] = sch.step(v[0, n_cond:].float(), t, latents[n_cond:], return_dict=False)[0]
        audio = asch.step(av[0].float(), at, audio, return_dict=False)[0]
    return latents[n_cond:]