"""
Custom diffusers pipeline for ZoomLDM CDM (DiT backbone).
"""

import math
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Optional

import torch
import torch.nn as nn
from diffusers import DDIMScheduler, DiffusionPipeline
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils import BaseOutput


def _modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor) -> torch.Tensor:
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


class Attention(nn.Module):
    """Minimal ViT-style self-attention with timm-compatible parameter names."""

    def __init__(self, dim: int, num_heads: int, qkv_bias: bool = True):
        super().__init__()
        if dim % num_heads != 0:
            raise ValueError(f"dim ({dim}) must be divisible by num_heads ({num_heads})")
        self.num_heads = num_heads
        self.head_dim = dim // num_heads
        self.scale = self.head_dim**-0.5
        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.proj = nn.Linear(dim, dim, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        bsz, tokens, dim = x.shape
        qkv = self.qkv(x)
        qkv = qkv.reshape(bsz, tokens, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # 3, B, H, T, D
        q, k, v = qkv.unbind(0)

        attn = (q @ k.transpose(-2, -1)) * self.scale
        attn = attn.softmax(dim=-1)

        x = attn @ v
        x = x.transpose(1, 2).reshape(bsz, tokens, dim)
        return self.proj(x)


class Mlp(nn.Module):
    """Minimal timm-like MLP block with matching names."""

    def __init__(self, in_features: int, hidden_features: int):
        super().__init__()
        self.fc1 = nn.Linear(in_features, hidden_features, bias=True)
        self.act = nn.GELU(approximate="tanh")
        self.fc2 = nn.Linear(hidden_features, in_features, bias=True)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.fc2(self.act(self.fc1(x)))


class DiTBlock(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int, mlp_ratio: float = 4.0):
        super().__init__()
        self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True)
        self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        mlp_hidden_dim = int(hidden_size * mlp_ratio)
        self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim)
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size, bias=True))

    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
        x = x + gate_msa.unsqueeze(1) * self.attn(_modulate(self.norm1(x), shift_msa, scale_msa))
        x = x + gate_mlp.unsqueeze(1) * self.mlp(_modulate(self.norm2(x), shift_mlp, scale_mlp))
        return x


class TimestepEmbedder(nn.Module):
    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
        super().__init__()
        self.mlp = nn.Sequential(
            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
            nn.SiLU(),
            nn.Linear(hidden_size, hidden_size, bias=True),
        )
        self.frequency_embedding_size = frequency_embedding_size

    @staticmethod
    def timestep_embedding(t: torch.Tensor, dim: int, max_period: int = 10000) -> torch.Tensor:
        half = dim // 2
        freqs = torch.exp(-math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half).to(
            device=t.device
        )
        args = t[:, None].float() * freqs[None]
        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
        if dim % 2:
            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
        return embedding

    def forward(self, t: torch.Tensor) -> torch.Tensor:
        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
        return self.mlp(t_freq)


class LabelEmbedder(nn.Module):
    def __init__(self, num_classes: int, hidden_size: int):
        super().__init__()
        self.embedding_table = nn.Embedding(num_classes, hidden_size)

    def forward(self, labels: torch.Tensor) -> torch.Tensor:
        return self.embedding_table(labels)


class FinalLayer(nn.Module):
    def __init__(self, hidden_size: int, out_channels: int):
        super().__init__()
        self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
        self.linear = nn.Linear(hidden_size, out_channels, bias=True)
        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size, bias=True))

    def forward(self, x: torch.Tensor, c: torch.Tensor) -> torch.Tensor:
        shift, scale = self.adaLN_modulation(c).chunk(2, dim=1)
        x = _modulate(self.norm_final(x), shift, scale)
        return self.linear(x)


