from qiskit import QuantumCircuit, execute, transpile, Aer
from qiskit.extensions import UnitaryGate, Initialize
from qiskit.quantum_info import Statevector
from qiskit.tools.monitor import job_monitor
from qiskit.compiler import assemble
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
from scipy.stats import unitary_group
import matplotlib.pyplot as plt

x = np.linspace(0.05, 1, 200)
c = np.cos(x)
c2 = (np.cos(x))**2
ex = np.exp(-x)
ex2 = np.exp(-x**2)
x_inv = (1/x)
x_inv2 = (1/x**2)
plt.figure(dpi=120)
plt.title("Cost Function plots", fontsize=20)
plt.xlabel("Theta values", fontsize=16)
plt.ylabel("Metric values", fontsize=16)
plt.plot(x, c, label="$cos(x)$")
plt.plot(x, c2, label="$cos^{2}(x)$")
plt.plot(x, ex, label="$e^{-x}$")
plt.plot(x, ex2, label="$e^{-x^{2}}$")
plt.grid()
plt.legend()

plt.figure(dpi=120)
plt.title("Cost Function plots", fontsize=20)
plt.xlabel("Theta values", fontsize=16)
plt.ylabel("Metric values", fontsize=16)
plt.plot(x, x_inv, label="$\\frac{1}{x}$")
plt.plot(x, x_inv2, label="$\\frac{1}{x^{2}}$")
plt.grid()
plt.legend()

class global_max_SPEA():
    def __init__(self, unitary, cost_function, max_metric_val, resolution=100, max_iters=20):

        # handle resolution
        if not isinstance(resolution, int):
            raise TypeError(
                "Please enter the number of intervals as an integer value")
        if resolution < 10 or resolution > 1e6:
            raise ValueError(
                "Resolution needs to be atleast 0.1 and greater than 0.000001")

        self.resolution = resolution

        # handle unitary
        if not isinstance(unitary, np.ndarray) and not isinstance(unitary, QuantumCircuit)\
                and not isinstance(unitary, UnitaryGate):
            raise TypeError(
                "A numpy array or Quantum Circuit or UnitaryGate needs to be passed as the unitary matrix")

        # convert circuit to numpy array for uniformity
        if isinstance(unitary, UnitaryGate):
            U = unitary.to_matrix()
        else:  # both QC and ndarray type
            U = unitary

        # note - the unitary here is not just a single qubit unitary
        if isinstance(U, np.ndarray):
            self.dims = U.shape[0]
        else:
            self.dims = 2**(U.num_qubits)

        if isinstance(U, np.ndarray):
            self.c_unitary_gate = UnitaryGate(data=U).control(
                num_ctrl_qubits=1, label='CU', ctrl_state='1')
        else:
            self.c_unitary_gate = U.control(
                num_ctrl_qubits=1, label='CU', ctrl_state='1')

            
            
            
        # handle metric
        if not isinstance(max_metric_val, int) and not isinstance(max_metric_val, float):
            raise TypeError(
                "The required metric value should be provided as an int or a float.")
        elif max_metric_val <= 0:
            raise ValueError(
                "The required metric must be finite and greater than 0.")
        elif cost_function in ['cos', 'cos2', 'exp', 'exp2'] and max_metric_val > 1:
            raise ValueError(
                "Maximum metric can't be greater than 1 in case of given cost function")

        self.max_metric = max_metric_val

        
        
        
        
        # handle max_iters
        if not isinstance(max_iters, int):
            raise TypeError("Max iterations must be of integer type")
        if max_iters <= 0 and max_iters > 1e5:
            raise ValueError(
                "Max iterations should be atleast 1 and less than 1e5")

        self.iterations = max_iters
        cost_fns = ['cos', 'cos2', 'exp', 'exp2', 'inv', 'inv2']
        # handle cost function
        if cost_function not in cost_fns:
            raise ValueError(
                "Cost function must be one of [cos, cos2, exp, exp2, inv, inv2]")

        self.cost_function = cost_function
        self.basis = []

    def __get_basis_vectors(self, randomize=True):
        # get the d dimensional basis for the unitary provided
        if randomize == True:
            UR = unitary_group.rvs(self.dims)
        else:
            UR = np.identity(self.dims)

        basis = []
        for k in UR:
            basis.append(np.array(k, dtype=complex))
        return basis

    def __get_unitary_circuit(self, backend):
        '''Return the pretranspiled circuit '''
        if backend is None:
            backend = Aer.get_backend('qasm_simulator')

        qc = QuantumCircuit(1 + int(np.log2(self.dims)))

        # make the circuit
        qc.h(0)
        qc = qc.compose(self.c_unitary_gate, qubits=range(
            1+int(np.log2(self.dims))))

        qc.barrier()
        qc = transpile(qc, backend=backend, optimization_level=3)

        return qc

    def __get_circuit(self, state, backend, shots, angle=None):
        '''Given an initial state ,
          return the circuit that is generated with 
          inverse rotation '''
        # all theta values are iterated over for the same state
        phi = Initialize(state)

        qc1 = QuantumCircuit(1 + int(np.log2(self.dims)), 1)
        # initialize the circuit
        qc1 = qc1.compose(phi, qubits=list(
            range(1, int(np.log2(self.dims))+1)))
        qc1.barrier()
        qc1 = transpile(qc1, backend=backend, optimization_level=1)

        # get the circuit2
        qc2 = self.unitary_circuit

        qc3 = QuantumCircuit(1 + int(np.log2(self.dims)), 1)
        if angle is not None:
            # add inverse rotation on the first qubit
            qc3.p(-2*np.pi*angle, 0)
        # add hadamard
        qc3.h(0)
        qc3 = transpile(qc3, backend=backend, optimization_level=1)

        # make final circuit
        qc = qc1 + qc2 + qc3

        # measure
        qc.measure([0], [0])

