"""Functions and classes for sample."""

import os
from abc import ABC, abstractmethod
import pickle
from collections import namedtuple
from typing import Union, Tuple, List, Callable

import numpy as np
import matplotlib.pyplot as plt

from tbplas.cython import sample as core
from tbplas.builder.super import SuperCell, SCInterHopping
from tbplas.builder.visual import ModelViewer


__all__ = ["Sample", "CachedSample"]


orb_color_type = Callable[[np.ndarray], List[str]]


class BaseSample(ABC):
    """
    Base class of Sample and CachedSample.

    Attributes
    ----------
    orb_eng: (num_orb_tot,) float64 array
        on-site energies of orbitals in the supercell in eV
    orb_pos: (num_orb_tot, 3) float64 array
        Cartesian coordinates of orbitals in nm
    hop_ind: (num_hop_tot, 2) int64 array
        row and column indices of hopping terms reduced by conjugate relation
    hop_eng: (num_hop_tot,) complex128 array
        energies of hopping terms in accordance with hop_ind in eV
    dr: (num_hop_tot, 3) float64 array
        distances of hopping terms in accordance with hop_ind in nm

    NOTES
    -----
    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.
    """
    def __init__(self) -> None:
        # The actual initialization should be done on-demand in init_array
        # method.
        self.orb_eng = np.array([], dtype=np.float64)
        self.orb_pos = np.array([], dtype=np.float64)
        self.hop_ind = np.array([], dtype=np.int64)
        self.hop_eng = np.array([], dtype=np.complex128)
        self.dr = np.array([], dtype=np.float64)

    @property
    def _size_prod(self) -> int:
        """Evaluate product of sizes of all arrays."""
        size_prod = (self.orb_eng.size * self.orb_pos.size *
                     self.hop_ind.size * self.hop_eng.size * self.dr.size)
        return size_prod

    @abstractmethod
    def init_array(self, force_init: bool = False) -> None:
        """
        Abstract interface for initializing the arrays.

        If any array is empty, build all arrays from scratch. Otherwise, build
        them when force_init is True.

        :param force_init: whether to force initializing the arrays from scratch
            even if they have already been initialized
        :return: None
        """
        pass

    def init_orb_eng(self, force_init: bool = False) -> None:
        """Reserved for compatibility. Use init_array instead."""
        self.init_array(force_init=force_init)

    def init_orb_pos(self, force_init: bool = False) -> None:
        """Reserved for compatibility. Use init_array instead."""
        self.init_array(force_init=force_init)

    def init_hop(self, force_init: bool = False) -> None:
        """Reserved for compatibility. Use init_array instead."""
        self.init_array(force_init=force_init)

    def reset_array(self) -> None:
        """
        Reset all arrays.
        :return: None
        """
        self.init_array(force_init=True)

    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,
            dr and sc0_lat_vec as attributes
        """
        self.init_array()
        Data = namedtuple("Data", ["orb_eng", "orb_pos", "hop_ind", "hop_eng",
                                   "dr", "sc0_lat_vec"])
        data = Data(self.orb_eng, self.orb_pos, self.hop_ind, self.hop_eng,
                    self.dr, self.sc0_lat_vec)
        return data

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

        The gauge-invariant Peierls phase from R_i to R_j is evaluated as
            exp[-1j * |e|/(2*h_bar*c) * (R_i - R_j) * (A_i + A_j)]
        or
            exp[1j * |e|/(2*h_bar*c) * (R_j - R_i) * (A_j + A_i)]
        For the reference, see eqn. 10-11 of ref [1] and eqn. 19-20 of ref [2].
        Be aware that ref [1] uses |e| while ref [2] uses e with sign.

        For perpendicular magnetic pointing to +z, we have A = (-By, 0, 0). Then
            (R_j - R_i) * (A_j + A_i) = -B * (y_j + y_i) * (x_j - x_i)
        However, there isn't the minus sign in the source code of TBPLaS. This is
        because we are considering a magnetic field pointing to -z.

        Note that the above formulae use Gaussian units. For SI units, the speed
        of light vanishes. This is verified by checking the dimension under SI
        units:
            [e]/[h_bar] = IT/(L^2 M T^-1) = L^-2 M^-1 I T^-2
            [R][A] = L * (M T^-2 I^-1 L) = L^2 M I^-1 T^-2
        which cancel out upon multiplication. Similarly, under Gaussian units
        we have:
        [e]/[c*h_bar] = L^(-3/2) M^(-1/2) T
        [R][A] = L^(3/2) M^(1/2) T
        which also cancel out upon multiplication.

        The analysis also inspires to calculate the scaling factor in SI units
        as:
            |e/2h_bar| = 759633724404755.2
        However, we use nm for lengths. So the scaling factor should be divided
        by 1e18, yielding 0.0007596337244047553. That is exactly the magic number
        pi / 4135.666734 = 0.0007596339008078771 in the source code of TBPLaS.

        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
        """
        self.init_array()
        assert gauge in (0, 1, 2), f"Illegal gauge {gauge}"
        core.apply_mag_field(self.hop_ind, self.hop_eng, self.dr,
                             self.orb_pos, intensity, gauge)

    @property
    def num_orb(self) -> int:
        """
        Get the total number of orbitals of the sample.
        :return: total number of orbitals
        """
        self.init_array()
        return self.orb_eng.shape[0]

    @property
    def num_hop(self) -> int:
        """
        Get the total number of hopping terms of the sample.
        :return: total number of hopping terms
        """
        self.init_array()
        return self.hop_ind.shape[0]

    @property
    @abstractmethod
    def sc0_lat_vec(self) -> np.ndarray:
        """
        Get lattice vectors of the first supercell.
        :return: (3, 3) float64 array
            Cartesian coordinates of lattice vectors in NM
        """
        pass

    @property
    @abstractmethod
    def sc0_origin(self) -> np.ndarray:
        """
        Get origin of the first supercell.
        :return: (3, ) float64 array
            Cartesian coordinates of origin in NM
        """
        pass


