"""Functions and classes for handling hopping terms."""
from typing import Tuple, Union, Iterable, Dict, List, Any
from collections import defaultdict, OrderedDict
import bisect

import numpy as np
from scipy.sparse import coo_matrix

from tbplas.builder.base import (rn_type, rn2_type, rn3_type, pair_type,
                                 hop_term_type, verify_rn, Lockable)


def check_conj(rn: rn3_type, pair: pair_type) -> bool:
    """
    Check whether to take the conjugate part of the hopping term.
    :param rn: (r_a, r_b, r_c) cell index
    :param pair: (orb_i, orb_j) orbital pair
    :return: whether to take conjugate
    """
    assert len(rn) == 3, f"Length of {rn} is not 3"
    assert len(pair) == 2, f"Length of {pair} is not 2"
    if rn[0] > 0:
        return False
    elif rn[0] < 0:
        return True
    elif rn[1] > 0:
        return False
    elif rn[1] < 0:
        return True
    elif rn[2] > 0:
        return False
    elif rn[2] < 0:
        return True
    else:
        return pair[0] > pair[1]


class BaseHopping(Lockable):
    """
    Base class for IntraHopping and InterHopping.

    Attributes
    ----------
    _hoppings: Dict[Tuple[int, int, int], Dict[Tuple[int, int], complex]]
        First level of keys are cell indices (rn).
        Second level of keys are orbital pairs.
        Values are hopping energies.
    """
    def __init__(self) -> None:
        super().__init__()
        self._hoppings = dict()

    def __hash__(self) -> int:
        return hash(tuple(self.to_list()))

    def remove_orbitals(self,
                        indices: Union[Iterable[int], np.ndarray]) -> None:
        """
        Remove the hopping terms corresponding to a list of orbitals and update
        remaining hopping terms.
        :param indices: indices of orbitals to remove
        :return: None
        """
        if __debug__:
            self.check_editable()
        for i in indices:
            assert i >= 0, f"Orbital index {i} should >= 0"
        indices = sorted(set(indices))

        for rn, hop_rn in self._hoppings.items():
            new_hop_rn = dict()
            for pair in hop_rn.keys():
                ii, jj = pair
                if ii in indices or jj in indices:
                    pass
                else:
                    ii -= bisect.bisect_left(indices, ii)
                    jj -= bisect.bisect_left(indices, jj)
                    new_hop_rn[(ii, jj)] = hop_rn[pair]
            self._hoppings[rn] = new_hop_rn

    def purge(self) -> None:
        """
        Remove empty cell indices.
        :return: None
        """
        if __debug__:
            self.check_editable()
        for rn in list(self._hoppings.keys()):
            if len(self._hoppings[rn]) == 0:
                self._hoppings.pop(rn)

    def to_list(self) -> List[hop_term_type]:
        """
        Flatten all hopping terms into a list.
        :return: list of hopping terms (rb, rb, rc, orb_i, orb_j, energy)
        """
        self.purge()
        hop_list = [(rn[0], rn[1], rn[2], pair[0], pair[1], energy)
                    for rn, hop_rn in self._hoppings.items()
                    for pair, energy in hop_rn.items()]
        return hop_list

    def to_array(self, use_int64: bool = False) -> Tuple[np.ndarray, np.ndarray]:
        """
        Convert hopping terms to array of 'hop_ind' and 'hop_eng', for
        constructing attributes of 'PrimitiveCell' or 'Sample'.
        :param use_int64: whether to use 64-bit integer for hop_ind, should be
            enabled for 'SuperCell' and 'SCInterHopping' classes
        :return: (hop_ind, hop_eng)
            hop_ind: (num_hop, 5) int32/int64 array, hopping indices
            hop_eng: (num_hop,) complex128 array, hopping energies
        """
        self.purge()
        num_hop = self.num_hop
        dtype = np.int64 if use_int64 else np.int32
        if num_hop == 0:
            hop_ind = np.array([], dtype=dtype)
            hop_eng = np.array([], dtype=np.complex128)
        else:
            hop_ind = np.zeros((num_hop, 5), dtype=dtype)
            hop_eng = np.zeros(num_hop, dtype=np.complex128)
            counter = 0
            for rn, hop_rn in self._hoppings.items():
                for pair, energy in hop_rn.items():
                    hop_ind[counter] = (rn[0], rn[1], rn[2], pair[0], pair[1])
                    hop_eng[counter] = energy
                    counter += 1
        return hop_ind, hop_eng

    @property
    def hoppings(self) -> Dict[rn3_type, Dict[rn2_type, complex]]:
        """Interface for the '_hoppings' attribute."""
        return self._hoppings

    @property
    def num_hop(self) -> int:
        """
        Count the number of hopping terms.
        :return: number of hopping terms
        """
        num_hop = 0
        for hop_rn in self._hoppings.values():
            num_hop += len(hop_rn)
        return num_hop


