"""Rigorous verification of the branch-and-bound certificate for the degree-3 switch lemma.
For each leaf box B (with feasible interior point x1): checks, in exact/ball arithmetic,
   G(x) <= G(x1) + (M - gc.x1) + sum gr*w <= V   for all x in B cap D,
where [gc-gr, gc+gr] encloses grad G on B, M >= max_{x in B cap D} gc.x (LP weak duality).
For each empty box: checks that one polytope constraint is violated on the whole box."""
import sys, json, time, numpy as np
from multiprocessing import Pool
from fractions import Fraction as Fr
from scipy.optimize import linprog
A=[[1,-1,-1],[-1,1,-1],[-1,-1,1],[1,1,1]]
def init():
    global arb, acb, ctx, arbG, grad, PI
    from flint import arb, acb, ctx
    import arbG as _a, grad as _g
    arbG=_a; grad=_g; ctx.prec=64; arbG.ctx.prec=64; PI=arb.pi()
def rvals(pat):
    return [{"33":2*arb(2).sqrt()/3,"34":2/(arb(6).root(4)+1/arb(6).root(4))}[p] for p in pat]
def bvec(): return [arb(0),arb(0),arb(0),2*PI]
def empty_ok(box):
    lo,hi=box["lo"],box["hi"]; b=bvec()
    for row,bb in zip(A,b):
        mn=sum((arb(lo[m]) if a>0 else arb(hi[m]))*a for m,a in enumerate(row))
        if mn>bb: return True
    return False
def lip_arb(pat,rs,lo,hi):
    s=[(1-r*r).sqrt() for r in rs]; L=[]
    for m,(i,j) in enumerate(((0,1),(0,2),(1,2))):
        a=rs[i]*rs[j]; b=s[i]*s[j]; Rmax=a+b*arb(lo[m]).cos(); Rmin=a+b*arb(hi[m]).cos()
        if pat[i]==pat[j]:
            L.append(s[i]*(1+arb(lo[m]).cos()).sqrt()/(2*PI*(1+Rmin).sqrt()))
        else:
            smax=arb(1) if lo[m]<=np.pi/2<=hi[m] else arb(lo[m]).sin().union(arb(hi[m]).sin())
            L.append(b*smax/(2*PI*(1-Rmax).sqrt()*(1+Rmin).sqrt()))
    return L
def dual_bound(gc,lo,hi):
    # max gc.x s.t. A x <= b, lo<=x<=hi ; weak-duality upper bound with exact recomputation
    Af=np.array(A,float); bf=np.array([0,0,0,2*np.pi])
    res=linprog(np.r_[bf,np.array(hi),-np.array(lo)],A_eq=np.hstack([Af.T,np.eye(3),-np.eye(3)]),b_eq=np.array(gc),
                bounds=[(0,None)]*10,method='highs')
    y=np.maximum(res.x[:4],0) if res.status==0 else np.zeros(4)
    M=arb(0)
    for k in range(4): M+=arb(float(y[k]))*bvec()[k]
    for m in range(3):
        aty=sum(arb(float(y[k]))*A[k][m] for k in range(4)); r=arb(gc[m])-aty
        u=r if r>0 else arb(0); v=-r if r<0 else arb(0)
        if not (r>0 or r<0 or r==0): u=abs(r); v=abs(r)  # undetermined sign: conservative
        M+=u*arb(hi[m])-v*arb(lo[m])
    return M
def interior_point(lo,hi):
    # maximise slack t: A x + t*|a| <= b, lo+t <= x <= hi-t
    Af=np.array(A,float); bf=np.array([0,0,0,2*np.pi]); nr=np.linalg.norm(Af,axis=1)
    Aub=np.vstack([np.hstack([Af,nr[:,None]]),np.hstack([-np.eye(3),np.ones((3,1))]),np.hstack([np.eye(3),np.ones((3,1))])])
    bub=np.r_[bf,-np.array(lo),np.array(hi)]
    res=linprog(np.r_[0,0,0,-1],A_ub=Aub,b_ub=bub,bounds=[(None,None)]*3+[(0,None)],method='highs')
    return (res.x[:3],res.x[3]) if res.status==0 else (None,0)
def strictly_ok(x,lo,hi):
    b=bvec()
    for row,bb in zip(A,b):
        if not (sum(arb(x[m])*a for m,a in enumerate(row))<bb): return False
    return all(lo[m]<=x[m]<=hi[m] and 0<x[m]<np.pi for m in range(3))
def make_x1(lo,hi,x0):
    p,t=interior_point(lo,hi)
    if p is None or t<=0: return None
    for lam in (1e-8,1e-7,1e-6,1e-5,1e-4,1e-3,1e-2,0.1,0.5,1.0):
        x=[float(x0[m]+lam*(p[m]-x0[m])) for m in range(3)]
        if strictly_ok(x,lo,hi): return x
    return None
CSTAR=[2.0,2.0,2.0]
def G_rig(rs,x,tol):
    eta=arb(10)**-8
    R=arbG.star(rs,x); Rp=[row[:] for row in R]
    for i in range(1,4): Rp[0][i]=-R[0][i]; Rp[i][0]=-R[i][0]
    def o4(Rm):
        v=acb.integral(arbG.integrand_t(Rm),0,1-eta,rel_tol=tol,abs_tol=tol,eval_limit=10**6,depth_limit=10**4)
        return arb(1)/16+v.real+arbG.tail(Rm,eta)
    return arb(1)/2-o4(R)+o4(Rp)
