# By orrp (http://www.wisdom.weizmann.ac.il/~orrp/).
# Draws the Margulis-Gabber-Galil construction.
# Click on nodes to highlight their neighbors and view the conductance
# of the resulting set.
import argparse

import networkx as nx
from matplotlib import pyplot as plt

GRAPH_COLOR = '#F1F1F1'
SET_COLOR = '#FF0000'
BOUNDARY_COLOR = '#FFB3B3'
NODE_SIZE = None  # Defaults to something nice.


def flatten(ll):
    """Flattens a list of lists into a 1-dimensional list."""
    return [x for l in ll for x in l]


class MGG:
    def __init__(self, n):
        self._n = n
        self._set = []

        # Construct MGG_n
        edges = [((x, y), v) for x in range(n)
                 for y in range(n) for v in self.get_neighbors((x, y))]
        self._graph = nx.Graph(edges)
        self.draw()

        # Listen for click
        self.cid = plt.gca().figure.canvas.mpl_connect('button_press_event', self)

    def set_set(self, new_set):
        self._set = new_set

    # The MGG construction

    def plus(self, a, b):
        return (a + b) % self._n

    def S(self, u):
        return u[0], self.plus(u[1], u[0])

    def S_inverse(self, u):
        return u[0], self.plus(u[1], -u[0])

    def T(self, u):
        return self.plus(u[0], u[1]), u[1]

    def T_inverse(self, u):
        return self.plus(u[0], -u[1]), u[1]

    def get_neighbors(self, u):
        """Returns list of all edges incident at u as a list of pairs"""
        nbs = [self.S(u), self.S_inverse(u), self.T(u), self.T_inverse(u)]
        nbs += [(u[0], self.plus(u[1], 1)), (u[0], self.plus(u[1], -1)),
                (self.plus(u[0], 1), u[1]), (self.plus(u[0], -1), u[1])]
        return nbs

    # Drawing

    def cross_edges(self, s, t):
        """Returns a list of edges that cross from s to t, where s and t are lists of nodes."""
        return [edge for edge in self._graph.edges()
                if (edge[0] in s and edge[1] in t) or (edge[0] in t and edge[1] in s)]

    def draw(self):
        # So that the graph is drawn as a grid
        pos = dict([(n, n) for n in self._graph.nodes()])
        # Draw rest of graph
        nx.draw_networkx(self._graph, node_color=GRAPH_COLOR, node_size=NODE_SIZE,
                         edge_color=GRAPH_COLOR, pos=pos, with_labels=False)

        # Draw set edges
        set_edges = self.cross_edges(self._set, self._set)
        nx.draw_networkx(self._graph, nodelist=self._set,
                         node_color=SET_COLOR, node_size=NODE_SIZE,
                         edgelist=set_edges, edge_color=SET_COLOR, pos=pos, with_labels=False)

        # Draw boundary edges
        set_complement = [u for u in self._graph.nodes() if u not in self._set]
        boundary_edges = self.cross_edges(self._set, set_complement)
        boundary_nodes = [u for u in flatten(
            boundary_edges) if u not in self._set]
        nx.draw_networkx(self._graph, nodelist=boundary_nodes,
                         node_color=BOUNDARY_COLOR, node_size=NODE_SIZE,
                         edgelist=boundary_edges, edge_color=BOUNDARY_COLOR,
                         pos=pos, with_labels=False)

        # Draw text
        text = "S="
        if self._set:
            text += "{" + str(self._set)[1:-1] + "}"
        else:
            text += "∅"
        set_density = min(len(self._set), len(set_complement))
        if set_density > 0:
            boundary_density = len(boundary_edges)
            text += "\nΦ(S)=" + str(boundary_density) + "/" + str(set_density) + \
                "≅" + str(boundary_density / float(set_density))
        plt.gca().set_title(text)

    # Handling user input

    def closest_node(self, point):
        return int(round(point[0])) % self._n, int(round(point[1])) % self._n

    def __call__(self, event):
        if event.inaxes != plt.gca().axes:
            return
        event_node = (self.closest_node((event.xdata, event.ydata)))
        if event_node in self._set:
            self._set = [u for u in self._set if u != event_node]
        else:
            self._set.append(event_node)

        self.draw()
        plt.gca().figure.canvas.draw()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Play around with the Margulis-Gabber-Galil construction")
    parser.add_argument(dest='n', type=int, metavar='n',
                        help='modulus on which the graph will be constructed (integer)')
    MGG(parser.parse_args().n)
    plt.show()