# !pip install qiskit

import math
import matplotlib.pyplot as plt

import numpy as np
import qiskit as q
from qiskit.visualization import plot_bloch_vector, plot_state_city

from google.colab import drive
drive.mount('/content/drive')

path = '/content/drive/My Drive/Colab Notebooks/Quantum_Compute/'


coords = [1, math.pi/2, math.pi]
plot_bloch_vector(coords, coord_type='spherical', figsize=(6, 6), )

## check density matrix from state vector 

initial_state1 = [0.j + 1./np.sqrt(2), 0, 0, 1/np.sqrt(2) + 0.j] 
initial_state2 = [0.j + 1./2., 0, 0, 0.j + np.sqrt(3)/2.]

m1 = q.quantum_info.Statevector(initial_state1)
print ('check the first statevector: ', m1)
m2 = q.quantum_info.Statevector(initial_state2)
print ('check the first statevector: ', m2)

rho_1 = q.quantum_info.DensityMatrix(m1, dims=4)
print ('check the density matrix corresponding to 1st statevector: ', '\n', rho_1)
rho_2 = q.quantum_info.DensityMatrix(m2, dims=4)
print ('check the density matrix corresponding to 2nd statevector: ', '\n', rho_2)



rho_1.draw('latex', prefix='\\rho_{1} = ')

plot_state_city(rho_1.data, title=r'Density Matrix for StateVector: $\frac{1}{\sqrt{2}}\left(|0\rangle + |1\rangle\right)$', 
                color=['magenta', 'purple'], alpha=0.7, figsize=(11, 7), )


rho_2.draw('latex', prefix='\\rho_{2} = ')

plot_state_city(rho_2.data, title=r'Density Matrix for StateVector: $\left(\frac{1}{2}|0\rangle + \frac{\sqrt{3}}{2}|1\rangle\right)$', 
                color=['magenta', 'purple'], alpha=0.7, figsize=(11, 7), )

rho_H_matrix = np.array([[9/20, 0, 0, np.sqrt(3)/20 + 8/20], [0, 0, 0, 0], [0, 0, 0, 0], [np.sqrt(3)/20 + 8/20, 0, 0, 11/20]])
rho_H = q.quantum_info.DensityMatrix(rho_H_matrix)
rho_H.draw('latex', prefix='\\rho_{mix} = ')

final_density_matrix = 0.8*rho_1 + 0.2*rho_2
check_rho_H = q.quantum_info.DensityMatrix(final_density_matrix)
check_rho_H.draw('latex', prefix='\\rho_{mix} = ')

### check the first property

rho_H_matrix_square = np.matmul(rho_H_matrix, rho_H_matrix)
rho_H_square = q.quantum_info.DensityMatrix(rho_H_matrix_square)

rho_H_matrix_square_Trace = np.trace(rho_H_matrix_square)

print ('Trace of the square of the density matrix: ', rho_H_matrix_square_Trace)

rho_H_square.draw('latex', prefix='\\rho_{mix}^2 = ')

### check the second property

rho_H_matrix_eigenvals = np.linalg.eigvals(rho_H_matrix)
print ('check the eigenvalues of the density matrix', rho_H_matrix_eigenvals[0], rho_H_matrix_eigenvals[1])
print ('check sum of the eigenvalues of the density matrix: ', np.sum(rho_H_matrix_eigenvals))

plot_state_city(check_rho_H.data, 
                title=r'Density Matrix for Mixed State: $0.8\times\frac{1}{\sqrt{2}}\left(|0\rangle + |1\rangle\right) + 0.2\times \left(\frac{1}{2}|0\rangle + \frac{\sqrt{3}}{2}|1\rangle\right)$', 
                color=['crimson', 'purple'], alpha=0.4, figsize=(11, 7), )

reduced_density_matrix = np.array([[9/20, np.sqrt(3)/20 + 8/20], [np.sqrt(3)/20 + 8/20, 11/20]])
red_rho_H = q.quantum_info.DensityMatrix(reduced_density_matrix)
red_rho_H.draw('latex', prefix='\\rho_{mix} = ')



from qiskit.visualization import plot_bloch_multivector
plot_bloch_multivector(state=red_rho_H.data, figsize=(6, 6), )

