from qiskit import QuantumCircuit, execute, transpile, Aer 
from qiskit.extensions import UnitaryGate,Initialize
from qiskit.quantum_info import Statevector 
from qiskit.tools.visualization import plot_bloch_vector
from qiskit.tools.visualization import plot_histogram,plot_bloch_multivector  
import numpy as np 
from time import sleep 
import sys 
sys.path.append("..")
from scipy.stats import unitary_group 
import matplotlib.pyplot as plt 
%matplotlib inline 

from Modules.normal_SPEA import SPEA 
from Modules.changed_SPEA import global_max_SPEA 

def generate_random_estimation(unitaries,SPE_type,resolution = 40,save = True):
    
    backend = Aer.get_backend('qasm_simulator')
    
    best_costs, errors_in_phases = [] , [] 
    for i,u_rand in enumerate(unitaries):
        print("FOR UNITARY ",i+1)
        # 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 
        if SPE_type == 'original':
            spea = SPEA(resolution = resolution, max_iters = 10, unitary = u_rand,error = 4)
        else:
            spea = global_max_SPEA(resolution = resolution, max_iters = 10, unitary = u_rand,error = 4)
        result = spea.get_eigen_pair(backend = backend, 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 :",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))
        
    if SPE_type == 'original':
        plt.title("Experiments over Random Matrices for original SPEA",fontsize= 16)
    else:
        plt.title("Experiments over Random Matrices for modified SPEA",fontsize= 16)
    
    plt.xlabel("Experiment Number")
    plt.ylabel("Metric value")
    plt.plot([i for i in range(len(unitaries))], best_costs, label = 'Best Costs', alpha = 0.5, color = 'r',marker='o',linewidth = 2)
    plt.plot([i for i in range(len(unitaries))], errors_in_phases, label = 'Corresponding Error in Phase',linewidth = 2, alpha = 0.5, color = 'b',marker='s') 
    plt.legend()
    plt.grid()
    if save:
        if SPE_type == 'original':
            plt.savefig("Random Estimation for Original SPE.jpg",dpi = 200)
        else:
            plt.savefig("Random Estimation for Modified SPE.jpg",dpi = 200)



unitaries = []
for i in range(6):
    u = unitary_group.rvs(2)
    unitaries.append(u)
print("Unitaries :",unitaries) 

generate_random_estimation(unitaries,'original',resolution=15)

generate_random_estimation(unitaries,'modified',resolution=15)

