"""Runtime adaptations of the rtcwm streaming sampler for ZeroGPU (RTX PRO 6000, sm_120). See app.py."""
import time

import torch
import torch.nn.functional as F
from torch import nn

from musubi_tuner.minimax_h3 import model as h3model
from musubi_tuner.minimax_h3.packing import _build_modulation_segments
from rtcwm.model import attention as rt_attention
from rtcwm.model import packing as rt_packing
from rtcwm.model import sampler as rt_sampler

STEPS, SHIFT = 4, 5.0

class FP8Linear(nn.Module):
    """rtcwm.model.fp8.FP8Linear (per-output-channel e4m3 weights, per-row e4m3 activations, torch._scaled_mm),
    with an emulated fallback should the rowwise-scaled fp8 GEMM be unavailable on this GPU."""
    use_scaled_mm = True

    def __init__(self, weight_bf16: torch.Tensor):
        super().__init__()
        w = weight_bf16.float()
        scale = (w.abs().amax(dim=1, keepdim=True) / 448.0).clamp_min(1e-12)
        self.register_buffer("weight_fp8", (w / scale).to(torch.float8_e4m3fn))
        self.register_buffer("weight_scale", scale.t().contiguous())
        self.out_features = w.shape[0]

    def forward(self, x):
        shape = x.shape
        x2 = x.reshape(-1, shape[-1])
        scale = (x2.abs().amax(dim=1, keepdim=True).float() / 448.0).clamp_min(1e-12)
        xq = (x2.float() / scale).to(torch.float8_e4m3fn)
        if FP8Linear.use_scaled_mm:
            try:
                y = torch._scaled_mm(xq, self.weight_fp8.t(), scale_a=scale, scale_b=self.weight_scale,
                                     out_dtype=torch.bfloat16)
                return y.reshape(*shape[:-1], self.out_features)
            except Exception as e:  # noqa: BLE001
                print(f"_scaled_mm unavailable ({e!r}); using the emulated FP8 GEMM", flush=True)
                FP8Linear.use_scaled_mm = False
        y = (xq.to(torch.bfloat16) @ self.weight_fp8.to(torch.bfloat16).t()).float() * scale * self.weight_scale
        return y.to(torch.bfloat16).reshape(*shape[:-1], self.out_features)


_ADALN_T = [None]          # the unique times of the forward in flight (set by run_blocks)


def _tkey(t: float) -> int:
    return int(round(float(t) * 1_000_000))


class CachedAdaln(nn.Module):
    """A DiT block's AdalnProj evaluated once per known time: same output as
    linear(silu(time_embedder(t).bf16)), looked up by t."""

    def __init__(self, table: torch.Tensor, keys: list, modalities: int, expand: int, hidden: int):
        super().__init__()
        self.register_buffer("table", table)                     # [n_t, expand*hidden*modalities] bf16
        self.index = {k: i for i, k in enumerate(keys)}
        self.modalities, self.expand, self.hidden = modalities, expand, hidden

    def forward(self, _temb):
        rows = [self.index[_tkey(t)] for t in _ADALN_T[0]]
        p = self.table[rows]
        p = p.reshape(p.shape[0] * self.modalities, self.expand * self.hidden)
        return p.chunk(self.expand, dim=-1)


def sampling_times() -> list:
    """Every model time a 4-step x0 stream evaluates: condition rows (1.0 text / audio / clean history, 0.999
    Picture 1 + proxy) and 1 - sigma of each step, computed like the sampler does (float32 rows)."""
    sig = rt_sampler.schedule(STEPS, "cpu", SHIFT)
    ts = [1.0, 0.999] + [1.0 - float(sig[i]) for i in range(STEPS)]
    return sorted({float(torch.tensor(t, dtype=torch.float32)) for t in ts})