class IntraHopping(BaseHopping):
    """
    Container for holding hopping terms in a primitive cell or modifications to
    hopping terms in a supercell.
    """
    @staticmethod
    def _norm_keys(rn: rn_type,
                   orb_i: int,
                   orb_j: int) -> Tuple[rn3_type, pair_type, bool]:
        """
        Normalize cell index and orbital pair into legal key, also check
        whether to take the conjugation and return the status in conj.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :return: (rn, pair, conj)
            where rn is the normalized cell index,
            pair is the normalized orbital pair,
            conj is the flag of whether to take the conjugate of hopping energy
        """
        assert orb_i >= 0, f"orb_i {orb_i} should >= 0"
        assert orb_j >= 0, f"orb_j {orb_j} should >= 0"
        pair = (orb_i, orb_j)
        conj = check_conj(rn, pair)
        if conj:
            rn = (-rn[0], -rn[1], -rn[2])
            pair = (orb_j, orb_i)
        return rn, pair, conj

    def add_hopping(self,
                    rn: rn_type,
                    orb_i: int,
                    orb_j: int,
                    energy: complex) -> None:
        """
        Add a new hopping term or update existing term.

        Conjugate terms are reduced by normalizing their cell indices and
        orbital pairs.

        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :param energy: hopping energy
        :return: None
        """
        if __debug__:
            self.check_editable()
        rn, pair, conj = self._norm_keys(rn, orb_i, orb_j)
        if conj:
            energy = energy.conjugate()
        try:
            hop_rn = self._hoppings[rn]
        except KeyError:
            hop_rn = self._hoppings[rn] = dict()
        hop_rn[pair] = energy

    def get_hopping(self,
                    rn: rn_type,
                    orb_i: int,
                    orb_j: int) -> Tuple[complex, bool]:
        """
        Get an existing hopping term.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :return: (energy, status)
            where energy is the hopping energy and status is whether the term
            has been found
        """
        rn, pair, conj = self._norm_keys(rn, orb_i, orb_j)
        try:
            energy = self._hoppings[rn][pair]
            status = True
        except KeyError:
            energy = None
            status = False
        if status and conj:
            energy = energy.conjugate()
        return energy, status

    def remove_hopping(self,
                       rn: rn_type,
                       orb_i: int,
                       orb_j: int) -> bool:
        """
        Remove an existing hopping term.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :return: where the hopping term is removed, False if not found
        """
        if __debug__:
            self.check_editable()
        rn, pair, conj = self._norm_keys(rn, orb_i, orb_j)
        try:
            self._hoppings[rn].pop(pair)
            status = True
        except KeyError:
            status = False
        return status

    def remove_rn(self, rn: rn_type) -> bool:
        """
        Remove all the hopping terms of given cell index.
        :param rn: (r_a, r_b, r_c), cell index
        :return: where the hopping terms are removed, False if not found
        """
        if __debug__:
            self.check_editable()
        rn, pair, conj = self._norm_keys(rn, 0, 0)
        try:
            self._hoppings.pop(rn)
            status = True
        except KeyError:
            status = False
        return status

    def count_pairs(self) -> Dict[pair_type, int]:
        """
        Count how many times each orbital pair emerges.
        :return: numbers of emergence of orbital pairs
        """
        pair_count = defaultdict(int)
        for hop_rn in self.hoppings.values():
            for pair, energy in hop_rn.items():
                # All terms are normalized to upper triangular
                if pair[0] <= pair[1]:
                    pair_count[pair] += 1
                else:
                    pair_conj = (pair[1], pair[0])
                    pair_count[pair_conj] += 1
        return pair_count


