"""
Functions and classes for managing the parameters of diagonalization-based
calculations.
"""

import pickle

import numpy as np


__all__ = ["write_item", "DiagConfig"]


def default_k_points() -> np.ndarray:
    return np.array([[0.0, 0.0, 0.0]])


def default_ka_array() -> np.ndarray:
    return np.linspace(-0.5, 0.5, 200)


def default_kb_array() -> np.ndarray:
    return np.linspace(0.0, 0.5, 100)


def write_item(f, key, value):
    """Write key, value pairs to file."""
    # Convert set/list/tuple/dict to numpy array
    if isinstance(value, set):
        value = np.array([_ for _ in value])
    if isinstance(value, list):
        value = np.array(value)
    if isinstance(value, tuple):
        value = np.array(value)
    if isinstance(value, dict):
        value = np.array([(_k, _v) for _k, _v in value.items()])

    # Bool must come before int, otherwise it will be written twice.
    if isinstance(value, bool):
        f.write(f"{key} {int(value)}\n")
    elif (isinstance(value, int) or
          isinstance(value, float) or
          isinstance(value, str)):
        f.write(f"{key} {value}\n")
    elif isinstance(value, np.ndarray):
        if len(value.shape) == 1:  # vector
            f.write(f"begin {key} {value.shape[0]} {1}\n")
            for d in value:
                f.write(f"  {str(d)}")
                f.write("\n")
        else:  # matrix
            f.write(f"begin {key} {value.shape[0]} {value.shape[1]}\n")
            for row in value:
                for d in row:
                    f.write(f"  {str(d)}")
                f.write("\n")
        f.write(f"end {key}\n")
    else:
        # Error message to indicate if we have missed any value types.
        raise AssertionError(f"Unknown data type {key}: {type(value)}")


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

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

        # dimension of mesh grid in 1st Brillouin zone
        # We set the default value to -1 to indicate that it has not been set
        # by the user.
        self.k_grid_size: tuple[int, int, int] = (-1, -1, -1)

        # Convention of Hamiltonian, should be either 1 or 2
        self.convention: int = 1

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

        # Number of occupied states, set the default value to -1 to indicate
        # that it has not been set by the user.
        self.num_occ: int = -1

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

        # Whether the orbitals are sorted in spin-major order
        self.spin_major: bool = True

        # Lower bound of the energy grid in eV
        self.e_min: float = -5.0

        # Upper bound of the energy grid in eV
        self.e_max: float = 5.0

        # Resolution of energy grid in eV
        self.e_step: float = 0.05

        # Broadening parameter in eV
        self.sigma: float = 0.05

        # Basis function to approximate the Delta function
        self.smear_func: str = "Gaussian"

        # Whether to calculate LDOS in DOS calculation
        self.ldos: bool = False

        # Site indices for LDOS calculation.
        self.ldos_orbital_indices: set[int] = {0}

        # FRACTIONAL coordinates of the loop along a-axis
        self.ka_array: np.ndarray = default_ka_array()

        # FRACTIONAL coordinates of the loop along b-axis
        self.kb_array: np.ndarray = default_kb_array()

        # FRACTIONAL coordinate of the loop along c-axis
        self.kc: float = 0.0

        # Berry specific settings

        # Size of extend Brillouin zone with respect to 1st Brillouin zone
        self.bz_size: tuple[int, int, int] = (1, 1, 1)

        # Fractional coordinate of shift for k-points in extended Brillouin zone
        self.bz_shift: np.ndarray = np.array([0.0, 0.0, 0.0])
    
        # whether to use analytical derivation of Hamiltonian
        self.ham_deriv_analytical: bool = False

        # displacement of k-points in 1/NM for numerical derivation
        self.ham_deriv_delta: float = 1.0e-6

        # Response functions specific settings

        # Infinitesimal in eV
        self.delta: float = 0.005

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

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

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

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

        # Component of AC conductivity
        # 1-XX, 2-XY, 3-XZ, 4-YX, 5-YY, 6-YZ, 7-ZX, 8-ZY, 9-ZZ
        self.ac_cond_component: int = 1

        # Algorithm settings

        # Algorithm for diagonalization
        # Candidates:
        # eigen: CPU solvers from eigen
        # lapacke: CPU solvers from LAPACKE
        # feast_dense: CPU solvers from feast for dense matrix
        # feast_sparse: CPU solvers from feast for sparse matrix
        self.diag_algo: str = "eigen"

        # Algorithm for post-processing
        # Candidates: cpu, gpu
        self.pp_algo: str = "cpu"

        # number of states
        # If set to 0, then it will be discarded and a reasonable default value
        #  will be used.
        self.feast_num_states: int = 0

        # lower bound of the feast searching range in eV
        self.feast_e_min: float = -5.0

        # upper bound of the feast searching range in eV
        self.feast_e_max: float = 5.0

        # criterion for convergence
        self.feast_eps: float = 1.0e-5

        # ratio of subspace size to number of states (m0/num_states)
        self.feast_m0_ratio: float = 1.0

        # whether to reuse results of last run as initial guess
        # NOTE:
        # 1. Reusing results should not be enabled from scratch, otherwise there
        #    will be convergence problems. So we cannot set fpm[4]
        #    (fortran fpm(5)) directly.
        #  2. This parameter sets feast_fpm[4] (fortran fpm(5)) to 1 after the
        #  first successful diagonalization, which is typically safe.
        self.feast_reuse_results: bool = False

        # fpm parameters
        # NOTE
        # 1. MKL has its own implementation of FEAST, which is SIGNIFICANTLY
        #    different from the original version. For example, it uses fpm[128]
        #    instead of fpm[64], and fpm[0] cannot be negative. See the
        #    following link for more details:
        #    https://www.intel.com/content/www/us/en/docs/onemkl/developer-reference-c/2023-0/extended-eigensolver-input-parameters.html
        # 2. Anaconda has MKL shipped with it. So when using Python from
        #    Anaconda, it is likely that MKL feast will shield the original
        #    version. Check that carefully.
        # 3. We only disable the verbose output by default (fpm[0] = 0), since
        #    carelessly changing the parameters may corrupt feast.
        # 4. We use C-style fpm with indices starting from 0.
        self.feast_fpm: dict = {0: 0}

    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
            k_points_bak = self.k_points.copy()
            ka_array_bak = self.ka_array.copy()
            kb_array_bak = self.kb_array.copy()
            self.k_points = None
            self.ka_array = None
            self.kb_array = None
            with open(f"{self.prefix}_config.pkl", "wb") as f:
                pickle.dump(self, f, pickle.HIGHEST_PROTOCOL)
            self.k_points = k_points_bak.copy()
            self.ka_array = ka_array_bak.copy()
            self.kb_array = kb_array_bak.copy()
        else:
            # Large arrays send as function arguments are not saved
            keys_exclude = {"k_points", "ka_array", "kb_array"}
            with open(f"{self.prefix}_config.dat", "w") as f:
                for key, value in self.__dict__.items():
                    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 k_points and k*_array 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.k_points = default_k_points()
        self.ka_array = default_ka_array()
        self.kb_array = default_kb_array()
        return data
