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

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


def covariance_cholesky(d, kind="power", gamma=.7, rho=.9):
    t = np.arange(d)
    if kind == "iid":
        c = np.zeros(d); c[0] = 1.
    elif kind == "ar":
        c = rho ** t
    elif kind == "power":
        # Positive definite stationary covariance: this is a finite Toeplitz
        # section of the mixture-like power-law kernel.
        c = (t + 1.) ** (-gamma)
    else:
        raise ValueError(kind)
    C = np.empty((d, d))
    for i in range(d):
        C[i, :] = c[np.abs(np.arange(d) - i)]
    # Small numerical jitter only protects Cholesky at large gamma/d.
    return np.linalg.cholesky(C + 1e-10 * np.eye(d)), c


def sample_weight(d, kind="power", gamma=.7, rho=.9, rows=None):
    rows = rows or d
    L, c = covariance_cholesky(d, kind, gamma, rho)
    z = rng.standard_normal((rows, d)) @ L.T
    return z / math.sqrt(d), c


def theoretical_m2(d, c):
    # Excludes lag zero, matching the paper's sum_{t=1}^N c_t^2.
    return float(np.sum(c[1:] ** 2))


def matrix_stats(W):
    s = np.linalg.svd(W, compute_uv=False)
    gram = W @ W.T
    m4 = float(np.trace(gram @ gram) / W.shape[0])
    return float(s[0]), m4, float(np.mean(W * W))


def verify_scaling(widths=(32, 64, 128, 256, 512), reps=12):
    out = {"power": {}, "ar": {}}
    for gamma in (.3, .5, .7):
        vals = []
        for d in widths:
            c = covariance_cholesky(d, "power", gamma)[1]
            vals.append(theoretical_m2(d, c))
        x, y = np.log(np.asarray(widths)), np.log(np.asarray(vals))
        slope = float(np.polyfit(x, y, 1)[0])
        # Boundary prediction uses a log-vs-width regression, not a power fit.
        log_slope = float(np.polyfit(np.log(widths), vals, 1)[0])
        out["power"][str(gamma)] = {
            "widths": list(widths), "m2": vals, "loglog_slope": slope,
            "predicted_slope": max(0., 1. - 2.*gamma),
            "boundary_log_slope": log_slope if gamma == .5 else None,
        }
    for rho in (.5, .9, .99):
        d = 512
        c = covariance_cholesky(d, "ar", rho=rho)[1]
        observed = theoretical_m2(d, c)
        predicted = rho*rho * (1-rho**(2*(d-1))) / (1-rho*rho)
        out["ar"][str(rho)] = {"observed_m2": observed,
                                "predicted_m2": float(predicted),
                                "relative_error": abs(observed-predicted)/predicted}
    return out


def spectral_sweep(widths=(64, 128, 256), reps=8):
    rows = []
    for d in widths:
        for kind, kwargs in [("iid", {}), ("power", {"gamma": .3}),
                             ("power", {"gamma": .7}), ("ar", {"rho": .9})]:
            stats = np.asarray([matrix_stats(sample_weight(d, kind, **kwargs)[0])
                                for _ in range(reps)])
            rows.append({"width": d, "kind": kind, **kwargs,
                         "max_sv_mean": float(stats[:,0].mean()),
                         "max_sv_std": float(stats[:,0].std()),
                         "m4_mean": float(stats[:,1].mean()),
                         "m4_std": float(stats[:,1].std()),
                         "entry_var": float(stats[:,2].mean())})
    return rows


def mlp_probe(width=128, steps=100, batch=96):
    # Same synthetic task and initialization seed for a small practical check.
    try:
        import torch
        torch.manual_seed(SEED)
        device = "cuda" if torch.cuda.is_available() else "cpu"
        X = torch.randn(1024, 32, device=device)
        y = ((X[:, :4].sum(1) + .5*X[:, 4:8].sum(1)) > 0).long()
        results = {}
        for name, kind, kw in [("iid", "iid", {}), ("safe_gamma07", "power", {"gamma": .7}),
                               ("unsafe_gamma03", "power", {"gamma": .3})]:
            torch.manual_seed(SEED + len(name))
            # Generate numpy correlated matrices, then use them as fixed initial values.
            w1, _ = sample_weight(width, kind, **kw, rows=width)
            w2, _ = sample_weight(width, kind, **kw, rows=2)
            net = torch.nn.Sequential(torch.nn.Linear(32, width), torch.nn.ReLU(),
                                      torch.nn.Linear(width, 2)).to(device)
            with torch.no_grad():
                # Input layer uses iid/structured rows with d=32; output is structured too.
                net[0].weight.copy_(torch.tensor(sample_weight(32, kind, **kw, rows=width)[0], device=device))
                net[0].bias.zero_(); net[2].weight.copy_(torch.tensor(w2, device=device)); net[2].bias.zero_()
            opt = torch.optim.SGD(net.parameters(), lr=.08)
            losses = []
            for step in range(steps):
                ix = torch.randint(0, len(X), (batch,), device=device)
                loss = torch.nn.functional.cross_entropy(net(X[ix]), y[ix])
                opt.zero_grad(); loss.backward(); opt.step()
                if step in (0, steps-1): losses.append(float(loss.detach().cpu()))
            results[name] = {"loss_step1": losses[0], "loss_step100": losses[1], "device": device}
        return results
    except Exception as e:
        return {"error": repr(e), "fallback": "spectral checks completed"}


def main():
    random.seed(SEED); np.random.seed(SEED)
    result = {"seed": SEED, "scaling_verification": verify_scaling(),
              "spectral_sweep": spectral_sweep(), "mlp_probe": mlp_probe()}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()