import sys
sys.path.append('..')

import socnet as sn

"""
Explanation: Encontro 02, Parte 4: Algoritmo de Dijkstra
Este guia foi escrito para ajudar você a atingir os seguintes objetivos:

implementar o algoritmo de Dijkstra;
praticar o uso da biblioteca da disciplina.

Primeiramente, vamos importar a biblioteca:
End of explanation
"""


sn.graph_width = 320
sn.graph_height = 180

"""
Explanation: A seguir, vamos configurar as propriedades visuais:
End of explanation
"""


g = sn.load_graph('4-dijkstra.gml', has_pos=True)

for n, m in g.edges():
    g.edge[n][m]['label'] = g.edge[n][m]['c']

sn.show_graph(g, elab=True)

"""
Explanation: Por fim, vamos carregar e visualizar um grafo:
End of explanation
"""


class Heap(object):
    def __init__(self, g):
        self.g = g
        self.h = []
        self.indices = {}

    def _parent(self, i):
        return (i - 1) // 2

    def _left(self, i):
        return 2 * i + 1

    def _right(self, i):
        return 2 * i + 2

    def _key(self, i):
        return self.g.node[self.h[i]]['d']

    def _swap(self, i, j):
        self.h[i], self.h[j] = self.h[j], self.h[i]
        self.indices[self.h[i]] = i
        self.indices[self.h[j]] = j

    def empty(self):
        return len(self.h) == 0

    def fix(self, n):
        i = self.indices[n]
        p = self._parent(i)
        while i > 0 and self._key(p) > self._key(i):
            self._swap(i, p)
            i = p
            p = self._parent(i)

    def put(self, n):
        self.indices[n] = len(self.h)
        self.h.append(n)
        self.fix(n)

    def get(self):
        n = self.h[0]
        self._swap(0, len(self.h) - 1)
        del self.h[-1]
        del self.indices[n]
        i = 0
        while True:
            l = self._left(i)
            r = self._right(i)
            c = i
            if l < len(self.h) and self._key(l) < self._key(c):
                c = l
            if r < len(self.h) and self._key(r) < self._key(c):
                c = r
            if i == c:
                break
            self._swap(i, c)
            i = c
        return n

"""
Explanation: Algoritmo de Dijkstra
A ideia do algoritmo de busca em largura é usar uma fila para visitar os nós em ordem de distância. A ideia do algoritmo de Bellman-Ford é atualizar a distância ponderada de um nó quando descobrimos uma aresta através da qual podemos melhorá-la. A busca em largura é rápida, mas não considera custos. O Bellman-Ford considera custos, mas não é rápido.
O algoritmo de Dijkstra usa um pouco de cada algoritmo para eficientemente calcular distâncias ponderadas no caso particular em que os custos são todos não-negativos. A ideia é usar uma fila de prioridade para visitar os nós em ordem de distância e também atualizar a distância ponderada de um nó não-visitado quando descobrimos uma aresta através da qual podemos melhorá-la. No pseudocódigo abaixo, "x.d" significa "atributo d do nó x".
inicialize uma fila de prioridade h, onde menor d significa maior prioridade

para cada nó n
    n.d = ∞
    coloque n em h

s.d = 0
conserte s em h, pois s.d mudou

enquanto h não está vazia
    retire n de h, onde n é o nó com maior prioridade

    para cada sucessor m de n
        d = n.d + c(n, m)

        se m.d &gt; d
            m.d = d
            conserte m em h, pois m.d mudou

Vamos especificar uma classe que representa a fila de prioridade necessária. Não é necessário entender todos os detalhes dela, apenas que o método fix deve ser usado para consertar, o método put deve ser usado para colocar e o método get deve ser usado para retirar.
End of explanation
"""


from math import inf, isinf

def snapshot(g, frames):
    for n in g.nodes():
        if isinf(g.node[n]['d']):
            g.node[n]['label'] = '∞'
        else:
            g.node[n]['label'] = str(g.node[n]['d'])

    frame = sn.generate_frame(g, nlab=True)

    frames.append(frame)

    
red = (255, 0, 0)  
blue = (0, 0, 255)
green = (0,255,0)
frames = [] 

sn.reset_node_colors(g)
sn.reset_edge_colors(g)
snapshot(g, frames)

h = Heap(g)
for n in g.nodes():
    g.node[n]['d'] = inf
    h.put(n)

s=0
g.node[s]['d'] = 0
h.fix(0)


while not h.empty():
    n = h.get()
    g.node[n]['color'] = green
    snapshot (g, frames)
    
    for m in g.successors(n):
        g.edge[n][m]['color'] = red
        d = g.node[n]['d'] + g.edge[n][m]['c']
        snapshot(g, frames)
        if g.node[m]['d'] > d:
            g.node[m]['d'] = d
            h.fix(m)
        g.edge[n][m]['color'] = sn.edge_color
        snapshot(g, frames)
    
    g.node[n]['color'] = blue
    snapshot(g, frames)


snapshot(g, frames)

sn.reset_edge_colors(g)
sn.reset_node_colors(g)
snapshot(g, frames)


sn.show_animation(frames)

"""
Explanation: Exercício
Monte uma visualização do algoritmo de Dijkstra. Use a classe Heap.
End of explanation
"""


