"""Checkpointed, resumable worker: DFS from a sub-root box; same bounds as bbK.run."""
import numpy as np, sys, json, time, os
import bbK, k4, dualK, bbC, gradx
from k4 import PAIRS
COMP={p:tuple(k for k in range(4) if k not in p) for p in PAIRS}
def slab_ub(rs,lo,hi):
    # Psi4 = const - asin(S_kl)/(4 pi) + E[sgnY_k sgnY_l; Y_pY_q<0]/4 <= const - asin(S_kl)/(4 pi) + acos(S_pq)/(4 pi)
    s=np.sqrt(1-rs**2); const=np.sum(np.arccos(rs))/(4*np.pi)+1/8; best=np.inf
    for m,(p,q) in enumerate(PAIRS):
        k,l=COMP[(p,q)]; mk=PAIRS.index((k,l))
        Skl=rs[k]*rs[l]+s[k]*s[l]*lo[mk]; Spq=rs[p]*rs[q]+s[p]*s[q]*lo[m]
        best=min(best,const-np.arcsin(Skl)/(4*np.pi)+np.arccos(min(Spq,1.0))/(4*np.pi))
    return best+1e-12
def worker(j,bstar,lo0,hi0,tag,ckpt_every=100000):
    s=0.1529003-0.13982212400800; V=bstar-j*s
    fn=f"par_j{j}_{tag}.json"
    if os.path.exists(fn):
        st=json.load(open(fn)); stack=[(np.array(a),np.array(b),c) for a,b,c in st["stack"]]
        leaves=st["leaves"]; empties=st["empties"]; n=st["n"]; done_vol=st["done_vol"]
    else:
        stack=[(np.array(lo0,float),np.array(hi0,float),[])]; leaves=[]; empties=[]; n=0; done_vol=0.0
    tot_vol=float(np.prod(np.array(hi0)-np.array(lo0)))
    pat=bbK.PATS[j]; 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]
    eq=[[pat[p]==pat[q] for q in range(4)] for p in range(4)]; A_sym=bbC.symA(j); t0=time.time(); best=-1
    def save(final=False):
        json.dump(dict(V=V,bstar=bstar,j=j,tag=tag,n=n,done_vol=done_vol,final=final,
                       stack=[(a.tolist(),b.tolist(),c) for a,b,c in stack],leaves=leaves,empties=empties),open(fn+".tmp","w"))
        os.replace(fn+".tmp",fn)
    while stack:
        lo,hi,cuts=stack.pop(); vol=float(np.prod(hi-lo))
        x,Acut,bcut=bbC.find_psd_point(lo,hi,cuts,A_sym)
        if x is None: empties.append(dict(lo=lo.tolist(),hi=hi.tolist(),cuts=[list(v) for v in cuts])); done_vol+=vol; continue
        n+=1
        if slab_ub(rs,lo,hi)<=V:
            leaves.append(dict(lo=lo.tolist(),hi=hi.tolist(),cuts=[list(v) for v in cuts],slab=True)); done_vol+=vol; continue
        sqmask=np.array([eq[p][q] and hi[m]>0.98 for m,(p,q) in enumerate(PAIRS)])
        need=[m for m in range(6) if not sqmask[m]]
        for refine in ((False,True) if (hi-lo).max()<=0.25 else (False,)):
            g=gradx.grad_best(rsa,list(lo),list(hi),need=need,refine_tol=1e-3,refine=refine)
            gc=np.zeros(6); gr=np.zeros(6); sq=np.zeros(6)
            for m,(p,q) in enumerate(PAIRS):
                if sqmask[m]: sq[m]=float(dualK.sqrt_K(rsa,p,q,lo[m]).upper())
                elif g[m] is not None and g[m].is_finite(): gc[m]=float(g[m].mid()); gr[m]=float(g[m].rad())
                else: gr[m]=bbK.Lglob(rs,p,q,lo[m],hi[m])
            x0,M=bbC.lp_max(gc,lo,hi,np.array(Acut),np.array(bcut))
            if x0 is None: x0=x; M=gc@x
            cands=[]
            for lam in (1e-6,1e-3,0.01,0.05,0.2,0.5,1.0):
                xa=x0+lam*(x-x0)
                if np.linalg.eigvalsh(k4.Cmat(xa)).min()>1e-10: cands.append(xa); break
            cands.append(x); ub=np.inf
            for xc in cands:
                Gc=k4.Psi(rs,xc); best=max(best,Gc)
                if Gc>V:
                    save(); print("REFUTED",j,tag,Gc,xc.tolist(),flush=True); return
                wd=np.maximum(xc-lo,hi-xc)
                sqt=np.sum(sq*2*np.maximum(np.sqrt(np.maximum(1-lo,0))-np.sqrt(np.maximum(1-xc,0)),np.sqrt(np.maximum(1-xc,0))-np.sqrt(np.maximum(1-hi,0))))
                u=Gc+(M-gc@xc)+np.sum(gr*wd)+sqt
                if u<ub: ub=u; x1=xc
            if ub<=V: break
        if ub<=V:
            leaves.append(dict(lo=lo.tolist(),hi=hi.tolist(),cuts=[list(v) for v in cuts],x0=x0.tolist(),x1=x1.tolist())); done_vol+=vol; continue
        score=(np.abs(gc)+gr)*(hi-lo)
        uw=np.array([np.sqrt(max(1-lo[k],0))-np.sqrt(max(1-hi[k],0)) for k in range(6)])
        score=np.where(sqmask,sq*2*uw+1e-12,score); m=int(np.argmax(score))
        if sqmask[m]: um=(np.sqrt(max(1-lo[m],0))+np.sqrt(max(1-hi[m],0)))/2; mid=1-um*um
        else: mid=(lo[m]+hi[m])/2
        _,v=np.linalg.eigh(k4.Cmat(x0)); newcuts=(cuts+[list(v[:,0])])[-4:]
        l2=hi.copy(); l2[m]=mid; h1=lo.copy(); h1[m]=mid
        stack.append((lo,l2,newcuts)); stack.append((h1,hi,newcuts))
        if n%ckpt_every==0:
            save(); print(f"j={j} {tag}: {n} boxes, {len(leaves)} leaves, {len(empties)} empty, done {done_vol/tot_vol*100:.3f}% vol, best {best:.6f}, {time.time()-t0:.0f}s",flush=True)
    save(final=True); print(f"FINAL j={j} {tag}: {n} boxes, {len(leaves)} leaves, {len(empties)} empty, best {best:.6f}, {time.time()-t0:.0f}s",flush=True)
if __name__=="__main__":
    j=int(sys.argv[1]); bstar=float(sys.argv[2]); part=int(sys.argv[3])
    cuts=[-1.0,-1/3,1/3,1.0]; lo=[-1.0]*6; hi=[1.0]*6; lo[0]=cuts[part]; hi[0]=cuts[part+1]
    worker(j,bstar,lo,hi,f"p{part}")