# Finite root-system checks for E6/F4 (A3, characteristic-p verification).
# Standard library only.  Bourbaki labels 1..6 (index 0..5): chain 1-3-4-5-6, node 2 on node 4.
from fractions import Fraction as Fr
from itertools import combinations
import sys

n = 6
edges = [(0,2),(2,3),(3,4),(4,5),(1,3)]
A = [[2 if i==j else 0 for j in range(n)] for i in range(n)]
for i,j in edges:
    A[i][j] = A[j][i] = -1
# simply laced: <alpha_i, alpha_j^vee> = A[i][j]; roots = coroots under identification.
def pair(r, c):            # <root r (simple-root coords), coroot c (simple-coroot coords)>
    return sum(r[i]*A[i][j]*c[j] for i in range(n) for j in range(n))
def refl(i, v):            # s_i on root lattice coordinates
    v = list(v); k = sum(v[j]*A[j][i] for j in range(n)); v[i] -= k; return tuple(v)
# positive roots
simple = [tuple(1 if j==i else 0 for j in range(n)) for i in range(n)]
roots = set(simple); frontier = list(simple)
while frontier:
    new = []
    for r in frontier:
        for i in range(n):
            s = refl(i, r)
            if s not in roots and all(x>=0 for x in s) or (s not in roots and all(x<=0 for x in s)):
                roots.add(s); new.append(s)
    frontier = new
allroots = set()
for r in list(roots):
    allroots.add(r); allroots.add(tuple(-x for x in r))
# close under reflections
changed = True
while changed:
    changed = False
    for r in list(allroots):
        for i in range(n):
            s = refl(i, r)
            if s not in allroots:
                allroots.add(s); changed = True
pos = sorted([r for r in allroots if all(x>=0 for x in r)])
print("number of roots", len(allroots), "positive", len(pos))
assert len(allroots) == 72 and len(pos) == 36
black = [1,2,3,4]   # nodes 2,3,4,5
PhiL = [r for r in allroots if r[0]==0 and r[5]==0]
print("|Phi_L| =", len(PhiL))
# theta = -w_L on root lattice (family page matrix, columns = images of simple roots)
thcols = {0:(-1,-1,-2,-2,-1,0), 1:(0,1,0,0,0,0), 2:(0,0,1,0,0,0), 3:(0,0,0,1,0,0), 4:(0,0,0,0,1,0), 5:(0,-1,-1,-2,-2,-1)}
def theta(v):
    out = [0]*n
    for j in range(n):
        for i in range(n):
            out[i] += thcols[j][i]*v[j]
    return tuple(out)
assert all(theta(theta(r)) == r for r in allroots)
assert all(theta(r) in allroots for r in allroots)
# theta preserves the pairing (A theta = theta^T A)
for r in allroots:
    for s in allroots:
        assert pair(theta(r), theta(s)) == pair(r, s)
fixed = [r for r in allroots if theta(r) == r]
neg = [r for r in allroots if theta(r) == tuple(-x for x in r)]
print("theta-fixed roots:", len(fixed), "equal to Phi_L:", set(fixed)==set(PhiL), "; roots with theta b = -b:", len(neg))
# beta + theta beta a root?
bt = [r for r in allroots if r not in PhiL and tuple(a+b for a,b in zip(r,theta(r))) in allroots]
print("roots outside L with beta+theta(beta) a root:", len(bt))
# restricted roots (c1,c6) multiplicities
from collections import Counter
res = Counter((r[0], r[5]) for r in allroots if r not in PhiL)
print("restricted-root multiplicities:", dict(res))
# theta maps Phi^+ \ Phi_L to negative roots
assert all(all(x<=0 for x in theta(r)) for r in pos if r not in PhiL)
# no root orthogonal to all black coroots  <=> Z_G(T^theta) = T
orth = [r for r in allroots if all(pair(r, simple[j])==0 for j in black)]
print("roots orthogonal to all black coroots:", len(orth))
# dim g^theta
t_fix = 4  # +1 eigenspace of -w_L on the coroot space = black span
pairs = len([r for r in allroots if r not in PhiL])//2
print("dim g^theta = t^theta + |Phi_L| + #pairs =", t_fix, "+", len(PhiL), "+", pairs, "=", t_fix+len(PhiL)+pairs)

# ---- weights: Dynkin-label coordinates
def wrefl(i, lam):   # s_i on weights in Dynkin labels: lam - lam_i alpha_i, alpha_i in Dynkin labels = row i of A
    lam = list(lam); k = lam[i]
    for j in range(n): lam[j] -= k*A[i][j]
    return tuple(lam)
def orbit(lam):
    seen = {lam}; fr = [lam]
    while fr:
        new = []
        for l in fr:
            for i in range(n):
                m = wrefl(i, l)
                if m not in seen: seen.add(m); new.append(m)
        fr = new
    return seen
