iterations = [i for i in range(5, 20, 2)]
iterations

from qiskit import IBMQ
import random 
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("../..")
import os
from scipy.stats import unitary_group
import matplotlib.pyplot as plt
%matplotlib inline

# IBMQ.load_account()

# provider = IBMQ.get_provider(hub='ibm-q-education')
# santiago = provider.get_backend('ibmq_santiago')
# casablanca = provider.get_backend('ibmq_casablanca')
# bogota = provider.get_backend('ibmq_bogota')
sim = Aer.get_backend('qasm_simulator')
# athens = provider.get_backend('ibmq_athens')

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

def generate_plots(unitary_size, costs, errors, overlaps, algorithm):

    colors = ['red', 'brown', 'cyan', 'green',
              'grey', 'blue', 'purple', 'black', 'orange']
    c1, c2, c3 = random.sample(colors, 3)

    # plot
    os.makedirs("Experiment_2/"+str(unitary_size) +
                "_qubit(random)/", exist_ok=True)
    # plot 1
    fig = plt.figure(figsize=(13, 6))
    ax1 = fig.add_subplot(1, 2, 1)
    ax1.set_title(str(unitary_size)+" qubit "+algorithm +
                  " Cost v/s Max iters", fontsize=16)
    ax1.set_xlabel("Maximum iterations ", fontsize=15)
    ax1.set_ylabel("Metrics Returned for unitary ", fontsize=15)
    ax1.plot(iterations, costs, label='Costs of Unitary',
             marker='o', color=c1, alpha=0.7)
    ax1.plot(iterations, overlaps, label='Average overlap from nearest eigenvector',
             marker='s', color=c2, alpha=0.6)
    ax1.legend(loc='best')
    ax1.grid()
    # plot 2
    ax2 = fig.add_subplot(1, 2, 2)
    ax2.set_title(str(unitary_size)+" qubit "+algorithm +
                  " % error v/s Max iters", fontsize=16)
    ax2.set_xlabel("Maximum iterations ", fontsize=15)
    ax2.set_ylabel("% error for nearest eigenvalue", fontsize=15)
    ax2.plot(iterations, errors, label='Average error from nearest eigenvalue',
             marker='o', color=c3, alpha=0.6)
    ax2.legend(loc='best')
    ax2.grid()
    # save axure
    fig.savefig("Experiment_2/"+str(unitary_size)+"_qubit(random)/" +
                algorithm+" Algorithm (alternate).JPG", dpi=200)

def get_results(eig_vals, eig_vect, bases, basis_indices, unitary, algorithm, experiments):
    '''Return the results of running the algorithm for this particular unitary matrix'''
    costs_g = []
    errors_g = []
    max_overlaps_g = []
    # find how the cost converges with increasing iterations
    for iters in iterations:
        costs = []
        errors = []
        overlaps = []
        
        i = 0 
        # run the experiments ...
        while len(costs) < experiments:
            if algorithm == 'original':
                spea = SPEA(unitary, resolution=30, error=3, max_iters=iters)
            else:
                spea = global_max_SPEA(
                    unitary, resolution=30, error=3, max_iters=iters)

            result = spea.get_eigen_pair(
                progress=False, backend=sim, basis=bases[i], basis_ind=basis_indices[i],
                algo='alternate', randomize=False,shots = 256)

            if result['cost'] < 0.65:
                continue
                
            i+=1 
            
            # find the costs
            costs.append(result['cost'])
            theta = result['theta']
            res_state = result['state']

            # find the abs difference in this theta with the closest eigenvalue
            # and append that to the errors ...
            min_error = 1e5
            for e in eig_vals:
                error = abs(e - theta)
                if error < min_error:
                    min_error = error
                    perc_error = ((error)/e)*100
            errors.append(perc_error)

            # find overlaps
            max_overlap = -1
            for k in eig_vect:
                dot = np.linalg.norm(np.dot(k, res_state.conjugate().T))**2
                max_overlap = max(max_overlap, dot)
            overlaps.append(max_overlap)

        print("Result with", iters, " iterations :")
        print("AVG. COST :", np.average(costs),
              "AVG. ERROR :", np.average(errors))
        # append the average result of your algorithm ...
        costs_g.append(np.average(costs))
        errors_g.append(np.average(errors))
        max_overlaps_g.append(np.average(overlaps))

    return costs_g, errors_g, max_overlaps_g

unit_2 = unitary_group.rvs(4)
unit_2

