#!/usr/bin/env python3
"""
Flow Match Relay — Full Analysis Toolkit
==========================================
Run after training. Analyzes:

  1. Relay diagnostics: drift, gates, anchor geometry
  2. CV measurement through the network at each layer
  3. Anchor utilization: which anchors are active per class?
  4. Generation quality: FID prep, per-class diversity
  5. The 0.29154 hunt: does drift converge to the binding constant?
  6. Feature map geometry: CV of bottleneck features
  7. Velocity field analysis: how does the relay affect v_pred?
  8. Gate dynamics: measure gate values at different timesteps
  9. Anchor constellation visualization
  10. Ablation: relay ON vs OFF generation comparison
"""

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

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
torch.manual_seed(42)

os.makedirs("analysis", exist_ok=True)


def compute_cv(points, n_samples=2000, 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, 10000), 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']

print("=" * 80)
print("FLOW MATCH RELAY — FULL ANALYSIS TOOLKIT")
print(f"  Device: {DEVICE}")
print("=" * 80)

# ── Load model ──
from transformers import AutoModel

model = AutoModel.from_pretrained(
    "AbstractPhil/geolip-diffusion-proto", trust_remote_code=True
).to(DEVICE)
model.eval()

n_params = sum(p.numel() for p in model.parameters())
n_relay = sum(p.numel() for n, p in model.named_parameters() if 'relay' in n)
print(f"  Params: {n_params:,} (relay: {n_relay:,}, {100*n_relay/n_params:.1f}%)")

# Find relay modules
relays = {}
for name, module in model.named_modules():
    if hasattr(module, 'drift') and hasattr(module, 'anchors'):
        relays[name] = module
print(f"  Relay modules: {len(relays)}")


# ══════════════════════════════════════════════════════════════════
# TEST 1: RELAY DIAGNOSTICS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 1: Relay Diagnostics — Drift, Gates, Anchor Geometry")
print(f"{'━'*80}")

for name, relay in relays.items():
    drift = relay.drift().detach().cpu()  # (P, A)
    gates = relay.gates.sigmoid().detach().cpu()  # (P,)
    home = F.normalize(relay.home, dim=-1).detach().cpu()
    anchors = F.normalize(relay.anchors, dim=-1).detach().cpu()

    P, A, d = home.shape

    print(f"\n  {name}:")
    print(f"    Patches: {P}, Anchors/patch: {A}, Patch dim: {d}")
    print(f"    Drift (rad):  mean={drift.mean():.6f}  std={drift.std():.6f}  "
          f"min={drift.min():.6f}  max={drift.max():.6f}")
    print(f"    Drift (deg):  mean={math.degrees(drift.mean()):.2f}°  "
          f"max={math.degrees(drift.max()):.2f}°")
    print(f"    Gates:        mean={gates.mean():.4f}  std={gates.std():.4f}  "
          f"min={gates.min():.4f}  max={gates.max():.4f}")

    # Anchor pairwise similarity within each patch
    for p in range(min(4, P)):
        sim = (anchors[p] @ anchors[p].T)
        sim.fill_diagonal_(0)
        print(f"    Patch {p}: anchor_cos mean={sim.mean():.4f} max={sim.max():.4f} "
              f"min={sim.min():.4f}")

    # Near 0.29154?
    near_029 = (drift - 0.29154).abs() < 0.05
    pct_near = near_029.float().mean().item()
    print(f"    Near 0.29154: {pct_near:.1%} of anchors within ±0.05")

    # Per-patch drift
    print(f"    Per-patch mean drift:")
    for p in range(P):
        d_p = drift[p].mean().item()
        marker = " ◄ 0.29" if abs(d_p - 0.29154) < 0.05 else ""
        print(f"      Patch {p:2d}: {d_p:.6f} rad ({math.degrees(d_p):.2f}°){marker}")


