"""PixelModel v6 — 155M-parameter text-to-image MMDiT, generates 256×256 images."""

import json
import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # MUST come before torch
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from PIL import Image
from safetensors.torch import load_file
from diffusers import AutoencoderKL
from transformers import CLIPTextModel, CLIPTokenizer, T5EncoderModel, T5TokenizerFast
import gradio as gr

# ─── Model architecture (copied from dit_v6.py, trust_remote_code equivalent) ──

import math
import torch.utils.checkpoint


def modulate(x, shift, scale):
    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)


def timestep_embedding(t, dim, max_period=10000):
    half = dim // 2
    freqs = torch.exp(-math.log(max_period) * torch.arange(half, device=t.device) / half)
    args = t[:, None].float() * freqs[None]
    emb = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
    if dim % 2:
        emb = torch.cat([emb, torch.zeros_like(emb[:, :1])], dim=-1)
    return emb


def rope_freqs(positions, dim, base=10000.0):
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    return torch.outer(positions.float(), inv_freq)


def rope_cos_sin(freqs):
    emb = torch.cat([freqs, freqs], dim=-1)
    return emb.cos(), emb.sin()


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x, cos, sin):
    return x * cos + rotate_half(x) * sin


def apply_rope_2d(x, row_cos, row_sin, col_cos, col_sin):
    x1, x2 = x.chunk(2, dim=-1)
    x1 = apply_rope(x1, row_cos, row_sin)
    x2 = apply_rope(x2, col_cos, col_sin)
    return torch.cat([x1, x2], dim=-1)


class RMSNormHead(nn.Module):
    def __init__(self, head_dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(head_dim))
        self.eps = eps

    def forward(self, x):
        n = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * n * self.weight


class SwiGLU(nn.Module):
    def __init__(self, dim, hidden):
        super().__init__()
        self.gate = nn.Linear(dim, hidden)
        self.up = nn.Linear(dim, hidden)
        self.down = nn.Linear(hidden, dim)

    def forward(self, x):
        return self.down(F.silu(self.gate(x)) * self.up(x))


class JointBlock(nn.Module):
    def __init__(self, dim, heads, mlp_hidden):
        super().__init__()
        self.heads = heads
        self.head_dim = dim // heads
        self.norm1_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm1_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.qkv_img = nn.Linear(dim, 3 * dim)
        self.qkv_txt = nn.Linear(dim, 3 * dim)
        self.qn_img = RMSNormHead(self.head_dim)
        self.kn_img = RMSNormHead(self.head_dim)
        self.qn_txt = RMSNormHead(self.head_dim)
        self.kn_txt = RMSNormHead(self.head_dim)
        self.proj_img = nn.Linear(dim, dim)
        self.proj_txt = nn.Linear(dim, dim)
        self.norm2_img = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.norm2_txt = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.mlp_img = SwiGLU(dim, mlp_hidden)
        self.mlp_txt = SwiGLU(dim, mlp_hidden)
        self.ada_img = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))
        self.ada_txt = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim))

    def forward(self, img, txt, c, rope_img, rope_txt, key_valid):
        s1i, sc1i, g1i, s2i, sc2i, g2i = self.ada_img(c).chunk(6, dim=-1)
        s1t, sc1t, g1t, s2t, sc2t, g2t = self.ada_txt(c).chunk(6, dim=-1)

        xi = modulate(self.norm1_img(img), s1i, sc1i)
        xt = modulate(self.norm1_txt(txt), s1t, sc1t)

        B, Ni, C = xi.shape
        Nt = xt.shape[1]
        H, D = self.heads, self.head_dim

        qi, ki, vi = self.qkv_img(xi).reshape(B, Ni, 3, H, D).permute(2, 0, 3, 1, 4)
        qt, kt, vt = self.qkv_txt(xt).reshape(B, Nt, 3, H, D).permute(2, 0, 3, 1, 4)

        qi, ki = self.qn_img(qi), self.kn_img(ki)
        qt, kt = self.qn_txt(qt), self.kn_txt(kt)

        row_cos, row_sin, col_cos, col_sin = rope_img
        qi = apply_rope_2d(qi, row_cos, row_sin, col_cos, col_sin)
        ki = apply_rope_2d(ki, row_cos, row_sin, col_cos, col_sin)

        t_cos, t_sin = rope_txt
        qt = apply_rope(qt, t_cos, t_sin)
        kt = apply_rope(kt, t_cos, t_sin)

        q = torch.cat([qi, qt], dim=2)
        k = torch.cat([ki, kt], dim=2)
        v = torch.cat([vi, vt], dim=2)

        mask = key_valid[:, None, None, :]
        o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
        o = o.transpose(1, 2).reshape(B, Ni + Nt, C)
        oi, ot = o[:, :Ni], o[:, Ni:]

        img = img + g1i.unsqueeze(1) * self.proj_img(oi)
        txt = txt + g1t.unsqueeze(1) * self.proj_txt(ot)

        img = img + g2i.unsqueeze(1) * self.mlp_img(modulate(self.norm2_img(img), s2i, sc2i))
        txt = txt + g2t.unsqueeze(1) * self.mlp_txt(modulate(self.norm2_txt(txt), s2t, sc2t))
        return img, txt


