"""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_bands", "load_dos", "load_states", "load_dyn_pol",
           "load_epsilon", "load_ac_cond", "load_z2", "load_berry", "load_spin"]


def load_bands(prefix: str,
               data_name: str = "bands") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load band structure.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (k_len, bands)
        k_len: (num_kpt,) float64 array, length of k-path
        bands: (num_kpt, num_bands) float64 array, band structure
    """
    # Columns of band data: k_len, band0, band1, ...
    # At least 2 columns.
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        k_len = data[:, 0]
        bands = data[:, 1:]
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        k_len = h5_load_double(h5_file, f"{data_name}.k_len",
                               transpose=True, flatten=True)
        bands = h5_load_double(h5_file, f"{data_name}.bands",
                               transpose=True)
    return k_len, bands


def load_dos(prefix: str,
             data_name: str = "dos") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load DOS.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :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
    """
    # Columns of dos data: energies, dos
    # 2 columns.
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        energies, dos = data[:, 0], data[:, 1]
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        energies = h5_load_double(h5_file, f"{data_name}.energies",
                                  transpose=True, flatten=True)
        dos = h5_load_double(h5_file, f"{data_name}.dos",
                             transpose=True, flatten=True)
    return energies, dos


def load_states(prefix: str) -> Tuple[np.ndarray, np.ndarray]:
    """
    Load eigenvalues and eigenstates.
    :param prefix: prefix of data files
    :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
    """
    # Eigenvalues data has the same format as band data
    data_name = "eigen_values"
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        bands = data[:, 1:]
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        bands = h5_load_double(h5_file, f"{data_name}.bands",
                               transpose=True)

    # Load states
    # Columns of eigenstates data:
    # (state0.real, state0.imag), (state1.real, state1.imag)
    # At least 2 columns.
    data_name = "eigen_states"
    states = []
    num_kpt = bands.shape[0]
    for ik in range(num_kpt):
        try:
            data = np_load_complex(f"{prefix}_{data_name}_k{ik}.dat",
                                   transpose=True)
        except FileNotFoundError:
            h5_file = f"{prefix}_{data_name}_k{ik}"
            data = h5_load_complex(h5_file, f"{data_name}", transpose=True)
        states.append(data)
    states = np.array(states)
    return bands, states


def load_dyn_pol(prefix: str,
                 data_name: str = "dyn_pol") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load dynamic polarizability from file.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (omegas, dyn_pol)
        omegas: (num_omega,) float64 array
            energies in eV
        dyn_pol: (num_qpt, num_omega) complex128 array
            dynamic polarizability in 1/(eV*nm**2) or 1/(eV*nm**3)
    """
    # Columns of dyn_pol data:
    # energies, (dp_q0.real, dp_q0.imag), (dp_q1.real, dp_q1.imag), ...
    # At least 3 columns.
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        omegas = data[:, 0]
        dyn_pol = (data[:, 1::2] + 1j * data[:, 2::2]).T
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        omegas = h5_load_double(h5_file, f"{data_name}.energies",
                                transpose=True, flatten=True)
        dyn_pol = h5_load_complex(h5_file, f"{data_name}.dyn_pol",
                                  transpose=True)
    return omegas, dyn_pol


def load_epsilon(prefix: str,
                 data_name: str = "epsilon") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load dielectric function from file.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (omegas, epsilon)
        omegas: (num_omega,) float64 array
            energies in eV
        epsilon: (num_qpt, num_omega) complex128 array
            relative dielectric function
    """
    # Columns of epsilon data:
    # energies, (eps_q0.real, eps_q0.imag), (eps_q1.real, eps_q1.imag), ...
    # At least 3 columns.
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        omegas = data[:, 0]
        epsilon = (data[:, 1::2] + 1j * data[:, 2::2]).T
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        omegas = h5_load_double(h5_file, f"{data_name}.energies",
                                transpose=True, flatten=True)
        epsilon = h5_load_complex(h5_file, f"{data_name}.epsilon",
                                  transpose=True)
    return omegas, epsilon


def load_ac_cond(prefix: str,
                 data_name: str = "ac_cond") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load AC Conductivity from file.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (omegas, ac_cond)
        omegas: (num_omega,) float64 array
            energies in eV
        ac_cond: (num_omega,) complex128 array
            AC conductivity in e**2/(hbar*nm) in 3d case and e**2/hbar in 2d case
    """
    # Columns of AC_cond data: energies, (ac_cond.real, ac_cond.imag)
    # 3 columns.
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        omegas = data[:, 0]
        # Similar as dyn_pol and epsilon, but for q0 only.
        ac_cond = (data[:, 1::2] + 1j * data[:, 2::2]).T[0]
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        omegas = h5_load_double(h5_file, f"{data_name}.energies",
                                transpose=True, flatten=True)
        ac_cond = h5_load_complex(h5_file, f"{data_name}.ac_cond",
                                  transpose=True, flatten=True)
    return omegas, ac_cond