# ══════════════════════════════════════════════════════════════════
# TEST 2: BOTTLENECK FEATURE GEOMETRY
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 2: Bottleneck Feature Geometry — CV at the relay point")
print(f"{'━'*80}")

# Load some real data
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
test_ds = datasets.CIFAR10('./data', train=False, download=True, transform=transform)
test_loader = torch.utils.data.DataLoader(test_ds, batch_size=256, shuffle=False)

# Hook to capture bottleneck features
bottleneck_features = {}

def hook_fn(name):
    def fn(module, input, output):
        if isinstance(output, torch.Tensor):
            bottleneck_features[name] = output.detach()
    return fn

# Register hooks ONLY on top-level mid blocks and relay modules (not submodules)
hooks = []
target_names = set(relays.keys()) | {'unet.mid_block1', 'unet.mid_block2', 'unet.mid_attn'}
for name, module in model.named_modules():
    if name in target_names:
        hooks.append(module.register_forward_hook(hook_fn(name)))

# Run a batch through at several timesteps
images, labels = next(iter(test_loader))
images = images.to(DEVICE)
labels_dev = labels.to(DEVICE)

print(f"\n  CV of bottleneck features at different timesteps:")
print(f"  {'t':>6} {'module':>40} {'CV':>8} {'eff_d':>8} {'norm':>8}")

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

    bottleneck_features.clear()
    with torch.no_grad():
        _ = model(x_t, t, labels_dev)

    for feat_name, feat in bottleneck_features.items():
        if feat.dim() == 4:
            # Feature map: pool spatial → (B, C)
            pooled = feat.mean(dim=(-2, -1))
        elif feat.dim() == 2:
            pooled = feat
        else:
            continue  # skip 1D or other odd shapes
        if pooled.dim() != 2 or pooled.shape[0] < 5 or pooled.shape[1] < 5:
            continue
        cv = compute_cv(pooled, n_samples=1000)
        ed = eff_dim(pooled)
        norm_mean = pooled.norm(dim=-1).mean().item()
        print(f"  {t_val:>6.2f} {feat_name:>40} {cv:>8.4f} {ed:>8.1f} {norm_mean:>8.2f}")

# Clean up hooks
for h in hooks:
    h.remove()


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

print(f"\n{'━'*80}")
print("TEST 3: Per-Class Anchor Utilization")
print(f"  Which anchors activate for each class?")
print(f"{'━'*80}")

# Collect bottleneck features per class
class_features = {c: [] for c in range(10)}

for images_batch, labels_batch in test_loader:
    images_batch = images_batch.to(DEVICE)
    labels_batch = labels_batch.to(DEVICE)
    B = images_batch.shape[0]

    t = torch.full((B,), 0.0, device=DEVICE)  # clean images (t=0)

    # Get features before relay
    bottleneck_features.clear()
    relay_name = list(relays.keys())[0]
    relay_mod = relays[relay_name]
    hook = relay_mod.register_forward_hook(hook_fn(relay_name))

    with torch.no_grad():
        _ = model(images_batch, t, labels_batch)

    hook.remove()

    if relay_name in bottleneck_features:
        feat = bottleneck_features[relay_name]
        if feat.dim() == 4:
            pooled = feat.mean(dim=(-2, -1))  # (B, C)
        else:
            pooled = feat
        for i in range(B):
            c = labels_batch[i].item()
            class_features[c].append(pooled[i].cpu())

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

# For each class, triangulate against the first relay's anchors
relay_mod = list(relays.values())[0]
anchors = F.normalize(relay_mod.anchors.detach(), dim=-1)  # (P, A, d)
P, A, d = anchors.shape

print(f"\n  Nearest anchor distribution per class (Patch 0):")
print(f"  {'class':>10}", end="")
for a in range(A):
    print(f" {a:>5}", end="")
print()

