"""Rigorous verification of the degree-4 switch-lemma certificate (sign-moment form).
Region R = box ∩ Sym ∩ E (E = 4x4 correlation matrices, c=(c12,c13,c14,c23,c24,c34)).
Leaf check:  Psi(x) <= Psi(x1) + (M - gc.x1) + sum_lin gr*w + sum_sqrt 2K|sqrt(1-x1)-sqrt(1-x)|  <  V  on R.
Empty check: Farkas certificate for box ∩ Sym ∩ cuts, else search harder / subdivide."""
import sys, json, time, numpy as np, itertools
from multiprocessing import Pool
from scipy.optimize import linprog
from fractions import Fraction as Fr
PAIRS=list(itertools.combinations(range(4),2))
def init():
    global arb, acb, ctx, arbK, dualK, k4, bbC, PI, gradx
    from flint import arb, acb, ctx
    import arbK as _a, dualK as _d, k4 as _k, bbC as _b, gradx as _g
    arbK=_a; dualK=_d; k4=_k; bbC=_b; gradx=_g; ctx.prec=64; PI=arb.pi()
def cut_rows(cuts):
    """tangent cuts v^T C v >= 0 as rows A x <= b, with exact rational coefficients (v taken exactly)"""
    A=[];b=[]
    for v in cuts:
        v=[Fr(float(t)) for t in v]; A.append([-2*v[p]*v[q] for p,q in PAIRS]); b.append(sum(t*t for t in v))
    return A,b
def system(jj,cuts):
    """exact rows of box-independent constraints: symmetry rows and tangent cuts"""
    A_sym=bbC.symA(jj); Ac,bc=cut_rows(cuts)
    return [[Fr(float(t)) for t in r] for r in A_sym]+Ac, [Fr(0)]*len(A_sym)+bc
def farb(q): return arb(q.numerator)/arb(q.denominator)
def exact_le(row,x,rhs):  # exact rational check row.x <= rhs
    return sum(Fr(a)*Fr(xi) for a,xi in zip(row,x))<=Fr(rhs)
def in_region(x,lo,hi,A_sym):
    if not all(lo[m]<=x[m]<=hi[m] for m in range(6)): return False
    if not all(exact_le(r,x,0.0) for r in A_sym): return False
    # strictly PD via leading principal minors in Arb
    C=[[arb(1) if p==q else None for q in range(4)] for p in range(4)]
    for (p,q),v in zip(PAIRS,x): C[p][q]=C[q][p]=arb(v)
    def det(M):
        n=len(M)
        if n==1: return M[0][0]
        return sum(((-1)**j)*M[0][j]*det([row[:j]+row[j+1:] for row in M[1:]]) for j in range(n))
    return all(det([row[:k] for row in C[:k]])>0 for k in range(2,5))
def dual_bound(gc,lo,hi,A,b):
    """weak duality: max gc.x over {A x<=b, lo<=x<=hi} <= b.y + hi.u - lo.v with A^T y + u - v = gc"""
    Ae,be=A,b; A=np.array([[float(t) for t in r] for r in Ae],float) if len(Ae) else np.zeros((0,6)); b=np.array([float(t) for t in be],float)
    m=len(A)
    res=linprog(np.r_[b,np.array(hi),-np.array(lo)],A_eq=np.hstack([A.T,np.eye(6),-np.eye(6)]) if m else np.hstack([np.eye(6),-np.eye(6)]),
                b_eq=np.array(gc),bounds=[(0,None)]*(m+12),method='highs')
    y=np.maximum(res.x[:m],0) if (res.status==0 and m) else np.zeros(m)
    M=arb(0)
    for k in range(m): M+=arb(float(y[k]))*farb(be[k])
    for i in range(6):
        aty=sum((arb(float(y[k]))*farb(Ae[k][i]) for k in range(m)),arb(0)); r=arb(gc[i])-aty
        if r>0: M+=r*arb(hi[i])
        elif r<0: M+=r*arb(lo[i])
        else: M+=abs(r)*arb(max(abs(lo[i]),abs(hi[i])))
    return M
