"""DEFINITIVE dim-2 census: all 2-dim subspaces Lambda <= Z_q^3 with a unit
coordinate in a normal vector, Lambda ∩ D = empty; exact MIS per quotient.
Usage: python census_dim2.py --q 11"""
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)
    seen=set(); normals=[]
    for nvt in itertools.product(range(q),repeat=3):
        if nvt==(0,0,0): continue
        nv=np.array(nvt,dtype=np.int64)%q
        upos=[i for i,c in enumerate(nv) if gcd(int(c),q)==1]
        if not upos: continue
        i=upos[0]; key=tuple((nv*pow(int(nv[i]),-1,q))%q)
        if key in seen: continue
        seen.add(key); normals.append(key)
    print(f"q={q}: {len(normals)} distinct normal directions", flush=True)
    out=[]; t0=time.time()
    for i,key in enumerate(normals):
        nv=np.array(key,dtype=np.int64)
        dots=allv@nv%q
        lam=allv[dots==0]
        lset=set(map(tuple,lam.tolist()))
        if not all(t not in Dset for t in lset if t!=(0,0,0)):
            out.append({"normal":list(key),"valid":False,"union":None,"status":"INVALID"})
            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*len(lam)==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
        out.append({"normal":list(key),"valid":True,"alpha":alpha,
                    "union":alpha*len(lam),"status":stn,"classes":n,"lambda_size":len(lam)})
        if (i+1)%20==0: print(f"  {i+1}/{len(normals)} [{time.time()-t0:.0f}s]",flush=True)
    print(f"done [{time.time()-t0:.0f}s]",flush=True)
    return out

if __name__=="__main__":
    ap=argparse.ArgumentParser(); ap.add_argument("--q",type=int,required=True)
    a=ap.parse_args()
    out=run(a.q)
    valid=[r for r in out if r.get("valid")]
    opt=[r for r in valid if r["status"]=="OPTIMAL"]
    print(f"\nvalid: {len(valid)} | OPTIMAL: {len(opt)}")
    if opt:
        mx=max(r["union"] for r in opt)
        print("max union:", mx, "attained by", sum(1 for r in opt if r["union"]==mx))
        from collections import Counter
        print("histogram:", dict(sorted(Counter(r["union"] for r in opt).items())))
    json.dump(out,open(rf"C:\Users\jack\recens-run\sc7\dim2_def_q{a.q}.json","w"),indent=1,default=int)