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)
        return config

    def draw(self, config):
        _pos = {(x,y):(x,y) for x,y in self.graph.nodes()}
        edge_colors = ['white' if config[i]==0 else 'black' for i,e in enumerate(self.graph.edges)]
        options = {
            'pos': _pos,
            'width':2,
            'edge_color':edge_colors,
            'node_size': 0
        }
        fig = plt.figure(figsize=(20,20))
        nx.draw(
        G=self.graph,
        **options
        )

        plt.savefig('./figures/my_fig.png',dpi=1000)
        plt.show()


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