import json, math, random
from pathlib import Path
import numpy as np
import torch

SEED = 3043

def seed_all(seed):
    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)


def kendall_tau_a(a, b):
    # Efficient Kendall tau over corresponding flattened distance values.
    from scipy.stats import kendalltau
    return float(kendalltau(np.asarray(a), np.asarray(b), variant="b").statistic)


def gaussian_verification():
    # P is the first m rows of a random orthogonal matrix. Repeated points
    # approximate the paper's independent Gaussian pair distribution.
    rng = np.random.default_rng(SEED)
    d, n = 64, 500
    rows = []
    for m in [2, 4, 8, 16, 32]:
        vals = []
        for rep in range(8):
            x = rng.normal(size=(n, d))
            q, _ = np.linalg.qr(rng.normal(size=(d, d)))
            z = x @ q[:, :m]
            D = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
            Dm = (d / m) * ((z[:, None, :] - z[None, :, :]) ** 2).sum(-1)
            iu = np.triu_indices(n, 1)
            vals.append(kendall_tau_a(D[iu], Dm[iu]))
        observed = float(np.mean(vals)); sd = float(np.std(vals))
        predicted = (2.0 / math.pi) * math.sqrt(m / d)
        rows.append({"m": m, "observed_tau": observed, "sd": sd, "prediction": predicted})
    # Also verify centering removes the large common distance baseline but
    # leaves exactly the same centered ordering / correlation.
    x = rng.normal(size=(n, d)); q, _ = np.linalg.qr(rng.normal(size=(d, d))); z = x @ q[:, :8]
    D = ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)
    Dm = (d / 8) * ((z[:, None, :] - z[None, :, :]) ** 2).sum(-1)
    iu = np.triu_indices(n, 1)
    raw_mean_gap = abs(float(D[iu].mean() - Dm[iu].mean()))
    centered_mean = float((D[iu] - D[iu].mean()).mean())
    centered_projected_mean = float((Dm[iu] - Dm[iu].mean()).mean())
    return {"ranking_scaling": rows, "mean_gap_before_centering": raw_mean_gap,
            "centered_means": [centered_mean, centered_projected_mean]}


def pair_dist(x):
    return ((x[:, None, :] - x[None, :, :]) ** 2).sum(-1)


def run_one(seed, centered, steps=500):
    seed_all(seed)
    device = "cuda" if torch.cuda.is_available() else "cpu"
    try:
        # A compact embedding problem: nuisance dimensions have much larger
        # variance than a low-dimensional class/task direction.
        g = torch.Generator().manual_seed(seed)
        n, d, m = 96, 32, 8
        labels = torch.arange(n) % 8
        centers = torch.randn(8, 3, generator=g) * 1.8
        x = torch.randn(n, d, generator=g) * 2.0
        x[:, :3] += centers[labels]
        x = x.to(device)
        W = torch.randn(m, d, generator=g, device=device) / math.sqrt(d)
        W = torch.nn.Parameter(W)
        opt = torch.optim.Adam([W], lr=0.035)
        # Fixed pair target and fixed normalization make comparisons matched.
        with torch.no_grad():
            target = pair_dist(x)
            target_mean = target.mean()
            target_std = target.std() + 1e-8
        for step in range(steps):
            opt.zero_grad()
            z = x @ W.t()
            pred = (d / m) * pair_dist(z)
            if centered:
                # stop-gradient EMA-like batch baseline; using the current
                # batch mean is the exact minibatch version in the proposal.
                a = (target - target.mean()) / (target.std() + 1e-8)
                b = (pred - pred.mean()) / (pred.std() + 1e-8)
                loss = ((a - b) ** 2).mean()
            else:
                loss = ((pred - target) / target_std).pow(2).mean()
            loss.backward(); opt.step()
        with torch.no_grad():
            pred = (d / m) * pair_dist(x @ W.t())
            iu = torch.triu_indices(n, n, offset=1, device=device)
            a = target[iu[0], iu[1]].detach().cpu().numpy()
            b = pred[iu[0], iu[1]].detach().cpu().numpy()
            tau = kendall_tau_a(a, b)
            centered_mse = float(np.mean(((a-a.mean())/a.std() - (b-b.mean())/b.std())**2))
            raw_mse = float(np.mean(((a-b)/a.std())**2))
            # Simple downstream retrieval: nearest projected neighbor shares label.
            zz = (x @ W.t()).detach().cpu().numpy()
            dd = ((zz[:,None,:]-zz[None,:,:])**2).sum(-1); np.fill_diagonal(dd, np.inf)
            recall = float(np.mean(labels.cpu().numpy()[dd.argmin(1)] == labels.cpu().numpy()))
        return {"tau": tau, "centered_mse": centered_mse, "raw_mse": raw_mse, "recall1": recall}
    except Exception as e:
        if device == "cuda":
            torch.cuda.empty_cache()
            # A CPU retry with the same code path is safer on shared GPUs.
            torch.cuda.is_available = lambda: False
            return run_one(seed, centered, steps)
        raise e


def main():
    verification = gaussian_verification()
    results = {}
    for mode, flag in [("raw", False), ("centered", True)]:
        vals = [run_one(s, flag) for s in [11, 22, 33, 44, 55]]
        results[mode] = {k: {"mean": float(np.mean([v[k] for v in vals])),
                             "sd": float(np.std([v[k] for v in vals]))} for k in vals[0]}
        results[mode]["runs"] = vals
    out = {"verification": verification, "mini_experiment": results}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()