from qiskit import QuantumRegister, QuantumCircuit
from qiskit.circuit import Parameter

#############################
#  Pauli channel on IBMQX2  #
#############################

# Quantum register
q = QuantumRegister(3, name="q")

# Quantum circuit
pauli = QuantumCircuit(q)

# Pauli channel acting on q_2
## Qubit identification
system = 0
a_0 = 1
a_1 = 2

# ## Define rotation angles
theta_1 = Parameter('θ1')
theta_2 = Parameter('θ2')
theta_3 = Parameter('θ3')

## Construct circuit
pauli.ry(theta_1, q[a_0])
pauli.cx(q[a_0], q[a_1])
pauli.ry(theta_2, q[a_0])
pauli.ry(theta_3, q[a_1])
pauli.cx(q[a_0], q[system])
pauli.cy(q[a_1], q[system])

# Draw circuit
pauli.draw(output='mpl')

def pauli_channel(q, p, system, pauli_ancillae):
    """
        Apply the Pauli channel to system with probabilities p
        
    Args:
    q (QuantumRegister): the quantum register for the circuit
    system (int): index of the system qubit
    pauli_ancillae (list): list of indices of the ancillary qubits
    p (list): list of probabilities [p_1, p_2, p_3] for the Pauli channel
    
    Returns:
    A QuantumCircuit implementing the Pauli channel
    """
    
    # Write code

# Suggested imports...
from qiskit.tools.qi.qi import entropy, partial_trace

def conditional_entropy(state, qubit_a, qubit_b):
    """Conditional entropy S(A|B) = S(AB) - S(B)
    
    Args:
        state: a vector or density operator
        qubit_a: 0-based index of the qubit A
        qubit_b: 0-based index of the qubit B
        
    Returns:
        int: the conditional entropy
    """
    # Write code here

def extractable_work(state, system_qubit, memory_qubit):
    """Extractable work from a two-qubit state
    =
    Cfr. Eq. (4) Bylicka et al., Sci. Rep. 6, 27989 (2016)
    
    Args:
        qubit_a: 0-based index of the system qubit S
        qubit_b: 0-based index of the memory qubit M
    """
    # Write code here

