# To clean enviroment variables 
%reset

import numpy as np
import pandas as pd
import folium
import matplotlib.pyplot as plt

try:
    import cplex
    from cplex.exceptions import CplexError
except:
    print("Warning: Cplex not found.")
import math

from qiskit.utils import algorithm_globals
from qiskit_algorithms import SamplingVQE
from qiskit_algorithms.optimizers import SPSA
from qiskit.circuit.library import RealAmplitudes
from qiskit.primitives import Sampler

df = pd.read_csv("uscities.csv")
columnsNeeded = ["city", "lat", "lng"]
# Inicialization of variables
locationsNumber = 15 # OJO que en local me crashea si sobrepaso 5

coordenatesDf = df[columnsNeeded].iloc[:locationsNumber]
n = coordenatesDf.shape[0]  # number of nodes + depot (n+1)
K = 3  # number of vehicles

print(coordenatesDf)

# Initialize instance values by calculate the squared Euclidean distances between a set of coordinates
#  represented in the dataframe.

def generate_instance(coordenatesDf):
    n = coordenatesDf.shape[0]
    
    xc = coordenatesDf["lat"]
    yc = coordenatesDf["lng"]
    loc = coordenatesDf["city"]
    instance = np.zeros([n, n])
    for ii in range(0, n):
        for jj in range(ii + 1, n):
            instance[ii, jj] = (xc[ii] - xc[jj]) ** 2 + (yc[ii] - yc[jj]) ** 2
            instance[jj, ii] = instance[ii, jj]
    return xc, yc, instance, loc

# Initialize the problem by randomly generating the instance
lat, lng, instance, loc = generate_instance(coordenatesDf)

print(lat, lng, loc, instance)

#print(instance)

class ClassicalOptimizer:
    def __init__(self, instance, n, K):

        self.instance = instance
        self.n = n  # number of nodes
        self.K = K  # number of vehicles

    def compute_allowed_combinations(self):
        f = math.factorial
        return f(self.n) / f(self.K) / f(self.n - self.K)

    def cplex_solution(self):

        # refactoring
        instance = self.instance
        n = self.n
        K = self.K

        my_obj = list(instance.reshape(1, n**2)[0]) + [0.0 for x in range(0, n - 1)]
        my_ub = [1 for x in range(0, n**2 + n - 1)]
        my_lb = [0 for x in range(0, n**2)] + [0.1 for x in range(0, n - 1)]
        my_ctype = "".join(["I" for x in range(0, n**2)]) + "".join(
            ["C" for x in range(0, n - 1)]
        )

        my_rhs = (
            2 * ([K] + [1 for x in range(0, n - 1)])
            + [1 - 0.1 for x in range(0, (n - 1) ** 2 - (n - 1))]
            + [0 for x in range(0, n)]
        )
        my_sense = (
            "".join(["E" for x in range(0, 2 * n)])
            + "".join(["L" for x in range(0, (n - 1) ** 2 - (n - 1))])
            + "".join(["E" for x in range(0, n)])
        )

        try:
            my_prob = cplex.Cplex()
            self.populatebyrow(my_prob, my_obj, my_ub, my_lb, my_ctype, my_sense, my_rhs)

            my_prob.solve()

        except CplexError as exc:
            print(exc)
            return

        x = my_prob.solution.get_values()
        x = np.array(x)
        cost = my_prob.solution.get_objective_value()

        return x, cost

    def populatebyrow(self, prob, my_obj, my_ub, my_lb, my_ctype, my_sense, my_rhs):

        n = self.n

        prob.objective.set_sense(prob.objective.sense.minimize)
        prob.variables.add(obj=my_obj, lb=my_lb, ub=my_ub, types=my_ctype)

        prob.set_log_stream(None)
        prob.set_error_stream(None)
        prob.set_warning_stream(None)
        prob.set_results_stream(None)

        rows = []
        for ii in range(0, n):
            col = [x for x in range(0 + n * ii, n + n * ii)]
            coef = [1 for x in range(0, n)]
            rows.append([col, coef])

        for ii in range(0, n):
            col = [x for x in range(0 + ii, n**2, n)]
            coef = [1 for x in range(0, n)]

            rows.append([col, coef])

        # Sub-tour elimination constraints:
        for ii in range(0, n):
            for jj in range(0, n):
                if (ii != jj) and (ii * jj > 0):

                    col = [ii + (jj * n), n**2 + ii - 1, n**2 + jj - 1]
                    coef = [1, 1, -1]

                    rows.append([col, coef])

        for ii in range(0, n):
            col = [(ii) * (n + 1)]
            coef = [1]
            rows.append([col, coef])

        prob.linear_constraints.add(lin_expr=rows, senses=my_sense, rhs=my_rhs)