def load_z2(prefix: str,
            data_name: str = "z2") -> Tuple[np.ndarray, np.ndarray]:
    """
    Load Z2 data.
    :param prefix: prefix of data file
    :param data_name: name of data file
    :return: (kb_array, phases)
        kb_array: (num_kb,) float64 array
            FRACTIONAL coordinates of the loop along b-axis
        phases: (num_kpt, num_occ) float64 array
            phases of WF centers
    """
    # Columns of Z2 data: k_b, phase0, phase1, ...
    # At least 2 columns.
    try:
        data = np_load_double(f"{prefix}_{data_name}.dat")
        kb_array = data[:, 0]
        phases = data[:, 1:]
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        kb_array = h5_load_double(h5_file, f"{data_name}.k_b",
                                  transpose=True, flatten=True)
        phases = h5_load_double(h5_file, f"{data_name}.phases",
                                transpose=True)
    return kb_array, phases


def load_berry(prefix: str,
               data_name: str = "berry") -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    Load Berry curvature.
    :param prefix: prefix of data files
    :param data_name: name of data file
    :return: (kpt_cart, bands, omega_xy)
        kpt_cart: (num_kpt, 3) float64 array
            Cartesian coordinates of k-points
        bands: (num_kpt, num_states) float64 array
            energies in eV
        omega_xy: (num_kpt, num_states) float64 array
            Berry curvature in NM^2
    """
    try:
        kpt_cart = np_load_double(f"{prefix}_berry_k_cart.dat")
        bands = np_load_double(f"{prefix}_berry_bands.dat")
        omega_xy = np_load_double(f"{prefix}_berry_omega_xy.dat")
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        kpt_cart = h5_load_double(h5_file, f"{data_name}.k_cart",
                                  transpose=True)
        bands = h5_load_double(h5_file, f"{data_name}.bands",
                               transpose=True)
        omega_xy = h5_load_double(h5_file, f"{data_name}.omega_xy",
                                  transpose=True)
    return kpt_cart, bands, omega_xy


def load_spin(prefix: str,
              data_name: str = "spin") -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """
    Load spin data.
    :param prefix: prefix of data files
    :param data_name: name of data file
    :return: (kpt_cart, bands, sigma_x, sigma_y, sigma_z)
        kpt_cart: (num_kpt, 3)
        bands: (num_kpt, num_states)
        sigma_* (num_kpt, num_states)
    """
    # Similar as load_berry.
    try:
        kpt_cart = np_load_double(f"{prefix}_spin_k_cart.dat")
        bands = np_load_double(f"{prefix}_spin_bands.dat")
        sigma_x = np_load_double(f"{prefix}_spin_sigma_x.dat")
        sigma_y = np_load_double(f"{prefix}_spin_sigma_y.dat")
        sigma_z = np_load_double(f"{prefix}_spin_sigma_z.dat")
    except FileNotFoundError:
        h5_file = f"{prefix}_diag"
        kpt_cart = h5_load_double(h5_file, f"{data_name}.k_cart",
                                  transpose=True)
        bands = h5_load_double(h5_file, f"{data_name}.bands",
                               transpose=True)
        sigma_x = h5_load_double(h5_file, f"{data_name}.sigma_x",
                                 transpose=True)
        sigma_y = h5_load_double(h5_file, f"{data_name}.sigma_y",
                                 transpose=True)
        sigma_z = h5_load_double(h5_file, f"{data_name}.sigma_z",
                                 transpose=True)
    return kpt_cart, bands, sigma_x, sigma_y, sigma_z
