from qiskit import QuantumRegister, QuantumCircuit
import numpy as np

#######################
#  ZZ pump on IBMQX2  #
#######################

# Quantum register
q = QuantumRegister(5, name='q')

# Quantum circuit  
zz = QuantumCircuit(q)

# ZZ pump acting on system qubits
## Qubit identification
system = [2, 1]
a_zz = 0

## Define pump efficiency 
## and corresponding rotation
p = 0.5
theta = 2 * np.arcsin(np.sqrt(p))

## Construct circuit
### Map information to ancilla
zz.cx(q[system[0]], q[system[1]])
zz.x(q[a_zz])
zz.cx(q[system[1]], q[a_zz])
    
### Conditional rotation
zz.cu3(theta, 0.0, 0.0, q[a_zz], q[system[1]])
    
### Inverse mapping
zz.cx(q[system[1]], q[a_zz])
zz.cx(q[system[0]], q[system[1]])

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

#######################
#  XX pump on IBMQX2  #
#######################

# Quantum register
q = QuantumRegister(5, name='q')

# Quantum circuit  
xx = QuantumCircuit(q)

# XX pump acting on system qubits
## Qubit identification
system = [2, 1]
a_xx = 4

## Define pump efficiency 
## and corresponding rotation
p = 0.5
theta = 2 * np.arcsin(np.sqrt(p))

## Construct circuit
### Map information to ancilla
xx.cx(q[system[0]], q[system[1]])
xx.h(q[system[0]])
xx.x(q[a_xx])
xx.cx(q[system[0]], q[a_xx])
    
### Conditional rotation
xx.cu3(theta, 0.0, 0.0, q[a_xx], q[system[0]])
    
### Inverse mapping
xx.cx(q[system[0]], q[a_xx])
xx.h(q[system[0]])
xx.cx(q[system[0]], q[system[1]])

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

###########################
#  ZZ-XX pumps on IBMQX2  #
###########################

# Quantum register
q = QuantumRegister(5, name='q')

# Quantum circuit  
zz_xx = QuantumCircuit(q)

# ZZ and XX pumps acting on system qubits
## Qubit identification
system = [2, 1]
a_zz = 0
a_xx = 4

## Define pump efficiency 
## and corresponding rotation
p = 0.5
theta = 2 * np.arcsin(np.sqrt(p))

## Construct circuit
## ZZ pump
### Map information to ancilla
zz_xx.cx(q[system[0]], q[system[1]])
zz_xx.x(q[a_zz])
zz_xx.cx(q[system[1]], q[a_zz])
    
### Conditional rotation
zz_xx.cu3(theta, 0.0, 0.0, q[a_zz], q[system[1]])
    
### Inverse mapping
zz_xx.cx(q[system[1]], q[a_zz])
#zz_xx.cx(q[system[0]], q[system[1]])

## XX pump
### Map information to ancilla
#zz_xx.cx(q[system[0]], q[system[1]])
zz_xx.h(q[system[0]])
zz_xx.x(q[a_xx])
zz_xx.cx(q[system[0]], q[a_xx])
    
### Conditional rotation
zz_xx.cu3(theta, 0.0, 0.0, q[a_xx], q[system[0]])
    
### Inverse mapping
zz_xx.cx(q[system[0]], q[a_xx])
zz_xx.h(q[system[0]])
zz_xx.cx(q[system[0]], q[system[1]])

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

def zz_pump(q, c, p, system, ancilla):
    """Returns a QuantumCircuit implementing the ZZ pump channel on the system qubits
    
    Args:
        q (QuantumRegister): the register to use for the circuit
        c (ClassicalRegister): the register to use for the measurement of the system qubits
        p (float): the efficiency for the channel, between 0 and 1
        system (list): list of indices for the system qubits
        ancilla (int): index for the ancillary qubit
    
    Returns:
        A QuantumCircuit object
    """

def xx_pump(q, c, p, system, ancilla):
    """Returns a QuantumCircuit implementing the XX pump channel on the system qubits
    
    Args:
        q (QuantumRegister): the register to use for the circuit
        c (ClassicalRegister): the register to use for the measurement of the system qubits
        p (float): the efficiency for the channel, between 0 and 1
        system (list): list of indices for the system qubits
        ancilla (int): index for the ancillary qubit
    
    Returns:
        A QuantumCircuit object
    """

def zz_xx_pump(q, c, p, system, ancillae):
    """Returns a QuantumCircuit implementing the composition channel on the system qubits
    
    Args:
        q (QuantumRegister): the register to use for the circuit
        c (ClassicalRegister): the register to use for the measurement of the system qubits
        p (float): the efficiency for both channels, between 0 and 1
        system (list): list of indices for the system qubits
        ancillae (list): list of indices for the ancillary qubits
    
    Returns:
        A QuantumCircuit object
    """