"""DEFINITIVE line census: all cyclic order-q subgroups <v> <= Z_q^3,
v having >=1 unit coordinate (necessary for order q); normalize by FIRST UNIT
coordinate (any position). Exact MIS per quotient. Usage: python census.py --q 9"""
import sys, time, itertools, json, argparse
from math import gcd
import numpy as np
from ortools.sat.python import cp_model

def run(q, tl=25):
    Dset=set()
    for d in itertools.product((-1,0,1),repeat=3):
        if any(x!=0 for x in d): Dset.add(tuple(x%q for x in d))
    Darr=np.array(sorted(Dset),dtype=np.int64)
    pw=(q**np.arange(3)).astype(np.int64)
    allv=np.array(list(itertools.product(range(q),repeat=3)),dtype=np.int64)
    Nall=len(allv)
    units=[u for u in range(1,q) if gcd(u,q)==1]
    seen=set(); dirs=[]
    for vt in itertools.product(range(q),repeat=3):
        v=np.array(vt,dtype=np.int64)%q
        if not v.any(): continue
        upos=[i for i,c in enumerate(v) if gcd(int(c),q)==1]
        if not upos: continue          # order drops below q
        i=upos[0]; key=tuple((v*pow(int(v[i]),-1,q))%q)
        if key in seen: continue
        seen.add(key); dirs.append(key)
    print(f"q={q}: {len(dirs)} distinct cyclic order-{q} lines", flush=True)
    out=[]; witnesses={}
    t0=time.time()
    for i,key in enumerate(dirs):
        v=np.array(key,dtype=np.int64)
        lam=np.array([(t*v)%q for t in range(q)],dtype=np.int64)
        if not all(tuple(l) not in Dset for l in lam if tuple(l)!=(0,0,0)):
            out.append({"v":[int(x) for x in v],"valid":False,"union":None,"status":"INVALID","classes":None})
            continue
        cand=(allv[:,None,:]+lam[None,:,:])%q
        rep=(cand@pw).min(axis=1)
        uniq,inv=np.unique(rep,return_inverse=True)
        n=len(uniq); assert n*q==Nall
        forb=set(); cls_of=inv
        for dvec in Darr:
            t=((allv+dvec[None,:])%q)
            candt=(t[:,None,:]+lam[None,:,:])%q
            mint=(candt@pw).min(axis=1)
            tgt=np.searchsorted(uniq,mint)
            m=cls_of!=tgt
            for x,y in zip(cls_of[m].tolist(),tgt[m].tolist()):
                forb.add((x,y))
        nbrs=[[j for j in range(n) if (j,ii) in forb or (ii,j) in forb] for ii in range(n)]
        mdl=cp_model.CpModel()
        xs=[mdl.NewBoolVar(str(z)) for z in range(n)]
        for uu in range(n):
            for vv in nbrs[uu]:
                if uu<vv: mdl.Add(xs[uu]+xs[vv]<=1)
        mdl.Maximize(sum(xs))
        sv=cp_model.CpSolver(); sv.parameters.max_time_in_seconds=tl
        sv.parameters.num_search_workers=8
        st=sv.Solve(mdl)
        stn=sv.StatusName(st)
        alpha=int(sv.ObjectiveValue()) if stn in ("OPTIMAL","FEASIBLE") else -1
        rec={"v":[int(x) for x in v],"alpha":alpha,"union":alpha*q,"status":stn,"classes":n}
        out.append(rec)
        if stn=="OPTIMAL":
            chosen=[z for z in range(n) if sv.Value(xs[z])==1]
            RP=[]
            for ci in chosen:
                idxs=np.nonzero(inv==ci)[0]
                codes=(allv[idxs]@pw)
                RP.append(allv[idxs[int(np.argmin(codes))]])
            RP=np.array(RP,dtype=np.int64)
            pts=[tuple((((r+lam[t])%q)).tolist()) for r in RP for t in range(q)]
            witnesses[key]=pts
        if (i+1)%20==0: print(f"  {i+1}/{len(dirs)} [{time.time()-t0:.0f}s]",flush=True)
    print(f"done [{time.time()-t0:.0f}s]",flush=True)
    return dirs,out,witnesses

if __name__=="__main__":
    ap=argparse.ArgumentParser(); ap.add_argument("--q",type=int,required=True)
    args=ap.parse_args()
    q=args.q
    dirs,out,wit=run(q)
    vals=[r for r in out if r["status"]=="OPTIMAL"]
    mx=max(r["union"] for r in vals)
    hot=[r for r in vals if r["union"]==mx]
    print(f"\nOPTIMAL instances: {len(vals)}/{len(out)} | max union = {mx} attained by {len(hot)} lines")
    json.dump(out,open(rf"C:\Users\jack\recens-run\sc7\lines_def_q{q}.json","w"),indent=1)
    # save one witness for the max
    for r in hot[:1]:
        pts=wit[tuple(r["v"])]
        A=np.array(pts,dtype=np.int64); bad=0
        for i in range(len(A)):
            dd=(A[i+1:]-A[i])%q
            bad+=int(((dd<=1)|(dd>=q-1)).all(axis=1).sum())
        print(f"witness for v={r['v']}: {len(A)} pts adjacent-pairs={bad}")
        assert bad==0 and len(A)==mx
        fn=rf"C:\Users\jack\recens-run\sc7\coset_c{q}3_{mx}.txt"
        with open(fn,"w") as f:
            for p in pts: f.write(" ".join(map(str,p))+"\n")
        print("saved",fn)