"""Functions and classes for supercell."""

from collections import namedtuple
from typing import Callable, Iterable, Tuple, Union, Set, List

import numpy as np
import matplotlib.pyplot as plt

from tbplas.base import lattice as lat
from tbplas.cython import super as core
from tbplas.builder.base import (verify_rn, verify_pbc, Lockable, rn_type,
                                 rn3_type, pbc_type)
from tbplas.builder.hopping import IntraHopping, InterHopping
from tbplas.builder.orbital_map import RegularOrbitalMapVac
from tbplas.builder.primitive import PrimitiveCell
from tbplas.builder.visual import ModelViewer


hop_type = Tuple[np.ndarray, np.ndarray, np.ndarray]


__all__ = ["SuperCell", "SCInterHopping"]


class SCCachedData:
    """
    Cached array attributes of SuperCell.

    NOTES
    -----
    1. DO NOT USE dataclass.

    Dataclass uses class attributes instead of instance attributes, i.e., there
    will only be one copy of the attributes, which will cause bugs.

    2. Why no vac_id_pc and orb_id_sc?

    It's unnecessary to have the vac_id_pc array, as it can be generated from
    vacancy_set on-the-fly. Similarly, orb_id_sc can be generated from
    orb_id_pc. So we do not keep it in memory to reduce RAM usage.

    However, it should be noted that vac_id_sc and orb_id_sc are generated via
    different approaches. We show it by an example of 2*2 supercell with 2
    orbitals per primitive cell. The indices of orbitals as well as vacancies
    in primitive cell representation are
               id_pc    id_sc    type
        (0, 0, 0, 0)        0     orb
        (0, 0, 0, 1)        1     vac
        (0, 1, 0, 0)        2     vac
        (0, 1, 0, 1)        3     orb
        (1, 0, 0, 0)        4     orb
        (1, 0, 0, 1)        5     vac
        (1, 0, 0, 0)        6     orb
        (1, 0, 0, 1)        7     orb

    The indices for vacancies in supercell representation are the original
    id_sc, i.e.1, 2, and 5. However, the indices for orbitals are re-ordered
    to be
               id_pc    id_sc    type
        (0, 0, 0, 0)        0     orb
        (0, 1, 0, 1)        1     orb
        (1, 0, 0, 0)        2     orb
        (1, 0, 0, 0)        3     orb
        (1, 0, 0, 1)        4     orb

    Attributes
    ----------
    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
    hash_dict: Dict[str, int]
        hashes of attributes to be used by 'sync_array' to update the arrays
    """
    def __init__(self):
        self.vac_id_sc = np.array([], dtype=np.int64)
        self.orb_id_pc = np.array([], dtype=np.int32)
        self.hash_dict = {"pc": None, "vac": None}


