#!/usr/bin/env python3
"""
RoPE Attention vs Constellation Relay
========================================
Two RoPE variants:
  1. Standard RoPE (Su et al.) — fixed base frequency 10000
  2. NTK-aware RoPE — scaled base frequency for longer contexts

Same battery of tests: single pass, depth stability, interleaved.
"""

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)


# ══════════════════════════════════════════════════════════════════
# RoPE IMPLEMENTATIONS
# ══════════════════════════════════════════════════════════════════

class RotaryEmbedding(nn.Module):
    """Standard RoPE — fixed sinusoidal rotation frequencies."""
    def __init__(self, dim, base=10000.0):
        super().__init__()
        self.dim = dim
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, seq_len, device):
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)  # (S, dim/2)
        emb = torch.cat([freqs, freqs], dim=-1)  # (S, dim)
        return emb.cos(), emb.sin()


class NTKRotaryEmbedding(nn.Module):
    """NTK-aware RoPE — scaled base for extended context."""
    def __init__(self, dim, base=10000.0, scale_factor=4.0):
        super().__init__()
        self.dim = dim
        # NTK scaling: base^(dim/(dim-2)) * scale_factor
        scaled_base = base * (scale_factor ** (dim / (dim - 2)))
        inv_freq = 1.0 / (scaled_base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq)

    def forward(self, seq_len, device):
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.einsum('i,j->ij', t, self.inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        return emb.cos(), emb.sin()


def apply_rotary(x, cos, sin):
    """Apply rotary embeddings to Q or K: (B, H, S, d)."""
    d = x.shape[-1]
    x1 = x[..., :d//2]
    x2 = x[..., d//2:]
    cos = cos[:x.shape[-2], :d//2].unsqueeze(0).unsqueeze(0)  # (1, 1, S, d/2)
    sin = sin[:x.shape[-2], :d//2].unsqueeze(0).unsqueeze(0)
    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)


# ══════════════════════════════════════════════════════════════════
# ATTENTION BLOCKS
# ══════════════════════════════════════════════════════════════════

class VanillaAttnBlock(nn.Module):
    """Standard self-attention — no position encoding."""
    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)


class RoPEAttnBlock(nn.Module):
    """Self-attention with Rotary Position Embeddings."""
    def __init__(self, dim, n_heads=4, rope_type='standard', rope_base=10000.0,
                 ntk_scale=4.0):
        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)

        if rope_type == 'standard':
            self.rope = RotaryEmbedding(self.head_dim, base=rope_base)
        elif rope_type == 'ntk':
            self.rope = NTKRotaryEmbedding(self.head_dim, base=rope_base,
                                            scale_factor=ntk_scale)
        self.rope_type = rope_type

    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)  # (3, B, H, S, hd)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # Apply RoPE to Q and K
        cos, sin = self.rope(S, x.device)
        q = apply_rotary(q, cos, sin)
        k = apply_rotary(k, cos, sin)

        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)


# ══════════════════════════════════════════════════════════════════
# CONSTELLATION RELAY (copy from v2)
# ══════════════════════════════════════════════════════════════════

class ConstellationRelay(nn.Module):
    def __init__(self, input_dim, patch_dim=16, n_anchors=16, n_phases=3,
                 pw_hidden=32, gate_init=-3.0):
        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

        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())

        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))
        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)
        self.gates = nn.Parameter(torch.full((P,), gate_init))
        self.norm = nn.LayerNorm(input_dim)

    def drift(self):
        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):
        h = F.normalize(self.home, dim=-1)
        c = F.normalize(self.anchors, dim=-1)
        omega = self.drift().unsqueeze(-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 forward(self, x):
        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)

        # Multi-phase triangulation
        phases = torch.linspace(0, 1, self.n_phases).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)
        tri = torch.cat(tris, dim=-1)

        # Patchwork
        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)

        gate = self.gates.sigmoid().unsqueeze(0).unsqueeze(-1)
        blended = gate * pw_out + (1 - gate) * patches
        out = blended.reshape(B, D)
        return x + out


