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

SEED = 3035
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
try:
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
except Exception:
    device = 'cpu'


def laplacian(A, psi, alpha):
    C = A * np.cos(psi[None, :] - psi[:, None] - alpha)
    return np.diag(C.sum(1)) - C


def core_verification():
    n = 6; K = 1.0
    # Symmetric positive graph: exact real spectrum and one gauge eigenvalue.
    A = np.full((n, n), 0.18); np.fill_diagonal(A, 0.0)
    psi = np.zeros(n); alpha = np.zeros((n, n))
    L = laplacian(A, psi, alpha)
    ev = np.linalg.eigvals(L)
    nonzero = np.sort(np.real(ev))[1:]
    lam_max = float(np.max(np.real(ev)))
    eta_c = 2.0 / (K * lam_max)
    # Test actual discrete perturbation amplification around the predicted boundary.
    q = np.random.default_rng(SEED).normal(size=n); q -= q.mean(); q /= np.linalg.norm(q)
    observed = []
    for ratio in [0.8, 1.0, 1.2]:
        eta = ratio * eta_c
        M = np.eye(n) - eta*K*L
        # maximum non-gauge one-step amplification is the relevant observed quantity
        vals = np.linalg.eigvals(M)
        amp = float(np.max(np.abs(vals[np.abs(vals-1)>1e-7])))
        observed.append({'ratio_to_eta_c': ratio, 'max_mode_amplification': amp})
    return {'eigenvalues': ev.tolist(), 'min_nonzero_real_eigenvalue': float(nonzero[0]),
            'max_eigenvalue': lam_max, 'predicted_eta_c': eta_c, 'euler_tests': observed,
            'boundary_crossed': observed[0]['max_mode_amplification'] < 1 and observed[2]['max_mode_amplification'] > 1}


class PhaseRNN(torch.nn.Module):
    def __init__(self, n=8, inp=3, dt=0.12):
        super().__init__(); self.n=n; self.dt=dt
        self.a = torch.nn.Parameter(torch.full((n,n), -1.2) + .08*torch.randn(n,n))
        self.raw_alpha = torch.nn.Parameter(.5*torch.randn(n,n))
        self.W = torch.nn.Parameter(.15*torch.randn(n, inp))
        self.b = torch.nn.Parameter(torch.zeros(n))
    def matrices(self):
        A = torch.nn.functional.softplus(self.a); A = A * (1-torch.eye(self.n, device=A.device))
        alpha = math.pi * torch.tanh(self.raw_alpha)
        return A, alpha
    def step(self, theta, x):
        A, alpha = self.matrices()
        # theta: batch,n; x: batch,input
        diff = theta[:,None,:] - theta[:,:,None] - alpha[None,:,:]
        coupling = (A[None,:,:] * torch.sin(diff)).sum(2)
        omega = x @ self.W.T + self.b
        return theta + self.dt*(omega + coupling)
    def rollout(self, x):
        th = torch.zeros(x.shape[0], self.n, device=x.device); out=[]
        for t in range(x.shape[1]):
            th = self.step(th, x[:,t]); out.append(th)
        return torch.stack(out, 1)
    def spectral_terms(self, psi, K=1.0, eta=None, gamma=.03):
        A, alpha = self.matrices()
        C = A * torch.cos(psi[None,:]-psi[:,None]-alpha)
        L = torch.diag(C.sum(1)) - C
        ev = torch.linalg.eigvals(L)
        # Remove gauge mode by excluding eigenvalue closest to zero.
        idx = torch.argmin(torch.abs(ev)); keep = torch.ones(self.n, dtype=torch.bool, device=ev.device); keep[idx]=False
        re = ev.real[keep]
        margin = torch.nn.functional.softplus(torch.as_tensor(gamma, device=psi.device)-re.min())
        if eta is None: eta=self.dt
        amp = torch.abs(1-eta*K*ev)
        # Gauge is exactly one; exclude it from stability penalty.
        excess = torch.relu(amp[keep]-1).pow(2).mean()
        return margin + excess, float(re.min().detach().cpu()), float(amp[keep].max().detach().cpu())


def train_case(use_spec, x, y, init_state, steps=350):
    model=PhaseRNN().to(device); model.load_state_dict(init_state)
    opt=torch.optim.Adam(model.parameters(), lr=.025)
    history=[]
    for it in range(steps):
        pred=model.rollout(x)
        task=torch.mean((torch.sin(pred)-y)**2)
        psi=pred.detach().mean((0,1)) if not use_spec else pred.mean((0,1))
        spec, margin, amp=model.spectral_terms(psi)
        loss=task + (.08*spec if use_spec else 0.0)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0); opt.step()
        if it in (0, steps-1): history.append((float(task.detach().cpu()), margin, amp))
    with torch.no_grad():
        pred=model.rollout(x); task=float(torch.mean((torch.sin(pred)-y)**2).cpu())
        psi=pred.mean((0,1)); spec, margin, amp=model.spectral_terms(psi)
        # finite perturbation test, quotienting out common phase by zero-mean perturbation
        th=pred[-1,-1].clone(); d=torch.randn_like(th); d-=d.mean(); d=d/d.norm()*1e-3
        th2=th+d
        for _ in range(12): th= model.step(th[None,:], torch.zeros(1,x.shape[2],device=device))[0]; th2=model.step(th2[None,:], torch.zeros(1,x.shape[2],device=device))[0]
        decay=float((th2-th).norm().cpu()/1e-3)
    return {'task_mse':task, 'spectral_loss':float(spec.cpu()), 'min_real_eigenvalue':float(margin), 'max_euler_amplification':float(amp), '12_step_perturbation_gain':decay, 'history':history}


def main():
    ver=core_verification()
    torch.manual_seed(SEED)
    T=24; B=64; inp=3; n=8
    x=torch.randn(B,T,inp,device=device)
    # Stable teacher target: smooth bounded phase observations.
    teacher=torch.zeros(B,n,device=device); ys=[]
    Wt=torch.randn(n,inp,device=device)*.18
    for t in range(T):
        teacher=teacher + .12*(x[:,t]@Wt.T - .35*torch.sin(teacher))
        ys.append(torch.sin(teacher))
    y=torch.stack(ys,1)
    torch.manual_seed(SEED+1); init=PhaseRNN(n,inp).state_dict()
    # Recreate same parameter initialization for both cases.
    baseline=train_case(False,x,y,init)
    idea=train_case(True,x,y,init)
    result={'device':device,'verification':ver,'baseline':baseline,'idea':idea,
            'claim_signal': idea['12_step_perturbation_gain'] < baseline['12_step_perturbation_gain'] and idea['max_euler_amplification'] <= baseline['max_euler_amplification']}
    Path('results.json').write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__=='__main__': main()