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)

from fractions import Fraction
import itertools

# Define Pauli multiplication ignoring phase
mult_table = {
    ('I','I'):'I', ('I','X'):'X', ('I','Y'):'Y', ('I','Z'):'Z',
    ('X','I'):'X', ('X','X'):'I', ('X','Y'):'Z', ('X','Z'):'Y',
    ('Y','I'):'Y', ('Y','X'):'Z', ('Y','Y'):'I', ('Y','Z'):'X',
    ('Z','I'):'Z', ('Z','X'):'Y', ('Z','Y'):'X', ('Z','Z'):'I',
}

def pauli_mult(a,b):
    return mult_table[(a,b)]

# Conjugation of a Pauli string by CNOT(c,t)
# Mapping on single-qubit letters ignoring phase:
# X_c -> X_c X_t; Z_t -> Z_c Z_t; others unchanged except Y derived from X/Z

single_to_pair = {
    'I': ('I','I'),
    'X': ('X','I'),
    'Y': ('Y','I'),  # treat as XZ product
    'Z': ('Z','I'),
}

# We'll implement conjugation by directly applying rules per letters on c and t

def conj_cnot(P, c, t):
    P = list(P)
    a = P[c]
    b = P[t]
    # We consider effect on c and t jointly by factoring Pauli letters into X^x Z^z, applying known mapping, then multiply
    # Represent letters as (x,z) bits
    def xz(letter):
        if letter=='I': return (0,0)
        if letter=='X': return (1,0)
        if letter=='Y': return (1,1)
        if letter=='Z': return (0,1)
    def letter_from_xz(x,z):
        if x==0 and z==0: return 'I'
        if x==1 and z==0: return 'X'
        if x==0 and z==1: return 'Z'
        if x==1 and z==1: return 'Y'
    xc, zc = xz(a)
    xt, zt = xz(b)
    # Under CNOT(c->t):
    # X_c -> X_c X_t => x_t ^= x_c
    # Z_t -> Z_c Z_t => z_c ^= z_t
    xt ^= xc
    zc ^= zt
    P[c] = letter_from_xz(xc, zc)
    P[t] = letter_from_xz(xt, zt)
    return ''.join(P)

# Build list of gates in forward order
# Gates: CNOT_12, CNOT_10, CNOT_23, CNOT_34, CNOT_04
# Using qubit indices 0..4

gates = [(1,2),(1,0),(2,3),(3,4),(0,4)]

# Build all 2-qubit Pauli errors per pair as strings of length 5
paulis = ['I','X','Y','Z']

# Helper to embed pair paulis on given qubits (i,j) into 5-qubit string

def embed_pair(i,j, a, b):
    P=['I']*5
    P[i]=a
    P[j]=b
    return ''.join(P)

# Precompute propagation of each pair Pauli inserted after gate k through subsequent gates to the end
E_final = [[None]*16 for _ in range(5)]
label_list = []
for idx, a in enumerate(paulis):
    for b_idx, b in enumerate(paulis):
        label_list.append((a,b))
# label_list order consistent: (I,I),(I,X),... But we want index mapping 0..15

for k,(ci,ti) in enumerate(gates):
    for idx,(a,b) in enumerate(label_list):
        P = embed_pair(ci,ti,a,b)
        # Propagate through gates k+1..end
        for (c,t) in gates[k+1:]:
            P = conj_cnot(P, c, t)
        E_final[k][idx] = P

# Probability polynomials for each 2-qubit error on a gate
# Represent polynomial as list of coefficients [c0..c5] for sum c_k p^k
class Poly:
    def __init__(self, deg=5):
        self.c = [Fraction(0) for _ in range(deg+1)]
    def add(self, other):
        for i in range(len(self.c)):
            self.c[i] += other.c[i]
    def mul(self, other):
        r = Poly(len(self.c)-1)
        for i,a in enumerate(self.c):
            for j,b in enumerate(other.c):
                if i+j < len(r.c):
                    r.c[i+j] += a*b
        return r
    def __repr__(self):
        return ' + '.join(f'{self.c[i]} p^{i}' for i in range(len(self.c)) if self.c[i]) or '0'

# For a single gate error option: II -> (1-p), non-identity -> p/15
single_probs = []
for idx,(a,b) in enumerate(label_list):
    P = Poly(5)
    if a=='I' and b=='I':
        P.c[0] = Fraction(1)
        P.c[1] = Fraction(-1)
    else:
        P.c[1] = Fraction(1,15)
    single_probs.append(P)

# Function: classify acceptance and logical goodness from total final Pauli
# Acceptance if: ancilla (qubit 4) letter in {I,Z} and code commutes with XXXX and ZZZZ.
# For logical goodness: data Pauli in the stabilizer of |00>_AB generated by XXXX, ZZZZ, ZZII, ZIZI

# Build group of stabilizers of |00>_AB as set of strings length 4 (data only)
# We'll generate by multiplying generators ignoring phase.

def pauli_mult_letter(a,b):
    return mult_table[(a,b)]

def mult_str4(A,B):
    return ''.join(pauli_mult_letter(a,b) for a,b in zip(A,B))

