"""MOSS-Image DiT for distillation.

One frozen bf16 base (the release MOSSImageTransformer) is shared by the teacher, the
data-space generator, the fake score and the discriminator backbone; each path only
differs by its LoRA, its time path and optional extras. The residual stream and the
modulation run in fp32 (the release trains fp32-master + autocast), matmuls in bf16.
"""
from __future__ import annotations

import importlib.util
import math
from dataclasses import dataclass
from pathlib import Path

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

BF16 = torch.bfloat16
PROJ = ('q', 'k', 'v', 'o', 'ffn_up_gate', 'ffn_down')


def load_release_module(model_dir):
    """Import the release pipeline_moss_image.py (read and reviewed; no side effects on import)."""
    path = Path(model_dir) / 'pipeline_moss_image.py'
    spec = importlib.util.spec_from_file_location('pipeline_moss_image', path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module


def load_base(model_dir, device):
    release = load_release_module(model_dir)
    tr = release.MOSSImageTransformer.from_pretrained(Path(model_dir) / 'transformer', torch_dtype=BF16,
                                                      local_files_only=True)
    return Moss(tr.to(device), release)


def rms(x, weight, eps=1e-6):
    y = x.float() * torch.rsqrt(x.float().square().mean(-1, keepdim=True) + eps)
    return y.to(x.dtype) * weight


def situ(hidden, gate):
    return (4.0 * torch.tanh(gate / 4.0) * torch.sigmoid(gate)) * (25.0 * torch.tanh(hidden / 25.0))


def proj_dims(dim, ffn_dim):
    return {'q': (dim, dim), 'k': (dim, dim), 'v': (dim, dim), 'o': (dim, dim),
            'ffn_up_gate': (dim, 2 * ffn_dim), 'ffn_down': (ffn_dim, dim)}


class LoRA(nn.Module):
    """Stacked LoRA for the projections of blocks [start, layers): A [n, r, in], B [n, out, r]; y += (alpha/r) B A x."""

    def __init__(self, dims, layers, rank, alpha, start=0):
        super().__init__()
        self.rank, self.alpha, self.scale, self.start = rank, alpha, alpha / rank, start
        layers -= start
        self.A = nn.ParameterDict()
        self.B = nn.ParameterDict()
        for name, (fan_in, fan_out) in dims.items():
            bound = 1 / math.sqrt(fan_in)
            self.A[name] = nn.Parameter(torch.empty(layers, rank, fan_in).uniform_(-bound, bound))
            self.B[name] = nn.Parameter(torch.zeros(layers, fan_out, rank))
        self._cache = None

    def delta(self, x, layer, name):
        layer -= self.start
        if self._cache is not None and not torch.is_grad_enabled():
            a, b = self._cache[name]
            a, b = a[layer], b[layer]
        else:
            a, b = self.A[name][layer].to(x.dtype), self.B[name][layer].to(x.dtype)
        return F.linear(F.linear(x, a), b) * self.scale

    def freeze_cache(self):
        """bf16 copies used by no-grad forwards only; drop_cache() before the next optimizer step."""
        with torch.no_grad():
            self._cache = {n: (self.A[n].to(BF16), self.B[n].to(BF16)) for n in self.A}

    def drop_cache(self):
        self._cache = None


class TimePath(nn.Module):
    """time_embedding MLP + both modulation biases (fp32); modulation weights stay the base's."""

    NAMES = ('time_embedding.0.weight', 'time_embedding.0.bias', 'time_embedding.2.weight',
             'time_embedding.2.bias', 'shared_block_modulation.bias', 'video_head_modulation.bias')

    def __init__(self, tensors, trainable):
        super().__init__()
        self.p = nn.ParameterDict({k.replace('.', '_'): nn.Parameter(tensors[k].detach().float().clone(),
                                                                     requires_grad=trainable)
                                   for k in self.NAMES})

    def __getitem__(self, name):
        return self.p[name.replace('.', '_')]

    def state(self):
        return {k: self[k].detach() for k in self.NAMES}


@dataclass
class Context:
    tokens: torch.Tensor  # [1, L, D] text_embedder output (bf16)
    length: int


class Moss(nn.Module):
    def __init__(self, tr, release):
        super().__init__()
        self.tr = tr.requires_grad_(False).eval()
        self.release = release
        cfg = tr.config
        self.dim, self.heads, self.depth = cfg.dim, cfg.num_heads, cfg.num_layers
        self.dims = proj_dims(cfg.dim, cfg.ffn_dim)
        sd = dict(tr.named_parameters())
        self.base_time = TimePath(sd, trainable=False)
        self.register_buffer('mod_w', tr.shared_block_modulation.weight.detach().float(), persistent=False)
        self.register_buffer('head_w', tr.video_head_modulation.weight.detach().float(), persistent=False)
        self._rope = {}

    # ---- conditioning ----
    def rope(self, h, w, length):
        key = (h, w, length)
        if key not in self._rope:
            if len(self._rope) > 64:
                self._rope.clear()
            with torch.no_grad():
                self._rope[key] = self.tr.rope(h, w, length, self.mod_w.device).to(BF16)
        return self._rope[key]

    @torch.no_grad()
    def context(self, emb):
        """emb [L, text_dim] text-encoder hidden states -> Context (text refinement is prompt-only)."""
        length = emb.shape[0]
        rope = self.tr.rope(1, 1, length, self.mod_w.device).to(BF16)[:, 1:]
        tokens = self.tr.text_embedder(emb[None].to(device=self.mod_w.device, dtype=BF16), rope)
        return Context(tokens, length)

    def time_mod(self, t, tp, extra=None):
        half = self.tr.config.freq_dim // 2
        t = torch.as_tensor(t, device=self.mod_w.device, dtype=torch.float32).reshape(-1)
        freq = 10000 ** (-torch.arange(half, device=t.device, dtype=torch.float64) / half)
        phase = t.double()[:, None] * freq
        emb = torch.cat((phase.cos(), phase.sin()), -1).float()
        time = F.linear(emb, tp['time_embedding.0.weight'], tp['time_embedding.0.bias'])
        time = F.linear(F.silu(time), tp['time_embedding.2.weight'], tp['time_embedding.2.bias'])
        if extra is not None:
            time = time + extra
        mod = F.linear(time, self.mod_w, tp['shared_block_modulation.bias']).view(-1, 6, self.dim)
        head = F.linear(time, self.head_w, tp['video_head_modulation.bias']).view(-1, 2, self.dim)
        return mod, head

    # ---- blocks ----
    def lin(self, y, i, name, lora):
        out = F.linear(y, self.tr.layers[i].proj[name].weight)
        if lora is not None and i >= lora.start:
            out = out + lora.delta(y, i, name)
        return out

    def block(self, i, x, rope, mod, lora):
        layer = self.tr.layers[i]
        heads = self.heads
        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.unbind(1)
        y = (rms(x, layer.norm1.weight) * (1 + scale_a[:, None]) + shift_a[:, None]).to(BF16)
        q = rms(self.lin(y, i, 'q', lora).unflatten(-1, (heads, -1)), layer.q_norm.weight).flatten(2)
        k = rms(self.lin(y, i, 'k', lora).unflatten(-1, (heads, -1)), layer.k_norm.weight).flatten(2)
        a = self.release._attention(q, k, self.lin(y, i, 'v', lora), rope, heads)
        x = x + self.lin(a, i, 'o', lora).float() * gate_a[:, None]
        y = (rms(x, layer.norm2.weight) * (1 + scale_f[:, None]) + shift_f[:, None]).to(BF16)
        hidden, gate = self.lin(y, i, 'ffn_up_gate', lora).chunk(2, -1)
        return x + self.lin(situ(hidden, gate), i, 'ffn_down', lora).float() * gate_f[:, None]

    def forward(self, latent, t, ctx, *, lora=None, tp=None, time_extra=None, skip=(), taps=None, ckpt=False):
        """latent [1, C, h, w] fp32, t on the 0..1000 scale. Returns velocity [1, C, h, w] fp32,
        or {block: image-token features [1, h*w, D]} when taps is given (stops after the last tap)."""
        _, _, h, w = latent.shape
        hw = h * w
        mod, head = self.time_mod(t, tp if tp is not None else self.base_time, time_extra)
        rope = self.rope(h, w, ctx.length)
        video = self.tr.video_embedder(latent.flatten(2).transpose(1, 2).to(BF16))
        x = torch.cat((video, ctx.tokens), 1).float()
        last = max(taps) if taps else self.depth - 1
        use_ckpt = ckpt and torch.is_grad_enabled()
        feats = {}
        for i in range(last + 1):
            if i not in skip:
                if use_ckpt:
                    x = checkpoint(self.block, i, x, rope, mod, lora, use_reentrant=False)
                else:
                    x = self.block(i, x, rope, mod, lora)
            if taps and i in taps:
                feats[i] = x[:, :hw]
        if taps:
            return feats
        x = rms(x[:, :hw], self.tr.final_norm_video.weight)
        x = x * (1 + head[:, 1, None]) + head[:, 0, None]
        x = F.linear(x.to(BF16), self.tr.final_linear_video.weight).float()
        return x.transpose(1, 2).reshape(1, -1, h, w)

    @torch.no_grad()
    def teacher(self, latent, t, pos, neg, cfg, skip_neg=(22,)):
        """Release guidance: CFG with block 22 skipped in the negative branch only."""
        vp = self(latent, t, pos)
        vn = self(latent, t, neg, skip=skip_neg)
        return vn + cfg * (vp - vn)


def release_sigmas(steps, shift=5.0, total=1000, training_shift=1.0):
    """MOSSImageScheduler.set_timesteps sigmas (shifted grid snapped to the training grid)."""
    def grid(count, s):
        lower = s / (total + s - 1)
        values = torch.linspace(1, lower, count)
        return s * values / (1 + (s - 1) * values)
    sigmas = grid(steps, shift)
    training = grid(total, training_shift).flip(0)
    upper = torch.searchsorted(training, sigmas).clamp(max=total - 1)
    lower = (upper - 1).clamp(min=0)
    nearest = torch.where((sigmas - training[lower]).abs() <= (training[upper] - sigmas).abs(), lower, upper)
    return torch.unique_consecutive(training[nearest]).tolist() + [0.0]


@torch.no_grad()
def sample_teacher(model, noise, pos, neg, steps=50, cfg=3.0, shift=5.0):
    x = noise.float()
    sigmas = release_sigmas(steps, shift)
    for s, s_next in zip(sigmas[:-1], sigmas[1:]):
        v = model.teacher(x, s * 1000, pos, neg, cfg)
        x = x + (s_next - s) * v
    return x