class CDMDiTModel(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(
        self,
        num_patches: int = 65,
        in_channels: int = 512,
        hidden_size: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        num_classes: int = 8,
        learn_sigma: bool = True,
    ):
        super().__init__()
        self.learn_sigma = learn_sigma
        self.in_channels = in_channels
        self.out_channels = in_channels * 2 if learn_sigma else in_channels
        self.num_patches = num_patches
        self.x_embedder = nn.Linear(in_channels, hidden_size)
        self.t_embedder = TimestepEmbedder(hidden_size)
        self.y_embedder = LabelEmbedder(num_classes, hidden_size)
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, hidden_size), requires_grad=False)
        self.blocks = nn.ModuleList([DiTBlock(hidden_size, num_heads, mlp_ratio=mlp_ratio) for _ in range(depth)])
        self.final_layer = FinalLayer(hidden_size, self.out_channels)

    def forward(self, x: torch.Tensor, t: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
        # x: (B, C, T), output: (B, out_channels, T)
        x = x.transpose(1, 2)
        x = self.x_embedder(x) + self.pos_embed
        t_emb = self.t_embedder(t)
        y_emb = self.y_embedder(y)
        c = t_emb + y_emb
        for block in self.blocks:
            x = block(x, c)
        x = self.final_layer(x, c)
        return x.transpose(1, 2)


@dataclass
class CDMPipelineOutput(BaseOutput):
    samples: torch.Tensor


class CDMDiTPipeline(DiffusionPipeline):
    def __init__(self, scheduler: DDIMScheduler, cdm: Optional[CDMDiTModel] = None):
        super().__init__()
        self.register_modules(scheduler=scheduler)
        self.cdm = cdm
        self._cdm_root = None
        scheduler_path = getattr(getattr(scheduler, "config", None), "_name_or_path", None)
        if scheduler_path:
            p = Path(scheduler_path)
            self._cdm_root = p.parent if p.name == "scheduler" else p

    @property
    def device(self) -> torch.device:
        self._load_cdm_if_needed()
        return next(self.cdm.parameters()).device

    def to(self, *args, **kwargs):
        self._load_cdm_if_needed()
        self.cdm.to(*args, **kwargs)
        return self

    def _load_cdm_if_needed(self):
        if self.cdm is not None:
            return
        if self._cdm_root is None:
            root_from_cfg = self.config.get("_name_or_path", None)
            if root_from_cfg:
                self._cdm_root = Path(root_from_cfg)
        if self._cdm_root is None:
            raise RuntimeError("Could not infer model root for loading CDM weights.")

        cdm_dir = self._cdm_root / "cdm"
        with open(cdm_dir / "config.json", encoding="utf-8") as f:
            cfg = json.load(f)
        cfg.pop("_class_name", None)
        cfg.pop("_diffusers_version", None)

        cdm = CDMDiTModel(**cfg)
        safetensors_path = cdm_dir / "diffusion_pytorch_model.safetensors"
        bin_path = cdm_dir / "diffusion_pytorch_model.bin"
        if safetensors_path.exists():
            from safetensors.torch import load_file

            state = load_file(str(safetensors_path))
        elif bin_path.exists():
            try:
                state = torch.load(bin_path, map_location="cpu", weights_only=True)
            except TypeError:
                state = torch.load(bin_path, map_location="cpu")
        else:
            raise FileNotFoundError(
                "No CDM weights found in cdm/ (expected diffusion_pytorch_model.safetensors or .bin)."
            )
        cdm.load_state_dict(state, strict=True)
        cdm.eval()
        self.cdm = cdm

    @torch.no_grad()
    def __call__(
        self,
        batch_size: int = 1,
        magnification: Optional[torch.Tensor] = None,
        num_inference_steps: int = 50,
        guidance_scale: float = 1.0,
        num_patches: Optional[int] = None,
        return_dict: bool = True,
    ):
        self._load_cdm_if_needed()
        device = self.device
        dtype = next(self.cdm.parameters()).dtype

        if magnification is None:
            magnification = torch.zeros(batch_size, dtype=torch.long, device=device)
        else:
            magnification = magnification.to(device=device, dtype=torch.long)
            if magnification.ndim == 0:
                magnification = magnification.view(1)

        batch_size = int(magnification.shape[0])
        tokens = num_patches or self.cdm.config.num_patches
        channels = self.cdm.config.in_channels

        latents = torch.randn((batch_size, channels, tokens), device=device, dtype=dtype)
        self.scheduler.set_timesteps(num_inference_steps, device=device)

        for t in self.progress_bar(self.scheduler.timesteps):
            model_in = torch.cat([latents, latents], dim=0)
            t_batch = t.expand(model_in.shape[0]).to(device)
            y_in = torch.cat([torch.zeros_like(magnification), magnification], dim=0)

            model_out = self.cdm(model_in, t_batch, y_in)
            eps, _sigma = model_out.chunk(2, dim=1) if self.cdm.config.learn_sigma else (model_out, None)
            eps_uncond, eps_cond = eps.chunk(2, dim=0)
            eps_guided = eps_uncond + guidance_scale * (eps_cond - eps_uncond)

            latents = self.scheduler.step(eps_guided, t, latents).prev_sample

        if not return_dict:
            return (latents,)
        return CDMPipelineOutput(samples=latents)