# %pip install qiskit 
# %pip install docplex
# %pip install qiskit_optimization
# %pip install networkx
# %pip install geopandas
# %pip install folium

import networkx as nx
import geopandas as gpd
import folium

import numpy as np 
import pandas as pd
import matplotlib.pyplot as plt
import itertools
import time
import random
from IPython.display import display

from qiskit import IBMQ
from qiskit.algorithms import QAOA, VQE
from qiskit.circuit.library import TwoLocal
from qiskit.algorithms.optimizers import COBYLA, SPSA # Classical simulator
from qiskit.providers.aer import AerSimulator

from docplex.mp.model import Model
from qiskit_optimization.runtime import QAOAClient, VQEClient
from qiskit_optimization.converters import QuadraticProgramToQubo 
from qiskit_optimization.translators import from_docplex_mp
from qiskit_optimization.algorithms import CplexOptimizer, MinimumEigenOptimizer

path = gpd.datasets.get_path('nybb')
df = gpd.read_file(path)

m = folium.Map(location=[40.70, -73.94], zoom_start=10, max_zoom=12, tiles='CartoDB positron')
    
# Project to NAD83 projected crs
df = df.to_crs(epsg=2263)

# Access the centroid attribute of each polygon
df['centroid'] = df.centroid

# Project to WGS84 geographic crs

# geometry (active) column
df = df.to_crs(epsg=4326)

# Centroid column
df['centroid'] = df['centroid'].to_crs(epsg=4326)

np.random.seed(7)
locations = []
ids = 0
for _, r in df.iterrows():
    lat, lon = r['centroid'].y, r['centroid'].x
    for i in range(2):
        lat_rand, lon_rand = lat + 0.2 * np.random.rand(), lon +0.1 * np.random.rand()
        locations.append((lon_rand, lat_rand))
        folium.Marker(location=[lat_rand, lon_rand], popup=f'Id: {ids}').add_to(m)
        ids += 1
center = np.array(locations).mean(axis=0)
locations = [(center[0], center[1])] + locations
folium.CircleMarker(location=[center[1], center[0]], radius=10, popup="<stong>Warehouse</stong>",
                    color="red",fill=True, fillOpacity=1, fillColor="tab:red").add_to(m)
m

# Normalizing the results to produce a graph in Graphx
companies = np.array(locations)
companies -= companies[0]
companies /= (np.max(np.abs(companies), axis=0))
r = list(np.sqrt(np.sum(companies ** 2, axis=1)))

threshold = 1 # Limit for to not consider the route from company i to j if the distance is larger than a threshold
n_companies = len(companies)
G = nx.Graph(name="VRP")
G.add_nodes_from(range(n_companies))
# G.add_weighted_edges_from([0, i, w] for i, w in zip(range(1, n_companies), r[1:]))
np.random.seed(2)
count = 0
for i in range(n_companies):
    for j in range(n_companies):
        if i != j:
            rij = np.sqrt(np.sum((companies[i] - companies[j])**2))
            if (rij < threshold) or (0 in [i, j]):
                count +=1
                G.add_weighted_edges_from([[i, j, rij]])
                r.append(rij)
colors = [plt.cm.get_cmap("coolwarm")(x) for x in r[1:]]
nx.draw(G, pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])
print(f"The number of edges of this problem is: {len(G.edges)}")

mdl = Model(name="VRP")
n_trucks = 3 # number of K trucks
x = {}
for i, j in G.edges():
    x[(i, j)] = mdl.binary_var(name=f"x_{i}_{j}") # Adding route from company i to company j as a binary variable
    x[(j, i)] = mdl.binary_var(name=f"x_{j}_{i}") # Adding route from company j to company i as a binary variable

print(f"The number of qubits needed to solve the problem is: {mdl.number_of_binary_variables}")
cost_func = mdl.sum(w["weight"] * x[(i, j)] for i, j, w in G.edges(data=True)) + mdl.sum(w["weight"] * x[(j, i)] for i, j, w in G.edges(data=True))
mdl.minimize(cost_func)

# Constraint 1a(yellow Fig. above): Only one truck goes out from company i 
for i in range(1, n_companies):
    mdl.add_constraint(mdl.sum(x[i, j] for j in range(n_companies) if (i, j) in x.keys()) == 1)
