import json, math, random, time
from pathlib import Path
import numpy as np
import torch
import torch.nn as nn

SEED = 619
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.set_num_threads(4)
device = "cuda" if torch.cuda.is_available() else "cpu"
try:
    if device == "cuda":
        torch.cuda.reset_peak_memory_stats()
except Exception:
    device = "cpu"


def make_accretive(d, delta=0.35, skew=1.8, seed=1):
    g = np.random.default_rng(seed)
    L = g.normal(size=(d, max(2, d // 2))) / math.sqrt(d)
    S = L @ L.T + delta * np.eye(d)
    U = g.normal(size=(d, d)) / math.sqrt(d)
    K = skew * (U - U.T) / 2
    return S + K


def exact_fractional(A, r):
    vals, vecs = np.linalg.eig(A.astype(np.complex128))
    return (vecs @ np.diag(vals ** r) @ np.linalg.inv(vecs)).real


def resolvent_fractional(A, r=0.5, m=10, eps=0.02, lo=1e-3, hi=100.0):
    # Log-grid midpoint quadrature for the stated integral after the
    # cancellation-safe identity lambda^-1 (lambda I+A)^-1 A.
    x = np.linspace(math.log(lo), math.log(hi), m + 1)
    lam = np.exp((x[:-1] + x[1:]) / 2)
    weights = lam * (x[1:] - x[:-1])
    Ae = A + eps * np.eye(A.shape[0])
    out = np.zeros_like(Ae, dtype=float)
    I = np.eye(A.shape[0])
    for l, w in zip(lam, weights):
        out += w * (l ** (r - 1.0)) * np.linalg.solve(l * I + Ae, Ae)
    return math.sin(math.pi * r) / math.pi * out


def math_check():
    rows = []
    for skew in [0.0, 1.8, 4.0]:
        A = make_accretive(8, skew=skew, seed=10)
        S = (A + A.T) / 2
        min_sym = np.linalg.eigvalsh(S).min()
        ref = exact_fractional(A, 0.5)
        approx = resolvent_fractional(A, m=10, eps=0.0, lo=1e-4, hi=300)
        rel = np.linalg.norm(approx-ref) / np.linalg.norm(ref)
        # Positive shifts should remain well-conditioned under accretivity.
        conds = [np.linalg.cond(l*np.eye(8)+A) for l in np.geomspace(1e-3, 100, 10)]
        rows.append({"skew": skew, "min_sym_eigenvalue": float(min_sym),
                     "relative_matrix_error": float(rel),
                     "max_shifted_condition": float(max(conds))})
    # Scalar claim: same quadrature approximates t^r over the sampled range.
    ts = np.geomspace(0.35, 30.0, 400)
    xs = np.linspace(math.log(1e-4), math.log(300), 11)
    lam = np.exp((xs[:-1]+xs[1:])/2)
    weights = lam*(xs[1:]-xs[:-1])
    pred = np.zeros_like(ts)
    for l,w in zip(lam,weights): pred += w*l**(-0.5)*ts/(l+ts)
    pred *= 1/math.pi
    scalar_rel = float(np.max(np.abs(pred-np.sqrt(ts))/np.sqrt(ts)))
    return {"matrices": rows, "scalar_max_relative_error": scalar_rel}


class ResolventBlock(nn.Module):
    def __init__(self, d, r=0.5, m=6, eps=0.03, lo=0.08, hi=8.0):
        super().__init__()
        self.d, self.r, self.m = d, r, m
        self.eps, self.lo, self.hi = eps, lo, hi
        self.L = nn.Parameter(torch.randn(d, d//2) / math.sqrt(d))
        self.U = nn.Parameter(torch.randn(d, d) / math.sqrt(d))
        self.gain = nn.Parameter(torch.tensor(0.15))
        self.register_buffer("eye", torch.eye(d))

    def operator(self):
        S = self.L @ self.L.T + 0.20*self.eye
        K = (self.U-self.U.T)/2
        return S + K + self.eps*self.eye

    def forward(self, x):
        A = self.operator()
        xx = x.T
        grid = torch.linspace(math.log(self.lo), math.log(self.hi), self.m+1, device=x.device, dtype=x.dtype)
        lam = torch.exp((grid[:-1]+grid[1:])/2)
        weights = lam*(grid[1:]-grid[:-1])
        z = torch.zeros_like(xx)
        for l,w in zip(lam, weights):
            rhs = A @ xx
            z = z + w*l**(self.r-1)*torch.linalg.solve(l*self.eye.to(x)+A, rhs)
        z = math.sin(math.pi*self.r)/math.pi*z
        return x + self.gain*z.T


class LinearBlock(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.W = nn.Parameter(torch.randn(d,d)/math.sqrt(d))
        self.gain = nn.Parameter(torch.tensor(0.15))
    def forward(self,x):
        return x + self.gain*(x @ self.W.T)


class Net(nn.Module):
    def __init__(self, block, d=12, classes=3):
        super().__init__()
        self.inp=nn.Linear(2,d)
        self.block=block
        self.out=nn.Linear(d,classes)
    def forward(self,x): return self.out(torch.tanh(self.block(torch.tanh(self.inp(x)))))


def training_check():
    g=torch.Generator().manual_seed(SEED)
    n=360
    centers=torch.tensor([[1.5,0.0],[-0.75,1.3],[-0.75,-1.3]])
    y=torch.randint(0,3,(n,),generator=g)
    x=centers[y]+0.42*torch.randn(n,2,generator=g)
    perm=torch.randperm(n,generator=g); tr,va=perm[:270],perm[270:]
    results={}
    for name, factory in [("baseline",lambda:LinearBlock(12)),("resolvent",lambda:ResolventBlock(12))]:
        torch.manual_seed(SEED+ (0 if name=='baseline' else 1))
        net=Net(factory()).to(device)
        opt=torch.optim.Adam(net.parameters(),lr=0.025)
        losses=[]; gradvars=[]; t0=time.time()
        for step in range(180):
            opt.zero_grad(set_to_none=True)
            logits=net(x[tr].to(device)); loss=nn.functional.cross_entropy(logits,y[tr].to(device))
            loss.backward()
            gn=torch.sqrt(sum((p.grad.detach()**2).sum() for p in net.parameters() if p.grad is not None))
            gradvars.append(float(gn)); opt.step(); losses.append(float(loss))
        with torch.no_grad():
            pred=net(x[va].to(device)).argmax(1).cpu(); acc=float((pred==y[va]).float().mean())
        results[name]={"final_train_loss":losses[-1],"loss_at_60":losses[59],"validation_accuracy":acc,
                       "gradient_norm_variance":float(np.var(gradvars)),"seconds":time.time()-t0,
                       "nan_or_inf":bool(not np.isfinite(losses).all())}
    return results


if __name__ == "__main__":
    out={"device":device,"math":math_check(),"training":training_check()}
    Path("results.json").write_text(json.dumps(out,indent=2))
    print(json.dumps(out,indent=2))