from itertools import product
from collections import defaultdict

# We will work with 5 qubits labeled 0..4 (4 is ancilla)
# Represent Pauli as symplectic (x|z) bitstrings of length 5

nq = 5

def pauli_to_symp(pauli_str):
    # pauli_str is a string of length nq over 'IXYZ'
    x = [0]*nq
    z = [0]*nq
    for i,ch in enumerate(pauli_str):
        if ch == 'I':
            x[i]=0; z[i]=0
        elif ch == 'X':
            x[i]=1; z[i]=0
        elif ch == 'Z':
            x[i]=0; z[i]=1
        elif ch == 'Y':
            x[i]=1; z[i]=1
        else:
            raise ValueError
    return tuple(x), tuple(z)

def symp_to_pauli(x,z):
    s = []
    for i in range(nq):
        if x[i]==0 and z[i]==0: s.append('I')
        elif x[i]==1 and z[i]==0: s.append('X')
        elif x[i]==0 and z[i]==1: s.append('Z')
        elif x[i]==1 and z[i]==1: s.append('Y')
        else: raise ValueError
    return ''.join(s)


def multiply_symp(x1,z1,x2,z2):
    # Pauli product modulo phase: (x1|z1) XOR (x2|z2)
    x = tuple((a^b) for a,b in zip(x1,x2))
    z = tuple((a^b) for a,b in zip(z1,z2))
    return x,z

# Conjugation by CNOT(c,t) in symplectic form

def conj_cnot(x,z,c,t):
    # x_t <- x_t XOR x_c
    # z_c <- z_c XOR z_t
    x = list(x); z = list(z)
    x[t] ^= x[c]
    z[c] ^= z[t]
    return tuple(x), tuple(z)

# Define gates in time order (right-to-left in original):
# 0: CNOT 1->2, 1: CNOT 1->0, 2: CNOT 2->3, 3: CNOT 3->4, 4: CNOT 0->4
cnot_list = [(1,2),(1,0),(2,3),(3,4),(0,4)]

# Build propagation maps: for each gate index g, for each 2-qubit Pauli on its support (16 incl. II),
# produce the final 5-qubit Pauli after propagating through subsequent gates (g+1..end)

# helper: enumerate 2-qubit Paulis in fixed order (I,X,Y,Z) for each qubit
single = ['I','X','Y','Z']

# Map: gate index -> list of 16 tuples (x,z)
prop_maps = []

for g,(c,t) in enumerate(cnot_list):
    # Build base 5-qubit Pauli for each choice on (c,t)
    finals = []
    for P_c in single:
        for P_t in single:
            # Create 5-qubit string
            s = ['I']*nq
            s[c] = P_c
            s[t] = P_t
            x,z = pauli_to_symp(''.join(s))
            # propagate forward through subsequent CNOTs
            for (c2,t2) in cnot_list[g+1:]:
                x,z = conj_cnot(x,z,c2,t2)
            finals.append((x,z))
    prop_maps.append(finals)

# Stabilizers on data qubits 0..3
# S_X = XXXX I, S_Z = ZZZZ I
SX_x, SX_z = pauli_to_symp('XXXXI')
SZ_x, SZ_z = pauli_to_symp('ZZZZI')

# Logical operators chosen:
XA_x, XA_z = pauli_to_symp('XIXII'.replace(' ','').strip())
XB_x, XB_z = pauli_to_symp('XXIII'.replace(' ','').strip())
ZA_x, ZA_z = pauli_to_symp('ZZIII'.replace(' ','').strip())
ZB_x, ZB_z = pauli_to_symp('ZIZII'.replace(' ','').strip())
# But above includes ancilla; ensure length 5: add I on qubit 4 automatically from pauli_to_symp with 'I' at end
# We'll just ensure strings are length 5 with ancilla I
XA_x, XA_z = pauli_to_symp('XIXII')
XB_x, XB_z = pauli_to_symp('XXIII')
ZA_x, ZA_z = pauli_to_symp('ZZIII')
ZB_x, ZB_z = pauli_to_symp('ZIZII')

# Commutation test: returns 0 if commute, 1 if anticommute

def symp_comm(x1,z1,x2,z2):
    s = 0
    for a,b in zip(x1,z2):
        s ^= (a & b)
    for a,b in zip(z1,x2):
        s ^= (a & b)
    return s  # 0 commute, 1 anticommute