class SuperCell(Lockable):
    """
    Class for representing a supercell from which the sample is constructed.

    Attributes
    ----------
    _prim_cell: 'PrimitiveCell' instance
        primitive cell from which the supercell is constructed
    _dim: (3, ) int32 array
        dimension of the supercell along a, b, and c directions
    _pbc: (3, ) int32 array
        whether to enable periodic condition along a, b, and c directions
        0 for False, 1 for True.
    _vacancy_set: Set[int]
        indices of vacancies in supercell representation
    _cached_data: 'SCCachedData'
        cached array attributes
    _hop_modifier: 'IntraHopping' instance
        modification to hopping terms in the supercell
    _orb_pos_modifier: function
        modification to orbital positions in the supercell

    NOTES
    -----
    1. Minimal supercell dimension of single-orbital primitive cell

    Assume that we have a primitive cell located at R=0. The furthest primitive
    cell between which hopping terms exist is located at R=N. It can be proved
    that if the dimension of supercell along that direction is less than N,
    the same matrix element hij will appear more than one time in hop_ind and
    hop_eng of 'SuperCell' class, as well as its conjugate counterpart. This
    will complicate the program and significantly slow it down, which situation
    we must avoid.

    Further, if the dimension of supercell falls in [N, 2*N], hij will appear
    only one time, yet its conjugate counterpart still exists. Although no
    problems have been found so far, we still need to avoid this situation.

    So the minimal dimension of supercell is 2*N+1, where N is the index of
    the furthest primitive cell between which hopping terms exists. Otherwise,
    the 'SuperCell' class, as well as the core functions of '_get_num_hop_sc'
    and 'build_hop' will not work properly.

    In the hr.dat file produced by Wannier90, there is an N_min and an N_max
    for the furthest primitive cell index. In that case, N should be the
    maximum of |N_max| and |N_min| as the result of translational symmetry.

    For a multiple-orbital primitive cell, the minimal supercell dimension may
    be not that large. We may even build a 1*1*1 supercell. Call the 'check_dim'
    method of 'PrimitiveCell' class to check if this is possible.

    2. hop_modifier

    We reduce hopping terms according to the conjugate relation
        <0, bra|H|R, ket> = <0, ket|H|-R, bra>*.
    So actually only half of hopping terms are stored.

    If the hopping terms claimed here are already included in the supercell,
    they will overwrite the existing terms. If the hopping terms or their
    conjugate counterparts are new to the supercell, they will be added to
    hop_* arrays. The dr array will also be updated accordingly.
    """
    def __init__(self,
                 prim_cell: PrimitiveCell,
                 dim: rn_type,
                 pbc: pbc_type = (False, False, False),
                 check_dim: bool = True,
                 vacancies: Union[Iterable[int], np.ndarray] = None,
                 orb_pos_modifier: Callable[[np.ndarray], None] = None) -> None:
        """
        NOTE: The indices of vacancies are always defined with respect to the
        pristine supercell, i.e., supercell without any vacancies. For example,
        for a pristine supercell with 18 orbitals, vacancies [0, 1] will remove
        the first 2 orbitals, yielding a supercell with 16 orbitals. For the 16
        orbital supercell, add_vacancies([0, 1]) or set_vacancies([0, 1]) will
        add any new vacancies, since [0, 1] are defined w.r.t pristine
        supercell, and have already been declared as vacancies. This rule also
        applies to trim().
        :param prim_cell: primitive cell from which the supercell is constructed
        :param dim: dimension of the supercell along a, b and c directions
        :param pbc: whether to enable periodic boundary condition along a, b,
            and c directions
        :param check_dim: whether to check supercell dimension using the
            single-orbital rule
        :param vacancies: list or equivalent int64 array of indices of vacancies
            in supercell representation
        :param orb_pos_modifier: modification to orbital positions in the
            supercell
        """
        super().__init__()

        # Set and lock the primitive cell
        self._prim_cell = prim_cell
        self._prim_cell.lock()
        self._prim_cell.add_subscriber(self)
        assert prim_cell.num_orb > 0, f"Num_orb_pc should > 0"
        assert prim_cell.num_hop > 0, f"Num_hop_pc should > 0"

        # Check and set dimension
        dim = verify_rn(dim, complete_item=1)
        if check_dim:
            hop_ind = self._prim_cell.hop_ind
            for i in range(3):
                rn_min = hop_ind[:, i].min()
                rn_max = hop_ind[:, i].max()
                dim_min = 2 * max(abs(rn_min), abs(rn_max)) + 1
                assert dim[i] >= dim_min, f"dim[{i}] should >= {dim_min}"
        # Use copy to ensure that it is not referenced to external data.
        self._dim = np.array(dim, dtype=np.int32).copy()

        # Check and set periodic boundary condition
        # Use copy to ensure that it is not referenced to external data.
        pbc = verify_pbc(pbc, complete_item=False)
        self._pbc = np.array([1 if _ else 0 for _ in pbc], dtype=np.int32).copy()

        # Initialize arrays assuming no vacancies
        self._vacancy_set = set()
        self._cached_data = SCCachedData()
        self._cached_data.orb_id_pc = \
            core.build_orb_id_pc(self._dim, self.num_orb_pc,
                                 self._cached_data.vac_id_sc)
        for key in ("pc", "vac"):
            self._cached_data.hash_dict[key] = self._get_attr_hash(key)

        # Set vacancies if any
        if vacancies is not None:
            self.add_vacancies(vacancies)

        # Set modifiers
        self._hop_modifier = IntraHopping()
        self._orb_pos_modifier = orb_pos_modifier

    def __hash__(self) -> int:
        fp = (self._prim_cell, tuple(self._dim), tuple(self._pbc),
              tuple(self._vacancy_set), self._hop_modifier,
              self._orb_pos_modifier)
        return hash(fp)

    def _get_attr_hash(self, attr: str) -> int:
        """
        Get hash of given attribute.

        For the primitive cell, we can use both its actual hash or the number
        of orbitals as the fingerprint, which is much faster.

        :param attr: name of the attribute
        :return: hash of the attribute
        """
        assert attr in ("pc", "vac")
        if attr == "pc":
            new_hash = hash(self._prim_cell)
        else:
            new_hash = hash(tuple(self._vacancy_set))
        return new_hash

    def _update_attr_hash(self, attr: str) -> bool:
        """
        Compare and update hash of given attribute.
        :param attr: name of the attribute
        :return: whether the hash has been updated
        """
        new_hash = self._get_attr_hash(attr)
        if self._cached_data.hash_dict[attr] != new_hash:
            self._cached_data.hash_dict[attr] = new_hash
            status = True
        else:
            status = False
        return status

    def _check_vacancies(self,
                         vacancies: Union[Iterable[int], np.ndarray]) -> None:
        """
        Check the list of vacancies.

        NOTE: see __init__ for the rule of indices of vacancies.

        :param vacancies: list or equivalent int64 array of indices of vacancies
            in supercell representation
        :return:
        """
        if isinstance(vacancies, np.ndarray):
            assert len(vacancies.shape) == 1, f"vacancies should be a vector"
        num_orb = self.num_orb_pc * np.prod(self._dim)
        for v in vacancies:
            assert 0 <= v < num_orb, f"vac_id {v} out of range(0, {num_orb})"

    def _sync_array(self, force_sync: bool = False) -> None:
        """
        Synchronize vac_id_sc and orb_id_pc according to primitive cell and
        vacancies.

        NOTE: The core function '_id_pc2sc_vac' requires vac_id_sc to be sorted
        in increasing order. Otherwise, it won't work properly! So we must sort
        it here.

        :param force_sync: whether to force synchronizing the arrays even if
            primitive cell and vacancy_set did not change
        :return: None
        """
        assert self._prim_cell.num_orb > 0, f"Num_orb_pc should > 0"
        assert self._prim_cell.num_hop > 0, f"Num_hop_pc should > 0"

        # Get update flags
        update_pc = self._update_attr_hash("pc")
        update_vac = self._update_attr_hash("vac")
        to_update = update_pc or update_vac

        if force_sync or to_update:
            if self.num_vac == 0:
                self._cached_data.vac_id_sc = np.array([], dtype=np.int64)
            else:
                # Check for invalid vacancies if primitive cell changes
                if update_pc:
                    self._check_vacancies(self._vacancy_set)
                vac_id_sc = np.array(list(self._vacancy_set), dtype=np.int64)
                sorted_idx = np.argsort(vac_id_sc, axis=0)
                self._cached_data.vac_id_sc = vac_id_sc[sorted_idx]
            self._cached_data.orb_id_pc = \
                core.build_orb_id_pc(self._dim, self.num_orb_pc,
                                     self._cached_data.vac_id_sc)

            # Re-lock primitive cell
            self._prim_cell.lock()

    def add_vacancies(self, vacancies: Union[Iterable[int], np.ndarray]) -> None:
        """
        Add a list of vacancies to the orbital set.

        NOTE: see __init__ for the rule of indices of vacancies.

        :param vacancies: list or equivalent int64 array of id_sc
            list of indices of vacancies in supercell representation
        :return: None
        """
        if __debug__:
            self.check_editable()
            self._check_vacancies(vacancies)
        # Force conversion to int, otherwise hash(self._vacancy_set) may fail.
        vacancies = set([int(_) for _ in vacancies])
        self._vacancy_set = self._vacancy_set.union(vacancies)

    def set_vacancies(self, vacancies: Union[Iterable[int], np.ndarray]) -> None:
        """
        Reset the set of vacancies.

        NOTE: see __init__ for the rule of indices of vacancies.

        :param vacancies: list or equivalent int64 array of id_sc
            list of indices of vacancies in supercell representation
        :return: None
        """
        if __debug__:
            self.check_editable()
        self._vacancy_set = set()
        self.add_vacancies(vacancies)

    def set_orb_pos_modifier(self, orb_pos_modifier: Callable = None) -> None:
        """
        Reset orb_pos_modifier.
        :param orb_pos_modifier: modifier to orbital positions
        :return: None
        """
        if __debug__:
            self.check_editable()
        self._orb_pos_modifier = orb_pos_modifier

    def _check_hop_index(self,
                         rn: rn3_type,
                         orb_i: int,
                         orb_j: int) -> None:
        """
        Check if the hop_index is legal.
        :param rn: cell index of the hopping term, i.e. R
        :param orb_i: index of orbital i in <i,0|H|j,R>
        :param orb_j: index of orbital j in <i,0|H|j,R>
        :return: None
        """
        num_orb = self.num_orb_sc
        assert 0 <= orb_i < num_orb, f"Orb_i {orb_i} out of range(0, {num_orb})"
        assert 0 <= orb_j < num_orb, f"Orb_j {orb_j} out of range(0, {num_orb})"
        error_msg = f"{rn}{orb_i, orb_i} is a diagonal term"
        assert not (rn == (0, 0, 0) and orb_i == orb_j), error_msg

    def add_hopping(self,
                    rn: rn_type,
                    orb_i: int,
                    orb_j: int,
                    energy: complex) -> None:
        """
        Add a new term to the hopping modifier.
        :param rn: cell index of the hopping term, i.e. R
        :param orb_i: index of orbital i in <i,0|H|j,R>
        :param orb_j: index of orbital j in <i,0|H|j,R>
        :param energy: hopping integral in eV
        :return: None
        """
        rn = verify_rn(rn)
        if __debug__:
            self.check_editable()
            self._check_hop_index(rn, orb_i, orb_j)
        self._hop_modifier.add_hopping(rn, orb_i, orb_j, energy)

    def trim(self) -> None:
        """
        Trim dangling orbitals and associated hopping terms.

        NOTE: see __init__ for the rule of indices of vacancies.

        :return: None.
        """
        if __debug__:
            self.check_editable()

        # Get indices of dangling orbitals
        hop_ind = self.get_hop()[0]
        id_sc_trim = core.get_orb_id_trim(self._dim, self.num_orb_pc,
                                          self._cached_data.orb_id_pc,
                                          self._cached_data.vac_id_sc,
                                          hop_ind)

        # Add vacancies
        self.add_vacancies(id_sc_trim)

        # Also trim hop_modifier
        self._hop_modifier.remove_orbitals(id_sc_trim)

    def plot(self,
             axes: plt.Axes,
             with_orbitals: bool = True,
             with_cells: bool = True,
             with_conj: bool = False,
             orb_color: Callable[[np.ndarray], List[str]] = None,
             hop_as_arrows: bool = True,
             hop_eng_cutoff: float = 1e-5,
             hop_color: str = "r",
             view: str = "ab") -> None:
        """
        Plot lattice vectors, orbitals, and hopping terms to axes.
        :param axes: axes on which the figure will be plotted
        :param with_orbitals: whether to plot orbitals as filled circles
        :param with_cells: whether to plot borders of primitive cells
        :param with_conj: whether to plot conjugate hopping terms as well
        :param orb_color: function for coloring the orbitals
        :param hop_as_arrows: whether to plot hopping terms as arrows
        :param hop_eng_cutoff: cutoff for showing hopping terms
        :param hop_color: color of hopping terms
        :param view: kind of view point, should be in 'ab', 'bc', 'ca', 'ba',
            'cb', 'ac'
        :return: None
        """
        viewer = ModelViewer(axes, self._prim_cell.lat_vec,
                             self._prim_cell.origin, view)
        orb_eng, orb_pos, hop_ind, hop_eng, dr = self.get_array()

        # Plot orbitals
        if orb_color is None:
            scatter_color = orb_eng
        else:
            scatter_color = orb_color(self._cached_data.orb_id_pc)
        if with_orbitals:
            viewer.scatter(orb_pos, c=scatter_color)

        # Plot hopping terms
        viewer.plot_hop(orb_pos, orb_pos, hop_ind, hop_eng, dr,
                        with_conj, hop_as_arrows, hop_eng_cutoff, hop_color)

        # Plot cells
        if with_cells:
            ra_max, rb_max, rc_max = self._dim
            viewer.plot_cells(view, (0, ra_max, 0, rb_max, 0, rc_max))

    def get_orb_map(self) -> RegularOrbitalMapVac:
        """
        Get the orbital map for indexing orbitals.

        NOTE: This function should be called everytime after the supercell
        changes. DO NOT reuse the result of last call, as its array attributes
        may be out of date.

        :return: orbital map
        """
        self._sync_array()
        return RegularOrbitalMapVac(self._dim, self.num_orb_pc,
                                    self._cached_data.vac_id_sc,
                                    self._cached_data.orb_id_pc)

    def _init_hop_gen(self, orb_pos: np.ndarray) -> hop_type:
        """
        Get initial hopping terms and distances using the general algorithm.
        :param orb_pos: (num_orb_sc, 3) float64 array
            Cartesian coordinates of orbitals in NM
        :return: (hop_ind, hop_eng, dr)
            initial hopping terms and distances
        """
        data = self._prim_cell.get_array()
        hop_ind, hop_eng, dr = \
            core.build_hop_gen(data.hop_ind, data.hop_eng,
                               self._dim, self._pbc, self.num_orb_pc,
                               self._cached_data.orb_id_pc,
                               self._cached_data.vac_id_sc,
                               self.sc_lat_vec, orb_pos)
        return hop_ind, hop_eng, dr

    def _init_hop_fast(self, orb_pos: np.ndarray) -> hop_type:
        """
        Get initial hopping terms and distance using the fast algorithm.

        NOTE: this algorithm works only for supercells without vacancies.

        :param orb_pos: (num_orb_sc, 3) float64 array
            Cartesian coordinates of orbitals in NM
        :return: (hop_ind, hop_eng, dr)
            initial hopping terms and distances
        """
        data = self._prim_cell.get_array()

        # Split pc hopping terms into free and periodic parts
        pc_ind_pbc, pc_eng_pbc, pc_ind_free, pc_eng_free = \
            core.split_pc_hop(data.hop_ind, data.hop_eng, self._pbc)

        # Build sc hopping terms from periodic parts
        # This is fast since we can predict the number of hopping terms.
        sc_ind_pbc, sc_eng_pbc, sc_dr_pbc = \
            core.build_hop_pbc(pc_ind_pbc, pc_eng_pbc,
                               self._dim, self._pbc, self.num_orb_pc,
                               self.sc_lat_vec, orb_pos)

        # Build hopping terms from free parts
        # Here we must call the general Cython function as we cannot predict
        # the number of hopping terms.
        sc_ind_free, sc_eng_free, sc_dr_free = \
            core.build_hop_gen(pc_ind_free, pc_eng_free,
                               self._dim, self._pbc, self.num_orb_pc,
                               self._cached_data.orb_id_pc,
                               self._cached_data.vac_id_sc,
                               self.sc_lat_vec, orb_pos)

        # Assemble hopping terms and distances
        if len(sc_ind_pbc) > 0:
            hop_ind = np.concatenate([sc_ind_pbc, sc_ind_free])
            hop_eng = np.concatenate([sc_eng_pbc, sc_eng_free])
            dr = np.concatenate([sc_dr_pbc, sc_dr_free])
        else:
            hop_ind = sc_ind_free
            hop_eng = sc_eng_free
            dr = sc_dr_free
        return hop_ind, hop_eng, dr

    def get_array(self,
                  with_orb_eng: bool = True,
                  with_orb_pos: bool = True,
                  with_hop: bool = True) -> namedtuple:
        """
        Get all arrays in a single call.
        :param with_orb_eng: whether to include orb_eng in result
        :param with_orb_pos: whether to include orb_pos in result
        :param with_hop: whether to include hop_ind, hop_end and dr in result
        :return: namedtuple with arrays as attributes
        """
        self._sync_array()

        # Default value
        orb_eng = np.array([], dtype=np.float64)
        orb_pos = np.array([], dtype=np.float64)
        hop_ind = np.array([], dtype=np.int64)
        hop_eng = np.array([], dtype=np.complex128)
        dr = np.array([], dtype=np.float64)

        # orb_eng
        if with_orb_eng:
            orb_eng = core.build_orb_eng(self._prim_cell.orb_eng,
                                         self._cached_data.orb_id_pc)

        # orb_pos
        if with_orb_pos or with_hop:  # dr requires orb_pos
            orb_pos = core.build_orb_pos(self._prim_cell.lat_vec,
                                         self._prim_cell.orb_pos,
                                         self._cached_data.orb_id_pc)
            orb_pos += self._prim_cell.origin
            if self._orb_pos_modifier is not None:
                self._orb_pos_modifier(orb_pos)

        # hop_*
        if with_hop:
            # Get initial hopping terms
            if self.num_vac == 0:
                hop_ind, hop_eng, dr = self._init_hop_fast(orb_pos)
            else:
                hop_ind, hop_eng, dr = self._init_hop_gen(orb_pos)

            # Apply hopping modifier
            if self._hop_modifier.num_hop != 0:
                ind_mod, eng_mod = self._hop_modifier.to_array(use_int64=True)
                hop_dict = core.CPPHopDict()
                hop_dict.add_hop_origin(hop_ind, hop_eng, dr)
                hop_dict.add_hop_modifier(ind_mod, eng_mod, self.sc_lat_vec,
                                          orb_pos)
                hop_ind, hop_eng, dr = hop_dict.to_array()

        # Assemble and return
        Data = namedtuple("Data",
                          ["orb_eng", "orb_pos", "hop_ind", "hop_eng", "dr"])
        data = Data(orb_eng, orb_pos, hop_ind, hop_eng, dr)
        return data

    def get_orb_eng(self) -> np.ndarray:
        """
        Get energies of all orbitals in the supercell.
        :return: (num_orb_sc,) float64 array
            on-site energies of orbitals in the supercell in eV
        """
        data = self.get_array(True, False, False)
        return data.orb_eng

    def get_orb_pos(self) -> np.ndarray:
        """
        Get positions of all orbitals in the supercell.
        :return: (num_orb_sc, 3) float64 array
            Cartesian coordinates of orbitals in the supercell in nm
        """
        data = self.get_array(False, True, False)
        return data.orb_pos

    def get_hop(self) -> hop_type:
        """
        Get hopping indices, energies and distances.

        The hopping terms will be reduced by conjugate relation. So only half
        of them will be returned as results.

        If periodic conditions are enabled, orbital indices in hop_ind[:, 1] may
        be wrapped back if it falls out of the supercell. Nevertheless, the
        distances in dr are still the ones before wrapping. This is essential
        for adding magnetic field, calculating band structure and many
        properties involving dr.

        :return: (hop_ind, hop_eng, dr)
            hop_ind: (num_hop_sc, 2) int64 array
            row and column indices of hopping terms reduced by conjugate relation
            hop_eng: (num_hop_sc,) complex128 array
            energies of hopping terms in accordance with hop_ind in eV
            dr: (num_hop_sc, 3) float64 array
            distances of hopping terms in accordance with hop_ind in nm
        """
        data = self.get_array(False, False, True)
        return data.hop_ind, data.hop_eng, data.dr

    def get_lattice_area(self, direction: str = "c") -> float:
        """
        Get the area formed by lattice vectors normal to given direction.
        :param direction: direction of area, e.g. "c" indicates the area formed
            by lattice vectors in the aOb plane.
        :return: area formed by lattice vectors in NM^2
        """
        assert direction in ("a", "b", "c"), f"Wrong direction {direction}"
        return lat.get_lattice_area(self.sc_lat_vec, direction)

    def get_lattice_volume(self) -> float:
        """
        Get the volume formed by all three lattice vectors in NM^3.
        :return: volume in NM^3
        """
        return lat.get_lattice_volume(self.sc_lat_vec)

    def get_reciprocal_vectors(self) -> np.ndarray:
        """
        Get the Cartesian coordinates of reciprocal lattice vectors in 1/NM.
        :return: (3, 3) float64 array
            reciprocal vectors in 1/NM
        """
        return lat.gen_reciprocal_vectors(self.sc_lat_vec)

    @property
    def num_orb_pc(self) -> int:
        """
        Get the number of orbitals of primitive cell.
        :return: number of orbitals in primitive cell.
        """
        return self._prim_cell.num_orb

    @property
    def num_orb_sc(self) -> int:
        """
        Get the number of orbitals of supercell.
        :return: number of orbitals in supercell
        """
        return self._cached_data.orb_id_pc.shape[0]

    @property
    def num_vac(self) -> int:
        """
        Get the number of vacancies of supercell.
        :return: number of vacancies in supercell
        """
        return len(self._vacancy_set)

    @property
    def prim_cell(self) -> PrimitiveCell:
        """
        Interface for the '_prim_cell' attribute.
        :return: the primitive cell
        """
        return self._prim_cell

    @property
    def sc_lat_vec(self) -> np.ndarray:
        """
        Get the lattice vectors of supercell.
        :return: (3, 3) float64 array
            lattice vectors of primitive cell in nm.
        """
        sc_lat_vec = self._prim_cell.lat_vec.copy()
        for i in range(3):
            sc_lat_vec[i] *= self._dim.item(i)
        return sc_lat_vec