# Constraint 1b (yellow Fig. above): Only one truck comes into company j 
for j in range(1, n_companies):
    mdl.add_constraint(mdl.sum(x[i, j] for i in range(n_companies) if (i, j) in x.keys()) == 1)
# Constraint 2: (orange Fig. above) For the warehouse
mdl.add_constraint(mdl.sum(x[i, 0] for i in range(1, n_companies)) == n_trucks)
mdl.add_constraint(mdl.sum(x[0, j] for j in range(1, n_companies)) == n_trucks)

# Constraint 3: (blue Fig. above) To eliminate sub-routes
companies_list = list(range(1, n_companies))
subroute_set = []
for i in range(2, len(companies_list) + 1):
    for comb in itertools.combinations(companies_list, i):
        subroute_set.append(list(comb)) #subset points

for subroute in subroute_set:
    constraint_3 = []
    for i, j in itertools.permutations(subroute, 2): #iterating over all the subset points
        if (i, j) in x.keys():
            constraint_3.append(x[(i,j)])
        elif i == j:
            pass
        else:
            constraint_3 = []
            break
    if len(constraint_3) != 0:
        mdl.add_constraint(mdl.sum(constraint_3) <= len(subroute) - 1)

quadratic_program = from_docplex_mp(mdl)
print(quadratic_program.export_as_lp_string())

sol = CplexOptimizer().solve(quadratic_program)
print(sol.prettyprint())

solution_cplex = sol.raw_results.as_name_dict()
G_sol = nx.Graph()
G_sol.add_nodes_from(range(n_companies))
for i in solution_cplex:
    nodes = i[2:].split("_")
    G_sol.add_edge(int(nodes[0]), int(nodes[1]))
nx.draw(G_sol, pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])

qubo = QuadraticProgramToQubo(penalty=15).convert(quadratic_program)
num_vars = qubo.get_num_binary_vars()
print(f"To represent the inital problem with {mdl.number_of_binary_variables} variables, the QUBO representation needs {num_vars} variables")

new_qubo = qubo.substitute_variables(sol.variables_dict)

start = time.time()
sol_slack = CplexOptimizer().solve(new_qubo)
end = time.time() - start
print(f"The solver needs {np.round(end, 3)} seconds to finish.")

qubo_no_slack = qubo.substitute_variables(sol_slack.variables_dict)
sol_no_slack = CplexOptimizer().solve(qubo_no_slack)

solution_slack = sol_no_slack.raw_results.as_name_dict()
G_sol_slack = nx.Graph()
G_sol_slack.add_nodes_from(range(n_companies))
for i in solution_slack:
    nodes = i[2:].split("_")
    G_sol.add_edge(int(nodes[0]), int(nodes[1]))
nx.draw(G_sol, pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])

index_1s = [k for k, v in sol_no_slack.variables_dict.items() if v == 1]
index_0s = [k for k, v in sol_no_slack.variables_dict.items() if v == 0]

def Optimization_QAOA(qubo, reps=1, optimizer=COBYLA(maxiter=1000), backend=None,
                      shots=1024, alpha=1.0, provider=None, local=False, error_mitigation=False):
    intermediate_info = {'nfev': [],
                         'parameters': [],
                         'stddev': [],
                         'mean': []
                             }
    
    def callback(nfev, parameters, mean, stddev):
#         display(f"Evaluation: {nfev}, Energy: {mean}, Std: {stddev}")
        intermediate_info['nfev'].append(nfev)
        intermediate_info['parameters'].append(parameters)
        intermediate_info['mean'].append(mean)
        intermediate_info['stddev'].append(stddev)
    
    if local:
        qaoa_mes = QAOA(optimizer=optimizer, reps=reps, quantum_instance=AerSimulator(),
                        callback=callback)
    else:
        qaoa_mes = QAOAClient(provider=provider, backend=backend, reps=reps, alpha=alpha,
                             shots=shots, callback=callback, optimizer=optimizer,
                             optimization_level=3,measurement_error_mitigation=error_mitigation)
    qaoa = MinimumEigenOptimizer(qaoa_mes)
    result = qaoa.solve(qubo)
    return result, intermediate_info

