# Importing standard Qiskit libraries
from qiskit import QuantumCircuit, execute, Aer, IBMQ
from qiskit.visualization import *
from qiskit.quantum_info import state_fidelity
# Numpy for numeric functions
import numpy as np
# Magic function to render plots in the notebook after the cell executing the plot command
%matplotlib inline

def run_on_statevector_simulator(quantum_circuit, decimals=6):
    """Takes a circuit, and runs it on the state vector simulator backend."""
    statevector_simulator = Aer.get_backend('statevector_simulator')
    job = execute(quantum_circuit, backend=statevector_simulator)
    result = job.result()    
    statevector = result.get_statevector(quantum_circuit, decimals=decimals)
    return statevector

# Generate a quantum circuit for two qubits 
qc = QuantumCircuit(2)
# Add the gates which generate |psi_0>
qc.x(1)
qc.h(0)
qc.h(1)
# Draw the quantum circuit
qc.draw(output='mpl')

# Run it on the state vector simulator
vec = run_on_statevector_simulator(qc)

# Visualize the resulting state vector
plot_bloch_multivector(vec)

# Add a CNOT gate to the circuit
qc.cnot(0,1)
# Draw the circuit
qc.draw(output='mpl')

# Run it on the state vector simulator
vec = run_on_statevector_simulator(qc)

# Visualize the resulting state vector
plot_bloch_multivector(vec)

# Generate a quantum circuit for two qubits which generates |psi_0>
qc2 = QuantumCircuit(2)
# Add the gates
qc2.x(1)
qc2.h(0)
qc2.h(1)
# Draw the quantum circuit
qc2.draw(output='mpl')

# Run it on the state vector simulator for various angles of Rx
# Number of steps
nsteps = 10
for i in range(nsteps):
    # We add a controlled R_x gate with different angle to our base circuit
    qc3 = QuantumCircuit(2)    
    qc3.compose(qc2, inplace=True)
    qc3.crx(i*4*np.pi/nsteps,0,1)
    # Run the resulting circuit on the state vector simulator
    vec = run_on_statevector_simulator(qc3)
    # Visualize the state vector
    h = plot_bloch_multivector(vec)    
    # Save the image to disk
    h.savefig(str(i)+'.png')    