from qiskit import QuantumCircuit
from qiskit.quantum_info import Statevector, Pauli, SparsePauliOp
from math import isclose

# Encoding circuits for |->, |+> and |0>
def encode_minus(qc,i):
    n = 5
    qc.h(range(i,i+n))
    qc.cz(0+i, 1+i)
    qc.cz(2+i, 3+i)
    qc.cz(1+i, 2+i)
    qc.cz(3+i, 4+i)
    qc.cz(0+i, 4+i)

def encode_plus(qc,i):
    n = 5
    encode_minus(qc,i)
    qc.z(range(i,i+n))

def apply_H(qc,i):
    qc.h(range(i,i+5))
    qc.swap(3+i,4+i)
    qc.swap(2+i,3+i)
    qc.swap(1+i,3+i)

def apply_HSdg(qc,i):
    """ Logical HSdg  (--[Sdg]--[H]--) 
    Action: X --> Y, Z --> X """
    qc.sdg(0+i)
    qc.h(0+i)
    qc.h(2+i)
    qc.s(2+i)
    qc.h(3+i)
    qc.s(3+i)
    qc.swap(1+i,2+i)
    qc.swap(0+i,2+i)

def encode_zero(qc,i):
    encode_plus(qc,i)
    apply_H(qc,i)

def apply_CNOT(qc,i):
    n = 5
    # Local Clifford rotation 
    qc.h(0+i)
    qc.s(0+i)
    qc.y(2+i)
    qc.h(4+i)
    qc.s(4+i)
    qc.h(0+n+i)
    qc.s(0+n+i)
    qc.y(2+n+i)
    qc.h(4+n+i)
    qc.s(4+n+i)

    # Logical CNOT
    qc.cx(0+i,5+i)
    qc.cx(2+i,7+i)
    qc.cx(4+i,9+i)
    qc.cx(0+i,9+i)
    qc.cx(2+i,5+i)
    qc.cx(4+i,7+i)
    qc.cx(0+i,7+i)
    qc.cx(2+i,9+i)
    qc.cx(4+i,5+i)

    # Logical Clifford rotation reversed 
    qc.sdg(0+i)
    qc.h(0+i)
    qc.y(2+i)
    qc.sdg(4+i)
    qc.h(4+i)
    qc.sdg(0+n+i)
    qc.h(0+n+i)
    qc.y(2+n+i)
    qc.sdg(4+n+i)
    qc.h(4+n+i)

def encode_bell_pair(qc,i):
    encode_plus(qc,i)
    encode_plus(qc,i+5)
    apply_H(qc,i+5)
    apply_CNOT(qc,i)

def verify_plus(qc,i,a):
    qc.cx(0+i,a)
    qc.h(a)
    qc.cx(a,2+i)
    qc.h(a)
    qc.cx(4+i,a)
    # qc.measure(a,a)

def verify_zero(qc,i,a):
    qc.sdg(0+i)
    qc.cx(0+i,a)
    qc.s(0+i)
    qc.cx(2+i,a)
    qc.sdg(4+i)
    qc.cx(4+i,a)
    qc.s(4+i)
    # qc.measure(a,a)

def knill_qec(qc,i,b):
    """i: first qubit index of data qubits
       b: first qubit index of bell pair"""
    for ind in range(5):
        qc.cx(i+ind,b+ind)
    for ind in range(5):
        qc.h(i+ind)
        # qc.measure(i+ind,i+ind)
        # qc.measure(b+ind,b+ind)