# Check acceptance and logical class
# acceptance conditions:
# - ancilla measurement in Z returns 0 <=> no X/Y on qubit 4 just before measurement: x4 == 0
# - code detects no error <=> E_data commutes with S_X and S_Z (i.e., symp_comm(E,SX)==0 and symp_comm(E,SZ)==0)

# For classification: compute (cZ_A,cX_A) and (cZ_B,cX_B)

from functools import lru_cache

# Precompute the 16 2-qubit options per gate as indices and mark IDs: 0..15
# We'll interpret option 0 as II, 1: IX, 2:IY, 3:IZ, 4:XI, ... standard product order

# Enumerate all 16^5 assignments: do a dynamic programming to accumulate counts by s (#non-identity),
# and classes: accepted? and within accepted: z-only logical or x-containing.

N_total = 16**5

N_acc_by_s = [0]*6
N_good_by_s = [0]*6  # good = accepted and logical action contains no X on A or B (i.e., z-only on both)
N_flip_by_s = [0]*6  # accepted but contains X on A or B (including Y)

# To accelerate, we can do nested loops or use recursion

# We'll enumerate using nested loops explicitly (5 levels), as 16^5 ~ 1.048M; manageable

# Provide functions to combine symplectic quickly

def add_symp(acc_x, acc_z, add_x, add_z):
    return tuple(a^b for a,b in zip(acc_x, add_x)), tuple(a^b for a,b in zip(acc_z, add_z))

# identity symp
I_x, I_z = pauli_to_symp('IIIII')

count = 0
for i0 in range(16):
    x0,z0 = prop_maps[0][i0]
    s0 = 0 if (i0==0) else 1
    for i1 in range(16):
        x01,z01 = add_symp(x0,z0, *prop_maps[1][i1])
        s1 = s0 + (0 if (i1==0) else 1)
        for i2 in range(16):
            x012,z012 = add_symp(x01,z01, *prop_maps[2][i2])
            s2 = s1 + (0 if (i2==0) else 1)
            for i3 in range(16):
                x0123,z0123 = add_symp(x012,z012, *prop_maps[3][i3])
                s3 = s2 + (0 if (i3==0) else 1)
                for i4 in range(16):
                    xfin,zfin = add_symp(x0123,z0123, *prop_maps[4][i4])
                    s = s3 + (0 if (i4==0) else 1)
                    # acceptance tests
                    # ancilla bitflip? reject if xfin[4]==1
                    if xfin[4]==1:
                        continue
                    # code detection: commute with SX and SZ
                    if symp_comm(xfin,zfin,SX_x,SX_z) or symp_comm(xfin,zfin,SZ_x,SZ_z):
                        continue
                    N_acc_by_s[s] += 1
                    # classify logical action on A,B using commutation with logical Z and X
                    cZA = symp_comm(xfin,zfin, ZA_x, ZA_z)
                    cXA = symp_comm(xfin,zfin, XA_x, XA_z)
                    cZB = symp_comm(xfin,zfin, ZB_x, ZB_z)
                    cXB = symp_comm(xfin,zfin, XB_x, XB_z)
                    # if contains any X (or Y) component on A or B -> flip class
                    if (cZA==1) or (cZB==1):
                        # contains X on A and/or B
                        N_flip_by_s[s] += 1
                    else:
                        # no X on either; either I or Z-only
                        N_good_by_s[s] += 1




# sanity check: accepted events split into good + flip
assert all(N_acc_by_s[k] == N_good_by_s[k] + N_flip_by_s[k] for k in range(6))

print('N_acc_by_s:', N_acc_by_s)
print('N_good_by_s:', N_good_by_s)
print('N_flip_by_s:', N_flip_by_s)

# Now construct polynomials: P_acc(p) = sum_s N_acc[s] * (1-p)^(5-s) * (p/15)^s
# P_good(p) similar, P_flip = P_acc - P_good

from fractions import Fraction

# We can express numerator polynomials in p of degree up to 5 by expanding (1-p)^(5-s) * p^s
# We will compute coefficients of P_acc(p) as rational Fractions with denominator 15^s factored separately? But we can include that as well.

# We'll return P_acc(p) as sum_{k=0..5} c_k p^k, rational coefficients.

