"""Functions and classes for manipulating the primitive cell."""

from typing import Hashable, Union, Iterable, Callable, List, Tuple, Dict
from collections import namedtuple, defaultdict
from math import cos, sin

import numpy as np
import matplotlib.pyplot as plt

from tbplas.base import consts
from tbplas.base import lattice as lat
from tbplas.builder.base import (Lockable, verify_rn, verify_pos, pos_type,
                                 rn_type, rn3_type, pbc3_type, pair_type)
from tbplas.builder.hopping import IntraHopping, InterHopping, HopDict
from tbplas.builder.visual import ModelViewer


# DO NOT use dataclass
Orbital = namedtuple("Orbital", ["position", "energy", "label"])


class PCCachedData:
    """
    Cached array attributes of PrimitiveCell.

    NOTE: DO NOT USE dataclass!

    Attributes
    ----------
    orb_eng: (num_orb, ) float64 array
        on-site energies of all orbitals in eV
    orb_pos: (num_orb, 3) float64 array
        FRACTIONAL positions of all orbitals
    hop_ind: (num_hop, 5) int32 array
        indices of all hopping terms
    hop_eng: (num_hop, ) complex128 array
        energies of all hopping terms in eV
    hash_dict: Dict[str, int]
        hashes of attributes to be used by 'sync_*' methods to update the arrays
    """
    def __init__(self):
        self.orb_eng = np.array([], dtype=np.float64)
        self.orb_pos = np.array([], dtype=np.float64)
        self.hop_ind = np.array([], dtype=np.int32)
        self.hop_eng = np.array([], dtype=np.complex128)
        self.hash_dict = {"orb": None, "hop": None}


class PCHopNotFoundError(Exception):
    """Exception for a non-existing Hopping term."""
    def __init__(self, hop_ind):
        super().__init__()
        self._hop_ind = hop_ind

    def __str__(self):
        return f"hopping term {self._hop_ind} not found"


def f_dot_s(f: np.ndarray, s: List[str], eps: float = 1.0e-5) -> str:
    """
    Calculate the inner product between a real vector and a list of symbols.
    :param f: real vector
    :param s: list of symbols
    :param eps: threshold for comparing floats
    :return: the symbolic inner product
    """
    assert len(f) == len(s)
    result = ""
    for i in range(len(f)):
        fi, si = f.item(i), s[i]
        f_abs = abs(fi)
        if f_abs > eps:
            # Get term without sign
            if abs(f_abs - 1) <= eps:
                term = f"{si}"
            else:
                term = f"{f_abs} * {si}"
            # Add sign depending on position
            sign = "-" if fi < 0 else "+"
            if result == "":
                result += f"{sign}{term}"
            else:
                result += f" {sign} {term}"
    result = result.lstrip().lstrip("+").lstrip()
    return result


