#!/usr/bin/env python3
"""
Constellation Diffusion
========================
Everything through the sphere. No skip projection. No attention.
The constellation IS the model's information bottleneck.

16384d encoder output → 256d sphere → 768d triangulation
→ conditioned patchwork → 16384d reconstruction

The patchwork must carry ALL information through 768 geometric
measurements. If it works, diffusion is solved through triangulation.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
import os
import time
from tqdm import tqdm
from torchvision import datasets, transforms
from torchvision.utils import save_image, make_grid

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True


# ══════════════════════════════════════════════════════════════════
# CONSTELLATION BOTTLENECK — NO SKIP
# ══════════════════════════════════════════════════════════════════

class ConstellationBottleneck(nn.Module):
    """
    Pure constellation bottleneck. No skip path.
    All information passes through S^15 triangulation.

    Flow:
      (B, spatial) → proj_in(spatial, embed) → LN → reshape → L2 norm
      → triangulate: P patches × A anchors × n_phases = tri_dim
      → concat(tri, cond)
      → deep patchwork with residual blocks
      → proj_out(hidden, spatial)
    """
    def __init__(
        self,
        spatial_dim,       # C*H*W from encoder
        embed_dim=256,
        patch_dim=16,
        n_anchors=16,
        n_phases=3,
        cond_dim=256,
        pw_hidden=1024,
        pw_depth=4,        # number of residual blocks in patchwork
    ):
        super().__init__()
        self.spatial_dim = spatial_dim
        self.embed_dim = embed_dim
        self.patch_dim = patch_dim
        self.n_patches = embed_dim // patch_dim
        self.n_anchors = n_anchors
        self.n_phases = n_phases

        P, A, d = self.n_patches, n_anchors, patch_dim

        # Encoder projection → sphere
        self.proj_in = nn.Sequential(
            nn.Linear(spatial_dim, embed_dim),
            nn.LayerNorm(embed_dim),
        )

        # Constellation anchors — home + learnable
        home = torch.empty(P, A, d)
        nn.init.xavier_normal_(home.view(P * A, d))
        home = F.normalize(home.view(P, A, d), dim=-1)
        self.register_buffer('home', home)
        self.anchors = nn.Parameter(home.clone())

        # Triangulation dimensions
        tri_dim = P * A * n_phases  # 16 * 16 * 3 = 768

        # Conditioning projection — align cond to patchwork input space
        pw_input = tri_dim + cond_dim
        self.input_proj = nn.Sequential(
            nn.Linear(pw_input, pw_hidden),
            nn.GELU(),
            nn.LayerNorm(pw_hidden),
        )

        # Deep patchwork — residual MLP blocks
        # This must carry ALL information. Make it deep enough.
        self.pw_blocks = nn.ModuleList()
        for _ in range(pw_depth):
            self.pw_blocks.append(nn.Sequential(
                nn.Linear(pw_hidden, pw_hidden),
                nn.GELU(),
                nn.LayerNorm(pw_hidden),
                nn.Linear(pw_hidden, pw_hidden),
                nn.GELU(),
                nn.LayerNorm(pw_hidden),
            ))

        # Reconstruction head
        self.proj_out = nn.Sequential(
            nn.Linear(pw_hidden, pw_hidden),
            nn.GELU(),
            nn.Linear(pw_hidden, spatial_dim),
        )

    def drift(self):
        h, c = F.normalize(self.home, dim=-1), F.normalize(self.anchors, dim=-1)
        return torch.acos((h * c).sum(-1).clamp(-1 + 1e-7, 1 - 1e-7))

    def at_phase(self, t):
        h, c = F.normalize(self.home, dim=-1), F.normalize(self.anchors, dim=-1)
        omega = self.drift().unsqueeze(-1)
        so = omega.sin().clamp(min=1e-7)
        return torch.sin((1-t)*omega)/so * h + torch.sin(t*omega)/so * c

    def triangulate(self, patches_n):
        """
        patches_n: (B, P, d) normalized on S^(d-1)
        Returns: (B, P*A*n_phases) full triangulation profile
        """
        phases = torch.linspace(0, 1, self.n_phases, device=patches_n.device).tolist()
        tris = []
        for t in phases:
            anchors_t = F.normalize(self.at_phase(t), dim=-1)
            cos = torch.einsum('bpd,pad->bpa', patches_n, anchors_t)
            tris.append(1.0 - cos)
        return torch.cat(tris, dim=-1).reshape(patches_n.shape[0], -1)

    def forward(self, x_flat, cond):
        """
        x_flat: (B, spatial_dim)
        cond:   (B, cond_dim)
        Returns: (B, spatial_dim)
        """
        # Project to sphere
        emb = self.proj_in(x_flat)  # (B, embed_dim)
        B = emb.shape[0]
        patches = emb.reshape(B, self.n_patches, self.patch_dim)
        patches_n = F.normalize(patches, dim=-1)  # on S^(d-1)

        # Triangulate — the geometric encoding
        tri = self.triangulate(patches_n)  # (B, tri_dim)

        # Inject conditioning
        pw_in = torch.cat([tri, cond], dim=-1)  # (B, tri_dim + cond_dim)

        # Deep patchwork with residual connections
        h = self.input_proj(pw_in)
        for block in self.pw_blocks:
            h = h + block(h)  # residual

        # Reconstruct spatial features
        return self.proj_out(h)


# ══════════════════════════════════════════════════════════════════
# UNET BUILDING BLOCKS
# ══════════════════════════════════════════════════════════════════

class SinusoidalPosEmb(nn.Module):
    def __init__(self, dim):
        super().__init__()
        self.dim = dim

    def forward(self, t):
        half = self.dim // 2
        emb = math.log(10000) / (half - 1)
        emb = torch.exp(torch.arange(half, device=t.device, dtype=t.dtype) * -emb)
        emb = t.unsqueeze(-1) * emb.unsqueeze(0)
        return torch.cat([emb.sin(), emb.cos()], dim=-1)


class AdaGroupNorm(nn.Module):
    def __init__(self, channels, cond_dim, n_groups=8):
        super().__init__()
        self.gn = nn.GroupNorm(min(n_groups, channels), channels, affine=False)
        self.proj = nn.Linear(cond_dim, channels * 2)
        nn.init.zeros_(self.proj.weight)
        nn.init.zeros_(self.proj.bias)

    def forward(self, x, cond):
        x = self.gn(x)
        s, sh = self.proj(cond).unsqueeze(-1).unsqueeze(-1).chunk(2, dim=1)
        return x * (1 + s) + sh


class ConvBlock(nn.Module):
    def __init__(self, channels, cond_dim):
        super().__init__()
        self.dw = nn.Conv2d(channels, channels, 7, padding=3, groups=channels)
        self.norm = AdaGroupNorm(channels, cond_dim)
        self.pw1 = nn.Conv2d(channels, channels * 4, 1)
        self.pw2 = nn.Conv2d(channels * 4, channels, 1)
        self.act = nn.GELU()

    def forward(self, x, cond):
        r = x
        x = self.dw(x)
        x = self.norm(x, cond)
        x = self.act(self.pw1(x))
        return r + self.pw2(x)


class Downsample(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.conv = nn.Conv2d(ch, ch, 3, stride=2, padding=1)
    def forward(self, x): return self.conv(x)


class Upsample(nn.Module):
    def __init__(self, ch):
        super().__init__()
        self.conv = nn.Conv2d(ch, ch, 3, padding=1)
    def forward(self, x):
        return self.conv(F.interpolate(x, scale_factor=2, mode='nearest'))


# ══════════════════════════════════════════════════════════════════
# CONSTELLATION DIFFUSION UNET
# ══════════════════════════════════════════════════════════════════

class ConstellationDiffusionUNet(nn.Module):
    """
    UNet where the middle block IS the constellation.
    No attention. No skip projection. Pure geometric bottleneck.
    """
    def __init__(
        self,
        in_ch=3,
        base_ch=64,
        ch_mults=(1, 2, 4),
        n_classes=10,
        cond_dim=256,
        embed_dim=256,
        n_anchors=16,
        n_phases=3,
        pw_hidden=1024,
        pw_depth=4,
    ):
        super().__init__()
        self.ch_mults = ch_mults

        # Conditioning
        self.time_emb = nn.Sequential(
            SinusoidalPosEmb(cond_dim),
            nn.Linear(cond_dim, cond_dim), nn.GELU(),
            nn.Linear(cond_dim, cond_dim))
        self.class_emb = nn.Embedding(n_classes, cond_dim)

        self.in_conv = nn.Conv2d(in_ch, base_ch, 3, padding=1)

        # Encoder
        self.enc = nn.ModuleList()
        self.enc_down = nn.ModuleList()
        ch = base_ch
        enc_channels = [base_ch]

        for i, m in enumerate(ch_mults):
            ch_out = base_ch * m
            self.enc.append(nn.ModuleList([
                ConvBlock(ch, cond_dim) if ch == ch_out
                else nn.Sequential(nn.Conv2d(ch, ch_out, 1), ConvBlock(ch_out, cond_dim)),
                ConvBlock(ch_out, cond_dim),
            ]))
            ch = ch_out
            enc_channels.append(ch)
            if i < len(ch_mults) - 1:
                self.enc_down.append(Downsample(ch))

        # Constellation bottleneck — NO SKIP
        mid_ch = ch
        H_mid = 32 // (2 ** (len(ch_mults) - 1))  # spatial size at bottleneck
        spatial_dim = mid_ch * H_mid * H_mid
        self.mid_spatial = (mid_ch, H_mid, H_mid)

        self.bottleneck = ConstellationBottleneck(
            spatial_dim=spatial_dim,
            embed_dim=embed_dim,
            patch_dim=16,
            n_anchors=n_anchors,
            n_phases=n_phases,
            cond_dim=cond_dim,
            pw_hidden=pw_hidden,
            pw_depth=pw_depth,
        )

        # Decoder
        self.dec_up = nn.ModuleList()
        self.dec_skip_proj = nn.ModuleList()
        self.dec = nn.ModuleList()

        for i in range(len(ch_mults) - 1, -1, -1):
            ch_out = base_ch * ch_mults[i]
            skip_ch = enc_channels.pop()
            self.dec_skip_proj.append(nn.Conv2d(ch + skip_ch, ch_out, 1))
            self.dec.append(nn.ModuleList([
                ConvBlock(ch_out, cond_dim),
                ConvBlock(ch_out, cond_dim),
            ]))
            ch = ch_out
            if i > 0:
                self.dec_up.append(Upsample(ch))

        self.out_norm = nn.GroupNorm(8, ch)
        self.out_conv = nn.Conv2d(ch, in_ch, 3, padding=1)
        nn.init.zeros_(self.out_conv.weight)
        nn.init.zeros_(self.out_conv.bias)

    def forward(self, x, t, class_labels):
        cond = self.time_emb(t) + self.class_emb(class_labels)
        h = self.in_conv(x)
        skips = [h]

        # Encoder
        for i in range(len(self.ch_mults)):
            for block in self.enc[i]:
                if isinstance(block, ConvBlock):
                    h = block(h, cond)
                elif isinstance(block, nn.Sequential):
                    h = block[0](h); h = block[1](h, cond)
            skips.append(h)
            if i < len(self.enc_down):
                h = self.enc_down[i](h)

        # ★ CONSTELLATION BOTTLENECK — everything through S^15 ★
        B = h.shape[0]
        h = self.bottleneck(h.reshape(B, -1), cond)
        h = h.reshape(B, *self.mid_spatial)

        # Decoder
        for i in range(len(self.ch_mults)):
            skip = skips.pop()
            if i > 0:
                h = self.dec_up[i - 1](h)
            h = torch.cat([h, skip], dim=1)
            h = self.dec_skip_proj[i](h)
            for block in self.dec[i]:
                h = block(h, cond)

        return self.out_conv(F.silu(self.out_norm(h)))


# ══════════════════════════════════════════════════════════════════
# SAMPLING
# ══════════════════════════════════════════════════════════════════

@torch.no_grad()
def sample(model, n=64, steps=50, cls=None, n_cls=10):
    model.eval()
    x = torch.randn(n, 3, 32, 32, device=DEVICE)
    labels = (torch.full((n,), cls, dtype=torch.long, device=DEVICE)
              if cls is not None else torch.randint(0, n_cls, (n,), device=DEVICE))
    dt = 1.0 / steps
    for s in range(steps):
        t = torch.full((n,), 1.0 - s * dt, device=DEVICE)
        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            v = model(x, t, labels)
        x = x - v.float() * dt
    return x.clamp(-1, 1), labels


# ══════════════════════════════════════════════════════════════════
# TRAINING
# ══════════════════════════════════════════════════════════════════

BATCH = 128
EPOCHS = 80
LR = 3e-4
SAMPLE_EVERY = 5

print("=" * 70)
print("CONSTELLATION DIFFUSION — PURE GEOMETRIC BOTTLENECK")
print(f"  No attention. No skip. Everything through S^15.")
print(f"  Device: {DEVICE}")
print("=" * 70)

transform = transforms.Compose([
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize((0.5,)*3, (0.5,)*3),
])
train_ds = datasets.CIFAR10('./data', train=True, download=True, transform=transform)
train_loader = torch.utils.data.DataLoader(
    train_ds, batch_size=BATCH, shuffle=True,
    num_workers=4, pin_memory=True, drop_last=True)

model = ConstellationDiffusionUNet(
    in_ch=3, base_ch=64, ch_mults=(1, 2, 4),
    n_classes=10, cond_dim=256, embed_dim=256,
    n_anchors=16, n_phases=3, pw_hidden=1024, pw_depth=4,
).to(DEVICE)

n_params = sum(p.numel() for p in model.parameters())
n_bn = sum(p.numel() for p in model.bottleneck.parameters())
n_enc = sum(p.numel() for n, p in model.named_parameters()
            if 'enc' in n or 'in_conv' in n)
n_dec = sum(p.numel() for n, p in model.named_parameters()
            if 'dec' in n or 'out' in n)
n_anchor = sum(p.numel() for n, p in model.named_parameters() if 'anchor' in n)

print(f"  Total:       {n_params:,}")
print(f"  Encoder:     {n_enc:,}")
print(f"  Bottleneck:  {n_bn:,} ({100*n_bn/n_params:.1f}%)")
print(f"    Anchors:   {n_anchor:,}")
print(f"  Decoder:     {n_dec:,}")
print(f"  Train:       {len(train_ds):,} images")

# Shape check
with torch.no_grad():
    d = torch.randn(2, 3, 32, 32, device=DEVICE)
    o = model(d, torch.rand(2, device=DEVICE), torch.randint(0, 10, (2,), device=DEVICE))
    print(f"  Shape:       {d.shape} → {o.shape} ✓")
    bn = model.bottleneck
    print(f"  Bottleneck:  {bn.spatial_dim}d → {bn.embed_dim}d sphere → "
          f"{bn.n_patches}p×{bn.patch_dim}d → "
          f"{bn.n_patches * bn.n_anchors * bn.n_phases} tri dims")
    print(f"  Patchwork:   {len(bn.pw_blocks)} residual blocks × {1024}d")
    print(f"  Compression: {bn.spatial_dim} → {bn.n_patches * bn.n_anchors * bn.n_phases} "
          f"({bn.spatial_dim / (bn.n_patches * bn.n_anchors * bn.n_phases):.1f}× ratio)")

optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=0.01)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
    optimizer, T_max=EPOCHS * len(train_loader), eta_min=1e-6)
scaler = torch.amp.GradScaler("cuda")

os.makedirs("samples_cd", exist_ok=True)
os.makedirs("checkpoints", exist_ok=True)

print(f"\n{'='*70}")
print(f"TRAINING — {EPOCHS} epochs, pure constellation diffusion")
print(f"{'='*70}")

best_loss = float('inf')

for epoch in range(EPOCHS):
    model.train()
    t0 = time.time()
    total_loss = 0
    n = 0

    pbar = tqdm(train_loader, desc=f"E{epoch+1:3d}/{EPOCHS}", unit="b")
    for images, labels in pbar:
        images = images.to(DEVICE, non_blocking=True)
        labels = labels.to(DEVICE, non_blocking=True)
        B = images.shape[0]

        t = torch.rand(B, device=DEVICE)
        eps = torch.randn_like(images)
        t_b = t.view(B, 1, 1, 1)
        x_t = (1 - t_b) * images + t_b * eps
        v_target = eps - images

        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
            v_pred = model(x_t, t, labels)
            loss = F.mse_loss(v_pred, v_target)

        optimizer.zero_grad(set_to_none=True)
        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(optimizer)
        scaler.update()
        scheduler.step()

        total_loss += loss.item()
        n += 1
        if n % 20 == 0:
            pbar.set_postfix(loss=f"{total_loss/n:.4f}", lr=f"{scheduler.get_last_lr()[0]:.1e}")

    elapsed = time.time() - t0
    avg_loss = total_loss / n

    mk = ""
    if avg_loss < best_loss:
        best_loss = avg_loss
        torch.save({
            'state_dict': model.state_dict(),
            'epoch': epoch + 1,
            'loss': avg_loss,
        }, 'checkpoints/constellation_diffusion_best.pt')
        mk = " ★"

    print(f"  E{epoch+1:3d}: loss={avg_loss:.4f} lr={scheduler.get_last_lr()[0]:.1e} "
          f"({elapsed:.0f}s){mk}")

    # Diagnostics
    if (epoch + 1) % 10 == 0:
        with torch.no_grad():
            drift = bn.drift().detach()
            near_029 = (drift - 0.29154).abs().lt(0.05).float().mean().item()
            print(f"  ★ drift: mean={drift.mean():.4f}rad ({math.degrees(drift.mean().item()):.1f}°) "
                  f"max={drift.max():.4f}rad ({math.degrees(drift.max().item()):.1f}°) "
                  f"near_0.29: {near_029:.1%}")

            # Anchor utilization quick check
            test_imgs = torch.randn(64, 3, 32, 32, device=DEVICE)
            t_test = torch.full((64,), 0.5, device=DEVICE)
            c_test = torch.randint(0, 10, (64,), device=DEVICE)
            cond = model.time_emb(t_test) + model.class_emb(c_test)
            h = model.in_conv(test_imgs)
            for i in range(len(model.ch_mults)):
                for block in model.enc[i]:
                    if isinstance(block, ConvBlock): h = block(h, cond)
                    elif isinstance(block, nn.Sequential): h = block[0](h); h = block[1](h, cond)
                if i < len(model.enc_down): h = model.enc_down[i](h)

            emb = bn.proj_in(h.reshape(64, -1))
            patches = F.normalize(emb.reshape(64, bn.n_patches, bn.patch_dim), dim=-1)
            anchors_n = F.normalize(bn.anchors, dim=-1)
            cos = torch.einsum('bpd,pad->bpa', patches, anchors_n)
            nearest = cos.argmax(dim=-1)  # (64, P)
            # Count unique anchors used across all patches
            unique = nearest.unique().numel()
            total = bn.n_patches * bn.n_anchors
            print(f"  ★ anchors: {unique}/{total} unique assignments "
                  f"({100*unique/total:.0f}% utilization)")

    # Sample
    if (epoch + 1) % SAMPLE_EVERY == 0 or epoch == 0:
        imgs, _ = sample(model, 64, 50)
        save_image(make_grid((imgs + 1) / 2, nrow=8), f'samples_cd/epoch_{epoch+1:03d}.png')
        print(f"  → samples_cd/epoch_{epoch+1:03d}.png")

        if (epoch + 1) % 20 == 0:
            names = ['plane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
            for c in range(10):
                cs, _ = sample(model, 8, 50, cls=c)
                save_image(make_grid((cs+1)/2, nrow=8),
                          f'samples_cd/epoch_{epoch+1:03d}_{names[c]}.png')
            print(f"  → per-class samples saved")


print(f"\n{'='*70}")
print(f"CONSTELLATION DIFFUSION — COMPLETE")
print(f"  Best loss: {best_loss:.4f}")
print(f"  Params: {n_params:,} (bottleneck: {n_bn:,})")
with torch.no_grad():
    drift = bn.drift().detach()
    print(f"  Final drift: mean={drift.mean():.4f} max={drift.max():.4f}")
    print(f"  Near 0.29154: {(drift - 0.29154).abs().lt(0.05).float().mean().item():.1%}")
print(f"{'='*70}")