class Sample(BaseSample):
    """
    Combination of SuperCell and SCInterHopping.

    Attributes
    ----------
    _sc_list: List[SuperCell]
        list of supercells within the sample
    _hop_list: List[SCInterHopping]
        list of inter-hopping sets between supercells within the sample
    """
    def __init__(self, *args: Union[SuperCell, SCInterHopping]) -> None:
        """
        :param args: supercells and inter-hopping sets within this sample
        :returns: None
        """
        super().__init__()

        # Parse supercells and inter-hopping terms
        assert len(args) > 0, f"No components assigned to sample"
        self._sc_list = []
        self._hop_list = []
        for i_arg, arg in enumerate(args):
            if isinstance(arg, SuperCell):
                self._sc_list.append(arg)
            elif isinstance(arg, SCInterHopping):
                self._hop_list.append(arg)
            else:
                raise AssertionError(f"component #{i_arg} should be instance of"
                                     f" SuperCell or SCInterHopping")
            arg.add_subscriber(self)

        # Check closure of inter-hopping instances
        for i_h, hop in enumerate(self._hop_list):
            assert hop.sc_bra in self._sc_list, (f"sc_bra of inter_hop #{i_h}"
                                                 f" not included in args")
            assert hop.sc_ket in self._sc_list, (f"sc_ket of inter_hop #{i_h}"
                                                 f" not included in args")

    def _get_num_orb_sc(self) -> List[int]:
        """
        Get numbers of orbitals in each supercell.
        :return: numbers of orbitals in each supercell
        """
        num_orb_sc = [sc.num_orb_sc for sc in self._sc_list]
        return num_orb_sc

    def _get_ind_start(self) -> List[int]:
        """
        Get starting indices of orbitals for each supercell for assembling
        hopping terms and distances.
        :return: starting indices of orbitals for each supercell
        """
        num_orb_sc = self._get_num_orb_sc()
        ind_start = [sum(num_orb_sc[:_]) for _ in range(len(num_orb_sc))]
        return ind_start

    @property
    def _is_monolithic(self) -> bool:
        """
        Check if the sample is composed with one supercell.
        :return: whether the sample is monolithic
        """
        return len(self._sc_list) == 1 and len(self._hop_list) == 0

    def init_array(self, force_init: bool = False) -> None:
        """
        Interface for initializing the arrays.

        If any array is empty, build all arrays from scratch. Otherwise, build
        them when force_init is True.

        :param force_init: whether to force initializing the arrays from scratch
            even if they have already been initialized
        :return: None
        """
        if force_init or self._size_prod == 0:
            if self._is_monolithic:
                data = self._sc_list[0].get_array()
                self.orb_eng = data.orb_eng
                self.orb_pos = data.orb_pos
                self.hop_ind = data.hop_ind
                self.hop_eng = data.hop_eng
                self.dr = data.dr
            else:
                sc_data = [_.get_array() for _ in self._sc_list]
                orb_eng_tot = [_.orb_eng for _ in sc_data]
                orb_pos_tot = [_.orb_pos for _ in sc_data]
                hop_ind_tot, hop_eng_tot, dr_tot = [], [], []
                ind_start = self._get_ind_start()

                # Collect hopping terms from each supercell
                for i_sc, data in enumerate(sc_data):
                    hop_ind, hop_eng, dr = data.hop_ind, data.hop_eng, data.dr
                    hop_ind += ind_start[i_sc]
                    hop_ind_tot.append(hop_ind)
                    hop_eng_tot.append(hop_eng)
                    dr_tot.append(dr)

                # Collect hopping terms from each inter hopping set
                for hop in self._hop_list:
                    idx_bra = self._sc_list.index(hop.sc_bra)
                    idx_ket = self._sc_list.index(hop.sc_ket)
                    data = hop.get_array(sc_data[idx_bra].orb_pos,
                                         sc_data[idx_ket].orb_pos)
                    hop_ind, hop_eng, dr = data.hop_ind, data.hop_eng, data.dr
                    hop_ind[:, 0] += ind_start[idx_bra]
                    hop_ind[:, 1] += ind_start[idx_ket]
                    hop_ind_tot.append(hop_ind)
                    hop_eng_tot.append(hop_eng)
                    dr_tot.append(dr)
                self.orb_eng = np.concatenate(orb_eng_tot)
                self.orb_pos = np.concatenate(orb_pos_tot)
                self.hop_ind = np.concatenate(hop_ind_tot)
                self.hop_eng = np.concatenate(hop_eng_tot)
                self.dr = np.concatenate(dr_tot)

    def reset_array(self) -> None:
        """
        Reset all modifications to arrays.
        :return: None
        """
        super().reset_array()
        for arg in self._sc_list:
            arg.lock()
        for arg in self._hop_list:
            arg.lock()

    def save(self, directory: str = "sample") -> None:
        """
        Save model to disk.
        :param directory: directory to which data will be saved
        :return: None
        """
        if not os.path.exists(directory):
            os.mkdir(directory)

        # Arrays
        self.init_array()
        np.save(f"{directory}/orb_eng", self.orb_eng)
        np.save(f"{directory}/orb_pos", self.orb_pos)
        np.save(f"{directory}/hop_ind", self.hop_ind)
        np.save(f"{directory}/hop_eng", self.hop_eng)
        np.save(f"{directory}/dr", self.dr)

        # Other attributes
        attr = {
            "sc0_lat_vec": self.sc0_lat_vec,
            "origin": self.sc0_origin,
        }
        with open(f"{directory}/attr.pkl", "wb") as f:
            pickle.dump(attr, f, pickle.HIGHEST_PROTOCOL)

    def plot(self,
             fig_name: str = None,
             fig_size: Tuple[float, float] = None,
             fig_dpi: int = 300,
             sc_orb_colors: List[orb_color_type] = None,
             sc_hop_colors: List[str] = None,
             inter_hop_colors: List[str] = None,
             **kwargs) -> None:
        """
        Plot lattice vectors, orbitals, and hopping terms.

        If figure name is give, 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: size of the figure
        :param fig_dpi: resolution of the figure file
        :param sc_orb_colors: colors for orbitals or each supercell
        :param sc_hop_colors: colors for the hopping terms of each supercell
        :param inter_hop_colors: colors for the hopping terms each inter-hopping
            container
        :param kwargs: arguments for the 'plot' method of 'Super' and
            'SCInterHopping' classes
        :return: None
        """
        fig, axes = plt.subplots(figsize=fig_size)
        axes.set_aspect('equal')

        if sc_orb_colors is None:
            sc_orb_colors = [None for _ in range(len(self._sc_list))]
        if sc_hop_colors is None:
            sc_hop_colors = ['r' for _ in range(len(self._sc_list))]
        if inter_hop_colors is None:
            inter_hop_colors = ['r' for _ in range(len(self._hop_list))]

        # Plot supercells and hopping terms
        for i, sc in enumerate(self._sc_list):
            sc.plot(axes, orb_color=sc_orb_colors[i],
                    hop_color=sc_hop_colors[i], **kwargs)
        for arg in ("with_orbitals", "with_cells"):
            if arg in kwargs.keys():
                kwargs.pop(arg)
        for i, hop in enumerate(self._hop_list):
            hop.plot(axes, hop_color=inter_hop_colors[i], **kwargs)

        # 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()

    @property
    def sc0_lat_vec(self) -> np.ndarray:
        """
        Get lattice vectors of the first supercell.
        :return: (3, 3) float64 array
            Cartesian coordinates of lattice vectors in NM
        """
        return self._sc_list[0].sc_lat_vec

    @property
    def sc0_origin(self) -> np.ndarray:
        """
        Get origin of the first supercell.
        :return: (3, ) float64 array
            Cartesian coordinates of origin in NM
        """
        return self._sc_list[0].prim_cell.origin