# 4 data qubits order [0,1,2,3]; ancilla 4 index 4.

# Generators:
G1 = 'XXXX'
G2 = 'ZZZZ'
G3 = 'ZZII'  # Z_A
G4 = 'ZIZI'  # Z_B

# Generate stabilizer group size 16
stab_set = set()
for a in [0,1]:
    for b in [0,1]:
        for c in [0,1]:
            for d in [0,1]:
                P = 'IIII'
                if a: P = mult_str4(P,G1)
                if b: P = mult_str4(P,G2)
                if c: P = mult_str4(P,G3)
                if d: P = mult_str4(P,G4)
                stab_set.add(P)

# Function to get x,z parity sums on data
XZ = {'I':(0,0),'X':(1,0),'Y':(1,1),'Z':(0,1)}

def commute_with_checks(P4):
    xsum = sum(XZ[ch][0] for ch in P4) % 2
    zsum = sum(XZ[ch][1] for ch in P4) % 2
    return (xsum==0) and (zsum==0)

# Evaluate
P_acc = Poly(5)
P_good = Poly(5)

# We will iterate over all patterns using nested loops; However for efficiency, we can apply dynamic programming: combine two gates at a time and convolve polynomials while tracking resulting Pauli. But let's keep nested loops; 16^5 ~ 1e6 iterations ok.

# We'll precompute for each gate the list of (final Pauli string, prob poly) for the 16 options
opts_per_gate = []
for k in range(5):
    lst=[]
    for idx in range(16):
        lst.append((E_final[k][idx], single_probs[idx]))
    opts_per_gate.append(lst)

# Multiply Pauli strings of 5 contributions ignoring phase

def mult5(Ps):
    Q = ['I']*5
    for P in Ps:
        for i in range(5):
            Q[i] = pauli_mult(Q[i], P[i])
    return ''.join(Q)

# Now enumerate
count=0
for opt0 in opts_per_gate[0]:
    for opt1 in opts_per_gate[1]:
        for opt2 in opts_per_gate[2]:
            for opt3 in opts_per_gate[3]:
                for opt4 in opts_per_gate[4]:
                    Ps = [opt0[0],opt1[0],opt2[0],opt3[0],opt4[0]]
                    Pfin = mult5(Ps)
                    # Acceptance conditions
                    anc_ok = (Pfin[4] in ('I','Z'))
                    if not anc_ok:
                        # no acceptance for ancilla=0
                        pass
                    dataP = Pfin[:4]
                    code_ok = commute_with_checks(dataP)
                    if anc_ok and code_ok:
                        # compute probability polynomial product
                        prob = opt0[1].mul(opt1[1]).mul(opt2[1]).mul(opt3[1]).mul(opt4[1])
                        P_acc.add(prob)
                        if dataP in stab_set:
                            P_good.add(prob)
                    count += 1

print('Enumerated patterns:', count)
print('P_acc coefficients:', P_acc.c)
print('P_good coefficients:', P_good.c)

from fractions import Fraction

# Precomputation reused: E_final, label_list, single_probs, etc.

# Build list of all patterns with at most 1 non-identity error
single_acc = 0
single_good = 0
single_total = 0
for k in range(5):
    for idx,(a,b) in enumerate(label_list):
        if a=='I' and b=='I':
            continue
        # pattern: gate k uses (a,b), others identity
        Ps = ['IIIII']*5
        Ps[k] = E_final[k][idx]
        Pfin = Ps[0]
        for j in range(1,5):
            # multiply
            Q=['I']*5
            for i in range(5):
                Q[i] = mult_table[(Pfin[i], Ps[j][i])]
            Pfin=''.join(Q)
        anc_ok = (Pfin[4] in ('I','Z'))
        dataP = Pfin[:4]
        code_ok = commute_with_checks(dataP)
        if anc_ok and code_ok:
            single_acc += 1
            if dataP in stab_set:
                single_good += 1
        single_total += 1

print('Single-error cases per gate accepted:', single_acc, 'out of', single_total)
print('Single-error accepted and good (should equal accepted):', single_good)

from fractions import Fraction
from decimal import Decimal, getcontext
getcontext().prec=30

acc = [Fraction(1,1), Fraction(-68,15), Fraction(704,75), Fraction(-32768,3375), Fraction(253952,50625), Fraction(-262144,253125)]
good = [Fraction(1,1), Fraction(-68,15), Fraction(656,75), Fraction(-29312,3375), Fraction(223232,50625), Fraction(-229376,253125)]

def poly_eval(coeffs, x):
    s = Decimal(0)
    xp=Decimal(1)
    for c in coeffs:
        s += Decimal(c.numerator)/Decimal(c.denominator) * xp
        xp *= Decimal(x)
    return s

for p in [0.0, 1e-3, 1e-2, 0.05, 0.1]:
    Pa = poly_eval(acc, p)
    Pg = poly_eval(good, p)
    F = Pg/Pa if Pa!=0 else Decimal('NaN')
    print(p, 'P_acc', Pa, 'P_good', Pg, 'F', F)