class PrimitiveCell(Lockable):
    """
    Class for representing the primitive cell.

    Attributes
    ----------
    _lat_vec: (3, 3) float64 array
        Cartesian coordinates of lattice vectors in NANO METER
        Each ROW corresponds to one lattice vector.
    _origin: (3, ) float64 array
        Cartesian coordinates of origin of lattice vectors in NANO METER
    _orbital_list: List[Orbital]
        list of orbitals in the primitive cell
    _hopping_dict: 'IntraHopping' instance
        container of hopping terms in the primitive cell
        Only half of the hopping terms are stored. Conjugate terms are added
        automatically when constructing the Hamiltonian.
    _cached_data: 'PCCachedData'
        cached array attributes
    """
    def __init__(self,
                 lat_vec: np.ndarray = np.eye(3, dtype=np.float64),
                 origin: np.ndarray = np.zeros(3, dtype=np.float64),
                 unit: float = consts.ANG) -> None:
        """
        :param lat_vec: (3, 3) float64 array
            Cartesian coordinates of lattice vectors in arbitrary unit
        :param origin: (3, ) float64 array
            Cartesian coordinate of lattice origin in arbitrary unit
        :param unit: conversion coefficient from arbitrary unit to NM
        :return: None
        """
        super().__init__()

        # Setup lattice vectors and origin
        # Use copy to ensure that they are not referenced to external data,
        # for consistence with C++ version.
        lat_vec = np.array(lat_vec, dtype=np.float64).copy()
        origin = np.array(origin, dtype=np.float64).copy()
        if __debug__:
            error_msg = f"Lat_vec {lat_vec} should have shape of (3, 3)"
            assert lat_vec.shape == (3, 3), error_msg
            volume = lat.get_lattice_volume(lat_vec)
            error_msg = f"Volume of lat_vec {volume} should > 0"
            assert volume > consts.INFINITESIMAL, error_msg
            error_msg = f"Length of origin {origin} should be 3"
            assert origin.shape == (3,), error_msg
            assert unit > 0.0, f"Unit {unit} should > 0"
        self._lat_vec = lat_vec * unit
        self._origin = origin * unit

        # Setup orbitals and hopping terms
        self._orbital_list = []
        self._hopping_dict = IntraHopping()

        # Setup cached arrays
        self._cached_data = PCCachedData()
        for key in ("orb", "hop"):
            self._cached_data.hash_dict[key] = self._get_attr_hash(key)

    def __hash__(self) -> int:
        lat_vec = tuple([tuple(_) for _ in self._lat_vec])
        fp = (lat_vec, tuple(self._origin), tuple(self._orbital_list),
              self._hopping_dict)
        return hash(fp)

    def _get_attr_hash(self, attr: str) -> int:
        """
        Get the hash of given attribute.
        :param attr: name of the attribute
        :return: hash of the attribute
        """
        assert attr in ("orb", "hop")
        if attr == "orb":
            new_hash = hash(tuple(self._orbital_list))
        else:
            new_hash = hash(self._hopping_dict)
        return new_hash

    def _update_attr_hash(self, attr: str) -> bool:
        """
        Compare and update the 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 _sync_array(self, force_sync: bool = False) -> None:
        """
        Synchronize orb_eng, orb_pos, hop_ind and hop_eng according to orbitals
        and hopping terms.
        :param force_sync: whether to force synchronizing the arrays even if the
            orbitals or hopping terms did not change
        :return: None
        """
        to_update = self._update_attr_hash('orb')
        if force_sync or to_update:
            self._cached_data.orb_eng = np.array(
                [orb.energy for orb in self._orbital_list], dtype=np.float64)
            self._cached_data.orb_pos = np.array(
                [orb.position for orb in self._orbital_list], dtype=np.float64)
        to_update = self._update_attr_hash('hop')
        if force_sync or to_update:
            hop_ind, hop_eng = self._hopping_dict.to_array(use_int64=False)
            self._cached_data.hop_ind = hop_ind
            self._cached_data.hop_eng = hop_eng

    def add_orbital(self,
                    position: pos_type,
                    energy: float = 0.0,
                    label: Hashable = "X") -> None:
        """
        Add a new orbital to the primitive cell.
        :param position: FRACTIONAL coordinate of the orbital
        :param energy: on-site energy of the orbital in eV
        :param label: orbital label
        :return: None
        """
        if __debug__:
            self.check_editable()
        position = verify_pos(position)
        self._orbital_list.append(Orbital(position, energy, label))

    def add_orbital_cart(self,
                         position: pos_type,
                         unit: float = consts.ANG,
                         energy: float = 0.0,
                         label: Hashable = "X") -> None:
        """
        Add a new orbital to the primitive cell in CARTESIAN coordinates.
        :param position: Cartesian coordinate of orbital in arbitrary unit
        :param unit: conversion coefficient from arbitrary unit to NM
        :param energy: on-site energy of the orbital in eV
        :param label: orbital label
        :return: None
        """
        if __debug__:
            self.check_editable()
            assert unit > 0.0, f"Unit {unit} should > 0"
        position = verify_pos(position)
        pos_cart = np.array(position) * unit
        pos_frac = lat.cart2frac_s(self._lat_vec, pos_cart, self._origin)
        position = tuple(pos_frac)
        self._orbital_list.append(Orbital(position, energy, label))

    @staticmethod
    def _check_orbital_index(orb_i: int, num_orb: int) -> None:
        """
        Checks if orbital index falls out of range.
        :param orb_i: orbital index
        :param num_orb: total number of orbitals
        :return: None
        """
        error_msg = f"Orbital index {orb_i} out of range(0, {num_orb})"
        assert 0 <= orb_i < num_orb, error_msg

    def set_orbital(self,
                    orb_i: int,
                    position: pos_type = None,
                    energy: float = None,
                    label: Hashable = None) -> None:
        """
        Modify the position, energy and label of an existing orbital.

        If position, energy or label is None, then the corresponding attribute
        will not be modified.

        :param orb_i: index of the orbital to modify
        :param position: new FRACTIONAL coordinate of the orbital
        :param energy: new on-site energy of the orbital in eV
        :param label: new orbital label
        :return: None
        """
        if __debug__:
            self.check_editable()
            self._check_orbital_index(orb_i, self.num_orb)
        orbital = self._orbital_list[orb_i]
        if position is None:
            position = orbital.position
        else:
            position = verify_pos(position)
        if energy is None:
            energy = orbital.energy
        if label is None:
            label = orbital.label
        self._orbital_list[orb_i] = Orbital(position, energy, label)

    def set_orbital_cart(self,
                         orb_i: int,
                         position: pos_type = None,
                         unit: float = consts.ANG,
                         energy: float = None,
                         label: Hashable = None) -> None:
        """
        Modify the position, energy and label of an existing orbital with
        Cartesian coordinates.

        If position, energy or label is None, then the corresponding attribute
        will not be modified.

        :param orb_i: index of the orbital to modify
        :param position: Cartesian coordinate of orbital in arbitrary unit
        :param unit: conversion coefficient from arbitrary unit to NM
        :param energy: new on-site energy of the orbital in eV
        :param label: new orbital label
        :return: None
        """
        if __debug__:
            self.check_editable()
            self._check_orbital_index(orb_i, self.num_orb)
            assert unit > 0.0, f"Unit {unit} should > 0"
        orbital = self._orbital_list[orb_i]
        if position is None:
            position = orbital.position
        else:
            position = verify_pos(position)
            pos_cart = np.array(position) * unit
            pos_frac = lat.cart2frac_s(self._lat_vec, pos_cart, self._origin)
            position = tuple(pos_frac)
        if energy is None:
            energy = orbital.energy
        if label is None:
            label = orbital.label
        self._orbital_list[orb_i] = Orbital(position, energy, label)

    def get_orbital(self, orb_i: int) -> Orbital:
        """
        Get given orbital instance.
        :param orb_i: index of the orbital
        :return: the i-th orbital
        """
        if __debug__:
            self.check_editable()
            self._check_orbital_index(orb_i, self.num_orb)
        return self._orbital_list[orb_i]

    def remove_orbitals(self, indices: Union[Iterable[int], np.ndarray]) -> None:
        """
        Remove given orbitals and associated hopping terms, then update
        remaining hopping terms.
        :param indices: indices of orbitals to remove
        :return: None
        """
        # Check arguments
        if __debug__:
            self.check_editable()
            num_orb = self.num_orb
            for orb_i in indices:
                self._check_orbital_index(orb_i, num_orb)

        # Delete the orbitals
        indices = sorted(set(indices))
        for i, orb_i in enumerate(indices):
            self._orbital_list.pop(orb_i - i)

        # Delete associated hopping terms
        self._hopping_dict.remove_orbitals(indices)

    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
        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 rn != (0, 0, 0) or 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 hopping term to the primitive cell, or update an existing
        hopping term.
        :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._hopping_dict.add_hopping(rn, orb_i, orb_j, energy)

    def add_hopping_dict(self, hop_dict: HopDict) -> None:
        """
        Add a matrix of hopping terms to the primitive cell, or update existing
        hopping terms.

        Reserved for compatibility with old version of TBPLaS.

        :param hop_dict: hopping dictionary
        :return: None
        """
        if __debug__:
            self.check_editable()
        hop_dict_sparse = hop_dict.to_spare()
        for rn, hop_mat in hop_dict_sparse.items():
            for i in range(hop_mat.nnz):
                orb_i = hop_mat.row.item(i)
                orb_j = hop_mat.col.item(i)
                hop_eng = hop_mat.data[i]
                self.add_hopping(rn, orb_i, orb_j, hop_eng)

    def get_hopping(self,
                    rn: rn_type,
                    orb_i: int,
                    orb_j: int) -> complex:
        """
        Get given hopping term or its conjugate counterpart.

        If not found, an exception will be raised. This is an on-purpose design,
        since the caller has no obligation to know what terms are in the cell.
        A LOT of functions and scripts rely on this behavior. DO NOT change it.

        :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: hopping energy in eV
        :raises PCHopNotFoundError: if the hop term or its conjugate counterpart
            is not found
        """
        rn = verify_rn(rn)
        if __debug__:
            self._check_hop_index(rn, orb_i, orb_j)
        energy, found = self._hopping_dict.get_hopping(rn, orb_i, orb_j)
        if found:
            return energy
        else:
            raise PCHopNotFoundError(rn + (orb_i, orb_j))

    def remove_hopping(self,
                       rn: rn_type,
                       orb_i: int,
                       orb_j: int) -> None:
        """
        Remove given hopping term or its conjugate counterpart.
        :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
        :raises PCHopNotFoundError: if the hop term or its conjugate counterpart
            is not found
        """
        rn = verify_rn(rn)
        if __debug__:
            self.check_editable()
            self._check_hop_index(rn, orb_i, orb_j)
        found = self._hopping_dict.remove_hopping(rn, orb_i, orb_j)
        if not found:
            raise PCHopNotFoundError(rn + (orb_i, orb_j))

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

        # Count the number of hopping terms of each orbital
        hop_count = defaultdict(int)
        for rn, hop_rn in self._hopping_dict.hoppings.items():
            for pair, energy in hop_rn.items():
                hop_count[pair[0]] += 1
                hop_count[pair[1]] += 1

        # Get indices of orbitals to remove
        orb_id_trim = [i_o for i_o, count in hop_count.items() if count <= 1]

        # Remove orbitals and hopping terms
        self.remove_orbitals(orb_id_trim)

    def apply_pbc(self, pbc: pbc3_type = (True, True, True)) -> None:
        """
        Apply periodic boundary conditions by removing hopping terms between
        cells along non-periodic direction.
        :param pbc: whether pbc is enabled along 3 directions
        :return: None
        """
        if __debug__:
            self.check_editable()
        assert len(pbc) == 3, f"Length of {pbc} is not 3"

        # Get the list of hopping terms to keep
        rn_remove = set()
        for rn in self._hopping_dict.hoppings.keys():
            to_keep = True
            for i_dim in range(3):
                if not pbc[i_dim] and rn[i_dim] != 0:
                    to_keep = False
                    break
            if not to_keep:
                rn_remove.add(rn)
        for rn in rn_remove:
            self._hopping_dict.remove_rn(rn)

    def reset_lattice(self,
                      lat_vec: np.ndarray = None,
                      origin: np.ndarray = None,
                      unit: float = consts.ANG,
                      fix_orb: bool = True) -> None:
        """
        Reset lattice vectors and origin.
        :param lat_vec: (3, 3) float64 array
            Cartesian coordinates of lattice vectors in arbitrary unit
        :param origin: (3, ) float64 array
            Cartesian coordinate of lattice origin in arbitrary unit
        :param unit: conversion coefficient from arbitrary unit to NM
        :param fix_orb: whether to keep Cartesian coordinates of orbitals
            unchanged after resetting lattice
        :return: None
        """
        if __debug__:
            self.check_editable()
            assert unit > 0.0, f"Unit {unit} should > 0"

        # Check lattice vectors and origin
        # Use copy to ensure that they are not referenced to external data,
        # for consistence with C++ version.
        if lat_vec is None:
            lat_vec = self._lat_vec.copy()
        else:
            lat_vec = np.array(lat_vec, dtype=np.float64).copy() * unit
        if origin is None:
            origin = self._origin.copy()
        else:
            origin = np.array(origin, dtype=np.float64).copy() * unit
        if __debug__:
            error_msg = f"Lat_vec {lat_vec} should have shape of (3, 3)"
            assert lat_vec.shape == (3, 3), error_msg
            volume = lat.get_lattice_volume(lat_vec)
            error_msg = f"Volume of lat_vec {volume} should > 0"
            assert volume > consts.INFINITESIMAL, error_msg
            error_msg = f"Length of origin {origin} should be 3"
            assert origin.shape == (3,), error_msg

        # Update orbital positions if needed
        if fix_orb and self.num_orb > 0:
            orb_pos = lat.cart2frac(lat_vec, self.orb_pos_nm, origin)
            for i, pos in enumerate(orb_pos):
                orbital = self._orbital_list[i]
                self._orbital_list[i] = Orbital(tuple(pos), orbital.energy,
                                                orbital.label)

        # Finally update the attributes
        # This must be done after updating orbital positions
        self._lat_vec = lat_vec
        self._origin = origin

    def rotate(self,
               angle: float = 0.0,
               center: np.ndarray = np.zeros(3),
               shift: float = 0.) -> None:
        """
        Rotate and shift primitive cell with respect to z-axis.
        :param angle: twisting angle in RADIANs, NOT degrees
        :param center: (3, ) float64 array
            Cartesian coordinates of the rotation center in NANOMETER
        :param shift: distance of shift in NANOMETER
        :return: None
        """
        # Get rotated lattice vectors and origin
        end_points = np.vstack((np.zeros(3), self._lat_vec)) + self._origin
        end_points = lat.rotate_coord(end_points, angle=angle, center=center)
        lat_vec = end_points[1:] - end_points[0]
        origin = end_points[0]

        # Reset lattice vectors and origin without fixing the orbitals
        self.reset_lattice(lat_vec, origin, unit=consts.NM, fix_orb=False)

        # Shift orbitals along z-axis
        orb_pos = self.orb_pos_nm + np.array([0, 0, shift])

        # Update orbital positions in fractional coordinates
        orb_pos = lat.cart2frac(lat_vec, orb_pos, origin)
        for i, pos in enumerate(orb_pos):
            self.set_orbital(i, position=tuple(pos))

    def check_dim(self) -> None:
        """
        Checks if a 1x1x1 supercell can be built from this primitive cell by
        searching for duplicate hopping terms.

        We require that the hopping terms should not contain diagonal terms
        (i==j) since they are redundant with the on-site terms in the sparse
        Hamiltonian and hopping distances. Yet, the off-diagonal terms should
        appear only once. Although these requirements are unphysical, they
        SIGNIFICANTLY simplify the functions for constructing the sparse
        matrices. For large models, these requirements are typically satisfied.

        :return: None
        """
        pair_count = self._hopping_dict.count_pairs()
        # Add contributions from on-site terms
        for i in range(self.num_orb):
            pair_count[(i, i)] += 1
        for pair, count in pair_count.items():
            assert count <= 1, f"Duplicate terms detected {pair}"

    def apply_magnetic_field(self, intensity: float, gauge: int = 0) -> None:
        """
        Apply magnetic field perpendicular to xOy-plane to -z direction via
        Peierls substitution.

        Reference:
        [1] https://journals.aps.org/prb/pdf/10.1103/PhysRevB.51.4940
        [2] https://journals.jps.jp/doi/full/10.7566/JPSJ.85.074709

        :param intensity: magnetic B field in Tesla
        :param gauge: gauge of vector potential of the magnetic field to -z
            0 for (By, 0, 0), 1 for (0, -Bx, 0), 2 for (0.5By, -0.5Bx, 0)
        :return: None
        """
        assert gauge in (0, 1, 2), f"Illegal gauge {gauge}"
        self._sync_array()
        orb_pos_nm = self.orb_pos_nm
        factor = consts.PEIERLS_FACTOR
        for rn, hop_rn in self._hopping_dict.hoppings.items():
            rn_cart = np.array(rn) * self._lat_vec
            for pair, energy in hop_rn.items():
                orb_i, orb_j = pair
                dr_cart = rn_cart + orb_pos_nm[orb_j] - orb_pos_nm[orb_i]
                dx = dr_cart.item(0)
                dy = dr_cart.item(1)
                sx = orb_pos_nm.item(orb_j, 0) + orb_pos_nm.item(orb_i, 0)
                sy = orb_pos_nm.item(orb_j, 1) + orb_pos_nm.item(orb_i, 1)
                if gauge == 0:
                    phase = factor * intensity * dx * sy
                elif gauge == 1:
                    phase = -factor * intensity * dy * sx
                else:
                    phase = factor * 0.5 * intensity * (dx * sx - dy * sy)
                exp_phase = cos(phase) + 1j * sin(phase)
                self._hopping_dict.add_hopping(rn, orb_i, orb_j, energy*exp_phase)

    def plot(self,
             fig_name: str = None,
             fig_size: Tuple[float, float] = None,
             fig_dpi: int = 300,
             with_orbitals: bool = True,
             with_cells: bool = True,
             with_conj: bool = True,
             orb_color: Callable[[List[Orbital]], 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.

        If figure name is given, save the figure to file. Otherwise, show it on
        the screen.

        :param fig_name: file name to which the figure will be saved
        :param fig_size: width and height of the figure
        :param fig_dpi: resolution of the figure file
        :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
            If true, hopping terms will be plotted as arrows using axes.arrow()
            method. Otherwise, they will be plotted as lines using
            LineCollection. The former is more intuitive but much slower.
        :param hop_eng_cutoff: cutoff for showing hopping terms
        :param hop_color: color of hopping terms
        :param view: kind of view point
        :returns: None
        """
        self._sync_array()
        fig, axes = plt.subplots(figsize=fig_size)
        axes.set_aspect('equal')
        viewer = ModelViewer(axes, self._lat_vec, self._origin, view)

        # Determine the range of rn
        rn_range = np.zeros((3, 2), dtype=np.int32)
        if self.num_hop > 0:
            for i in range(3):
                ri_min = self._cached_data.hop_ind[:, i].min()
                ri_max = self._cached_data.hop_ind[:, i].max()
                if with_conj:
                    rn_range[i, 0] = min([ri_min, ri_max, -ri_min, -ri_max])
                    rn_range[i, 1] = max([ri_min, ri_max, -ri_min, -ri_max])
                else:
                    rn_range[i, 0] = ri_min
                    rn_range[i, 1] = ri_max
        ra_min, ra_max = rn_range.item(0, 0), rn_range.item(0, 1)
        rb_min, rb_max = rn_range.item(1, 0), rn_range.item(1, 1)
        rc_min, rc_max = rn_range.item(2, 0), rn_range.item(2, 1)

        # Plot orbitals
        orb_pos = self.orb_pos_nm
        if orb_color is None:
            scatter_color = self._cached_data.orb_eng
        else:
            scatter_color = orb_color(self._orbital_list)
        if self.num_orb > 0:
            if with_orbitals:
                for i_a in range(ra_min, ra_max+1):
                    for i_b in range(rb_min, rb_max+1):
                        for i_c in range(rc_min, rc_max+1):
                            center = np.matmul((i_a, i_b, i_c), self._lat_vec)
                            pos_rn = orb_pos + center
                            viewer.scatter(pos_rn, s=100, c=scatter_color)

        # Plot hopping terms
        if self.num_hop > 0:
            hop_ind = self._cached_data.hop_ind[:, 3:]
            hop_eng = self._cached_data.hop_eng
            dr = self.dr_nm
            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:
            viewer.plot_cells(view, (ra_min, ra_max+1,
                                     rb_min, rb_max+1,
                                     rc_min, rc_max+1))

        # Hide spines and ticks.
        for key in ("top", "bottom", "left", "right"):
            axes.spines[key].set_visible(False)
        axes.set_xticks([])
        axes.set_yticks([])
        fig.tight_layout()
        plt.autoscale()
        if fig_name is not None:
            plt.savefig(fig_name, dpi=fig_dpi)
        else:
            plt.show()
        plt.close()

    def print(self) -> None:
        """
        Print orbital and hopping information.
        :return: None
        """
        print("Lattice origin (nm):")
        for v in self._origin:
            print(f"{v:10.5f}", end="")
        print()
        print("Lattice vectors (nm):")
        for v in self._lat_vec:
            print(f"{v[0]:10.5f}{v[1]:10.5f}{v[2]:10.5f}")
        print("Orbitals:")
        for orbital in self._orbital_list:
            pos = orbital.position
            output = f"{pos[0]:10.5f}{pos[1]:10.5f}{pos[2]:10.5f}"
            output += f"{orbital.energy:10.5f}"
            print(output)
        print("Hopping terms:")
        for rn, hop_rn in self._hopping_dict.hoppings.items():
            for pair, energy in hop_rn.items():
                print(" ", rn, pair, energy)

    def print_hk(self, convention: int = 1) -> None:
        """
        Print analytical Hamiltonian as the function of k-point.
        :param convention: convention for setting up the Hamiltonian
        :return: None
        """
        self._sync_array()
        assert self.num_orb > 0, f"Num_orb should > 0"
        assert self.num_hop > 0, f"Num_hop should > 0"
        assert convention in (1, 2), f"Illegal convention {convention}"

        # Function to add a hopping term to the Hamiltonian
        def _add_hop(_pair, _energy, _kdr):
            if _kdr == "":
                term = f"{_energy} * 1\n"
            else:
                term = f"{_energy} * exp_i({_kdr})\n"
            try:
                hk[_pair] += f" + {term}"
            except KeyError:
                hk[_pair] = term

        # Collect on-site energies
        hk = dict()
        for i, energy in enumerate(self._cached_data.orb_eng):
            hk[(i, i)] = str(energy)

        # Collect hopping terms
        dr = self.dr if convention == 1 else self._cached_data.hop_ind[:, :3]
        k_labels = ["ka", "kb", "kc"]
        for ih, hop in enumerate(self._cached_data.hop_ind):
            pair = (hop.item(3), hop.item(4))
            energy = self._cached_data.hop_eng.item(ih)
            _add_hop(pair, energy, f_dot_s(dr[ih], k_labels))
            # Conjugate terms should be added explicitly, otherwise the
            # diagonal terms will be wrong!
            pair = (pair[1], pair[0])
            energy = energy.conjugate()
            _add_hop(pair, energy, f_dot_s(-dr[ih], k_labels))

        # Print formulae
        reduced_pairs = [p for p in hk.keys() if p[0] <= p[1]]
        for pair in reduced_pairs:
            ham_ij = f"ham[{pair[0]}, {pair[1]}]"
            formula = hk[pair].rstrip("\n")
            print(f"{ham_ij} = ({formula})")
        for pair in reduced_pairs:
            if pair[0] != pair[1]:
                ham_ij = f"ham[{pair[0]}, {pair[1]}]"
                ham_ji = f"ham[{pair[1]}, {pair[0]}]"
                formula = f"{ham_ij}.conjugate()"
                print(f"{ham_ji} = {formula}")
        print("with exp_i(x) := cos(2 * pi * x) + 1j * sin(2 * pi * x)")

    def print_py(self) -> None:
        """
        Print python code for constructing the primitive cell.
        :return: None
        """
        # Lattice vectors and origin
        print("# Lattice vectors and origin.")
        print("lat_vec = np.array([")
        for i in range(3):
            print("[", end="")
            for j in range(3):
                print(" ", self._lat_vec[i, j], ",", end="")
            print("],")
        print("])")
        print(f"origin = np.array("
              f"[{self._origin[0]}, {self._origin[1]}, {self._origin[2]}]"
              f")\n")

        # Create the primitive cell and set orbitals
        # Since TBPLaS uses NM for lattice vectors and origin, the unit is
        # always 1.0.
        print("# Create the primitive cell and set orbitals.")
        print("prim_cell = PrimitiveCell(lat_vec, origin, unit=1.0)")
        for i, orbital in enumerate(self._orbital_list):
            pos = orbital.position
            eng = orbital.energy
            label = orbital.label
            print(f"prim_cell.add_orbital("
                  f"({pos[0]}, {pos[1]}, {pos[2]}), {eng}, '{label}'"
                  f")")
        print()

        # Add hopping terms
        print("# Add hopping terms.")
        for rn, hop_rn in self._hopping_dict.hoppings.items():
            for pair, eng in hop_rn.items():
                if abs(eng) > 1.0e-5:
                    print(f"prim_cell.add_hopping("
                          f"({rn[0]}, {rn[1]}, {rn[2]}),"
                          f" {pair[0]}, {pair[1]}, {eng}"
                          f")")
        print("\nreturn prim_cell")

    def print_cxx(self) -> None:
        """
        Print c++ code for constructing the primitive cell.
        :return: None
        """
        # Lattice vectors and origin
        print("// Lattice vectors and origin.")
        print("Eigen::Matrix3d lat_vec{")
        for i in range(3):
            print("{", end="")
            for j in range(2):
                print(" ", self._lat_vec[i, j], ",", end="")
            if i < 2:
                print(" ", self._lat_vec[i, 2], "},")
            else:
                print(" ", self._lat_vec[i, 2], "}")
        print("};")
        print(f"Eigen::Vector3d origin("
              f"{self._origin[0]}, {self._origin[1]}, {self._origin[2]}"
              f");\n")

        # Create the primitive cell and set orbitals
        # Since TBPLaS uses NM for lattice vectors and origin, the unit is
        # always 1.0.
        print("// Create the primitive cell and set orbitals.")
        print(f"PrimitiveCell<complex_t> prim_cell({self.num_orb}, lat_vec,"
              f" origin, 1.0);")
        label_table = []
        for orbital in self._orbital_list:
            if orbital.label not in label_table:
                label_table.append(orbital.label)
        for i, orbital in enumerate(self._orbital_list):
            pos = orbital.position
            eng = orbital.energy
            label_idx = label_table.index(orbital.label)
            print(f"prim_cell.set_orbital("
                  f"{i}, {pos[0]}, {pos[1]}, {pos[2]}, {eng}, {label_idx}"
                  f");")
        print()

        # Add hopping terms
        print("// Add hopping terms.")
        for rn, hop_rn in self._hopping_dict.hoppings.items():
            for pair, eng in hop_rn.items():
                if abs(eng) > 1.0e-5:
                    print("prim_cell.add_hopping(", end="")
                    print(f"{rn[0]}, {rn[1]}, {rn[2]}, ", end="")
                    print(f"{pair[0]}, {pair[1]}, ", end="")
                    print("{", eng.real, ", ", eng.imag, "}", end="")
                    print(");")
        print("\nreturn prim_cell;")

    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._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._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._lat_vec)

    def get_array(self) -> namedtuple:
        """
        Get all array attributes for sending to C++ extensions.
        :return: name tuple with orb_eng, orb_pos, hop_ind, hop_eng, lat_vec and
            origin as attributes
        """
        self._sync_array()
        Data = namedtuple("Data", ["orb_eng", "orb_pos", "hop_ind", "hop_eng",
                                   "lat_vec", "origin"])
        data = Data(self._cached_data.orb_eng, self._cached_data.orb_pos,
                    self._cached_data.hop_ind, self._cached_data.hop_eng,
                    self._lat_vec, self._origin)
        return data

    @property
    def lat_vec(self) -> np.ndarray:
        """
        Interface for the '_lat_vec' attribute.
        :return: (3, 3) float64 array
            Cartesian coordinates of lattice vectors in NM
        """
        return self._lat_vec

    @property
    def origin(self) -> np.ndarray:
        """
        Interface for the '_origin' attribute.
        :return: (3, ) float64 array
            Cartesian coordinate of origin in NM
        """
        return self._origin

    @property
    def orbitals(self) -> List[Orbital]:
        """
        Interface for the '_orbital_list' attribute.
        :return: list of orbitals
        """
        return self._orbital_list

    @property
    def hoppings(self) -> Dict[rn3_type, Dict[pair_type, complex]]:
        """
        Interface for the hopping terms.
        :return: hopping terms
        """
        return self._hopping_dict.hoppings

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

    @property
    def num_hop(self) -> int:
        """
        Get the number of hopping terms WITHOUT considering conjugate relation.
        :return: number of hopping terms
        """
        return self._hopping_dict.num_hop

    @property
    def orb_eng(self) -> np.ndarray:
        """
        Interface for the '_orb_eng' attribute.
        :return: (num_orb,) float64 array
            on-site energies in eV
        """
        self._sync_array()
        return self._cached_data.orb_eng

    @property
    def orb_pos(self) -> np.ndarray:
        """
        Interface for the '_orb_pos' attribute.
        :return: (num_orb, 3) float64 array
            fractional coordinates or orbitals
        """
        self._sync_array()
        return self._cached_data.orb_pos

    @property
    def hop_ind(self) -> np.ndarray:
        """
        Interface for the '_hop_ind' attribute.
        :return: (num_hop, 5) int32 array
            hopping indices
        """
        self._sync_array()
        return self._cached_data.hop_ind

    @property
    def hop_eng(self) -> np.ndarray:
        """
        Interface for the '_hop_eng' attribute.
        :return: (num_hop,) complex128 array
            hopping energies
        """
        self._sync_array()
        return self._cached_data.hop_eng

    @property
    def orb_pos_nm(self) -> np.ndarray:
        """
        Get the Cartesian coordinates of orbitals in NANOMETER.
        :return: (num_orb, 3) float64 array
            Cartesian coordinates of orbitals in NANOMETER
        """
        return lat.frac2cart(self._lat_vec, self.orb_pos, self._origin)

    @property
    def orb_pos_ang(self) -> np.ndarray:
        """
        Get the Cartesian coordinates of orbitals in ANGSTROM.
        :return: (num_orb, 3) float64 array
            Cartesian coordinates of orbitals in arbitrary NANOMETER
        """
        return self.orb_pos_nm * 10

    @property
    def dr(self) -> np.ndarray:
        """
        Get the hopping distances in FRACTIONAL coordinates.
        :return: (num_hop, 3) float64 array
            hopping distances in FRACTIONAL coordinates
        """
        self._sync_array()
        dr = np.zeros((self.num_hop, 3), dtype=np.float64)
        orb_pos = self._cached_data.orb_pos
        for i_h, ind in enumerate(self._cached_data.hop_ind):
            orb_i, orb_j = ind.item(3), ind.item(4)
            dr[i_h] = orb_pos[orb_j] + ind[0:3] - orb_pos[orb_i]
        return dr

    @property
    def dr_nm(self) -> np.ndarray:
        """
        Get the hopping distances in NM.
        :return: (num_hop, 3) float64 array
            hopping distances in NM
        """
        return lat.frac2cart(self._lat_vec, self.dr)

    @property
    def dr_ang(self) -> np.ndarray:
        """
        Get the hopping distances in ANGSTROM.
        :return: (num_hop, 3) float64 array
            hopping distances in ANGSTROM
        """
        return self.dr_nm * 10


