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


def concentration_matrix(T, W):
    i = np.arange(T)
    d = i[:, None] - i[None, :]
    C = np.empty((T, T), dtype=float)
    nz = d != 0
    C[nz] = np.sin(2.0 * np.pi * W * d[nz]) / (np.pi * d[nz])
    C[~nz] = 2.0 * W
    return (C + C.T) * 0.5


def fourier_basis(T, r):
    # Complex unitary DFT columns, ordered by distance from zero frequency.
    n = np.arange(T)
    freqs = np.fft.fftfreq(T)
    order = np.argsort(np.abs(freqs), kind="stable")[:r]
    U = np.exp(2j * np.pi * np.outer(n, freqs[order])) / np.sqrt(T)
    return U


def projection_error(X, U):
    # X: samples x T; U: T x r orthonormal (possibly complex)
    Z = X @ U
    Xhat = Z @ U.conj().T
    return float(np.mean(np.sum(np.abs(X - Xhat) ** 2, axis=1) /
                     np.maximum(np.sum(np.abs(X) ** 2, axis=1), 1e-12)))


def main():
    rng = np.random.default_rng(2368)
    T, W = 96, 0.22
    delta = 0.1
    C = concentration_matrix(T, W)
    evals, Uall = np.linalg.eigh(C)
    order = np.argsort(evals)[::-1]
    evals, Uall = evals[order], Uall[:, order]

    # The finite concentration operator should be PSD, have eigenvalues in [0,1],
    # and have effective dimension approximately 2WT.
    c = 2.0 * W * T
    barL = math.log((1.0 - delta) / delta)
    asym_rank = c + barL / math.pi**2 * math.log(max(4.0 * math.pi**2 * c / barL, 1.000001))
    r_emp = int(np.sum(evals > delta))
    r_asym = int(np.clip(math.ceil(asym_rank), 1, T))

    # Draw finite-window samples from the band-limited covariance C. In this
    # distribution, the DPSS eigenvectors are the optimal rank-r coordinates.
    # A tiny jitter makes the square root robust to roundoff only.
    pos = np.maximum(evals, 0.0)
    X = rng.standard_normal((1200, T)) @ (Uall * np.sqrt(pos)).T

    ranks = [8, 12, 16, 20, 24, 28, 32, 40, 48]
    rows = []
    for r in ranks:
        Ud = Uall[:, :r]
        Uf = fourier_basis(T, r)
        Ur = np.linalg.qr(rng.standard_normal((T, r)))[0]
        rows.append({
            "rank": r,
            "dpss_error": projection_error(X, Ud),
            "fourier_error": projection_error(X.astype(complex), Uf),
            "random_error": projection_error(X, Ur),
        })

    # At the theorem-inspired threshold, also report the retained energy and
    # the predicted attention quadratic-work reduction.
    r = r_asym
    retained = float(np.sum(evals[:r]) / np.sum(evals))
    threshold_retained = float(np.sum(evals[:r_emp]) / np.sum(evals))
    result = {
        "config": {"T": T, "W": W, "delta": delta, "samples": len(X)},
        "math_check": {
            "min_eigenvalue": float(evals[-1]),
            "max_eigenvalue": float(evals[0]),
            "eigenvalues_in_unit_interval": bool(evals[0] <= 1.0 + 1e-10 and evals[-1] >= -1e-10),
            "trace": float(np.trace(C)),
            "time_bandwidth_2WT": c,
            "asymptotic_rank": float(asym_rank),
            "asymptotic_rank_clipped_ceiling": r_asym,
            "empirical_rank_above_delta": r_emp,
            "empirical_rank_relative_error": float(abs(r_emp - c) / c),
            "energy_retained_at_asym_rank": retained,
            "energy_retained_above_delta_rank": threshold_retained,
        },
        "curves": rows,
        "comparison_at_asym_rank": {
            "rank": r,
            "dpss_error": rows[[q["rank"] for q in rows].index(r)]["dpss_error"] if r in ranks else projection_error(X, Uall[:, :r]),
            "fourier_error": rows[[q["rank"] for q in rows].index(r)]["fourier_error"] if r in ranks else projection_error(X.astype(complex), fourier_basis(T, r)),
            "random_error": rows[[q["rank"] for q in rows].index(r)]["random_error"] if r in ranks else None,
            "quadratic_attention_work_fraction": float((r / T) ** 2),
            "quadratic_attention_work_reduction": float(1.0 - (r / T) ** 2),
        },
    }
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()