"""Tag-conditioned DeCo transformer over SAME-L music latents (1.19B params)."""

from __future__ import annotations

import math
from pathlib import Path

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

WEIGHTS = Path(__file__).parent / "audio_dit.safetensors"

CHANNELS = 256
FRAMES = 1024
DIM = 1536
HEADS = 16
HEAD_DIM = DIM // HEADS
MLP_HIDDEN = 4 * DIM
DECODER_DIM = 2048
NUM_TAGS = 1470
MAX_TAGS = 8
ROPE_THETA = 10000.0

ENCODER_BLOCKS = 2
MIDDLE_BLOCKS = 22
DECODER_BLOCKS = 2
OUTPUT_BLOCKS = 3


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        dtype = x.dtype
        y = x.float()
        y = y * torch.rsqrt(y.square().mean(dim=-1, keepdim=True) + self.eps)
        return (self.weight * y).to(dtype)


def rope_tables() -> tuple[torch.Tensor, torch.Tensor]:
    inv_freq = 1.0 / (
        ROPE_THETA ** (torch.arange(0, HEAD_DIM, 2, dtype=torch.float32) / HEAD_DIM)
    )
    freqs = torch.outer(torch.arange(FRAMES, dtype=torch.float32), inv_freq)
    return freqs.cos()[None, None], freqs.sin()[None, None]


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    cos = cos[..., : x.shape[-2], :].to(device=x.device, dtype=x.dtype)
    sin = sin[..., : x.shape[-2], :].to(device=x.device, dtype=x.dtype)
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class FeedForward(nn.Module):
    def __init__(self, dim: int = DIM, hidden: int = MLP_HIDDEN):
        super().__init__()
        self.w12 = nn.Linear(dim, 2 * hidden, bias=False)
        self.w3 = nn.Linear(hidden, dim, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        gate, value = self.w12(x).chunk(2, dim=-1)
        return self.w3(F.silu(gate) * value)


class Attention(nn.Module):
    """Self-attention with optional tag keys and values appended to the sequence."""

    def __init__(self, use_tags: bool):
        super().__init__()
        self.use_tags = use_tags
        self.qkv_x = nn.Linear(DIM, 3 * DIM, bias=False)
        if use_tags:
            self.kv_y = nn.Linear(DIM, 2 * DIM, bias=False)
        self.q_norm = RMSNorm(HEAD_DIM)
        self.k_norm = RMSNorm(HEAD_DIM)
        self.proj = nn.Linear(DIM, DIM)

    def forward(
        self,
        x: torch.Tensor,
        tags: torch.Tensor,
        tag_padding: torch.Tensor,
        cos: torch.Tensor,
        sin: torch.Tensor,
    ) -> torch.Tensor:
        batch, frames, _ = x.shape
        q, k, v = self.qkv_x(x).view(
            batch, frames, 3, HEADS, HEAD_DIM
        ).permute(2, 0, 3, 1, 4)
        q = apply_rope(self.q_norm(q), cos, sin)
        k = apply_rope(self.k_norm(k), cos, sin)

        attn_mask = None
        if self.use_tags:
            kt, vt = self.kv_y(tags).view(
                batch, MAX_TAGS, 2, HEADS, HEAD_DIM
            ).permute(2, 0, 3, 1, 4)
            k = torch.cat([k, self.k_norm(kt)], dim=2)
            v = torch.cat([v, vt], dim=2)
            attn_mask = torch.zeros(
                batch, 1, 1, frames + MAX_TAGS, device=x.device, dtype=x.dtype
            )
            attn_mask[..., frames:].masked_fill_(
                tag_padding[:, None, None, :], float("-inf")
            )

        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)
        return self.proj(out.transpose(1, 2).reshape(batch, frames, DIM))


