import json
from pathlib import Path
import numpy as np

SEED=3157

def projector(A,horizon=6):
    if not np.all(np.isfinite(A)): return None
    O=np.vstack([np.linalg.matrix_power(A,t) for t in range(horizon)])
    if not np.all(np.isfinite(O)): return None
    Q,_=np.linalg.qr(O,mode='reduced')
    return Q@Q.T

def opnorm(M):
    if M is None or not np.all(np.isfinite(M)): return float('inf')
    try: return float(np.linalg.svd(M,compute_uv=False)[0])
    except np.linalg.LinAlgError: return float('inf')

def gap(A,B,horizon=6):
    P,Q=projector(A,horizon),projector(B,horizon)
    return opnorm(None if P is None or Q is None else P-Q)

def make_problem(seed=SEED,d=3,n=4):
    rng=np.random.default_rng(seed); targets=[]
    for _ in range(n):
        X=rng.normal(size=(d,d)); X=X/(1.8*max(1.,opnorm(X)))
        targets.append(X+.12*np.diag(rng.normal(size=d)))
    return np.array([T+.20*rng.normal(size=T.shape) for T in targets]),np.array(targets)

def fd_check():
    rng=np.random.default_rng(11); A=rng.normal(size=(3,3)); A=A/(2*opnorm(A)); D=rng.normal(size=(3,3)); D=D/opnorm(D)
    P=projector(A); href=1e-7; ref=(projector(A+href*D)-P)/href; rows=[]
    for h in [1e-1,3e-2,1e-2,3e-3,1e-3,3e-4]: rows.append((h,opnorm((projector(A+h*D)-P)/h-ref)))
    B=A+.03*rng.normal(size=A.shape); Delta=.04*D
    pred=opnorm((P+(projector(A+Delta)-P))-projector(B)); actual=gap(A+Delta,B)
    return rows,pred,actual

def run(trust,seed,lr,epsilon,steps=100,horizon=6):
    A,T=make_problem(seed); initial=A.copy(); losses=[]; gaps=[]; alphas=[]; spikes=0; diverged=False
    for step in range(steps):
        if not np.all(np.isfinite(A)) or np.max(np.abs(A))>1e100: diverged=True; break
        grads=2*(A-T); proposed=np.zeros_like(A)
        for inds in ([0,1],[2,3]):
            D=-lr*np.mean(grads[list(inds)],axis=0)
            for i in inds: proposed[i]=D
        aa=[]
        for inds in ([0,1],[2,3]):
            inds=list(inds); leader=inds[0]
            p=[projector(A[i]+proposed[i],horizon) for i in inds]
            mg=opnorm(None if p[0] is None or p[1] is None else p[1]-p[0])
            alpha=min(1.,epsilon/(mg+1e-12)) if trust else 1.
            if trust:
                while alpha>1e-6 and gap(A[inds[1]]+alpha*proposed[inds[1]],A[leader]+alpha*proposed[leader],horizon)>epsilon: alpha*=.5
            aa.append(alpha)
            for i in inds: A[i]+=alpha*proposed[i]
        g=max(gap(A[1],A[0],horizon),gap(A[3],A[2],horizon))
        loss=float(np.mean((A-T)**2)) if np.all(np.isfinite(A)) else float('inf')
        if not np.isfinite(loss): diverged=True; break
        if step and loss>1.5*losses[-1]: spikes+=1
        losses.append(loss); gaps.append(g); alphas.append(min(aa))
    return dict(final_loss=losses[-1] if losses else float('inf'),min_loss=min(losses) if losses else float('inf'),max_gap=max(gaps) if gaps else float('inf'),mean_gap=np.mean(gaps) if gaps else float('inf'),clipped=sum(a<.999999 for a in alphas),mean_alpha=np.mean(alphas) if alphas else 0.,spikes=spikes,diverged=diverged,initial_gap=max(gap(initial[1],initial[0],horizon),gap(initial[3],initial[2],horizon)))

def summarize(vals):
    keys=['final_loss','min_loss','max_gap','mean_gap','clipped','mean_alpha','spikes','diverged','initial_gap']
    return {k:[float(v[k]) for v in vals] for k in keys}

def main():
    fd,pred,actual=fd_check(); print('FINITE_DIFFERENCE')
    for h,e in fd: print(f'h={h:.1e} error={e:.6e}')
    print(f'PREDICTED_ACTUAL_GAP pred={pred:.8f} actual={actual:.8f} abs_error={abs(pred-actual):.3e}')
    out={'fd':fd,'predicted_gap':pred,'actual_gap':actual,'runs':{}}
    for eps in [.90,.95,1.00]:
      for lr in [.8,1.0,1.2]:
       for trust in [False,True]:
        vals=[run(trust,s,lr,eps) for s in [3157,3158,3159,3160,3161]]
        key=f"{'trust' if trust else 'baseline'}_eps{eps}_lr{lr}"; out['runs'][key]=summarize(vals)
        print('\n'+key)
        for k in ['final_loss','min_loss','max_gap','mean_alpha','clipped','spikes','diverged']:
            x=np.array([v[k] for v in vals]); print(f'{k}: mean={np.mean(x):.6g} std={np.std(x):.6g}')
    Path('results.json').write_text(json.dumps(out,indent=2))
if __name__=='__main__': main()