"""Exercises on Young tableaux and the symmetric group, with an answer key computed by methods independent of the apps:
dimensions by removing corners, characters by the Frobenius formula (not Murnaghan-Nakayama), Kostka numbers by listing tableaux,
fixed tabloids by brute force, class sizes by counting permutations.  Standard library only.
Usage:  python3 rep_exercises.py [seed] [how_many]      prints a worksheet and then the answer key."""
import sys, random
from itertools import permutations
from math import factorial
from functools import lru_cache

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 corners(lam):
    return [i for i in range(len(lam)) if i == len(lam) - 1 or lam[i] > lam[i + 1]]

@lru_cache(None)
def dim(lam):
    """standard Young tableaux of shape lam, by removing the corner that holds the largest entry"""
    if sum(lam) <= 1: return 1
    t = 0
    for i in corners(lam):
        q = list(lam); q[i] -= 1
        t += dim(tuple(x for x in q if x))
    return t

def hook(lam, i, j):
    return lam[i] - j - 1 + sum(1 for r in range(i + 1, len(lam)) if lam[r] > j) + 1

def rsk(w):
    P, Q = [], []
    for t, x in enumerate(w, 1):
        r = 0
        while True:
            if r == len(P): P.append([x]); Q.append([t]); break
            row = P[r]; c = next((k for k, v in enumerate(row) if v > x), None)
            if c is None: row.append(x); Q[r].append(t); break
            row[c], x = x, row[c]; r += 1
    return P, Q

def frobenius(lam, mu):
    """chi^lam(mu) = coefficient of x^(lam+delta) in (Vandermonde) * p_mu, with k = len(lam) variables"""
    k = len(lam); delta = list(range(k - 1, -1, -1)); target = [lam[i] + delta[i] for i in range(k)]
    @lru_cache(None)
    def ways(rem, idx):
        if idx == len(mu): return 1 if all(x == 0 for x in rem) else 0
        t = 0
        for j in range(k):
            if rem[j] >= mu[idx]:
                r = list(rem); r[j] -= mu[idx]; t += ways(tuple(r), idx + 1)
        return t
    total = 0
    for s in permutations(range(k)):
        sign = 1
        for a in range(k):
            for b in range(a + 1, k):
                if s[a] > s[b]: sign = -sign
        rem = tuple(target[i] - delta[s[i]] for i in range(k))
        if min(rem) >= 0: total += sign * ways(rem, 0)
    return total

def kostka(lam, mu):
    cells = [(i, j) for i, r in enumerate(lam) for j in range(r)]
    T = {}; used = [0] * len(mu)
    def rec(k):
        if k == len(cells): return 1
        i, j = cells[k]; t = 0
        for v in range(1, len(mu) + 1):
            if used[v - 1] >= mu[v - 1]: continue
            if j and T[(i, j - 1)] > v: continue
            if i and T[(i - 1, j)] >= v: continue
            T[(i, j)] = v; used[v - 1] += 1; t += rec(k + 1); used[v - 1] -= 1
        return t
    return rec(0)

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

def fixed_tabloids(mu, nu):
    n = sum(mu); p = perm_of_type(nu); a = [0] * n; c = [0] * len(mu)
    def rec(k):
        if k == n: return 1 if all(a[p[x]] == a[x] for x in range(n)) else 0
        t = 0
        for r in range(len(mu)):
            if c[r] < mu[r]: a[k] = r; c[r] += 1; t += rec(k + 1); c[r] -= 1
        return t
    return rec(0)

def weyl_dim(lam, m):
    if len(lam) > m: return 0
    l = list(lam) + [0] * (m - len(lam)); num = den = 1
    for i in range(m):
        for j in range(i + 1, m): num *= l[i] - l[j] + j - i; den *= j - i
    return num // den

def cycle_type(p):
    seen, out = set(), []
    for i in range(len(p)):
        if i not in seen:
            c, j = 0, i
            while j not in seen: seen.add(j); j = p[j]; c += 1
            out.append(c)
    return tuple(sorted(out, reverse=True))

def class_size_count(mu):
    return sum(1 for p in permutations(range(sum(mu))) if cycle_type(p) == tuple(mu))

def exercise(rnd, kind):
    if kind == 0:
        n = rnd.randint(4, 8); lam = rnd.choice(list(partitions(n)))
        return f"How many standard Young tableaux have shape {lam}?", dim(lam)
    if kind == 1:
        n = rnd.randint(4, 8); lam = rnd.choice(list(partitions(n))); i = rnd.randrange(len(lam)); j = rnd.randrange(lam[i])
        return f"Hook length of the box in row {i+1}, column {j+1} of the diagram {lam}?", hook(lam, i, j)
    if kind == 2:
        n = rnd.randint(5, 7); w = rnd.sample(range(1, n + 1), n); P, Q = rsk(w)
        return f"Row insertion of the word {w}: last entry of the first row of P?", P[0][-1]
    if kind == 3:
        n = rnd.randint(4, 6); lam = rnd.choice(list(partitions(n))); mu = rnd.choice(list(partitions(n)))
        return f"Character value chi^{lam} on the cycle type {mu}?", frobenius(lam, mu)
    if kind == 4:
        n = rnd.randint(4, 6); mu = rnd.choice([p for p in partitions(n) if len(p) > 1]); lam = rnd.choice(list(partitions(n)))
        return f"Kostka number K_({lam},{mu})?", kostka(lam, mu)
    if kind == 5:
        n = rnd.randint(4, 7); lam = rnd.choice(list(partitions(n)))
        return f"Irreducible constituents of the restriction of S^{lam} from S_{n} to S_{n-1}?", len(corners(lam))
    if kind == 6:
        n = rnd.randint(3, 6); m = rnd.randint(2, 4); lam = rnd.choice(list(partitions(n)))
        return f"Dimension of the GL_{m} module V_{lam} (0 if more than {m} rows)?", weyl_dim(lam, m)
    if kind == 7:
        n = rnd.randint(4, 6); mu = rnd.choice(list(partitions(n)))
        return f"Permutations in S_{n} of cycle type {mu}?", class_size_count(mu)
    n = rnd.randint(4, 6); mu = rnd.choice([p for p in partitions(n) if len(p) > 1]); nu = rnd.choice(list(partitions(n)))
    return f"Tabloids of shape {mu} fixed by a permutation of cycle type {nu}?", fixed_tabloids(mu, nu)

def self_test():
    for n in range(1, 7):
        P = list(partitions(n)); fact = factorial(n)
        for lam in P: assert frobenius(lam, (1,) * n) == dim(lam)
        assert sum(dim(l) ** 2 for l in P) == fact
        for a in P:
            for b in P:
                s = sum(class_size_count(mu) * frobenius(a, mu) * frobenius(b, mu) for mu in P)
                assert s == (fact if a == b else 0), (a, b)
    print("self-test passed: Frobenius characters are orthogonal for n <= 6, dimensions agree with corner removal")

if __name__ == "__main__":
    seed = int(sys.argv[1]) if len(sys.argv) > 1 else 1; count = int(sys.argv[2]) if len(sys.argv) > 2 else 9
    self_test(); rnd = random.Random(seed); items = [exercise(rnd, k % 9) for k in range(count)]
    print(f"\nWorksheet (seed {seed})")
    for k, (q, a) in enumerate(items, 1): print(f"{k:2d}. {q}")
    print("\nAnswer key")
    for k, (q, a) in enumerate(items, 1): print(f"{k:2d}. {a}")