#!/usr/bin/env python3
"""Rationalize the L=128 box-SDP dual into an exact SOS certificate."""
from pathlib import Path
from fractions import Fraction
import json, numpy as np

ROOT=Path('/root/snapshot')
z=np.load('/tmp/dual128.npz')
tri=json.load(open('/tmp/tri128.json'))
L=128;D=10**8;R=10**8
SLACK=100_000       # 0.001 on every diagonal
FACTOR_GAP=50_000   # factor H-0.0005 I, leaving an exact DD residual
A=z['A'].astype(np.int64)
a=np.rint(z['diag']*D).astype(np.int64)+SLACK
tris=[]
Hnum=2*D*A.copy()
for i in range(L):Hnum[i,i]+=2*int(a[i])
for lam,typ,i,j,k in tri:
    m=max(0,round(lam*D))
    if not m:continue
    tris.append((typ,i,j,k,m))
    signs=([1,1,1] if typ==0 else [1,-1,-1] if typ==1 else [-1,1,-1] if typ==2 else [-1,-1,1])
    for (u,v),sg in zip(((i,j),(i,k),(j,k)),signs):
        Hnum[u,v]-=m*sg;Hnum[v,u]-=m*sg
H=Hnum.astype(float)/(2*D)
e=np.linalg.eigvalsh(H)
print('H min eigenvalue',e[0])
HF=H-np.eye(L)*(FACTOR_GAP/D)
P=np.rint(np.linalg.cholesky(HF)*R).astype(np.int64)
# exact residual Q = R^2 Hnum - 2D P P^T, via Python integers
Po=P.astype(object);Q=(R*R)*Hnum.astype(object)-(2*D)*(Po@Po.T)
margins=[]
for i in range(L):margins.append(int(Q[i,i])-sum(abs(int(Q[i,j])) for j in range(L) if i!=j))
print('triangles',len(tris),'min residual margin',min(margins))
assert min(margins)>=0
S=int(a.sum())+sum(t[4] for t in tris)
qbound=Fraction(-S,D)
delta=Fraction(1,16)-Fraction(1,8*L)+qbound/Fraction(8*L*L)
print('sum',S,'q',qbound,float(qbound),'delta',delta,float(delta))
out=ROOT/'workspace/experiments/experiment_1_ap_density/sdp/l128_box_sos.txt'
with out.open('w') as f:
    print('AP_BOX_SOS_LOWER_V1',file=f);print('L',L,file=f);print('denom',D,file=f);print('factor_scale',R,file=f)
    print('diag',*map(int,a),file=f);print('triangles',len(tris),file=f)
    for t in tris:print(*t,file=f)
    print('factor',file=f)
    for i in range(L):print(*map(int,P[i,:i+1]),file=f)
    print('claimed_multiplier_sum',S,file=f);print('claimed_q_bound',qbound.numerator,qbound.denominator,file=f)
    print('claimed_delta_bound',delta.numerator,delta.denominator,file=f)
print('wrote',out,out.stat().st_size)
