"""
Constellation Bottleneck — Full Analysis
==========================================
Paste directly after the training cell.
Uses `model` already in memory.
"""

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

DEVICE = "cuda"
os.makedirs("analysis_bn", exist_ok=True)

def compute_cv(points, n_samples=1500, n_points=5):
    N = points.shape[0]
    if N < n_points: return float('nan')
    points = F.normalize(points.to(DEVICE).float(), dim=-1)
    vols = []
    for _ in range(n_samples):
        idx = torch.randperm(min(N, 5000), device=DEVICE)[:n_points]
        pts = points[idx].unsqueeze(0)
        gram = torch.bmm(pts, pts.transpose(1, 2))
        norms = torch.diagonal(gram, dim1=1, dim2=2)
        d2 = norms.unsqueeze(2) + norms.unsqueeze(1) - 2 * gram
        d2 = F.relu(d2)
        cm = torch.zeros(1, 6, 6, device=DEVICE, dtype=torch.float32)
        cm[:, 0, 1:] = 1; cm[:, 1:, 0] = 1; cm[:, 1:, 1:] = d2
        v2 = -torch.linalg.det(cm) / 9216
        if v2[0].item() > 1e-20:
            vols.append(v2[0].sqrt().cpu())
    if len(vols) < 50: return float('nan')
    vt = torch.stack(vols)
    return (vt.std() / (vt.mean() + 1e-8)).item()

def eff_dim(x):
    x_c = x - x.mean(0, keepdim=True)
    n = min(512, x.shape[0])
    _, S, _ = torch.linalg.svd(x_c[:n].float(), full_matrices=False)
    p = S / S.sum()
    return p.pow(2).sum().reciprocal().item()

CLASS_NAMES = ['plane','auto','bird','cat','deer','dog','frog','horse','ship','truck']

model.eval()
bn = model.bottleneck

print("=" * 80)
print("CONSTELLATION BOTTLENECK — FULL ANALYSIS")
print(f"  Params: {sum(p.numel() for p in model.parameters()):,}")
print(f"  Bottleneck: {sum(p.numel() for p in bn.parameters()):,}")
print("=" * 80)

# Load test data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,)*3, (0.5,)*3),
])
test_ds = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_ds, batch_size=256, shuffle=False)
images_test, labels_test = next(iter(test_loader))
images_test = images_test.to(DEVICE)
labels_test = labels_test.to(DEVICE)


# ══════════════════════════════════════════════════════════════════
# TEST 1: BOTTLENECK DIAGNOSTICS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 1: Bottleneck Diagnostics")
print(f"{'━'*80}")

drift = bn.drift().detach()
home = F.normalize(bn.home, dim=-1).detach()
curr = F.normalize(bn.anchors, dim=-1).detach()
P, A, d = home.shape

print(f"  Patches: {P}, Anchors/patch: {A}, Patch dim: {d}")
print(f"  Drift: mean={drift.mean():.6f} rad ({math.degrees(drift.mean()):.2f}°)")
print(f"         std={drift.std():.6f}  min={drift.min():.6f}  max={drift.max():.6f}")
print(f"         max degrees: {math.degrees(drift.max()):.2f}°")
print(f"  Skip gate: {bn.skip_gate.sigmoid().item():.4f}")
print(f"  Near 0.29154: {(drift - 0.29154).abs().lt(0.05).float().mean().item():.1%}")

# Per-patch drift
print(f"\n  Per-patch drift:")
for p in range(P):
    d_p = drift[p].mean().item()
    d_max = drift[p].max().item()
    marker = " ◄ 0.29" if abs(d_p - 0.29154) < 0.05 else ""
    marker2 = " ◄ MAX near 0.29" if abs(d_max - 0.29154) < 0.05 else ""
    print(f"    P{p:2d}: mean={d_p:.4f} ({math.degrees(d_p):.1f}°) "
          f"max={d_max:.4f} ({math.degrees(d_max):.1f}°){marker}{marker2}")

