"""Functions and classes for TBPM calculation."""

import os
import re
from typing import Tuple, Union

import numpy as np

from tbplas.cython import interface_tbpm as interface
from tbplas.base.parallel import MPIEnv
from tbplas.builder.primitive import PrimitiveCell
from tbplas.builder.sample import Sample
from tbplas.tbpm.config import TBPMConfig, JobInfo
import tbplas.tbpm.io as tbpm_io


__all__ = ["get_qe_indices", "TBPMSolver"]


def get_qe_indices(config: Union[TBPMConfig, JobInfo],
                   energies: np.ndarray) -> np.ndarray:
    """
    Get the indices of quasi-eigenstates in energies, for calculating and
    analyzing DC conductivity.
    :param config: 'Config' instance
        parameters controlling TBPM calculation
    :param energies: (num_eng,) float64 array
        energies corresponding to TBPM steps, for indexing the quasi-eigenstates
    :return: indices of the quasi-eigenstates
    """
    en_limit = config.dc_cond_energy_limits
    assert len(en_limit) % 2 == 0
    e_min, e_max = en_limit[0::2], en_limit[1::2]
    qe_indices = set()
    for i in range(len(e_min)):
        indices = np.where((energies >= e_min[i]) & (energies <= e_max[i]))[0]
        qe_indices = qe_indices.union(indices)
    qe_indices = np.array(sorted(list(qe_indices)))
    return qe_indices


