# run this cell if you're executing this notebook in your browser
!pip install qiskit

from IPython.display import clear_output
clear_output()

from qiskit import *
import numpy as np


n = 2 #number of qubits used for the example in this notebook

from qiskit.visualization import array_to_latex


target_state = [ # state vector of the target state
    1/np.sqrt(n),
    0,
    (complex(0,1))/np.sqrt(n),
    0,
    ]



array_to_latex(target_state, prefix = '|\\psi_{ts}> = ' ) # a nifty, clean latex representation

target_state_mathematical = [ 
    [1/np.sqrt(n)],
    [0],
    [(complex(0,1))/np.sqrt(n)],
    [0],
    ]



array_to_latex(target_state_mathematical, prefix = 'Mathematical: |\\psi_{m}> = ' )

qr_init = QuantumRegister(n)

qc_init = QuantumCircuit(qr_init)
qc_init.initialize(target_state) # the initialization

qc_init.draw('mpl')

qr_gate = QuantumRegister(n)

qc_gate = QuantumCircuit(qr_gate)

qc_gate.h(qr_gate[1])
qc_gate.s(qr_gate[1])

qc_gate.draw('mpl')

from qiskit.quantum_info import Statevector

qc_gate_statevector = Statevector(qc_gate) # alternative: qc_gate_statevector = quantum_info.Statevector.from_instruction(qc_gate)

qc_gate_statevector.draw('latex', prefix = '|\\psi_{qc-gate-statevector}> = ')

backend = Aer.get_backend('statevector_simulator') #alternative: from qiskit.providers.aer import StatevectorSimulator | backend = StatevectorSimulator()
job = execute(qc_gate, backend)
qc_gate_state_sim = job.result().get_statevector()

array_to_latex(qc_gate_state_sim, prefix = '|\\psi_{qc-gate-sim}> = ') # Note .draw() isn't used here, because .get_statevector() returns an array

from qiskit.quantum_info import state_fidelity

print('Target state and qc_gate_statevector: ')
display(state_fidelity(target_state, qc_gate_statevector))

print('\nTarget state and qc_gate_state_sim: ')
display(state_fidelity(target_state, qc_gate_state_sim))


qc_gate_statevector.draw('qsphere')

from qiskit.visualization import plot_state_qsphere

plot_state_qsphere(qc_gate_statevector)

qc_gate_statevector.draw('bloch')

from qiskit. visualization import plot_bloch_multivector

plot_bloch_multivector(qc_gate_statevector)

from qiskit.visualization import plot_state_paulivec

plot_state_paulivec(qc_gate_statevector)

from qiskit.quantum_info import DensityMatrix

rho = DensityMatrix(qc_gate) #alternative: rho = quantum_info.DensityMatrix.from_instruction(qc_gate)

rho.draw('latex', prefix = '\\rho = ')

rho_m_matrix = np.array([[1/2,0,complex(0,-1/2),0],[0,0,0,0],[complex(0,1/2),0,1/2,0],[0,0,0,0]])

rho_m = DensityMatrix(rho_m_matrix)
display(rho_m.draw('latex', prefix='\\rho = '))

ket_target_state = np.matrix(target_state).T
bra_target_state  = ket_target_state.conj().T

display(array_to_latex(ket_target_state, prefix = '|\\psi_m> = ' ))
display(array_to_latex(bra_target_state, prefix = '<\\psi_m| = ' ))

target_state_matrix = np.matmul(ket_target_state, bra_target_state)

rho = DensityMatrix(target_state_matrix)
display(rho.draw('latex', prefix='\\rho = '))

from qiskit.visualization import plot_state_city

plot_state_city(rho, color = ['teal', 'black']) # colors cuz I felt like prettying it up

rho.draw('hinton')

from qiskit.visualization import plot_state_hinton

plot_state_hinton(rho)

from qiskit.quantum_info import Statevector

init_state = Statevector(qc_init)
print('\n Initial state: ')
display(init_state.draw('latex', prefix = '|\\psi_{ts}> ='))

# create circuit with Hadamards operating on each qubit
qc_hadamards = QuantumCircuit(n) 
qc_hadamards.h(range(n))  

display(qc_hadamards.draw('mpl'))

new_state = init_state.evolve(qc_hadamards)

print('\n Final state:  ')
display(new_state.draw('latex', prefix= '|\\psi_{f}> ='))




qc_combined = qc_hadamards.compose(qc_init, range(n), front = True)

display(qc_combined.draw('mpl'))

qc_combined_state = Statevector(qc_combined)

print('\n State of circuit qc_init + qc_hadamards = ')
display(qc_combined_state.draw('latex'))

from qiskit.quantum_info import Operator

display(rho.draw('latex', prefix='Initial State, \\rho = '))

H = Operator.from_label('H') #fetches the Hadamard
display(array_to_latex(H.data, prefix ='H = '))



HH = H.expand(H) # tensor product of H with itself
display(array_to_latex(HH.data, prefix ='H \\otimes H = '))

rho_HH = rho.evolve(HH)
display(rho_HH.draw('latex', prefix='\\rho_f = H\\rho H^\\dagger = '))


ket_combined_state = np.matrix(qc_combined_state).T
bra_combined_state  = ket_combined_state.conj().T

display(array_to_latex(ket_combined_state, prefix = '|\\psi_f> = ' ))
display(array_to_latex(bra_combined_state, prefix = '<\\psi_f| = ' ))

combined_state_matrix = np.matmul(ket_combined_state, bra_combined_state)

rho_comb = DensityMatrix(combined_state_matrix)
display(rho_comb.draw('latex', prefix='\\rho_f = '))

