"""Fundamentals of Solvers based on exact diagonalization."""

import re
import os
from typing import Tuple

import numpy as np
import scipy.linalg.lapack as lapack

from tbplas.cython import interface_diag as interface
from tbplas.base.parallel import MPIEnv
from tbplas.base.kpoints import gen_kdist
from tbplas.builder.primitive import PrimitiveCell
from tbplas.diag.config import DiagConfig
from tbplas.diag.io import load_bands, load_dos, load_states


__all__ = ["DiagSolver", "calc_bands", "calc_dos"]


class DiagSolver:
    """
    Base class for solvers based on exact diagonalization.

    Attributes
    ----------
    _model: 'PrimitiveCell' instance
        model for which properties will be calculated
    _overlap: 'PrimitiveCell' instance
        model holding the overlap of basis functions
    _cpp_model: 'CPPPrimitiveCell' instance
        c++ interface of model
    _cpp_overlap: 'CPPPrimitiveCell' instance
        c++ interface of overlap
    _cpp_solver: 'CppDiagSolver' or other class instance
        C++ interface of solvers
    _cpp_model_data: namedtuple
        container of data to which _cpp_model is referenced to avoid dangling
        pointers
    _cpp_overlap_data: namedtuple
        container of data to which _cpp_overlap is referenced to avoid dangling
        pointers
    _mpi_env: 'MPIEnv' instance
        mpi parallel environment
    config: 'DiagConfig' instance
        parameters controlling the calculations
    """
    def __init__(self,
                 model: PrimitiveCell,
                 overlap: PrimitiveCell = None,
                 echo_details: bool = True) -> None:
        """
        :param model: model for which properties will be calculated
        :param overlap: model holding the overlap of basis functions
        :param echo_details: whether to output parallelization details
        """
        self._model = model
        self._overlap = overlap
        self._cpp_model = None
        self._cpp_overlap = None
        self._cpp_solver = None
        self._cpp_model_data = None
        self._cpp_overlap_data = None
        self._mpi_env = MPIEnv(echo_details=echo_details)
        self.config = DiagConfig()

    def _check_k_points(self) -> None:
        """Check correctness of self.config.k_points."""
        assert len(self.config.k_points.shape) == 2
        assert self.config.k_points.shape[0] >= 1
        assert self.config.k_points.shape[1] == 3

    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 _set_cpp_model(self) -> None:
        """
        Update self._cpp_model.

        We must create a new overlap for each run, since the model may change
        during two adjacent runs, so are the addresses of array attributes.
        """
        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)

    def _set_cpp_overlap(self) -> None:
        """
        Update self._cpp_overlap.

        We must create a new overlap for each run, since the model may change
        during two adjacent runs, so are the addresses of array attributes.
        """
        if self._overlap is not None:
            msg = f"Num_orb of model and overlap should be equal"
            assert self._overlap.num_orb == self._model.num_orb, msg
            assert self._overlap.num_hop > 0, f"Num_hop should > 0"
            self._cpp_overlap_data = self._overlap.get_array()
            self._cpp_overlap = interface.CPPPrimitiveCell(
                self._cpp_overlap_data.orb_eng,
                self._cpp_overlap_data.orb_pos,
                self._cpp_overlap_data.hop_ind,
                self._cpp_overlap_data.hop_eng,
                self._cpp_overlap_data.lat_vec,
                self._cpp_overlap_data.origin)
        else:
            # Reset self._cpp_overlap for each new calculation.
            self._cpp_overlap = None

    def _set_cpp_solver(self, job_type: str) -> None:
        """
        Update self._cpp_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.

        :param job_type: type of calculation
        :return: None
        """
        solver_map = {
            "diag": interface.CPPDiagSolver,
            "lindhard": interface.CPPLindhard,
            "spin": interface.CPPSpinTexture,
            "z2": interface.CPPZ2,
            "berry": interface.CPPBerry,
        }
        solver_type = solver_map[job_type]
        self._cpp_solver = solver_type(
            self._cpp_model,
            self._cpp_overlap,
            f"{self.config.prefix}_config.dat".encode())

    def _pre_run(self, job_type: str) -> None:
        """
        Prepare essentials for calculation.
        :param job_type: type of calculation
        :return: None
        """
        self._set_output()
        self._set_cpp_model()
        self._set_cpp_overlap()
        self._set_cpp_solver(job_type)
        # 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_bands(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate band structure along given k_path.
        :return: (k_len, bands)
            k_len: (num_kpt,) float64 array, length of k-path
            bands: (num_kpt, num_bands) float64 array, band structure
        """
        if __debug__:
            self._check_k_points()
        self._pre_run("diag")
        self._cpp_solver.calc_bands(self.config.k_points)
        return load_bands(self.config.prefix)

    def calc_bands_scipy(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate band structure along given k_path.

        This function calls C++ interface to set up the Hamiltonian. However,
        the diagonalization is done using scipy. For now this function is
        intended to be called by the 'ParamFit' class. It also works as an
        example for obtaining the Hamiltonian for post-process.

        :return: (k_len, bands)
            k_len: (num_kpt,) float64 array, length of k-path
            bands: (num_kpt, num_bands) float64 array, band structure
        """
        if __debug__:
            self._check_k_points()
        self._set_cpp_model()
        self._set_cpp_overlap()

        # Initialize Hamiltonian, which must be column-major (FORTRAN order).
        num_orb = self._model.num_orb
        ham = np.zeros((num_orb, num_orb), dtype=np.complex128, order='F')
        with_overlap = (self._overlap is not None)
        overlap = np.zeros_like(ham) if with_overlap else None

        # Calculate eigenvalues
        num_kpt = self.config.k_points.shape[0]
        bands = np.zeros((num_kpt, num_orb), dtype=np.float64)
        for i_k, kpt in enumerate(self.config.k_points):
            self._cpp_model.build_ham_dense(kpt, ham, 1)
            if with_overlap:
                self._cpp_overlap.build_ham_dense(kpt, overlap, 1)
                w, v, info = lapack.zhegv(ham, overlap, itype=1, jobz="N")
            else:
                w, v, info = lapack.zheev(ham, compute_v=0)
            if info != 0:
                raise RuntimeError("Diagonalization failed")
            bands[i_k] = w
        k_len = gen_kdist(self._model.lat_vec, self.config.k_points)
        return k_len, bands

    def calc_dos(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate density of states for given energy range and step.
        :return: (energies, dos)
            energies: (num_eng,) float64 array
            energies determined by e_min, e_max and e_step
            dos: (num_eng,) float64 array
            density of states in states/eV
        """
        if __debug__:
            self._check_k_points()
        self._pre_run("diag")
        self._cpp_solver.calc_dos(self.config.k_points)
        return load_dos(self.config.prefix)

    def calc_states(self) -> Tuple[np.ndarray, np.ndarray]:
        """
        Calculate bands as well as eigenstates.

        :return: (bands, states)
            bands: (num_kpt, num_bands) float64 array
                band structure
            states: (num_kpt, num_bands, num_orb) complex128 array
                eigenstates at each k-point, with each ROW being a state
        """
        if __debug__:
            self._check_k_points()
        self._pre_run("diag")
        self._cpp_solver.calc_states(self.config.k_points)
        return load_states(self.config.prefix)

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


def calc_bands(model: PrimitiveCell,
               k_path: np.ndarray,
               **kwargs) -> Tuple[np.ndarray, np.ndarray]:
    """
    Function for quick calculation of band structure.
    :param model: model for which band structure will be calculated
    :param k_path: k-path in the Brillouin zone
    :param kwargs: settings in DiagConfig
    :return: (k_len, bands) length of k-path and band structure
    """
    solver = DiagSolver(model)
    solver.config.k_points = k_path
    for key, value in kwargs.items():
        setattr(solver.config, key, value)
    return solver.calc_bands()


def calc_dos(model: PrimitiveCell,
             k_points: np.ndarray,
             **kwargs) -> Tuple[np.ndarray, np.ndarray]:
    """
    Function for quick calculation of dos.
    :param model: model for which band structure will be calculated
    :param k_points: k-points in the Brillouin zone
    :param kwargs: settings in DiagConfig
    :return: (energy, dos) energy and dos
    """
    solver = DiagSolver(model)
    solver.config.k_points = k_points
    for key, value in kwargs.items():
        setattr(solver.config, key, value)
    return solver.calc_dos()
