import numpy as np
from itertools import product, combinations

I=np.array([[1,0],[0,1]],complex)
X=np.array([[0,1],[1,0]],complex)
Y=np.array([[0,-1j],[1j,0]],complex)
Z=np.array([[1,0],[0,-1]],complex)
paulis=[I,X,Y,Z]
labels=['I','X','Y','Z']

def kron_all(ops):
    out=ops[0]
    for op in ops[1:]: out=np.kron(out,op)
    return out

def graph_state(n, edges):
    psi=np.ones(2**n,complex)/np.sqrt(2**n)
    # basis index bits qubit 0..n-1 as MSB? phase invariant consistent
    for idx in range(2**n):
        bits=[(idx>>(n-1-q))&1 for q in range(n)]
        phase=1
        for a,b in edges:
            if bits[a]*bits[b]: phase*=-1
        psi[idx]*=phase
    return psi

def corr_tensor(psi,n):
    coeff={}
    for inds in product(range(4), repeat=n):
        op=kron_all([paulis[i] for i in inds])
        val=np.vdot(psi, op@psi)
        if abs(val)>1e-9:
            coeff[inds]=float(np.real_if_close(val))
    return coeff

def matrix_for_cut(coeff,S,n, source_full=False, target_full=False):
    Sc=[i for i in range(n) if i not in S]
    row=[]; col=[]
    for inds in product(range(4), repeat=len(S)):
        if all(i==0 for i in inds): continue
        if source_full and any(i==0 for i in inds): continue
        row.append(inds)
    for inds in product(range(4), repeat=len(Sc)):
        if all(i==0 for i in inds): continue
        if target_full and any(i==0 for i in inds): continue
        col.append(inds)
    M=np.zeros((len(col),len(row))) # target rows, source cols
    for r,tinds in enumerate(col):
        for c,sinds in enumerate(row):
            full=[0]*n
            for q,ind in zip(S,sinds): full[q]=ind
            for q,ind in zip(Sc,tinds): full[q]=ind
            M[r,c]=coeff.get(tuple(full),0.0)
    return M,row,col

def partial_rho(psi, keep, n):
    rho=np.outer(psi, psi.conj()).reshape([2]*n*2)
    trace=[i for i in range(n) if i not in keep]
    # trace out from high to low axes
    for q in sorted(trace, reverse=True):
        rho=np.trace(rho, axis1=q, axis2=q+rho.ndim//2)
    d=2**len(keep)
    return rho.reshape(d,d)

n=4
edges=[(0,1),(1,2),(2,3),(3,0)]
psi=graph_state(n,edges)
coeff=corr_tensor(psi,n)
print('nonzero coeffs', len(coeff))
for cutname,S in [('adjacent12|34',[0,1]),('diagonal13|24',[0,2]),('adjacent14|23',[0,3])]:
    M,_,_=matrix_for_cut(coeff,S,n)
    G,_,_=matrix_for_cut(coeff,S,n,source_full=True,target_full=True)
    norm=np.linalg.svd(M/3, compute_uv=False).sum()
    rawn=np.linalg.svd(M, compute_uv=False).sum()
    gnorm=np.linalg.svd(G/3, compute_uv=False).sum()
    graw=np.linalg.svd(G, compute_uv=False).sum()
    print('\n',cutname)
    print('full normalized',norm,'raw',rawn,'sing',np.linalg.svd(M/3,compute_uv=False))
    print('genuine normalized',gnorm,'raw',graw,'sing',np.linalg.svd(G/3,compute_uv=False))
    rho=partial_rho(psi,S,n)
    print('rho eigen',np.linalg.eigvalsh(rho),'max mixed dist',np.linalg.norm(rho-np.eye(4)/4))
    # list nonzero stabilizers labels
    # print rows/cols maybe