import numpy as np

from collections import deque
from functools import partials

class WollMonteCarlo:
    """ Wolff Monte Carlo simulator for the Ising model """
    def __init__(self, L, T, method = None):
        self._L = L
        self._T = T
        self._K = 2. / T
        self._method = method

        # set up the initial state
        self._state = np.random.randint(0, 2, size = [L, L])

    @property
    def L(self):
        return self._L

    @property
    def T(self):
        return self._T
    
    @property
    def K(self):
        return self._K

    @property
    def state(self):
        return self._state

    def probability_add_bonds(self, x1, y1, x2, y2, state):
        """ The probability for adding a bond. """
        E = 1.0 if state[x1, y1] == state[x2, y2] else 0.0
        return 1.0 - np.exp(-self.K * E)

    def wolff_iterative(self, state):
        """ Iterative Wolff algorithm.
            This algorithm uses a doubly ended queue (deque), which provides O(1) operations for adding and removing items from both ends of a list.
        """

        # Convenient lists for indexing
        # Below, we'll use these to get the left (or 'above', etc) neighbors of a site.
        # Includes periodic boundaries! Another option would be to replace
        #    left[x1]   by   (x1 - 1) % L
        # but that does an addition and a module every step.
        left = [self.L - 1] + list(range(self.L - 1))
        right = list(range(1, self.L)) + [0]

        # Book-keeping containers
        sites_to_consider = deque()
        sites_to_flip = set()
        bonds_considered = set()

        # Initial queue of sites to consider, just consisting of a single (x,y) location
        sites_to_consider.append((
            np.random.randint(0, self.L),
            np.random.randint(0, self.L)
        ))

        # As long as there are sites to consider
        while sites_to_consider:
            # Pick a new site to consider from the queue, either using
            # breadth first or depth first
            if self._method == "BFS":
                x1, y1 = sites_to_consider.popleft()
            if self._method == "DFS":
                x1, y1 = sites_to_consider.pop()

            # For the neighbors of this site
            for x2, y2 in zip([left[x1], right[x1], x1, x1],
                              [y1, y1, left[y1], right[y1]]):
                
                # Check if we have not already considered this pair
                if not (x1, y1, x2, y2) in bonds_considered:
                    # Add the pair so that we don't flip it twice
                    bonds_considered.add((x1, y1, x2, y2))
                    bonds_considered.add((x2, y2, x1, y1))

                    if np.random.rand() < self.probability_add_bond(x1, y1, x2, y2, state):
                        sites_to_consider.append((x2, y2))
                        sites_to_flip.add((x1, y1))
                        sites_to_flip.add((x2, y2))
                        
        return sites_to_flip

    def step(self):
        """Use Wolff and perform update."""

        # Get a list of sites to flip...
        to_flip = self.wolff_iterative(self._state)

        # ...and flip them
        for (x, y) in to_flip:
            self._state[x, y] = 1 - self._state[x, y]

        # Return the list of the flipped sites
        return to_flip