class TBPMSolver:
    """
    Wrapper class over C++ TBPM subroutines.

    Attributes
    ----------
    _model: 'PrimitiveCell' or 'Sample' instance
        model for which TBPM calculations will be performed
    _cpp_model: 'CPPPrimitiveCell' or 'CPPSample' instance
        c++ interface of model
    _cpp_solver: 'CPPTBPMSolverForPrimitiveCell' or 'CPPTBPMSolverForSample' instance
        c++ interface of solver
    _cpp_model_data: namedtuple
        container of data to which _cpp_model is referenced to avoid dangling
        pointers
    _mpi_env: 'MPIEnv' instance
        mpi parallel environment
    config: 'TBPMConfig' instance
        parameters controlling TBPM calculation
    """
    def __init__(self,
                 model: Union[PrimitiveCell, Sample],
                 echo_details: bool = True) -> None:
        """
        :param model: model for which TBPM calculations will be performed
        :param echo_details: whether to output parallelization details
        """
        self._model = model
        self._cpp_model = None
        self._cpp_solver = None
        self._cpp_model_data = None
        self._mpi_env = MPIEnv(echo_details=echo_details)
        self.config = TBPMConfig()

    def _set_output(self) -> None:
        """
        Setup directory for output.
        :return: None
        """
        if re.match(r"^\w+$", self.config.prefix) is not None:
            directory, prefix = "./", self.config.prefix
        elif re.match(r"^\w+/\w+$", self.config.prefix) is not None:
            directory, prefix = self.config.prefix.split("/")
        else:
            raise AssertionError(f"Illegal prefix {self.config.prefix}")
        # Create directory if not exits
        # NOTE: On some file systems, e.g. btrfs, creating directory may be
        # slow. Other processes may try to access the directory before it is
        # created and run into I/O errors. So we need to put a barrier here.
        if self._mpi_env.is_master:
            if not os.path.exists(directory):
                os.mkdir(directory)
        self._mpi_env.barrier()

        # Print output details
        self._mpi_env.print("Output details:")
        self._mpi_env.print(f"{'':2s}{'Directory':11s}: {directory}")
        self._mpi_env.print(f"{'':2s}{'Prefix':11s}: {prefix}")
        self._mpi_env.print()

    def _pre_run(self) -> None:
        """
        Prepare essentials for TBPM calculation.
        :return: None
        """
        self._set_output()

        # Initialize solver
        # We must create a new solver for each run, since the model may change
        # during two adjacent runs, so are the addresses of array attributes.
        if isinstance(self._model, PrimitiveCell):
            assert self._model.num_orb > 0, f"Num_orb should > 0"
            assert self._model.num_hop > 0, f"Num_hop should > 0"
            self._cpp_model_data = self._model.get_array()
            self._cpp_model = interface.CPPPrimitiveCell(
                self._cpp_model_data.orb_eng,
                self._cpp_model_data.orb_pos,
                self._cpp_model_data.hop_ind,
                self._cpp_model_data.hop_eng,
                self._cpp_model_data.lat_vec,
                self._cpp_model_data.origin)
            self._cpp_solver = interface.CPPTBPMSolverForPrimitiveCell(
                self._cpp_model,
                f"{self.config.prefix}_config.dat".encode())
        else:
            self._cpp_model_data = self._model.get_array()
            self._cpp_model = interface.CPPSample(
                self._cpp_model_data.orb_eng,
                self._cpp_model_data.orb_pos,
                self._cpp_model_data.hop_ind,
                self._cpp_model_data.hop_eng,
                self._cpp_model_data.dr,
                self._cpp_model_data.sc0_lat_vec)
            self._cpp_solver = interface.CPPTBPMSolverForSample(
                self._cpp_model,
                f"{self.config.prefix}_config.dat".encode())

        # Master process saves config while other processes waiting
        if self._mpi_env.is_master:
            self.config.save(fmt="pkl")
            self.config.save(fmt="dat")
        self._mpi_env.barrier()

    def calc_ldos_haydock(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate local density of states (LDOS) using Haydock recursion method.

        CAUTION: this method works for only one site. Although it can be adopted
        to deal with multiple sites, the time usage will be unaffordable.
        Use TBPM instead if you want to calculate LDOS for multiple sites.

        Ref: https://journals.jps.jp/doi/10.1143/JPSJ.80.054710

        :return: (energies, ldos)
            energies: (2*num_steps+1, ) float64 array
            energies in eV
            ldos: (2*num_steps+1, ) float64 array
            LDOS value to corresponding energies in 1/eV
        """
        self._pre_run()
        self._cpp_solver.calc_ldos_haydock()
        return tbpm_io.load_ldos_haydock(self.config.prefix)

    def calc_corr_dos(self) -> np.ndarray:
        """
        Calculate correlation function of density of states (DOS).
        :return: (num_steps+1, ) complex128 array
            dimensionless DOS correlation function
        """
        self._pre_run()
        self._cpp_solver.calc_corr_dos()
        return tbpm_io.load_corr_dos(self.config.prefix)

    def calc_corr_ac_cond(self) -> np.ndarray:
        """
        Calculate correlation function of optical (AC) conductivity.
        :return: (4, num_steps) complex128 array
            AC correlation function along xx, yx, xy, yy directions
            Unit is e^2/hbar^2 * (eV)^2 * nm^2.
        """
        self._pre_run()
        self._cpp_solver.calc_corr_ac_cond()
        return tbpm_io.load_corr_ac_cond(self.config.prefix)

    def calc_corr_dyn_pol(self) -> np.ndarray:
        """
        Calculate correlation function of dynamical polarization.
        :return: (num_qpt, num_steps) float64 array
            dimensionless Dynamical polarization correlation function.
        """
        self._pre_run()
        self._cpp_solver.calc_corr_dyn_pol()
        return tbpm_io.load_corr_dyn_pol(self.config.prefix)

    def calc_quasi_eigenstates(self) -> np.ndarray:
        """
        Calculate squared quasi-eigenstates.
        :return: (num_qe, num_orb) float64 array
            Each row is a squared quasi-eigenstate at energy
            of config.qe_energies.
        """
        self._pre_run()
        self._cpp_solver.calc_quasi_eigenstates()
        return tbpm_io.load_quasi_eigen(self.config.prefix)

    def calc_corr_dc_cond(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate correlation function of electronic (DC) conductivity.
        :return: (corr_dos, corr_dc)
            corr_dos: (num_steps,) complex128 array
            dimensionless DOS correlation function
            corr_dc: (2, num_eng, num_steps_dc) complex128 array
            DC conductivity correlation function in e^2/hbar^2 * (eV)^2 * nm^2 in
            accordance with config.dckb_energies.
        """
        self._pre_run()
        self._cpp_solver.calc_corr_dc_cond()
        return tbpm_io.load_corr_dc_cond(self.config.prefix)

    def calc_hall_cond(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate Hall conductivity according to Kubo-Bastin formula.

        Reference:
        https://journals.aps.org/prl/pdf/10.1103/PhysRevLett.114.116602

        The unit of hall_cond in 2d case from eqn. 1 of the reference follows:
        [hall_cond] = [hbar*e^2/Omega] * [dE * Trace<...>]
                    = hbar*e^2/nm^2 * 1/eV * nm^2/hbar^2 * eV
                    = e^2/hbar
        which is consistent with AC and DC conductivity.
        Note that the delta function in the formula has the unit of 1/eV.

        The unit can also be determined from enq. 4 as:
        [hall_cond] = [e^2*hbar/(Omega*delaE^2)] * [mu]
                    = hbar/(nm^2*V^2) * nm^2 * (eV)^2 / hbar^2
                    = e^2/hbar
        Note that the scaled energy is dimensionless.

        :return: (energies, conductivity)
            energies: (num_eng, ) float64 array
            chemical potentials specified in config.dckb_energies
            conductivity: (num_eng, ) float64 array
            Hall conductivity according to energies
        """
        self._pre_run()
        self._cpp_solver.calc_hall_cond()
        return tbpm_io.load_hall_cond(self.config.prefix)

    def calc_wft(self) -> np.ndarray:
        """
        Calculate propagation of wave function from given initial state.
        :return: (num_time, num_orb) complex128 array
            Each row is time-dependent wave function at config.wft_time_save.
        """
        assert self._mpi_env.size == 1, ("Using more than 1 mpi process is not"
                                         " allowed for wft calculation")
        assert len(self.config.wft_wf0.shape) == 1, "wft_wf0 should be 1d"
        assert self.config.wft_wf0.shape[0] == self._model.num_orb, \
            "wft_wf0 should have the same size of num_orb"
        self._pre_run()
        self._cpp_solver.calc_wft(self.config.wft_wf0)
        return tbpm_io.load_wft(self.config.prefix)

    def calc_corr_bands(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate correlation function of band structure.
        :return: (k_len, corr_bands)
            k_len: (num_kpt,) float64 array
            length of k-path in 1/nm
            corr_bands: (num_kpt, num_steps+1) complex128 array
            correlation functions for each k-point
        """
        self._pre_run()
        self._cpp_solver.calc_corr_bands()
        return tbpm_io.load_corr_bands(self.config.prefix)

    @property
    def is_master(self) -> bool:
        """Wrapper for determine the master process for compatibility."""
        return self._mpi_env.is_master
