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

SEED = 1729
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)

def kron_all(xs):
    z = xs[0]
    for x in xs[1:]: z = torch.kron(z, x)
    return z

def numpy_kron_all(xs):
    z = xs[0]
    for x in xs[1:]: z = np.kron(z, x)
    return z

def verify_majorization(trials=100, n=3, d=3):
    worst = -np.inf; failures = 0; max_k_gap = np.zeros(numpy_kron_all([np.eye(d)]*n).shape[0])
    tight_ratios=[]
    for t in range(trials):
        As=[]; Bs=[]
        for i in range(n):
            x=np.random.randn(d,d); As.append(x@x.T + .05*np.eye(d))
            x=np.random.randn(d,d); Bs.append(x@x.T + .05*np.eye(d))
        M=numpy_kron_all(As)+numpy_kron_all(Bs)
        lm=np.linalg.eigvalsh(M)[::-1]
        av=[np.linalg.eigvalsh(x)[::-1] for x in As]
        bv=[np.linalg.eigvalsh(x)[::-1] for x in Bs]
        c=numpy_kron_all(av)+numpy_kron_all(bv); c=np.sort(c)[::-1]
        gaps=np.cumsum(lm)-np.cumsum(c)
        worst=max(worst, float(gaps.max()))
        max_k_gap=np.maximum(max_k_gap, gaps)
        failures += int(np.any(gaps > 2e-9))
        tight_ratios.append(lm[0]/c[0])
    return {'trials':trials, 'failures':failures, 'worst_cumulative_gap':worst,
            'max_gap_by_k':max_k_gap.tolist(), 'max_eigenvalue_ratio_mean':float(np.mean(tight_ratios)),
            'max_eigenvalue_ratio_min':float(np.min(tight_ratios)), 'max_eigenvalue_ratio_max':float(np.max(tight_ratios))}

class SepPSD(torch.nn.Module):
    def __init__(self, n=3, d=3, eps=1e-3):
        super().__init__(); self.n=n; self.d=d; self.eps=eps
        self.LA=torch.nn.ParameterList([torch.nn.Parameter(.15*torch.randn(d,d)) for _ in range(n)])
        self.LB=torch.nn.ParameterList([torch.nn.Parameter(.15*torch.randn(d,d)) for _ in range(n)])
    def factors(self):
        I=torch.eye(self.d, device=self.LA[0].device)
        A=[L@L.T+self.eps*I for L in self.LA]
        B=[L@L.T+self.eps*I for L in self.LB]
        return A,B
    def bound(self, k=1):
        A,B=self.factors()
        ae=[torch.linalg.eigvalsh(x).flip(0) for x in A]
        be=[torch.linalg.eigvalsh(x).flip(0) for x in B]
        c=kron_all(ae)+kron_all(be)
        return torch.sort(c,descending=True).values[:k].sum()/k
    def matrix(self):
        A,B=self.factors(); return kron_all(A)+kron_all(B)
    def forward(self,x, control=None):
        M=self.matrix()
        if control is not None:
            r=self.bound(1)
            scale=torch.clamp(torch.as_tensor(control,device=x.device)/r, max=1.0)
            M=M*scale
        return x@M.T

def train(mode, seed=1729, steps=250, lr=.12):
    torch.manual_seed(seed)
    dev='cuda' if torch.cuda.is_available() else 'cpu'
    try:
        model=SepPSD().to(dev)
        g=torch.Generator(device=dev); g.manual_seed(seed+1)
        X=torch.randn(512,27,device=dev,generator=g)
        # A fixed, well-conditioned PSD target gives a nontrivial operator fitting task.
        q,_=torch.linalg.qr(torch.randn(27,27,device=dev,generator=g))
        target=q@torch.diag(torch.linspace(.15,1.0,27,device=dev))@q.T
        Y=X@target.T
        opt=torch.optim.Adam(model.parameters(),lr=lr)
        losses=[]; norms=[]; bounds=[]; spikes=0
        t0=time.perf_counter()
        for step in range(steps):
            ix=torch.arange((step*64)%448,(step*64)%448+64,device=dev)
            xb,yb=X[ix],Y[ix]
            pred=model(xb, control=2.0 if mode=='controller' else None)
            loss=((pred-yb)**2).mean()
            if mode=='penalty': loss=loss + .01*model.bound(1)
            if not torch.isfinite(loss): spikes += 1; break
            opt.zero_grad(); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1000.0); opt.step()
            losses.append(float(loss.detach().cpu()))
            with torch.no_grad():
                M=model.matrix(); norms.append(float(torch.linalg.eigvalsh(M)[-1].cpu())); bounds.append(float(model.bound(1).cpu()))
        elapsed=time.perf_counter()-t0
        final_loss=losses[-1] if losses else float('inf')
        return {'mode':mode,'device':dev,'final_loss':final_loss,'best_loss':min(losses) if losses else float('inf'),
                'steps_completed':len(losses),'nonfinite_steps':spikes,'max_exact_norm':max(norms) if norms else None,
                'final_exact_norm':norms[-1] if norms else None,'final_R1':bounds[-1] if bounds else None,
                'norm_over_R1':norms[-1]/bounds[-1] if bounds else None,'seconds':elapsed}
    except Exception as e:
        if dev=='cuda':
            torch.cuda.empty_cache()
            # retry CPU by temporarily hiding CUDA is awkward; report error and caller reruns subprocess-free CPU path
        return {'mode':mode,'error':repr(e)}

def main():
    verification=verify_majorization()
    results=[]
    for mode in ['baseline','penalty','controller']:
        r=train(mode)
        if 'error' in r and r.get('device')=='cuda':
            # CPU fallback in-process
            old=torch.cuda.is_available
            torch.cuda.is_available=lambda: False
            r=train(mode)
            torch.cuda.is_available=old
        results.append(r)
    out={'seed':SEED,'verification':verification,'training':results}
    with open('results.json','w') as f: json.dump(out,f,indent=2)
    print(json.dumps(out,indent=2))
if __name__=='__main__': main()