# Anchor pairwise spread
print(f"\n  Anchor spread per patch:")
for p in range(min(8, P)):
    sim = (curr[p] @ curr[p].T)
    sim.fill_diagonal_(0)
    print(f"    P{p}: mean_cos={sim.mean():.4f} max={sim.max():.4f} min={sim.min():.4f}")

# Anchor effective dimensionality
print(f"\n  Anchor effective dimensionality:")
for p in range(min(8, P)):
    _, S, _ = torch.linalg.svd(curr[p].float(), full_matrices=False)
    pr = S / S.sum()
    ed = pr.pow(2).sum().reciprocal().item()
    print(f"    P{p}: eff_dim={ed:.1f} / {A}")

# Drift histogram — where do anchors cluster?
all_drifts = drift.flatten().cpu().numpy()
print(f"\n  Drift distribution:")
bins = [0.0, 0.05, 0.10, 0.15, 0.20, 0.25, 0.30, 0.35, 0.40]
hist, _ = np.histogram(all_drifts, bins=bins)
for i in range(len(bins)-1):
    bar = "█" * hist[i]
    print(f"    {bins[i]:.2f}-{bins[i+1]:.2f}: {hist[i]:3d} {bar}")


# ══════════════════════════════════════════════════════════════════
# TEST 2: SPHERE REPRESENTATION — CV OF BOTTLENECK EMBEDDINGS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 2: Sphere Representation — CV of bottleneck embeddings")
print(f"  These live on S^15. Does CV approach 0.20?")
print(f"{'━'*80}")

# Hook to capture sphere embeddings
sphere_embeddings = {}
tri_profiles = {}

def hook_sphere(module, input, output):
    # The forward method: proj_in → norm → reshape → normalize
    # We need to grab AFTER L2 norm. Hook the full bottleneck
    # and manually compute the sphere embedding.
    pass

# Manually extract sphere embeddings at different timesteps
print(f"\n  {'t':>6} {'CV_sphere':>10} {'CV_tri':>10} {'eff_d_sph':>10} "
      f"{'eff_d_tri':>10} {'sph_norm':>10}")

for t_val in [0.0, 0.1, 0.25, 0.5, 0.75, 0.9, 1.0]:
    B = images_test.shape[0]
    t = torch.full((B,), t_val, device=DEVICE)
    eps = torch.randn_like(images_test)
    t_b = t.view(B, 1, 1, 1)
    x_t = (1 - t_b) * images_test + t_b * eps

    with torch.no_grad():
        # Run encoder manually
        cond = model.time_emb(t) + model.class_emb(labels_test)
        h = model.in_conv(x_t)
        skips = [h]
        for i in range(len(model.channel_mults)):
            for block in model.enc[i]:
                if isinstance(block, nn.Sequential):
                    h = block[0](h); h = block[1](h, cond)
                else:
                    h = block(h, cond)
            skips.append(h)
            if i < len(model.enc_down):
                h = model.enc_down[i](h)

        # Get sphere embedding
        h_flat = h.reshape(B, -1)
        emb = bn.proj_in(h_flat)
        emb = bn.proj_in_norm(emb)
        patches = emb.reshape(B, bn.n_patches, bn.patch_dim)
        patches_n = F.normalize(patches, dim=-1)

        # CV of sphere embeddings (flatten patches back to one vector)
        sphere_flat = patches_n.reshape(B, -1)  # (B, 256) on product of spheres
        cv_sphere = compute_cv(sphere_flat, n_samples=1000)
        ed_sphere = eff_dim(sphere_flat)
        norm_sph = sphere_flat.norm(dim=-1).mean().item()

        # Triangulation profile
        tri = bn.triangulate(patches_n)  # (B, 768)
        cv_tri = compute_cv(tri, n_samples=1000)
        ed_tri = eff_dim(tri)

        # Per-patch CV
        if t_val == 0.0:
            print(f"\n  Per-patch CV at t=0 (should be ≈0.20 if d=16):")
            for p in range(min(8, bn.n_patches)):
                patch_p = patches_n[:, p, :]  # (B, 16) on S^15
                cv_p = compute_cv(patch_p, n_samples=1000)
                print(f"    Patch {p}: CV={cv_p:.4f}")
            print()

    print(f"  {t_val:>6.2f} {cv_sphere:>10.4f} {cv_tri:>10.4f} {ed_sphere:>10.1f} "
          f"{ed_tri:>10.1f} {norm_sph:>10.4f}")


