"""HoLo-FuSe inference: UNet + frozen-HSL label conditioning + respaced DDIM sampler.
Matches the training code in https://github.com/Woojiggun/HoLo-FuSe (experiment.py).
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F
import hsl_embedding_zero as hslz

FEAT_DIM = hslz.FEAT_DIM  # 27


def timestep_emb(t, dim):
    half = dim // 2
    freqs = torch.exp(-math.log(10000) * torch.arange(half, device=t.device) / (half - 1))
    a = t.float()[:, None] * freqs[None]
    return torch.cat([torch.sin(a), torch.cos(a)], dim=-1)


class ResBlock(nn.Module):
    def __init__(self, cin, cout, emb_dim, groups=8):
        super().__init__()
        self.n1 = nn.GroupNorm(groups, cin); self.c1 = nn.Conv2d(cin, cout, 3, padding=1)
        self.emb = nn.Linear(emb_dim, cout)
        self.n2 = nn.GroupNorm(groups, cout); self.c2 = nn.Conv2d(cout, cout, 3, padding=1)
        self.skip = nn.Conv2d(cin, cout, 1) if cin != cout else nn.Identity()

    def forward(self, x, emb):
        h = self.c1(F.silu(self.n1(x)))
        h = h + self.emb(emb)[:, :, None, None]
        h = self.c2(F.silu(self.n2(h)))
        return h + self.skip(x)


class AttnBlock(nn.Module):
    def __init__(self, ch, groups=8):
        super().__init__()
        self.norm = nn.GroupNorm(groups, ch)
        self.qkv = nn.Conv2d(ch, ch * 3, 1)
        self.proj = nn.Conv2d(ch, ch, 1)

    def forward(self, x):
        B, C, H, W = x.shape
        q, k, v = self.qkv(self.norm(x)).chunk(3, dim=1)
        q = q.reshape(B, C, H * W).permute(0, 2, 1)
        k = k.reshape(B, C, H * W)
        v = v.reshape(B, C, H * W).permute(0, 2, 1)
        a = torch.softmax(q @ k * (C ** -0.5), dim=-1)
        o = (a @ v).permute(0, 2, 1).reshape(B, C, H, W)
        return x + self.proj(o)


class Downsample(nn.Module):
    def __init__(self, ch):
        super().__init__(); self.op = nn.Conv2d(ch, ch, 3, stride=2, padding=1)

    def forward(self, x, emb=None): return self.op(x)


class Upsample(nn.Module):
    def __init__(self, ch):
        super().__init__(); self.op = nn.Conv2d(ch, ch, 3, padding=1)

    def forward(self, x, emb=None): return self.op(F.interpolate(x, scale_factor=2, mode="nearest"))


class UNet(nn.Module):
    def __init__(self, c=3, base=128, ch_mults=(1, 2, 2, 2), num_res=2, attn_res=(16,),
                 emb_dim=256, cond_dim=128, img_size=128):
        super().__init__()
        self.emb_dim = emb_dim
        attn_res = set(attn_res)
        self.tproj = nn.Sequential(nn.Linear(emb_dim, emb_dim), nn.SiLU(), nn.Linear(emb_dim, emb_dim))
        self.cproj = nn.Linear(cond_dim, emb_dim)
        self.in_conv = nn.Conv2d(c, base, 3, padding=1)
        L = len(ch_mults)
        chs = [base]; now = base
        self.downs = nn.ModuleList()
        for i, m in enumerate(ch_mults):
            res = img_size >> i; out = base * m
            for _ in range(num_res):
                grp = [ResBlock(now, out, emb_dim)]; now = out
                if res in attn_res: grp.append(AttnBlock(out))
                self.downs.append(nn.ModuleList(grp)); chs.append(now)
            if i != L - 1:
                self.downs.append(Downsample(now)); chs.append(now)
        self.mid1 = ResBlock(now, now, emb_dim); self.mid_attn = AttnBlock(now); self.mid2 = ResBlock(now, now, emb_dim)
        self.ups = nn.ModuleList()
        for i, m in reversed(list(enumerate(ch_mults))):
            res = img_size >> i; out = base * m
            for _ in range(num_res + 1):
                grp = [ResBlock(now + chs.pop(), out, emb_dim)]; now = out
                if res in attn_res: grp.append(AttnBlock(out))
                self.ups.append(nn.ModuleList(grp))
            if i != 0:
                self.ups.append(Upsample(now))
        self.out = nn.Sequential(nn.GroupNorm(8, now), nn.SiLU(), nn.Conv2d(now, c, 3, padding=1))

    def forward(self, x, t, cond):
        emb = self.tproj(timestep_emb(t, self.emb_dim))
        if cond is not None:
            emb = emb + self.cproj(cond)
        h = self.in_conv(x); hs = [h]
        for stage in self.downs:
            if isinstance(stage, Downsample):
                h = stage(h)
            else:
                for blk in stage:
                    h = blk(h, emb) if isinstance(blk, ResBlock) else blk(h)
            hs.append(h)
        h = self.mid1(h, emb); h = self.mid_attn(h); h = self.mid2(h, emb)
        for stage in self.ups:
            if isinstance(stage, Upsample):
                h = stage(h)
            else:
                h = torch.cat([h, hs.pop()], dim=1)
                for blk in stage:
                    h = blk(h, emb) if isinstance(blk, ResBlock) else blk(h)
        return self.out(h)


class HSLLabelCond(nn.Module):
    """label bytes -> frozen 27-D HSL frame (0 learned params) -> small learned readout."""
    def __init__(self, cond_dim, K=8):
        super().__init__()
        self.K = K
        self.zero = hslz.ZeroInput(K=K, dim=512)
        self.readout = nn.Sequential(nn.Linear(K * FEAT_DIM, cond_dim), nn.SiLU(),
                                     nn.Linear(cond_dim, cond_dim))

    @torch.no_grad()
    def _hsl(self, label_strs):
        ids = []
        for s in label_strs:
            b = s.encode("utf-8", "replace")[:self.K]
            b = b + b"\x00" * (self.K - len(b))
            ids.append(list(b))
        ids = torch.tensor(ids, dtype=torch.long)
        return self.zero.features(ids).reshape(len(label_strs), -1)

    def forward(self, label_strs):
        return self.readout(self._hsl(label_strs).float().to(next(self.parameters()).device))


class LearnedLabelCond(nn.Module):
    """Control arm: same-budget learned embedding (definition matches training exactly)."""
    def __init__(self, cond_dim, class_list):
        super().__init__()
        self.idx = {c: i for i, c in enumerate(class_list)}
        self.emb = nn.Embedding(len(class_list), cond_dim)
        self.readout = nn.Sequential(nn.Linear(cond_dim, cond_dim), nn.SiLU(),
                                     nn.Linear(cond_dim, cond_dim))

    def forward(self, label_strs):
        ii = torch.tensor([self.idx.get(s, 0) for s in label_strs],
                          device=self.emb.weight.device)
        return self.readout(self.emb(ii))


def make_cosine_acp(T=250):
    s = 0.008
    steps = torch.arange(T + 1, dtype=torch.float64)
    ac = torch.cos(((steps / T) + s) / (1 + s) * math.pi / 2) ** 2
    ac = ac / ac[0]
    betas = (1 - ac[1:] / ac[:-1]).clamp(1e-4, 0.999).float()
    return torch.cumprod(1 - betas, 0)


@torch.no_grad()
def ddim_sample(model, cond, uncond, n=1, c=3, hw=128, T=250, steps=16, cfg=1.6,
                dyn=0.99, seed=None, progress=None):
    """Respaced DDIM (eta=0) over the trained T=250 cosine schedule, with CFG + dynamic thresholding."""
    dev = next(model.parameters()).device
    if seed is not None and seed >= 0:
        torch.manual_seed(seed)
    acp = make_cosine_acp(T)
    ts = torch.linspace(0, T - 1, steps).round().long().unique().tolist()  # ascending
    x = torch.randn(n, c, hw, hw, device=dev)
    for k in range(len(ts) - 1, -1, -1):
        t = ts[k]
        tt = torch.full((n,), t, dtype=torch.long, device=dev)
        if cond is not None and cfg != 1.0:
            eps = model(x, tt, uncond) if uncond is not None else model(x, tt, None)
            eps = eps + cfg * (model(x, tt, cond) - eps)
        else:
            eps = model(x, tt, cond)
        a = acp[t]
        x0 = (x - (1 - a).sqrt() * eps) / a.sqrt()
        if dyn and dyn > 0:
            s = torch.quantile(x0.reshape(n, -1).abs(), dyn, dim=1).clamp(min=1.0).view(-1, 1, 1, 1)
            x0 = x0.clamp(-s, s) / s
        else:
            x0 = x0.clamp(-1, 1)
        if k == 0:
            x = x0
        else:
            ap = acp[ts[k - 1]]
            x = ap.sqrt() * x0 + (1 - ap).sqrt() * eps
        if progress is not None:
            progress((len(ts) - k) / len(ts))
    return x.clamp(-1, 1)


def load_holofuse(arm="hsl", ckpt_path=None, device="cpu"):
    """One-call loader: returns (model, cond_enc) with EMA weights applied.
    arm: 'hsl' | 'learned' | 'none'. Downloads from ggunio/HoLo-FuSe unless ckpt_path is given."""
    if ckpt_path is None:
        from huggingface_hub import hf_hub_download
        ckpt_path = hf_hub_download("ggunio/HoLo-FuSe", f"holofuse_{arm}_128.pt")
    st = torch.load(ckpt_path, map_location=device, weights_only=False)
    model = UNet().to(device).eval()
    if arm == "hsl":
        cond_enc = HSLLabelCond(128).to(device).eval()
    elif arm == "learned":
        cond_enc = LearnedLabelCond(128, ["Cat", "Dog"]).to(device).eval()
    else:
        cond_enc = None
    params = list(model.parameters()) + (list(cond_enc.parameters()) if cond_enc else [])
    with torch.no_grad():
        for p, e in zip(params, st["ema"]):
            p.copy_(e.to(device))
    return model, cond_enc


def generate(label="Cat", arm="hsl", steps=16, cfg=1.6, seed=0, n=1, device="cpu",
             model=None, cond_enc=None):
    """One-call generation -> list of PIL images (128px)."""
    from PIL import Image
    if model is None:
        model, cond_enc = load_holofuse(arm, device=device)
    with torch.no_grad():
        cond = cond_enc([label] * n) if cond_enc else None
        uncond = torch.zeros_like(cond) if cond is not None else None
        x = ddim_sample(model, cond, uncond, n=n, steps=steps, cfg=cfg, dyn=0.99, seed=seed)
    arrs = ((x.clamp(-1, 1) + 1) * 127.5).to(torch.uint8).cpu().numpy().transpose(0, 2, 3, 1)
    return [Image.fromarray(a, "RGB") for a in arrs]