#importa todas as funções do PuLP
from pulp import *


import networkx as nx
import numpy as np
import random

#criação da instância

#conj das cidades
cidades = ["A", "B", "C", "D", "E", "F"]

#matriz de custos
matriz_custos = [[0,3,5,19,2,4],
                 [3,0,14,15,11,7],
                 [5,14,0,20,1,9],
                 [19,15,20,0,9,10],
                 [2,11,1,9,0,13],
                 [4,7,9,10,13,0]]

#correspondência entre os valores do vetor
#cidades e a matriz de custos
custos = makeDict([cidades, cidades],
                  matriz_custos)

n = len(cidades)

#estrutura que ignora entrada (i,i)
arcos = [(i,j) for i in cidades 
         for j in cidades if i!=j]

#fim de criação da instância

#início da criação do modelo

#variavel que guarda a formulação para o problema que queremos resolver
#recebe o nome do problems e o objetivo
formulacaoMTZ = LpProblem("Formulacao MTZ para TSP", LpMinimize)

#criar as variáveis
#LpVariable.dicts(nome variavel, conj. onde esta definida, limiteinferior, limite superior, tipo)
x = LpVariable.dicts("x", (cidades, cidades), 0, 1, LpInteger)

u = LpVariable.dicts("u", cidades, 1, n, LpContinuous)


#adicionar as restrições
#----a primeira a adicionar é a FO
formulacaoMTZ += (lpSum([x[i][j]*custos[i][j] for (i,j) in arcos]),
                  "FO",
                  )

#----adicionar as restrições
for i in cidades:
    formulacaoMTZ += (
        lpSum([x[i][j] for j in cidades if (i,j) in arcos]) == 1,
        f"sai_{i}",
        )
    
    formulacaoMTZ += (
        lpSum([x[j][i] for j in cidades if (j,i) in arcos]) == 1,
        f"entra_{i}",
        )
    
for (i,j) in arcos:
    if i != cidades[0] and j != cidades[0]: 
        #assumindo que o depósito está na posição 0 do vetor cidades
        formulacaoMTZ += (
            u[i] + x[i][j] <= u[j] + (n - 2)*(1-x[i][j]),
            f"posicao_{i}_{j}",
            )


#extrair o modelo para um modelo para um ficheiro
formulacaoMTZ.writeLP("MTZ.lp")

#resolve o modelo
formulacaoMTZ.solve()

#estado da resolução
print("Estado:", LpStatus[formulacaoMTZ.status])

#valor otimo
print("Custo total = ", value(formulacaoMTZ.objective))

#escrever a solução
for v in formulacaoMTZ.variables():
    if v.varValue > 0:
        print(v.name, "=", v.varValue)


# Create a graph representing the TSP tour
G = nx.DiGraph()
G.add_nodes_from(cidades)
 
graph_edges = [(i,j) for (i, j) in arcos if x[i][j].varValue == 1]
G.add_edges_from(graph_edges)
 
# Create a dictionary to store the edge values (weights) based on x[i][j].varValue
edge_values = {(i, j): x[i][j].varValue for (i, j) in arcos if x[i][j].varValue == 1}
 
 
# Draw the graph with nodes and edges
pos = nx.spring_layout(G)  # Positioning of nodes
nx.draw(G, pos, with_labels=True, node_color="red", node_size=2500)
 
# Draw the edge labels using the x[i][j].varValue values
nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_values, font_color='black')
 
#output diferente