def farkas_empty(lo,hi,A,b):
    """certify {A x<=b, lo<=x<=hi} empty: y,u,v>=0, A^T y+u-v=0, b.y+hi.u-lo.v<0 (checked exactly)"""
    Ae,be=A,b; m=len(Ae)
    if m==0: return False
    A=np.array([[float(t) for t in r] for r in Ae],float); b=np.array([float(t) for t in be],float)
    res=linprog(np.r_[b,np.array(hi),-np.array(lo)],A_eq=np.hstack([A.T,np.eye(6),-np.eye(6)]),b_eq=np.zeros(6),
                bounds=[(0,1)]*m+[(0,None)]*12,method='highs')
    if res.status!=0 or res.fun>-1e-12: return False
    y=[Fr(max(float(t),0.0)) for t in res.x[:m]]
    tot=sum(y[k]*be[k] for k in range(m))
    for i in range(6):
        r=-sum(y[k]*Ae[k][i] for k in range(m))   # need u-v = r with u,v>=0 minimal cost
        tot+= r*Fr(hi[i]) if r>0 else r*Fr(lo[i])
    return tot<0
COMP={p:tuple(k for k in range(4) if k not in p) for p in PAIRS}
def slab_ok(pat,V,lo,hi):
    """Pair bound: for every pair (p,q) with complement (k,l), sgnY_p sgnY_q = 1 - 2*1[Y_pY_q<0] gives
    Psi4 = sum_i acos(r_i)/(4 pi) + 1/8 - asin(S_kl)/(4 pi) + E[sgnY_k sgnY_l; Y_pY_q<0]/4,
    and the last term is at most P(Y_pY_q<0)/4 = acos(S_pq)/(4 pi). Both terms are largest at lo."""
    rsa=[{"34":k4.R34A,"44":k4.R44A}[x] for x in pat]; s=[(1-r*r).sqrt() for r in rsa]
    const=sum((r.acos() for r in rsa),arb(0))/(4*PI)+arb(1)/8
    for m,(p,q) in enumerate(PAIRS):
        k,l=COMP[(p,q)]; mk=PAIRS.index((k,l))
        Skl=rsa[k]*rsa[l]+s[k]*s[l]*arb(lo[mk]); Spq=rsa[p]*rsa[q]+s[p]*s[q]*arb(lo[m])
        if not (Skl>-1 and Skl<1 and Spq>-1 and Spq<1): continue
        ub=const-Skl.asin()/(4*PI)+Spq.acos()/(4*PI)
        if ub<arb(V): return float(ub.upper())
    return None
def upper_face(lo,hi,A,b):
    """Certify that every x in {A x <= b, lo <= x <= hi} has x_m >= hi_m for some single m with lo_m<hi_m.
    Weak duality, exactly: for y >= 0, max(-x_m) <= b.y + sum_i max(r_i*hi_i, r_i*lo_i), r = -e_m - A^T y.
    If this is <= -hi_m, the half-open box [lo,hi) contains no point of the region."""
    Ae,be=A,b; mrows=len(Ae)
    if mrows==0: return None
    Af=np.array([[float(t) for t in r] for r in Ae],float); bf=np.array([float(t) for t in be],float)
    for m in range(6):
        if not lo[m]<hi[m]: continue
        gc=np.zeros(6); gc[m]=-1.0
        res=linprog(np.r_[bf,np.array(hi),-np.array(lo)],A_eq=np.hstack([Af.T,np.eye(6),-np.eye(6)]),b_eq=gc,
                    bounds=[(0,None)]*(mrows+12),method='highs')
        if res.status!=0 or res.fun>-hi[m]+1e-9: continue
        for den in (None,2**10,2**20,2**30):
            if den is None: y=[Fr(max(float(t),0.0)) for t in res.x[:mrows]]
            else: y=[Fr(round(max(float(t),0.0)*den),den) for t in res.x[:mrows]]
            tot=sum(y[k]*be[k] for k in range(mrows))
            for i in range(6):
                r=(Fr(-1) if i==m else Fr(0))-sum(y[k]*Ae[k][i] for k in range(mrows))
                tot+=max(r*Fr(hi[i]),r*Fr(lo[i]))
            if tot<=-Fr(hi[m]): return m
    return None
