from qiskit import *
import matplotlib.pyplot as plt
import random
from scipy.stats import unitary_group
import os
import sys 
sys.path.append("../..")
import seaborn as sns
import numpy as np

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

unit = unitary_group.rvs(4)
unit

# get the eigenvalues and eigenstates
eig_v, eig_vect = np.linalg.eig(unit)
eig_v = np.angle(eig_v)
eig = []
for k in eig_v:
    if k < 0:
        e = (k+2*np.pi)/(2*np.pi)
    else:
        e = (k)/(2*np.pi)
    eig.append(e)
eig_v = np.array(eig)
print("Eigenvalues :", eig_v)
print("Eigenstates :", eig_vect)

sample = unitary_group.rvs(4)
basis = []
for k in sample:
    basis.append(k)
basis_ind = np.random.choice(range(4))
print("Basis :", basis)
print("Basis index :", basis_ind)

def generate_plot1(actual_E, returned_E, iters):
    colors = ['blue', 'orange', 'red',
              'green', 'brown', 'magenta', 'pink']

    plt.figure(figsize=(9, 7))
    plt.scatter(range(iters), returned_E, marker='o', edgecolors='grey', color=np.random.choice(colors),
                alpha=0.8,)
    for i, k in enumerate(actual_E):
        if i == 0:
            plt.plot([0, iters], [k, k], color='black',
                     linewidth=2, label='Actual Values')
        else:
            plt.plot([0, iters], [k, k], color='black',
                     linewidth=2)
    plt.xlabel("Number of iterations", fontsize=14)
    plt.ylabel("Eigenvalues", fontsize=14)
    plt.title("Scatter plot for returned eigenvalues", fontsize=17)
    plt.legend()
    plt.grid()

def generate_plot2(actual_E, returned_E, size, experiments):
    colors = ['blue', 'orange', 'red',
              'green', 'brown', 'magenta', 'pink']

    fig = plt.figure(figsize=(16, 13))
    for i in range(size):
        ax = fig.add_subplot(size/2, 2, i+1)

        ax.scatter(range(experiments), returned_E[i], marker='o', edgecolors='grey', color=np.random.choice(colors),
                   alpha=0.8,)
        for i, k in enumerate(actual_E):
            if i == 0:
                ax.plot([0, experiments], [k, k], color='black',
                        linewidth=2, label='Actual Values')
            else:
                ax.plot([0, experiments], [k, k], color='black',
                        linewidth=2)
        ax.set_xlabel("Experiment Number", fontsize=14)
        ax.set_ylabel("Eigenvalues", fontsize=14)
        ax.set_title("Scatter plot for returned eigenvalues", fontsize=17)
        ax.legend(loc='best')
        ax.grid()

simulator = Aer.get_backend('qasm_simulator')


spea1 = SPEA(unit,resolution=40,error = 2, max_iters=15)

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 100:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate', basis=basis,
        basis_ind=basis_ind)
    if res['cost'] < 0.75:
        continue
    print(res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 100)

spea2 = global_max_SPEA(unit, resolution=40, error=2, max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret2 = []
while len(eigen_vals_ret2) != 100:
    res = spea2.get_eigen_pair(
        backend=simulator, algo='alternate', basis=basis,
        basis_ind=basis_ind)
    if res['cost'] < 0.75:
        continue
    print(res)
    eigen_vals_ret2.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret2, 100)

unit = unitary_group.rvs(4)
# get the eigenvalues and eigenstates
eig_v, eig_vect = np.linalg.eig(unit)
eig_v = np.angle(eig_v)
eig = []
for k in eig_v:
    if k < 0:
        e = (k+2*np.pi)/(2*np.pi)
    else:
        e = (k)/(2*np.pi)
    eig.append(e)
eig_v = np.array(eig)
print(eig_v)


y = np.array([[0,-1j],
            [1j,0]])
yy= np.kron(y,y)
basis = []
for k in yy:
    basis.append(k)

basis  

eig_ret = []

