import json, time, math, random
import numpy as np
import torch

SEED = 429
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(min(8, torch.get_num_threads()))
try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
except Exception:
    device = torch.device("cpu")

# Finite Caratheodory-kernel quantities. The block K is assembled for scalar output.
def ck_quantities(A, B, C, lambdas, eps=1e-5):
    n = A.shape[0]
    dtype = torch.complex64 if A.dtype == torch.float32 else torch.complex128
    Ac = A.to(dtype); Bc = B.to(dtype); Cc = C.to(dtype)
    eye = torch.eye(n, dtype=dtype, device=A.device)
    phis = []
    for z in lambdas:
        z = torch.as_tensor(z, dtype=dtype, device=A.device)
        X = torch.linalg.solve(eye - z * Ac, Bc[:, None])
        phis.append((1.0 + (Cc[None, :] @ X).squeeze()).reshape(()))
    phis = torch.stack(phis)
    H = phis.real
    # For scalar output K is M by M; this is the exact block formula.
    z = torch.as_tensor(lambdas, dtype=dtype, device=A.device)
    K = (phis[:, None] + phis.conj()[None, :]) / (2 * (1 - z[:, None] * z.conj()[None, :]))
    K = (K + K.conj().T) / 2
    hmin = H.min()
    kmin = torch.linalg.eigvalsh(K).real.min()
    penalty = torch.relu(eps - H).square().mean() + torch.relu(eps - kmin).square()
    return penalty, hmin, kmin

def make_lambdas(M=12, rho=.93):
    # fixed points, with a modest radial spread as suggested in the idea
    return np.array([(rho + .035*np.sin(2*np.pi*j/M))*np.exp(2j*np.pi*j/M) for j in range(M)])

# Cheap direct verification: a passive constant Phi has PSD H and PSD kernel;
# a deliberately non-passive constant Phi violates both, and gradient descent repairs it.
def math_check():
    lam = make_lambdas(12)
    A = torch.zeros(1); B = torch.ones(1); Cgood = torch.ones(1)
    _, hg, kg = ck_quantities(A, B, Cgood, lam)
    Cbad = torch.tensor([-1.4], requires_grad=True)
    before = ck_quantities(A, B, Cbad, lam)
    opt = torch.optim.SGD([Cbad], lr=.25)
    for _ in range(30):
        opt.zero_grad(); loss, _, _ = ck_quantities(A, B, Cbad, lam); loss.backward(); opt.step()
    after = ck_quantities(A, B, Cbad, lam)
    return dict(good_hmin=float(hg), good_kmin=float(kg), bad_hmin=float(before[1]),
                bad_kmin=float(before[2]), repaired_hmin=float(after[1]),
                repaired_kmin=float(after[2]), repaired_C=float(Cbad.detach()))

class LinearSSM(torch.nn.Module):
    def __init__(self, n, init_scale=.9):
        super().__init__()
        self.A = torch.nn.Parameter(init_scale * torch.randn(n, n) / math.sqrt(n))
        self.B = torch.nn.Parameter(torch.randn(n) / math.sqrt(n))
        self.C = torch.nn.Parameter(torch.randn(n) / math.sqrt(n))
    def forward(self, x):
        # x: batch,time; zero initial state
        h = torch.zeros(x.shape[0], self.A.shape[0], device=x.device)
        ys=[]
        for t in range(x.shape[1]):
            h = h @ self.A.T + x[:, t:t+1] * self.B[None, :]
            ys.append((h * self.C[None, :]).sum(1))
        return torch.stack(ys, 1)

def target_data(A, B, C, batch, length, device):
    x = torch.randn(batch, length, device=device)
    h = torch.zeros(batch, A.shape[0], device=device); ys=[]
    for t in range(length):
        h = h @ A.T + x[:, t:t+1] * B[None, :]
        ys.append((h*C[None,:]).sum(1))
    return x, torch.stack(ys, 1)

def run_one(kind, seed, device):
    torch.manual_seed(seed); np.random.seed(seed)
    n=8; T=24; Tlong=96
    # Stable teacher; train data and long evaluation use independent fixed draws.
    diag = torch.linspace(.55,.88,n,device=device)
    At = torch.diag(diag); Bt = torch.linspace(.2,.8,n,device=device)/math.sqrt(n); Ct = torch.sin(torch.arange(n,device=device)+1)/math.sqrt(n)
    xtr,ytr=target_data(At,Bt,Ct,96,T,device); xev,yev=target_data(At,Bt,Ct,96,Tlong,device)
    model=LinearSSM(n).to(device)
    opt=torch.optim.Adam(model.parameters(),lr=.025)
    lam=make_lambdas()
    t0=time.perf_counter()
    for step in range(350):
        opt.zero_grad(); pred=model(xtr); task=(pred-ytr).square().mean()
        if kind == 'ck':
            ck,_,_=ck_quantities(model.A,model.B,model.C,lam,eps=1e-4)
            loss=task + .15*ck
        else: loss=task
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
    elapsed=time.perf_counter()-t0
    with torch.no_grad():
        train=((model(xtr)-ytr)**2).mean().item(); long=((model(xev)-yev)**2).mean().item()
        ck,hmin,kmin=ck_quantities(model.A,model.B,model.C,lam,eps=1e-4)
        rho=torch.linalg.eigvals(model.A).abs().max().item()
    return dict(train_mse=train,long_mse=long,hmin=float(hmin),kmin=float(kmin),ck_penalty=float(ck),spectral_radius=rho,seconds=elapsed)

def main():
    check=math_check(); results={}
    for kind in ['baseline','ck']:
        runs=[run_one(kind,s,device) for s in [SEED,SEED+1,SEED+2]]
        results[kind]={k:float(np.mean([r[k] for r in runs])) for k in runs[0]}
        results[kind]['runs']=runs
    results['device']=str(device); results['math_check']=check
    results['overhead_ratio']=results['ck']['seconds']/results['baseline']['seconds']
    with open('results.json','w') as f: json.dump(results,f,indent=2)
    print(json.dumps(results,indent=2))

if __name__ == '__main__': main()