class InterHopping(BaseHopping):
    """
    Container for holding hopping terms between two models.

    Attributes
    ----------
    __model_bra: 'PrimitiveCell' or 'SuperCell' instance
        the 'bra' model from which the hopping terms exist
    __model_ket: 'PrimitiveCell' or 'SuperCell' instance
        the 'ket' model to which the hopping terms exist
    """
    def __init__(self, model_bra: Any, model_ket: Any) -> None:
        """
        :param model_bra: the 'bra' model from which the hopping terms exist
        :param model_ket: the 'ket' model to which the hopping terms exist
        """
        super().__init__()
        self.__model_bra = model_bra
        self.__model_ket = model_ket

    def __hash__(self) -> int:
        fp = (tuple(self.to_list()), self.__model_bra, self.__model_ket)
        return hash(fp)

    @staticmethod
    def _norm_keys(rn: rn_type,
                   orb_i: int,
                   orb_j: int) -> Tuple[rn3_type, pair_type]:
        """
        Normalize cell index and orbital pair into permitted keys, also check if
        rn is legal since the class is exposed to the user.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :return: (rn, pair)
            where rn is the normalized cell index,
            pair is the normalized orbital pair,
        """
        assert orb_i >= 0, f"orb_i {orb_i} should >= 0"
        assert orb_j >= 0, f"orb_j {orb_j} should >= 0"
        rn = verify_rn(rn)
        pair = (orb_i, orb_j)
        return rn, pair

    def add_hopping(self,
                    rn: rn_type,
                    orb_i: int,
                    orb_j: int,
                    energy: complex) -> None:
        """
        Add a new hopping term or update existing term.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :param energy: hopping energy
        :return: None
        """
        if __debug__:
            self.check_editable()
        rn, pair = self._norm_keys(rn, orb_i, orb_j)
        try:
            hop_rn = self._hoppings[rn]
        except KeyError:
            hop_rn = self._hoppings[rn] = dict()
        hop_rn[pair] = energy

    def get_hopping(self,
                    rn: rn_type,
                    orb_i: int,
                    orb_j: int) -> Tuple[complex, bool]:
        """
        Get an existing hopping term.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :return: (energy, status)
            where energy is the hopping energy and status is whether the term
            has been found
        """
        rn, pair = self._norm_keys(rn, orb_i, orb_j)
        try:
            energy = self._hoppings[rn][pair]
            status = True
        except KeyError:
            energy = None
            status = False
        return energy, status

    def remove_hopping(self,
                       rn: rn_type,
                       orb_i: int,
                       orb_j: int) -> bool:
        """
        Remove an existing hopping term.
        :param rn: (r_a, r_b, r_c), cell index
        :param orb_i: orbital index or bra
        :param orb_j: orbital index of ket
        :return: where the hopping term is removed, False if not found
        """
        if __debug__:
            self.check_editable()
        rn, pair = self._norm_keys(rn, orb_i, orb_j)
        try:
            self._hoppings[rn].pop(pair)
            status = True
        except KeyError:
            status = False
        return status

    def remove_rn(self, rn: rn_type) -> bool:
        """
        Remove all the hopping terms of given cell index.
        :param rn: (r_a, r_b, r_c), cell index
        :return: where the hopping terms are removed, False if not found
        """
        if __debug__:
            self.check_editable()
        rn, pair = self._norm_keys(rn, 0, 0)
        try:
            self._hoppings.pop(rn)
            status = True
        except KeyError:
            status = False
        return status

    def count_pairs(self) -> Dict[pair_type, int]:
        """
        Count how many times each orbital pair emerges.
        :return: numbers of emergence of orbital pairs
        """
        pair_count = defaultdict(int)
        for rn, hop_rn in self.hoppings.items():
            for pair, energy in hop_rn.items():
                pair_count[pair] += 1
        return pair_count


