"""Representations of the symmetric group S_n: Specht modules by exact linear algebra, characters as traces, and a check against the
Murnaghan-Nakayama rule and the orthogonality relations.  Standard library only (fractions, itertools, math)."""
from fractions import Fraction as F
from itertools import permutations, product
from math import factorial

def partitions(n, m=None):
    m = n if m is None else m
    if n == 0: yield (); return
    for k in range(min(n, m), 0, -1):
        for rest in partitions(n - k, k): yield (k,) + rest

def standard_tableaux(lam):
    n = sum(lam); out = []; T = [[0] * x for x in lam]; cnt = [0] * len(lam)
    def put(k):
        if k > n: out.append([r[:] for r in T]); return
        for i in range(len(lam)):
            if cnt[i] < lam[i] and (i == 0 or cnt[i - 1] > cnt[i]):
                T[i][cnt[i]] = k; cnt[i] += 1; put(k + 1); cnt[i] -= 1; T[i][cnt[i]] = 0
    put(1); return out

def tabloid(rows): return tuple(tuple(sorted(r)) for r in rows)

def polytabloid(t):
    """e_t = sum over the column group of sgn(sigma) * sigma{t}, as a dict tabloid -> coefficient"""
    cols = [[t[i][j] for i in range(len(t)) if len(t[i]) > j] for j in range(len(t[0]))]
    def sign(p): return (-1) ** sum(1 for a in range(len(p)) for b in range(a) if p[b] > p[a])
    vec = {}
    for choice in product(*[list(permutations(range(len(c)))) for c in cols]):
        rows = [list(r) for r in t]; s = 1
        for j, (c, p) in enumerate(zip(cols, choice)):
            s *= sign(p)
            for i in range(len(c)): rows[i][j] = c[p[i]]
        key = tabloid(rows); vec[key] = vec.get(key, 0) + s
    return {k: v for k, v in vec.items() if v}

def act(p, vec):
    """p is a tuple with p[i-1] = image of i"""
    out = {}
    for tab, c in vec.items():
        key = tabloid([[p[x - 1] for x in r] for r in tab]); out[key] = out.get(key, 0) + c
    return out

def solve_coordinates(basis, w):
    """coordinates of w in the span of the basis vectors (dicts), exact; raises if w is not in the span"""
    keys = sorted({k for b in basis for k in b} | set(w)); m = len(basis)
    rows = [[F(b.get(k, 0)) for b in basis] + [F(w.get(k, 0))] for k in keys]
    piv = []; r = 0
    for c in range(m):
        p = next((i for i in range(r, len(rows)) if rows[i][c] != 0), None)
        if p is None: raise ValueError("dependent basis")
        rows[r], rows[p] = rows[p], rows[r]; rows[r] = [x / rows[r][c] for x in rows[r]]
        for i in range(len(rows)):
            if i != r and rows[i][c] != 0: rows[i] = [x - rows[i][c] * y for x, y in zip(rows[i], rows[r])]
        r += 1
    assert all(rows[i][m] == 0 for i in range(r, len(rows))), "vector outside the span"
    return [rows[i][m] for i in range(m)]

def character(lam, p):
    """trace of p on the Specht module S^lam, in the basis of standard polytabloids"""
    basis = [polytabloid(t) for t in standard_tableaux(lam)]
    tr = F(0)
    for j, b in enumerate(basis): tr += solve_coordinates(basis, act(p, b))[j]
    assert tr.denominator == 1; return int(tr)

def perm_of_type(mu):
    p = []; s = 0
    for c in mu: p += [s + (k + 1) % c + 1 for k in range(c)]; s += c
    return tuple(p)

def class_size(mu):
    n = sum(mu); d = 1
    for c in set(mu): d *= c ** mu.count(c) * factorial(mu.count(c))
    return factorial(n) // d

def murnaghan_nakayama(lam, mu):
    """characters from removing border strips (beta-numbers); independent of the construction above"""
    if not mu: return 1 if sum(lam) == 0 else 0
    r, rest = mu[0], mu[1:]; beta = [x + len(lam) - 1 - i for i, x in enumerate(lam)]; total = 0
    for i, b in enumerate(beta):
        nb = b - r
        if nb < 0 or nb in beta: continue
        sign = (-1) ** sum(1 for x in beta if nb < x < b)
        new = sorted(beta[:i] + [nb] + beta[i + 1:], reverse=True)
        total += sign * murnaghan_nakayama(tuple(x - (len(new) - 1 - k) for k, x in enumerate(new) if x - (len(new) - 1 - k) > 0), rest)
    return total

if __name__ == "__main__":
    for n in range(2, 7):
        parts = list(partitions(n)); table = [[character(l, perm_of_type(m)) for m in parts] for l in parts]
        for a, l in enumerate(parts):
            for b, m in enumerate(parts): assert table[a][b] == murnaghan_nakayama(l, m), (l, m)
        sizes = [class_size(m) for m in parts]
        for a in range(len(parts)):
            for b in range(len(parts)):
                ip = sum(s * table[a][k] * table[b][k] for k, s in enumerate(sizes))
                assert ip == (factorial(n) if a == b else 0)
        print(f"n = {n}: {len(parts)} irreducible characters; traces on Specht modules = Murnaghan-Nakayama; rows orthogonal")
    print("S_3 table (rows (3), (2,1), (1,1,1); columns e, (12), (123)):")
    for l in partitions(3): print(l, [character(l, perm_of_type(m)) for m in [(1, 1, 1), (2, 1), (3,)]])