def Lglob_arb(rsa,p,q,lo,hi):
    s=[(1-r*r).sqrt() for r in rsa]; a=rsa[p]*rsa[q]; bb=s[p]*s[q]
    Smax=a+bb*arb(hi); Smin=a+bb*arb(lo); mm=abs(Smax).union(abs(Smin)).upper()
    return bb/(4*PI)/(1-mm*mm).sqrt()
def pd_near(lo,hi):
    """a strictly PD correlation point near the box centre (may lie outside the box)"""
    c=(np.array(lo)+np.array(hi))/2; C=k4.Cmat(c); w,v=np.linalg.eigh(C); w=np.maximum(w,1e-3)
    C=(v*w)@v.T; d=np.sqrt(np.diag(C)); C=C/np.outer(d,d); C=0.999*C+0.001*np.eye(4)
    return [float(C[p,q]) for p,q in PAIRS]
def strictly_pd(x):
    C=[[arb(1) if p==q else None for q in range(4)] for p in range(4)]
    for (p,q),v in zip(PAIRS,x): C[p][q]=C[q][p]=arb(v)
    def det(M):
        n=len(M)
        if n==1: return M[0][0]
        return sum(((-1)**j)*M[0][j]*det([row[:j]+row[j+1:] for row in M[1:]]) for j in range(n))
    return all(det([row[:k] for row in C[:k]])>0 for k in range(2,5)) and all(-1<v<1 for v in x)
def dist(x1,lo,hi):
    """rigorous upper bound for max over y in [lo,hi] of |y-x1|"""
    return abs(arb(x1)-arb(lo)).union(abs(arb(hi)-arb(x1))).upper()
def find_psd_cuts(lo,hi,cuts,A_sym,iters=40):
    """bbC.find_psd_point, but returning the cut vectors so that emptiness can be certified exactly"""
    vecs=[list(map(float,v)) for v in cuts]
    A=list(A_sym); b=[0.0]*len(A_sym)
    for v in vecs: r,rb=bbC.c4.cut_row(np.array(v)); A.append(r); b.append(rb)
    x=(lo+hi)/2
    for _ in range(iters):
        ok_sym=np.all(A_sym@x<=1e-12) if len(A_sym) else True
        w,v=bbC.c4.min_eig_vec(x)
        if w>1e-9 and ok_sym: return x,vecs
        if w<=1e-9:
            r,rb=bbC.c4.cut_row(v); A.append(r); b.append(rb); vecs.append(list(map(float,v)))
        Aa=np.array(A); nr=np.linalg.norm(Aa,axis=1)+1e-12
        Aub=np.vstack([np.hstack([Aa,nr[:,None]]),np.hstack([-np.eye(6),np.ones((6,1))]),np.hstack([np.eye(6),np.ones((6,1))])])
        bub=np.r_[np.array(b),-lo,hi]
        res=linprog(np.r_[np.zeros(6),-1],A_ub=Aub,b_ub=bub,bounds=[(None,None)]*7,method='highs')
        if res.status!=0 or res.x[6]<0: return None,vecs
        x=res.x[:6]
    w,_=bbC.c4.min_eig_vec(x)
    return (x,vecs) if (w>1e-9 and (np.all(A_sym@x<=1e-12) if len(A_sym) else True)) else (None,vecs)