# ══════════════════════════════════════════════════════════════════
# TEST 3: PER-CLASS ANCHOR ROUTING
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 3: Per-Class Anchor Routing")
print(f"{'━'*80}")

# Collect per-class nearest anchors across all patches
class_nearest = {c: [] for c in range(10)}
anchors_n = F.normalize(bn.anchors.detach(), dim=-1)

for images_b, labels_b in test_loader:
    images_b = images_b.to(DEVICE)
    labels_b = labels_b.to(DEVICE)
    B = images_b.shape[0]
    t = torch.zeros(B, device=DEVICE)  # clean images

    with torch.no_grad():
        cond = model.time_emb(t) + model.class_emb(labels_b)
        h = model.in_conv(images_b)
        for i in range(len(model.channel_mults)):
            for block in model.enc[i]:
                if isinstance(block, nn.Sequential):
                    h = block[0](h); h = block[1](h, cond)
                else:
                    h = block(h, cond)
            if i < len(model.enc_down):
                h = model.enc_down[i](h)

        h_flat = h.reshape(B, -1)
        emb = bn.proj_in_norm(bn.proj_in(h_flat))
        patches = F.normalize(emb.reshape(B, bn.n_patches, bn.patch_dim), dim=-1)

        # Nearest anchor per patch
        cos = torch.einsum('bpd,pad->bpa', patches, anchors_n)  # (B, P, A)
        nearest = cos.argmax(dim=-1)  # (B, P)

        for i in range(B):
            c = labels_b[i].item()
            class_nearest[c].append(nearest[i].cpu())

    if sum(len(v) for v in class_nearest.values()) > 5000:
        break

# Show routing for first 4 patches
for p_idx in range(min(4, bn.n_patches)):
    print(f"\n  Patch {p_idx} — nearest anchor per class:")
    print(f"  {'class':>10}", end="")
    for a in range(A):
        print(f" {a:>4}", end="")
    print()

    for c in range(10):
        if not class_nearest[c]:
            continue
        nearest_all = torch.stack(class_nearest[c])  # (N, P)
        nearest_p = nearest_all[:, p_idx]
        counts = torch.bincount(nearest_p, minlength=A).float()
        counts = counts / counts.sum()
        row = f"  {CLASS_NAMES[c]:>10}"
        for a in range(A):
            pct = counts[a].item()
            if pct > 0.15:
                row += f" {pct:>3.0%}█"
            elif pct > 0.05:
                row += f" {pct:>3.0%}░"
            else:
                row += f"  {pct:>3.0%}"
            #row += f" {pct:>3.0%}"
        print(row)

# Are anchor patterns class-specific?
print(f"\n  Anchor routing entropy per class (lower = more concentrated):")
for c in range(10):
    if not class_nearest[c]:
        continue
    nearest_all = torch.stack(class_nearest[c])
    # Average across patches
    total_entropy = 0
    for p_idx in range(bn.n_patches):
        counts = torch.bincount(nearest_all[:, p_idx], minlength=A).float()
        counts = counts / counts.sum()
        entropy = -(counts * (counts + 1e-8).log()).sum().item()
        total_entropy += entropy
    avg_entropy = total_entropy / bn.n_patches
    max_entropy = math.log(A)
    print(f"    {CLASS_NAMES[c]:>10}: H={avg_entropy:.3f} / {max_entropy:.3f} "
          f"({avg_entropy/max_entropy:.1%} of max)")


# ══════════════════════════════════════════════════════════════════
# TEST 4: SKIP GATE ANALYSIS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 4: Skip Gate — how much goes through constellation vs skip?")
print(f"{'━'*80}")

