import json
import math
import os
import numpy as np
from scipy.linalg import hadamard

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


def signed_perm(d, rng):
    p = rng.permutation(d)
    s = rng.choice(np.array([-1.0, 1.0]), size=d)
    # Sigma z = s * z[p].
    S = np.zeros((d, d))
    S[np.arange(d), p] = s
    return S


def make_u(d, rng):
    H = hadamard(d).astype(float) / math.sqrt(d)
    A = H
    B = H
    return A @ signed_perm(d, rng) @ B


def coherence_sweep(dims=(16, 32, 64, 128, 256), trials=200):
    rows = []
    for d in dims:
        vals = []
        ortho = []
        for _ in range(trials):
            U = make_u(d, rng)
            vals.append(math.sqrt(d) * np.max(np.abs(U)))
            ortho.append(np.linalg.norm(U.T @ U - np.eye(d), ord='fro'))
        vals = np.asarray(vals)
        predicted = math.sqrt(2.0 * math.log(d))
        rows.append({
            'd': d,
            'coherence_mu_mean': float(vals.mean()),
            'coherence_mu_p95': float(np.quantile(vals, .95)),
            'predicted_sqrt_2logd': predicted,
            'ratio_mean_to_prediction': float(vals.mean() / predicted),
            'orthogonality_fro_mean': float(np.mean(ortho)),
        })
    return rows


def quantize(z, bits=4):
    q = 2 ** (bits - 1) - 1
    delta = np.max(np.abs(z), axis=-1, keepdims=True) / max(q, 1)
    delta = np.maximum(delta, 1e-12)
    return np.clip(np.rint(z / delta), -q, q) * delta


def make_vectors(d, n, sparsity, rng):
    # Each vector has k active coordinates with Gaussian values and one controlled
    # outlier scale. This stresses block-max quantization while retaining random orientation.
    k = max(1, int(round(d * sparsity)))
    X = np.zeros((n, d))
    for i in range(n):
        idx = rng.choice(d, size=k, replace=False)
        X[i, idx] = rng.normal(size=k)
    # Normalize every block so comparisons have identical input energy.
    X /= np.linalg.norm(X, axis=1, keepdims=True)
    return X


def quantization_sweep(dims=(64, 128, 256), sparsities=(1/256, 1/64, 1/16, 1/4, 1.0), n=1000, bits=4):
    rows = []
    for d in dims:
        # one transform per dimension, frozen as proposed
        U = make_u(d, rng)
        for sp in sparsities:
            X = make_vectors(d, n, sp, rng)
            # Baseline: coordinate quantization; permutation-only is mathematically
            # identical to this max-scale quantizer and is included as a check.
            Y0 = quantize(X, bits)
            perm = rng.permutation(d)
            Xperm = X[:, perm]
            Yp = quantize(Xperm, bits)[:, np.argsort(perm)]
            # Full preconditioner and exact inverse.
            Z = X @ U.T
            Zq = quantize(Z, bits)
            Xhat = Zq @ U
            mse0 = float(np.mean((Y0 - X) ** 2))
            msep = float(np.mean((Yp - X[:, rng.permutation(d)]) ** 2))
            msem = float(np.mean((Xhat - X) ** 2))
            # normalized by signal energy, which is 1 per row
            rows.append({
                'd': d, 'sparsity': sp, 'baseline_nmse': mse0,
                'perm_nmse': msep, 'mixed_nmse': msem,
                'mixed_gain_vs_baseline': mse0 / max(msem, 1e-30),
                'baseline_max_abs_mean': float(np.mean(np.max(np.abs(X), axis=1))),
                'mixed_max_abs_mean': float(np.mean(np.max(np.abs(Z), axis=1))),
            })
    return rows


def sampling_check(d=128, fractions=(1/32, 1/8, 1/2), trials=3000):
    # For unbiased coordinate sampling, theory predicts E||xhat-x||^2 = d/m - 1
    # for every unit vector. Mixing should additionally reduce the largest sampled
    # coordinate energy for sparse inputs, as predicted by incoherence.
    U = make_u(d, rng)
    x = np.zeros(d); x[0] = 1.0
    z = U @ x
    rows = []
    for frac in fractions:
        m = max(1, int(round(d * frac)))
        errs0, errsm = [], []
        for _ in range(trials):
            I = rng.choice(d, size=m, replace=False)
            y0 = np.zeros(d); y0[I] = x[I] * d / m
            yz = np.zeros(d); yz[I] = z[I] * d / m
            xr = U.T @ yz
            errs0.append(np.sum((y0-x)**2))
            errsm.append(np.sum((xr-x)**2))
        rows.append({'d': d, 'm': m, 'predicted_nmse': d/m-1,
                     'identity_mean_nmse': float(np.mean(errs0)),
                     'mixed_mean_nmse': float(np.mean(errsm)),
                     'mixed_p95_nmse': float(np.quantile(errsm,.95)),
                     'identity_max_coordinate_energy': 1.0,
                     'mixed_max_coordinate_energy': float(np.max(z*z))})
    return rows


def main():
    coh = coherence_sweep()
    quant = quantization_sweep()
    # A simple quantitative prediction from incoherence: for a k-sparse unit vector,
    # a dense incoherent transform should have max coordinate O(sqrt(log d/d)),
    # while identity has max=1. Report the observed ratio at the sparsest setting.
    prediction = []
    for d in (64, 128, 256):
        r = next(x for x in quant if x['d'] == d and abs(x['sparsity'] - 1/256) < 1e-12)
        predicted_max = math.sqrt(2 * math.log(d) / d)
        prediction.append({
            'd': d,
            'predicted_mixed_max_scale': predicted_max,
            'observed_mixed_max_abs': r['mixed_max_abs_mean'],
            'observed_identity_max_abs': r['baseline_max_abs_mean'],
            'predicted_identity_max_abs': 1.0,
            'observed_mixed_to_identity': r['mixed_max_abs_mean'] / r['baseline_max_abs_mean'],
        })
    # Prediction 1: coherence is approximately sqrt(2 log d / d), up to a
    # slowly varying constant; estimate the observed log-log slope after removing log.
    ds = np.array([x['d'] for x in coh], float)
    ys = np.array([x['coherence_mu_mean'] for x in coh]) / np.sqrt(ds)
    slope = float(np.polyfit(np.log(ds), np.log(ys), 1)[0])
    sampling = sampling_check()
    result = {'seed': SEED, 'coherence': coh, 'coherence_scaled_loglog_slope': slope,
              'quantization': quant, 'sparse_spreading_prediction': prediction,
              'sampling': sampling}
    with open('results.json', 'w') as f:
        json.dump(result, f, indent=2)
    print(json.dumps(result, indent=2))


if __name__ == '__main__':
    main()