"""Verify one final search part (par_j<j>_p<k>.json) in Arb, resumably.
Items are processed in fixed chunks; each finished chunk is appended to <part>.vchunks.jsonl, and a
restart skips the chunks already recorded. A summary is printed and written to <part>.verified.json.
Usage: python verify_part.py par_j1_p0.json NPROC [SHARD NSHARDS]
With shards, instance SHARD handles the chunks with index = SHARD mod NSHARDS and writes its own
<part>.vchunks.s<SHARD>.jsonl; all chunk files of a part are read when resuming and summarizing."""
import sys, json, time, os, glob
from multiprocessing import Pool
import verify4
CH=2000
def work_chunk(args):
    ci,items=args
    tot=dict(ok=0,unresolved=0,sub_empty=0,sub_leaf=0,max_ub=-1.0,bad=[])
    for it in items:
        st,info=verify4.work(it)
        tot[st]+=1; tot["sub_empty"]+=info["empty"]; tot["sub_leaf"]+=info["leaf"]
        tot["max_ub"]=max(tot["max_ub"],info["max_ub"])
        if st!="ok": tot["bad"].append(info["unresolved"])
    return ci,tot
if __name__=="__main__":
    fn=sys.argv[1]; nproc=int(sys.argv[2])
    d=json.load(open(fn)); assert d["final"], fn+" not final"
    j=d["j"]; pat=list(__import__("bbC").PATS[j]); V=d["V"]
    items=[(pat,V,e) for e in d["empties"]]+[(pat,V,l) for l in d["leaves"]]
    chunks=[(i,items[i*CH:(i+1)*CH]) for i in range((len(items)+CH-1)//CH)]
    shard,nsh=(int(sys.argv[3]),int(sys.argv[4])) if len(sys.argv)>4 else (None,None)
    base=fn.replace(".json",".vchunks")
    out=base+(".jsonl" if shard is None else f".s{shard}.jsonl"); done={}
    for cf in glob.glob(base+"*.jsonl"):
        for line in open(cf):
            r=json.loads(line); done[r["chunk"]]=r
    todo=[c for c in chunks if c[0] not in done and (shard is None or c[0]%nsh==shard)]
    print(f"{fn}: V={V} items {len(items)} chunks {len(chunks)} done {len(done)} todo {len(todo)}",flush=True)
    t0=time.time()
    with Pool(nproc,initializer=verify4.init) as pool, open(out,"a") as f:
        for k,(ci,tot) in enumerate(pool.imap_unordered(work_chunk,todo)):
            tot["chunk"]=ci; f.write(json.dumps(tot,default=str)+"\n"); f.flush(); done[ci]=tot
            if (k+1)%10==0: print(f"  {len(done)}/{len(chunks)} chunks, {time.time()-t0:.0f}s",flush=True)
    done={}
    for cf in glob.glob(base+"*.jsonl"):
        for line in open(cf):
            r=json.loads(line); done[r["chunk"]]=r
    if len(done)<len(chunks):
        print(f"shard finished; {len(done)}/{len(chunks)} chunks recorded so far",flush=True); sys.exit(0)
    S=dict(ok=0,unresolved=0,sub_empty=0,sub_leaf=0); worst=-1.0; bad=[]
    for r in done.values():
        for key in S: S[key]+=r[key]
        worst=max(worst,r["max_ub"]); bad+=r["bad"]
    summ=dict(part=fn,j=j,V=V,items=len(items),counts=S,max_ub=worst,bad=bad)
    json.dump(summ,open(fn.replace(".json",".verified.json"),"w"),default=str)
    print("DONE",fn,"V=",V,S,"max ub",worst,f"{time.time()-t0:.0f}s",flush=True)