from qiskit import QuantumCircuit, execute, transpile, Aer 
from qiskit.extensions import UnitaryGate,Initialize
from qiskit.quantum_info import Statevector 
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 
import os 
from scipy.stats import unitary_group 
import matplotlib.pyplot as plt 
%matplotlib inline 

from qiskit import IBMQ 
IBMQ.load_account()

provider = IBMQ.get_provider(hub='ibm-q-education')
simulator = Aer.get_backend('qasm_simulator')
athens = provider.get_backend('ibmq_athens')


class global_max_SPEA():
    def __init__(self,unitary,resolution = 100, error = 3, 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 error 
        if not isinstance(error,int):
            raise TypeError("The allowable error should be provided as an int. Interpreted as 10**(-error)")
        if error <= 0:
            raise ValueError("The error threshold must be finite and greater than 0.")
            
        self.error = error 
        
        # 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 
        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, 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)
        shots = 512

        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 = 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])
        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
        phi = Initialize(state)
        circuits = []

        for theta in angles:
            qc = self.get_circuit(state,backend,theta)
            circuits.append(qc)
        
        circuits = assemble(circuits)
        # 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
        phi = Initialize(state)

        qc = self.get_circuit(state,backend)
        qc = assemble(qc)
        # 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, result['theta'])
        qc = assemble(qc)
        counts = backend.run(qc, shots=shots).result().get_counts()
        
        try:
            result['cost'] = counts['0']/shots
        except:
            result['cost'] = 0
        # no 0 counts presenta

        # return the result
        return result
    
    def get_eigen_pair(self,backend,algo = 'alternate',progress = False,randomize = True,
                      basis=None,basis_ind = None, shots=512):
        '''Finding the eigenstate pair for the unitary'''
        #handle algorithm...
        
        self.unitary_circuit = self.get_unitary_circuit(backend)
        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' ")
        
        # 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 
        # define resolution of angles and precision 
        precision = 1/10**self.error 
        samples = self.resolution 
        
        # initialization of range 
        left,right = 0,1
        
        # 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 

        # the range upto which theta extends iin each iteration 
        angle_range = 0.5
        # a parameter 
        a = 1 
        # start algorithm        
        iters = 0 
        found = True
        plus = (1/np.sqrt(2))*np.array([1,1])
        minus = (1/np.sqrt(2))*np.array([1,-1])
        
        while 1 - cost >= precision:
            # get angles, note if theta didn't change, then we need to 
            # again generate the same range again 
            right = min(1,theta_max + angle_range/2)
            left = max(0,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 = [],[],[] 
            
            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*(1 - cost)*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']
                
                # append these parameters 
                
                # bundle the circuits together ...
                
                if curr_cost > cost: # 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)
                    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 
                
                # O(n) , would update this though 
                index = np.argmax(costs)
                # update the parameters of the model 
                cost = costs[index]
                theta_max = thetas[index]
                best_phi = states[index]
                if progress:
                    print("Best Phi is :",best_phi)
                    print("Theta estimate :",theta_max)
                    print("Current cost :",cost) 
                angle_range /= 2 # updated phi and thus theta too -> refine theta range
            
            # update the iterations 
            iters+=1 
            if progress:
                print("\nCOST :",cost)
                print("THETA :",theta_max)
            
            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 
        return results
            

simulator = Aer.get_backend('qasm_simulator')

U = QuantumCircuit(2)
U.cp(2*np.pi*(1/3),0,1)
U.draw('mpl')

spe = global_max_SPEA(U,resolution= 40,error = 3,max_iters=15)

# O(max_iters*(shots+resolution)*|Bm|)
result = spe.get_eigen_pair(progress = False,algo = 'alternate',backend=simulator,shots=512)

result 

U = QuantumCircuit(3)
U.cp(2*np.pi*(1/3),1,2)
U.draw('mpl')

errors = []
eigvals = np.array([0.0,1/3,1.0])

for resolution in range(10,80,10):
    spe = global_max_SPEA(U,resolution= resolution,error = 4,max_iters=15)    
    res = spe.get_eigen_pair(backend = simulator)
    theta = res['theta']
    min_error = 1e5 
    # get percentage error 
    for e in eigvals:
        error = abs(e-theta) 
        if error < min_error:
            min_error = error 
            if e != 0:
                perc_error = (error/e)*100 
            else:
                perc_error = error*100
    errors.append(perc_error)

plt.title("Eigenphase estimation for theta = 1/3, 0, 1 (3-qubit unitary)",fontsize = 15)
plt.grid()
plt.plot(list(range(10,80,10)),errors,marker = 'o',color='g',label = 'Estimates',alpha=0.7)
# plt.plot([10,80],[0.25,0.25],color='black',label = "True")
# plt.plot([10,80],[0,0],color='black',label = "True")
# plt.plot([10,80],[1,1],color='black',label = "True")
plt.legend()
plt.xlabel("Resolution of theta ")
plt.ylabel("Errors from closest eigvals")
plt.savefig("Alternate_SPE_PLOT_1.jpg",dpi = 200)

U = QuantumCircuit(2)
U.cp(2*np.pi*(1/2),0,1)
U.draw('mpl')

eigvals = np.array([0.0,1/2,1.0])

errors = []
for resolution in range(10,80,10):
    spe = global_max_SPEA(U,resolution= resolution,error = 3,max_iters=12)    
    #run 5 experiments of this 
    p_errors = []
    
    for exp in range(4):
        res = spe.get_eigen_pair(progress = False,backend = simulator)
        theta = res['theta']
        min_error = 1e5 
        min_val = -1 
        # get min error 
        
        for e in eigvals:
            error = abs(e-theta) 
            if error < min_error:
                min_error = error 
                min_val = e 
                
        # percent error in this experiment 
        if min_val!=0:
            p_errors.append((min_error/min_val)*100)
        else:
            p_errors.append((min_error)*100)
    print("Percentage Errors in current iteration:",p_errors)
    errors.append(np.average(p_errors))
    print("Errors :",errors) 

plt.title("Eigenphase estimation for theta = 1/2, 0.0, 1.0 (2-qubit unitary)",fontsize = 15)
plt.grid()
plt.plot(list(range(10,80,10)),errors,marker = 'o',color='orange',label = 'Estimates',alpha=0.7)
# plt.plot([10,80],[0.25,0.25],color='black',label = "True")
# plt.plot([10,80],[0,0],color='black',label = "True")
# plt.plot([10,80],[1,1],color='black',label = "True")
plt.legend()
plt.xlabel("Resolution of theta ")
plt.ylabel("Errors from closest eigvals")
plt.savefig("Alternate_SPE_PLOT_2(simulator).jpg",dpi = 200)

