import json
import math
import random
import time
import numpy as np

SEED = 2261
np.random.seed(SEED)
random.seed(SEED)


def projection(d, r, rng):
    return rng.normal(0.0, 1.0 / math.sqrt(r), size=(r, d)).astype(np.float64)


def pairwise_ratios(X, P, max_pairs=200000, rng=None):
    n = len(X)
    rng = np.random.default_rng(SEED) if rng is None else rng
    total = n * (n - 1) // 2
    if total <= max_pairs:
        ii, jj = np.triu_indices(n, 1)
    else:
        ii = rng.integers(0, n, size=max_pairs)
        jj = rng.integers(0, n, size=max_pairs)
        keep = ii != jj
        ii, jj = ii[keep], jj[keep]
    dx = X[ii] - X[jj]
    dz = (X[ii] - X[jj]) @ P.T
    return np.linalg.norm(dz, axis=1) / np.maximum(np.linalg.norm(dx, axis=1), 1e-12)


def width_rule(d, n, eps, C=1.0):
    raw = min(d, n - 1, math.log(2.0 + eps * eps * n) / (eps * eps))
    return max(1, min(d, int(math.ceil(C * raw)))), raw


def toy_math_verification():
    rng = np.random.default_rng(SEED + 1)
    d, n = 128, 180
    X = rng.normal(size=(n, d))
    # Prediction 1: Gaussian scaling is unbiased: E[||Pv||^2/||v||^2] = 1.
    v = rng.normal(size=d)
    norm_trials = []
    for _ in range(300):
        P = projection(d, 32, rng)
        norm_trials.append(np.sum((P @ v) ** 2) / np.sum(v ** 2))
    norm_mean = float(np.mean(norm_trials))
    # Prediction 2: for a fixed pair, ratio has SD approximately 1/sqrt(2r).
    fixed_v = X[0] - X[1]
    fixed_ratios = []
    for _ in range(500):
        P = projection(d, 32, rng)
        fixed_ratios.append(np.linalg.norm(P @ fixed_v) / np.linalg.norm(fixed_v))
    fixed_sd = float(np.std(fixed_ratios, ddof=1))
    fixed_sd_pred = 1.0 / math.sqrt(2 * 32)
    # Prediction 3: maximum of many pair errors scales as sqrt(log(number pairs)/r).
    max_rows = []
    for r in [8, 16, 32, 64, 128]:
        P = projection(d, r, rng)
        ratios = pairwise_ratios(X, P, rng=rng)
        max_rows.append({
            "r": r,
            "max_abs_error": float(np.max(np.abs(ratios - 1))),
            "p95_abs_error": float(np.percentile(np.abs(ratios - 1), 95)),
            "predicted_max_scale": math.sqrt(2.0 * math.log(max(2, n * (n - 1) / 2)) / r),
            "predicted_fixed_sd": 1.0 / math.sqrt(2 * r),
        })
    # Prediction 4: sharp schedule transitions with n and saturates at d or n-1.
    schedule = []
    for eps in [0.1, 0.2, 0.3]:
        for nn in [8, 32, 128, 512]:
            rr, raw = width_rule(128, nn, eps, C=1)
            schedule.append({"epsilon": eps, "n": nn, "r": rr, "raw": raw})
    return {
        "unbiased_norm": {"observed_mean": norm_mean, "predicted": 1.0, "relative_error": abs(norm_mean - 1.0)},
        "fixed_pair_sd": {"observed": fixed_sd, "predicted": fixed_sd_pred, "ratio_observed_to_predicted": fixed_sd / fixed_sd_pred},
        "max_error_scaling": max_rows,
        "sharp_width_schedule": schedule,
    }


def softmax_loss(W, X, y):
    logits = X @ W
    logits -= logits.max(axis=1, keepdims=True)
    exp = np.exp(logits)
    probs = exp / exp.sum(axis=1, keepdims=True)
    loss = -np.log(np.maximum(probs[np.arange(len(y)), y], 1e-12)).mean()
    grad = X.T @ (probs - np.eye(W.shape[1])[y]) / len(y)
    return float(loss), grad


def train_linear(Xtr, ytr, Xva, yva, steps=250, lr=0.4):
    W = np.zeros((Xtr.shape[1], 2), dtype=np.float64)
    for _ in range(steps):
        _, g = softmax_loss(W, Xtr, ytr)
        W -= lr * g
    loss, _ = softmax_loss(W, Xva, yva)
    acc = float(np.mean(np.argmax(Xva @ W, axis=1) == yva))
    return loss, acc


def downstream_experiment():
    rng = np.random.default_rng(SEED + 2)
    ntr, nva, d, n_tokens = 2400, 800, 96, 24
    # Two classes differ in a low-dimensional token mean; unrelated token noise
    # makes the full-width representation deliberately width-sensitive.
    ytr = rng.integers(0, 2, ntr)
    yva = rng.integers(0, 2, nva)
    def make(y):
        x = rng.normal(size=(len(y), n_tokens, d))
        x[:, :, :8] += (2 * y[:, None, None] - 1) * 0.9
        return x
    Xtr0, Xva0 = make(ytr), make(yva)
    full_tr, full_va = Xtr0.mean(axis=1), Xva0.mean(axis=1)
    full_loss, full_acc = train_linear(full_tr, ytr, full_va, yva)
    rows = []
    for eps in [0.1, 0.2, 0.3]:
        for C in [1, 2, 4]:
            r, raw = width_rule(d, n_tokens, eps, C)
            P = projection(d, r, rng)
            ztr, zva = Xtr0 @ P.T, Xva0 @ P.T
            loss, acc = train_linear(ztr.mean(axis=1), ytr, zva.mean(axis=1), yva)
            ratios = pairwise_ratios(Xva0.reshape(-1, d)[:160], P, rng=rng)
            rows.append({"epsilon": eps, "C": C, "r": r, "raw": raw, "val_loss": loss, "val_accuracy": acc,
                         "relative_accuracy_loss": (full_acc - acc) / max(full_acc, 1e-12),
                         "p95_abs_distortion": float(np.percentile(np.abs(ratios - 1), 95)),
                         "activation_width_ratio": r / d})
    return {"full_width": {"d": d, "val_loss": full_loss, "val_accuracy": full_acc, "activation_width": d}, "compressed": rows}


def main():
    t0 = time.time()
    out = {"seed": SEED, "math": toy_math_verification(), "downstream": downstream_experiment()}
    out["runtime_sec"] = time.time() - t0
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()