#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import abc
from enum import Enum, IntEnum
import itertools


class Spin(IntEnum):
    UP = 0
    DOWN = 1


class SpinPairs(Enum):
    """The spin orbitals corresponding to a pair of spatial orbitals."""

    SAME = 1
    ALL = 0
    DIFF = -1


class HubbardLattice(metaclass=abc.ABCMeta):
    """Base class for a Hubbard model lattice.

    Subclasses must define the following properties:
        n_dofs (int): The number of degrees of freedom per site (and spin if
            applicable).
        n_sites (int): The number of sites in the lattice.
        spinless (bool): Whether or not the fermion has spin (False if so).
        edge_types (Tuple[Hashable, ...]): The types of edges that a term could
            correspond to. Examples include 'onsite', 'neighbor',
            'diagonal_neighbor', etc.
        onsite_edge_types (Sequence[Hashable]): The edge types that connect
            sites to themselves.

    And the following methods:
        site_pairs_iter(edge_type: Hashable) \\
            -> Iterable[Tuple[int, int]]: Iterable
            over pairs of sites corresponding to the given edge type.

    For 'spinful' lattices, the ``spin_indices`` ``0`` and ``1`` correspond to
    'up' and 'down', respectively.
    """

    @abc.abstractproperty
    def n_dofs(self):
        """The number of degrees of freedom per site
        (and spin if applicable)."""

    @abc.abstractproperty
    def n_sites(self):
        """The number of sites in the lattice."""

    @abc.abstractproperty
    def spinless(self):
        """Whether or not the fermion has spin (False if so)."""

    @abc.abstractproperty
    def edge_types(self):
        """The types of edges that a term could correspond to.

        Examples include 'onsite', 'neighbor', 'diagonal_neighbor', etc.
        """

    @abc.abstractproperty
    def onsite_edge_types(self):
        """The edge types that connect sites to themselves."""

    @abc.abstractmethod
    def site_pairs_iter(self, edge_type, ordered=True):
        """Iterable over pairs of sites corresponding to the given edge type."""

    # properties

    @property
    def n_spin_values(self):
        return 1 if self.spinless else 2

    @property
    def n_spin_orbitals_per_site(self):
        return self.n_dofs * self.n_spin_values

    @property
    def n_spin_orbitals(self):
        return self.n_sites * self.n_spin_orbitals_per_site

    # indexing

    def site_index_offset(self, site_index):
        return site_index * self.n_spin_orbitals_per_site

    def dof_index_offset(self, dof_index):
        return dof_index * self.n_spin_values

    def to_spin_orbital_index(self, site_index, dof_index, spin_index):
        """The index of the spin orbital."""
        return (self.site_index_offset(site_index) +
                self.dof_index_offset(dof_index) + spin_index)

    def from_spin_orbital_index(self, spin_orbital_index):
        site_index, offset = divmod(spin_orbital_index,
                                    self.n_spin_orbitals_per_site)
        dof_index, spin_index = divmod(offset, self.n_spin_values)
        return site_index, dof_index, spin_index

    # iteration

    @property
    def site_indices(self):
        return range(self.n_sites)

    @property
    def dof_indices(self):
        return range(self.n_dofs)

    def dof_pairs_iter(self, exclude_same=False):
        return ((a, b)
                for a in range(self.n_dofs)
                for b in range(a + exclude_same, self.n_dofs))

    @property
    def spin_indices(self):
        return range(self.n_spin_values)

    def spin_pairs_iter(self, spin_pairs=SpinPairs.ALL, ordered=True):
        if spin_pairs == SpinPairs.ALL:
            return (itertools.product(self.spin_indices, repeat=2)
                    if ordered else itertools.combinations_with_replacement(
                        self.spin_indices, 2))
        elif spin_pairs == SpinPairs.SAME:
            return ((s, s) for s in self.spin_indices)
        elif spin_pairs == SpinPairs.DIFF:
            return (itertools.permutations(self.spin_indices, 2) if ordered else
                    itertools.combinations(self.spin_indices, 2))
        raise ValueError(
            '{} not a valid SpinPairs specification.'.format(spin_pairs))

    # validation

    def validate_edge_type(self, edge_type):
        if edge_type not in self.edge_types:
            raise ValueError('{} not a valid edge type {}.'.format(
                edge_type, self.edge_types))

    def validate_dof(self, dof, length=None):
        if not (0 <= dof < self.n_dofs):
            raise ValueError('not (0 <= {} < n_dofs = {})'.format(
                dof, self.n_dofs))

    def validate_dofs(self, dofs, length=None):
        for dof in dofs:
            self.validate_dof(dof)
        if (length is not None) and (len(dofs) != length):
            raise ValueError('len({}) != {}'.format(dofs, length))


