"""Orbital maps for converting orbital indices."""

import numpy as np
from typing import Union, List

from tbplas.builder.base import id_pc_type, rn3_type
from tbplas.cython import orbital_map as core


__all__ = ["GeneralOrbitalMap", "RegularOrbitalMap", "RegularOrbitalMapVac"]


class GeneralOrbitalMap:
    """
    Helper class for converting orbital indices between pc and sc
    representations in general cases.

    NOTE: The orbital indices passed to 'add' method must not be duplicate.
    Otherwise, bugs will be definitely raised.

    Attributes
    ----------
    _id_pc: List[Tuple[int, int, int, int]]
        orbital indices in pc representation
    _id_sc: Dict[Tuple[int, int, int, int], int]
        map of orbital indices from pc to sc representation
    _counter: int
        counter of the orbitals
    """
    def __init__(self) -> None:
        self._id_pc = []
        self._id_sc = dict()
        self._counter = 0

    def add(self, id_pc: id_pc_type) -> None:
        """
        Add an orbital from index in pc representation.
        :param id_pc: index of the orbital
        :return: None
        """
        assert len(id_pc) == 4, f"Length of {id_pc} should be 4"
        assert id_pc[3] >= 0, f"{id_pc[3]} out of range"
        self._id_pc.append(id_pc)
        self._id_sc[id_pc] = self._counter
        self._counter += 1

    def pc2sc(self, id_pc: id_pc_type) -> int:
        """
        Convert orbital index from pc to sc representation.
        :param id_pc: index of the orbital
        :return: corresponding index in sc representation
        """
        assert len(id_pc) == 4, f"Length of {id_pc} should be 4"
        assert id_pc[3] >= 0, f"{id_pc[3]} out of range"
        return self._id_sc[id_pc]

    def sc2pc(self, id_sc: int) -> id_pc_type:
        """
        Convert orbital index from sc to pc representation.
        :param id_sc: index of the orbital
        :return: corresponding index in pc representation
        """
        assert id_sc >= 0, f"{id_sc} should >= 0"
        return self._id_pc[id_sc]

    @property
    def counter(self) -> int:
        """
        Get the current number of orbitals added to the map.
        :return: current counter starting from 0
        """
        return self._counter


class RegularOrbitalMap:
    """
    Similar to GeneralOrbitalMap, but the orbitals are generated using a 4-fold
    loops over dim_a, dim_b, dim_c and num_orb_pc.

    Attributes
    ----------
    _dim: Tuple[int, int, int]
        dimension of supercell along a, b, c directions
    _dim_bc: int
        cached dim_b * dim_c, for speed boost
    _num_orb_pc: int
        number of orbitals in primitive cell
    """
    def __init__(self, dim: rn3_type, num_orb_pc: int) -> None:
        """
        :param dim: dimension of supercell along a, b, c directions
        :param num_orb_pc: number of orbitals in primitive cell
        """
        assert len(dim) == 3, f"Length of {dim} should be 3"
        for i_d, d in enumerate(dim):
            assert d >= 0, f"Dim[{i_d}] {d} should >= 1"
        assert num_orb_pc >= 1, f"num_orb_pc {num_orb_pc} should >= 1"
        self._dim = dim
        self._dim_bc = dim[1] * dim[2]
        self._num_orb_pc = num_orb_pc

    def pc2sc(self, i_a: int, i_b: int, i_c: int, i_o: int) -> int:
        """
        Convert orbital index from pc to sc representation.

        NOTE: For efficiency we do not accept id_pc_t as input argument as in
        GeneralOrbitalMap.

        :param i_a: cell index along a direction
        :param i_b: cell index along a direction
        :param i_c: cell index along a direction
        :param i_o: orbital index
        :return: corresponding index in sc representation
        """
        assert 0 <= i_a < self._dim[0], f"{i_a} out of range"
        assert 0 <= i_b < self._dim[1], f"{i_b} out of range"
        assert 0 <= i_c < self._dim[2], f"{i_c} out of range"
        assert 0 <= i_o < self._num_orb_pc, f"{i_o} out of range"
        num_pc = i_a * self._dim_bc + i_b * self._dim[2] + i_c
        id_sc = num_pc * self._num_orb_pc + i_o
        return id_sc

    def sc2pc(self, id_sc: int) -> id_pc_type:
        """
        Convert orbital index from sc to pc representation.
        :param id_sc: index of the orbital
        :return: corresponding index in pc representation
        """
        assert id_sc >= 0, f"{id_sc} should >= 0"
        i_o, n1 = id_sc % self._num_orb_pc, id_sc // self._num_orb_pc
        i_c, n1 = n1 % self._dim[2], n1 // self._dim[2]
        i_b, i_a = n1 % self._dim[1], n1 // self._dim[1]
        return i_a, i_b, i_c, i_o


