"""
LDM.py
======
Old ATLAS-WDS backend latent diffusion model definitions.

Use these with the old backend checkpoints:
  LDM_VAE.pt   <- vae_best.pt
  LDM_UNET.pt  <- unet_best.pt
Optional:
  ctrl_best.pt can initialize SpecCtrl, but for the new ICWDS Stage-3 path the
  ControlNet/condition adapter should normally be retrained.

Main classes:
  SpecVAE   : 47x72 spectrum VAE, latent shape (B, 8, 6, 9)
  SpecUNet  : diffusion U-Net operating on latent z
  SpecCtrl  : old ControlNet architecture
  Diff      : DDPM/DDIM helper
"""
import math
import torch
import torch.nn as nn
import torch.nn.functional as F


class ResBlock(nn.Module):
    def __init__(self, ci, co):
        super().__init__()
        self.net = nn.Sequential(
            nn.GroupNorm(min(32, ci), ci), nn.SiLU(), nn.Conv2d(ci, co, 3, padding=1),
            nn.GroupNorm(min(32, co), co), nn.SiLU(), nn.Conv2d(co, co, 3, padding=1),
        )
        self.skip = nn.Conv2d(ci, co, 1) if ci != co else nn.Identity()
    def forward(self, x):
        return self.net(x) + self.skip(x)


class DS(nn.Module):
    def __init__(self, c):
        super().__init__(); self.c = nn.Conv2d(c, c, 3, stride=2, padding=1)
    def forward(self, x):
        return self.c(x)


class PSUp(nn.Module):
    def __init__(self, c):
        super().__init__(); self.c = nn.Conv2d(c, c * 4, 3, padding=1); self.p = nn.PixelShuffle(2)
    def forward(self, x):
        return self.p(self.c(x))


class SinEmb(nn.Module):
    def __init__(self, d):
        super().__init__(); self.d = d
    def forward(self, t):
        h = self.d // 2
        e = math.log(10000) / max(h - 1, 1)
        e = torch.exp(torch.arange(h, device=t.device) * -e)
        e = t[:, None].float() * e[None, :]
        return torch.cat([e.sin(), e.cos()], -1)


class URB(nn.Module):
    def __init__(self, ci, co, td):
        super().__init__()
        self.n1 = nn.GroupNorm(min(32, ci), ci); self.c1 = nn.Conv2d(ci, co, 3, padding=1)
        self.n2 = nn.GroupNorm(min(32, co), co); self.c2 = nn.Conv2d(co, co, 3, padding=1)
        self.tp = nn.Linear(td, co * 2); self.sk = nn.Conv2d(ci, co, 1) if ci != co else nn.Identity()
        self.a = nn.SiLU()
    def forward(self, x, te):
        h = self.a(self.n1(x)); h = self.c1(h)
        ss = self.tp(self.a(te))[:, :, None, None]
        sc, sh = ss.chunk(2, dim=1)
        h = self.n2(h) * (1 + sc) + sh
        h = self.a(h)
        return self.c2(h) + self.sk(x)


