import numpy as np
import itertools
import sys


def s_n(n): 
    return list(itertools.permutations(range(1,n+1)))

def canon_tab(n):
    return list(range(1,n+1))

def split_rows(part, lst):
    assert(sorted(part, reverse=True) == part)
    n = sum(part)
    xs = lst.copy()
    split = []
    for s in part:
        split.append(xs[0:s])
        xs = xs[s:]
    return split

def is_row_stab(part, perm):
    n = sum(part)
    assert(len(perm) == n)
    return list(map(sorted, split_rows(part, canon_tab(n)))) == list(map(sorted, split_rows(part, list(perm))))

def cycles(perm):
    cycles = []
    xs = list(range(1, len(perm)+1))
    while len(xs) > 0:
        c = [xs[0]]
        k = xs[0]
        while (perm[k-1] not in c):
            k = perm[k-1]
            c.append(k)
        xs = [ x for x in xs if x not in c ]
        cycles.append(c)
    return cycles

def cycle_type(perm):
    cs = cycles(perm)
    types = {}
    for c in cs:
        if len(c) not in types:
            types[len(c)] = 1
        else:
            types[len(c)] += 1
    t = []
    for i in range(1, max(perm)+1):
        t.append(types.get(i, 0))
    return tuple(t)

def conjugacy_classes(n):
    cs = {}
    for perm in s_n(n):
        ct = cycle_type(perm)
        if ct not in cs:
            cs[ct] = [perm]
        else:
            cs[ct].append(perm)
    l = list(cs.items())
    return sorted(l, key=lambda x: x[0], reverse=True)

def partitions(n, conj_classes=None):
    parts = []
    if conj_classes == None:
        conj_classes = conjugacy_classes(n)
    for cycle_type, conj in conj_classes:
        part = []
        for i, c in enumerate(cycle_type):
            for s in range(0, c):
                part.append(i+1)
        parts.append(sorted(part, reverse=True))
        #parts.append([(i+1) * cycle_type[i] for i in range(0, n)])
    return list(sorted(parts, reverse=True))

def psi(n):
    cs = conjugacy_classes(n)
    ps = partitions(n, cs)
    m = np.zeros((len(cs), len(cs)))
    for i, part in enumerate(ps):
        row_stabs = set(filter(lambda sigma: is_row_stab(part, sigma), s_n(n)))
        for j, (cycle_type, conj) in enumerate(cs):
            m[i, j] = len(row_stabs & set(conj)) * np.math.factorial(n) / (len(row_stabs) * len(conj))
    return m

def sigma(n):
    cs = conjugacy_classes(n)
    m = np.zeros((len(cs), len(cs)))
    for i, (cycle_type, conj) in enumerate(cs):
        m[i, i] = len(conj) / np.math.factorial(n)
    return m

if __name__ == '__main__':
    if len(sys.argv) >= 2:
        n = int(sys.argv[1])
    else:
        n = 3

    psi = psi(n)
    sigma = sigma(n)

    a = np.dot(psi, np.dot(sigma, np.transpose(psi)))
    k = np.linalg.cholesky(a)
    x = np.dot(np.linalg.inv(k), psi)
    print("Character table for S_n where n =", n)
    print("Rows are indexed by partitions and columns by conjugacy classes")
    print(x.round())