#         qc = assemble(qc,shots = shots)
        return qc

    def __get_standard_cost(self, angles, state, backend, shots):
        '''Given an initial state and a set of angles,
          return the best cost and the associated angle
          state is a normalized state in ndarray form'''
        result = {'cost': -1, 'theta': -1}
        # all theta values are iterated over for the same state
        circuits = []

        for theta in angles:
            qc = self.__get_circuit(state, backend, shots, theta)
            circuits.append(qc)

        # execute only once...
        counts = backend.run(circuits, shots=shots).result().get_counts()
        # get the cost for this theta
        for k, theta in zip(counts, angles):
            # for all experiments you ran
            try:
                C_val = (k['0'])/shots
            except:
                C_val = 0

            if C_val > result['cost']:
                # means this is a better theta value
                result['theta'] = theta
                result['cost'] = C_val
        return result

    def __get_alternate_cost(self, angles, state, backend, shots):
        '''Given an initial state and a set of angles,
          return the best cost and the associated angle
          state is a normalized state in ndarray form'''
        result = {'cost': -1, 'theta': -1}
        # all theta values are iterated over for the same state
        # run the circuit once
        qc = self.__get_circuit(state, backend, shots)

        # execute only once...
        counts = backend.run(qc, shots=shots).result().get_counts()

        # generate experimental probabilities
        try:
            p0 = counts['0']/shots
        except:
            p0 = 0
        try:
            p1 = counts['1']/shots
        except:
            p1 = 0

        # now, find the best theta as specified by the
        # alternate method classically
        min_s = 1e5
        for theta in angles:
            # generate theoretical probabilities
            c0 = (np.cos(np.pi*theta))**2
            c1 = (np.sin(np.pi*theta))**2

            # generate s value
            s = (p0-c0)**2 + (p1-c1)**2
            if s < min_s:
                result['theta'] = theta
                min_s = s

        # now , we have the best theta stored in phi
        # run circuit once again to get the value of C*
        qc = self.__get_circuit(state, backend, shots, result['theta'])
        counts = backend.run(qc, shots=shots).result().get_counts()

        try:
            result['cost'] = counts['0']/shots
        except:
            result['cost'] = 0
        # no 0 counts present

        # return the result
        return result

    
    
    
    
    
    def __get_metric_val(self, cost, theta):
        '''Generate optimization metric value 
        based on the given cost function '''
        func = self.cost_function

        if func == 'cos':
            return cost*(np.cos(theta))
        elif func == 'cos2':
            return cost*((np.cos(theta))**2)
        elif func == 'exp':
            return cost*(np.exp(-theta))
        elif func == 'exp2':
            return cost*(np.exp(-theta**2))
        elif func == 'inv':
            try:
                return cost/theta
            except:
                return 1e7
        elif func == 'inv2':
            try:
                return cost/(theta**2)
            except:
                return 1e7
            
            
            
            
            

    def get_eigen_pair(self, backend, algo='alternate', theta_left=0, theta_right=1,
                       progress=False, basis=None, basis_ind=None, randomize=True, target_metric=None, shots=512
                       ):
        '''Finding the eigenstate pair for the unitary'''
        # handle algorithm...

        self.unitary_circuit = self.__get_unitary_circuit(backend)

        # handle theta bounds
        if(theta_left > theta_right):
            raise ValueError(
                "Left bound for theta should be smaller than the right bound")
        elif (theta_left < 0) or (theta_right > 1):
            raise ValueError("Bounds of theta are [0,1].")

        # handle algorithm
        if not isinstance(algo, str):
            raise TypeError(
                "Algorithm must be mentioned as a string from the values {alternate,standard}")
        elif algo not in ['alternate', 'standard']:
            raise ValueError(
                "Algorithm must be specified as 'alternate' or 'standard' ")

        if target_metric is not None:
            if (target_metric <= 0):
                raise ValueError("Target metric must be a real value > 0")

        # handle progress...
        if not isinstance(progress, bool):
            raise TypeError("Progress must be a boolean variable")
        if not isinstance(randomize, bool):
            raise Exception("Randomize must be a boolean variable")

        results = dict()

        # first initialize the state phi
        if basis is None:
            self.basis = self.__get_basis_vectors(randomize)
        else:
            # is basis is specified, given as array of vectors...
            self.basis = basis

        # choose a random index
        if basis_ind is None:
            ind = np.random.choice(self.dims)
        else:
            # choose the index given in that basis
            ind = basis_ind

        phi = self.basis[ind]
        # doing the method 1 of our algorithm

        
        
        
        # new
        if target_metric is None:
            max_metric = self.max_metric
        else:
            if self.cost_function in ['cos','cos2','exp','exp2']:
                assert target_metric <= 1 , "For given cost function, target cost can't be greater than 1"
            max_metric = target_metric

            
            
            
            
        samples = self.resolution

        # initialization of range
        left, right = theta_left, theta_right

        # generate the angles
        angles = np.linspace(left, right, samples)

        # iterate once
        if algo == 'alternate':
            result = self.__get_alternate_cost(angles, phi, backend, shots)
        else:
            result = self.__get_standard_cost(angles, phi, backend, shots)

        # get initial estimates
        cost = result['cost']
        theta_max = result['theta']
        best_phi = phi
        global_metric = self.__get_metric_val(cost, theta_max)

        # the range upto which theta extends iin each iteration
        angle_range = (right - left)/2
        # a parameter
        a = 1
        # start algorithm
        iters = 0
        found = True

        while global_metric < max_metric:
            # get angles, note if theta didn't change, then we need to
            # again generate the same range again
            right = min(theta_right, theta_max + angle_range/2)
            left = max(theta_left, theta_max - angle_range/2)
            if progress:
                print("Right :", right)
                print("Left :", left)
            # generate the angles only if the theta has been updated
            if found == True:
                angles = np.linspace(left, right, samples)

            found = False  # for this iteration
            if progress:
                print("ITERATION NUMBER", iters+1, "...")

            # generate a cost dict for each of the iterations

            thetas, costs, states, metrics = [], [], [], []

            for i in range((2*self.dims)):
                # everyone is supplied with the same range of theta in one iteration
                # define z
                if i < self.dims:
                    z = 1
                else:
                    z = 1j
                    
                    
                    
                    

                # alter and normalise phi
                curr_phi = best_phi + z*a * \
                    (max_metric - global_metric)*self.basis[i % self.dims]
                curr_phi = curr_phi / np.linalg.norm(curr_phi)

                
                
                
                
                
                # iterate (angles would be same until theta is changed)
                if algo == 'alternate':
                    res = self.__get_alternate_cost(
                        angles, curr_phi, backend, shots)
                else:
                    res = self.__get_standard_cost(
                        angles, curr_phi, backend, shots)

                curr_cost = res['cost']
                curr_theta = res['theta']
                # new
                curr_metric = self.__get_metric_val(curr_cost, curr_theta)

                # append these parameters

                if curr_metric > global_metric:  # then only add this cost in the cost and states list
                    thetas.append(float(curr_theta))
                    costs.append(float(curr_cost))
                    states.append(curr_phi)
                    # new
                    metrics.append(curr_metric)
                    found = True

                    # now each iteration would see the same state as the best phi
                    # is updated once at the end of the iteration

                    # also, the cost is also updated only once at the end of the iteration

                if progress:
                    sys.stdout.write('\r')
                    sys.stdout.write("%f %%completed" %
                                     (100*(i+1)/(2*self.dims)))
                    sys.stdout.flush()

            # 1 iteration completes

            if found == False:
                # phi was not updated , change a
                a = a/2
                if progress:
                    print("\nNo change, updating a...")
            else:
                # if found is actually true, then only update

                # new
                # we need to update the best phi , theta max and cost on
                # the basis of cost / theta_max value

                
                
                
                index = np.argmax(metrics)

                
                
                
                
                # update the parameters of the model
                cost = costs[index]
                theta_max = thetas[index]
                best_phi = states[index]
                global_metric = metrics[index]

                angle_range /= 2  # updated phi and thus theta too -> refine theta range

            # update the iterations
            iters += 1
            if progress:
                print("Best Phi is :", best_phi)
                print("Theta estimate :", theta_max)
                print("Current cost :", cost)
                print("Current Metric :", global_metric)

            if iters >= self.iterations:
                print(
                    "Maximum iterations reached for the estimation.\nTerminating algorithm...")
                break

        # add cost, eigenvector and theta to the dict
        results['cost'] = cost
        results['theta'] = theta_max
        results['state'] = best_phi
        results['metric'] = global_metric

        return results