spea1 = SPEA(unit, resolution=30, error=2, max_iters=10)
simulator = Aer.get_backend('qasm_simulator')

for i in range(4):
    basis_index = i
    eigen_vals_ret1 = []
    while len(eigen_vals_ret1) != 15:
        res = spea1.get_eigen_pair(
            backend=simulator, algo='alternate', basis=basis, basis_ind=basis_index)
        if res['cost'] < 0.75:
            continue
        print((len(eigen_vals_ret1)/15)*100, "% done...")
        eigen_vals_ret1.append(res['theta'])
    eig_ret.append(eigen_vals_ret1)

generate_plot2(eig_v, eig_ret, 4, 15)

eig_ret2 = []

spea2 = global_max_SPEA(unit, resolution=30, error=2, max_iters=10)
simulator = Aer.get_backend('qasm_simulator')

for i in range(4):
    basis_index = i
    eigen_vals_ret2 = []
    while len(eigen_vals_ret2) != 40:
        res = spea2.get_eigen_pair(
            backend=simulator, algo='alternate', basis=basis, basis_ind=basis_index)
        if res['cost'] < 0.75:
            continue
        print((len(eigen_vals_ret2)/40)*100, "% done...")
        eigen_vals_ret2.append(res['theta'])
    eig_ret2.append(eigen_vals_ret2)

generate_plot2(eig_v, eig_ret2, 4, 40)

plt.title("Original Algorithm ")
for i in range(4):
    sns.kdeplot(eig_ret[i], shade=True, label='Basis '+str(i))
plt.xlabel("Eigenvalues Returned")
plt.legend()

plt.title("Modified Algorithm ")
for i in range(4):
    sns.kdeplot(eig_ret2[i], shade=True, label='Basis '+str(i), palette='red')
plt.xlabel("Eigenvalues Returned")
plt.legend()

def min_eigensolver(unitary, algo, experiments=20, min_threshold=0.9):

    dims = unitary.shape[0]
    simulator = Aer.get_backend('qasm_simulator')
    step_size = (1/dims)

    # generate the steps
    steps = np.arange(0, 1 + 1e-12, step_size)
    if steps[-1] != 1:
        steps = np.append(steps, 1)
    print("Steps :", steps)
    if algo == 0:
        spe = SPEA(unitary, resolution=40, error=2, max_iters=12)
    else:
        spe = global_max_SPEA(unitary, resolution=40, error=2, max_iters=12)
    # start the experiments

    max_cost = -1
    for i in range(len(steps) - 1):
        # define left and right bounds
        left = steps[i]
        right = steps[i+1]
        
        costs, eigs = [], []  # lists to store costs
        for _ in range(experiments):
            res = spe.get_eigen_pair(
                backend=simulator, theta_left=left, theta_right=right,
                randomize=True)
            costs.append(res['cost'])
            eigs.append(res['theta'])
        
        # if the cost is above threshold, return them
        if np.average(costs) > min_threshold:
            return (np.average(costs), eigs)
        if np.average(costs) > max_cost:
            max_cost = np.average(costs)
            best_eigs = eigs

    # return the cost with the max average cost
    return (max_cost, best_eigs)

unit = unitary_group.rvs(4)
unit 

# get the eigenvalues and eigenstates
eig_v, eig_vect = np.linalg.eig(unit)
eig_v = np.angle(eig_v)
eig = []
for k in eig_v:
    if k < 0:
        e = (k+2*np.pi)/(2*np.pi)
    else:
        e = (k)/(2*np.pi)
    eig.append(e)
eig_v = np.array(eig)
print("Eigenvalues :", eig_v)
print("Eigenstates :", eig_vect)

res1 = min_eigensolver(unit, algo=0, experiments=40, min_threshold=0.85)
res2 = min_eigensolver(unit, algo=1, experiments=40, min_threshold=0.85)

res1[1], res2[1]

generate_plot1(eig_v,res1[1],40)

generate_plot1(eig_v,res2[1],40)