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

SEED = 137
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)

def discrete_gramians(A, B, C, T):
    n = A.shape[0]
    Wo = np.zeros((n, n)); Phi = np.eye(n)
    for _ in range(T):
        Wo += Phi.T @ C.T @ C @ Phi
        Phi = A @ Phi
    Wr = np.zeros((n, n))
    for k in range(T):
        P = np.linalg.matrix_power(A, T - 1 - k)
        Wr += P @ B @ B.T @ P.T
    return (Wo + Wo.T) / 2, (Wr + Wr.T) / 2

def norm_min(W):
    tr = np.trace(W)
    return float(np.min(np.linalg.eigvalsh(W)) / (tr / W.shape[0])) if tr > 1e-12 else 0.0

def mechanism_checks():
    # With A=aI and C=diag(1,r), Wo is proportional to diag(1,r^2).
    # Therefore normalized lambda_min is exactly 2r^2/(1+r^2), a quantitative
    # parameter-scaling prediction. The same construction applies to Wr.
    A = np.diag([.82, .82]); T = 30
    ratios = [0.0, .05, .1, .2, .4, .7, 1.0]
    obs = []
    reach = []
    for r in ratios:
        C = np.diag([1., r]); Wo, _ = discrete_gramians(A, np.eye(2), C, T)
        _, Wr = discrete_gramians(A, np.diag([1., r]), C, T)
        pred = 2*r*r/(1+r*r) if r else 0.0
        obs.append({'coupling':r, 'normalized_lambda_min':norm_min(Wo), 'predicted':pred})
        reach.append({'coupling':r, 'normalized_lambda_min':norm_min(Wr), 'predicted':pred})
    # Direct quadrature verifies the continuous-time formula for a scalar mode.
    a=.18; c=.7; b=.55; dt=.002; horizon=2.
    times=np.arange(0, horizon, dt)
    wo_num=float(np.sum(dt*(c*np.exp(-a*times))**2))
    wr_num=float(np.sum(dt*(b*np.exp(-a*(horizon-times)))**2))
    wo_exact=c*c*(1-np.exp(-2*a*horizon))/(2*a)
    wr_exact=b*b*(1-np.exp(-2*a*horizon))/(2*a)
    # Measurement noise prediction: least-squares covariance is sigma^2 Wo^-1.
    noise=.03; samples=12000; est=[]; rng=np.random.default_rng(SEED)
    for r in [.08,.12,.2,.3,.5]:
        C=np.diag([1., r]); Wo,_=discrete_gramians(A,np.zeros((2,2)),C,T)
        H=np.vstack([C @ np.linalg.matrix_power(A,k) for k in range(T)])
        z=rng.standard_normal((2,samples)); ys=H@z+noise*rng.standard_normal((2*T,samples))
        zhat=np.linalg.solve(H.T@H, H.T@ys).T
        mse=float(np.mean((zhat-z.T)**2)); lam=float(np.min(np.linalg.eigvalsh(Wo)))
        est.append({'coupling':r, 'lambda_min':lam, 'mse':mse, 'predicted_sigma2_over_lambda':noise*noise/lam})
    slope=float(np.polyfit(np.log([x['lambda_min'] for x in est]), np.log([x['mse'] for x in est]), 1)[0])
    return {'observation_scaling':obs, 'reachability_scaling':reach,
            'scaling_max_abs_error_observation':max(abs(x['normalized_lambda_min']-x['predicted']) for x in obs),
            'scaling_max_abs_error_reachability':max(abs(x['normalized_lambda_min']-x['predicted']) for x in reach),
            'integral_check':{'wo_numeric':wo_num,'wo_exact':wo_exact,'wr_numeric':wr_num,'wr_exact':wr_exact},
            'noise_inverse_scaling':est, 'log_mse_vs_log_lambda_slope':slope}

class LinearSSM(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.A=torch.nn.Parameter(torch.tensor([[.75,.02],[.01,.65]]))
        self.B=torch.nn.Parameter(torch.randn(2,1)*.1)
        self.C=torch.nn.Parameter(torch.randn(1,2)*.1)
    def forward(self,u):
        z=torch.zeros(u.shape[0],2,device=u.device); ys=[]
        for k in range(u.shape[1]):
            z=z@self.A.T+u[:,k,0:1]*self.B.T; ys.append(z@self.C.T)
        return torch.stack(ys,1)
    def grams(self,T):
        phi=torch.eye(2,device=self.A.device); wo=torch.zeros((2,2),device=self.A.device)
        for _ in range(T): wo=wo+phi.T@self.C.T@self.C@phi; phi=self.A@phi
        wr=torch.zeros((2,2),device=self.A.device)
        for k in range(T):
            p=torch.linalg.matrix_power(self.A,T-1-k); wr=wr+p@self.B@self.B.T@p.T
        return (wo+wo.T)/2,(wr+wr.T)/2

def train(reg, device):
    torch.manual_seed(SEED + int(reg*1000)); n=256; T=20
    A0=torch.tensor([[.90,0.],[0.,.72]]); B0=torch.tensor([[.8],[.35]]); C0=torch.tensor([[1.,.12]])
    u=torch.randn(n,T,1); z=torch.zeros(n,2); ys=[]
    for k in range(T): z=z@A0.T+u[:,k,0:1]*B0.T; ys.append(z@C0.T)
    y=torch.stack(ys,1).to(device); u=u.to(device); model=LinearSSM().to(device)
    opt=torch.optim.Adam(model.parameters(),lr=.025)
    for _ in range(500):
        pred=model(u); task=((pred-y)**2).mean(); wo,wr=model.grams(T)
        def penalty(w):
            ev=torch.linalg.eigvalsh(w); ratio=ev[0]/(torch.trace(w)/2+1e-8)
            return torch.relu(torch.tensor(.08,device=device)-ratio), ratio
        po,ro=penalty(wo); pr,rr=penalty(wr); loss=task+reg*(po+pr)
        opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(),2.); opt.step()
    return {'task_mse':float(task.detach().cpu()), 'normalized_observability':float(ro.detach().cpu()), 'normalized_reachability':float(rr.detach().cpu())}

def main():
    checks=mechanism_checks(); device='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        baseline=train(0.,device); idea=train(.08,device)
    except Exception:
        device='cpu'; baseline=train(0.,device); idea=train(.08,device)
    out={'seed':SEED,'device':device,'checks':checks,'training':{'baseline':baseline,'gramian_regularized':idea}}
    Path('results.json').write_text(json.dumps(out,indent=2)); print(json.dumps(out,indent=2))
if __name__=='__main__': main()