import json, math, time
import numpy as np
from scipy.linalg import eigh

SEED = 1226


def unitary(A, t):
    w, V = eigh(A)
    return (V * np.exp(-1j * t * w)) @ V.T


def dyadic_exact(A, lam, K):
    n = A.shape[0]
    I = np.eye(n, dtype=complex)
    R = 1j * np.linalg.inv(I - math.exp(-lam) * unitary(A, 1.0))
    for k in range(1, K + 1):
        t = 2.0 ** (-k)
        R -= 1j * t * np.linalg.inv(I + math.exp(-lam * t) * unitary(A, t))
    return R


def dyadic_neumann(A, X, lam, K, J):
    Y = np.zeros_like(X, dtype=complex)
    U = unitary(A, 1.0)
    Z = X.astype(complex)
    for j in range(J + 1):
        Y += 1j * math.exp(-j * lam) * Z
        Z = U @ Z
    for k in range(1, K + 1):
        t = 2.0 ** (-k)
        c = math.exp(-lam * t)
        U = unitary(A, t)
        Z = X.astype(complex)
        for j in range(J + 1):
            Y -= 1j * t * ((-c) ** j) * Z
            Z = U @ Z
    return Y


def scalar_verification():
    vals = np.linspace(-2.0, 2.0, 17)
    A = np.diag(vals)
    out = {}
    for lam in [0.2, 0.5, 1.0, 2.0]:
        target = np.diag(1.0 / (vals - 1j * lam))
        errs = []
        for K in [2, 4, 6, 8, 10, 12]:
            errs.append(float(np.max(np.abs(dyadic_exact(A, lam, K) - target))))
        # Exact resolvent tail is bounded by sum_{k>K} 2^-k = 2^-K,
        # since ||(I+cU)^-1|| <= 1/(1-c) is not uniformly useful; empirically
        # the measured tail should decrease with K and remain below a constant*2^-K.
        ratios = [errs[i] / errs[i + 1] for i in range(len(errs)-1)]
        # Neumann prediction: at scale t, remainder is at most c^(J+1)/(1-c).
        nerrs = []
        for J in [1, 2, 4, 8, 12, 16]:
            X = np.eye(len(vals))
            nerrs.append(float(np.max(np.abs(dyadic_neumann(A, X, lam, 10, J) - dyadic_exact(A, lam, 10)))))
        # Unitary norm prediction over random vectors and times.
        x = np.random.default_rng(7).normal(size=len(vals))
        norm_errors = [abs(np.linalg.norm(unitary(A, t) @ x) / np.linalg.norm(x) - 1) for t in [1, .5, .125, .01]]
        out[str(lam)] = {"K": [2,4,6,8,10,12], "exact_errors": errs,
                         "K_error_ratios": ratios, "J": [1,2,4,8,12,16],
                         "neumann_errors": nerrs, "unitary_relative_norm_errors": norm_errors}
    return out


def graph_data(n=80, p_in=.16, p_out=.035, d=8):
    rng = np.random.default_rng(SEED)
    y = np.repeat([0, 1], n // 2)
    G = np.zeros((n, n), float)
    for i in range(n):
        for j in range(i+1, n):
            p = p_in if y[i] == y[j] else p_out
            if rng.random() < p: G[i,j] = G[j,i] = 1
    deg = G.sum(1)
    L = np.eye(n) - (G / np.sqrt(np.maximum(deg, 1)[:,None] * np.maximum(deg, 1)[None,:]))
    # Features include weak signal; held-out nodes use same fixed feature matrix.
    X = rng.normal(0, 1, (n, d))
    X[:, 0] += (2*y-1) * .65
    perm = rng.permutation(n)
    tr, va = perm[:n//2], perm[n//2:]
    return L, X, y, tr, va


def train_models():
    # Tiny ridge classifier compares standard low-degree polynomial filtering to dyadic filtering.
    L, X, y, tr, va = graph_data()
    # Rescale normalized Laplacian to [-1,1] (spectrum is approximately [0,2]).
    A = L - np.eye(len(L))
    rng = np.random.default_rng(SEED)
    results = {}
    configs = [("baseline_poly", 6), ("dyadic_K2_J2", 6), ("dyadic_K4_J2", 6)]
    for name, budget in configs:
        t0 = time.perf_counter()
        if name == "baseline_poly":
            Z = X.astype(complex); F = [Z]
            U = A @ Z
            for _ in range(budget):
                F.append(U); U = A @ U
            H = np.concatenate([z.real for z in F], axis=1)
        else:
            K = int(name.split('K')[1].split('_')[0]); J = int(name.split('J')[1])
            Hc = dyadic_neumann(A, X, .8, K, J)
            H = np.concatenate([X, Hc.real, Hc.imag], axis=1)
        # closed-form ridge, no training randomness
        lam = 1.0
        W = np.linalg.solve(H[tr].T @ H[tr] + lam*np.eye(H.shape[1]), H[tr].T @ (2*y[tr]-1))
        pred = (H @ W > 0).astype(int)
        acc = float(np.mean(pred[va] == y[va]))
        norm = float(np.linalg.norm(H) / np.linalg.norm(X))
        results[name] = {"val_accuracy": acc, "feature_norm_ratio": norm,
                         "seconds": time.perf_counter()-t0, "sparse_op_budget": budget}
    return results


def main():
    np.random.seed(SEED)
    report = {"seed": SEED, "verification": scalar_verification(), "mini_experiment": train_models()}
    with open("results.json", "w") as f: json.dump(report, f, indent=2)
    print(json.dumps(report, indent=2))

if __name__ == "__main__": main()