"""All-in-one module for parsing data produced by cpp."""

from typing import Tuple

import numpy as np

from tbplas.base.io import (np_load_double, np_load_complex,
                            h5_load_double, h5_load_complex)


__all__ = ["load_ldos_haydock", "load_corr_dos", "load_corr_ac_cond",
           "load_corr_dyn_pol", "load_quasi_eigen", "load_corr_dc_cond",
           "load_hall_cond", "load_wft", "load_corr_bands"]


def load_ldos_haydock(prefix: str, data_name: str = "ldos") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load LDOS data produced by Haydock method.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :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
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_double(file_name)
        energies, ldos = data[:, 0], data[:, 1]
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        energies = h5_load_double(h5_file, f"{data_name}.energies",
                                  flatten=True)
        ldos = h5_load_double(h5_file, f"{data_name}.ldos",
                              flatten=True)
    return energies, ldos


def load_corr_dos(prefix: str, data_name: str = "corr_dos") -> np.ndarray:
    """
    Load DOS correlation function.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (num_steps+1, ) complex128 array
        dimensionless DOS correlation function
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_complex(file_name, flatten=True)
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        data = h5_load_complex(h5_file, data_name, flatten=True)
    return data


def load_corr_ac_cond(prefix: str, data_name: str = "corr_ac_cond") -> np.ndarray:
    """
    Load AC Conductivity correlation function.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (4, num_steps) complex128 array
        AC correlation function along xx, yx, xy, yy directions
        Unit is e^2/h_bar^2 * (eV)^2 * nm^2.
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_complex(file_name)
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        data = h5_load_complex(h5_file, data_name)
    return data


def load_corr_dyn_pol(prefix: str, data_name: str = "corr_dyn_pol") -> np.ndarray:
    """
    Load dynamic polarizability correlation function.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (num_qpt, num_steps) float64 array
        dimensionless Dynamical polarization correlation function.
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_double(file_name, transpose=True)
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        data = h5_load_double(h5_file, data_name, transpose=True)
    return data


def load_quasi_eigen(prefix: str, data_name: str = "quasi_eigen") -> np.ndarray:
    """
    Load squared quasi-eigenstates.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (num_qe, num_orb) float64 array
        Each row is a squared quasi-eigenstate at energy
        of config.qe_energies.
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_double(file_name, transpose=True)
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        data = h5_load_double(h5_file, data_name, transpose=True)
    return data


def load_corr_dc_cond(prefix: str, data_name: str = "corr_dc_cond") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load DC conductivity correlation function.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (corr_dos, corr_dc_cond)
        corr_dos: (num_steps,) complex128 array
        dimensionless DOS correlation function
        corr_dc_cond: (2, num_eng, num_steps_dc) complex128 array
        DC conductivity correlation function in e^2/h_bar^2 * (eV)^2 * nm^2 in
        accordance with config.dckb_energies.
    """
    try:
        # DO NOT use load_corr_dos since it will cause GROUP_NOT_FOUND error.
        corr_dos = np_load_complex(f"{prefix}_{data_name}_dos.dat",
                                   flatten=True)
        corr_dc_cond_xx = np_load_complex(f"{prefix}_{data_name}_xx.dat")
        corr_dc_cond_yy = np_load_complex(f"{prefix}_{data_name}_yy.dat")
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        corr_dos = h5_load_complex(h5_file, f"{data_name}.dos")
        corr_dc_cond_xx = h5_load_complex(h5_file, f"{data_name}.xx")
        corr_dc_cond_yy = h5_load_complex(h5_file, f"{data_name}.yy")
    # assemble corr_dc
    num_eng = corr_dc_cond_xx.shape[0]
    num_step = corr_dc_cond_xx.shape[1]
    corr_dc_cond = np.zeros((2, num_eng, num_step), dtype=np.complex128)
    corr_dc_cond[0] = corr_dc_cond_xx
    corr_dc_cond[1] = corr_dc_cond_yy
    return corr_dos, corr_dc_cond


def load_hall_cond(prefix: str, data_name: str = "hall_cond") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load LDOS data produced by Haydock method.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :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
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_double(file_name)
        energies, cond = data[:, 0], data[:, 1]
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        energies = h5_load_double(h5_file, f"{data_name}.energies",
                                  flatten=True)
        cond = h5_load_double(h5_file, f"{data_name}.hall_cond",
                              flatten=True)
    return energies, cond


def load_wft(prefix: str, data_name: str = "wft") -> np.ndarray:
    """
    Load wave function.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (num_time, num_orb) complex128 array
        Each row is time-dependent wave function at config.wft_time_save.
    """
    try:
        file_name = f"{prefix}_{data_name}.dat"
        data = np_load_complex(file_name, transpose=True)
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        data = h5_load_complex(h5_file, data_name, transpose=True)
    return data


def load_corr_bands(prefix: str, data_name: str = "corr_bands") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load band structure correlation function.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (k_len, corr_bands)
        k_len: (num_kpt, ) float64 array
            length of kpath in 1/nm
        corr_bands: (num_kpt, num_steps+1) complex128 array
            correlation for each k-point
    """
    try:
        k_len = np_load_double(f"{prefix}_k_len.dat", flatten=True)
        corr_bands = np_load_complex(f"{prefix}_{data_name}.dat",
                                     transpose=True)
    except FileNotFoundError:
        h5_file = f"{prefix}_tbpm"
        k_len = h5_load_double(h5_file, f"{data_name}.k_len",
                               flatten=True)
        corr_bands = h5_load_complex(h5_file, f"{data_name}.corr_bands",
                                     transpose=True)
    return k_len, corr_bands
