"""Functions and classes for managing the parameters of TBPM calculation."""

import pickle
import re
from typing import List

import numpy as np

from tbplas.diag.config import write_item


__all__ = ["TBPMConfig", "ConfigParser", "JobInfo"]


def default_wft_wf0() -> np.ndarray:
    return np.array([0], dtype=np.complex128)


class TBPMConfig:
    """Class for representing TBPM parameters."""
    def __init__(self):
        # Prefix prepended to output files
        self.prefix: str = "sample"

        # Whether to save results to data files, currently affects quasi-eigenstates
        # and time-dependent wave functions
        # NOTE: This flag is for benchmarking purposes, where we need to disable
        # output to avoid expensive I/O.
        self.save_data: bool = True

        # Dimension of the model, should be 2 or 3
        self.dimension: int = 3

        # Temperature in Kelvin.
        self.temperature: float = 300.0

        # Chemical potential in eV
        self.mu: float = 0.0

        # Spin degeneracy, should be 1 or 2
        self.num_spin: int = 1

        # Algorithm for TBPM/Haydock/KuboBastin calculations
        # Candidates: cpu/cpu_fast/gpu
        self.algo: str = "cpu"

        # Pool size for fast TBPM algorithms
        self.pool_size: int = 4

        # Rescaling factor of Hamiltonian in eV
        # If the factor is negative, then it will be discarded and a reasonable
        # default value will be used.
        self.rescale: float = -1.0

        # Bessel function precision cut-off.
        self.bessel_precision: float = 1.0e-14

        # Number of random initial wave functions.
        self.num_random_samples: int = 1

        # Number of time steps.
        self.num_time_steps: int = 1024

        # Seed for random wave function generation.
        self.seed: int = 1337

        # Check the wave function for divergence after this number of steps.
        self.wfn_check_time_steps: int = 128

        # Threshold for checking divergence of wave function.
        self.wfn_check_threshold: float = 1.0e-9

        # Precision cut-off of Fermi-Dirac distribution.
        self.fermi_cheb_precision: float = 1.0e-12

        # Number of Fermi-Dirac distribution FFT steps, must be power of two.
        self.num_fermi_fft_steps: int = 32768

        # Whether to calculate LDOS in TBPM DOS calculation
        # For the Haydock method, we always calculate LDOS.
        self.ldos: bool = False

        # Site indices for LDOS calculation.
        # There is no limit for the number of sites for TBPM.
        # However, the Haydock recursion method accepts only one site.
        self.ldos_orbital_indices: set[int] = {0}

        # Infinitesimal in eV
        self.ldos_delta: float = 0.01

        # Recursion depth of Haydock method.
        self.ldos_recursion_depth: int = 2000

        # FRACTIONAL coordinates of k-points in ROW-WISE order
        # dimension: (num_kpt, 3)
        self.k_points: np.ndarray = np.array([[0.0, 0.0, 0.0]])

        # Relative background dielectric constant
        self.dyn_pol_back_epsilon: float = 1.0

        # Scaling factor for Coulomb potential
        self.dyn_pol_coulomb_factor: float = 1.0

        # CARTESIAN coordinates of q-points in 1/NM in ROW-WISE order
        # dimension: (num_qpt, 3)
        self.dyn_pol_q_points: np.ndarray = np.array([[1.0, 0.0, 0.0]])

        # Lists of minimum and maximum of energy window for DC conductivity in eV
        self.dc_cond_energy_limits: list[tuple] = [(-0.5, 0.5)]

        # Number of time steps for calculating dc correlation.
        self.dc_cond_num_time_steps: int = 1024

        # List of chemical potentials to calculate Hall conductivity in eV.
        self.dckb_energies: np.ndarray = np.linspace(-0.2, 0.2, 40)

        # Number of kernels in Kernel Polynomial Method(KPM).
        # If the model is not large enough, a large dckb_num_kernel will lead to
        # erroneous results, which have large absolute values.
        self.dckb_num_kernel: int = 2048

        # Number of integral steps.
        self.dckb_num_integ_steps: int = 2048

        # Component of Hall conductivity, 1 gives XX, 2 gives XY.
        self.dckb_component: int = 1

        # Unit of Hall conductivity, should be hbar or h
        self.dckb_unit: str = "hbar"

        # List of energies of quasi-eigenstates in eV.
        self.qe_energies: np.ndarray = np.array([-0.1, 0.0, 0.1])

        # steps on which time the time-dependent wave function will be saved
        self.wft_time_save: np.ndarray = np.array([0])

        # scale factor for the time step with respect to the default value
        self.wft_dt_scale: float = 1.0

        # Initial state of time-dependent wave function
        # Must have data even we do not need it, otherwise c++ backend will fail.
        self.wft_wf0: np.ndarray = default_wft_wf0()

    def save(self, fmt: str = "pkl") -> None:
        """
        Save configuration to file.
        :param fmt: format of output
        :return: None
        """
        assert fmt in ("pkl", "dat"), f"Unknown format {fmt}"
        if fmt == "pkl":
            # Large arrays are backed up and restored, but not saved
            wf0_bak = self.wft_wf0.copy()
            self.wft_wf0 = None
            with open(f"{self.prefix}_config.pkl", "wb") as f:
                pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
            self.wft_wf0 = wf0_bak.copy()
        else:
            # Large arrays send as function arguments are not saved
            keys_exclude = {"wft_wf0"}
            with open(f"{self.prefix}_config.dat", "w") as f:
                for key, value in self.__dict__.items():
                    if key == "dc_cond_energy_limits":
                        for v in value:
                            assert len(v) == 2, f"Illegal energy limit {v}"
                    if key not in keys_exclude:
                        write_item(f, key, value)
                        f.write("\n")

    def load(self, prefix: str, fmt: str = "pkl") -> None:
        """
        Read configuration from file.

        NOTE: Large arrays like wft_wf0 are not saved and restored.

        :param prefix: prefix of config file
        :param fmt: format of config file
        """
        assert fmt in ("pkl",), f"Unknown format {fmt}"
        with open(f"{prefix}_config.pkl", 'rb') as f:
            data = pickle.load(f)
        for key, value in data.__dict__.items():
            setattr(self, key, value)
        self.wft_wf0 = default_wft_wf0()