def leaf_ok(pat,V,lo,hi,cuts,x0,refine=True,x1s=None):
    rs=np.array([{"34":k4.r34,"44":k4.r44}[x] for x in pat]); rsa=[{"34":k4.R34A,"44":k4.R44A}[x] for x in pat]
    j=sum(1 for x in pat if x=="34"); jj={("44","44","44","44"):0,("34","44","44","44"):1,("34","34","44","44"):2,("44","34","34","34"):3,("34","34","34","34"):4}[tuple(pat)]
    A_sym=bbC.symA(jj); A,b=system(jj,cuts)
    eq=[[pat[p]==pat[q] for q in range(4)] for p in range(4)]
    sqm=[eq[p][q] and hi[m]>0.98 for m,(p,q) in enumerate(PAIRS)]
    g=gradx.grad_best(rsa,list(lo),list(hi),need=[m for m in range(6) if not sqm[m]],refine_tol=1e-3,refine=refine)
    gc=[0.0]*6; gr=[arb(0)]*6; K=[arb(0)]*6
    for m,(p,q) in enumerate(PAIRS):
        if sqm[m]: K[m]=dualK.sqrt_K(rsa,p,q,lo[m])
        elif g[m] is not None and g[m].is_finite():
            gc[m]=float(g[m].mid()); gr[m]=abs(g[m]-arb(gc[m])).upper()
        else: gr[m]=Lglob_arb(rsa,p,q,lo[m],hi[m])
    M=dual_bound(gc,lo,hi,A,b)
    # candidate evaluation points: LP point moved toward an interior point, and the box-centre PD point
    cands=[]
    xin,_,_=bbC.find_psd_point(np.array(lo),np.array(hi),cuts,A_sym,iters=30)
    for base in ([x0] if x0 is not None else []):
        if xin is None: break
        for lam in (1e-6,1e-4,1e-2,0.05,0.2,0.5,1.0):
            xa=[float(base[m]+lam*(xin[m]-base[m])) for m in range(6)]
            if in_region(xa,lo,hi,A_sym): cands.append(xa); break
    if xin is not None and in_region(list(map(float,xin)),lo,hi,A_sym): cands.append(list(map(float,xin)))
    for extra in ([x1s] if x1s is not None else [])+[pd_near(lo,hi)]:
        if extra is not None and in_region(list(map(float,extra)),lo,hi,A_sym): cands.append(list(map(float,extra)))
    if not cands:
        xo=pd_near(lo,hi)
        if strictly_pd(xo): cands.append(xo)
    if not cands: return ("no-point",None)
    base=(gc,gr,K,M)
    def prep(x1):
        gc,gr,K,M=base
        if not all(lo[m]<=x1[m]<=hi[m] for m in range(6)):
            hlo=[min(lo[m],x1[m]) for m in range(6)]; hhi=[max(hi[m],x1[m]) for m in range(6)]
            sqh=[eq[p][q] and hhi[m]>0.98 for m,(p,q) in enumerate(PAIRS)]
            if sqh!=sqm: return None
            gh=gradx.grad_best(rsa,hlo,hhi,need=[m for m in range(6) if not sqm[m]],refine_tol=1e-3,refine=refine); gc=[0.0]*6; gr=[arb(0)]*6; K=[arb(0)]*6
            for m,(p,q) in enumerate(PAIRS):
                if sqm[m]: K[m]=dualK.sqrt_K(rsa,p,q,hlo[m])
                elif gh[m] is not None and gh[m].is_finite():
                    gc[m]=float(gh[m].mid()); gr[m]=abs(gh[m]-arb(gc[m])).upper()
                else: gr[m]=Lglob_arb(rsa,p,q,hlo[m],hhi[m])
            M=dual_bound(gc,lo,hi,A,b)
        # float estimate of the bound (chooses the point and the integration tolerance only)
        Pf=float(k4.Psi(rs,np.array(x1)))
        est=Pf+float((M-sum((arb(gc[m])*arb(x1[m]) for m in range(6)),arb(0))).upper())
        est+=sum(float(gr[m].upper())*max(x1[m]-lo[m],hi[m]-x1[m],0.0) for m in range(6) if not sqm[m])
        for m in range(6):
            if sqm[m]:
                est+=2*float(K[m].upper())*max(abs(np.sqrt(1-lo[m])-np.sqrt(max(1-x1[m],0))),abs(np.sqrt(max(1-x1[m],0))-np.sqrt(max(1-hi[m],0))))
        lam=float(np.linalg.eigvalsh(k4.SigmaY(rs,np.array(x1))).min())
        return (gc,gr,K,M,est,lam)
    prepared=[]
    for x1 in cands:
        P=prep(x1)
        if P is not None and P[4]<V: prepared.append((x1,P))
    # prefer well-conditioned evaluation points (the Arb integral is much cheaper there)
    prepared.sort(key=lambda z:-z[1][5])
    best=None
    for x1,(gc,gr,K,M,est,lam) in prepared:
        margin=V-est
        tol=min(2.0**-20,max(2.0**-36,margin/8))
        Pv=arbK.Psi(rsa,x1,tol)
        if not Pv.is_finite(): continue
        lin=M-sum((arb(gc[m])*arb(x1[m]) for m in range(6)),arb(0))
        quad=sum((gr[m]*dist(x1[m],lo[m],hi[m]) for m in range(6) if not sqm[m]),arb(0))
        sq=arb(0)
        for m in range(6):
            if sqm[m]:
                u1=(1-arb(x1[m])).sqrt(); ulo=(1-arb(lo[m])).sqrt(); uhi=(1-arb(min(hi[m],1.0))).sqrt() if hi[m]<1 else arb(0)
                sq+=2*K[m]*abs(ulo-u1).union(abs(u1-uhi)).upper()
        ub=Pv+lin+quad+sq
        if best is None or float(ub.upper())<float(best.upper()): best=ub
        if ub<arb(V): return ("ok",float(ub.upper()))
    return ("fail",None if best is None else float(best.upper()))