class HubbardSquareLattice(HubbardLattice):
    r"""A square lattice for a Hubbard model.

    Valid edge types are:
        * 'onsite'
        * 'horizontal_neighbor'
        * 'vertical_neighbor'
        * 'neighbor': union of 'horizontal_neighbor' and 'vertical_neighbor'
        * 'diagonal_neighbor'
    """

    def __init__(self,
                 x_dimension,
                 y_dimension,
                 n_dofs=1,
                 spinless=False,
                 periodic=True):
        """
        Args:
            x_dimension (int): The width of the grid.
            y_dimension (int): The height of the grid.
            n_dofs (int, optional): The number of degrees of freedom per site
                (and spin if applicable). Defaults is 1.
            periodic (bool, optional): If True, add periodic boundary
                conditions. Default is True.
            spinless (bool, optional): If True, return a spinless Fermi-Hubbard
                model. Default is False.
        """

        self.x_dimension = x_dimension
        self.y_dimension = y_dimension
        self._n_dofs = n_dofs
        self._spinless = spinless
        self.periodic = periodic

    @property
    def n_dofs(self):
        return self._n_dofs

    @property
    def spinless(self):
        return self._spinless

    @property
    def n_sites(self):
        return self.x_dimension * self.y_dimension

    @property
    def edge_types(self):
        return ('onsite', 'neighbor', 'diagonal_neighbor',
                'horizontal_neighbor', 'vertical_neighbor')

    @property
    def onsite_edge_types(self):
        return ('onsite',)

    def site_pairs_iter(self, edge_type, ordered=True):
        if edge_type == 'onsite':
            return ((i, i) for i in self.site_indices)
        elif edge_type == 'neighbor':
            return self.neighbors_iter(ordered)
        elif edge_type == 'horizontal_neighbor':
            return self.horizontal_neighbors_iter(ordered)
        elif edge_type == 'vertical_neighbor':
            return self.vertical_neighbors_iter(ordered)
        elif edge_type == 'diagonal_neighbor':
            return self.diagonal_neighbors_iter(ordered)
        raise ValueError('Edge type {} is not valid.'.format(edge_type))

    def __repr__(self):
        return '{}({})'.format(
            self.__class__.__name__, ', '.join(
                (('x_dimension={}'.format(self.x_dimension)),
                 ('y_dimension={}'.format(self.y_dimension)),
                 ('n_dofs={}'.format(self.n_dofs)), ('spinless={}'.format(
                     self.spinless)), ('periodic={}'.format(self.periodic)))))

    # site indexing

    def to_site_index(self, site):
        """The index of a site."""
        x, y = site
        return x + y * self.x_dimension

    def from_site_index(self, site_index):
        return divmod(site_index, self.x_dimension)[::-1]

    # neighbor counting and iteration

    def n_horizontal_neighbor_pairs(self, ordered=True):
        """Number of horizontally neighboring (unordered) pairs of sites."""
        n_horizontal_edges_per_y = (
            self.x_dimension - (self.x_dimension <= 2 or not self.periodic))
        return (self.y_dimension * n_horizontal_edges_per_y *
                (2 if ordered else 1))

    def n_vertical_neighbor_pairs(self, ordered=True):
        """Number of vertically neighboring (unordered) pairs of sites."""
        n_vertical_edges_per_x = (self.y_dimension -
                                  (self.y_dimension <= 2 or not self.periodic))
        return (self.x_dimension * n_vertical_edges_per_x *
                (2 if ordered else 1))

    def n_neighbor_pairs(self, ordered=True):
        """Number of neighboring (unordered) pairs of sites."""
        return (self.n_horizontal_neighbor_pairs(ordered) +
                self.n_vertical_neighbor_pairs(ordered))

    def neighbors_iter(self, ordered=True):
        return itertools.chain(self.horizontal_neighbors_iter(ordered),
                               self.vertical_neighbors_iter(ordered))

    def diagonal_neighbors_iter(self, ordered=True):
        n_sites_per_y = (self.x_dimension -
                         (self.x_dimension <= 2 or not self.periodic))
        n_sites_per_x = (self.y_dimension -
                         (self.y_dimension <= 2 or not self.periodic))
        for x in range(n_sites_per_y):
            for y in range(n_sites_per_x):
                for dy in (-1, 1):
                    i = self.to_site_index((x, y))
                    j = self.to_site_index(((x + 1) % self.x_dimension,
                                            (y + dy) % self.y_dimension))
                    yield (i, j)
                    if ordered:
                        yield (j, i)

    def horizontal_neighbors_iter(self, ordered=True):
        n_horizontal_edges_per_y = (
            self.x_dimension - (self.x_dimension <= 2 or not self.periodic))
        for x in range(n_horizontal_edges_per_y):
            for y in range(self.y_dimension):
                i = self.to_site_index((x, y))
                j = self.to_site_index(((x + 1) % self.x_dimension, y))
                yield (i, j)
                if ordered:
                    yield (j, i)

    def vertical_neighbors_iter(self, ordered=True):
        n_vertical_edges_per_x = (self.y_dimension -
                                  (self.y_dimension <= 2 or not self.periodic))
        for y in range(n_vertical_edges_per_x):
            for x in range(self.x_dimension):
                i = self.to_site_index((x, y))
                j = self.to_site_index((x, (y + 1) % self.y_dimension))
                yield (i, j)
                if ordered:
                    yield (j, i)

    # square-specific geometry

    @property
    def shape(self):
        return (self.x_dimension, self.y_dimension)

    def delta_mag(self, X, Y, by_index=False):
        """The distance between sites X and Y in each dimension."""
        if by_index:
            return self.delta_mag(self.from_site_index(X),
                                  self.from_site_index(Y))
        if self.periodic:
            return tuple(
                min(abs((s * (x - y)) % d)
                    for s in (-1, 1))
                for d, x, y in zip(self.shape, X, Y))
        return tuple(abs(x - xx) for x, xx in zip(X, Y))

    def manhattan_distance(self, X, Y, by_index=False):
        return sum(self.delta_mag(X, Y, by_index))