eig_vals2, eig_vect2 = np.linalg.eig(unit_2)
eig_vals2 = np.angle(eig_vals2)
e = []
for k in eig_vals2:
    if k < 0:
        v = (k + 2*np.pi)/(2*np.pi)
    else:
        v = (k)/(2*np.pi)
    e.append(v)
eig_vals2 = np.array(e)
print("Eigenstates :", eig_vect2)
print("Eigenvalues :", eig_vals2)

bases2 , basis_indices2 = [], []
for _ in range(4):
    sample = unitary_group.rvs(4)
    basis = []
    for k in sample:
        basis.append(np.array(k, dtype=complex))
#     print("2 qubit basis :", basis2)
    ind = np.random.choice(range(4))
#     print(ind)
    bases2.append(basis)
    basis_indices2.append(ind)
# print("Basis set :",bases2)
print("Basis indices :",basis_indices2)

costs_2qubit_b, errors_2qubit_b, max_overlaps_2qubit_b = get_results(eig_vals2,
                                                                     eig_vect2, bases2, basis_indices2,
                                                                     unit_2, "original", 4)

generate_plots(2, costs_2qubit_b, errors_2qubit_b,
               max_overlaps_2qubit_b, "Original")

costs_2qubit_c, errors_2qubit_c, max_overlaps_2qubit_c = get_results(
    eig_vals2, eig_vect2, bases2, basis_indices2, unit_2, "modified", 4)

generate_plots(2, costs_2qubit_c, errors_2qubit_c,
               max_overlaps_2qubit_c, "Modified")

unit_3 = unitary_group.rvs(8)
unit_3

eig_vals3, eig_vect3 = np.linalg.eig(unit_3)
eig_vals3 = np.angle(eig_vals3)
e = []
for k in eig_vals3:
    if k < 0:
        v = (k + 2*np.pi)/(2*np.pi)
    else:
        v = (k)/(2*np.pi)
    e.append(v)
eig_vals3 = np.array(e)
print("Eigenstates :", eig_vect3)
print("Eigenvalues :", eig_vals3)

bases3 , basis_indices3 = [], []
for _ in range(4):
    sample = unitary_group.rvs(8)
    basis = []
    for k in sample:
        basis.append(np.array(k, dtype=complex))
#     print("3 qubit basis :", basis3)
    ind = np.random.choice(range(8))
#     print(ind)
    bases3.append(basis)
    basis_indices3.append(ind)
# print("Basis set :",bases3)
print("Basis indices :",basis_indices3)

costs_3qubit_b, errors_3qubit_b, max_overlaps_3qubit_b = get_results(
    eig_vals3, eig_vect3,bases3, basis_indices3, unit_3, "original", 4)

generate_plots(3, costs_3qubit_b, errors_3qubit_b,
               max_overlaps_3qubit_b, "Original")

costs_3qubit_c, errors_3qubit_c, max_overlaps_3qubit_c = get_results(
    eig_vals3, eig_vect3, bases3, basis_indices3, unit_3, "modified", 4)
generate_plots(3, costs_3qubit_c, errors_3qubit_c,
               max_overlaps_3qubit_c, "Modified")

unit_4 = unitary_group.rvs(16)
# unit_4

eig_vals4, eig_vect4 = np.linalg.eig(unit_4)
eig_vals4 = np.angle(eig_vals4)
e = []
for k in eig_vals4:
    if k < 0:
        v = (k + 2*np.pi)/(2*np.pi)
    else:
        v = (k)/(2*np.pi)
    e.append(v)
eig_vals4 = np.array(e)
print("Eigenstates :", eig_vect4)
print("Eigenvalues :", eig_vals4)

bases4 , basis_indices4 = [], []
for _ in range(4):
    sample = unitary_group.rvs(16)
    basis = []
    for k in sample:
        basis.append(np.array(k, dtype=complex))
    ind = np.random.choice(range(16))
    bases4.append(basis)
    basis_indices4.append(ind)
print("Basis indices :",basis_indices4)

costs_4qubit_b, errors_eig_4qubit_b, max_overlaps_4qubit_b = get_results(
    eig_vals4, eig_vect4, bases4, basis_indices4, unit_4, 'original', 4)
generate_plots(4, costs_4qubit_b, 
               errors_eig_4qubit_b, max_overlaps_4qubit_b, "Original")

costs_4qubit_c, errors_eig_4qubit_c,  max_overlaps_4qubit_c = get_results(
    eig_vals4, eig_vect4, bases4, basis_indices4, unit_4, 'modified', 4)
generate_plots(4, costs_4qubit_c, 
               errors_eig_4qubit_c, max_overlaps_4qubit_c, "Modified")

