# E6 root data in Bourbaki labelling; cocharacters in coroot coordinates (basis alpha_i^vee).
# Pairing <alpha_j, lambda> = sum_i A[j][i]*lambda_i where A[i][j] = <alpha_i, alpha_j^vee> (Cartan matrix).
from fractions import Fraction as Fr
from itertools import product
A = [[2,0,-1,0,0,0],[0,2,0,-1,0,0],[-1,0,2,-1,0,0],[0,-1,-1,2,-1,0],[0,0,0,-1,2,-1],[0,0,0,0,-1,2]]
n=6
def pair_root(j, lam):            # <alpha_j, lam>, lam in coroot coords (A symmetric so row=col)
    return sum(A[j][i]*lam[i] for i in range(n))
def refl(j, lam):                 # s_j(lam) = lam - <alpha_j,lam> alpha_j^vee
    c = pair_root(j, lam); l=list(lam); l[j]-=c; return tuple(l)
def norm2(lam):                   # (lam,lam) with (alpha_i^vee,alpha_j^vee)=A_ij (simply laced)
    return sum(A[i][j]*lam[i]*lam[j] for i in range(n) for j in range(n))
def dominant(lam): return all(pair_root(j,lam)>=0 for j in range(n))
def orbit(lam):
    seen={lam}; frontier=[lam]
    while frontier:
        nf=[]
        for v in frontier:
            for j in range(n):
                w=refl(j,v)
                if w not in seen: seen.add(w); nf.append(w)
        frontier=nf
    return seen
# theta from families/e6-f4.html: column j = theta(alpha_j^vee)
TH=[[-1,0,0,0,0,0],[-1,1,0,0,0,-1],[-2,0,1,0,0,-1],[-2,0,0,1,0,-2],[-1,0,0,0,1,-2],[0,0,0,0,0,-1]]
def theta(lam): return tuple(sum(TH[i][j]*lam[j] for j in range(n)) for i in range(n))
def phi(lam): return tuple(lam[i]-theta(lam)[i] for i in range(n))
def ab(lam):   # coordinates (a,b) of phi(lam) = a w1^vee + b w6^vee : a=<alpha_1,phi lam>, b=<alpha_6,phi lam>
    p=phi(lam); return (pair_root(0,p), pair_root(5,p))
# checks
I=tuple(tuple(1 if i==j else 0 for j in range(n)) for i in range(n))
assert all(theta(theta(e))==e for e in I)
assert all(norm2(theta(e))==norm2(e) for e in I)
for e in I: assert pair_root(1,phi(e))==0 and pair_root(2,phi(e))==0 and pair_root(3,phi(e))==0 and pair_root(4,phi(e))==0
# positive roots of L = D4 on nodes 2,3,4,5 (indices 1..4), as root-coordinate vectors; use coroot orbit of alpha_j^vee restricted
black=[1,2,3,4]
def orbit_sub(lam, gens):
    seen={lam}; fr=[lam]
    while fr:
        nf=[]
        for v in fr:
            for j in gens:
                w=refl(j,v)
                if w not in seen: seen.add(w); nf.append(w)
        fr=nf
    return seen
PhiL=set()
for j in black: PhiL|=orbit_sub(I[j],black)
PhiL_pos=[r for r in PhiL if all(c>=0 for c in r)]
assert len(PhiL)==24 and len(PhiL_pos)==12
# 2 rho_L as a coroot-coordinate vector (identify roots and coroots, simply laced): sum of positive roots
two_rhoL=tuple(sum(r[i] for r in PhiL_pos) for i in range(n))
def two_rhoL_pair(lam):  # <2 rho_L, lam> = (2rho_L, lam) via the form
    return sum(A[i][j]*two_rhoL[i]*lam[j] for i in range(n) for j in range(n))
if __name__=='__main__':
    print('2rho_L =',two_rhoL, ' <2rho_L, alpha_i^vee> =',[two_rhoL_pair(e) for e in I])
    print('phi(alpha_i^vee) in (a,b):',[ab(e) for e in I])
    Phi=orbit(I[0]); print('|Phi|=',len(Phi))
    # restricted roots: ab of roots
    from collections import Counter
    print('ab of roots:',Counter(ab(r) for r in Phi))