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 collections import Counter, defaultdict
from fractions import Fraction as Fr
import itertools, pickle
def s1(a,b): return (-a,a+b)
def s6(a,b): return (a+b,-b)
def wx_dom(k):
    a,b=k
    while a<0 or b<0:
        if a<0: a,b=s1(a,b)
        elif b<0: a,b=s6(a,b)
    return (a,b)
def wx_orbit(k):
    seen={k}; fr=[k]
    while fr:
        nf=[]
        for v in fr:
            for w in (s1(*v),s6(*v)):
                if w not in seen: seen.add(w); nf.append(w)
        fr=nf
    return seen
def dominant_vectors(maxnorm):
    # enumerate dominant lam in Z^6 (coroot coords) with norm2<=maxnorm by enumerating fundamental-weight coords c_i>=0 and converting
    # lam = sum c_i omega_i^vee ; omega_i^vee = A^{-1} e_i ; norm = c^T A^{-1} c
    import fractions
    # inverse Cartan via fractions
    M=[[Fr(A[i][j]) for j in range(n)]+[Fr(1 if i==j else 0) for j in range(n)] for i in range(n)]
    for c in range(n):
        p=[r for r in range(c,n) if M[r][c]!=0][0]; M[c],M[p]=M[p],M[c]
        pv=M[c][c]; M[c]=[x/pv for x in M[c]]
        for r in range(n):
            if r!=c and M[r][c]!=0:
                f=M[r][c]; M[r]=[x-f*y for x,y in zip(M[r],M[c])]
    Ainv=[[M[i][n+j] for j in range(n)] for i in range(n)]
    out=[]
    import math
    S=int(math.isqrt(3*maxnorm//2))+1   # norm >= (2/3)(sum c)^2 since all entries of A^{-1} are >= 2/3
    def comps(total, parts):
        if parts==1: yield (total,); return
        for x in range(total+1):
            for rest in comps(total-x, parts-1): yield (x,)+rest
    cands=[c for tot in range(1,S+1) for c in comps(tot,n)]
    for c in cands:
        lam=[sum(Ainv[i][j]*c[j] for j in range(n)) for i in range(n)]
        if all(x.denominator==1 for x in lam):
            lam=tuple(int(x) for x in lam)
            if 0<norm2(lam)<=maxnorm: out.append((c,lam))
    return out
def restrict_orbit_sum(mu):
    """r(m_mu) as dict: W_X-dominant kappa -> Counter{exponent e : count}, coefficient of M_kappa = sum count * t^e, t=q^{1/2}, e=<2rho_L,nu>."""
    res=defaultdict(Counter)
    for nu in orbit(mu):
        k=ab(nu)
        if k==wx_dom(k):
            res[k][two_rhoL_pair(nu)]+=1
    return dict(res)
if __name__=='__main__':
    import sys
    maxnorm=int(sys.argv[1]) if len(sys.argv)>1 else 12
    doms=dominant_vectors(maxnorm)
    doms.sort(key=lambda x: norm2(x[1]))
    data={}
    for c,lam in doms:
        R=restrict_orbit_sum(lam)
        data[lam]=(c,R)
        print('mu(fund coords)=',c,'norm',norm2(lam),'phi(mu)=',ab(lam),'|orbit|',len(orbit(lam)))
        for k in sorted(R, key=lambda k: -(k[0]**2+k[0]*k[1]+k[1]**2)):
            print('    M',k,':',dict(sorted(R[k].items())))
    pickle.dump(data, open('e6_restrict_%d.pkl'%maxnorm,'wb'))