# S1: leading coefficients a_J of Sakellaridis's P_lambda, three ways.
# (1) CT formula  a_J = CT( prod_{Theta+} (1 - sigma y^{2r} e^theta) * prod_{Phi_J+} (1 - e^{-gamma}) )
# (2) inversion formula  a_J = sum_{u in W_J} (-1)^{l(u)} prod_{theta in N_Theta(u)} (-sigma y^{2r})
# (3) direct: coefficient of e^lambda in Pi_lambda = prod_{Theta+}(1 - sigma y^{2r} T_theta) s_lambda,
#     s_nu = sum_w (-1)^w e^{rho-w rho+w nu} / prod_{gamma>0}(1-e^gamma)  (Kostant partition function)
# Polynomials in y = q^{-1/2} are dicts {exponent: int}.  Standard library only.
from fractions import Fraction as Fr
from itertools import combinations, product
import sys

def padd(a, b, s=1):
    c = dict(a)
    for k, v in b.items():
        c[k] = c.get(k, 0) + s * v
        if c[k] == 0: del c[k]
    return c
def pmul(a, b):
    c = {}
    for k1, v1 in a.items():
        for k2, v2 in b.items():
            c[k1 + k2] = c.get(k1 + k2, 0) + v1 * v2
    return {k: v for k, v in c.items() if v}
def pstr(a):
    if not a: return "0"
    return " + ".join(("%d" % v if k == 0 else ("%s y^%d" % ("" if v == 1 else ("-" if v == -1 else str(v)), k))) for k, v in sorted(a.items()))

class RootSystem:
    def __init__(self, cor, pair):
        # cor: simple coroots (vectors); pair(v, j) = <v, gamma_j>
        self.cor = [tuple(Fr(x) for x in c) for c in cor]
        self.pair = pair; self.n = len(cor); self.dim = len(cor[0])
        # generate positive coroots by closure, and W as list of (word, matrix action via function)
        self.pos = self._positive()
        self.rho = tuple(sum(c[i] for c in self.pos) / 2 for i in range(self.dim))
        self.W = self._weyl()
    def s(self, j, v):
        p = self.pair(v, j)
        return tuple(v[i] - p * self.cor[j][i] for i in range(self.dim))
    def act(self, word, v):
        for j in reversed(word): v = self.s(j, v)
        return v
    def _positive(self):
        seen = set(self.cor); frontier = list(self.cor)
        while frontier:
            new = []
            for v in frontier:
                for j in range(self.n):
                    w = self.s(j, v)
                    if w not in seen and self.is_pos(w):
                        seen.add(w); new.append(w)
            frontier = new
        return sorted(seen)
    def coords(self, v):
        # coordinates in simple coroots (assumes simple coroots are a basis of their span; solve by least squares on rationals)
        n, d = self.n, self.dim
        # Gaussian elimination on augmented system
        M = [[self.cor[j][i] for j in range(n)] + [v[i]] for i in range(d)]
        r = 0; piv = []
        for c in range(n):
            pr = next((i for i in range(r, d) if M[i][c] != 0), None)
            if pr is None: continue
            M[r], M[pr] = M[pr], M[r]
            M[r] = [x / M[r][c] for x in M[r]]
            for i in range(d):
                if i != r and M[i][c] != 0:
                    f = M[i][c]; M[i] = [a - f * b for a, b in zip(M[i], M[r])]
            piv.append(c); r += 1
        for i in range(r, d):
            if M[i][n] != 0: return None  # not in span
        x = [Fr(0)] * n
        for i, c in enumerate(piv): x[c] = M[i][n]
        return tuple(x)
    def is_pos(self, v):
        c = self.coords(v)
        return c is not None and all(x >= 0 for x in c) and any(x > 0 for x in c)
    def _weyl(self):
        # elements as reduced words, found by BFS on the orbit of a regular vector
        reg = tuple(sum(Fr(k + 1) * self.cor[k][i] for k in range(self.n)) for i in range(self.dim))
        # use a generic vector to separate elements
        gen = tuple(Fr(1 + 7 * i, 3 + i) for i in range(self.dim))
        key = lambda word: tuple(self.act(word, c) for c in self.cor) + (self.act(word, gen),)
        elems = {key(()): ()}; frontier = [()]
        while frontier:
            new = []
            for w in frontier:
                for j in range(self.n):
                    w2 = (j,) + w
                    k = key(w2)
                    if k not in elems:
                        elems[k] = w2; new.append(w2)
            frontier = new
        return list(elems.values())  # BFS => words are reduced, len = length

def antidom_rep(R, nu):
    # dot action u*(nu) = u(nu - rho) + rho ; find u with mu := u^{-1}*nu antidominant-regular shifted; return (sign, mu) or (0,None)
    v = tuple(nu[i] - R.rho[i] for i in range(R.dim))
    sign = 1
    for _ in range(1000):
        js = [j for j in range(R.n) if R.pair(v, j) > 0]
        if not js: break
        v = R.s(js[0], v); sign = -sign
    if any(R.pair(v, j) == 0 for j in range(R.n)):
        return 0, None
    return sign, tuple(v[i] + R.rho[i] for i in range(R.dim))