class RegularOrbitalMapVac:
    """
    Similar to RegularOrbitalMap, but can handle cases with vacancies.

    Attributes
    ----------
    _dim: (3, ) int32 array
        dimension of the supercell along a, b, and c directions
    _num_orb_pc: int
        number of orbitals in primitive cell
    _vac_id_sc: (num_vac,) int64 array
        indices of vacancies in supercell representation
    _orb_id_pc: (num_orb_sc, 4) int32 array
        indices of orbitals in primitive cell representation
    """
    def __init__(self,
                 dim: np.ndarray,
                 num_orb_pc: int,
                 vac_id_sc: np.ndarray,
                 orb_id_pc: np.ndarray) -> None:
        """
        NOTE: Currently, this class is applied in SuperCell class only. It is
        assumed this class will check the parameters passed to this class. No
        assertions will be made here.
        :param dim: dimension of the supercell along a, b and c directions
        :param num_orb_pc: number of orbitals in primitive cell
        :param vac_id_sc: indices of vacancies in sc representation
        :param orb_id_pc: indices of orbitals in pr representation
        """
        self._dim = dim
        self._num_orb_pc = num_orb_pc
        self._vac_id_sc = vac_id_sc
        self._orb_id_pc = orb_id_pc

    def _check_id_pc(self, id_pc: Union[id_pc_type, np.ndarray]) -> None:
        """
        Checks if orbital or vacancy index in primitive cell representation
        is legal.

        A legal id_pc should have cell indices falling within 0 <= rn < dim and
        orbital index falling within 0 <= ib < num_orb_pc.

        :param id_pc: (ia, ib, ic, io) or equivalent int32 array
            orbital or vacancy index in primitive cell representation
        :return: None
        """
        assert len(id_pc) == 4, f"Length of {id_pc} should be 4"
        if isinstance(id_pc, tuple):
            assert 0 <= id_pc[0] < self._dim.item(0), f"{id_pc}[0] out of range"
            assert 0 <= id_pc[1] < self._dim.item(1), f"{id_pc}[1] out of range"
            assert 0 <= id_pc[2] < self._dim.item(2), f"{id_pc}[2] out of range"
            assert 0 <= id_pc[3] < self._num_orb_pc,  f"{id_pc}[3] out of range"
        elif isinstance(id_pc, np.ndarray):
            assert 0 <= id_pc.item(0) < self._dim.item(0), f"{id_pc}[0] out of range"
            assert 0 <= id_pc.item(1) < self._dim.item(1), f"{id_pc}[1] out of range"
            assert 0 <= id_pc.item(2) < self._dim.item(2), f"{id_pc}[2] out of range"
            assert 0 <= id_pc.item(3) < self._num_orb_pc,  f"{id_pc}[3] out of range"
        else:
            raise AssertionError(f"Illegal id_pc type {type(id_pc)}")

    def sc2pc(self, id_sc: int) -> id_pc_type:
        """
        Convert orbital (NOT VACANCY) index from sc representation to pc
        representation.

        NOTE: This method is safe, but EXTREMELY SLOW. If you are going to
        call this method many times, use orb_id_sc2pc_array instead.

        :param id_sc: index of orbital in supercell representation
        :return: index of orbital in primitive cell representation
        """
        assert 0 <= id_sc < self._orb_id_pc.shape[0], f"id_sc {id_sc} out of range"
        return tuple(self._orb_id_pc[id_sc])

    def pc2sc(self, id_pc: Union[id_pc_type, np.ndarray]) -> int:
        """
        Convert orbital (NOT VACANCY) index from pc representation to sc
        representation.

        NOTE: This method is safe, but EXTREMELY SLOW. If you are going to
        call this method many times, use orb_id_pc2sc_array instead.

        :param id_pc: (ia, ib, ic, io), or equivalent int32 array
            index of orbital in primitive cell representation
        :return: index of orbital in supercell representation
        """
        if __debug__:
            self._check_id_pc(id_pc)
        id_pc = np.array(id_pc, dtype=np.int32)
        orb_id_sc = core.id_pc2sc(self._dim, self._num_orb_pc, id_pc,
                                  self._vac_id_sc)
        assert orb_id_sc != -1, f"id_pc {id_pc} is a vacancy"
        return orb_id_sc

    def sc2pc_array(self, id_sc_array: Union[np.ndarray, List[int]]) -> np.ndarray:
        """
        Convert an array of orbital (NOT VACANCY) indices from sc
        representation to pc representation.
        :param id_sc_array: (num_orb, ) int64 array or list of int64
            orbital indices in supercell representation
        :return: (num_orb, 4) int32 array
            orbital indices in primitive cell representation
        """
        id_sc_array = np.array(id_sc_array, dtype=np.int64)
        assert len(id_sc_array.shape) == 1, f"id_sc_array should be a vector"
        num_orb_sc = self._orb_id_pc.shape[0]
        status = core.check_id_sc_array(num_orb_sc, id_sc_array)
        error_code, error_idx = status
        assert error_code == 0, f"id_sc {id_sc_array[error_idx]} out of range"
        id_pc_array = core.id_sc2pc_array(self._orb_id_pc, id_sc_array)
        return id_pc_array

    def pc2sc_array(self, id_pc_array: Union[np.ndarray, List[id_pc_type]]) -> np.ndarray:
        """
        Convert an array of orbital (NOT VACANCY) indices from pc
        representation to sc representation.
        :param id_pc_array: (num_orb, 4) int32 array or list of id_pc_t
            orbital indices in primitive cell representation
        :return: (num_orb, ) int64 array
            orbital indices in supercell representation
        """
        id_pc_array = np.array(id_pc_array, dtype=np.int32)
        if __debug__:
            error_msg = f"Illegal shape of id_pc_array {id_pc_array.shape}"
            assert len(id_pc_array.shape) == 2, error_msg
            assert id_pc_array.shape[1] == 4, error_msg
        status = core.check_id_pc_array(self._dim, self._num_orb_pc,
                                        id_pc_array,
                                        self._vac_id_sc)
        error_code, error_idx, error_jdx = status
        if error_code == -2:
            raise AssertionError(f"index {error_jdx} of {id_pc_array[error_idx]}"
                                 f" out of range")
        if error_code == -1:
            raise AssertionError(f"id_pc {id_pc_array[error_idx]} is a vacancy")
        id_sc_array = core.id_pc2sc_array(self._dim, self._num_orb_pc,
                                          id_pc_array,
                                          self._vac_id_sc)
        return id_sc_array
