import numpy as np
import networkx as nx
import matplotlib.pyplot as plt

class Lattice():
    def __init__(self, n):
        self.n = n
        self.graph = self._generate_lattice(self.n)

    def _generate_lattice(self, n):
        lx = range(-n ,n+1, 1)
        ly = range(-n ,n+1, 1)
        lattice = nx.grid_2d_graph(lx, ly)
        return lattice

    def generate_configuration(self, p):
        nb_edges = len(self.graph.edges())
        config = np.random.binomial(1, p, size=nb_edges)
        opening = {e: val for e, val in zip(self.graph.edges, config)}
        nx.set_edge_attributes(self.graph, values=opening, name='is_open')
        return config

    def draw(self, config):
        _pos = {(x,y):(x,y) for x,y in self.graph.nodes()}


        edge_colors = ['white' if d['is_open']==0 else 'black' for (u, v, d) in self.graph.edges(data=True)]
        options = {
            'pos': _pos,
            'width':1,
            'edge_color':edge_colors,
            'node_size': .1,
            'node_color': 'black',
            # 'linewidth': 0, 
        }
        fig = plt.figure(figsize=(30,30))
        nx.draw(
        G=self.graph,
        **options
        )
        plt.gca().set_aspect('equal')
        # plt.savefig('./figures/my_fig.png',dpi=1000)
        plt.show()


if __name__ == "__main__":
    lattice = Lattice(50)
    latt_conf = lattice.generate_configuration(.2)
    lattice.draw(latt_conf)