def kostant(R, kappa, memo):
    c = R.coords(kappa)
    if c is None or any(x < 0 for x in c) or any(x.denominator != 1 for x in c): return 0
    if all(x == 0 for x in c): return 1
    key = c
    if key in memo: return memo[key]
    # count multisets of positive coroots summing to kappa: recursion on ordered list
    pos = [R.coords(p) for p in R.pos]
    def rec(i, rem):
        if all(x == 0 for x in rem): return 1
        if i == len(pos): return 0
        tot = 0; p = pos[i]; r = rem
        while all(x >= 0 for x in r):
            tot += rec(i + 1, r)
            r = tuple(a - b for a, b in zip(r, p))
        return tot
    memo[key] = rec(0, c)
    return memo[key]

def mult(R, mu, lam, memo):
    # coefficient of e^lam in s_mu (lowest-weight Weyl character)
    tot = 0
    for w in R.W:
        sgn = (-1) ** len(w)
        num = tuple(R.rho[i] - R.act(w, R.rho)[i] + R.act(w, mu)[i] for i in range(R.dim))
        tot += sgn * kostant(R, tuple(lam[i] - num[i] for i in range(R.dim)), memo)
    return tot

def xpoly(sig, r):
    return {int(2 * r): sig}

def a_direct(R, Thp, lam):
    memo = {}; total = {}
    for k in range(len(Thp) + 1):
        for I in combinations(range(len(Thp)), k):
            cI = {0: 1}
            vec = list(lam)
            for i in I:
                th, sig, r = Thp[i]
                cI = pmul(cI, {int(2 * r): -sig})
                vec = [vec[t] + th[t] for t in range(R.dim)]
            eps, mu = antidom_rep(R, tuple(vec))
            if eps == 0: continue
            m = mult(R, mu, lam, memo)
            if m: total = padd(total, {e: eps * m * v for e, v in cI.items()})
    return total

def stab_J(R, lam):
    return [j for j in range(R.n) if R.pair(lam, j) == 0]

def sub_W(R, J):
    return [w for w in R.W if all(j in J for j in w)] if J is not None else R.W

def a_inversion(R, Thp, J):
    tot = {}
    for u in sub_W(R, J):
        term = {0: (-1) ** len(u)}
        for th, sig, r in Thp:
            img = R.act(u, th)
            if R.is_neg_theta(img):
                term = pmul(term, {int(2 * r): -sig})
        tot = padd(tot, term)
    return tot

def a_CT(R, Thp, J):
    # f = prod (1 - sig y^{2r} e^th); g = prod_{Phi_J+}(1 - e^{-gamma}); CT(fg)
    f = {tuple(Fr(0) for _ in range(R.dim)): {0: 1}}
    for th, sig, r in Thp:
        nf = {}
        for v, c in f.items():
            nf[v] = padd(nf.get(v, {}), c)
            v2 = tuple(v[i] + th[i] for i in range(R.dim))
            nf[v2] = padd(nf.get(v2, {}), pmul(c, {int(2 * r): -sig}))
        f = {k: v for k, v in nf.items() if v}
    PJ = [p for p in R.pos if all((R.coords(p)[k] == 0) or (k in J) for k in range(R.n))]
    g = {tuple(Fr(0) for _ in range(R.dim)): 1}
    for p in PJ:
        ng = {}
        for v, c in g.items():
            ng[v] = ng.get(v, 0) + c
            v2 = tuple(v[i] - p[i] for i in range(R.dim))
            ng[v2] = ng.get(v2, 0) - c
        g = {k: v for k, v in ng.items() if v}
    tot = {}
    for v, c in f.items():
        mv = tuple(-x for x in v)
        if mv in g: tot = padd(tot, {e: g[mv] * x for e, x in c.items()})
    return tot

def check(name, R, Thp, lams):
    pos_set = Thp
    def is_neg_theta(v):
        return any(tuple(-x for x in v) == t[0] for t in Thp)
    R.is_neg_theta = is_neg_theta
    print("==", name, "|W| =", len(R.W), "|Theta+| =", len(Thp))
    for lam in lams:
        lam = tuple(Fr(x) for x in lam)
        J = stab_J(R, lam)
        assert all(R.pair(lam, j) <= 0 for j in range(R.n)), "lambda not antidominant"
        d = a_direct(R, Thp, lam); c1 = a_CT(R, Thp, J); c2 = a_inversion(R, Thp, J)
        print("  lambda", [str(x) for x in lam], "J", J, "| direct:", pstr(d), "| CT:", pstr(c1), "| inversion:", pstr(c2), "| agree:", d == c1 == c2)