def Optimization_VQE(qubo, ansatz, optimizer=SPSA(maxiter=50), backend=None,
                     shots=1024, provider=None, local=False, error_mitigation=False):

    intermediate_info = {'nfev': [],
                         'parameters': [],
                         'stddev': [],
                         'mean': []
                         }

    def callback(nfev, parameters, mean, stddev):
#         display(f"Evaluation: {nfev}, Energy: {mean}, Std: {stddev}")
        intermediate_info['nfev'].append(nfev)
        intermediate_info['parameters'].append(parameters)
        intermediate_info['mean'].append(mean)
        intermediate_info['stddev'].append(stddev)
        
    if local:
        vqe_mes = VQE(ansatz=ansatz, quantum_instance=AerSimulator(),
                        callback=callback, optimizer=optimizer)
    else:
        vqe_mes = VQEClient(ansatz=ansatz, provider=provider, backend=backend, shots=shots,
                        callback=callback, optimizer=optimizer,measurement_error_mitigation=error_mitigation)
    vqe = MinimumEigenOptimizer(vqe_mes)
    result = vqe.solve(qubo)
    return result, intermediate_info

def graph(solution, solution_not_kept):
    G = nx.Graph()
    G.add_nodes_from(range(n_companies))
    for k, v in solution.items():
        if v == 1:
            nodes = k[2:].split("_")
            G.add_edge(int(nodes[0]), int(nodes[1]))
    for k, v in solution_not_kept.items():
        if v == 1:
            nodes = k[2:].split("_")
            G.add_edge(int(nodes[0]), int(nodes[1]))
    return G

num_1s = 3
num_0s = 4
random.seed(1)
keep_1s = random.sample(index_1s, num_1s)
keep_0s = random.sample(index_0s, num_0s)
sol_qubo = sol.variables_dict
solution_not_kept7 = {i:sol_qubo[i] for i in sol_qubo.keys() if i not in keep_1s + keep_0s} 

qubo_7vars = qubo_no_slack.substitute_variables(solution_not_kept7)

print(qubo_7vars.export_as_lp_string())

sol7_local_qaoa_cobyla = Optimization_QAOA(qubo_7vars, reps=2, optimizer=COBYLA(maxiter=100), local=True)
sol7_local_qaoa_spsa = Optimization_QAOA(qubo_7vars, reps=2, optimizer=SPSA(maxiter=100), local=True)

IBMQ.load_account()
provider = IBMQ.get_provider(hub='ibm-q-community', group='cdl-hackathon', project='main')
provider2 = IBMQ.get_provider(hub='ibm-q-research', group='guanajuato-1', project='main')

sol7_oslo_qaoa_cobyla = Optimization_QAOA(qubo_7vars, reps=2, local=False, backend=provider2.backend.ibm_oslo,
                                       provider=provider2, optimizer=COBYLA(maxiter=50))

sol7_oslo_qaoa_spsa = Optimization_QAOA(qubo_7vars, reps=2, local=False, backend=provider2.backend.ibm_oslo,
                                       provider=provider2, optimizer=SPSA(maxiter=50))

sol7_oslo_qaoa_spsa_mitig = Optimization_QAOA(qubo_7vars, reps=2, local=False, backend=provider2.backend.ibm_oslo,
                                       provider=provider2, optimizer=SPSA(maxiter=50), error_mitigation=True)

ansatz = TwoLocal(qubo_7vars.get_num_vars(), rotation_blocks='ry', entanglement_blocks='cz')

sol7_local_vqe_spsa = Optimization_VQE(qubo_7vars, ansatz, local=True, optimizer=SPSA(maxiter=200))
sol7_local_vqe_cobyla = Optimization_VQE(qubo_7vars, ansatz, local=True, optimizer=COBYLA(maxiter=200))

sol7_lagos_vqe_cobyla = Optimization_VQE(qubo_7vars, ansatz, local=False, optimizer=COBYLA(maxiter=100),
                                 backend=provider2.backend.ibm_lagos, provider=provider2)

sol7_lagos_vqe_spsa = Optimization_VQE(qubo_7vars, ansatz, local=False, optimizer=SPSA(maxiter=25),
                                 backend=provider2.backend.ibm_lagos, provider=provider2)

sol7_lagos_vqe_spsa2 = Optimization_VQE(qubo_7vars, ansatz, local=False, optimizer=SPSA(maxiter=25),
                                 backend=provider2.backend.ibm_lagos, provider=provider2, error_mitigation=True)