class CachedSample(BaseSample):
    """
    Sample with attributes loaded from disk.

    Attributes
    ----------
    _directory: str
        directory within which the data is saved
    _attr: dict
        dictionary containing the key-value type attributes

    Note
    ----
    There are two cases where we need to save/load the sample. Most of the
    time we are just doing post-processing, and need only the essential
    attributes of the sample. We may also load the sample for another full
    TBPM run. For the former case, the attributes are loaded during
    initialization and reloaded in init_*. For the latter case, the attributes
    and arrays are loaded in init_*.
    """
    def __init__(self, directory: str = "sample") -> None:
        """
        :param directory: directory within which the data is saved
        """
        super().__init__()
        self._directory = directory
        with open(f"{self._directory}/attr.pkl", "rb") as f:
            self._attr = pickle.load(f)

    def init_array(self, force_init: bool = False) -> None:
        """
        Interface for initializing the arrays.

        If any array is empty, build all arrays from scratch. Otherwise, build
        them when force_init is True.

        :param force_init: whether to force initializing the arrays from scratch
            even if they have already been initialized
        :return: None
        """
        if force_init or self._size_prod == 0:
            self.orb_eng = np.load(f"{self._directory}/orb_eng.npy")
            self.orb_pos = np.load(f"{self._directory}/orb_pos.npy")
            self.hop_ind = np.load(f"{self._directory}/hop_ind.npy")
            self.hop_eng = np.load(f"{self._directory}/hop_eng.npy")
            self.dr = np.load(f"{self._directory}/dr.npy")

    @property
    def sc0_lat_vec(self) -> np.ndarray:
        """
        Get lattice vectors of the first supercell.
        :return: (3, 3) float64 array
            Cartesian coordinates of lattice vectors in NM
        """
        return self._attr["sc0_lat_vec"]

    @property
    def sc0_origin(self) -> np.ndarray:
        """
        Get origin of the first supercell.
        :return: (3, ) float64 array
            Cartesian coordinates of origin in NM
        """
        return self._attr["origin"]

    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 = False,
             orb_color: orb_color_type = 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 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
        :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
        """
        self.init_array()
        fig, axes = plt.subplots(figsize=fig_size)
        axes.set_aspect('equal')
        viewer = ModelViewer(axes, self.sc0_lat_vec, self.sc0_origin, view)

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

        # Plot hopping terms
        hop_ind, hop_eng, dr = self.hop_ind, self.hop_eng, self.dr
        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, (0, 1, 0, 1, 0, 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()