for c in range(10):
    if not class_features[c]:
        continue
    feats = torch.stack(class_features[c]).to(DEVICE)  # (N, C)
    # Chunk into patches
    patches = feats.reshape(-1, P, d)
    patch0 = F.normalize(patches[:, 0], dim=-1)  # (N, d)
    # Find nearest anchor
    cos = patch0 @ anchors[0].T  # (N, A)
    nearest = cos.argmax(dim=-1)  # (N,)
    counts = torch.bincount(nearest, minlength=A).float()
    counts = counts / counts.sum()
    row = f"  {CLASS_NAMES[c]:>10}"
    for a in range(A):
        pct = counts[a].item()
        marker = "█" if pct > 0.15 else "▓" if pct > 0.10 else "░" if pct > 0.05 else " "
        row += f" {pct:>4.0%}{marker}"
    print(row)


# ══════════════════════════════════════════════════════════════════
# TEST 4: GATE DYNAMICS ACROSS TIMESTEPS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 4: Gate Dynamics — do relay gates respond to timestep?")
print(f"{'━'*80}")

# The gates are parameters (not input-dependent), so they're constant.
# But we can measure the relay's EFFECTIVE contribution at each t.
print(f"  Note: gates are learned parameters, not t-dependent.")
print(f"  Measuring relay output magnitude at different t instead.\n")

relay_name = list(relays.keys())[0]
relay_mod = relays[relay_name]

relay_in = {}
relay_out = {}

def hook_in(module, input, output):
    if isinstance(input, tuple):
        relay_in['x'] = input[0].detach()
    else:
        relay_in['x'] = input.detach()
    relay_out['x'] = output.detach()

hook = relay_mod.register_forward_hook(hook_in)

images_small = images[:64]
labels_small = labels_dev[:64]

print(f"  {'t':>6} {'relay_Δ_norm':>14} {'relay_Δ_cos':>14} {'input_norm':>12} {'output_norm':>12}")

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

    relay_in.clear(); relay_out.clear()
    with torch.no_grad():
        _ = model(x_t, t, labels_small)

    if 'x' in relay_in and 'x' in relay_out:
        x_in = relay_in['x']
        x_out = relay_out['x']
        delta = (x_out - x_in)
        # Flatten everything beyond batch dim for norm
        delta_flat = delta.reshape(delta.shape[0], -1)
        in_flat = x_in.reshape(x_in.shape[0], -1)
        out_flat = x_out.reshape(x_out.shape[0], -1)
        delta_norm = delta_flat.norm(dim=-1).mean().item()
        in_norm = in_flat.norm(dim=-1).mean().item()
        out_norm = out_flat.norm(dim=-1).mean().item()

        cos_change = 1 - F.cosine_similarity(in_flat, out_flat).mean().item()
        print(f"  {t_val:>6.2f} {delta_norm:>14.4f} {cos_change:>14.8f} "
              f"{in_norm:>12.2f} {out_norm:>12.2f}")

hook.remove()


# ══════════════════════════════════════════════════════════════════
# TEST 5: GENERATION QUALITY — PER-CLASS DIVERSITY
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 5: Generation Quality — Per-Class Diversity")
print(f"{'━'*80}")

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

all_generated = []
for c in range(10):
    with torch.no_grad():
        imgs = model.sample(n_samples=64, class_label=c)  # (64, 3, 32, 32) in [0,1]
    all_generated.append(imgs)

    flat = imgs.reshape(64, -1)  # (64, 3072)
    flat_n = F.normalize(flat, dim=-1)

    # Intra-class cosine similarity
    sim = flat_n @ flat_n.T
    mask = ~torch.eye(64, device=DEVICE, dtype=torch.bool)
    intra_cos = sim[mask].mean().item()
    intra_std = sim[mask].std().item()

    cv = compute_cv(flat, n_samples=500)
    norm_mean = flat.norm(dim=-1).mean().item()

    print(f"  {CLASS_NAMES[c]:>10} {intra_cos:>10.4f} {intra_std:>10.4f} "
          f"{cv:>8.4f} {norm_mean:>8.2f}")

