# We'll enumerate all possible two-qubit Pauli errors inserted after each of 5 CNOT gates
# and compute the exact analytic (symbolic-by-counting) fidelity under post-selection.

from collections import defaultdict

# Qubits: 0,1,2,3 are data; 4 is ancilla
nq = 5

# Represent Paulis by (X_bits, Z_bits), each a list of 0/1 of length nq. Phase ignored.

def pauli_id():
    return [0]*nq, [0]*nq

# Multiply two Paulis modulo phase: XOR X and Z bits

def pauli_mul(p, q):
    X1, Z1 = p
    X2, Z2 = q
    return [x1 ^ x2 for x1, x2 in zip(X1, X2)], [z1 ^ z2 for z1, z2 in zip(Z1, Z2)]

# Conjugate Pauli by CNOT(control -> target)

def conj_cnot(p, c, t):
    X, Z = p
    # Copy to avoid in-place hazards
    Xc, Xt = X[c], X[t]
    Zc, Zt = Z[c], Z[t]
    # Mapping (Aaronson-Gottesman):
    # X_c -> X_c, X_t -> X_t ^ X_c
    # Z_c -> Z_c ^ Z_t, Z_t -> Z_t
    X[t] = Xt ^ Xc
    Z[c] = Zc ^ Zt
    # X[c] and Z[t] unchanged otherwise
    return X, Z

# Build single-qubit Pauli given symbol: 'I','X','Y','Z' on a given qubit

def single_pauli(sym, q):
    X = [0]*nq
    Z = [0]*nq
    if sym == 'I':
        return X, Z
    if sym == 'X':
        X[q] = 1
        return X, Z
    if sym == 'Z':
        Z[q] = 1
        return X, Z
    if sym == 'Y':
        X[q] = 1
        Z[q] = 1
        return X, Z
    raise ValueError('bad sym')

# Build two-qubit Pauli on qubits (qa, qb) as tensor product Pa*Pb

def twoq_pauli(sym_a, qa, sym_b, qb):
    Xa, Za = single_pauli(sym_a, qa)
    Xb, Zb = single_pauli(sym_b, qb)
    return pauli_mul((Xa, Za), (Xb, Zb))

# CNOT gates sequence indices and their qubits in time order after H_1
# Gate list: C1: (1->2), C2: (1->0), C3: (2->3), C4: (3->4), C5: (0->4)
cnots = [(1,2), (1,0), (2,3), (3,4), (0,4)]
ng = len(cnots)

# Build the list of 16 two-qubit Paulis for a given pair of qubits: identity + 15 non-identity.
syms = ['I','X','Y','Z']

def twoq_pauli_list(qa, qb):
    plist = []
    for a in syms:
        for b in syms:
            if a == 'I' and b == 'I':
                continue
            plist.append(twoq_pauli(a, qa, b, qb))
    # identity as index 0 separately
    return plist

# Precompute for each gate i, the 16 possibilities (0=identity, 1..15 non-identity) propagated to the end
precomp = []
for i, (c, t) in enumerate(cnots):
    # Build 16 entries
    arr = []
    # idx 0: identity
    arr.append(pauli_id())
    # 15 non-identity two-qubit Paulis on qubits c,t
    nonids = twoq_pauli_list(c, t)
    for P in nonids:
        X, Z = [row[:] for row in P]
        # Conjugate by subsequent CNOTs j>i
        for j in range(i+1, ng):
            cj, tj = cnots[j]
            X, Z = conj_cnot((X, Z), cj, tj)
        arr.append((X, Z))
    precomp.append(arr)

# Helper: compute commutation of Pauli P with Q given by (XQ,ZQ): returns 0 if commute, 1 if anticommute
# Using symplectic product over all qubits.

def commutes(P, Q):
    XP, ZP = P
    XQ, ZQ = Q
    s = 0
    for j in range(nq):
        s ^= (XP[j] & ZQ[j]) ^ (ZP[j] & XQ[j])
    return (s == 0)