class MMDiT(nn.Module):
    def __init__(self, latent_ch=4, latent_size=32, patch=2, dim=512, depth=16, heads=8,
                 t5_dim=768, clip_dim=512, t5_len=32, mlp_hidden=1408,
                 repa_dim=384, repa_layer=8):
        super().__init__()
        self.latent_ch = latent_ch
        self.latent_size = latent_size
        self.patch = patch
        self.grid = latent_size // patch
        self.patch_dim = latent_ch * patch * patch
        self.dim = dim
        self.depth = depth
        self.heads = heads
        self.head_dim = dim // heads
        self.t5_len = t5_len
        self.repa_layer = repa_layer

        self.x_embed = nn.Linear(self.patch_dim, dim)
        self.t_mlp = nn.Sequential(nn.Linear(dim, dim), nn.SiLU(), nn.Linear(dim, dim))
        self.clip_proj = nn.Linear(clip_dim, dim)
        self.t5_proj = nn.Linear(t5_dim, dim)

        self.blocks = nn.ModuleList([JointBlock(dim, heads, mlp_hidden) for _ in range(depth)])

        self.norm_out = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
        self.ada_out = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim))
        self.head = nn.Linear(dim, self.patch_dim)

        self.repa_head = nn.Sequential(nn.Linear(dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, repa_dim))

        hd2 = self.head_dim // 2
        rows = torch.arange(self.grid).repeat_interleave(self.grid)
        cols = torch.arange(self.grid).repeat(self.grid)
        row_cos, row_sin = rope_cos_sin(rope_freqs(rows, hd2))
        col_cos, col_sin = rope_cos_sin(rope_freqs(cols, hd2))
        self.register_buffer("row_cos", row_cos, persistent=False)
        self.register_buffer("row_sin", row_sin, persistent=False)
        self.register_buffer("col_cos", col_cos, persistent=False)
        self.register_buffer("col_sin", col_sin, persistent=False)
        t_cos, t_sin = rope_cos_sin(rope_freqs(torch.arange(t5_len), self.head_dim))
        self.register_buffer("t_cos", t_cos, persistent=False)
        self.register_buffer("t_sin", t_sin, persistent=False)

        self._init()

    def _init(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
        for b in self.blocks:
            nn.init.zeros_(b.ada_img[-1].weight); nn.init.zeros_(b.ada_img[-1].bias)
            nn.init.zeros_(b.ada_txt[-1].weight); nn.init.zeros_(b.ada_txt[-1].bias)
        nn.init.zeros_(self.ada_out[-1].weight); nn.init.zeros_(self.ada_out[-1].bias)
        nn.init.zeros_(self.head.weight); nn.init.zeros_(self.head.bias)

    def patchify(self, x):
        B, C, H, W = x.shape
        p = self.patch
        x = x.reshape(B, C, H // p, p, W // p, p)
        x = x.permute(0, 2, 4, 1, 3, 5).reshape(B, (H // p) * (W // p), C * p * p)
        return x

    def unpatchify(self, x):
        B, N, _ = x.shape
        p = self.patch
        g = self.grid
        C = self.latent_ch
        x = x.reshape(B, g, g, C, p, p).permute(0, 3, 1, 4, 2, 5)
        return x.reshape(B, C, g * p, g * p)

    def forward(self, x, t, t5_seq, t5_mask, clip_pool, return_repa=False, use_checkpoint=False):
        B = x.shape[0]
        img = self.x_embed(self.patchify(x))
        txt = self.t5_proj(t5_seq)
        c = self.t_mlp(timestep_embedding(t, self.dim)) + self.clip_proj(clip_pool)

        key_valid = torch.cat([
            torch.ones(B, img.shape[1], dtype=torch.bool, device=x.device),
            t5_mask.bool(),
        ], dim=1)

        rope_img = (self.row_cos, self.row_sin, self.col_cos, self.col_sin)
        rope_txt = (self.t_cos, self.t_sin)

        repa_hidden = None
        for i, blk in enumerate(self.blocks):
            if use_checkpoint and self.training:
                img, txt = torch.utils.checkpoint.checkpoint(
                    blk, img, txt, c, rope_img, rope_txt, key_valid, use_reentrant=False)
            else:
                img, txt = blk(img, txt, c, rope_img, rope_txt, key_valid)
            if return_repa and i == self.repa_layer:
                repa_hidden = img

        shift, scale = self.ada_out(c).chunk(2, dim=-1)
        img = modulate(self.norm_out(img), shift, scale)
        out = self.unpatchify(self.head(img))

        if return_repa:
            return out, self.repa_head(repa_hidden)
        return out


# ─── Load everything at module scope (ZeroGPU rule 2) ──────────────────────────

MODEL_REPO = "bench-labs/PixelModel-v6"
VAE_REPO = "madebyollin/sdxl-vae-fp16-fix"
CLIP_REPO = "openai/clip-vit-base-patch32"
T5_REPO = "google/flan-t5-base"
T5_LEN = 32
CLIP_LEN = 40

_config = json.load(open("config.json"))["dit"]
model = MMDiT(
    dim=_config["dim"], depth=_config["depth"], heads=_config["heads"],
    mlp_hidden=_config["mlp_hidden"], t5_len=_config["t5_len"],
).to("cuda").eval()
# strict=False because the released safetensors omits the repa_head
# (training-only auxiliary projection head, dropped from published weights)
model.load_state_dict(load_file("model.safetensors"), strict=False)

vae = AutoencoderKL.from_pretrained(VAE_REPO).to("cuda").half().eval()
vae_scale = vae.config.scaling_factor

t5_tok = T5TokenizerFast.from_pretrained(T5_REPO)
t5 = T5EncoderModel.from_pretrained(T5_REPO).to("cuda").eval()
clip_tok = CLIPTokenizer.from_pretrained(CLIP_REPO)
clip_txt = CLIPTextModel.from_pretrained(CLIP_REPO).to("cuda").eval()


def _encode(strings):
    """Encode text into T5 sequence + CLIP pooled vector (matches main.py exactly)."""
    te = t5_tok(strings, padding="max_length", max_length=T5_LEN, truncation=True,
                return_tensors="pt").to("cuda")
    seq = t5(input_ids=te["input_ids"], attention_mask=te["attention_mask"]).last_hidden_state.float()
    ce = clip_tok(strings, padding="max_length", max_length=CLIP_LEN, truncation=True,
                   return_tensors="pt").to("cuda")
    pool = clip_txt(input_ids=ce["input_ids"]).pooler_output.float()
    return seq, te["attention_mask"].float(), pool


# Null (unconditional) embedding — lazily computed inside GPU context
# (can't run encoders at module scope — no GPU attached until @spaces.GPU)
_null_cache = None


# ─── Inference (matches main.py sampling loop exactly) ─────────────────────────

@spaces.GPU(duration=30)
def generate(prompt: str, cfg: float = 3.0, steps: int = 50, seed: int = 0,
             progress: gr.Progress = gr.Progress(track_tqdm=True)):
    """Generate a 256x256 image from a text prompt using PixelModel v6.

    Args:
        prompt: The text prompt describing what to generate.
        cfg: Classifier-free guidance scale (3.0 is the model's sweet spot).
        steps: Number of rectified-flow sampling steps.
        seed: RNG seed for reproducibility (0 = random each time).
    """
    global _null_cache

    if seed != 0:
        torch.manual_seed(seed)

    seq, mask, pool = _encode([prompt])

    if _null_cache is None:
        _null_cache = _encode([""])
    null_seq, null_mask, null_pool = _null_cache

    B = seq.shape[0]
    x = torch.randn(B, 4, 32, 32, device="cuda")
    ns = null_seq.expand(B, -1, -1)
    nm = null_mask.expand(B, -1)
    npo = null_pool.expand(B, -1)
    dt = 1.0 / steps

    with torch.no_grad():
        for i in range(steps):
            t = torch.full((B,), i * dt, device="cuda")
            with torch.autocast("cuda", dtype=torch.bfloat16):
                vc = model(x, t, seq, mask, pool)
                vu = model(x, t, ns, nm, npo)
            x = x + (vu + cfg * (vc - vu)).float() * dt
            progress((i + 1) / steps)

        img = vae.decode((x / vae_scale).half()).sample.float()
    img = ((img.clamp(-1, 1) + 1) / 2)[0].permute(1, 2, 0).cpu().numpy()
    return Image.fromarray((img * 255).round().astype(np.uint8))


# ─── UI ────────────────────────────────────────────────────────────────────────

CSS = """
#col-container { max-width: 900px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks() as demo:
    gr.Markdown("# PixelModel v6\n155M-parameter MMDiT text-to-image model generating 256×256 images.")
    with gr.Column(elem_id="col-container"):
        with gr.Row():
            prompt = gr.Textbox(
                show_label=False,
                placeholder="Describe an image…",
                container=False,
                scale=4,
            )
            run = gr.Button("Generate", variant="primary", scale=1)

        output = gr.Image(label="Generated image", height=320)

        with gr.Accordion("Advanced settings", open=False):
            cfg = gr.Slider(1.0, 10.0, value=3.0, step=0.5,
                           label="CFG (guidance scale)", info="3.0 is the model's optimal value")
            steps = gr.Slider(10, 100, value=50, step=5,
                             label="Steps", info="50 steps recommended")
            seed = gr.Number(label="Seed (0 = random)", value=0, precision=0)

        gr.Examples(
            examples=[
                ["a bowl of ramen with a soft boiled egg"],
                ["a red fox sitting in a snowy forest"],
                ["a lighthouse on a cliff at sunset"],
                ["a golden retriever running on a beach"],
                ["a city street at night with neon signs"],
                ["a cup of coffee on a wooden table"],
            ],
            inputs=[prompt],
            outputs=output,
            fn=generate,
            cache_examples=True,
            cache_mode="lazy",
        )

    run.click(
        generate,
        inputs=[prompt, cfg, steps, seed],
        outputs=output,
        api_name="generate",
    )

    demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)