import numpy as np
from itertools import permutations
import gzip


from qiskit import*

import time

from qiskit.aqua.algorithms import VQE
from qiskit.aqua.algorithms import QAOA

from qiskit.aqua.components.optimizers import SPSA
# from qiskit.aqua.components.variational_forms import RY
from qiskit.aqua import QuantumInstance
from qiskit.aqua.components.optimizers import COBYLA

from qiskit import IBMQ

from qiskit import ClassicalRegister, QuantumRegister, QuantumCircuit

from qiskit.optimization.applications.ising.tsp import TspData
import qiskit.optimization.applications.ising.tsp as tsp
from qiskit.circuit.library import TwoLocal
from qiskit.circuit.library import RealAmplitudes

# # from qiskit import IBMQ
# # # Load account from disk
IBMQ.load_account()
IBMQ.providers()

def readInData():
    
    """
    Output G : N by N distance matrix from Matrices1a.txt file. 
    
    """
    
    G = []

    p = [3,4,5,6,7,8,9,10,11]
    q = [i**(2) for i in p ]
    m = 0
    
    v = open("Matrices.txt"  , "r")
    w = v.read().split()
    for i in range (len(w)):
        w[i] = int(float(w[i]))
    for i in range (len(q)):
        G.append(np.reshape(w[m:m+q[i]] , (p[i] , p[i])))
        m = m + q[i]
        
    return G

distanceMatrix =  readInData() #Array of different sized matrices


def determineIfFeasible(result):
    
    """
    Determines if eigenstate is feasible or infeasible.
    
    Output: arr = Infeasible if eiegenstate is infeasible or arr = binary array of feasible solution
    
    """
    
    
    data = sorted(result['eigenstate'].items(), key=lambda item: item[1])[::-1]
    for i in range(len(data)):
        a = tsp.tsp_feasible(data[i][0])
        arr = 'Infeasible'
        if a == True:
            b = str(data[i][0])
            arr = [b , data[i][1]]
            break
    return arr


def optimal(a,b,c,f,u):
    
    """
    Read in data of initial optimal point that will be used in the quantum algorithm
    
    """
    openfile = open("optimal.txt" , "r")
    readFile = openfile.read().split()
    t = []
    for i in readFile:
        if i != ',':
            q = len(i)
            t.append(float(i[0:q-1]))

    v, r, o, d, z = np.array(t[0:a]), np.array(t[a:a+b]), np.array(t[a+b : a+b+c]), np.array(t[a+b+c:a+b+c+f]), np.array(t[a+b+c+f:a+b+c+f+u])
    
    return [v,r,o,d,z]

R = optimal(54,96,100,216,294) #Array of corresponding initial points

def quantumApproximateOptimizationAlgorithm(numIter, numShots, distanceMatrix,pValue, deviceName, initialPoint):
    
    """
    Implementation of the QAOA
    
    Output: classial TSP solution (total length of tour), time taken to execute algorithm
    
    """
    # Map problem to isining hamiltonian
    x = TspData('tmp',len(distanceMatrix),np.zeros((3,3)),distanceMatrix)
    qubitOp  = tsp.get_operator(x)
    
    seed = 10598
    spsa = SPSA(maxiter = numIter)
        
    qaoa =  QAOA(qubitOp, spsa, pValue, include_custom = False, initialPoint = initialPoint)
    
    
    my_provider = IBMQ.get_provider(ibm_hub) #Replace ibm_hub with appropriate qiskit hub name
    
    device = my_provider.get_backend(deviceName) # deviceName is the Device of IBM qunatum device in a string
    
    quantum_instance = QuantumInstance(device, seed_simulator=seed, seed_transpiler=seed,shots = numShots,
                                       skip_qobj_validation = False)

    #Convert quantum result into its classical form and determine if feasible or infeasible
    result = qaoa.run(quantum_instance)
    
    answer = determineIfFeasible(result)
    
    if answer == 'Infeasible':
        solution = -1
    else:
        binarry = [int(p) for p in answer[0]]
        route = tsp.get_tsp_solution(binarry)
        solution = tsp.tsp_value(route,distanceMatrix)
        
    return  solution, result['optimizer_time']

## Example for 3 by 3 instance implemented using QAOA:

numIter = 1
numShots = 8192
distanceMatrix = distanceMatrix[0] 
pValue = 3
deviceName = 'ibmq_manhattan'
initialPoint = R[0]

finalResult  quantumApproximateOptimizationAlgorithm(numIter, numShots, distanceMatrix, pValue, deviceName, initialPoint)