#!/usr/bin/env python3
"""
CV Spectrum Empirical Test
===========================
Is CV ≈ 0.20-0.23 intrinsic to S^(d-1), or does it require training?

Tests:
  1. Uniform random on S^(d-1) — varying d from 8 to 1024
  2. Gaussian projected to sphere — same dims
  3. Clustered (10, 100, 1000 clusters) — does structure change CV?
  4. Post-contrastive simulation — push apart + attract to anchors
  5. Varying n_points (1K to 100K) — does sample size matter?
  6. Varying n_pentachora samples — convergence check

If CV band is intrinsic to the sphere, test 1 should already show it.
If it requires training dynamics, only test 4 should show it.
"""

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

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


def compute_cv(points, n_samples=2000, n_points=5):
    """
    Compute pentachoron volume CV from a set of unit-normed points.
    points: (N, D) on unit sphere
    """
    N = points.shape[0]
    if N < n_points:
        return float('nan')

    points = points.to(DEVICE).float()
    vols = []

    for _ in range(n_samples):
        idx = torch.randperm(min(N, 10000), device=DEVICE)[:n_points]
        pts = points[idx].unsqueeze(0)  # (1, 5, D)

        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)
    cv = (vt.std() / (vt.mean() + 1e-8)).item()
    return cv


def uniform_sphere(n, d):
    """Random uniform on S^(d-1)."""
    x = torch.randn(n, d)
    return F.normalize(x, dim=-1)


def clustered_sphere(n, d, n_clusters):
    """Points clustered around n_clusters random centroids on S^(d-1)."""
    centroids = F.normalize(torch.randn(n_clusters, d), dim=-1)
    # Assign each point to a random cluster, then perturb
    assignments = torch.randint(0, n_clusters, (n,))
    base = centroids[assignments]
    noise = torch.randn(n, d) * 0.3  # perturbation strength
    return F.normalize(base + noise, dim=-1)


def contrastive_sim(n, d, n_clusters=10, n_steps=200, tau=0.07):
    """
    Simulate contrastive learning dynamics on the sphere.
    Start uniform, iteratively:
      - Attract same-cluster points
      - Repel different-cluster points
    """
    # Start uniform
    points = F.normalize(torch.randn(n, d), dim=-1).to(DEVICE)
    labels = torch.randint(0, n_clusters, (n,), device=DEVICE)

    for step in range(n_steps):
        # Use a subset for efficiency
        idx = torch.randperm(n, device=DEVICE)[:min(2000, n)]
        subset = points[idx]
        sub_labels = labels[idx]

        # Compute pairwise cosine
        sim = subset @ subset.T / tau
        same_mask = sub_labels.unsqueeze(0) == sub_labels.unsqueeze(1)

        # Soft contrastive gradient
        probs = sim.softmax(dim=-1)
        # Push away from different class, pull toward same class
        attract = (same_mask.float() * probs).sum(dim=-1, keepdim=True)
        grad = subset * attract  # simplified attraction

        # Compute same-class centroid pull
        for c in range(n_clusters):
            c_mask = sub_labels == c
            if c_mask.sum() < 2:
                continue
            centroid = F.normalize(subset[c_mask].mean(0, keepdim=True), dim=-1)
            # Pull toward centroid
            points[idx[c_mask]] = F.normalize(
                points[idx[c_mask]] + 0.1 * (centroid - points[idx[c_mask]]),
                dim=-1)

        # Spread all points apart slightly
        if step % 10 == 0:
            sub2 = points[torch.randperm(n, device=DEVICE)[:500]]
            sim2 = sub2 @ sub2.T
            sim2.fill_diagonal_(-2)
            nn_idx = sim2.argmax(dim=1)
            points[torch.randperm(n, device=DEVICE)[:500]] = F.normalize(
                sub2 - 0.02 * sub2[nn_idx], dim=-1)

    return points.cpu()


def anchored_sphere(n, d, n_anchors=64):
    """
    Points attracted toward fixed anchors on sphere.
    Simulates post-constellation embedding space.
    """
    anchors = F.normalize(torch.randn(n_anchors, d), dim=-1)
    # Each point near a random anchor
    assignments = torch.randint(0, n_anchors, (n,))
    base = anchors[assignments]
    noise = torch.randn(n, d) * 0.2
    return F.normalize(base + noise, dim=-1)


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

print("=" * 70)
print("CV SPECTRUM — EMPIRICAL GEOMETRY TEST")
print(f"  Device: {DEVICE}")
print(f"  Pentachoron samples: 2000 per measurement")
print("=" * 70)

results = []

# ── TEST 1: Uniform sphere, varying dimension ──
print(f"\n{'─'*70}")
print("TEST 1: Uniform random on S^(d-1)")
print(f"{'─'*70}")
print(f"  {'dim':>6} {'n_pts':>8} {'CV':>8} {'time':>6}")

for d in [8, 16, 32, 64, 128, 256, 384, 512, 768, 1024]:
    n = 10000
    t0 = time.time()
    pts = uniform_sphere(n, d)
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {d:>6} {n:>8} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('uniform', d, n, cv))


# ── TEST 2: Varying sample count at fixed dim ──
print(f"\n{'─'*70}")
print("TEST 2: Sample count sweep (d=128)")
print(f"{'─'*70}")
print(f"  {'dim':>6} {'n_pts':>8} {'CV':>8} {'time':>6}")