unit = unitary_group.rvs(8)
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)

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

    plt.figure(figsize=(9, 7))
    plt.scatter(range(experiments), 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, experiments], [k, k], color='black',
                     linewidth=2, label='Actual Values')
        else:
            plt.plot([0, experiments], [k, k], color='black',
                     linewidth=2)
    plt.xlabel("Experiment Number", fontsize=14)
    plt.ylabel("Eigenvalues", fontsize=14)
    plt.title("Scatter plot for returned eigenvalues with cost function "+str(function), fontsize=17)
    plt.legend()
    plt.grid()

spea1 = global_max_SPEA(unit, resolution=30, max_metric_val=0.92,cost_function='cos', max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 40:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate')
    if res['metric'] < 0.8:
        continue
    print("Result :", res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 40, "cos")

spea1 = global_max_SPEA(unit, resolution=30, max_metric_val=0.92,cost_function='cos2', max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 40:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate')
    if res['metric'] < 0.8:
        continue
#     print("Result :", res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 40, "cos2")

spea1 = global_max_SPEA(unit, resolution=30, max_metric_val=0.92,cost_function='exp', max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 40:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate')
    if res['metric'] < 0.8:
        continue
#     print("Result :", res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 40, "exp")

spea1 = global_max_SPEA(unit, resolution=30, max_metric_val=0.92,cost_function='exp2', max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 40:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate')
    if res['metric'] < 0.8:
        continue
    print("Result :", res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 40, "exp2")

spea1 = global_max_SPEA(unit, resolution=30, max_metric_val=4,cost_function='inv', max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 40:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate')
    if res['metric'] < 3:
        continue
    print("Result :", res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 40, "inv")

spea1 = global_max_SPEA(unit, resolution=30, max_metric_val=20,cost_function='inv2', max_iters=15)
simulator = Aer.get_backend('qasm_simulator')

eigen_vals_ret1 = []
while len(eigen_vals_ret1) != 40:
    res = spea1.get_eigen_pair(
        backend=simulator, algo='alternate')
    if res['metric'] < 15:
        continue
    print("Result :", res)
    eigen_vals_ret1.append(res['theta'])

generate_plot1(eig_v, eigen_vals_ret1, 40, "inv2")

