# The floor's LP (scipy/HiGHS): one exact 4x4 window of 0/1 pixels (all 65,536 configurations) plus
# 2-point windows at the longer lags; the dual certificate, refined on a fine frequency grid. See check_floor.py
# for what a certificate proves and how it is checked without a solver.
# Window LP floor, all configurations of a w x w window at once (sparse), plus 2-point windows at long lags.
import sys, numpy as np, itertools, time
import os; sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), '..'))
import filters as F
from scipy.optimize import linprog
import scipy.sparse as sp
taps=F.kuwahara(5)
def build(wn, Llong, M):
    Q=[(i,j) for j in range(wn) for i in range(wn)]; nq=len(Q)
    pairs=[(a,b) for a in range(nq) for b in range(a,nq)]; P=len(pairs)
    lag=[(Q[b][0]-Q[a][0],Q[b][1]-Q[a][1]) for a,b in pairs]
    S=[(i,j) for j in range(0,Llong+1) for i in range(-Llong,Llong+1) if i*i+j*j<=Llong*Llong and not (j==0 and i<=0)]
    QQ=set((dx,dy) for dx in range(-wn+1,wn) for dy in range(-wn+1,wn))
    S=[r for r in S if r not in QQ]   # long lags only (QQ is symmetric)
    nS=len(S)
    fy,fx=np.mgrid[0:M,0:M]/M; FX=fx.ravel(); FY=fy.ravel()
    H2=F.grid_response(taps,M).ravel()
    iw=0; irho=P; is0=P+nq; ib=is0+1; nv=ib+4*nS
    Axi=np.zeros((M*M,nv))
    for k,((a,b),r) in enumerate(zip(pairs,lag)):
        Axi[:,iw+k]=(1 if a==b else 2)*np.cos(2*np.pi*(FX*r[0]+FY*r[1]))
    for k,r in enumerate(S):
        Axi[:,ib+4*k]=2*np.cos(2*np.pi*(FX*r[0]+FY*r[1])); Axi[:,ib+4*k+1]=2
    X=np.array(list(itertools.product((0,1),repeat=nq)),dtype=np.int8)
    PA=np.array([a for a,b in pairs]); PB=np.array([b for a,b in pairs]); mult=np.where(PA==PB,1,2)
    prod=(X[:,PA]*X[:,PB]).astype(float)*mult
    Acfg=sp.hstack([sp.csr_matrix(-prod), sp.csr_matrix(X.astype(float)), sp.csr_matrix(np.ones((len(X),1))), sp.csr_matrix((len(X),4*nS))]).tocsr()
    r2=[];c2=[];v2=[];row=0
    for k in range(nS):
        for x0,x1 in itertools.product((0,1),repeat=2):
            for col,val in ((ib+4*k+3,1),(ib+4*k,-2*x0*x1),(ib+4*k+1,-(x0+x1)),(ib+4*k+2,(x0+x1))):
                if val: r2.append(row);c2.append(col);v2.append(val)
            row+=1
    A2=sp.csr_matrix((v2,(r2,c2)),shape=(row,nv))
    A=sp.vstack([sp.csr_matrix(Axi),Acfg,A2]).tocsr()
    b=np.concatenate([H2,np.zeros(Acfg.shape[0]+row)])
    return dict(A=A,b=b,nv=nv,P=P,nq=nq,pairs=pairs,mult=mult,iw=iw,irho=irho,is0=is0,ib=ib,nS=nS,S=S,lag=lag,X=X)
def objective(f,B):
    c=np.zeros(B['nv'])
    c[B['iw']:B['iw']+B['P']]=-f*f*B['mult']; c[B['irho']:B['irho']+B['nq']]=f; c[B['is0']]=1
    for k in range(B['nS']):
        i=B['ib']+4*k; c[i]=-2*f*f; c[i+1]=-2*f*f; c[i+2]=2*f; c[i+3]=1
    return c
def xi_rows(B, FX, FY):
    """lam_hat rows at arbitrary frequencies (same column layout as the grid rows)."""
    R=np.zeros((len(FX),B['nv']))
    for k,((a,b),r) in enumerate(zip(B['pairs'],B['lag'])):
        R[:,B['iw']+k]=(1 if a==b else 2)*np.cos(2*np.pi*(FX*r[0]+FY*r[1]))
    for k,r in enumerate(B['S']):
        R[:,B['ib']+4*k]=2*np.cos(2*np.pi*(FX*r[0]+FY*r[1])); R[:,B['ib']+4*k+1]=2
    return R
FINE=384
def solve(f,B,refine=True,verbose=False):
    c=objective(f,B); A=B['A']; b=B['b']
    fy,fx=np.mgrid[0:FINE,0:FINE]/FINE; FX=fx.ravel(); FY=fy.ravel()
    Hf=F.grid_response(taps,FINE).ravel()
    for it in range(30):
        res=linprog(-c,A_ub=A,b_ub=b,bounds=[(None,None)]*B['nv'],method='highs')
        if not refine: break
        g=Hf-xi_rows(B,FX,FY)@res.x
        bad=np.argsort(g)[:300]; bad=bad[g[bad]<-2e-5]
        if verbose: print('  refine',it,'min g',g.min(),'adding',len(bad),flush=True)
        if len(bad)==0: break
        A=sp.vstack([A,sp.csr_matrix(xi_rows(B,FX[bad],FY[bad]))]).tocsr(); b=np.concatenate([b,Hf[bad]])
    return -res.fun, res
if __name__=='__main__':
    wn,Ll,M=int(sys.argv[1]),int(sys.argv[2]),int(sys.argv[3])
    t=time.time(); B=build(wn,Ll,M); print('built',B['A'].shape,B['A'].nnz,round(time.time()-t),'s',flush=True)
    m2=0.12672903767193516**2
    for f in [float(v) for v in sys.argv[4].split(',')]:
        t=time.time(); lb,res=solve(f,B)
        print(f'f {f}: LB {lb:.6f} parseval {m2*f*(1-f):.6f} ratio {lb/(m2*f*(1-f)):.3f} status {res.status} {time.time()-t:.0f}s',flush=True)