# code by Jang Soo Kim
def delete_elements(P,A):
    "Return the poset obtained from P by deleting the elements in A"
    elms = [x for x in  P.list() if x not in A]
    rels = [R for R in P.relations() if R[0] not in A and R[1] not in A]
    return Poset([elms,rels], cover_relations=False)

def add_bottom(P):
    "Return the poset obtained from P by adding the minimum element 0."
    elms = ["zero_hat"] + P.list()
    rels = [["zero_hat",x] for x in P.list()] + P.relations()
    return Poset([elms,rels], cover_relations=False)

def Poin(P,t,omega=0):
    """
    Input:

    P: a poset or a list of numbers [a,b,c,...].
    If P = [a,b,c,...], then P is set to be the cell poset of the partition [a,b,c,...].

    t: a variable

    omega: a labeling of P which is given by default to the first lin ext of P.

    Output: The Poincare polynomial of the poset P.
    The bijection from linear extensions to P-transverse permutations is used.
    """
    if type(P) == list:
        P = Partition(P).cell_poset().dual()
    if omega == 0:
        omega = P.linear_extensions()[0]
    poly = 0
    for L in P.linear_extensions():
        poly += t^(len(P)-len(get_P_transverse_permutation(P,omega,L)))
    return poly

def get_P_transverse_permutation(P,omega,L):
    """
    Input:
    P: a poset
    omega: a labeling of P
    L: a linear extension of P

    Output: The P-transverse permutation corresponding to L with respect to omega.
    """
    Q = add_bottom(P)
    pi = []
    while len(L)>0:
        N = next_level(Q,P,omega,L)
        pi = pi + N[-1]
        [Q, P, omega, L] = N[:-1]
    return pi

def next_level(Q,P,omega,L):
    """
    INPUT: Q,P,omega,L
    Q: poset containing P
    omega: labeling
    L: a linear extension of P

    OUTPUT: [Q1, P1, omega, L1, pi]
    Q1 = P
    P1 = P - minimal blocks
    omega = omega
    L1 = L restricted to P1
    pi = P-transverse permutation of the minimal blocks in P.
    """
    L1 = L
    pi = []
    C = []   # the list of current level elements in P
    D = []   # the list of current level elements in P that are not minimal in Q
    B = []
    for x in L:
        if x not in P.minimal_elements():  # If x is not a minimal element in P this level is done.
            break
        C.append(x)
        if x not in Q.minimal_elements():
            D.append(x)
            if omega.index(x) == min([omega.index(d) for d in D]):
                if len(B)>0:
                    pi.append(B)
                B = [x]
        else:
            B.append(x)
    pi.append(B)
    Q1 = P
    P1 = delete_elements(P,C)
    L1 = L[len(C):]
    return [Q1, P1, omega, L1, pi]

statistic = lambda P: Poin(P, QQ["q"].gen()).leading_coefficient()