for n in [500, 1000, 2000, 5000, 10000, 25000, 50000, 100000]:
    d = 128
    t0 = time.time()
    pts = uniform_sphere(n, d)
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {d:>6} {n:>8} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('n_sweep', d, n, cv))


# ── TEST 3: Clustered points ──
print(f"\n{'─'*70}")
print("TEST 3: Clustered on sphere (d=128, n=10000)")
print(f"{'─'*70}")
print(f"  {'clusters':>10} {'spread':>8} {'CV':>8} {'time':>6}")

for n_clusters in [2, 5, 10, 20, 50, 100, 500, 1000]:
    d = 128
    n = 10000
    t0 = time.time()
    pts = clustered_sphere(n, d, n_clusters)
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {n_clusters:>10} {'0.3':>8} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('clustered', n_clusters, n, cv))


# ── TEST 4: Varying cluster spread ──
print(f"\n{'─'*70}")
print("TEST 4: Cluster spread sweep (d=128, 10 clusters)")
print(f"{'─'*70}")
print(f"  {'spread':>8} {'CV':>8} {'time':>6}")

for spread in [0.01, 0.05, 0.1, 0.2, 0.3, 0.5, 0.7, 1.0, 2.0, 5.0]:
    d = 128
    n = 10000
    centroids = F.normalize(torch.randn(10, d), dim=-1)
    assignments = torch.randint(0, 10, (n,))
    base = centroids[assignments]
    noise = torch.randn(n, d) * spread
    pts = F.normalize(base + noise, dim=-1)
    t0 = time.time()
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {spread:>8.3f} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('spread', spread, n, cv))


# ── TEST 5: Anchor-attracted points ──
print(f"\n{'─'*70}")
print("TEST 5: Anchor-attracted (d=128, n=10000)")
print(f"{'─'*70}")
print(f"  {'n_anchors':>10} {'CV':>8} {'time':>6}")

for na in [4, 8, 16, 32, 64, 128, 256]:
    d = 128
    n = 10000
    t0 = time.time()
    pts = anchored_sphere(n, d, n_anchors=na)
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {na:>10} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('anchored', na, n, cv))


# ── TEST 6: Simulated contrastive dynamics ──
print(f"\n{'─'*70}")
print("TEST 6: Simulated contrastive training (d=128)")
print(f"{'─'*70}")
print(f"  {'n_clusters':>10} {'steps':>6} {'CV':>8} {'time':>6}")

for nc in [5, 10, 50, 100]:
    d = 128
    n = 10000
    t0 = time.time()
    pts = contrastive_sim(n, d, n_clusters=nc, n_steps=200)
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {nc:>10} {'200':>6} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('contrastive', nc, n, cv))


# ── TEST 7: Dimension sweep for contrastive ──
print(f"\n{'─'*70}")
print("TEST 7: Contrastive sim across dimensions (10 clusters, 200 steps)")
print(f"{'─'*70}")
print(f"  {'dim':>6} {'CV':>8} {'time':>6}")

for d in [16, 32, 64, 128, 256, 512]:
    n = 10000
    t0 = time.time()
    pts = contrastive_sim(n, d, n_clusters=10, n_steps=200)
    cv = compute_cv(pts, n_samples=2000)
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {d:>6} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('contrastive_dim', d, n, cv))


# ── TEST 8: The 100K point test ──
print(f"\n{'─'*70}")
print("TEST 8: 100K points — uniform vs clustered vs contrastive")
print(f"{'─'*70}")

configs = [
    ("uniform_100k", lambda: uniform_sphere(100000, 128)),
    ("clustered_10_100k", lambda: clustered_sphere(100000, 128, 10)),
    ("clustered_100_100k", lambda: clustered_sphere(100000, 128, 100)),
    ("anchored_64_100k", lambda: anchored_sphere(100000, 128, 64)),
]

print(f"  {'name':>25} {'CV':>8} {'time':>6}")
for name, gen_fn in configs:
    t0 = time.time()
    pts = gen_fn()
    cv = compute_cv(pts, n_samples=5000)  # more samples for 100K
    dt = time.time() - t0
    tag = "◄ BAND" if 0.18 <= cv <= 0.27 else ""
    print(f"  {name:>25} {cv:>8.4f} {dt:>5.1f}s {tag}")
    results.append(('100k', name, 100000, cv))


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

print(f"\n{'='*70}")
print("SUMMARY")
print(f"{'='*70}")

# Count how many hit the band
in_band = sum(1 for r in results if isinstance(r[3], float)
              and not math.isnan(r[3]) and 0.18 <= r[3] <= 0.27)
total = sum(1 for r in results if isinstance(r[3], float)
            and not math.isnan(r[3]))

print(f"  Tests in CV band [0.18, 0.27]: {in_band}/{total}")
print(f"  Tests outside band: {total - in_band}/{total}")

# Group by test type
from collections import defaultdict
by_type = defaultdict(list)
for r in results:
    if isinstance(r[3], float) and not math.isnan(r[3]):
        by_type[r[0]].append(r[3])

print(f"\n  CV by test type:")
for t in sorted(by_type.keys()):
    vals = by_type[t]
    print(f"    {t:<20}: mean={np.mean(vals):.4f} "
          f"std={np.std(vals):.4f} "
          f"range=[{min(vals):.4f}, {max(vals):.4f}] "
          f"in_band={sum(0.18 <= v <= 0.27 for v in vals)}/{len(vals)}")

print(f"\n{'='*70}")
print("DONE")
print(f"{'='*70}")