def poly_add(a,b):
    # add polynomials (list of Fractions)
    n = max(len(a),len(b))
    res = [Fraction(0) for _ in range(n)]
    for i in range(len(a)): res[i]+=a[i]
    for i in range(len(b)): res[i]+=b[i]
    return res

def poly_scale(a, s):
    return [coeff*s for coeff in a]

# Precompute binomial expansions of (1-p)^(t)
from math import comb

binoms = {}
for t in range(6):
    # (1-p)^t = sum_{j=0..t} C(t,j) * (-p)^j
    coeffs = [Fraction(0) for _ in range(t+1)]
    for j in range(t+1):
        coeffs[j] = Fraction(((-1)**j)*comb(t,j), 1)
    binoms[t] = coeffs

# function to multiply polynomials; degrees small

def poly_mul(a,b):
    res = [Fraction(0) for _ in range(len(a)+len(b)-1)]
    for i,ai in enumerate(a):
        for j,bj in enumerate(b):
            res[i+j] += ai*bj
    return res

# Build P_acc, P_good as polynomials in p
P_acc = [Fraction(0)]
P_good = [Fraction(0)]
P_flip = [Fraction(0)]

for s in range(6):
    # weight factor: (p/15)^s * (1-p)^(5-s)
    # = (1/15^s) * p^s * sum_{j=0..5-s} C(5-s,j) (-p)^j
    a = [Fraction(0) for _ in range(s+1)]
    a[s] = Fraction(1, 15**s)  # p^s / 15^s
    b = binoms[5-s]  # (1-p)^(5-s)
    term = poly_mul(a,b)  # polynomial in p
    P_acc = poly_add(P_acc, poly_scale(term, Fraction(N_acc_by_s[s], 1)))
    P_good = poly_add(P_good, poly_scale(term, Fraction(N_good_by_s[s], 1)))
    P_flip = poly_add(P_flip, poly_scale(term, Fraction(N_flip_by_s[s], 1)))

# Simplify: reduce by total number of assignments? Not needed, we already sum actual probabilities.

# Let's display polynomials; Convert to simplified Fractions

def poly_to_str(P):
    terms = []
    for k,ck in enumerate(P):
        if ck == 0: continue
        terms.append(f'({ck})*p^{k}')
    return ' + '.join(terms) if terms else '0'

print('P_acc(p) =', poly_to_str(P_acc))
print('P_good(p) =', poly_to_str(P_good))
print('P_flip(p) =', poly_to_str(P_flip))

# Finally, we want the conditional fidelity F(p) = P_good(p) / P_acc(p)
# For presentation, we can also expand F(p) as series around p=0 up to p^3

import decimal

# Evaluate at a few p values to sanity check monotonicity
for p in [0, 1e-4, 1e-3, 1e-2, 5e-2, 0.1]:
    # compute numerical evaluation
    def eval_poly(P,p):
        s = 0.0
        pw = 1.0
        for ck in P:
            s += float(ck)*pw
            pw *= p
        return s
    acc = eval_poly(P_acc,p)
    good = eval_poly(P_good,p)
    flip = eval_poly(P_flip,p)
    F = good/acc if acc>0 else 1.0
    print(f'p={p}: P_acc={acc:.10e}, P_good={good:.10e}, P_flip={flip:.10e}, Fidelity={F:.10e}')

# Also compute series expansion to O(p^2) or O(p^3)
from math import isclose

# We'll compute rational series expansion by dividing polynomials formally to some order

max_order = 5

# F(p) = (G0 + G1 p + G2 p^2 + ...)/(A0 + A1 p + A2 p^2 + ...)
# Compute series coefficients by solving (A * F) = G.

A = P_acc
G = P_good

# Ensure lists length
A += [Fraction(0)]*(max_order+1 - len(A))
G += [Fraction(0)]*(max_order+1 - len(G))

Fseries = [Fraction(0) for _ in range(max_order+1)]
# F0 = G0/A0
Fseries[0] = G[0]/A[0]
for n in range(1, max_order+1):
    # sum_{k=0..n} A[k] F[n-k] = G[n]
    s = Fraction(0)
    for k in range(1,n+1):
        s += A[k]*Fseries[n-k]
    Fseries[n] = (G[n] - s)/A[0]

print('Fidelity series coefficients up to p^5:')
for n,cn in enumerate(Fseries):
    print(n, cn)