class SCInterHopping(InterHopping):
    """
    Container class for hopping terms between different supercells within the
    sample.

    Attributes
    ----------
    _sc_bra: 'SuperCell' instance
        the 'bra' supercell from which the hopping terms exist
    _sc_ket: 'SuperCell' instance
        the 'ket' supercell to which the hopping terms exist

    NOTES
    -----
    1. Reduction

    Since inter-hopping terms exist within different super-cells, there is no
    need to reduce them according to the conjugate relation.

    2. Rules

    We restrict hopping terms to be from the (0, 0, 0) 'bra' supercell to any
    'ket' supercell. The counterparts are restored via the conjugate relation:
        <bra, R0, i|H|ket, Rn, j> = <ket, R0, j|H|bra, -Rn, i>*
    """
    def __init__(self, sc_bra: SuperCell, sc_ket: SuperCell) -> None:
        """
        :param sc_bra: the 'bra' supercell from which the hopping terms exist
        :param sc_ket: the 'ket' supercell to which the hopping terms exist
        """
        super().__init__(sc_bra, sc_ket)
        self._sc_bra = sc_bra
        self._sc_ket = sc_ket

    def get_array(self,
                  pos_bra: np.ndarray = None,
                  pos_ket: np.ndarray = None,
                  check_dup: bool = False) -> namedtuple:
        """
        Get all arrays in a single call.

        NOTE: If this function is called by Sample.init_array(), then pos_bra
        and pos_ket are external and passed as arguments. In other cases, they
        are generated in place.

        :param pos_bra: (num_orb_bra, 3) float64 array
            Cartesian coordinates of sc_bra in NM
        :param pos_ket: (num_orb_ket, 3) float64 array
            Cartesian coordinates of sc_ket in NM
        :param check_dup: whether to check for duplicate hopping terms
        :return: namedtuple with arrays as attributes
        """
        assert self.num_hop > 0, f"Num_hop should > 0"

        # Check for duplicate hopping terms
        if check_dup:
            pair_count = self.count_pairs()
            for pair, count in pair_count.items():
                assert count <= 1, f"Duplicate terms detected {pair}"

        # Build hop_*
        hop_ind, hop_eng = self.to_array(use_int64=True)

        # Build dr
        if pos_bra is None:
            pos_bra = self._sc_bra.get_orb_pos()
        if pos_ket is None:
            pos_ket = self._sc_ket.get_orb_pos()
        dr = core.build_inter_dr(hop_ind, pos_bra, pos_ket,
                                 self._sc_ket.sc_lat_vec)

        # Drop rn in hop_ind
        hop_ind = np.copy(hop_ind[:, 3:5], order="C")
        Data = namedtuple("Data", ["hop_ind", "hop_eng", "dr"])
        data = Data(hop_ind, hop_eng, dr)
        return data

    def get_hop(self, check_dup: bool = False) -> hop_type:
        """
        Get hopping indices, energies and distances.

        If periodic conditions are enabled, orbital indices in hop_ind[:, 1] may
        be wrapped back if it falls out of the supercell. Nevertheless, the
        distances in dr are still the ones before wrapping. This is essential
        for adding magnetic field, calculating band structure and many
        properties involving dr.

        :param check_dup: whether to check for duplicate hopping terms
        :return: (hop_ind, hop_eng, dr)
            hop_ind: (num_hop_sc, 2) int64 array
            row and column indices of hopping terms
            hop_eng: (num_hop_sc,) complex128 array
            energies of hopping terms in accordance with hop_ind in eV
            dr: (num_hop_sc, 3) float64 array
            distances of hopping terms in accordance with hop_ind in nm
        """
        data = self.get_array(check_dup=check_dup)
        return data.hop_ind, data.hop_eng, data.dr

    def plot(self,
             axes: plt.Axes,
             with_conj: bool = False,
             hop_as_arrows: bool = True,
             hop_eng_cutoff: float = 1e-5,
             hop_color: str = "r",
             view: str = "ab") -> None:
        """
        Plot hopping terms to axes.
        :param axes: axes on which the figure will be plotted
        :param with_conj: whether to plot conjugate hopping terms as well
        :param hop_as_arrows: whether to plot hopping terms as arrows
        :param hop_eng_cutoff: cutoff for showing hopping terms.
        :param hop_color: color of hopping terms
        :param view: kind of view point, should be in 'ab', 'bc', 'ca', 'ba',
            'cb', 'ac'
        :return: None
        """
        viewer = ModelViewer(axes, self._sc_bra.prim_cell.lat_vec,
                             self._sc_bra.prim_cell.origin, view)

        # Plot hopping terms
        orb_pos_bra = self._sc_bra.get_orb_pos()
        orb_pos_ket = self._sc_ket.get_orb_pos() if with_conj else None
        hop_ind, hop_eng, dr = self.get_hop()
        viewer.plot_hop(orb_pos_bra, orb_pos_ket, hop_ind, hop_eng, dr,
                        with_conj, hop_as_arrows, hop_eng_cutoff, hop_color)

    @property
    def sc_bra(self) -> SuperCell:
        """Interface for the '_sc_bra' attribute."""
        return self._sc_bra

    @property
    def sc_ket(self) -> SuperCell:
        """Interface for the '_sc_ket' attribute."""
        return self._sc_ket
