import json, math, time
from pathlib import Path
import numpy as np

SEED = 2454
rng = np.random.default_rng(SEED)


def make_rotations(n, T, seed):
    r = np.random.default_rng(seed)
    ij = r.integers(0, n, size=(T, 2))
    # Resample equal pairs; sorting gives valid i<j.
    eq = ij[:, 0] == ij[:, 1]
    while np.any(eq):
        ij[eq, 1] = r.integers(0, n, size=int(eq.sum()))
        eq = ij[:, 0] == ij[:, 1]
    a = np.minimum(ij[:, 0], ij[:, 1])
    b = np.maximum(ij[:, 0], ij[:, 1])
    th = r.uniform(0.0, 2.0 * np.pi, size=T)
    return a.astype(np.int32), b.astype(np.int32), np.cos(th), np.sin(th)


def kac_apply(x, rotations):
    # x has shape (batch,n) or (n,); this is the stated streamed update.
    a, b, c, s = rotations
    z = np.array(x, dtype=np.float64, copy=True)
    for i, j, cc, ss in zip(a, b, c, s):
        u = z[..., i].copy()
        v = z[..., j].copy()
        z[..., i] = cc * u + ss * v
        z[..., j] = -ss * u + cc * v
    return z


def haar_projection(n, m, seed):
    r = np.random.default_rng(seed)
    g = r.normal(size=(n, n))
    q, rr = np.linalg.qr(g)
    # Make the QR convention deterministic up to column signs; either is Haar.
    q *= np.sign(np.diag(rr))[None, :]
    return np.sqrt(n / m) * q[:m, :]


def gaussian_projection(n, m, seed):
    return np.random.default_rng(seed).normal(size=(m, n)) / math.sqrt(m)


def metrics(values):
    e = values - 1.0
    return {
        "mean": float(values.mean()),
        "mean_abs_error": float(np.mean(np.abs(e))),
        "std": float(values.std()),
        "p95_abs_error": float(np.quantile(np.abs(e), .95)),
        "mse": float(np.mean(e * e)),
    }


def main():
    n = 64
    m = 16
    x = rng.normal(size=n)
    x /= np.linalg.norm(x)
    # Haar benchmark uses independent Haar directions, with the same fixed x norm.
    trials = 500
    haar_vals = np.empty(trials)
    for k in range(trials):
        q = haar_projection(n, m, 10000 + k)
        haar_vals[k] = np.sum((q @ x) ** 2)
    haar = metrics(haar_vals)

    # Prediction 1: every finite Kac product is orthogonal, independently of T.
    orth = []
    # Prediction 2: mean scaled subset energy is approximately one and approaches one.
    # Prediction 3: quadratic-statistic spread approaches the Haar spread as T grows.
    sweep = [0, n // 4, n // 2, n, 2*n, 4*n, 8*n, 16*n]
    rows = []
    for T in sweep:
        vals = np.empty(trials)
        defects = np.empty(30)
        for k in range(trials):
            rot = make_rotations(n, T, 20000 + 1000 * k + T)
            z = kac_apply(x, rot)
            vals[k] = (n / m) * np.sum(z[:m] ** 2)
            if k < len(defects):
                # Applying the same product to a few basis vectors estimates Q^T Q.
                Q = np.column_stack([kac_apply(np.eye(n)[j], rot) for j in range(n)])
                defects[k] = np.max(np.abs(Q.T @ Q - np.eye(n)))
        mm = metrics(vals)
        pred_mean = 1.0 + (float((n/m) * np.sum(x[:m]**2)) - 1.0) * ((n-1.0)/n)**T
        rows.append({"T": T, "predicted_mean": pred_mean, "mean_prediction_abs_error": abs(mm["mean"] - pred_mean), **mm, "std_over_haar_std": float(mm["std"] / haar["std"]),
                     "orthogonality_max_abs": float(defects.max())})
        orth.append(float(defects.max()))

    # Distortion benchmark on a common batch of pair differences.
    batch = 128
    X = rng.normal(size=(batch, n))
    X /= np.linalg.norm(X, axis=1, keepdims=True)
    pairs = X[::2] - X[1::2]
    # Compare the same fixed Kac transform to dense references at practical T.
    Tbench = 8 * n
    rot = make_rotations(n, Tbench, 777)
    t0 = time.perf_counter(); Yk = np.sqrt(n/m) * kac_apply(X, rot)[:, :m]; kac_time = time.perf_counter()-t0
    G = gaussian_projection(n, m, 778)
    t0 = time.perf_counter(); Yg = X @ G.T; gauss_time = time.perf_counter()-t0
    H = haar_projection(n, m, 779)
    t0 = time.perf_counter(); Yh = X @ H.T; haar_time = time.perf_counter()-t0
    def pair_dist(Y):
        return np.sum((Y[::2]-Y[1::2])**2, axis=1) / np.sum(pairs*pairs, axis=1)
    distortion = {"Kac": metrics(pair_dist(Yk)), "Gaussian": metrics(pair_dist(Yg)), "Haar": metrics(pair_dist(Yh))}
    # Repeat timings enough to avoid timer noise, while keeping the literal stream honest.
    reps = 5
    t0=time.perf_counter()
    for _ in range(reps): kac_apply(X, rot)
    kac_time=(time.perf_counter()-t0)/reps
    t0=time.perf_counter()
    for _ in range(reps): X @ G.T
    gauss_time=(time.perf_counter()-t0)/reps
    explicit_triple_bytes = int(Tbench * (4 + 4 + 8))
    dense_orthogonal_bytes = int(n * n * 8)
    storage = {"Kac_runtime_arrays_bytes": int(sum(a.nbytes for a in rot)),
               "Kac_explicit_i_j_theta_bytes": explicit_triple_bytes,
               "dense_projection_bytes": int(G.nbytes),
               "dense_full_orthogonal_bytes": dense_orthogonal_bytes,
               "dense_full_over_explicit_triples": float(dense_orthogonal_bytes / explicit_triple_bytes),
               "dense_projection_over_runtime_arrays": float(G.nbytes / sum(a.nbytes for a in rot))}
    result = {
        "seed": SEED, "n": n, "m": m, "trials": trials,
        "predictions": {
            "P1": "orthogonality defect remains at floating-point roundoff for every T",
            "P2": "scaled quadratic norm has mean near 1",
            "P3": "quadratic-statistic spread moves toward the Haar reference as T increases"
        },
        "haar_reference": haar, "sweep": rows,
        "prediction_checks": {
            "P1_max_defect_all_T": float(max(orth)),
            "P2_mean_range": [float(min(r["mean"] for r in rows)), float(max(r["mean"] for r in rows))],
            "P2_max_abs_error_vs_analytic_mean": float(max(r["mean_prediction_abs_error"] for r in rows)),
            "P3_std_ratio_T0": rows[0]["std_over_haar_std"],
            "P3_std_ratio_T8n": next(r["std_over_haar_std"] for r in rows if r["T"] == 8*n),
            "P3_std_ratio_T16n": rows[-1]["std_over_haar_std"]
        },
        "distortion": distortion,
        "timing_seconds_per_batch": {"Kac_literal": kac_time, "Gaussian_dense": gauss_time, "Haar_dense_not_timed": haar_time},
        "storage": storage
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()