gate = bn.skip_gate.sigmoid().item()
print(f"  Skip gate value: {gate:.4f}")
print(f"  Skip path:          {gate:.1%}")
print(f"  Constellation path: {1-gate:.1%}")
print(f"  Skip proj params:   {sum(p.numel() for p in [bn.skip_proj.weight, bn.skip_proj.bias]):,}")
print(f"  Patchwork params:   {sum(p.numel() for p in bn.patchwork.parameters()):,}")
print(f"\n  ⚠ skip_proj is Linear(16384, 16384) = "
      f"{bn.skip_proj.weight.numel():,} params")
print(f"  ⚠ This single layer is {bn.skip_proj.weight.numel()/1e6:.0f}M params — "
      f"larger than the rest of the model combined")


# ══════════════════════════════════════════════════════════════════
# TEST 5: GENERATION — PER CLASS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 5: Generation Quality")
print(f"{'━'*80}")

print(f"  {'class':>10} {'intra_cos':>10} {'std':>8} {'CV':>8} {'norm':>8}")

all_gen = []
for c in range(10):
    imgs, _ = sample(model, 64, 50, class_label=c)
    imgs = (imgs + 1) / 2  # to [0,1]
    all_gen.append(imgs)

    flat = imgs.reshape(64, -1)
    flat_n = F.normalize(flat, dim=-1)
    sim = flat_n @ flat_n.T
    mask = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
    intra = sim[mask].mean().item()
    std = sim[mask].std().item()
    cv = compute_cv(flat, 500)
    norm = flat.norm(dim=-1).mean().item()
    print(f"  {CLASS_NAMES[c]:>10} {intra:>10.4f} {std:>8.4f} {cv:>8.4f} {norm:>8.2f}")

    save_image(make_grid(imgs[:16], nrow=4), f"analysis_bn/class_{CLASS_NAMES[c]}.png")

# All classes grid
all_grid = torch.cat([g[:4] for g in all_gen])
save_image(make_grid(all_grid, nrow=10), "analysis_bn/all_classes.png")


# ══════════════════════════════════════════════════════════════════
# TEST 6: ABLATION — SKIP ONLY vs CONSTELLATION ONLY
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 6: Ablation — Skip-only vs Constellation-only")
print(f"{'━'*80}")

original_gate = bn.skip_gate.data.clone()

# A) Full model (as trained)
torch.manual_seed(999)
with torch.no_grad():
    imgs_full, _ = sample(model, 32, 50, class_label=3)

# B) Skip only (gate → +100, sigmoid ≈ 1.0)
bn.skip_gate.data.fill_(100.0)
torch.manual_seed(999)
with torch.no_grad():
    imgs_skip, _ = sample(model, 32, 50, class_label=3)

# C) Constellation only (gate → -100, sigmoid ≈ 0.0)
bn.skip_gate.data.fill_(-100.0)
torch.manual_seed(999)
with torch.no_grad():
    imgs_const, _ = sample(model, 32, 50, class_label=3)

# Restore
bn.skip_gate.data.copy_(original_gate)

imgs_full_01 = (imgs_full + 1) / 2
imgs_skip_01 = (imgs_skip + 1) / 2
imgs_const_01 = (imgs_const + 1) / 2

# Compare
for name, imgs in [('skip_only', imgs_skip), ('const_only', imgs_const)]:
    delta = (imgs_full - imgs).abs()
    pixel_diff = delta.mean().item()
    cos = F.cosine_similarity(
        imgs_full.reshape(32, -1), imgs.reshape(32, -1)).mean().item()
    print(f"  {name:>15}: pixel_Δ={pixel_diff:.6f}  cos_sim={cos:.6f}  "
          f"max_Δ={delta.max():.4f}")

# Save comparison: top=full, mid=skip_only, bot=constellation_only
comparison = torch.cat([imgs_full_01[:8], imgs_skip_01[:8], imgs_const_01[:8]])
save_image(make_grid(comparison, nrow=8), "analysis_bn/ablation_skip_vs_const.png")
print(f"  ✓ Saved (top=full, mid=skip_only, bot=constellation_only)")


# ══════════════════════════════════════════════════════════════════
# TEST 7: VELOCITY FIELD
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 7: Velocity Field Quality")
print(f"{'━'*80}")

print(f"  {'t':>6} {'v_norm':>10} {'v·target':>10} {'mse':>10}")

