#!/usr/bin/env python3
"""
Constellation Relay v2 — Vectorized
=====================================
All patches triangulated in a single batched matmul.
Gates init near zero (relay off at init, learns to open).

Architecture:
  input (B, D) → LN → reshape (B, P, d)
  → L2 normalize each patch to S^(d-1)
  → batched tri: (B, P, d) × (P, d, A*phases) → (B, P, A*phases)
  → patchwork MLP per patch (shared or independent)
  → gated residual (gate init ≈ 0.05)
  → reshape (B, D) + skip connection
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
import time

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

HAS_FP8 = hasattr(torch, 'float8_e4m3fn')


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)
    _, S, _ = torch.linalg.svd(x_c[:512].float(), full_matrices=False)
    p = S / S.sum()
    return p.pow(2).sum().reciprocal().item()


def uniform_sphere(n, d):
    return F.normalize(torch.randn(n, d), dim=-1)


def uniform_hypersphere_init(n, d):
    if n <= d:
        M = torch.randn(d, n)
        Q, _ = torch.linalg.qr(M)
        return Q.T.contiguous()
    else:
        M = torch.randn(d, d)
        Q, _ = torch.linalg.qr(M)
        basis = Q.T
        extra = F.normalize(torch.randn(n - d, d), dim=-1)
        vecs = torch.cat([basis, extra], dim=0)
        for _ in range(200):
            sim = vecs @ vecs.T
            sim.fill_diagonal_(-2.0)
            nn_idx = sim.argmax(dim=1)
            vecs = F.normalize(vecs - 0.05 * vecs[nn_idx], dim=-1)
        return vecs


# ══════════════════════════════════════════════════════════════════
# VECTORIZED CONSTELLATION RELAY
# ══════════════════════════════════════════════════════════════════

class ConstellationRelay(nn.Module):
    """
    Fully vectorized constellation relay.

    All P patches triangulated in one einsum.
    Gates init near zero — relay is transparent at init.
    """
    def __init__(
        self,
        input_dim,
        patch_dim=16,
        n_anchors=16,
        n_phases=3,
        pw_hidden=32,
        gate_init=-3.0,      # sigmoid(-3) ≈ 0.047
    ):
        super().__init__()
        assert input_dim % patch_dim == 0
        self.input_dim = input_dim
        self.patch_dim = patch_dim
        self.n_patches = input_dim // patch_dim
        self.n_anchors = n_anchors
        self.n_phases = n_phases

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

        # All anchors: (P, A, d) — Xavier init, then normalize to sphere
        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())

        # Patchwork: shared architecture, independent weights per patch
        # Input: n_phases * n_anchors per patch
        tri_dim = n_phases * A
        self.pw_w1 = nn.Parameter(torch.empty(P, tri_dim, pw_hidden))
        self.pw_b1 = nn.Parameter(torch.zeros(1, P, pw_hidden))
        self.pw_w2 = nn.Parameter(torch.empty(P, pw_hidden, d))
        self.pw_b2 = nn.Parameter(torch.zeros(1, P, d))
        # Xavier init each patch's weights
        for p in range(P):
            nn.init.xavier_normal_(self.pw_w1.data[p])
            nn.init.xavier_normal_(self.pw_w2.data[p])
        self.pw_norm = nn.LayerNorm(d)

        # Gates: init near zero so relay is transparent
        self.gates = nn.Parameter(torch.full((P,), gate_init))

        self.norm = nn.LayerNorm(input_dim)

    def drift(self):
        """Per-anchor drift: (P, A) in radians."""
        h = F.normalize(self.home, dim=-1)
        c = F.normalize(self.anchors, dim=-1)
        cos = (h * c).sum(dim=-1).clamp(-1 + 1e-7, 1 - 1e-7)
        return torch.acos(cos)

    def at_phase(self, t):
        """Slerp all anchors: (P, A, d)."""
        h = F.normalize(self.home, dim=-1)
        c = F.normalize(self.anchors, dim=-1)
        omega = self.drift().unsqueeze(-1)  # (P, A, 1)
        sin_omega = omega.sin().clamp(min=1e-7)
        return (torch.sin((1 - t) * omega) / sin_omega * h +
                torch.sin(t * omega) / sin_omega * c)

    def _fp8_or_native(self, patches_n, anchors_n):
        """
        Batched triangulation: patches (B, P, d) × anchors (P, A, d) → (B, P, A)
        Uses fp8 if available, otherwise native.
        """
        if HAS_FP8:
            # Per-tensor scaling for fp8
            amax_p = patches_n.abs().amax().clamp(min=1e-12)
            amax_a = anchors_n.abs().amax().clamp(min=1e-12)
            sp = torch.finfo(torch.float8_e4m3fn).max / amax_p
            sa = torch.finfo(torch.float8_e4m3fn).max / amax_a
            p8 = (patches_n * sp).to(torch.float8_e4m3fn).to(patches_n.dtype) / sp
            a8 = (anchors_n * sa).to(torch.float8_e4m3fn).to(anchors_n.dtype) / sa
            # einsum: (B, P, d) × (P, A, d) → (B, P, A)
            return torch.einsum('bpd,pad->bpa', p8, a8)
        else:
            return torch.einsum('bpd,pad->bpa', patches_n, anchors_n)

    def triangulate_multiphase(self, patches_n):
        """
        Multi-phase stroboscope, fully vectorized.
        patches_n: (B, P, d) normalized
        Returns: (B, P, A * n_phases)
        """
        phases = torch.linspace(0, 1, self.n_phases).tolist()
        tris = []
        for t in phases:
            anchors_t = F.normalize(self.at_phase(t), dim=-1)  # (P, A, d)
            cos = self._fp8_or_native(patches_n, anchors_t)
            tris.append(1.0 - cos)
        return torch.cat(tris, dim=-1)  # (B, P, A*phases)

    def forward(self, x, return_diagnostics=False):
        B, D = x.shape
        P, A, d = self.n_patches, self.n_anchors, self.patch_dim

        x_n = self.norm(x)
        patches = x_n.reshape(B, P, d)
        patches_n = F.normalize(patches, dim=-1)  # each patch on S^(d-1)

        # Vectorized multi-phase triangulation
        tri = self.triangulate_multiphase(patches_n)  # (B, P, A*phases)

        # Vectorized patchwork: (B, P, tri_dim) → (B, P, d)
        # pw_w1: (P, tri_dim, hidden), pw_w2: (P, hidden, d)
        h = torch.einsum('bpt,pth->bph', tri, self.pw_w1) + self.pw_b1
        h = F.gelu(h)
        pw_out = torch.einsum('bph,phd->bpd', h, self.pw_w2) + self.pw_b2
        pw_out = self.pw_norm(pw_out)

        # Gated residual per patch: gate ≈ 0.05 at init
        gate = self.gates.sigmoid().unsqueeze(0).unsqueeze(-1)  # (1, P, 1)
        blended = gate * pw_out + (1 - gate) * patches

        out = blended.reshape(B, D)
        result = x + out  # global skip

        if return_diagnostics:
            # fp64 addressing for diagnostics
            anchors_n = F.normalize(self.anchors, dim=-1)
            cos_all = torch.einsum('bpd,pad->bpa', patches_n, anchors_n)
            cos64 = cos_all.to(torch.float64)
            top2 = cos64.topk(2, dim=-1)
            nearest_cos = top2.values[:, :, 0].float().mean().item()
            margin = (top2.values[:, :, 0] - top2.values[:, :, 1]).float().mean().item()
            drift = self.drift()
            diag = {
                'nearest_cos': nearest_cos,
                'margin': margin,
                'drift_mean': drift.mean().item(),
                'drift_max': drift.max().item(),
                'gate_mean': self.gates.sigmoid().mean().item(),
                'gate_min': self.gates.sigmoid().min().item(),
                'gate_max': self.gates.sigmoid().max().item(),
            }
            return result, diag
        return result

    def get_diagnostics(self):
        drift = self.drift().detach().cpu()
        gates = self.gates.sigmoid().detach().cpu()
        return {
            'drift_mean': drift.mean().item(),
            'drift_std': drift.std().item(),
            'gate_mean': gates.mean().item(),
            'gate_std': gates.std().item(),
        }


# ══════════════════════════════════════════════════════════════════
# SELF-ATTENTION BLOCK (comparison)
# ══════════════════════════════════════════════════════════════════

class SelfAttnBlock(nn.Module):
    def __init__(self, dim, n_heads=4):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = dim // n_heads
        self.qkv = nn.Linear(dim, 3 * dim, bias=False)
        self.out_proj = nn.Linear(dim, dim, bias=False)
        self.norm = nn.LayerNorm(dim)

    def forward(self, x):
        B, S, D = x.shape
        x_n = self.norm(x)
        qkv = self.qkv(x_n).reshape(B, S, 3, self.n_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)
        q, k, v = qkv[0], qkv[1], qkv[2]
        attn = (q @ k.transpose(-2, -1)) * (self.head_dim ** -0.5)
        attn = attn.softmax(dim=-1)
        out = (attn @ v).transpose(1, 2).reshape(B, S, D)
        return x + self.out_proj(out)


# ══════════════════════════════════════════════════════════════════
# TESTS
# ══════════════════════════════════════════════════════════════════

N = 2000
D = 128
N_CV = 2000

print("=" * 80)
print("CONSTELLATION RELAY v2 — VECTORIZED + COLD GATES")
print(f"  Input dim: {D} = {D//16} patches × 16")
print(f"  Device: {DEVICE}")
print(f"  fp8: {'hardware' if HAS_FP8 else 'simulated'}")
print("=" * 80)


# ── TEST 1: Single pass ──
print(f"\n{'━'*80}")
print("TEST 1: Single relay pass — CV preservation + gate state")
print(f"{'━'*80}")

relay = ConstellationRelay(D, patch_dim=16, n_anchors=16, n_phases=3, pw_hidden=32).to(DEVICE)
n_params_relay = sum(p.numel() for p in relay.parameters())
print(f"  Relay params: {n_params_relay:,}")
print(f"  Gate init: sigmoid({relay.gates[0].item():.1f}) = {relay.gates[0].sigmoid().item():.4f}")

pts = uniform_sphere(N, D).to(DEVICE)
cv_before = compute_cv(pts, n_samples=N_CV)
ed_before = eff_dim(pts)

with torch.no_grad():
    out, diag = relay(pts, return_diagnostics=True)

out_n = F.normalize(out, dim=-1)
cv_after = compute_cv(out, n_samples=N_CV)
cv_normed = compute_cv(out_n, n_samples=N_CV)
ed_after = eff_dim(out)
cos_orig = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()

print(f"  Before:  CV={cv_before:.4f}  eff_dim={ed_before:.1f}")
print(f"  After:   CV={cv_after:.4f}  eff_dim={ed_after:.1f}")
print(f"  Normed:  CV={cv_normed:.4f}")
print(f"  cos_to_orig: {cos_orig:.6f}")
print(f"  Diagnostics: cos={diag['nearest_cos']:.4f} margin={diag['margin']:.4f} "
      f"drift={diag['drift_mean']:.4f} gate={diag['gate_mean']:.4f}")


# ── TEST 2: Repeated cycles — relay vs attention vs interleaved ──
print(f"\n{'━'*80}")
print("TEST 2: Depth stability — relay vs attention vs interleaved")
print(f"{'━'*80}")

pts = uniform_sphere(N, D).to(DEVICE)

# A) Relay only — 16 cycles
print(f"\n  A) Relay-only:")
print(f"  {'cycle':>6} {'CV':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10} {'gate':>6}")

relay_stack = nn.ModuleList([
    ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) for _ in range(16)])

x = pts.clone()
for i, r in enumerate(relay_stack):
    with torch.no_grad():
        x = r(x)
    if (i + 1) in [1, 2, 4, 8, 12, 16]:
        cv = compute_cv(x, N_CV)
        cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
        ed = eff_dim(x)
        cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
        g = r.gates.sigmoid().mean().item()
        print(f"  {i+1:>6} {cv:>8.4f} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f} {g:>6.3f}")

# B) Attention only — 16 cycles
print(f"\n  B) Attention-only:")
print(f"  {'cycle':>6} {'CV':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}")

attn_stack = nn.ModuleList([
    SelfAttnBlock(D, n_heads=8).to(DEVICE) for _ in range(16)])

x = pts.clone()
for i, layer in enumerate(attn_stack):
    with torch.no_grad():
        x = layer(x.unsqueeze(0)).squeeze(0)
    if (i + 1) in [1, 2, 4, 8, 12, 16]:
        cv = compute_cv(x, N_CV)
        cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
        ed = eff_dim(x)
        cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
        print(f"  {i+1:>6} {cv:>8.4f} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")

# C) Interleaved: attn → relay → attn → relay
print(f"\n  C) Interleaved (attn → relay → ...):")
print(f"  {'step':>6} {'type':>6} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}")

attn_il = nn.ModuleList([SelfAttnBlock(D, 8).to(DEVICE) for _ in range(8)])
relay_il = nn.ModuleList([ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) for _ in range(8)])

x = pts.clone()
step = 0
for i in range(8):
    with torch.no_grad():
        x = attn_il[i](x.unsqueeze(0)).squeeze(0)
    step += 1
    if step in [1, 2, 4, 8, 12, 16]:
        cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
        ed = eff_dim(x)
        cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
        print(f"  {step:>6} {'attn':>6} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")

    with torch.no_grad():
        x = relay_il[i](x)
    step += 1
    if step in [1, 2, 4, 8, 12, 16]:
        cv_n = compute_cv(F.normalize(x, dim=-1), N_CV)
        ed = eff_dim(x)
        cos = (F.normalize(pts, dim=-1) * F.normalize(x, dim=-1)).sum(-1).mean().item()
        print(f"  {step:>6} {'relay':>6} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")


# ── TEST 3: Clustered input ──
print(f"\n{'━'*80}")
print("TEST 3: Clustered input — relay vs attention")
print(f"{'━'*80}")

centroids = F.normalize(torch.randn(10, D), dim=-1).to(DEVICE)
assignments = torch.randint(0, 10, (N,))

print(f"  {'spread':>8} {'CV_before':>10} {'CV_relay':>10} {'CV_attn':>10} "
      f"{'Δ_relay':>8} {'Δ_attn':>8}")

for spread in [0.1, 0.3, 0.5, 1.0]:
    pts_c = F.normalize(centroids[assignments] +
                        torch.randn(N, D, device=DEVICE) * spread, dim=-1)
    cv_b = compute_cv(pts_c, N_CV)

    r = ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE)
    with torch.no_grad():
        out_r = r(pts_c)
    cv_r = compute_cv(F.normalize(out_r, dim=-1), N_CV)

    a = SelfAttnBlock(D, 8).to(DEVICE)
    with torch.no_grad():
        out_a = a(pts_c.unsqueeze(0)).squeeze(0)
    cv_a = compute_cv(F.normalize(out_a, dim=-1), N_CV)

    print(f"  {spread:>8.1f} {cv_b:>10.4f} {cv_r:>10.4f} {cv_a:>10.4f} "
          f"{cv_r-cv_b:>+8.4f} {cv_a-cv_b:>+8.4f}")


# ── TEST 4: Throughput ──
print(f"\n{'━'*80}")
print("TEST 4: Throughput — vectorized relay vs attention")
print(f"{'━'*80}")

relay_b = ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE)
attn_b = SelfAttnBlock(D, 8).to(DEVICE)
pts_b = uniform_sphere(N, D).to(DEVICE)

# Warmup
for _ in range(10):
    with torch.no_grad():
        _ = relay_b(pts_b)
        _ = attn_b(pts_b.unsqueeze(0))
torch.cuda.synchronize()

t0 = time.time()
for _ in range(200):
    with torch.no_grad():
        _ = relay_b(pts_b)
torch.cuda.synchronize()
relay_ms = (time.time() - t0) / 200 * 1000

t0 = time.time()
for _ in range(200):
    with torch.no_grad():
        _ = attn_b(pts_b.unsqueeze(0))
torch.cuda.synchronize()
attn_ms = (time.time() - t0) / 200 * 1000

n_params_attn = sum(p.numel() for p in attn_b.parameters())
print(f"  Relay:     {relay_ms:.2f}ms  ({n_params_relay:,} params)")
print(f"  Attention: {attn_ms:.2f}ms  ({n_params_attn:,} params)")
print(f"  Speedup:   {attn_ms/relay_ms:.2f}× {'◄ relay faster' if relay_ms < attn_ms else '◄ attn faster'}")


# ── TEST 5: Patch count sweep ──
print(f"\n{'━'*80}")
print("TEST 5: Patch count sweep")
print(f"{'━'*80}")

pts = uniform_sphere(N, D).to(DEVICE)
print(f"  {'patches':>8} {'d_patch':>8} {'params':>10} {'CV_out':>8} {'cos_orig':>10} {'ms':>8}")

for n_patches in [1, 2, 4, 8, 16]:
    pd = D // n_patches
    if D % n_patches != 0: continue
    r = ConstellationRelay(D, pd, min(pd, 16), 3, 32).to(DEVICE)
    np_r = sum(p.numel() for p in r.parameters())

    with torch.no_grad(): _ = r(pts)  # warmup
    torch.cuda.synchronize()

    t0 = time.time()
    for _ in range(100):
        with torch.no_grad():
            out = r(pts)
    torch.cuda.synchronize()
    ms = (time.time() - t0) / 100 * 1000

    cv = compute_cv(F.normalize(out, dim=-1), N_CV)
    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
    print(f"  {n_patches:>8} {pd:>8} {np_r:>10,} {cv:>8.4f} {cos:>10.6f} {ms:>8.2f}")


# ── TEST 6: Scaling to large dim ──
print(f"\n{'━'*80}")
print("TEST 6: Scaling — relay at d=256, 512, 768, 1024")
print(f"{'━'*80}")

print(f"  {'D':>6} {'patches':>8} {'params':>10} {'CV_in':>8} {'CV_out':>8} {'cos':>10} {'ms':>8}")

for D_test in [256, 512, 768, 1024]:
    pts_t = uniform_sphere(N, D_test).to(DEVICE)
    r = ConstellationRelay(D_test, 16, 16, 3, 32).to(DEVICE)
    np_r = sum(p.numel() for p in r.parameters())
    cv_in = compute_cv(pts_t, N_CV)

    with torch.no_grad(): _ = r(pts_t)
    torch.cuda.synchronize()
    t0 = time.time()
    for _ in range(50):
        with torch.no_grad():
            out = r(pts_t)
    torch.cuda.synchronize()
    ms = (time.time() - t0) / 50 * 1000

    cv_out = compute_cv(F.normalize(out, dim=-1), N_CV)
    cos = (F.normalize(pts_t, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
    print(f"  {D_test:>6} {D_test//16:>8} {np_r:>10,} {cv_in:>8.4f} {cv_out:>8.4f} "
          f"{cos:>10.6f} {ms:>8.2f}")


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

print(f"\n{'='*80}")
print("SUMMARY")
print(f"{'='*80}")
print(f"""
  v2 changes:
    - Vectorized: all patches in single einsum, no Python loop
    - Gates init at sigmoid(-3.0) ≈ 0.047 (was 0.622)
    - Patchwork via batched einsum (P, tri, hidden) × (P, hidden, d)

  Architecture: {D}d → {D//16}p × 16d, {16}A/patch, 3-phase
  Params: {n_params_relay:,}
  
  Key comparison at depth 16:
    Relay-only cos_to_orig:      (see Test 2A, cycle 16)
    Attention-only cos_to_orig:  (see Test 2B, cycle 16)
    Interleaved cos_to_orig:     (see Test 2C, step 16)
""")
print(f"{'='*80}")
print("DONE")
print(f"{'='*80}")