import itertools

# Define Pauli multiplication ignoring phase; represent as strings 'I','X','Y','Z'
# Provide function to compute commutation: returns True if commute, False if anticommute.

paulis = ['I','X','Y','Z']

# Single-qubit commutation table: anticommute pairs are X-Z, X-Y, Y-Z
anticommute_pairs = {('X','Z'),('Z','X'),('X','Y'),('Y','X'),('Y','Z'),('Z','Y')}

def single_anticommute(a,b):
    if a=='I' or b=='I' or a==b: return False
    return (a,b) in anticommute_pairs

# For tensor products represented as strings over length 4

def anticommutes(P,Q):
    # Count number of positions where local Paulis anticommute
    count = 0
    for a,b in zip(P,Q):
        if single_anticommute(a,b):
            count += 1
    # anticommute iff odd number of local anticommutes
    return (count % 2)==1

def commutes(P,Q):
    return not anticommutes(P,Q)

# Code stabilizers
S1 = 'XXXX'
S2 = 'ZZZZ'

# Logical operators
X_A = 'XIXI'
Z_A = 'ZZII'
X_B = 'XXII'
Z_B = 'ZIZI'

# Allowed pair Paulis on (2,1) and (0,3)
# Note our order is (0,1,2,3), so to place a 2-qubit Pauli on pair (2,1)
# we fill positions: 0:'I',1:P1,2:P2,3:'I' with P2 on qubit2, P1 on qubit1.

def embed_pair_21(p1,p2):
    # p1 acts on qubit1, p2 acts on qubit2
    return 'I' + p1 + p2 + 'I'

def embed_pair_03(p0,p3):
    # p0 acts on qubit0, p3 acts on qubit3
    return p0 + 'I' + 'I' + p3

# Function to multiply Pauli strings ignoring phase
# We just need for combining P21 and P03, which act on disjoint supports, so multiplication is just overlay

def combine(P,Q):
    # they act on disjoint supports in our construction; assert no conflicts
    out = []
    for a,b in zip(P,Q):
        out.append(a if b=='I' else (b if a=='I' else None))
    # None should not appear
    return ''.join(out)

# Construct all allowed two-qubit Paulis on a pair, with probabilities depending on p
p = 0.123 # placeholder, probabilities symbolic later

pair_nonidentity = [x for x in itertools.product(paulis, paulis) if not (x[0]=='I' and x[1]=='I')]

# We'll classify types per pair: parity (x,z)
XZ = {'I':(0,0),'X':(1,0),'Y':(1,1),'Z':(0,1)}

def pair_type(p,q):
    x = (XZ[p][0] + XZ[q][0]) % 2
    z = (XZ[p][1] + XZ[q][1]) % 2
    return (x,z)

# Verify classes
from collections import defaultdict
classes = defaultdict(list)
for a,b in itertools.product(paulis,paulis):
    classes[pair_type(a,b)].append((a,b))

classes_count = {k:len(v) for k,v in classes.items()}
print('Class sizes per pair type (x,z):', classes_count)

# Now enumerate all P_total = P21 * P03 with types equal for acceptance and compute logical effect

all_pairs = list(itertools.product(paulis, paulis))

# define commutation with logicals -> determine logical effect Ebar as a pair of single-qubit Paulis on (A,B)

def logical_action(P):
    # Determine x,z bits for A and B via commutation
    # x_bit = anticommutes with Z, z_bit = anticommutes with X
    xA = 1 if anticommutes(P, Z_A) else 0
    zA = 1 if anticommutes(P, X_A) else 0
    xB = 1 if anticommutes(P, Z_B) else 0
    zB = 1 if anticommutes(P, X_B) else 0
    # map bits to char
    inv = {(0,0):'I',(1,0):'X',(0,1):'Z',(1,1):'Y'}
    return inv[(xA,zA)] + inv[(xB,zB)]

# Now compute probabilities symbolically in terms of p: we'll tally weights as linear/quadratic polynomials in p and 1-p
from fractions import Fraction