class Attn(nn.Module):
    def __init__(self, c, nh=4):
        super().__init__(); self.nh = nh; self.norm = nn.GroupNorm(min(32, c), c)
        self.qkv = nn.Conv2d(c, c * 3, 1); self.proj = nn.Conv2d(c, c, 1)
    def forward(self, x):
        B, C, H, W = x.shape
        h = self.norm(x)
        qkv = self.qkv(h).reshape(B, 3, self.nh, C // self.nh, H * W)
        q, k, v = qkv[:, 0], qkv[:, 1], qkv[:, 2]
        a = (torch.einsum('bhdn,bhdm->bhnm', q, k) * (C // self.nh) ** -0.5).softmax(-1)
        out = torch.einsum('bhnm,bhdm->bhdn', a, v).reshape(B, C, H, W)
        return x + self.proj(out)


class SpecVAE(nn.Module):
    def __init__(self, C=64, zc=8):
        super().__init__()
        self.enc = nn.Sequential(
            nn.Conv2d(1, C, 3, padding=1), ResBlock(C, C), ResBlock(C, C), DS(C),
            ResBlock(C, 2 * C), ResBlock(2 * C, 2 * C), DS(2 * C),
            ResBlock(2 * C, 4 * C), ResBlock(4 * C, 4 * C), DS(4 * C),
            ResBlock(4 * C, 4 * C), ResBlock(4 * C, 4 * C),
        )
        self.en = nn.GroupNorm(32, 4 * C); self.ea = nn.SiLU()
        self.mu = nn.Conv2d(4 * C, zc, 1); self.lv = nn.Conv2d(4 * C, zc, 1)
        self.di = nn.Conv2d(zc, 4 * C, 1)
        self.dec = nn.Sequential(
            ResBlock(4 * C, 4 * C), ResBlock(4 * C, 4 * C), PSUp(4 * C),
            ResBlock(4 * C, 2 * C), ResBlock(2 * C, 2 * C), PSUp(2 * C),
            ResBlock(2 * C, C), ResBlock(C, C), PSUp(C),
            ResBlock(C, C), nn.GroupNorm(32, C), nn.SiLU(), nn.Conv2d(C, 1, 3, padding=1), nn.Tanh(),
        )

    @torch.no_grad()
    def enc_lat(self, x):
        x = F.pad(x, (0, 0, 0, 1), value=-1)
        h = self.ea(self.en(self.enc(x)))
        return self.mu(h)

    def encode(self, x):
        x = F.pad(x, (0, 0, 0, 1), value=-1)
        h = self.ea(self.en(self.enc(x)))
        return self.mu(h), self.lv(h)

    def decode(self, z):
        return self.dec(self.di(z))[:, :, :47, :]


class SpecUNet(nn.Module):
    def __init__(self, C=128, zc=8):
        super().__init__(); td = C * 4
        self.te = nn.Sequential(SinEmb(C), nn.Linear(C, td), nn.SiLU(), nn.Linear(td, td))
        self.ci = nn.Conv2d(zc, C, 3, padding=1)
        self.e1a = URB(C, C, td); self.e1b = URB(C, C, td); self.d1 = nn.Conv2d(C, C, 3, stride=2, padding=1)
        self.e2a = URB(C, 2 * C, td); self.e2b = URB(2 * C, 2 * C, td)
        self.m1 = URB(2 * C, 2 * C, td); self.ma = Attn(2 * C); self.m2 = URB(2 * C, 2 * C, td)
        self.d2a = URB(4 * C, 2 * C, td); self.d2b = URB(2 * C, 2 * C, td)
        self.u1 = nn.ConvTranspose2d(2 * C, 2 * C, 4, stride=2, padding=1)
        self.d1a = URB(2 * C + C, C, td); self.d1b = URB(C, C, td)
        self.out = nn.Sequential(nn.GroupNorm(min(32, C), C), nn.SiLU(), nn.Conv2d(C, zc, 3, padding=1))

    def forward(self, x, t, cf=None):
        te = self.te(t)
        h = self.ci(x)
        s1 = self.e1b(self.e1a(h, te), te)
        h = self.d1(s1)
        s2 = self.e2b(self.e2a(h, te), te)
        h = self.m2(self.ma(self.m1(s2, te)), te)
        if cf:
            h = h + cf.get('mid', 0)
            s1 = s1 + cf.get('s1', 0)
        h = self.d2b(self.d2a(torch.cat([h, s2], 1), te), te)
        h = self.u1(h)[:, :, :6, :9]
        h = self.d1b(self.d1a(torch.cat([h, s1], 1), te), te)
        return self.out(h)


class ZC(nn.Module):
    def __init__(self, ci, co):
        super().__init__(); self.c = nn.Conv2d(ci, co, 1)
        nn.init.zeros_(self.c.weight); nn.init.zeros_(self.c.bias)
    def forward(self, x):
        return self.c(x)


class SpecCtrl(nn.Module):
    def __init__(self, C=128, zc=8):
        super().__init__(); td = C * 4
        self.ci = nn.Conv2d(zc, zc, 1)
        self.te = nn.Sequential(SinEmb(C), nn.Linear(C, td), nn.SiLU(), nn.Linear(td, td))
        self.cin = nn.Conv2d(zc * 2, C, 3, padding=1)
        self.e1a = URB(C, C, td); self.e1b = URB(C, C, td); self.d1 = nn.Conv2d(C, C, 3, stride=2, padding=1)
        self.e2a = URB(C, 2 * C, td); self.e2b = URB(2 * C, 2 * C, td)
        self.m1 = URB(2 * C, 2 * C, td); self.ma = Attn(2 * C); self.m2 = URB(2 * C, 2 * C, td)
        self.zs1 = ZC(C, C); self.zm = ZC(2 * C, 2 * C)

    def forward(self, zn, t, zc):
        te = self.te(t)
        c = self.ci(zc)
        h = self.cin(torch.cat([zn, c], 1))
        s1 = self.e1b(self.e1a(h, te), te)
        h = self.d1(s1)
        h = self.e2b(self.e2a(h, te), te)
        h = self.m2(self.ma(self.m1(h, te)), te)
        return {'s1': self.zs1(s1), 'mid': self.zm(h)}


class Diff:
    def __init__(self, T=1000, bs=1e-4, be=0.02, dev='cpu'):
        b = torch.linspace(bs, be, T)
        a = 1 - b
        ac = torch.cumprod(a, 0)
        self.T = T
        self.ac = ac.to(dev)
        self.sac = ac.sqrt().to(dev)
        self.somc = (1 - ac).sqrt().to(dev)
        self.dev = dev

    def q_sample(self, z0, t, noise=None):
        if noise is None:
            noise = torch.randn_like(z0)
        return self.sac[t].view(-1, 1, 1, 1) * z0 + self.somc[t].view(-1, 1, 1, 1) * noise, noise

    @torch.no_grad()
    def ddim(self, model, shape, steps=80, cfn=None):
        si = torch.linspace(0, self.T - 1, steps, dtype=torch.long, device=self.dev)
        x = torch.randn(shape, device=self.dev)
        for i in reversed(range(len(si))):
            t = si[i].expand(shape[0])
            at = self.ac[t[0]]
            ap = self.ac[si[i - 1]] if i > 0 else torch.tensor(1.0, device=self.dev)
            c = cfn(x, t) if cfn else None
            np_ = model(x, t, cf=c)
            x0p = ((x - np_ * (1 - at).sqrt()) / at.sqrt()).clamp(-3, 3)
            x = ap.sqrt() * x0p + np_ * (1 - ap).sqrt()
        return x


# Optional aliases.
VAE = SpecVAE
UNet = SpecUNet
ControlNet = SpecCtrl
Model = SpecUNet