def prepare_transformer(tr, lora: dict, fp8: bool = True, cache_adaln: bool = True):
    """Merge the streaming LoRA into the blocks' linears (musubi merge_to, multiplier 1, bf16), put them in FP8
    (rtcwm.model.fp8), precompute the blocks' AdaLN outputs for the sampling times, and install the rtcwm
    attention (with the SDPA stand-ins)."""
    t0 = time.time()
    tr.requires_grad_(False)
    used = set()
    times = sampling_times()
    with torch.no_grad():
        temb = tr.time_embedder(torch.tensor(times, dtype=torch.float32)).to(torch.bfloat16)
        act = F.silu(temb).float()
        keys = [_tkey(t) for t in times]
        for i, blk in enumerate(tr.blocks):
            # LoRA merge (musubi LoRAModule.merge_to, multiplier 1, bf16) + FP8 (rtcwm.model.fp8.quantize_blocks)
            for parent, prefix, names in ((blk.attn, "attn", ("qkv_proj", "out_proj")), (blk.mlp, "mlp", ("fc1", "fc2"))):
                for n in names:
                    lin = getattr(parent, n)
                    assert lin.bias is None
                    key = f"lora_unet_blocks_{i}_{prefix}_{n}"
                    down = lora[f"{key}.lora_down.weight"].float()
                    up = lora[f"{key}.lora_up.weight"].float()
                    scale = float(lora[f"{key}.alpha"].float()) / down.shape[0]
                    w = (lin.weight.float() + (up @ down) * scale).to(torch.bfloat16)
                    if fp8:
                        setattr(parent, n, FP8Linear(w))
                    else:
                        lin.weight.data = w
                    used.update({f"{key}.lora_down.weight", f"{key}.lora_up.weight", f"{key}.alpha"})
                    del w, lin
            if not cache_adaln:
                continue
            ad = blk.adaln_proj
            table = (act @ ad.linear.weight.float().t() + ad.linear.bias.float()).to(torch.bfloat16)
            blk.adaln_proj = CachedAdaln(table, keys, ad.modalities, ad.expand, ad.hidden_size)
            del ad
            if i % 10 == 0:
                print(f"block {i} prepared ({time.time() - t0:.0f}s)", flush=True)
    assert used == set(lora), f"unused LoRA keys: {sorted(set(lora) - used)[:5]}"
    print(f"merged {len(used) // 3} LoRA modules, FP8 blocks ready in {time.time() - t0:.0f}s", flush=True)
    install(tr)
    return tr.eval()


def install(tr):
    rt_attention.install(tr)
    h3model.Attention.forward = _attention_forward


# --------------------------------------------------------------------------------------------- attention
def _sdpa(q, k, v):
    """FlashAttention-3 stand-in: q, k, v [B,S,H,D] -> [B,S,H,D], unmasked."""
    return F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)).transpose(1, 2)


rt_attention._fa3 = _sdpa


class DenseMask:
    """rtcwm.model.masks.df_mask as prefix attention calls (rtcwm.model.attention.dense_df) -- for the condition
    rows alone (n_chunks 0) or together with chunk 0 (n_chunks 1)."""

    def __init__(self, head_rows, duv_rows, audio_rows, n_chunks, static_sees_first):
        self.head_rows, self.duv_rows, self.audio_rows = head_rows, duv_rows, audio_rows
        self.n_chunks, self.static_sees_first = n_chunks, static_sees_first


def _df_mask(head_rows, duv_rows, audio_rows, n_chunks, device, first_chunk=0, static_sees_first=True):
    assert first_chunk == 0
    return DenseMask(head_rows, duv_rows, audio_rows, n_chunks, static_sees_first)


rt_sampler.df_mask = _df_mask
_rtcwm_forward = rt_attention._forward


def _attention_forward(self, hidden_states, rotation_table=None):
    mask, kv = rt_attention._state["mask"], rt_attention._state["kv"]
    if not isinstance(mask, DenseMask):
        return _rtcwm_forward(self, hidden_states, rotation_table)
    b, s, _ = hidden_states.shape
    q, k, v = self.qkv_proj(hidden_states).split(self.inner_dim, dim=-1)
    q = self.q_norm(q.reshape(b, s, self.num_heads, self.head_dim))
    k = self.k_norm(k.reshape(b, s, self.num_heads, self.head_dim))
    v = v.reshape(b, s, self.num_heads, self.head_dim)
    if rotation_table is not None:
        q = rt_attention.rope(q, rotation_table)
        k = rt_attention.rope(k, rotation_table)
    q, k, v = q.to(torch.bfloat16), k.to(torch.bfloat16), v.to(torch.bfloat16)
    if kv is not None and kv.get("store") is not None:
        kv["store"].append((k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous()))
    out = rt_attention.dense_df(q, k, v, mask)
    return self.out_proj(out.reshape(b, s, self.inner_dim))


def run_blocks(model, hidden, t_rows, tags, positions, *, mask=None, kv=None, train=False, offload=False):
    """rtcwm.model.packing.run_blocks (inference), publishing the forward's times to CachedAdaln."""
    uniq, inv = torch.unique(t_rows, sorted=True, return_inverse=True)
    segments = (_build_modulation_segments(3 * inv + tags),)
    temb = model.time_embedder(uniq.to(hidden.device)).to(model.dtype)
    rot = model._rotation_table(positions.unsqueeze(0).to(hidden.device), hidden.dtype)
    _ADALN_T[0] = uniq.tolist()
    for block in model.blocks:
        with rt_attention.active(mask, kv):
            hidden = block(hidden, temb, segments, rot)
    return hidden, temb, inv


rt_packing.run_blocks = run_blocks
rt_sampler.run_blocks = run_blocks