import sys; sys.path.insert(0, "/private/tmp/claude-501/-Users-miltonlin-Projects-Trees/75c2c3d9-ef06-4644-b59b-d12583740b0c/scratchpad/solve/A6")
from e6_setup import *
from e6_restrict import dominant_vectors, restrict_orbit_sum, wx_dom
from fractions import Fraction as Fr
import pickle, os
maxnorm=int(sys.argv[1]); target=tuple(int(x) for x in sys.argv[2].split(',')) if len(sys.argv)>2 else (3,0)
fn='e6_restrict_%d.pkl'%maxnorm
if os.path.exists(fn): data=pickle.load(open(fn,'rb'))
else:
    data={}
    for c,lam in dominant_vectors(maxnorm): data[lam]=(c,restrict_orbit_sum(lam))
    pickle.dump(data,open(fn,'wb'))
mus=sorted(data, key=norm2)
kappas=sorted({k for mu in mus for k in data[mu][1]})
print('window: %d orbit sums (norm<=%d), %d target orbit sums M_kappa'%(len(mus),maxnorm,len(kappas)))
def solve(tval):
    t=Fr(tval)
    # matrix rows = kappas, cols = mus ; entries = coefficient of M_kappa in r(m_mu) at t
    Mx=[[sum(cnt*t**e for e,cnt in data[mu][1].get(k,{}).items()) for mu in mus]+[Fr(1) if k==target else Fr(0)] for k in kappas]
    # Gaussian elimination; return rank of M, rank of [M|b]
    rows=[r[:] for r in Mx]; ncol=len(mus); piv=[]
    r=0
    for c in range(ncol+1):
        p=next((i for i in range(r,len(rows)) if rows[i][c]!=0),None)
        if p is None: continue
        rows[r],rows[p]=rows[p],rows[r]; pv=rows[r][c]; rows[r]=[x/pv for x in rows[r]]
        for i in range(len(rows)):
            if i!=r and rows[i][c]!=0:
                f=rows[i][c]; rows[i]=[x-f*y for x,y in zip(rows[i],rows[r])]
        piv.append(c); r+=1
    rankM=len([c for c in piv if c<ncol]); rankAug=len(piv)
    return rankM, rankAug
for tv in [2,3,5,7,Fr(11,3)]:
    rM,rA=solve(tv)
    print('t=%s: rank M=%d, rank [M|b]=%d ->'%(tv,rM,rA), 'M_%s IN image-span'%(target,) if rM==rA else 'NOT in span of this window')