# ══════════════════════════════════════════════════════════════════
# TEST SUITE
# ══════════════════════════════════════════════════════════════════

N = 2000
D = 128
N_CV = 2000

print("=" * 90)
print("RoPE ATTENTION vs CONSTELLATION RELAY")
print(f"  Input dim: {D}, Sequence length: {N}")
print(f"  Device: {DEVICE}")
print("=" * 90)

pts = uniform_sphere(N, D).to(DEVICE)
cv_base = compute_cv(pts, N_CV)
ed_base = eff_dim(pts)
print(f"  Baseline: CV={cv_base:.4f}  eff_dim={ed_base:.1f}")

# Build all architectures
configs = {
    'vanilla':    lambda: VanillaAttnBlock(D, 8).to(DEVICE),
    'rope_std':   lambda: RoPEAttnBlock(D, 8, 'standard', 10000).to(DEVICE),
    'rope_ntk':   lambda: RoPEAttnBlock(D, 8, 'ntk', 10000, 4.0).to(DEVICE),
    'relay':      lambda: ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE),
}


# ── TEST 1: Single pass comparison ──
print(f"\n{'━'*90}")
print("TEST 1: Single pass — all architectures")
print(f"{'━'*90}")
print(f"  {'arch':>12} {'params':>8} {'CV_out':>8} {'CV_norm':>8} "
      f"{'cos_orig':>10} {'eff_dim':>8}")

for name, builder in configs.items():
    module = builder()
    np_ = sum(p.numel() for p in module.parameters())
    with torch.no_grad():
        if name == 'relay':
            out = module(pts)
        else:
            out = module(pts.unsqueeze(0)).squeeze(0)
    cv = compute_cv(out, N_CV)
    cv_n = compute_cv(F.normalize(out, dim=-1), N_CV)
    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
    ed = eff_dim(out)
    print(f"  {name:>12} {np_:>8,} {cv:>8.4f} {cv_n:>8.4f} {cos:>10.6f} {ed:>8.1f}")


# ── TEST 2: Depth sweep — 16 layers each ──
print(f"\n{'━'*90}")
print("TEST 2: Depth sweep — 16 layers, all architectures")
print(f"{'━'*90}")

checkpoints = [1, 2, 4, 8, 12, 16]

for name, builder in configs.items():
    print(f"\n  {name}:")
    print(f"  {'depth':>6} {'CV':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}")

    stack = nn.ModuleList([builder() for _ in range(16)])
    x = pts.clone()
    for i, layer in enumerate(stack):
        with torch.no_grad():
            if name == 'relay':
                x = layer(x)
            else:
                x = layer(x.unsqueeze(0)).squeeze(0)
        if (i + 1) in checkpoints:
            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}")


# ── TEST 3: Interleaved — each attention type + relay ──
print(f"\n{'━'*90}")
print("TEST 3: Interleaved — [attn type] → relay → [attn type] → relay → ...")
print(f"{'━'*90}")

for attn_name in ['vanilla', 'rope_std', 'rope_ntk']:
    print(f"\n  {attn_name} + relay interleaved:")
    print(f"  {'step':>6} {'type':>8} {'CV_n':>8} {'eff_d':>8} {'cos_orig':>10}")

    attn_builder = configs[attn_name]
    attn_layers = nn.ModuleList([attn_builder() for _ in range(8)])
    relay_layers = nn.ModuleList([
        ConstellationRelay(D, 16, 16, 3, 32).to(DEVICE) for _ in range(8)])

    x = pts.clone()
    step = 0
    for i in range(8):
        # Attention step
        with torch.no_grad():
            x = attn_layers[i](x.unsqueeze(0)).squeeze(0)
        step += 1
        if step in checkpoints:
            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':>8} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")

        # Relay step
        with torch.no_grad():
            x = relay_layers[i](x)
        step += 1
        if step in checkpoints:
            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':>8} {cv_n:>8.4f} {ed:>8.1f} {cos:>10.6f}")


# ── TEST 4: Throughput comparison ──
print(f"\n{'━'*90}")
print("TEST 4: Throughput")
print(f"{'━'*90}")