w1 = (1,0,0,0,0,0); w6 = (0,0,0,0,0,1)
V27 = orbit(w1)
print("|weights of V(omega1)| =", len(V27))
inX = sorted([l for l in V27 if all(l[j]==0 for j in black)])
print("weights of V(omega1) in Z omega1 + Z omega6 (Dynkin labels):", inX)
# zero-sum triples
V27l = sorted(V27)
triads = [t for t in combinations(V27l, 3) if all(sum(x[k] for x in t)==0 for k in range(n))]
rep = [t for t in combinations(V27l,2) if all(2*t[0][k]+t[1][k]==0 for k in range(n)) or all(t[0][k]+2*t[1][k]==0 for k in range(n))]
print("number of zero-sum triples of distinct weights:", len(triads), "; pairs with 2mu+nu=0:", len(rep))
T0 = tuple(sorted([(1,0,0,0,0,0),(-1,0,0,0,0,1),(0,0,0,0,0,-1)]))
assert T0 in [tuple(sorted(t)) for t in triads]
# W-orbit of the triad T0
seen = {T0}; fr = [T0]
while fr:
    new = []
    for t in fr:
        for i in range(n):
            u = tuple(sorted(wrefl(i, x) for x in t))
            if u not in seen: seen.add(u); new.append(u)
    fr = new
print("W-orbit of the triad {omega1, omega6-omega1, -omega6} has size", len(seen))
# |W| via orbit of a regular weight
rho = (1,1,1,1,1,1)
print("|W(E6)| =", len(orbit(rho)))
# W-orbit of the set Phi_L (as frozenset of roots)
PL = frozenset(PhiL); seen = {PL}; fr = [PL]
while fr:
    new = []
    for S in fr:
        for i in range(n):
            U = frozenset(refl(i, r) for r in S)
            if U not in seen: seen.add(U); new.append(U)
    fr = new