plt.figure()
plt.plot(sol7_lagos_vqe_spsa[1]["mean"], linestyle=":", color="seagreen", label="VQE-Lagos-SPSA")
plt.plot(sol7_lagos_vqe_spsa2[1]["mean"], linestyle="-.", color="indigo", label="VQE-Lagos-SPSA-Mitig")
plt.plot(sol7_lagos_vqe_cobyla[1]["mean"], color="olive", label="VQE-Lagos-COBYLA-Mitig")
plt.grid()
plt.savefig("./Images/inset.png")

plt.figure()
plt.plot(sol7_oslo_qaoa_spsa[1]["mean"], linestyle=":", color="seagreen", label="QAOA-Oslo-SPSA")
plt.plot(sol7_oslo_qaoa_cobyla[1]["mean"], color="olive", label="QAOA-Oslo-COBYLA")
plt.grid()
plt.savefig("./Images/inset2.png")

sol7 = {"local_qaoa_cobyla":sol7_local_qaoa_cobyla, "local_vqe_cobyla":sol7_local_vqe_cobyla,
       "local_qaoa_spsa":sol7_local_qaoa_spsa, "local_vqe_spsa":sol7_local_vqe_spsa,
       "oslo_qaoa_cobyla":sol7_oslo_qaoa_cobyla, "oslo_qaoa_spsa":sol7_oslo_qaoa_spsa,
       "lagos_vqe_spsa":sol7_lagos_vqe_spsa, "lagos_vqe_cobyla_mitig":sol7_lagos_vqe_cobyla}
np.save("./Data/sol7.npy", sol7)

sol7 = [sol7_local_qaoa_cobyla, sol7_local_vqe_cobyla, sol7_local_qaoa_spsa, sol7_local_vqe_spsa]
q_alg = ["QAOA", "VQE"]
fig, ax = plt.subplots(3,2, figsize=(15,18))
ax[0,0].plot(sol7_local_qaoa_cobyla[1]["mean"], color="darkblue", label="QAOA-COBYLA")
ax[0,0].plot(sol7_local_qaoa_spsa[1]["mean"], linestyle="--", color="lightcoral", label="QAOA-SPSA")
ax[0,0].plot(sol7_oslo_qaoa_spsa[1]["mean"], linestyle=":", color="seagreen", label="QAOA-Oslo-SPSA")
ax[0,0].plot(sol7_oslo_qaoa_cobyla[1]["mean"], color="olive", label="QAOA-Oslo-COBYLA")

ax[0,1].plot(sol7_local_vqe_cobyla[1]["mean"], color="darkblue", label="VQE-COBYLA")
ax[0,1].plot(sol7_local_vqe_spsa[1]["mean"], linestyle="--", color="lightcoral", label="VQE-SPSA")
ax[0,1].plot(sol7_lagos_vqe_spsa[1]["mean"], linestyle=":", color="seagreen", label="VQE-Lagos-SPSA")
ax[0,1].plot(sol7_lagos_vqe_spsa2[1]["mean"], linestyle="-.", color="indigo", label="VQE-Lagos-SPSA-Mitig")
ax[0,1].plot(sol7_lagos_vqe_cobyla[1]["mean"], color="olive", label="VQE-Lagos-COBYLA-Mitig")

for i in range(2):
    ax[0,i].set_xlabel("Iterations", fontsize=14)
    ax[0,i].set_ylabel("Cost", fontsize=14)
    ax[0,i].legend()
    ax[0,i].grid()
    plt.sca(ax[1,i])
    nx.draw(graph(sol7[i][0].variables_dict,solution_not_kept7), pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])
    ax[1,i].set_title(f"{q_alg[i]}-COBYLA", fontsize=14)
    plt.sca(ax[2,i])
    nx.draw(graph(sol7[i+2][0].variables_dict,solution_not_kept7), pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])
    ax[2,i].set_title(f"{q_alg[i]}-SPSA", fontsize=14)
    
plt.savefig("./Images/Sol7Q.png")