# Save per-class grid
for c in range(10):
    grid = make_grid(all_generated[c][:16], nrow=4)
    save_image(grid, f"analysis/class_{CLASS_NAMES[c]}.png")

# All classes grid
all_grid = torch.cat([imgs[:4] for imgs in all_generated])
save_image(make_grid(all_grid, nrow=10), "analysis/all_classes.png")
print(f"\n  ✓ Saved per-class grids to analysis/")


# ══════════════════════════════════════════════════════════════════
# TEST 6: VELOCITY FIELD ANALYSIS
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 6: Velocity Field — how does v_pred behave across t?")
print(f"{'━'*80}")

images_v = images[:128]
labels_v = labels_dev[:128]

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

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

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

    v_norm = v_pred.reshape(128, -1).norm(dim=-1).mean().item()
    v_std = v_pred.std().item()
    # Cosine between predicted and target velocity
    v_cos = F.cosine_similarity(
        v_pred.reshape(128, -1), v_target.reshape(128, -1)).mean().item()
    # MSE
    mse = F.mse_loss(v_pred, v_target).item()

    print(f"  {t_val:>6.2f} {v_norm:>10.2f} {v_std:>10.4f} "
          f"{v_cos:>10.4f} {mse:>10.4f}")


# ══════════════════════════════════════════════════════════════════
# TEST 7: ABLATION — RELAY ON vs OFF
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 7: Ablation — Relay ON vs OFF during generation")
print(f"  Disable relay gates, measure generation difference")
print(f"{'━'*80}")

# Save original gate values
original_gates = {}
for name, relay in relays.items():
    original_gates[name] = relay.gates.data.clone()

# Generate with relay ON
torch.manual_seed(123)
with torch.no_grad():
    imgs_on = model.sample(n_samples=32, class_label=3)

# Disable relays (set gates to -100 → sigmoid ≈ 0)
for name, relay in relays.items():
    relay.gates.data.fill_(-100.0)

# Generate with relay OFF (same seed)
torch.manual_seed(123)
with torch.no_grad():
    imgs_off = model.sample(n_samples=32, class_label=3)

# Restore gates
for name, relay in relays.items():
    relay.gates.data.copy_(original_gates[name])

# Compare
delta = (imgs_on - imgs_off)
pixel_diff = delta.abs().mean().item()
cos_diff = F.cosine_similarity(
    imgs_on.reshape(32, -1), imgs_off.reshape(32, -1)).mean().item()

print(f"  Relay ON  — mean pixel: {imgs_on.mean():.4f}  std: {imgs_on.std():.4f}")
print(f"  Relay OFF — mean pixel: {imgs_off.mean():.4f}  std: {imgs_off.std():.4f}")
print(f"  Pixel diff:    {pixel_diff:.6f}")
print(f"  Cosine sim:    {cos_diff:.6f}")
print(f"  Max pixel Δ:   {delta.abs().max():.6f}")

# Save comparison
comparison = torch.cat([imgs_on[:8], imgs_off[:8]], dim=0)
save_image(make_grid(comparison, nrow=8), "analysis/relay_ablation.png")
print(f"  ✓ Saved analysis/relay_ablation.png (top=ON, bottom=OFF)")


# ══════════════════════════════════════════════════════════════════
# TEST 8: ANCHOR CONSTELLATION STRUCTURE
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 8: Anchor Constellation Structure")
print(f"{'━'*80}")