print("W-orbit of Phi_L has size", len(seen), "=> |N_W(Phi_L)| =", 51840//len(seen))

# ---- matrices for Weyl group elements acting on the root lattice
def mat_of_word(word):   # word applied right-to-left: word=[i1,...,ik] means s_i1 ... s_ik
    cols = []
    for j in range(n):
        v = simple[j]
        for i in reversed(word): v = refl(i, v)
        cols.append(v)
    return cols
def apply(cols, v):
    out = [0]*n
    for j in range(n):
        for i in range(n): out[i] += cols[j][i]*v[j]
    return tuple(out)
def refl_root(beta, v):   # s_beta(v) = v - <v, beta^vee> beta
    k = pair(v, beta); return tuple(v[i]-k*beta[i] for i in range(n))
a1 = simple[0]; a6 = simple[5]
b1 = tuple(-x for x in theta(a1)); b6 = tuple(-x for x in theta(a6))
print("beta1 =", b1, " beta6 =", b6, " (alpha1,beta1) =", pair(a1,b1), " (alpha6,beta6) =", pair(a6,b6))
print("alpha1+-beta1 roots?", tuple(x+y for x,y in zip(a1,b1)) in allroots, tuple(x-y for x,y in zip(a1,b1)) in allroots,
      "; alpha6+-beta6 roots?", tuple(x+y for x,y in zip(a6,b6)) in allroots, tuple(x-y for x,y in zip(a6,b6)) in allroots)
def w1op(v): return refl_root(a1, refl_root(b1, v))
def w6op(v): return refl_root(a6, refl_root(b6, v))
print("s_a1 s_b1 stabilizes Phi_L:", set(w1op(r) for r in PhiL)==set(PhiL), "; s_a6 s_b6 stabilizes Phi_L:", set(w6op(r) for r in PhiL)==set(PhiL))
# action on X(X)=Z omega1+Z omega6: use root-lattice coordinates of 3 omega_i
inv3 = {}  # 3*omega_i in simple-root coordinates
inv3[0] = (4,3,5,6,4,2); inv3[5] = (2,3,4,6,5,4)
def coords_in_w(v):   # v in X(X) (rational), return (a,b) with v = a w1 + b w6 using Dynkin labels <v, alpha_i^vee>
    lab = [sum(v[k]*A[k][i] for k in range(n)) for i in range(n)]
    assert all(lab[j]==0 for j in black), lab
    return (Fr(lab[0],3), Fr(lab[5],3))  # since v = 3*(...) we divide
for name, op in [("s_a1 s_b1", w1op), ("s_a6 s_b6", w6op)]:
    im1 = coords_in_w(op(inv3[0])); im6 = coords_in_w(op(inv3[5]))
    print(name, ": omega1 ->", im1, " omega6 ->", im6)
# order of <w1, w6> acting on the full root lattice
def key(op):
    return tuple(op(s) for s in simple)
elems = {key(lambda v: v)}; gens = [w1op, w6op]
cur = [lambda v: v]
ops = {key(lambda v: v): (lambda v: v)}
frontier = [(lambda v: v)]
while frontier:
    new = []
    for f in frontier:
        for g in gens:
            h = (lambda f=f, g=g: (lambda v: g(f(v))))()
            k = key(h)
            if k not in ops: ops[k] = h; new.append(h)
    frontier = new
print("order of <s_a1 s_b1, s_a6 s_b6> in W:", len(ops))

# ---- Smith normal form of (1 - theta^*) on the weight lattice (torsion of X^*(T)/(1-theta)X^*(T))
# theta on weights in Dynkin coordinates: theta^*(lam)_i = <theta lam, alpha_i^vee> = <lam, theta alpha_i^vee>
def theta_w(lam):
    return tuple(sum(lam[k]*thcols[i][k] for k in range(n)) for i in range(n))
M = []
for j in range(n):
    e = tuple(1 if k==j else 0 for k in range(n))
    im = theta_w(e)
    M.append([e[k]-im[k] for k in range(n)])
def smith_diag(M):
    import copy
    M = [list(r) for r in M]; rws = len(M); cl = len(M[0]); diag = []
    t = 0
    while t < min(rws, cl):
        piv = None
        for i in range(t, rws):
            for j in range(t, cl):
                if M[i][j] != 0 and (piv is None or abs(M[i][j]) < abs(M[piv[0]][piv[1]])): piv = (i,j)
        if piv is None: break
        i,j = piv; M[t],M[i] = M[i],M[t]
        for r in M: r[t],r[j] = r[j],r[t]
        done = False
        while not done:
            done = True
            for i in range(t+1, rws):
                q = M[i][t]//M[t][t]
                for j in range(cl): M[i][j] -= q*M[t][j]
                if M[i][t] != 0:
                    done = False; M[t],M[i] = M[i],M[t]
            for j in range(t+1, cl):
                q = M[t][j]//M[t][t]
                for i in range(rws): M[i][j] -= q*M[i][t]
                if M[t][j] != 0:
                    done = False
                    for r in M: r[t],r[j] = r[j],r[t]
        diag.append(abs(M[t][t])); t += 1
    return diag
print("elementary divisors of (1-theta^*) on X^*(T):", smith_diag(M))

# ---- (A3c) bounds
mu1 = (2,0,1,1,1,1)     # coroot coordinates
def theta_cor(c): return theta(c)   # simply laced, same matrix on coroots
phimu1 = tuple(a-b for a,b in zip(mu1, theta_cor(mu1)))
print("phi(mu1) =", phimu1, " (3 omega1^vee =", inv3[0], ")")
UQ = [r for r in pos if r[0] > 0]
UPM = [r for r in pos if r[0]==0 and r[5] > 0]
print("|Phi(U_Q)| =", len(UQ), " min <b,mu1> on Phi(U_Q) =", min(pair(r,mu1) for r in UQ),
      "; |Phi(U_{P cap M})| =", len(UPM), " min on it =", min(pair(r,mu1) for r in UPM))
zeta = inv3[0]
print("<b, 3omega1^vee> on Phi_M:", set(pair(r,zeta) for r in allroots if r[0]==0), " on Phi(U_Q):", set(pair(r,zeta) for r in UQ))

# ---- orbit dimensions B n_v x0 and the Brion path for w_gamma1
def word_op(word):
    def f(v):
        for i in reversed(word): v = refl(i, v)
        return v
    return f
def inv_word(word): return list(reversed(word))
def dimorbit(word):    # dim B n_v x0, v = s_word
    vinv = word_op(inv_word(word))
    Pv = set(vinv(r) for r in pos)                     # positive roots of B^v = n_v^{-1} B n_v
    R = [r for r in Pv if theta(r) in Pv]
    fixedR = [r for r in R if r in PhiL]
    return 42 - (4 + len(fixedR) + (len(R)-len(fixedR))//2)
def typ(word, alpha):  # type of (B n_v x0, alpha)
    vinv = word_op(inv_word(word))
    return 'G' if vinv(simple[alpha]) in PhiL else 'U'
path = [[], [0], [2,0], [3,2,0], [1,3,2,0], [4,3,2,0]]
for w in path:
    print("orbit s_", [i+1 for i in w], "dim", dimorbit(w))
for w, a in [([],0),([0],2),([2,0],3),([3,2,0],1),([3,2,0],4)]:
    print("type of (Y_", [i+1 for i in w], ", alpha", a+1, ") =", typ(w,a))
wbr = [0,2,3,1,4,3,2,0]   # s1 s3 s4 s2 s5 s4 s3 s1
op = word_op(wbr)
length = len([r for r in pos if all(x<=0 for x in op(r))])
print("w = s1 s3 s4 s2 s5 s4 s3 s1: length", length, "; stabilizes Phi_L:", set(op(r) for r in PhiL)==set(PhiL),
      "; omega1 ->", coords_in_w(op(inv3[0])), " omega6 ->", coords_in_w(op(inv3[5])))
# rho - rho_L pairing with alpha1^vee
two_rho_L = [0]*n
for r in pos:
    if r in PhiL:
        for i in range(n): two_rho_L[i] += r[i]
print("<rho_L, alpha_j^vee> =", [Fr(pair(tuple(two_rho_L), simple[j]),2) for j in range(n)])