import numpy as np

# Pauli matrices
I = np.eye(2)
X = np.array([[0,1],[1,0]], dtype=complex)
Y = np.array([[0,-1j],[1j,0]], dtype=complex)
Z = np.array([[1,0],[0,-1]], dtype=complex)
H = (1/np.sqrt(2))*np.array([[1,1],[1,-1]], dtype=complex)

# Tensor product utility
from functools import reduce

def kron(*ops):
    return reduce(np.kron, ops)

# 4-qubit basis: qubits ordered (0,1,2,3)

# Build CNOT_ij unitary on 4 qubits

def CNOT_ij(i,j):
    # Build projector form: |0><0|_i ⊗ I + |1><1|_i ⊗ X_j
    P0 = np.array([[1,0],[0,0]], dtype=complex)
    P1 = np.array([[0,0],[0,1]], dtype=complex)
    ops0 = [I,I,I,I]
    ops1 = [I,I,I,I]
    ops0[i] = P0
    ops1[i] = P1
    # X on j for second term
    ops1[j] = X
    return kron(*ops0) + kron(*ops1)

# Hadamard on qubit q (others identity)

def H_q(q):
    ops = [I,I,I,I]
    ops[q] = H
    return kron(*ops)

# Two-qubit Pauli labels and matrices
Psingle = [I, X, Y, Z]
Pnames = ['I','X','Y','Z']

# Build two-qubit Pauli on a specified pair (a,b) with order qubit index 0..3

def two_qubit_pauli_on_pair(a,b, pa, pb):
    ops = [I,I,I,I]
    ops[a] = Psingle[pa]
    ops[b] = Psingle[pb]
    return kron(*ops)

# Circuit unitaries (ideal): U = CNOT_03 * H0 * CNOT_21 * H2
U = CNOT_ij(0,3) @ H_q(0) @ CNOT_ij(2,1) @ H_q(2)

# Initial |0000>
zero = np.array([[1],[0]], dtype=complex)
psi0 = kron(zero, zero, zero, zero)
psi_ideal = U @ psi0
rho_ideal = psi_ideal @ psi_ideal.conj().T

# Build the noisy rho(p) symbolically as sum over errors with weights
# But we'll just compute the coefficient of fidelity as a function of p by summing contributions

# Enumerate two-qubit paulis on pair (2,1) and (0,3)
non_id_pairs = []
for a in range(4):
    for b in range(4):
        name = Pnames[a]+Pnames[b]
        non_id_pairs.append((a,b,name))

# Build the list of 16 two-qubit paulies and classify which are identity
pairs = non_id_pairs

# Let's compute the probability weight that maps to fidelity 1: we check for each pair of errors whether
# E03 * U * E21 * ... acting on |0000> equals psi_ideal up to global phase. We can test overlap squared ~1.

# Build E for pair (2,1) and (0,3)

import math

def fidelity_weight():
    # We'll compute count of (pauli on 03, on 21) that leave the state invariant up to phase.
    survivors = []
    for a03 in range(4):
        for b03 in range(4):
            E03 = two_qubit_pauli_on_pair(0,3, a03, b03)
            for a21 in range(4):
                for b21 in range(4):
                    E21 = two_qubit_pauli_on_pair(2,1, a21, b21)
                    # Build final state with this error model: errors applied after CNOTs
                    # That is: first H2, then CNOT21, then error E21, then H0, then CNOT03, then error E03
                    psi = E03 @ CNOT_ij(0,3) @ H_q(0) @ E21 @ CNOT_ij(2,1) @ H_q(2) @ psi0
                    # Overlap magnitude
                    overlap = (psi_ideal.conj().T @ psi)[0,0]
                    if np.allclose(abs(overlap), 1.0, atol=1e-12):
                        survivors.append(((a03,b03),(a21,b21)))
    return survivors

survivors = fidelity_weight()
len_survivors, survivors[:10]

print(len(survivors))
print(survivors[:16])

import numpy as np
I = np.eye(2); X = np.array([[0,1],[1,0]], complex); Y = np.array([[0,-1j],[1j,0]], complex); Z = np.array([[1,0],[0,-1]], complex)
H = (1/np.sqrt(2))*np.array([[1,1],[1,-1]], complex)
from functools import reduce

def kron(*ops):
    return reduce(np.kron, ops)

def CNOT_ij(i,j):
    P0 = np.array([[1,0],[0,0]], complex)
    P1 = np.array([[0,0],[0,1]], complex)
    ops0 = [I,I,I,I]
    ops1 = [I,I,I,I]
    ops0[i] = P0
    ops1[i] = P1
    ops1[j] = X
    return kron(*ops0) + kron(*ops1)

def H_q(q):
    ops = [I,I,I,I]
    ops[q] = H
    return kron(*ops)

Psingle = [I,X,Y,Z]
U = CNOT_ij(0,3) @ H_q(0) @ CNOT_ij(2,1) @ H_q(2)
zero = np.array([[1],[0]], complex)
psi0 = kron(zero,zero,zero,zero)
psi_ideal = U @ psi0
rho_ideal = psi_ideal @ psi_ideal.conj().T

# Build final density matrix under the channel parameter p

def rho_of_p(p):
    rho = np.zeros((16,16), complex)
    # enumerate E21 and E03 errors and accumulate weights
    for a03 in range(4):
        for b03 in range(4):
            weight03 = (1-p) if (a03==0 and b03==0) else (p/15)
            E03 = kron(Psingle[a03], I, I, Psingle[b03])
            for a21 in range(4):
                for b21 in range(4):
                    weight21 = (1-p) if (a21==0 and b21==0) else (p/15)
                    E21 = kron(I, Psingle[b21], Psingle[a21], I)
                    psi = E03 @ CNOT_ij(0,3) @ H_q(0) @ E21 @ CNOT_ij(2,1) @ H_q(2) @ psi0
                    rho += weight03 * weight21 * (psi @ psi.conj().T)
    return rho

for p in [0.0, 0.01, 0.1, 0.2, 0.5, 1.0]:
    rho = rho_of_p(p)
    F = (psi_ideal.conj().T @ rho @ psi_ideal)[0,0].real
    print(p, F, 'vs formula', (1 - 4*p/5)**2)