class ConfigParser:
    """
    Parser for *_config.dat and *_info.dat
    """
    def __init__(self, file_name: str) -> None:
        """
        :param file_name: file name of *_config.dat or *_info.dat
        """
        with open(file_name, "r") as inf:
            self._content = inf.readlines()

    def parse_line(self, key: str) -> List[str]:
        """
        Extract and split line matching "key value".
        :param key: keyword of the line
        :return: matched line split with spaces
        """
        pattern = re.compile(rf"^\s*{key}\s+\S+", re.IGNORECASE)
        for line in self._content:
            if re.search(pattern, line) is not None:
                return line.split()
        raise RuntimeError(f"Key {key} not found")

    def parse_block(self, block_name: str) -> np.ndarray:
        """
        Extract array defined by the block matching:
            begin block_name M N
                x, x, x
                x, x, x
                ..., ...
            end block_name

        :param block_name: name of the block
        :return: extracted array
        """
        pattern = re.compile(rf"^\s*begin\s+{block_name}\s+\d+", re.IGNORECASE)
        for nl, line in enumerate(self._content):
            if re.search(pattern, line) is not None:
                s = line.split()
                nl_start = nl + 1
                num_row = int(s[2])
                block_content = self._content[nl_start:nl_start + num_row]
                array = [[float(_) for _ in line.split()] for line in block_content]
                return np.array(array)
        raise RuntimeError(f"Block {block_name} not found")


class JobInfo:
    """Container of parameters produced by TBPM calculations."""
    def __init__(self, file_name: str) -> None:
        """
        :param file_name: file name of *_info.dat
        """
        self._file_name = file_name

        # Model info
        self.energy_range = None
        self.num_orbitals = None
        self.lattice_area = None
        self.lattice_volume = None

        # TBPM configuration
        self.dimension = None
        self.num_time_steps = None
        self.num_spin = None
        self.temperature = None
        self.dyn_pol_coulomb_factor = None
        self.dyn_pol_back_epsilon = None
        self.dyn_pol_q_points = None
        self.dc_cond_energy_limits = None
        self.dc_cond_num_time_steps = None

    def load(self) -> None:
        """Load *_info.dat."""
        parser = ConfigParser(self._file_name)

        # Model info
        self.energy_range = float(parser.parse_line("energy_range")[1])
        self.num_orbitals = float(parser.parse_line("num_orbitals")[1])
        self.lattice_area = float(parser.parse_line("lattice_area")[1])
        self.lattice_volume = float(parser.parse_line("lattice_volume")[1])

        # TBPM configuration
        self.dimension = int(parser.parse_line("dimension")[1])
        self.num_time_steps = int(parser.parse_line("num_time_steps")[1])
        self.num_spin = int(parser.parse_line("num_spin")[1])
        self.temperature = float(parser.parse_line("temperature")[1])
        self.dyn_pol_coulomb_factor = float(parser.parse_line("dyn_pol_coulomb_factor")[1])
        self.dyn_pol_back_epsilon = float(parser.parse_line("dyn_pol_back_epsilon")[1])
        self.dyn_pol_q_points = parser.parse_block("dyn_pol_q_points")
        self.dc_cond_energy_limits = parser.parse_block("dc_cond_energy_limits").flatten()
        self.dc_cond_num_time_steps = int(parser.parse_line("dc_cond_num_time_steps")[1])
