"""Brute-force check of Lemma B.1 (factor graphs as tensor networks) on 200 random factor graphs:
the size, degree and bipartiteness bounds hold, and the contraction of the network equals the partition
function of the factor graph, up to the factor 2^u for the u variables that occur in no scope.
Prints the number of failures, which should be 0.
Usage: python check_tensor_networks.py
"""
import itertools, random
random.seed(3)
def build(n,factors):
    # factors: list of (scope tuple, table dict)
    V=[]; E=[]  # vertices: (name, ports list of edge ids, tensor fn)
    eid=0
    ports={v:[] for v in range(n)}
    verts=[]
    for fi,(sc,T) in enumerate(factors):
        es=[]
        for v in sc:
            es.append(eid); ports[v].append(eid); eid+=1
        verts.append(('F',es,(lambda T: lambda a: T[tuple(a)])(T)))
    eq=lambda a: 1 if len(set(a))<=1 else 0
    for v in range(n):
        d=len(ports[v])
        if d==0: continue
        if d==1: verts.append(('U',[ports[v][0]],lambda a:1))
        elif d==2: verts.append(('E',ports[v][:],eq))
        else:
            ext=ports[v][:]
            # chain tree: d-2 ternary nodes
            prev=None
            nodes=[]
            # first node takes ext[0],ext[1], internal
            cur_edges=[ext[0],ext[1]]
            for t in range(d-2):
                if t<d-3:
                    ie=eid; eid+=1
                    verts.append(('E',cur_edges+[ie],eq)); cur_edges=[ie,ext[t+2]]
                else:
                    verts.append(('E',cur_edges+[ext[d-1]],eq))
    nE=eid
    # subdivide: each edge e -> two edges e_a (old side) e_b (new vertex side)
    newverts=[]
    edge_ct=0
    mapping={}
    for e in range(nE):
        mapping[e]=(2*e,2*e+1)
    # each edge has two endpoints among verts; assign first occurrence -> 2e, second -> 2e+1
    seen={}
    for (typ,es,fn) in verts:
        nes=[]
        for e in es:
            if e in seen: nes.append(2*e+1)
            else: seen[e]=1; nes.append(2*e)
        newverts.append((typ,nes,fn))
    for e in range(nE): newverts.append(('S',[2*e,2*e+1],eq))
    return newverts,2*nE
bad=0
for trial in range(200):
    n=random.randint(1,5); a=random.randint(1,3); nf=random.randint(1,4)
    factors=[]
    for _ in range(nf):
        r=random.randint(1,min(a,n)); sc=tuple(random.sample(range(n),r))
        T={x:random.randint(0,3) for x in itertools.product([0,1],repeat=r)}
        factors.append((sc,T))
    A=sum(len(sc) for sc,_ in factors)
    used=set(v for sc,_ in factors for v in sc); u=n-len(used)
    Z=sum( __import__('math').prod(T[tuple(x[v] for v in sc)] for sc,T in factors) for x in itertools.product([0,1],repeat=n))
    verts,nE=build(n,factors)
    assert nE<=4*A and len(verts)<=4*A, (nE,len(verts),A)
    assert max(len(es) for _,es,_ in verts)<=max(a,3)
    # bipartite: S vertices vs others; every edge has one S endpoint
    cnt={}
    for typ,es,_ in verts:
        for e in es: cnt.setdefault(e,[]).append(typ)
    assert all(len(v)==2 and (v.count('S')==1) for v in cnt.values())
    if nE>22: continue
    con=0
    for x in itertools.product([0,1],repeat=nE):
        p=1
        for typ,es,fn in verts:
            p*=fn([x[e] for e in es])
            if p==0: break
        con+=p
    if Z!=2**u*con: bad+=1
print("fg2tn failures:",bad)