from qiskit import QuantumCircuit,  Aer 
from qiskit.extensions import UnitaryGate,Initialize
from qiskit.quantum_info import Statevector 
from qiskit.tools.visualization import plot_bloch_multivector  
import numpy as np 
import sys 
from scipy.stats import unitary_group 
import matplotlib.pyplot as plt 
from qiskit import IBMQ 
%matplotlib inline 


sys.path.append("..")
from Modules.changed_SPEA import global_max_SPEA

# IBMQ.load_account()
# provider = IBMQ.get_provider(hub='??replace_with_your_hub_name??')

sim = Aer.get_backend('qasm_simulator')

U1 = UnitaryGate(data=np.array([[0,1],
                               [1,0]]))

spe = global_max_SPEA(U1,resolution= 30,error = 4,max_iters=12)

result = spe.get_eigen_pair(progress = False,backend=sim)

 result 

t = []
for resolution in range(10,50,5):
    spe = global_max_SPEA(U1,resolution= resolution,error = 4,max_iters=10)    
    res = spe.get_eigen_pair(backend = sim)
    theta = res['theta']
    t.append(theta)

plt.figure(figsize = (8,9))
plt.title("Eigenphase estimation for X gate, modified algo",fontsize = 15)
plt.grid()
plt.plot(list(range(10,50,5)),t,marker = 'o',color='g',label = 'Estimates',alpha=0.7)
plt.plot([10,50],[0.5,0.5],color='black',label = "True")
plt.plot([10,50],[0,0],color='black')
plt.plot([10,50],[1,1],color='black')

plt.legend()
plt.xlabel("Resolution of theta ")
plt.ylabel("Estimates")
plt.savefig("Plots/SPE_PLOT1_global_max.jpg",dpi = 200)

u2 = UnitaryGate(data=np.array([[1,0],
                               [0, np.exp(2*np.pi*1j*(1/5))]]))

t = []
for resolution in range(10,50,5):
    spe = global_max_SPEA(u2,resolution= resolution,error = 4,max_iters=10)    
    res = spe.get_eigen_pair(backend = sim)
    theta = res['theta']
    t.append(theta)

plt.figure(figsize = (9,9))
plt.title("Eigenphase estimation for theta = 1/5, 0, 1, modified algo",fontsize = 15)
plt.grid()
plt.plot(list(range(10,50,5)),t,marker = 'o',color='g',label = 'Estimates',alpha=0.7)
plt.plot([10,50],[0.2,0.2],color='black',label = "True")
plt.plot([10,50],[0,0],color='black')
plt.plot([10,50],[1,1],color='black')

plt.legend()
plt.xlabel("Resolution of theta ")
plt.ylabel("Estimates")
plt.savefig("Plots/SPE_PLOT2_global_max.jpg",dpi = 200)

unitary_group.rvs()

u_rand = unitary_group.rvs(2)
print("Random Unitary :",u_rand)

eigen_phases, eigen_vectors = np.linalg.eig(u_rand)
print("Eigen states of Unitary :",eigen_vectors)
eigen_phases = np.angle(eigen_phases)
ep = []
for k in eigen_phases:
    if k < 0:
        ep.append(k + 2*np.pi)
    else:
        ep.append(k)
eigen_phases = ep
print("Eigen phases of unitary :",eigen_phases)

ev1 , ev2 = eigen_vectors[0] , eigen_vectors[1]
ev1 = ev1 / np.linalg.norm(ev1)
ev2 = ev2 / np.linalg.norm(ev2)

qc1 = QuantumCircuit(1)
qc2 = QuantumCircuit(1)
sv1 = Initialize(ev1)
sv2 = Initialize(ev2)
qc1 = qc1.compose(sv1,qubits = [0])
qc2 = qc2.compose(sv2,qubits = [0])
qc1.draw('mpl')

sv = Statevector.from_instruction(qc1)
plot_bloch_multivector(sv)

sv = Statevector.from_instruction(qc2)
plot_bloch_multivector(sv)

spea = global_max_SPEA(resolution = 50, max_iters = 10, unitary = u_rand,error = 4)
result = spea.get_eigen_pair(progress = False, randomize = True)
print("Result of our estimation :",result)

res_state = result['state']
res_state = res_state/ np.linalg.norm(res_state)
qc_res = QuantumCircuit(1)
sv_init = Initialize(res_state)
qc_res = qc_res.compose(sv_init, qubits = [0])
qc_res.draw('mpl')

statevector_res = Statevector.from_instruction(qc_res)
plot_bloch_multivector(statevector_res)

def generate_random_estimation(experiments=5,resolution = 40):
    sim = Aer.get_backend('qasm_simulator')
    best_costs, errors_in_phases = [] , [] 
    for exp in range(experiments):
        u_rand = unitary_group.rvs(2)

        # generate the phases and vectors 
        eigen_phases, eigen_vectors = np.linalg.eig(u_rand)
        eigen_phases = np.angle(eigen_phases)
        ep = []
        # doing this as phase maybe be negative 
        for k in eigen_phases:
            if k < 0:
                ep.append((k + 2*np.pi)/(2*np.pi))
            else:
                ep.append(k/(2*np.pi))
        eigen_phases = ep
        
        ev1 , ev2 = eigen_vectors[0] , eigen_vectors[1]
        ev1 = ev1 / np.linalg.norm(ev1)
        ev2 = ev2 / np.linalg.norm(ev2)
        
        # generate their corresponding init statevectors 
        sv1 = Initialize(ev1)
        sv2 = Initialize(ev2)
        print("Eigenvectors",ev1,ev2)
        print("Eigenphases",eigen_phases)
        
        # run the algorithm 
        spea = global_max_SPEA(resolution = resolution, max_iters = 10, unitary = u_rand,error = 4)
        result = spea.get_eigen_pair(backend = sim, progress = False, randomize = True)
        
        # get the results 
        res_state = result['state']
        res_theta = result['theta']
        sv_res = Initialize(res_state)
        print("Result",result)
        print("Phase returned(0-> 2*pi) :",res_theta)
        
        # get the dot products 
        d1 = np.linalg.norm(np.dot(ev1, res_state.conjugate().T))**2 
        d2 = np.linalg.norm(np.dot(ev2, res_state.conjugate().T))**2 
    
        # make a bloch sphere
        qc = QuantumCircuit(2)
        qc = qc.compose(sv_res,qubits = [0])
        if d1 > d2:
            print(" Best overlap :",d1)
            # it is closer to the first
            qc = qc.compose(sv1,qubits = [1])
            best_costs.append(result['cost'])
            errors_in_phases.append(abs(res_theta - eigen_phases[0]))
        else:
            # it is closer to the second
            print(" Best overlap :",d2)
            qc = qc.compose(sv2,qubits = [1])
            best_costs.append(result['cost'])
            errors_in_phases.append(abs(res_theta - eigen_phases[1]))
        print("Bloch Sphere for the states...")
        s = Statevector.from_instruction(qc)
        display(plot_bloch_multivector(s))
    
    plt.title("Experiments for Random Matrices",fontsize= 16)
    plt.xlabel("Experiment Number")
    plt.ylabel("Metric value")
    plt.plot([i for i in range(experiments)], best_costs, label = 'Best Costs', alpha = 0.5, color = 'g',marker='o')
    plt.plot([i for i in range(experiments)], errors_in_phases, label = 'Corresponding Error in Phase', alpha = 0.5, color = 'b',marker='s') 
    plt.legend()
    plt.grid()
    return 

generate_random_estimation()



