import json
import math
import random
import numpy as np
import torch

SEED = 58
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
DT = torch.float64
EPS = 1e-10


def pairs(n, d, seed):
    g = torch.Generator().manual_seed(seed)
    u = torch.randn(n, d, generator=g, dtype=DT)
    v = torch.randn(n, d, generator=g, dtype=DT)
    # Rank-one PSD features X=uu^T, Y=vv^T, normalized to remove trivial scale.
    u = u / u.norm(dim=1, keepdim=True)
    v = v / v.norm(dim=1, keepdim=True)
    return u, v


def ratios(A, u, v):
    # A contains sensing vectors a_i; A_i=a_i a_i^T is PSD.
    zu = (u @ A.T).square()
    zv = (v @ A.T).square()
    numerator = (zu - zv).abs().sum(dim=1)
    # ||uu^T-vv^T||_F, computed exactly for unit u,v.
    den = torch.sqrt(2.0 - 2.0 * (u * v).sum(dim=1).square())
    return numerator / (den + EPS)


def stats(q):
    q = q.detach().cpu().numpy()
    lo, hi = np.quantile(q, [.05, .95])
    return {"L_q05": float(lo), "U_q95": float(hi),
            "beta": float(hi / max(lo, EPS)), "median": float(np.median(q)),
            "mean": float(np.mean(q))}


def random_A(m, d, seed, gamma=1.0):
    g = torch.Generator().manual_seed(seed)
    A = torch.randn(m, d, generator=g, dtype=DT)
    # A controlled anisotropy: amplify one coordinate before row normalization.
    A[:, 0] *= gamma
    return A / A.norm(dim=1, keepdim=True)


def optimize_A(m, d, train_u, train_v, test_u, test_v, steps=500):
    g = torch.Generator().manual_seed(9000 + m)
    W = torch.randn(m, d, generator=g, dtype=DT, requires_grad=True)
    opt = torch.optim.Adam([W], lr=.035)
    history = []
    for step in range(steps):
        A = W / (W.norm(dim=1, keepdim=True) + EPS)
        q = ratios(A, train_u, train_v)
        lo = torch.quantile(q, .05)
        hi = torch.quantile(q, .95)
        loss = torch.log(hi + EPS) - torch.log(lo + EPS)
        opt.zero_grad(); loss.backward(); opt.step()
        if step in (0, steps - 1): history.append(float(loss.detach()))
    A = (W / (W.norm(dim=1, keepdim=True) + EPS)).detach()
    return A, history


def main():
    d = 8
    train_u, train_v = pairs(2400, d, 10)
    test_u, test_v = pairs(10000, d, 11)
    out = {"seed": SEED, "d": d, "predictions": {}}

    # Prediction 1: common positive scaling changes L and U equally, hence beta is invariant.
    A = random_A(64, d, 21)
    base = stats(ratios(A, test_u, test_v))
    scaled = stats(ratios(3.7 * A, test_u, test_v))
    out["scale_invariance"] = {"base": base, "scaled_3.7": scaled,
        "observed_beta_relative_change": scaled["beta"] / base["beta"] - 1.0,
        "observed_L_scale": scaled["L_q05"] / base["L_q05"],
        "observed_U_scale": scaled["U_q95"] / base["U_q95"]}
    out["predictions"]["scale"] = "beta unchanged; L and U multiply by 3.7^2 because A_i=a_i a_i^T"

    # Prediction 2: more independent nonnegative PSD measurements narrow the empirical ratio spread.
    widths = [8, 16, 32, 64, 128, 256]
    width_rows = []
    for m in widths:
        s = stats(ratios(random_A(m, d, 100 + m), test_u, test_v))
        width_rows.append({"m": m, **s})
    out["width_sweep"] = width_rows
    out["predictions"]["width"] = "beta should generally fall as independent measurements average fluctuations"

    # Prediction 3: directional anisotropy creates collapsed directions and increases beta.
    anis_rows = []
    for gamma in [1., 2., 4., 8., 16.]:
        s = stats(ratios(random_A(128, d, 333, gamma), test_u, test_v))
        anis_rows.append({"gamma": gamma, **s})
    out["anisotropy_sweep"] = anis_rows
    out["predictions"]["anisotropy"] = "increasing directional imbalance raises U/L"

    # Mini experiment: same width, random unregularized PSD sensing vs conditioned sensing.
    m = 32
    baseline_A = random_A(m, d, 700)
    cond_A, hist = optimize_A(m, d, train_u, train_v, test_u, test_v)
    out["mini_experiment"] = {
        "m": m, "baseline_random": stats(ratios(baseline_A, test_u, test_v)),
        "conditioned_optimized": stats(ratios(cond_A, test_u, test_v)),
        "train_log_condition_start_end": hist,
        "mean_row_norm_baseline": float(baseline_A.norm(dim=1).mean()),
        "mean_row_norm_conditioned": float(cond_A.norm(dim=1).mean())}

    # Simple pass/fail checks stated quantitatively.
    scale = out["scale_invariance"]
    scale_ok = abs(scale["observed_beta_relative_change"]) < 1e-8 and abs(scale["observed_L_scale"]-3.7**2) < 1e-8 and abs(scale["observed_U_scale"]-3.7**2) < 1e-8
    b = [x["beta"] for x in width_rows]
    width_ok = b[-1] < b[0]
    a = [x["beta"] for x in anis_rows]
    anis_ok = a[-1] > a[0]
    mini_ok = out["mini_experiment"]["conditioned_optimized"]["beta"] < out["mini_experiment"]["baseline_random"]["beta"]
    out["checks"] = {"scale_exact": scale_ok, "width_endpoint": width_ok,
                      "anisotropy_endpoint": anis_ok, "optimization_test_beta": mini_ok,
                      "mechanism_manifestations": int(scale_ok) + int(width_ok) + int(anis_ok)}
    with open("results.json", "w") as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()