import numpy as np

from qiskit import QuantumCircuit, assemble, Aer, execute, BasicAer
from qiskit.visualization import *
from qiskit.quantum_info import Statevector, random_statevector
from qiskit.tools.jupyter import *
from qiskit.extensions import Initialize

bell_0 = QuantumCircuit(2,2)
bell_0.h(0)
bell_0.cx(0,1)

bell_0.draw('latex')

#we define the basis
sv = Statevector.from_label("00")

#we evolve this initial state through our circuit
sv_bell_0 = sv.evolve(bell_0)
sv_bell_0.draw('latex')

bell_1 = QuantumCircuit(2,2)
bell_1.h(0)
bell_1.cx(0,1)
bell_1.z(1)

#mpl = matplotlib
#latex_source also possible, but really hard to see what is going on
bell_1.draw('mpl')

sv_bell_1 = sv.evolve(bell_1)
sv_bell_1.draw('latex_source')

bell_2 = QuantumCircuit(2,2)
bell_2.h(0)
bell_2.x(1)
bell_2.cx(0,1)

bell_2.draw('text') 

sv_bell_2 = sv.evolve(bell_2)
sv_bell_2.draw('text') #text by default

bell_3 = QuantumCircuit(2,2)
bell_3.h(0)
bell_3.x(1)
bell_3.cx(0,1)
bell_3.z(1)

bell_3.draw() #text by default

sv_bell_3 = sv.evolve(bell_3)
sv_bell_3.draw()

ghz = QuantumCircuit(3)
ghz.h(0)
ghz.cx(0,1)
ghz.cx(1,2)

ghz.draw('mpl')

### drawing the statevector
sv_ghz= Statevector.from_label("000")
state_ghz = sv_ghz.evolve(ghz)
state_ghz.draw('latex')

plot_state_qsphere(bell_1)

plot_state_hinton(bell_1)

plot_state_paulivec(bell_3)

plot_state_city(bell_2) 

plot_bloch_multivector(bell_2)

# It is not a separable state!!

ex_sep_state = QuantumCircuit(2)
ex_sep_state.h(0)
ex_sep_state.h(1)

ex_sv = Statevector.from_label("00")
ex_state = ex_sv.evolve(ex_sep_state)
plot_bloch_multivector(ex_state)

plot_bloch_vector([0.7,0.5,1]) #here for [x,y,z]; x=Tr[Xρ] and similar for y and z

###Part1: circuit part
qc = QuantumCircuit(3)
qc.h(0)
qc.cx(0,1)
qc.cx(1,2)

#it saves the current simulator quantum state as a statevector.
qc.save_statevector()  

###Parts 2 and 3 using Aer_simulator
qobj = assemble(qc)     
backend1 = Aer.get_backend('aer_simulator') 

###Parts 4 and 5
result = backend1.run(qobj).result()
sv_qc0 = result.get_statevector()
sv_qc0.draw('latex')

sv_qc0.draw('qsphere')

######## PART2: statevetor_simulator

###Part1: circuit part

qc1 = QuantumCircuit(3)
qc1.h(0)
qc1.cx(0,1)
qc1.cx(1,2)

### backend part

backend2 = BasicAer.get_backend('statevector_simulator') 

### result part

# by default shots = 1024
job = execute(qc1, backend2, shots=1024) 
result = job.result()

#finally, we get the statevector from the result
sv_qc1 = result.get_statevector(qc1)

sv_qc1.draw('qsphere')

plot_state_city(sv_qc1)

N = 1/np.sqrt(3)
desired_state = [N,np.sqrt(1-N**2)]

#adding it to a circuit
qc_custom0 = QuantumCircuit(1)
qc_custom0.initialize(desired_state,0) #as simple as this!
qc_custom0.draw('mpl')

meas = QuantumCircuit(1,1)
meas.measure(0,0)

#we add, through the composite function, the measurement to the custom state
qc_custom = meas.compose(qc_custom0, front=True)
qc_custom.draw('mpl')

#To get histograms, we use qasm_simulator. 
#We will see this in the notebook qasm_simulator_and_visualization.ipynb
backend_qasm = BasicAer.get_backend('qasm_simulator')

job0 = execute(qc_custom, backend_qasm, shots=1000)
counts = job0.result().get_counts()
plot_histogram(counts)

# random_statevector of dimension 2
psi = random_statevector(2)

#Initialize qubits in a specific state.
init_gate = Initialize(psi)

# defining a name to the state. By default is |psi\rangle
init_gate.label = "initial state"

qc_random = QuantumCircuit(1)

# random state for the first qibit
qc_random.append(init_gate, [0])

qc_random.draw('mpl')

qc_random.measure_all()
qc_random.draw('mpl')

job1 = execute(qc_random, backend_qasm, shots=1000)
counts1 = job1.result().get_counts()
plot_histogram(counts1)

decomp = qc_custom0.decompose()
decomp.draw('mpl')

