import json, math, os
import numpy as np
from scipy.optimize import minimize, brentq
from itertools import product

rng = np.random.default_rng(7)

def mons(n, deg):
    out=[]
    def rec(rem,k,p):
        if k==n-1: out.append(tuple(p+[rem])); return
        for a in range(rem+1): rec(rem-a,k+1,p+[a])
    rec(deg,0,[]); return out

def poly_add(a,b,scale=1.):
    c=dict(a)
    for e,v in b.items(): c[e]=c.get(e,0.)+scale*v
    return {e:v for e,v in c.items() if abs(v)>1e-14}

def poly_mul(a,b):
    c={}
    for x,v in a.items():
      for y,w in b.items():
        e=tuple(i+j for i,j in zip(x,y)); c[e]=c.get(e,0.)+v*w
    return c

def power(p,k):
    q={(0,)*len(next(iter(p))):1.}
    for _ in range(k): q=poly_mul(q,p)
    return q

def gram_system(poly,n,deg):
    mm=mons(n,deg); pairs=[(i,j) for i in range(len(mm)) for j in range(i,len(mm))]
    exps=sorted(set(tuple(mm[i][k]+mm[j][k] for k in range(n)) for i,j in pairs)|set(poly))
    rows=[]; rhs=[]
    for e in exps:
      row=np.zeros(len(pairs))
      for k,(i,j) in enumerate(pairs):
        ee=tuple(mm[i][z]+mm[j][z] for z in range(n))
        if ee==e: row[k]=1 if i==j else 2
      rows.append(row); rhs.append(poly.get(e,0.))
    return mm,pairs,np.asarray(rows),np.asarray(rhs)

def solve_sos(poly,n,r,tries=1):
    # Fast approximate SDP: alternate projection onto coefficient affine space and PSD cone.
    norm2={tuple(2 if i==j else 0 for i in range(n)):1. for j in range(n)}
    lifted=poly_mul(power(norm2,r),poly)
    deg=max(sum(e) for e in lifted)//2
    mm,pairs,C,b=gram_system(lifted,n,deg); m=len(mm)
    # map symmetric matrix to the upper-triangular coefficient vector
    x=np.linalg.lstsq(C,b,rcond=None)[0]
    def unpack(v):
      Q=np.zeros((m,m))
      for a,(i,j) in zip(v,pairs): Q[i,j]=Q[j,i]=a
      return Q
    def pack(Q): return np.array([Q[i,j] for i,j in pairs])
    # affine projection, with a PSD projection between corrections
    G=np.linalg.pinv(C@C.T,rcond=1e-10)
    for _ in range(100):
      Q=unpack(x); w,U=np.linalg.eigh((Q+Q.T)/2); Q=(U*np.maximum(w,0.))@U.T
      x=x+C.T@(G@(b-C@pack(Q)))
    Q=unpack(x); ev=np.linalg.eigvalsh((Q+Q.T)/2)
    return {'success':bool(ev.min()>=-2e-5 and np.max(np.abs(C@x-b))<2e-5),
            'margin':float(ev.min()),'residual':float(np.max(np.abs(C@x-b))),
            'status':'alternating-projection'}

def normpow(n,k):
    return power({tuple(2 if i==j else 0 for i in range(n)):1. for j in range(n)},k)

def shifted_motzkin(eps):
    # Motzkin form plus eps*||x||_2^6; coefficients are accumulated explicitly.
    p={(4,2,0):1.,(2,4,0):1.,(0,0,6):1.,(2,2,2):-3.}
    return poly_add(p,normpow(3,3),scale=eps)

def threshold(n,r):
    # Coarse threshold scan avoids expensive repeated nonlinear solves.
    grid=[0.,.05,.15,.35,1.0]
    vals=[]
    for e in grid:
      rec=solve_sos(shifted_motzkin(e),n,r,tries=1)
      vals.append((e,rec['margin']))
      if rec['margin']>=-2e-5: return e,rec['margin']
    return grid[-1],vals[-1][1]

def rollout_boundary():
    # scalar z+=gamma*a*z; assess stability from 250-step bounded rollout.
    rows=[]
    for a in [0.5,0.8,1.0,1.2,1.5]:
      g=np.linspace(.1,2.2/a,401); empirical=[]
      for gamma in g:
        z=1.;
        for _ in range(250): z*=gamma*a
        empirical.append(np.isfinite(z) and abs(z)<=1.0001)
      stable=[x for x,ok in zip(g,empirical) if ok]
      last=max(stable,default=0.)
      rows.append({'a':a,'predicted_gamma':1/a,'observed_gamma':float(last),
                   'relative_error':abs(last-1/a)/(1/a)})
    return rows

def main():
    # Mechanism test: Motzkin is nonnegative, but not level-0 SOS; multiplier levels improve margin.
    levels=[]
    for r in range(2):
      rec=solve_sos(shifted_motzkin(0.),3,r,tries=1); levels.append({'r':r,**rec})
    thresholds=[]
    for r in range(2):
      th,mg=threshold(3,r); thresholds.append({'r':r,'epsilon':th,'margin':mg})
    boundary=rollout_boundary()
    result={'levels':levels,'thresholds':thresholds,'boundary':boundary,
      'predictions':{'threshold_monotone':all(thresholds[i]['epsilon']>=thresholds[i+1]['epsilon']-2e-3 for i in range(len(thresholds)-1)),
                     'level_improves':levels[-1]['margin']>levels[0]['margin']+1e-4,
                     'boundary_relative_error_max':max(abs(x['observed_gamma']-x['predicted_gamma'])/x['predicted_gamma'] for x in boundary)}}
    open('results.json','w').write(json.dumps(result,indent=2))
    print(json.dumps(result,indent=2))
if __name__=='__main__': main()