class HopDict:
    """
    Class for holding hopping terms in matrix form, reserved for compatibility
    with old version of TBPlaS.

    DO NOT try to rewrite this class based on IntraHopping. This class is
    intended for compatibility reasons and follows a different philosophy than
    IntraHopping. Hopping terms will not be reduced by conjugate relation even
    for the (0, 0, 0) cell, as some users may utilize H(R) and H(-R) for
    different kinds of hopping terms. However, this makes this class bug-prone.
    So we do not recommend common users to use this class.

    Since the hopping terms are not reduced, the order of H(R) and H(-R) may
    lead to random results. To avoid that, we use OrderedDict instead of dict
    for holding the hopping terms.

    Attributes
    ----------
    _hoppings: OrderedDict[Tuple[int, int, int], np.ndarray]
        Keys should be cell indices and values should be complex matrices.
    _num_orb: int
        number of orbitals
    _mat_shape: Tuple[int, int]
        shape of hopping matrices
    """
    def __init__(self, num_orb: int) -> None:
        """
        :param num_orb: number of orbitals
        """
        assert num_orb > 0, f"Number of orbital {num_orb} should be > 0"
        self._hoppings = OrderedDict()
        self._num_orb = num_orb
        self._mat_shape = (self._num_orb, self._num_orb)

    def __setitem__(self, rn: rn_type, hop_mat: np.ndarray) -> None:
        """
        Add or update a hopping matrix according to cell index.
        :param rn: cell index of hopping matrix
        :param hop_mat: (num_orb, num_orb) complex128 array hopping matrix
        :return: None
        """
        rn = verify_rn(rn)
        if not isinstance(hop_mat, np.ndarray):
            hop_mat = np.array(hop_mat)
        error_msg = (f"Shape of hopping matrix {hop_mat.shape} does not match"
                     f" {self._mat_shape}")
        assert hop_mat.shape == self._mat_shape,  error_msg

        # Set hopping matrix
        # We copy the hopping matrix to avoid changing it accidentally.
        self._hoppings[rn] = hop_mat.copy()

        # Check for diagonal terms
        if __debug__ and rn == (0, 0, 0):
            self._check_diag()

    def __getitem__(self, rn: rn_type) -> np.ndarray:
        """
        Get the hopping matrix according to cell index.
        :param rn: cell index of hopping matrix
        :return: (num_orb, num_orb) complex128 array hopping matrix of rn
        """
        rn = verify_rn(rn)
        try:
            hop_mat = self._hoppings[rn]
        except KeyError:
            hop_mat = self._hoppings[rn] = np.zeros(self._mat_shape,
                                                    dtype=np.complex128)
        return hop_mat

    def _check_diag(self) -> None:
        """
        Check for diagonal hopping terms (on-site energies).
        :return: None
        """
        rn = (0, 0, 0)
        if rn in self._hoppings.keys():
            hop_mat = self._hoppings[rn]
            for i, energy in enumerate(np.abs(np.diag(hop_mat))):
                assert energy < 1e-5, f"Hop term {rn}{i, i} is diagonal"

    def to_spare(self) -> Dict[rn3_type, coo_matrix]:
        """
        Return an ordered dictionary with values being sparse matrices.
        :return: ordered sparse hopping dictionary
        """
        self._check_diag()
        sparse_hop_dict = OrderedDict()
        for rn, hop_mat in self._hoppings.items():
            sparse_hop_dict[rn] = coo_matrix(hop_mat)
        return sparse_hop_dict

    @property
    def num_orb(self) -> int:
        """Interface for the '_num_orb' attribute."""
        return self._num_orb

    @property
    def mat_shape(self) -> Tuple[int, int]:
        """Interface for the 'mat_shape' attribute."""
        return self._mat_shape

    @property
    def hoppings(self) -> Dict[rn3_type, np.ndarray]:
        """Interface for the 'hoppings' attribute."""
        return self._hoppings
