import networkx as nx
from matplotlib import pyplot as plt
from island.match import Match
from island.matches import Matches


class Draw:
    def __init__(self, show):
        self.matches = Matches.from_profile_expr(lambda r: True)
        self.show = show

    def generate_graph(self, match):
        g = nx.empty_graph()
        for r in match.query('player', 'join').raw_data:
            g.add_node(r['pid'])
        for r in match.query('action', 'done').raw_data:
            if r['a'] > r['b']:
                f, t = r['b'], r['a']
            else:
                f, t = r['a'], r['b']
            if g.has_edge(f, t):
                g.edges[f, t]['weight'] += 1
            else:
                g.add_edge(f, t, weight=1, t=1)

        for r in match.query('neighbor', 'create').raw_data:
            if r['a'] > r['b']:
                f, t = r['b'], r['a']
            else:
                f, t = r['a'], r['b']
            if g.has_edge(f, t):
                g.edges[f, t]['t'] |= 2
            else:
                g.add_edge(f, t, weight=1, t=2)
        return g

    def draw_graph(self, g, m):
        social_edges = []
        game_edges = []
        weight = []
        for e in g.edges():
            if g[e[0]][e[1]]['t'] & 1 == 1:
                game_edges.append(e)
                weight.append(g[e[0]][e[1]]['weight'])
            else:
                social_edges.append(e)

        pos = nx.spring_layout(g, scale=0.3)
        nx.draw_networkx_edges(g, edgelist=social_edges, pos=pos, edge_color='gray', alpha=0.8, style='dashed')
        nx.draw_networkx_edges(g, edgelist=game_edges, pos=pos, width=weight, alpha=0.8, edge_color='#d63031')
        nx.draw_networkx_nodes(g, pos=pos, node_color='#0984e3', node_size=300)

        if self.show:
            plt.show()
        else:
            plt.savefig("graph/%s-cluster.eps"%m.name)


    def draw_a_graph(self, name):
        for m in self.matches.data:
            if m.name == name:
                g = self.generate_graph(m)
                self.draw_graph(g, m)

if __name__ == '__main__':
    d = Draw(True)
    d.draw_a_graph('G398')