class Block(nn.Module):
    def __init__(self, use_tags: bool, local_context: bool = False):
        super().__init__()
        self.norm1 = RMSNorm(DIM)
        self.attn = Attention(use_tags)
        self.norm2 = RMSNorm(DIM)
        self.mlp = FeedForward()
        if local_context:
            self.local_scale = nn.Parameter(torch.empty(()))
            self.local_conv = nn.Conv2d(DIM, DIM, 3, padding=1, groups=DIM)
        else:
            self.local_scale = None
            self.local_conv = None

    def forward(self, x, tags, tag_padding, mod, cos, sin):
        if self.local_conv is not None:
            local = self.local_conv(x.transpose(1, 2).unsqueeze(2))
            local = local.squeeze(2).transpose(1, 2)
            x = x + self.local_scale * local
        shift1, scale1, gate1, shift2, scale2, gate2 = mod.chunk(6, dim=-1)
        x = x + gate1 * self.attn(
            self.norm1(x) * (1 + scale1) + shift1,
            tags,
            tag_padding,
            cos,
            sin,
        )
        return x + gate2 * self.mlp(self.norm2(x) * (1 + scale2) + shift2)


class TagAttention(nn.Module):
    def __init__(self):
        super().__init__()
        self.qkv = nn.Linear(DIM, 3 * DIM, bias=False)
        self.q_norm = RMSNorm(HEAD_DIM)
        self.k_norm = RMSNorm(HEAD_DIM)
        self.proj = nn.Linear(DIM, DIM)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        batch, length, _ = x.shape
        q, k, v = self.qkv(x).view(
            batch, length, 3, HEADS, HEAD_DIM
        ).permute(2, 0, 3, 1, 4)
        out = F.scaled_dot_product_attention(self.q_norm(q), self.k_norm(k), v)
        return self.proj(out.transpose(1, 2).reshape(batch, length, DIM))


class TagRefiner(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm1 = RMSNorm(DIM)
        self.attn = TagAttention()
        self.norm2 = RMSNorm(DIM)
        self.mlp = FeedForward()
        self.adaLN_modulation = nn.Sequential(nn.Linear(DIM, 6 * DIM))

    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
        shift1, scale1, gate1, shift2, scale2, gate2 = self.adaLN_modulation(
            condition
        ).chunk(6, dim=-1)
        x = x + gate1 * self.attn(self.norm1(x) * (1 + scale1) + shift1)
        return x + gate2 * self.mlp(self.norm2(x) * (1 + scale2) + shift2)


class Timestep(nn.Module):
    def __init__(self):
        super().__init__()
        self.mlp = nn.Sequential(nn.Linear(256, DIM), nn.SiLU(), nn.Linear(DIM, DIM))

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        freqs = torch.exp(
            -math.log(10000.0)
            * torch.arange(128, device=t.device, dtype=torch.float32)
            / 128
        )
        args = t.float()[:, None] * 1000.0 * freqs[None]
        return self.mlp(torch.cat([args.cos(), args.sin()], dim=-1))


class OutputResBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.in_ln = nn.LayerNorm(DECODER_DIM, eps=1e-6)
        self.mlp = nn.Sequential(
            nn.Linear(DECODER_DIM, DECODER_DIM),
            nn.SiLU(),
            nn.Linear(DECODER_DIM, DECODER_DIM),
        )
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(), nn.Linear(DECODER_DIM, 3 * DECODER_DIM)
        )

    def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        shift, scale, gate = self.adaLN_modulation(y).chunk(3, dim=-1)
        return x + gate * self.mlp(self.in_ln(x) * (1 + scale) + shift)