if __name__ == "__main__":
    # E6/F4: W_X = W(A2), Lambda_X = coroot lattice, Theta+ = positive coroots, sigma=+, r=4  (y^8 = t = q^-4)
    A2 = RootSystem([(1, 0), (0, 1)], lambda v, j: (2 * v[0] - v[1]) if j == 0 else (-v[0] + 2 * v[1]))
    Th = [((Fr(1), Fr(0)), 1, Fr(4)), ((Fr(0), Fr(1)), 1, Fr(4)), ((Fr(1), Fr(1)), 1, Fr(4))]
    check("E6/F4 (A2, Theta+=Phi+, t=y^8)", A2, Th, [(0, 0), (-1, -2), (-2, -1), (-1, -1), (-3, -3), (-2, -4)])
    # T\PGL2 split: Lambda = Z varpi, gamma^vee = 2 varpi, Theta+ = {varpi, varpi}, r = 1/2
    A1 = RootSystem([(2,)], lambda v, j: v[0])
    check("T\\PGL2 split", A1, [((Fr(1),), 1, Fr(1, 2)), ((Fr(1),), 1, Fr(1, 2))], [(0,), (-1,), (-2,), (-3,)])
    check("T\\PGL2 non-split", A1, [((Fr(1),), 1, Fr(1, 2)), ((Fr(1),), -1, Fr(1, 2))], [(0,), (-1,), (-2,)])
    # Sakellaridis Example 7.2.4: PGL2\PGL2^3, W = (Z/2)^3, gamma^vee_i = 2 e_i, colours e1+e2-e3 etc., r = 1/2
    A13 = RootSystem([(2, 0, 0), (0, 2, 0), (0, 0, 2)], lambda v, j: v[j])
    Th3 = [((Fr(1), Fr(1), Fr(-1)), 1, Fr(1, 2)), ((Fr(1), Fr(-1), Fr(1)), 1, Fr(1, 2)),
           ((Fr(-1), Fr(1), Fr(1)), 1, Fr(1, 2)), ((Fr(1), Fr(1), Fr(1)), 1, Fr(1, 2))]
    check("Ex 7.2.4 PGL2^3/PGL2", A13, Th3, [(0, 0, 0), (-1, 0, 0), (0, -2, 0), (-1, -1, 0), (-2, 0, -1), (-1, -1, -1), (-2, -3, -1)])
    # B2 (simple coroots: short gc1, long gc2 in the coroot system C2), mixed types with proportional colours:
    # gamma_1 type G (r=a), gamma_2 type T split (gamma^vee_2/2 twice, r + r' = 1/2 + 1/2)
    # coordinates: coroot basis; pairing via Cartan matrix of B2 roots: <gc1,g1>=2,<gc1,g2>=-1,<gc2,g1>=-2,<gc2,g2>=2
    B2 = RootSystem([(1, 0), (0, 1)], lambda v, j: (2 * v[0] - 2 * v[1]) if j == 0 else (-v[0] + 2 * v[1]))
    print("  B2 positive coroots (coroot coords):", [tuple(str(x) for x in p) for p in B2.pos])
    ThB = []
    for p in B2.pos:
        # W-class of p: the two classes are distinguished by squared length; determine class via coords
        c = B2.coords(p)
        orbit_of_1 = (B2.act((1,), (Fr(1), Fr(0))), (Fr(1), Fr(0)))
    # classes: orbit of gc1
    def orbit(v):
        return {B2.act(w, v) for w in B2.W}
    O1 = orbit((Fr(1), Fr(0))); O2 = orbit((Fr(0), Fr(1)))
    for p in B2.pos:
        if p in O1: ThB.append((p, 1, Fr(3)))
        else:
            h = tuple(x / 2 for x in p); ThB.append((h, 1, Fr(1, 2))); ThB.append((h, 1, Fr(1, 2)))
    check("B2: gamma1 type G (r=3), gamma2 type T split (r=r'=1/2)", B2, ThB, [(0, 0), (-1, -1), (-2, -1), (-3, -2), (-4, -2), (-2, -2)])
    # Xi-lemma check: |N_Theta(v)| = |N_Phi(v)| for all v (vector sums), in each example
    for name, R, T in [("A2", A2, Th), ("A1^3", A13, Th3), ("B2", B2, ThB)]:
        ok = True
        for w in R.W:
            sT = [Fr(0)] * R.dim; sP = [Fr(0)] * R.dim
            for th, sig, r in T:
                img = R.act(w, th)
                if any(tuple(-x for x in img) == t[0] for t in T):
                    sT = [a + b for a, b in zip(sT, th)]
            for p in R.pos:
                if R.is_pos(tuple(-x for x in R.act(w, p))):
                    sP = [a + b for a, b in zip(sP, p)]
            ok = ok and sT == sP
        print("  Xi-lemma |N_Theta(w)| = |N_Phi(w)| for all w in", name, ":", ok)