class PCInterHopping(InterHopping):
    """
    Class for holding hopping terms between different primitive cells
    in hetero-structure.

    Attributes
    ----------
    _pc_bra: 'PrimitiveCell' instance
        the 'bra' primitive cell from which the hopping terms exist
    _pc_ket: 'PrimitiveCell' instance
        the 'ket' primitive cell from which the hopping terms exist

    NOTES
    -----
    We assume hopping terms to be from (0, 0, 0) cell of pc_bra to any cell of
    pc_ket. The counterparts can be restored via the conjugate relation:
        <pc_bra, R0, i|H|pc_ket, Rn, j> = <pc_ket, R0, j|H|pc_bra, -Rn, i>*
    """
    def __init__(self, pc_bra: PrimitiveCell, pc_ket: PrimitiveCell) -> None:
        """
        :param pc_bra: 'bra' primitive cell from which the hopping terms exist
        :param pc_ket: 'ket' primitive cell from which the hopping terms exist
        """
        super().__init__(pc_bra, pc_ket)
        self._pc_bra = pc_bra
        self._pc_ket = pc_ket

    @property
    def pc_bra(self) -> PrimitiveCell:
        """Interface for the '_pc_bra' attribute."""
        return self._pc_bra

    @property
    def pc_ket(self) -> PrimitiveCell:
        """Interface for the '_pc_ket' attribute."""
        return self._pc_ket