class OutputLayer(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm_final = nn.LayerNorm(DECODER_DIM, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(DECODER_DIM, CHANNELS)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.linear(self.norm_final(x))


class OutputDecoder(nn.Module):
    def __init__(self):
        super().__init__()
        self.cond_embed = nn.Linear(DIM, DECODER_DIM)
        self.input_proj = nn.Linear(CHANNELS, DECODER_DIM)
        self.res_blocks = nn.ModuleList(OutputResBlock() for _ in range(OUTPUT_BLOCKS))
        self.final_layer = OutputLayer()

    def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
        x = self.input_proj(x)
        y = self.cond_embed(condition).unsqueeze(1)
        for block in self.res_blocks:
            x = block(x, y)
        return self.final_layer(x)


class AudioDiT(nn.Module):
    def __init__(self):
        super().__init__()
        self.mask_token = nn.Parameter(torch.empty(1, DIM))
        self.fusion_proj = nn.Linear(2 * DIM, DIM)
        self.x_embedder = nn.Linear(CHANNELS, DIM)
        self.t_embedder = Timestep()
        self.tag_embedding = nn.Embedding(NUM_TAGS, DIM, padding_idx=0)
        self.tag_norm = nn.LayerNorm(DIM)
        self.adaLN_shared = nn.Linear(DIM, 6 * DIM)
        self.encoder = nn.ModuleList(
            Block(use_tags=i % 2 == 0, local_context=True)
            for i in range(ENCODER_BLOCKS)
        )
        self.middle = nn.ModuleList(
            Block(use_tags=(i + ENCODER_BLOCKS) % 2 == 0)
            for i in range(MIDDLE_BLOCKS)
        )
        self.decoder = nn.ModuleList(
            Block(
                use_tags=(i + ENCODER_BLOCKS + MIDDLE_BLOCKS) % 2 == 0,
                local_context=True,
            )
            for i in range(DECODER_BLOCKS)
        )
        self.tag_refiner = TagRefiner()
        self.output_decoder = OutputDecoder()
        cos, sin = rope_tables()
        self.register_buffer("cos", cos, persistent=False)
        self.register_buffer("sin", sin, persistent=False)

    def _condition(self, t: float | torch.Tensor, tags: torch.Tensor):
        batch = tags.shape[0]
        if not torch.is_tensor(t):
            t = torch.full((batch,), t, device=tags.device, dtype=torch.float32)
        else:
            t = t.to(tags.device).reshape(-1)
            if t.numel() == 1:
                t = t.expand(batch)
        c = self.t_embedder(t).unsqueeze(1)
        condition = F.silu(c)
        y = self.tag_norm(self.tag_embedding(tags))
        tag_padding = tags == 0
        # Match the training embedder exactly: its float mask promotes the
        # product before it is cast back to the model's embedding dtype.
        y = (y * (~tag_padding).unsqueeze(-1).float()).to(c.dtype)
        y = self.tag_refiner(y, condition)
        return c, self.adaLN_shared(condition), y, tag_padding

    def _run(self, blocks, x, tags, tag_padding, mod):
        for block in blocks:
            x = block(x, tags, tag_padding, mod, self.cos, self.sin)
        return x

    def _predict(self, x: torch.Tensor, t, tags: torch.Tensor, shallow: bool):
        if x.ndim != 3 or x.shape[1:] != (CHANNELS, FRAMES):
            raise ValueError(f"expected latents shaped (B, {CHANNELS}, {FRAMES})")
        c, mod, y, tag_padding = self._condition(t, tags)
        raw = x.transpose(1, 2)
        f = self._run(self.encoder, self.x_embedder(raw), y, tag_padding, mod)
        if shallow:
            g = self.mask_token[:, None, :].expand_as(f)
        else:
            g = self._run(self.middle, f, y, tag_padding, mod)
        h = self.fusion_proj(torch.cat([f, g], dim=-1))
        h = self._run(self.decoder, h, y, tag_padding, mod)
        h = F.silu(c + h)
        batch = x.shape[0]
        out = self.output_decoder(
            raw.reshape(batch * FRAMES, 1, CHANNELS), h.reshape(batch * FRAMES, DIM)
        )
        return out.reshape(batch, FRAMES, CHANNELS).transpose(1, 2)

    def forward(self, x: torch.Tensor, t: float | torch.Tensor, tags: torch.Tensor):
        """Predict clean latents from noisy `(B, 256, 1024)` inputs."""
        return self._predict(x, t, tags, shallow=False)

    def shallow(self, x: torch.Tensor, t: float | torch.Tensor, tags: torch.Tensor):
        """Path-drop branch used as the weak model for alternating PDG."""
        return self._predict(x, t, tags, shallow=True)


def load_model(
    device: str = "cuda", dtype: torch.dtype | None = None
) -> AudioDiT:
    if dtype is None:
        dtype = torch.bfloat16 if device == "cuda" else torch.float32
    state = load_file(WEIGHTS)
    model = AudioDiT()
    model.load_state_dict(state, strict=True)
    return model.to(device=device, dtype=dtype).eval().requires_grad_(False)