import sys
sys.path.append('../')

from expressibility import expressibility
from entanglement import Ent

import numpy as np
import kaleidoscope.qiskit
from kaleidoscope import bloch_sphere
import matplotlib.pyplot as plt

from qiskit import QuantumCircuit
from qiskit.quantum_info import DensityMatrix

def getDensityMatrix(circuit):
    return DensityMatrix(circuit).data

import warnings
warnings.filterwarnings('ignore')

def sampleI(layer=1, qubits=1):
    circuit = QuantumCircuit(1)
    circuit.i(0)
    
    return circuit

def sampleHU1(layer=1, qubits=1):
    circuit = QuantumCircuit(1)
    
    circuit.h(0)
    circuit.u1(np.random.uniform(0, np.pi * 2), 0)
    
    return circuit

def sampleHU2(layer=1, qubits=1):
    circuit = QuantumCircuit(1)
    
    circuit.h(0)
    circuit.u2(np.random.uniform(0, np.pi * 2), np.random.uniform(0, np.pi * 2), 0)
    
    return circuit

def sampleHU3(layer=1, qubits=1):
    circuit = QuantumCircuit(1)
    
    circuit.h(0)
    circuit.u3(np.random.uniform(0, np.pi * 2), np.random.uniform(0, np.pi * 2), np.random.uniform(0, np.pi * 2), 0)
    
    return circuit

X = np.matrix([
    [0, 1], [1, 0]
])
Y = np.matrix([
    [0, -1j],
    [1j, 0]
])
Z = np.matrix([
    [1, 0],
    [0, -1]
])

def point(sample):
    rou = getDensityMatrix(sample())
    
    return [
        np.trace(rou @ X).real,
        np.trace(rou @ Y).real,
        np.trace(rou @ Z).real
    ]

color = "#FF9933"
bins = 100

sampleI().draw(output='mpl')

pointsI = [ point(sampleI) for i in range(100) ]

bloch_sphere(points=pointsI, points_color=color)

h_i, p_i = expressibility(1, sampleI, bins=bins, return_detail=True)

fig, ax = plt.subplots()
plt.bar([ i/100 for i in range(bins) ], h_i, color="#9999FF", label="Haar", width=0.01)
plt.bar([ i/100 for i in range(bins) ], p_i, color="#FF9933AA", label="VQC", width=0.01)
plt.ylabel("Probability", fontsize=15)
plt.xlabel("Fidelity", fontsize=15)
plt.title("Expr")
ax.legend(loc="upper center")

sampleHU1().draw(output='mpl')

pointsHU1 = [ point(sampleHU1) for i in range(2000) ]

bloch_sphere(points=pointsHU1, points_color=color)

h_hu1, p_hu1 = expressibility(1, sampleHU1, bins=bins, return_detail=True)

fig, ax = plt.subplots()
plt.bar([ i/100 for i in range(bins) ], h_hu1, color="#9999FF", label="Haar", width=0.01)
plt.bar([ i/100 for i in range(bins) ], p_hu1, color="#FF9933AA", label="VQC", width=0.01)
plt.ylabel("Probability", fontsize=15)
plt.xlabel("Fidelity", fontsize=15)
plt.title("Expr")
ax.legend(loc="upper center")

sampleHU2().draw(output='mpl')

pointsHU2 = [ point(sampleHU2) for i in range(2000) ]

bloch_sphere(points=pointsHU2, points_color=color)

h_hu2, p_hu2 = expressibility(1, sampleHU2, bins=bins, return_detail=True)

fig, ax = plt.subplots()
plt.bar([ i/100 for i in range(bins) ], h_hu2, color="#9999FF", label="Haar", width=0.01)
plt.bar([ i/100 for i in range(bins) ], p_hu2, color="#FF9933AA", label="VQC", width=0.01)
plt.ylabel("Probability", fontsize=15)
plt.xlabel("Fidelity", fontsize=15)
plt.title("Expr")
ax.legend(loc="upper center")

sampleHU3().draw(output='mpl')

pointsHU3 = [ point(sampleHU3) for i in range(2000) ]

bloch_sphere(points=pointsHU3, points_color=color)

h_hu3, p_hu3 = expressibility(1, sampleHU3, bins=bins, return_detail=True)

fig, ax = plt.subplots()
plt.bar([ i/100 for i in range(bins) ], h_hu3, color="#9999FF", label="Haar", width=0.01)
plt.bar([ i/100 for i in range(bins) ], p_hu3, color="#FF9933AA", label="VQC", width=0.01)
plt.ylabel("Probability", fontsize=15)
plt.xlabel("Fidelity", fontsize=15)
plt.title("Expr")
ax.legend(loc="upper center")