# Instantiate the classical optimizer class
classical_optimizer = ClassicalOptimizer(instance, n, K)

# Print number of feasible solutions
print("Number of feasible solutions = " + str(classical_optimizer.compute_allowed_combinations()))


# Solve the problem in a classical fashion via CPLEX
x = None
z = None
try:
    x, classical_cost = classical_optimizer.cplex_solution()
    # Put the solution in the z variable
    z = [x[ii] for ii in range(n**2) if ii // n != ii % n]
    # Print the solution
    print(z)
except:
    print("CPLEX may be missing.")


m = folium.Map(location=[39.487660, -97.594333], zoom_start=0)
marker_icon1 = folium.Icon(color = "red")
for i in range(len(lat)):
    if (i == 0):
        folium.Marker(location=[lat[i], lng[i]], tooltip=f"Location: {loc[i]}, Order: {i}", icon=marker_icon1).add_to(m)
    else:
        folium.Marker(location=[lat[i], lng[i]], tooltip=f"Location: {loc[i]}, Order: {i}").add_to(m)

for ii in range(0, n**2):
    if x[ii] > 0:
        ix = ii // n
        iy = ii % n
        
        folium.PolyLine([(lat[ix], lng[ix]), (lat[iy], lng[iy])], color="blue").add_to(m)

m

from qiskit_optimization import QuadraticProgram
from qiskit_optimization.algorithms import MinimumEigenOptimizer
from qiskit.algorithms.minimum_eigensolvers import QAOA, NumPyMinimumEigensolver

class QuantumOptimizer:
    def __init__(self, instance, n, K):

        self.instance = instance
        self.n = n
        self.K = K

    def binary_representation(self, x_sol=0):

        instance = self.instance
        n = self.n
        K = self.K

        A = np.max(instance) * 100  # A parameter of cost function

        # Determine the weights w
        instance_vec = instance.reshape(n**2)
        w_list = [instance_vec[x] for x in range(n**2) if instance_vec[x] > 0]
        w = np.zeros(n * (n - 1))
        for ii in range(len(w_list)):
            w[ii] = w_list[ii]

        # Some variables I will use
        Id_n = np.eye(n)
        Im_n_1 = np.ones([n - 1, n - 1])
        Iv_n_1 = np.ones(n)
        Iv_n_1[0] = 0
        Iv_n = np.ones(n - 1)
        neg_Iv_n_1 = np.ones(n) - Iv_n_1

        v = np.zeros([n, n * (n - 1)])
        for ii in range(n):
            count = ii - 1
            for jj in range(n * (n - 1)):

                if jj // (n - 1) == ii:
                    count = ii

                if jj // (n - 1) != ii and jj % (n - 1) == count:
                    v[ii][jj] = 1.0

        vn = np.sum(v[1:], axis=0)

        # Q defines the interactions between variables
        Q = A * (np.kron(Id_n, Im_n_1) + np.dot(v.T, v))

        # g defines the contribution from the individual variables
        g = (
            w
            - 2 * A * (np.kron(Iv_n_1, Iv_n) + vn.T)
            - 2 * A * K * (np.kron(neg_Iv_n_1, Iv_n) + v[0].T)
        )

        # c is the constant offset
        c = 2 * A * (n - 1) + 2 * A * (K**2)

        try:
            max(x_sol)
            # Evaluates the cost distance from a binary representation of a path
            fun = (
                lambda x: np.dot(np.around(x), np.dot(Q, np.around(x)))
                + np.dot(g, np.around(x))
                + c
            )
            cost = fun(x_sol)
        except:
            cost = 0

        return Q, g, c, cost

    def construct_problem(self, Q, g, c) -> QuadraticProgram:
        qp = QuadraticProgram()
        for i in range(n * (n - 1)):
            qp.binary_var(str(i))
        qp.objective.quadratic = Q
        qp.objective.linear = g
        qp.objective.constant = c
        return qp

    def solve_problem(self, qp):
        algorithm_globals.random_seed = 10598
        #vqe = SamplingVQE(sampler=Sampler(), optimizer=SPSA(), ansatz=RealAmplitudes())
        #optimizer = MinimumEigenOptimizer(min_eigen_solver=vqe)

        meo = MinimumEigenOptimizer(min_eigen_solver=NumPyMinimumEigensolver())
        result = meo.solve(qp)

        # compute cost of the obtained result
        _, _, _, level = self.binary_representation(x_sol=result.x)
        return result.x, level

# Instantiate the quantum optimizer class with parameters:
quantum_optimizer = QuantumOptimizer(instance, n, K)

# Check if the binary representation is correct
try:
    if z is not None:
        Q, g, c, binary_cost = quantum_optimizer.binary_representation(x_sol=z)
        print("Binary cost:", binary_cost, "classical cost:", classical_cost)
        if np.abs(binary_cost - classical_cost) < 0.01:
            print("Binary formulation is correct")
        else:
            print("Error in the binary formulation")
    else:
        print("Could not verify the correctness, due to CPLEX solution being unavailable.")
        Q, g, c, binary_cost = quantum_optimizer.binary_representation()
        print("Binary cost:", binary_cost)
except NameError as e:
    print("Warning: Please run the cells above first.")
    print(e)

qp = quantum_optimizer.construct_problem(Q, g, c)
print(qp)

#quantum_solution, quantum_cost = quantum_optimizer.solve_problem(qp)
quantum_solution, quantum_cost = quantum_optimizer.solve_problem(qp)

print(quantum_solution, quantum_cost)

print(classical_cost)

m = folium.Map(location=[39.487660, -97.594333], zoom_start=0)
marker_icon1 = folium.Icon(color = "red")
for i in range(len(lat)):
    if (i == 0):
        folium.Marker(location=[lat[i], lng[i]], tooltip=f"Location: {loc[i]}, Order: {i}", icon=marker_icon1).add_to(m)
    else:
        folium.Marker(location=[lat[i], lng[i]], tooltip=f"Location: {loc[i]}, Order: {i}").add_to(m)

for ii in range(0, n**2):
    if x[ii] > 0:
        ix = ii // n
        iy = ii % n
        
        folium.PolyLine([(lat[ix], lng[ix]), (lat[iy], lng[iy])], color="blue").add_to(m)

m

print(quantum_cost)

x_quantum = np.zeros(n**2)
kk = 0

for ii in range(n**2):
    if ii // n != ii % n:
        x_quantum[ii] = quantum_solution[kk]
        kk += 1

m = folium.Map(location=[39.487660, -97.594333], zoom_start=0)
marker_icon1 = folium.Icon(color = "red")
for i in range(len(lat)):
    if (i == 0):
        folium.Marker(location=[lat[i], lng[i]], tooltip=f"Location: {loc[i]}, Order: {i}", icon=marker_icon1).add_to(m)
    else:
        folium.Marker(location=[lat[i], lng[i]], tooltip=f"Location: {loc[i]}, Order: {i}").add_to(m)

for ii in range(0, n**2):
    if x_quantum[ii] > 0:
        ix = ii // n
        iy = ii % n
        
        folium.PolyLine([(lat[ix], lng[ix]), (lat[iy], lng[iy])], color="blue").add_to(m)

m

algorithms = ("Classic", "Quantum")
data = {
    'K = 1': (2249.2068134000006, 1706.2245994000696),
    'k = 2': (2771.940853740001, 1845.127222779207),
    'K = 3': (3981.1556002800016, 3981.155600280501),
}

x = np.arange(len(algorithms))  # the label locations
width = 0.25  # the width of the bars
multiplier = 0

fig, ax = plt.subplots(layout='constrained')

for attribute, measurement in data.items():
    offset = width * multiplier
    rects = ax.bar(x + offset, measurement, width, label=attribute)
    ax.bar_label(rects, padding=3)
    multiplier += 1

# Add some text for labels, title and custom x-axis tick labels, etc.
ax.set_ylabel('Length (mm)')
ax.set_title('Comparision of Quantum and Classical Cost')
ax.set_xticks(x + width, algorithms)
ax.legend(loc='upper left', ncols=3)
ax.set_ylim(0, 5000)

plt.show()