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

import socnet as sn

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

implementar o algoritmo de Kruskal;
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('5-kruskal.gml', has_pos=True)

for e in g.edges_iter():
    g.edge[e[0]][e[1]]['label'] = g.edge[e[0]][e[1]]['c']

sn.show_graph(g, elab=True)

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


class Forest(object):
    def __init__(self, g):
        self.g = g
        self.f = set()
        for n in g.nodes():
            self._make_set(n)

    def _make_set(self, x):
        g.node[x]['p'] = x
        g.node[x]['rank'] = 0

    def _union(self, x, y):
        self._link(self._find_set(x), self._find_set(y))

    def _link(self, x, y):
        if g.node[x]['rank'] > g.node[y]['rank']:
            g.node[y]['p'] = x
        else:
            g.node[x]['p'] = y
            if g.node[x]['rank'] == g.node[y]['rank']:
                g.node[y]['rank'] = g.node[y]['rank'] + 1

    def _find_set(self, x):
        if x != g.node[x]['p']:
            g.node[x]['p'] = self._find_set(g.node[x]['p'])
        return g.node[x]['p']

    def adding_does_not_form_circuit(self, n, m):
        return self._find_set(n) != self._find_set(m)

    def add(self, n, m):
        self.f.add((n, m))
        self._union(n, m)

"""
Explanation: Árvores geradoras mínimas
Dizemos que:
* um passeio $\langle n_0, n_1, \ldots, n_{k-1} \rangle$ é um circuito se $\langle n_0, n_1, \ldots, n_{k-2} \rangle$ é um caminho e $n_0 = n_{k-1}$;
* um conjunto de arestas $F$ é uma floresta se não existem circuitos no grafo $(N, F)$;
* um grafo é conexo se para quaisquer nós $s$ e $t$ existe um caminho de $s$ a $t$;
* uma floresta $T$ é uma árvore geradora se o grafo $(N, T)$ é conexo.
O custo de uma árvore geradora $T$ é
$\sum_{{n, m} \in T} c(n, m)$.
Uma árvore geradora é mínima se não existe outra árvore geradora de custo menor. Note que podem existir múltiplas árvores geradoras mínimas.
Algoritmo de Kruskal
Podemos eficientemente obter uma árvore geradora mínima usando o algoritmo de Kruskal. A ideia desse algoritmo é simples: inicializamos uma floresta $F$ como o conjunto vazio e verificamos todas as arestas em ordem não-decrescente de custo. Para cada aresta, adicionamos ela a $F$ se essa adição não formar circuito no grafo $(N, F)$.
Vamos especificar uma classe que representa a floresta. Não é necessário entender todos os detalhes dela, apenas que o atributo f é o conjunto das arestas e os dois últimos métodos são auto-explicativos.
End of explanation
"""


from math import inf, isinf

def snapshot(g, frames):
    frame = sn.generate_frame(g, nlab=False, elab=True)
    frames.append(frame)

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

f = Forest (g)
edges = []

e = g.edges_iter()

for i in e:
    edges.append((i[0],i[1],g.get_edge_data(i[0],i[1])['c']))

edges.sort(reverse = True, key=lambda x: (-x[2],x[0]))

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

for n,m,c in edges:
    g.edge[m][n]['color'] = green
    snapshot(g, frames)
    
    if(f.adding_does_not_form_circuit(n,m)):
        g.edge[m][n]['color'] = blue
        snapshot(g, frames)  
        f.add(n,m)
    else:
        g.edge[m][n]['color'] = sn.edge_color
        snapshot(g, frames)


sn.show_animation(frames)

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