for t_val in [0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95]:
    B = 128
    imgs_v = images_test[:B]
    labs_v = labels_test[:B]
    t = torch.full((B,), t_val, device=DEVICE)
    eps = torch.randn_like(imgs_v)
    t_b = t.view(B, 1, 1, 1)
    x_t = (1 - t_b) * imgs_v + t_b * eps
    v_target = eps - imgs_v

    with torch.no_grad():
        v_pred = model(x_t, t, labs_v)

    v_norm = v_pred.reshape(B, -1).norm(dim=-1).mean().item()
    v_cos = F.cosine_similarity(
        v_pred.reshape(B, -1), v_target.reshape(B, -1)).mean().item()
    mse = F.mse_loss(v_pred, v_target).item()
    print(f"  {t_val:>6.2f} {v_norm:>10.2f} {v_cos:>10.4f} {mse:>10.4f}")


# ══════════════════════════════════════════════════════════════════
# TEST 8: ODE TRAJECTORY — CV THROUGH GENERATION
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 8: ODE Trajectory — geometry through generation")
print(f"{'━'*80}")

n_steps = 50
B_traj = 256
x = torch.randn(B_traj, 3, 32, 32, device=DEVICE)
labels_traj = torch.randint(0, 10, (B_traj,), device=DEVICE)
dt = 1.0 / n_steps

print(f"  {'step':>6} {'t':>6} {'x_norm':>10} {'x_std':>10} {'CV':>8}")

for step in range(n_steps):
    t_val = 1.0 - step * dt
    t = torch.full((B_traj,), t_val, device=DEVICE)
    with torch.no_grad(), torch.amp.autocast("cuda", dtype=torch.bfloat16):
        v = model(x, t, labels_traj)
    x = x - v.float() * dt

    if step in [0, 1, 5, 10, 20, 30, 40, 49]:
        xf = x.reshape(B_traj, -1)
        print(f"  {step:>6} {t_val:>6.2f} {xf.norm(dim=-1).mean().item():>10.2f} "
              f"{x.std().item():>10.4f} {compute_cv(xf, 500):>8.4f}")


# ══════════════════════════════════════════════════════════════════
# TEST 9: INTER vs INTRA CLASS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 9: Inter vs Intra Class Separation")
print(f"{'━'*80}")

intra_sims = []
inter_sims = []
for c in range(10):
    flat = F.normalize(all_gen[c].reshape(64, -1), dim=-1)
    sim = flat @ flat.T
    mask = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
    intra_sims.append(sim[mask].mean().item())

for i in range(10):
    for j in range(i+1, 10):
        fi = F.normalize(all_gen[i].reshape(64, -1), dim=-1)
        fj = F.normalize(all_gen[j].reshape(64, -1), dim=-1)
        inter_sims.append((fi @ fj.T).mean().item())

print(f"  Intra-class cos: {np.mean(intra_sims):.4f} ± {np.std(intra_sims):.4f}")
print(f"  Inter-class cos: {np.mean(inter_sims):.4f} ± {np.std(inter_sims):.4f}")
ratio = np.mean(intra_sims) / (np.mean(inter_sims) + 1e-8)
print(f"  Separation ratio: {ratio:.3f}×")


# ══════════════════════════════════════════════════════════════════
# SUMMARY
# ══════════════════════════════════════════════════════════════════

print(f"\n{'='*80}")
print("ANALYSIS COMPLETE")
print(f"{'='*80}")
print(f"""
  Files in analysis_bn/:
    class_*.png              per-class samples
    all_classes.png          4 per class grid
    ablation_skip_vs_const.png  top=full, mid=skip, bot=constellation

  Key questions answered:
    1. Does per-patch CV ≈ 0.20? (Test 2)
       → If yes, the bottleneck lives at the natural S^15 dimension
    2. Is anchor routing class-specific? (Test 3)
       → If entropy varies by class, constellation routes differently
    3. Does the skip path dominate? (Tests 4 & 6)
       → If skip_only ≈ full, the 268M skip_proj IS the model
    4. Does constellation-only work at all? (Test 6)
       → The real test of whether geometric encoding carries signal
""")
print("=" * 80)