import json
import time
import numpy as np
from scipy.sparse import csr_matrix
from scipy.sparse.linalg import spsolve


def shift_matrix(n, k=1):
    P = np.zeros((n, n))
    for i in range(n):
        P[(i + k) % n, i] = 1.0
    return P


def cycle_laplacian(n):
    return 2*np.eye(n) - shift_matrix(n, 1) - shift_matrix(n, -1)


def symmetry_basis(n):
    j = np.arange(n)[:, None]
    k = np.arange(n)[None, :]
    return np.exp(2j*np.pi*j*k/n) / np.sqrt(n)


def projector_checks(n):
    Ps = []
    for k in range(n):
        P = np.zeros((n, n), dtype=complex)
        for g in range(n):
            rho = shift_matrix(n, g)
            chi = np.exp(2j*np.pi*k*g/n)
            P += np.conj(chi) * rho / n
        Ps.append(P)
    return {
        "resolution_error": float(np.linalg.norm(sum(Ps)-np.eye(n))),
        "idempotence_error": float(max(np.linalg.norm(P@P-P) for P in Ps)),
        "orthogonality_error": float(max(np.linalg.norm(Ps[i]@Ps[j]) for i in range(n) for j in range(n) if i != j)),
    }


def matrices(n, theta, tau):
    L = cycle_laplacian(n)
    # Symmetry-compatible learned combination of identity, Laplacian, and L^2.
    H = 0.7*np.eye(n) + 0.4*L + 0.08*(L@L)
    K = np.eye(n) + tau * theta * H
    return L, K


def one_case(n, theta, tau, rng):
    L, K = matrices(n, theta, tau)
    Q = symmetry_basis(n)
    rho = shift_matrix(n, 1)
    comm = np.linalg.norm(rho@L-L@rho) / np.linalg.norm(L)
    Kh = Q.conj().T @ K @ Q
    offdiag = Kh - np.diag(np.diag(Kh))
    block_error = np.linalg.norm(offdiag) / np.linalg.norm(Kh)
    b = rng.normal(size=n)
    t0 = time.perf_counter()
    x_sparse = spsolve(csr_matrix(K), b)
    sparse_ms = 1000*(time.perf_counter()-t0)
    t0 = time.perf_counter()
    bh = Q.conj().T @ b
    # C_n irreps are one-dimensional: independent scalar block solves.
    xh = bh / np.diag(Kh)
    x_fourier = np.real_if_close(Q @ xh).real
    fourier_ms = 1000*(time.perf_counter()-t0)
    residual = np.linalg.norm(K@x_fourier-b)/np.linalg.norm(b)
    equivalence = np.linalg.norm(x_fourier-x_sparse)/np.linalg.norm(x_sparse)
    return comm, block_error, equivalence, residual, sparse_ms, fourier_ms


def main():
    rng = np.random.default_rng(1234)
    ns = [16, 32, 64, 128, 256]
    thetas = [0.0, 0.1, 1.0, 10.0]
    taus = [0.01, 0.1, 1.0]
    proj = projector_checks(16)
    rows = []
    for n in ns:
        for theta in thetas:
            for tau in taus:
                vals = one_case(n, theta, tau, rng)
                rows.append({"n":n,"theta":theta,"tau":tau,"comm":vals[0],"block_error":vals[1],"equivalence":vals[2],"residual":vals[3],"sparse_ms":vals[4],"fourier_ms":vals[5]})
    by_n = []
    for n in ns:
        rr = [r for r in rows if r["n"] == n]
        by_n.append({"n": n,
                     "median_sparse_ms": float(np.median([r["sparse_ms"] for r in rr])),
                     "median_fourier_ms": float(np.median([r["fourier_ms"] for r in rr])),
                     "fourier_over_sparse": float(np.median([r["fourier_ms"] for r in rr]) / np.median([r["sparse_ms"] for r in rr]))})
    tol = 1e-12
    predictions = [
        {"prediction": "commutator remains at roundoff for every theta,tau", "predicted_bound": tol,
         "observed_max": max(r["comm"] for r in rows), "confirmed": bool(max(r["comm"] for r in rows) < tol)},
        {"prediction": "symmetry basis removes off-block entries for every theta,tau", "predicted_bound": tol,
         "observed_max": max(r["block_error"] for r in rows), "confirmed": bool(max(r["block_error"] for r in rows) < tol)},
        {"prediction": "transformed solve equals original solve for every theta,tau", "predicted_bound": tol,
         "observed_max": max(r["equivalence"] for r in rows), "confirmed": bool(max(r["equivalence"] for r in rows) < tol)},
    ]
    summary = {
        "sweep": {"n": ns, "theta": thetas, "tau": taus},
        "predictions": predictions,
        "timing_by_n": by_n,
        "projectors": proj,
        "max_commutator": max(r["comm"] for r in rows),
        "max_block_error": max(r["block_error"] for r in rows),
        "max_solution_relative_error": max(r["equivalence"] for r in rows),
        "max_residual": max(r["residual"] for r in rows),
        "median_sparse_ms": float(np.median([r["sparse_ms"] for r in rows])),
        "median_fourier_ms": float(np.median([r["fourier_ms"] for r in rows])),
        "rows": rows,
    }
    with open("results.json", "w") as f:
        json.dump(summary, f, indent=2)
    print(json.dumps({k:v for k,v in summary.items() if k != "rows"}, indent=2))


if __name__ == "__main__":
    main()