def strictly_inD(x):
    b=bvec()
    for row,bb in zip(A,b):
        if not (sum(arb(x[m])*a for m,a in enumerate(row))<bb): return False
    return all(0<x[m]<np.pi for m in range(3))
def bound_at(pat,rs,lo,hi,x1,gc,gr,M,Gval):
    w=[abs(arb(x1[m])-arb(lo[m])).union(abs(arb(hi[m])-arb(x1[m]))).upper() for m in range(3)]
    return Gval+M-sum(arb(gc[m])*arb(x1[m]) for m in range(3))+sum(gr[m]*w[m] for m in range(3))
def grad_bounds(pat,rs,lo,hi):
    g=grad.grad_box(rs,lo,hi)
    if all(x.is_finite() for x in g):
        gc=[float(x.mid()) for x in g]
        gr=[(abs(x-arb(c))).upper() for x,c in zip(g,gc)]
        return gc,gr
    return [0.0,0.0,0.0],lip_arb(pat,rs,lo,hi)
def leaf_ok(args):
    import bb as fb
    pat,V,leaf=args
    rs=rvals(pat); lo,hi,x0=leaf["lo"],leaf["hi"],leaf["x0"]
    p,t=interior_point(lo,hi)
    target=p if (p is not None and t>0) else CSTAR
    cands=[]
    for lam in (1e-6,1e-4,1e-3,1e-2,0.05,0.2,0.5,1.0):
        x=[float(x0[m]+lam*(target[m]-x0[m])) for m in range(3)]
        if strictly_inD(x): cands.append(x)
    if not cands: return ("no-point",leaf)
    rs_f=np.array([{"33":fb.r33,"34":fb.r34}[q] for q in pat])
    scored=[]
    for x in cands:
        hlo=[min(lo[m],x[m]) for m in range(3)]; hhi=[max(hi[m],x[m]) for m in range(3)]
        gc,gr=grad_bounds(pat,rs,hlo,hhi); M=dual_bound(gc,lo,hi)
        est=bound_at(pat,rs,lo,hi,x,gc,gr,M,arb(float(fb.G(rs_f,np.array(x)))))
        scored.append((float(est.mid()),x,gc,gr,M))
    scored.sort(key=lambda z:z[0])
    last=None
    for est,x,gc,gr,M in scored[:3]:
        for tol in (arb(2)**-30,arb(2)**-44):
            Gv=G_rig(rs,x,tol)
            if not Gv.is_finite(): continue
            ub=bound_at(pat,rs,lo,hi,x,gc,gr,M,Gv)
            if ub<arb(V): return ("ok",float(ub.upper()))
            last=("fail",dict(leaf=leaf,ub=str(ub),est=est))
    return last if last else ("arb-nan",leaf)
def lp_feasible(lo,hi):
    Af=np.array(A,float); bf=np.array([0,0,0,2*np.pi])
    res=linprog(np.zeros(3),A_ub=Af,b_ub=bf,bounds=list(zip(lo,hi)),method='highs'); return res.status==0
def lp_argmax(c,lo,hi):
    Af=np.array(A,float); bf=np.array([0,0,0,2*np.pi])
    res=linprog(-np.array(c),A_ub=Af,b_ub=bf,bounds=list(zip(lo,hi)),method='highs'); return res.x if res.status==0 else None
def verify_box(pat,V,lo,hi,depth,stats):
    if empty_ok(dict(lo=lo,hi=hi)): stats["empty"]+=1; return True
    if lp_feasible(lo,hi):
        rs=rvals(pat); gc,gr=grad_bounds(pat,rs,lo,hi); x0=lp_argmax(gc,lo,hi)
        if x0 is not None:
            r=leaf_ok((pat,V,dict(lo=lo,hi=hi,x0=list(x0))))
            if r[0]=="ok": stats["leaf"]+=1; stats["max_ub"]=max(stats["max_ub"],r[1]); return True
    if depth>=8: stats["unresolved"].append(dict(lo=lo,hi=hi)); return False
    w=[hi[m]-lo[m] for m in range(3)]; 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,lo,h1,depth+1,stats); b=verify_box(pat,V,l2,hi,depth+1,stats); return a and b
def work(item):
    kind,pat,V,obj=item
    stats=dict(empty=0,leaf=0,max_ub=-1.0,unresolved=[])
    ok=verify_box(pat,V,obj["lo"],obj["hi"],0,stats)
    return ("ok" if ok else "unresolved",stats)
if __name__=="__main__":
    fn=sys.argv[1]; pat=fn.split("_")[1].split(","); V=float(fn.split("_")[2].rsplit(".json",1)[0])
    data=json.load(open(fn)); items=[("E",pat,V,e) for e in data["empties"]]+[("L",pat,V,l) for l in data["leaves"]]
    t0=time.time(); counts={}; bad=[]; worst=-1
    with Pool(int(__import__("os").environ.get("NPROC","15")),initializer=init) as pool:
        for k,(st,info) in enumerate(pool.imap_unordered(work,items,chunksize=8)):
            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)%2000==0: print(f"{k+1}/{len(items)} {counts} {time.time()-t0:.0f}s",flush=True)
    print("DONE",fn,"V=",V,counts,"max ub over leaves",worst,f"{time.time()-t0:.0f}s",flush=True)
    json.dump(dict(V=V,counts=counts,max_ub=worst,bad=bad[:50]),open("verified_"+fn,"w"),default=str)