def verify_box(pat,V,lo,hi,cuts,x0,depth,stats,x1s=None):
    jj={("44","44","44","44"):0,("34","44","44","44"):1,("34","34","44","44"):2,("44","34","34","34"):3,("34","34","34","34"):4}[tuple(pat)]
    A_sym=bbC.symA(jj); A,b=system(jj,cuts)
    if farkas_empty(lo,hi,A,b): stats["empty"]+=1; return True
    if upper_face(lo,hi,A,b) is not None: stats["face"]=stats.get("face",0)+1; return True
    sv=slab_ok(pat,V,lo,hi)
    if sv is not None: stats["slab"]=stats.get("slab",0)+1; stats["max_ub"]=max(stats["max_ub"],sv); return True
    st,val=leaf_ok(pat,V,lo,hi,cuts,x0,refine=False,x1s=x1s)
    if st!="ok": st,val=leaf_ok(pat,V,lo,hi,cuts,x0,refine=True,x1s=x1s)
    if st=="ok": stats["leaf"]+=1; stats["max_ub"]=max(stats["max_ub"],val); return True
    # try to certify emptiness with extra cuts from a cutting-plane loop
    xin,vecs=find_psd_cuts(np.array(lo),np.array(hi),cuts,A_sym,iters=40)
    if xin is None:
        Ax,bx=system(jj,vecs)
        if farkas_empty(lo,hi,Ax,bx): stats["empty"]+=1; return True
    if depth>=20: stats["unresolved"].append(dict(lo=list(lo),hi=list(hi),why=st)); return False
    w=[hi[m]-lo[m] for m in range(6)]; m=int(np.argmax(w)); mid=(lo[m]+hi[m])/2
    h1=list(hi); h1[m]=mid; l2=list(lo); l2[m]=mid
    a=verify_box(pat,V,list(lo),h1,cuts,None,depth+1,stats); bq=verify_box(pat,V,l2,list(hi),cuts,None,depth+1,stats); return a and bq
def work(item):
    pat,V,box=item
    stats=dict(empty=0,leaf=0,max_ub=-1.0,unresolved=[])
    ok=verify_box(pat,V,box["lo"],box["hi"],box.get("cuts",[]),box.get("x0"),0,stats,x1s=box.get("x1"))
    return ("ok" if ok else "unresolved",stats)
if __name__=="__main__":
    fn=sys.argv[1]; j=int(fn.split("_j")[1].split("_")[0]); pat=list(bbC.PATS[j])
    data=json.load(open(fn)); V=data["V"]
    items=[(pat,V,e) for e in data["empties"]]+[(pat,V,l) for l in data["leaves"]]
    t0=time.time(); counts={}; bad=[]; worst=-1
    with Pool(15,initializer=init) as pool:
        for k,(st,info) in enumerate(pool.imap_unordered(work,items,chunksize=16)):
            counts[st]=counts.get(st,0)+1
            for key in ("empty","leaf"): counts["sub_"+key]=counts.get("sub_"+key,0)+info[key]
            worst=max(worst,info["max_ub"])
            if st!="ok": bad.append(info["unresolved"])
            if (k+1)%20000==0: print(f"{k+1}/{len(items)} {counts} {time.time()-t0:.0f}s",flush=True)
    print("DONE",fn,"V=",V,counts,"max ub",worst,f"{time.time()-t0:.0f}s",flush=True)
    json.dump(dict(V=V,counts=counts,max_ub=worst,bad=bad[:100]),open("verified4_"+fn,"w"),default=str)