# Build stabilizers and logical operators on data only, promoted to 5 qubits (ancilla idle)
# Sx = XXXX on qubits 0..3
Sx = ([1,1,1,1,0], [0,0,0,0,0])
Sz = ([0,0,0,0,0], [1,1,1,1,0])
# Z_A = ZZII (qubits 0,1)
ZA = ([0,0,0,0,0], [1,1,0,0,0])
# Z_B = ZIZI (qubits 0,2)
ZB = ([0,0,0,0,0], [1,0,1,0,0])

# Acceptance and correctness counters by number of faulty gates k
c_acc = [0]*(ng+1)
c_cor = [0]*(ng+1)

# Enumerate over all 16 choices per gate: 0=identity, 1..15 non-identity
# To speed up, we'll do 5 nested loops but with integers; Should be fine.

choices = [range(16) for _ in range(ng)]

def build_total(idx_tuple):
    # XOR the contributions
    Xtot = [0]*nq
    Ztot = [0]*nq
    for i, idx in enumerate(idx_tuple):
        Xp, Zp = precomp[i][idx]
        Xtot = [a ^ b for a,b in zip(Xtot, Xp)]
        Ztot = [a ^ b for a,b in zip(Ztot, Zp)]
    return (Xtot, Ztot)

# We'll iterate and count.

import itertools

for idx_tuple in itertools.product(*choices):
    # Count faulty gates
    k = sum(1 for idx in idx_tuple if idx != 0)
    # Quickly skip if too many? No, count all.
    Xtot, Ztot = build_total(idx_tuple)
    Ptot = (Xtot, Ztot)
    # Ancilla acceptance: X bit on qubit 4 must be 0 (no X/Y on ancilla)
    if Xtot[4] != 0:
        continue
    # Data acceptance: commute with Sx and Sz
    if not commutes(Ptot, Sx):
        continue
    if not commutes(Ptot, Sz):
        continue
    # Accepted
    c_acc[k] += 1
    # Correct if commute with ZA and ZB as well
    if commutes(Ptot, ZA) and commutes(Ptot, ZB):
        c_cor[k] += 1

c_acc, c_cor

print('Computing...')
from collections import defaultdict
nq=5

def pauli_id():
    return [0]*nq, [0]*nq

def pauli_mul(p, q):
    X1, Z1 = p
    X2, Z2 = q
    return [x1 ^ x2 for x1, x2 in zip(X1, X2)], [z1 ^ z2 for z1, z2 in zip(Z1, Z2)]

def conj_cnot(p, c, t):
    X, Z = p
    Xc, Xt = X[c], X[t]
    Zc, Zt = Z[c], Z[t]
    X[t] = Xt ^ Xc
    Z[c] = Zc ^ Zt
    return X, Z

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

def single_pauli(sym, q):
    X = [0]*nq
    Z = [0]*nq
    if sym == 'I':
        return X, Z
    if sym == 'X':
        X[q] = 1
    elif sym == 'Z':
        Z[q] = 1
    elif sym == 'Y':
        X[q] = 1
        Z[q] = 1
    else:
        raise ValueError
    return X, Z

def twoq_pauli(sa, qa, sb, qb):
    Xa, Za = single_pauli(sa, qa)
    Xb, Zb = single_pauli(sb, qb)
    return pauli_mul((Xa, Za), (Xb, Zb))

cnots = [(1,2), (1,0), (2,3), (3,4), (0,4)]
ng = len(cnots)


def twoq_pauli_list(qa, qb):
    plist = []
    for a in syms:
        for b in syms:
            if a=='I' and b=='I':
                continue
            plist.append(twoq_pauli(a, qa, b, qb))
    return plist

precomp = []
for i,(c,t) in enumerate(cnots):
    arr = []
    arr.append(pauli_id())
    for P in twoq_pauli_list(c,t):
        X,Z = [row[:] for row in P]
        for j in range(i+1, ng):
            cj, tj = cnots[j]
            X,Z = conj_cnot((X,Z), cj, tj)
        arr.append((X,Z))
    precomp.append(arr)