for name, relay in relays.items():
    home = F.normalize(relay.home.detach().cpu(), dim=-1)
    curr = F.normalize(relay.anchors.detach().cpu(), dim=-1)
    P, A, d = home.shape

    print(f"\n  {name}:")

    # Home vs current — did training move them?
    home_curr_cos = (home * curr).sum(dim=-1)  # (P, A)
    print(f"    Home↔Current cos: mean={home_curr_cos.mean():.6f}  "
          f"min={home_curr_cos.min():.6f}")

    # Anchor spread — how well-distributed?
    for p in range(min(4, P)):
        cos_matrix = curr[p] @ curr[p].T  # (A, A)
        cos_matrix.fill_diagonal_(0)
        print(f"    Patch {p} anchor spread: "
              f"mean_cos={cos_matrix.mean():.4f}  "
              f"max_cos={cos_matrix.max():.4f}  "
              f"min_cos={cos_matrix.min():.4f}")

    # Effective anchor dimensionality
    for p in range(min(4, P)):
        _, S, _ = torch.linalg.svd(curr[p].float(), full_matrices=False)
        pr = S / S.sum()
        anchor_eff_dim = pr.pow(2).sum().reciprocal().item()
        print(f"    Patch {p} anchor eff_dim: {anchor_eff_dim:.1f} / {A}")


# ══════════════════════════════════════════════════════════════════
# TEST 9: SAMPLING TRAJECTORY — TRACK CV THROUGH ODE
# ══════════════════════════════════════════════════════════════════

print(f"\n{'━'*80}")
print("TEST 9: Sampling Trajectory — CV through ODE steps")
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_pixel':>10}")

checkpoints = [0, 1, 5, 10, 20, 30, 40, 49]
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 checkpoints:
        x_flat = x.reshape(B_traj, -1)
        norm = x_flat.norm(dim=-1).mean().item()
        std = x.std().item()
        cv = compute_cv(x_flat, n_samples=500)
        print(f"  {step:>6} {t_val:>6.2f} {norm:>10.2f} {std:>10.4f} {cv:>10.4f}")


# ══════════════════════════════════════════════════════════════════
# TEST 10: INTER-CLASS vs INTRA-CLASS GEOMETRY
# ══════════════════════════════════════════════════════════════════

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

# Use generated images
class_means = []
for c in range(10):
    flat = all_generated[c].reshape(64, -1)
    class_means.append(F.normalize(flat.mean(dim=0, keepdim=True), dim=-1))

class_means = torch.cat(class_means, dim=0)  # (10, 3072)
inter_sim = class_means @ class_means.T

print(f"  Inter-class cosine similarity matrix:")
print(f"  {'':>8}", end="")
for c in range(10):
    print(f" {CLASS_NAMES[c][:4]:>5}", end="")
print()

for i in range(10):
    print(f"  {CLASS_NAMES[i]:>8}", end="")
    for j in range(10):
        val = inter_sim[i, j].item()
        if i == j:
            print(f"   1.0", end="")
        else:
            print(f" {val:>5.2f}", end="")
    print()

# Intra vs inter
intra_sims = []
inter_sims = []
for c in range(10):
    flat = F.normalize(all_generated[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):
        flat_i = F.normalize(all_generated[i].reshape(64, -1), dim=-1)
        flat_j = F.normalize(all_generated[j].reshape(64, -1), dim=-1)
        cross = (flat_i @ flat_j.T).mean().item()
        inter_sims.append(cross)

print(f"\n  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}")
print(f"  Separation ratio: {np.mean(intra_sims) / (np.mean(inter_sims) + 1e-8):.2f}×")


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

print(f"\n{'='*80}")
print("ANALYSIS COMPLETE")
print(f"{'='*80}")
print(f"""
  Files saved to analysis/:
    - class_*.png:           per-class generated samples
    - all_classes.png:       4 samples per class, 10 columns
    - relay_ablation.png:    relay ON (top) vs OFF (bottom)

  Key metrics to look for:
    1. Anchor drift → did any converge near 0.29154?
    2. Gate values → did they learn to open from init (0.047)?
    3. Per-class anchor utilization → class-specific routing?
    4. Relay ablation → does turning off the relay change generation?
    5. Intra/inter-class ratio → > 1.0 means classes are separable
    6. Velocity cosine → higher = better flow matching
    7. CV through ODE → how does geometry evolve during generation?
""")
print(f"{'='*80}")