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

SEED = 2843
np.random.seed(SEED); random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

def jacobian(W, U, us, b, x0, C):
    n = W.shape[0]
    P = torch.eye(n, dtype=W.dtype, device=W.device)
    x = x0
    blocks = [C @ P]
    masks = []
    for u in us:
        z = W @ x + U @ u + b
        d = (z > 0).to(W.dtype)
        masks.append(d)
        x = torch.relu(z)
        P = torch.diag(d) @ W @ P
        blocks.append(C @ P)
    return torch.cat(blocks, dim=0), masks

def metrics(J, eps=1e-8):
    s = torch.linalg.svdvals(J)
    gram = J.T @ J + eps * torch.eye(J.shape[1], device=J.device)
    return {'rank': int((s > 1e-6).sum()), 'smin': float(s[-1]),
            'smax': float(s[0]), 'logdet': float(torch.linalg.slogdet(gram)[1]),
            'singular_values': s.detach().cpu().numpy().tolist()}

def autograd_check():
    n, m, T = 6, 3, 3
    g = torch.Generator(device=device).manual_seed(SEED + 99)
    W = torch.randn(n,n,generator=g,device=device)/math.sqrt(n)
    U = torch.randn(n,2,generator=g,device=device)
    b = torch.randn(n,generator=g,device=device)
    x0 = torch.randn(n,generator=g,device=device,requires_grad=True)
    us = [torch.randn(2,generator=g,device=device) for _ in range(T)]
    C = torch.randn(m,n,generator=g,device=device)
    J,_ = jacobian(W,U,us,b,x0,C)
    def obs(q):
        x=q; ys=[C@x]
        for u in us:
            x=torch.relu(W@x+U@u+b); ys.append(C@x)
        return torch.cat(ys)
    A=torch.autograd.functional.jacobian(obs,x0)
    return float((J-A).abs().max())

def setup(seed, n=16):
    g=torch.Generator(device=device).manual_seed(seed)
    W=torch.randn(n,n,generator=g,device=device)/math.sqrt(n)
    U=torch.zeros(n,1,device=device)
    b=torch.ones(n,device=device)*0.7
    x0=torch.randn(n,generator=g,device=device)
    return W,U,b,x0

def Cselect(m,n):
    C=torch.zeros(m,n,device=device)
    C[torch.arange(m),torch.arange(m)]=1
    return C

def boundary_sweep():
    n=16; W,U,b,x0=setup(SEED); u=[torch.zeros(1,device=device)]
    out=[]
    for m in [2,4,6,7,8,9,12,16]:
        q=metrics(jacobian(W,U,u,b,x0,Cselect(m,n))[0])
        out.append({'m':m,'counting_max_rank':min(n,2*m),'rank':q['rank'],'smin':q['smin'],'logdet':q['logdet']})
    return out

def horizon_sweep():
    n,m=16,4; W,U,b,x0=setup(SEED+1); us=[torch.zeros(1,device=device) for _ in range(8)]
    out=[]
    for T in range(9):
        q=metrics(jacobian(W,U,us[:T],b,x0,Cselect(m,n))[0])
        out.append({'T':T,'counting_max_rank':min(n,m*(T+1)),'rank':q['rank'],'smin':q['smin']})
    return out

def regularizer_trial(m, regularized, steps=300):
    n=16; W,U,b,x0=setup(SEED+10); W.requires_grad_(); b.requires_grad_()
    us=[torch.zeros(1,device=device)]
    C=Cselect(m,n)
    initial=metrics(jacobian(W,U,us,b,x0,C)[0])
    if regularized:
        opt=torch.optim.Adam([W,b],lr=0.01)
        for _ in range(steps):
            opt.zero_grad()
            J,_=jacobian(W,U,us,b,x0,C)
            G=J.T@J+1e-3*torch.eye(n,device=device)
            loss=-torch.linalg.slogdet(G)[1]+0.01*(W.square().mean()+b.square().mean())
            loss.backward(); opt.step()
    final=metrics(jacobian(W,U,us,b,x0,C)[0])
    return {'initial':initial,'final':final}

def main():
    result={'seed':SEED,'device':device,
      'autograd_max_abs_error':autograd_check(),
      'predictions':{
       'P1':'At T=1, rank(J)<=2m; m<8 cannot have rank 16, while m=8 is only a necessary counting threshold.',
       'P2':'At m=4, rank 16 is first possible at T=3 because 4(T+1)>=16.',
       'P3':'A log-det regularizer should increase finite-horizon Gramian logdet, but cannot overcome structural rank loss.'},
      'boundary_T1':boundary_sweep(),'horizon_m4':horizon_sweep(),
      'regularizer_m9':regularizer_trial(9,True),
      'fixed_control_m9':regularizer_trial(9,False)}
    Path('results.json').write_text(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))

if __name__=='__main__': main()