# Build operators
Sx = ([1,1,1,1,0], [0,0,0,0,0])
Sz = ([0,0,0,0,0], [1,1,1,1,0])
ZA = ([0,0,0,0,0], [1,1,0,0,0])
ZB = ([0,0,0,0,0], [1,0,1,0,0])


def commutes(P,Q):
    XP,ZP = P
    XQ,ZQ = Q
    s=0
    for j in range(nq):
        s ^= (XP[j] & ZQ[j]) ^ (ZP[j] & XQ[j])
    return (s==0)

import itertools

c_acc = [0]*(ng+1)
c_cor = [0]*(ng+1)

for idx_tuple in itertools.product(*([range(16)]*ng)):
    k = sum(1 for idx in idx_tuple if idx!=0)
    # Build total
    Xtot=[0]*nq
    Ztot=[0]*nq
    for i,idx in enumerate(idx_tuple):
        Xp,Zp = precomp[i][idx]
        Xtot = [a^b for a,b in zip(Xtot,Xp)]
        Ztot = [a^b for a,b in zip(Ztot,Zp)]
    # acceptance checks
    if Xtot[4]!=0:
        continue
    Ptot=(Xtot,Ztot)
    if not commutes(Ptot,Sx):
        continue
    if not commutes(Ptot,Sz):
        continue
    c_acc[k]+=1
    if commutes(Ptot,ZA) and commutes(Ptot,ZB):
        c_cor[k]+=1

print('c_acc:', c_acc)
print('c_cor:', c_cor)

from fractions import Fraction

c_acc = [1, 7, 282, 4222, 31637, 94923]
c_cor = [1, 7, 138, 1198, 7397, 24027]

# Build polynomials N(p), D(p) as lists of coefficients for p^0..p^5
# N(p) = sum_k c_cor[k]*(p/15)^k*(1-p)^{5-k}
# (1-p)^{5-k} = sum_{j=0}^{5-k} binom(5-k,j) * (1)^{5-k-j} * (-p)^j
# So the coefficient for p^{k+j} accumulates c_cor[k]*(1/15)^k * binom(5-k,j) * (-1)^j

import math

def poly_from_counts(counts):
    coeffs = [Fraction(0,1) for _ in range(6)]
    for k, ck in enumerate(counts):
        for j in range(0, 6-k):
            deg = k + j
            coeffs[deg] += Fraction(ck, 15**k) * Fraction(math.comb(5-k, j)) * ((-1)**j)
    return coeffs

N = poly_from_counts(c_cor)
D = poly_from_counts(c_acc)

print('N coefficients (p^0..p^5):')
print([str(c) for c in N])
print('D coefficients (p^0..p^5):')
print([str(c) for c in D])

# Small-p expansion of F(p) up to p^2 or p^3 by series division
from decimal import Decimal, getcontext
getcontext().prec = 50

# Let's also compute numeric coefficients rational to show maybe exact

# Series division to degree 3
F = [Fraction(0,1) for _ in range(6)]
# We'll compute F such that N = D * F modulo p^5; standard power series division
for m in range(6):
    # compute F[m]
    s = N[m]
    for j in range(m):
        s -= D[m-j] * F[j]
    F[m] = s / D[0]

print('Series F coeffs (p^0..p^5):', [str(f) for f in F])

from decimal import Decimal, getcontext
getcontext().prec = 50

c_acc = [1, 7, 282, 4222, 31637, 94923]
c_cor = [1, 7, 138, 1198, 7397, 24027]

def Acc(p):
    from math import comb
    # direct formula using counts
    res_acc = Decimal(0)
    res_cor = Decimal(0)
    for k in range(6):
        w = (Decimal(p)/Decimal(15))**k * (Decimal(1)-Decimal(p))**(5-k)
        res_acc += Decimal(c_acc[k]) * w
        res_cor += Decimal(c_cor[k]) * w
    return res_acc, res_cor, (res_cor/res_acc if res_acc!=0 else Decimal(0))

for pv in [Decimal('0.001'), Decimal('0.01'), Decimal('0.05'), Decimal('0.1')]:
    A, C, F = Acc(pv)
    print('p=', pv, ' Acc=', A, ' Cor=', C, ' F=', F)