from qiskit import QuantumRegister, ClassicalRegister, QuantumCircuit
import numpy as np

###########################################
#  Amplitude damping channel on IBMQ_VIGO #
###########################################

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

# Quantum circuit
ad = QuantumCircuit(q)

# Amplitude damping channel acting on system qubit
## Qubit identification
system = 1
environment = 2

# Define rotation angle
theta = 0.0

# Construct circuit
ad.x(q[system])
# Notice the extra factor 2 due to how qiskit defines the unitary angles
ad.cu(theta, 0.0, 0.0, 0.0, q[system], q[environment])
ad.cx(q[environment], q[system])

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

def c1(R,t):
    """Returns the coherence factor in the amplitude damping channel
    
    Args:
        R (float): value of R = \gamma_0/\lambda
        t (float): value of the time variable
    
    Returns:
        A float number
    """

def amplitude_damping_channel(q, c, sys, env, R, t):
    """Returns a QuantumCircuit implementing the amplitude damping channel on the system qubit
    
    Args:
        q (QuantumRegister): the register to use for the circuit
        c (ClassicalRegister): the register to use for the measurement of the system qubit
        sys (int): index for the system qubit
        env (int): index for the environment qubit
        R (float): value of R = \gamma_0/\lambda
        t (float): value of the time variable
    
    Returns:
        A QuantumCircuit object
    """

#######################################
#  Amplitude damping channel          #
#   with non-M. witness on IBMQ_VIGO  #
#######################################

# Quantum and classical register
q = QuantumRegister(5, name='q')
c = ClassicalRegister(2, name='c')

# Quantum circuit
ad = QuantumCircuit(q, c)

# Amplitude damping channel acting on system qubit
# with non-Markovianity witness
## Qubit identification
system = 1
environment = 2
ancilla = 3

# Define rotation angle
theta = 0.0

# Construct circuit
## Bell state between system and ancilla
ad.h(q[system])
ad.cx(q[system], q[ancilla])

## Channel acting on system qubit
ad.cu3(theta, 0.0, 0.0, q[system], q[environment])
ad.cx(q[environment], q[system])

## Local measurement for the witness
### Choose observable
observable = 'YY'
### Change to the corresponding basis
if observable == 'XX':
    ad.h(q[system])
    ad.h(q[ancilla])
elif observable == 'YY':
    ad.sdg(q[system])
    ad.h(q[system])
    ad.sdg(q[ancilla])
    ad.h(q[ancilla])
### Measure
ad.measure(q[system], c[0])
ad.measure(q[ancilla], c[1])

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

def amplitude_damping_channel_witness(q, c, sys, env, anc, observable, R, t):
    """Returns a QuantumCircuit implementing the amplitude damping channel on the system qubit with non-Markovianity witness
    
    Args:
        q (QuantumRegister): the register to use for the circuit
        c (ClassicalRegister): the register to use for the measurement of the system and ancilla qubits
        sys (int): index for the system qubit
        env (int): index for the environment qubit
        anc (int): index for the ancillary qubit
        observable (str): the observable to be measured. Possible values "XX", "YY", "ZZ"
        R (float): value of R = \gamma_0/\lambda
        t (float): value of the time variable
    
    Returns:
        A QuantumCircuit object
    """