#!/usr/bin/env python3
"""Rationalize the reflection-symmetric continuum RLT certificate."""
from pathlib import Path
import numpy as np
z=np.load('/tmp/quadrltsym.npz');U=z['U'].copy();V=z['V'].copy();NA=z['NA'].copy();NB=z['NB'].copy();eta=z['eta'].copy();H=z['H'];m=np.rint(z['m']).astype(np.int64);sl=np.rint(z['sl']).astype(np.int64)
L=len(m);n=L//2;D=10**6;Rscale=10**8;M=L//2
# Enforce all nonnegativity and leave strict row-budget slack before rounding down.
U=np.maximum(U,0);V=np.maximum(V,0);eta=np.maximum(eta,0);NA=np.maximum((NA+NA.T)/2,0);NB=np.maximum((NB+NB.T)/2,0)
for i in range(n):
    b=U[i].sum()+V[i].sum()+M*eta[i]
    if m[i]==0: U[i]=0;V[i]=0;eta[i]=0
    elif b>m[i]-.02:
        f=(m[i]-.02)/b;U[i]*=f;V[i]*=f;eta[i]*=f
Ui=np.floor(U*D+1e-10).astype(np.int64);Vi=np.floor(V*D+1e-10).astype(np.int64);ei=np.floor(eta*D+1e-10).astype(np.int64)
NAi=np.rint(NA*D).astype(np.int64);NBi=np.rint(NB*D).astype(np.int64)
# Exact integer budgets.
budget=Ui.sum(1)+Vi.sum(1)+M*ei
assert np.all(budget<=m[:n]*D) and min(Ui.min(),Vi.min(),NAi.min(),NBi.min(),ei.min())>=0
# Reconstruct integer 2D-scaled sector matrices. H has half-integer diagonal.
perm=np.r_[np.arange(n),np.arange(L-1,n-1,-1)];Hp=H[np.ix_(perm,perm)];HA=Hp[:n,:n];HB=Hp[:n,n:]
# Scale by 2D to make every base coefficient integral.
Ps=np.rint(2*D*(HA+HB)).astype(object)-2*(NAi.astype(object)+NBi.astype(object))+(Ui+Ui.T+Vi+Vi.T).astype(object)+2*(ei[:,None]+ei[None,:]).astype(object)
Pa=np.rint(2*D*(HA-HB)).astype(object)-2*NAi.astype(object)+2*NBi.astype(object)+(Ui+Ui.T-Vi-Vi.T).astype(object)
Psf=np.array(Ps,dtype=float)/(2*D);Paf=np.array(Pa,dtype=float)/(2*D)
print('budget slack',np.min(m[:n]-budget/D),'eigen',np.linalg.eigvalsh(Psf)[:3],np.linalg.eigvalsh(Paf)[:3])
# Exact PSD factors, with plenty of existing spectral margin.
def factor_and_margin(Pint,Pfloat):
    gap=.05
    F=np.rint(np.linalg.cholesky(Pfloat-gap*np.eye(n))*Rscale).astype(np.int64)
    Fo=F.astype(object);Q=(Rscale**2)*Pint-(2*D)*(Fo@Fo.T)
    margins=[int(Q[i,i])-sum(abs(int(Q[i,j])) for j in range(n) if i!=j) for i in range(n)]
    print('margin',min(margins));assert min(margins)>=0
    return F
Fs=factor_and_margin(Ps,Psf);Fa=factor_and_margin(Pa,Paf)
out=Path('/root/snapshot/workspace/experiments/experiment_1_ap_density/sdp/continuum_rlt_548.txt')
with out.open('w') as f:
 print('AP_CONTINUUM_RLT_V1',file=f);print('L',L,file=f);print('denom',D,file=f);print('factor_scale',Rscale,file=f)
 print('U',file=f)
 for row in Ui:print(*map(int,row),file=f)
 print('V',file=f)
 for row in Vi:print(*map(int,row),file=f)
 print('NA',file=f)
 for i in range(n):print(*map(int,NAi[i,:i+1]),file=f)
 print('NB',file=f)
 for i in range(n):print(*map(int,NBi[i,:i+1]),file=f)
 print('eta',*map(int,ei),file=f)
 print('factor_sym',file=f)
 for i in range(n):print(*map(int,Fs[i,:i+1]),file=f)
 print('factor_anti',file=f)
 for i in range(n):print(*map(int,Fa[i,:i+1]),file=f)
print('wrote',out,out.stat().st_size)