print(f"  {'arch':>12} {'ms':>8} {'params':>10}")

for name, builder in configs.items():
    module = builder()
    np_ = sum(p.numel() for p in module.parameters())

    # Warmup
    for _ in range(10):
        with torch.no_grad():
            if name == 'relay':
                _ = module(pts)
            else:
                _ = module(pts.unsqueeze(0))
    torch.cuda.synchronize()

    t0 = time.time()
    for _ in range(200):
        with torch.no_grad():
            if name == 'relay':
                _ = module(pts)
            else:
                _ = module(pts.unsqueeze(0))
    torch.cuda.synchronize()
    ms = (time.time() - t0) / 200 * 1000
    print(f"  {name:>12} {ms:>8.2f} {np_:>10,}")


# ── TEST 5: Clustered input — all architectures ──
print(f"\n{'━'*90}")
print("TEST 5: Clustered input (10 clusters, d=128)")
print(f"{'━'*90}")

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

print(f"  {'spread':>8} {'CV_base':>8} {'vanilla':>8} {'rope_std':>8} "
      f"{'rope_ntk':>8} {'relay':>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)

    row = f"  {spread:>8.1f} {cv_b:>8.4f}"
    for name, builder in configs.items():
        module = builder()
        with torch.no_grad():
            if name == 'relay':
                out = module(pts_c)
            else:
                out = module(pts_c.unsqueeze(0)).squeeze(0)
        cv = compute_cv(F.normalize(out, dim=-1), N_CV)
        row += f" {cv:>8.4f}"
    print(row)


# ── TEST 6: RoPE frequency analysis ──
print(f"\n{'━'*90}")
print("TEST 6: RoPE base frequency sweep")
print(f"  Does the rotation frequency affect geometric preservation?")
print(f"{'━'*90}")

print(f"  {'base':>10} {'CV_n':>8} {'cos_orig':>10} {'eff_dim':>8}")

for base in [100, 500, 1000, 5000, 10000, 50000, 100000, 500000]:
    module = RoPEAttnBlock(D, 8, 'standard', base).to(DEVICE)
    with torch.no_grad():
        out = module(pts.unsqueeze(0)).squeeze(0)
    cv_n = compute_cv(F.normalize(out, dim=-1), N_CV)
    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
    ed = eff_dim(out)
    print(f"  {base:>10} {cv_n:>8.4f} {cos:>10.6f} {ed:>8.1f}")


# ── TEST 7: NTK scale factor sweep ──
print(f"\n{'━'*90}")
print("TEST 7: NTK scale factor sweep (base=10000)")
print(f"{'━'*90}")

print(f"  {'scale':>8} {'CV_n':>8} {'cos_orig':>10} {'eff_dim':>8}")

for scale in [1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 64.0]:
    module = RoPEAttnBlock(D, 8, 'ntk', 10000, scale).to(DEVICE)
    with torch.no_grad():
        out = module(pts.unsqueeze(0)).squeeze(0)
    cv_n = compute_cv(F.normalize(out, dim=-1), N_CV)
    cos = (F.normalize(pts, dim=-1) * F.normalize(out, dim=-1)).sum(-1).mean().item()
    ed = eff_dim(out)
    print(f"  {scale:>8.1f} {cv_n:>8.4f} {cos:>10.6f} {ed:>8.1f}")


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

print(f"\n{'='*90}")
print("SUMMARY — cos_to_orig at depth 16")
print(f"{'='*90}")
print(f"""
  Compare the depth-16 cos_to_orig from Test 2 across all architectures:
  
    vanilla attention:     (see Test 2)
    RoPE standard:         (see Test 2)
    RoPE NTK:              (see Test 2)
    constellation relay:   (see Test 2)
    
  And the interleaved results from Test 3:
    vanilla + relay:       (see Test 3)
    rope_std + relay:      (see Test 3)
    rope_ntk + relay:      (see Test 3)
""")
print(f"{'='*90}")
print("DONE")
print(f"{'='*90}")