# Probabilities per 2-qubit Pauli on a pair:
# II: 1-p; each other: p/15

def pair_prob(a,b):
    if a=='I' and b=='I':
        return ('lin', Fraction(1,1), Fraction(-1,1))  # Represent as (type, A, B) meaning A + B p, but we want (1-p)
    else:
        return ('constp', Fraction(1,15)) # means (p/15)

# We will accumulate probabilities for acceptance and for GHZ-stabilizing logical actions

from collections import Counter

# We'll symbolically compute: total weights as functions: W = c0 + c1 p + c2 p^2
# Because two independent errors -> polynomial up to p^2.

class Poly:
    def __init__(self,c0=Fraction(0),c1=Fraction(0),c2=Fraction(0)):
        self.c0=Fraction(c0); self.c1=Fraction(c1); self.c2=Fraction(c2)
    def __iadd__(self, other):
        self.c0 += other.c0; self.c1 += other.c1; self.c2 += other.c2
        return self
    def __add__(self, other):
        out=Poly(self.c0,self.c1,self.c2); out+=other; return out
    def __mul__(self, other):
        # multiply polynomials up to degree 2 only; but we will only multiply up to deg1 x deg1 -> deg2 (fine)
        out=Poly()
        out.c0 = self.c0*other.c0
        out.c1 = self.c0*other.c1 + self.c1*other.c0
        out.c2 = self.c0*other.c2 + self.c1*other.c1 + self.c2*other.c0
        return out
    def __repr__(self):
        return f"{self.c0} + {self.c1} p + {self.c2} p^2"

# Build polynomial probability for each (a,b) pair on each pair (21 or 03)

def prob_poly_pair(a,b):
    if a=='I' and b=='I':
        # 1-p = 1 + (-1) p
        return Poly(1, -1, 0)
    else:
        # p/15 = 0 + (1/15) p
        return Poly(0, Fraction(1,15), 0)

# Now enumerate and sum accepted cases.

# For acceptance, type equality across pairs:

def pair_type_key(pair):
    return pair_type(pair[0], pair[1])

# GHZ logical stabilizer set on (A,B): {II, XX, ZZ, YY}
GHZ_stab = {'II', 'XX', 'ZZ', 'YY'}

P_accept = Poly()
P_good = Poly()
# Also tally distribution over logical actions for accepted events
logdist = {}
for a in all_pairs:
    P21 = embed_pair_21(a[1], a[0])  # careful: a is (p0, p1) ??? For pair21 we used p1 on qubit1, p2 on qubit2: here let a=(p1,p2)
    # Actually our 'a' will be (p1,p2) for pair (1,2); we set P21 accordingly
    P21 = embed_pair_21(a[0], a[1])
    prob21 = prob_poly_pair(a[0], a[1])
    type21 = pair_type(a[0], a[1])
    for b in all_pairs:
        P03 = embed_pair_03(b[0], b[1])
        prob03 = prob_poly_pair(b[0], b[1])
        type03 = pair_type(b[0], b[1])
        if type03 != type21:
            continue
        Ptot = combine(P21, P03)
        # Confirm commute with stabilizers
        if not (commutes(Ptot,S1) and commutes(Ptot,S2)):
            # Shouldn't happen: type equality implies commute with both
            pass
        act = logical_action(Ptot)
        logdist.setdefault(act, Poly())
        weight = prob21 * prob03
        P_accept += weight
        logdist[act] += weight
        if act in GHZ_stab:
            P_good += weight

print('P_accept polynomial:', P_accept)
print('P_good polynomial:', P_good)
print('Distribution over logical actions (polynomials):')
for k,v in sorted(logdist.items()):
    print(k, v)

# Simplify to rational coefficients
acc = P_accept
good = P_good

# Compute logical fidelity F_log(p) = good / acc (as rational functions)
# For general p, F(p) = (c0g + c1g p + c2g p^2)/(c0a + c1a p + c2a p^2)

num = (good.c0, good.c1, good.c2)
den = (acc.c0, acc.c1, acc.c2)
print('num:', num)
print('den:', den)