num_1s = 5
num_0s = 10
random.seed(1)
keep_1s = random.sample(index_1s, num_1s)
keep_0s = random.sample(index_0s, num_0s)
sol_qubo = sol.variables_dict
solution_not_kept = {i:sol_qubo[i] for i in sol_qubo.keys() if i not in keep_1s + keep_0s} 

qubo_15vars = qubo_no_slack.substitute_variables(solution_not_kept)

sol15_qasm_qaoa_cobyla = Optimization_QAOA(qubo_15vars, reps=3, local=False, backend=provider2.backend.ibmq_qasm_simulator,
                                       provider=provider2, optimizer=COBYLA(maxiter=100))

solution15_qaoa =  sol15_qasm_qaoa_cobyla[0].variables_dict
G_sol15_qaoa = nx.Graph()
G_sol15_qaoa.add_nodes_from(range(n_companies))
for k, v in solution15_qaoa.items():
    if v == 1:
        nodes = k[2:].split("_")
        G_sol15_qaoa.add_edge(int(nodes[0]), int(nodes[1]))
for k, v in solution_not_kept.items():
    if v == 1:
        nodes = k[2:].split("_")
        G_sol15_qaoa.add_edge(int(nodes[0]), int(nodes[1]))
nx.draw(G_sol15_qaoa, pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])

sol15_qasm_qaoa_spsa = Optimization_QAOA(qubo_15vars, reps=3, local=False, backend=provider2.backend.ibmq_qasm_simulator,
                                       provider=provider2, optimizer=SPSA(maxiter=25))

ansatz15 = TwoLocal(qubo_15vars.get_num_vars(), rotation_blocks='ry', entanglement_blocks='cz')
sol15_qasm_vqe_cobyla = Optimization_VQE(qubo_15vars,ansatz=ansatz15, local=False, backend=provider2.backend.ibmq_qasm_simulator,
                                       provider=provider2, optimizer=COBYLA(maxiter=100))

solution15_vqe =  sol15_qasm_vqe_cobyla[0].variables_dict
G_sol15_vqe = nx.Graph()
G_sol15_vqe.add_nodes_from(range(n_companies))
for k, v in solution15_vqe.items():
    if v == 1:
        nodes = k[2:].split("_")
        G_sol15_vqe.add_edge(int(nodes[0]), int(nodes[1]))
for k, v in solution_not_kept.items():
    if v == 1:
        nodes = k[2:].split("_")
        G_sol15_vqe.add_edge(int(nodes[0]), int(nodes[1]))
nx.draw(G_sol15_vqe, pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])

sol15_qasm = {"qaoa_cobyla":sol15_qasm_qaoa_cobyla, "vqe_cobyla":sol15_qasm_vqe_cobyla}
np.save("./Data/sol15_qasm.npy", sol15_qasm)

solution15 = [G_sol15_qaoa, G_sol15_vqe]
fig, ax = plt.subplots(2,2, figsize=(15,10))
ax[0,0].plot(sol15_qasm_qaoa_cobyla[1]["mean"], color="darkblue", label="QAOA-COBYLA")
ax[0,1].plot(sol15_qasm_vqe_cobyla[1]["mean"], color="slateblue", label="VQE-COBYLA")
for i in range(2):
    ax[0,i].grid()
    ax[0,i].set_xlabel("Iterations", fontsize=14)
    ax[0,i].set_ylabel("Cost", fontsize=14)
    ax[0,i].legend()
    plt.sca(ax[1,i])
    nx.draw(solution15[i], pos=companies, with_labels=True, node_size=500,
        edge_color=colors, width=1, font_color="white",font_size=14,
        node_color = ["tab:red"] + (n_companies-1)*["darkblue"])
fig.savefig("./Images/Sol15Q.png")

num_1s = 8
num_0s = 12
random.seed(1)
keep_1s = random.sample(index_1s, num_1s)
keep_0s = random.sample(index_0s, num_0s)
sol_qubo = sol.variables_dict
solution_not_kept = {i:sol_qubo[i] for i in sol_qubo.keys() if i not in keep_1s + keep_0s} 

qubo_20vars = qubo_no_slack.substitute_variables(solution_not_kept)

sol20_qasm_qaoa_cobyla_r1 = Optimization_QAOA(qubo_20vars, reps=1, local=False, backend=provider2.backend.ibmq_qasm_simulator,
                                       provider=provider2, optimizer=COBYLA(maxiter=10))

