/**
 * @brief Classes for results of diagonalization calculations.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2024, TBPLaS develop team.
 */

#ifndef TBPLAS_DIAG_RESULTS_H
#define TBPLAS_DIAG_RESULTS_H

#include <eigen3/Eigen/Dense>

namespace tbplas::diag {

/**
 * @brief Base class for data files.
 *
 */
class BaseData {
public:
    /**
     * @brief Construct a new BaseData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit BaseData(const std::string& prefix = "sample")
        : prefix_(prefix)
    {
    }

protected:
    /// @brief Prefix of data files
    std::string prefix_ = "sample";
};

/**
 * @brief Class for holding band structure
 *
 */
class BandData : public BaseData {
public:
    /// @brief length of k-path in 1/NM, dimension: (num_kpt, 1)
    Eigen::VectorXd k_len;

    /// @brief energies along k-path in COLUMN-WISE order in eV
    /// dimension: (num_states, num_kpt)
    Eigen::MatrixXd bands;

    /**
     * @brief Construct a new BandData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit BandData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.k_len and ${data_name}.bands
     *
     * 2. Data format of dat output:
     *      dimension: (num_kpt, num_states+1)
     *      columns: k_len, band0, band1, ...
     */
    void save(const std::string& data_name = "bands") const;
};

/**
 * @brief Class for holding DOS
 *
 */
class DOSData : public BaseData {
public:
    /// @brief energy grid corresponding to e_min, e_max and e_step in eV
    /// dimension: (num_eng, 1)
    Eigen::VectorXd energies;

    /// @brief normalized density of states in states/eV
    /// dimension: (num_eng, 1)
    Eigen::VectorXd dos;

    /**
     * @brief Construct a new DOSData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit DOSData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.energies and ${data_name}.dos
     *
     * 2. Data format of dat output:
     *      dimension: (num_eng, 2)
     *      columns: energies, dos
     */
    void save(const std::string& data_name = "dos") const;
};

/**
 * @brief Class for holding eigenvalues and eigenstates
 *
 */
class EigenStateData : public BandData {
public:
    /**
     * @brief Construct a new EigenStateData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit EigenStateData(const std::string& prefix = "sample")
        : BandData(prefix)
    {
    }

    /**
     * @brief Save eigenvalues of all k-points to file.
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.k_len and ${data_name}.bands
     *
     * 2. This function is a wrapper over BandData.save(). See that function
     *    for the format of output.
     */
    void save_eigen_values(const std::string& data_name = "eigen_values") const;

    /**
     * @brief Save eigenstates of given k-point to file.
     *
     * @param[in] ik index of k-point
     * @param[in] eigen_states eigenstates to save
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}_kN.dat (no HDF5)
     *      ${prefix}_${data_name}_kN.h5 (HDF5), with dataset eigen_states
     *    We must use separate files since all MPI processes have data to write.
     *
     * 2. Data format of dat output:
     *      dimension: (num_orb, num_states*2)
     *      columns: (state0.real, state0.imag), (state1.real, state1.imag), ...
     */
    void save_eigen_states(
        const int& ik,
        const Eigen::MatrixXcd& eigen_states,
        const std::string& data_name = "eigen_states") const;
};

/**
 * @brief Class for holding the results of spin-texture calculation
 *
 */
class SpinData : public BaseData {
public:
    /// @brief Cartesian coordinates of k-points in 1/NM, dimension: (3, num_kpt)
    Eigen::Matrix3Xd k_cart;

    /// @brief Energies on k-points in eV, dimension: (num_states, num_kpt)
    Eigen::MatrixXd bands;

    /// @brief Expectation of sigma_x, dimension: (num_states, num_kpt)
    Eigen::MatrixXd sigma_x;

    /// @brief Expectation of sigma_y, dimension: (num_states, num_kpt)
    Eigen::MatrixXd sigma_y;

    /// @brief Expectation of sigma_z, dimension: (num_states, num_kpt)
    Eigen::MatrixXd sigma_z;

    /**
     * @brief Construct a new SpinData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit SpinData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}_k_cart.dat, ${prefix}_${data_name}_bands.dat, ${prefix}_${data_name}_sigma_*.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.k_cart, ${data_name}.bands, ${data_name}.sigma_*
     *
     * 2. Data format of dat output of k_cart:
     *      dimension: (num_kpt, 3)
     *      columns: kx, ky, kz of k_cart
     *
     *   Data format of dat output of bands:
     *     dimension: (num_kpt, num_states)
     *     columns: band0, band1, ...
     *
     *   Data format of dat output of sigma_*:
     *     dimension (num_kpt, num_states)
     *     columns: sigma0, sigma1, ...
     */
    void save(const std::string& data_name = "spin") const;
};

/**
 * @brief Class for holding the results of Berry curvature calculation
 *
 */
class BerryData : public BaseData {
public:
    /// @brief Cartesian coordinates of k-points in 1/NM, dimension: (3, num_kpt)
    Eigen::Matrix3Xd k_cart;

    /// @brief Energies on k-points in eV, dimension: (num_states, num_kpt)
    Eigen::MatrixXd bands;

    /// @brief xy component of Berry curvature in NM^2, dimension: (num_states, num_kpt)
    Eigen::MatrixXd omega_xy;

    /**
     * @brief Construct a new BerryData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit BerryData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}_k_cart.dat, ${prefix}_${data_name}_bands.dat, ${prefix}_${data_name}_omega_xy.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.k_cart, ${data_name}.bands, ${data_name}.omega_xy
     *
     * 2. Data format of dat output of k_cart:
     *      dimension: (num_kpt, 3)
     *      columns: kx, ky, kz of k_cart
     *
     *   Data format of dat output of bands:
     *      dimension: (num_kpt, num_states)
     *      columns: band0, band1, ...
     *
     *   Data format of dat output of omega_xy:
     *      dimension (num_kpt, num_states)
     *      columns: omega_k0, omega_k1, omega_k2, ...
     */
    void save(const std::string& data_name = "berry") const;
};

/**
 * @brief Data class for holding the phase factor of Wannier centers
 *
 */
class Z2Data : public BaseData {
public:
    /// @brief FRACTIONAL coordinates of the loop along a-axis, dimension: (num_ka, 1)
    Eigen::VectorXd k_a;

    /// @brief FRACTIONAL coordinates of the loop along b-axis, dimension: (num_kb, 1)
    Eigen::VectorXd k_b;

    /// @brief phase factor of WF centers, dimension: (num_occ, num_kpt)
    Eigen::MatrixXd phases;

    /**
     * @brief Construct a new Z2Data object
     *
     * @param[in] prefix prefix of data files
     */
    explicit Z2Data(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.k_b and ${data_name}.phases
     *
     * 2. Data format of dat output:
     *      dimension: (num_kpt, num_occ+1)
     *      columns: k_b, phase0, phase1, ...
     */
    void save(const std::string& data_name = "z2") const;
};

/**
 * @brief Data class for holding dynamic polarizability.
 *
 */
class DynPolData : public BaseData {
public:
    /// @brief energy grid corresponding to e_min, e_max and e_step in eV
    /// dimension: (num_eng, 1)
    Eigen::VectorXd energies;

    /// @brief dynamic polarizability in 1/(eV*nm**2) in 2d case 1/(eV*nm**3) in 3d case,
    /// dimension: (num_eng, num_qpt)
    Eigen::MatrixXcd dyn_pol;

    /**
     * @brief Construct a new DynPolData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit DynPolData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.energies and ${data_name}.dyn_pol
     *
     * 2. Data format of dat output:
     *      dimension (num_eng, 1+2*num_qpt)
     *      columns: energies, (dp_q0.real, dp_q0.imag), (dp_q1.real, dp_q1.imag), ...
     */
    void save(const std::string& data_name = "dyn_pol") const;
};

/**
 * @brief Data class for holding dielectric function.
 *
 */
class EpsilonData : public BaseData {
public:
    /// @brief energy grid corresponding to e_min, e_max and e_step in eV
    /// dimension: (num_eng, 1)
    Eigen::VectorXd energies;
    /// @brief relative dielectric function, dimension: (num_eng, num_qpt)
    Eigen::MatrixXcd epsilon;

    /**
     * @brief Construct a new EpsilonData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit EpsilonData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.energies and ${data_name}.epsilon
     *
     *2. Data format of dat output:
     *     dimension (num_eng, 1+2*num_qpt)
     *     columns: energies, (eps_q0.real, eps_q0.imag), (eps_q1.real, eps_q1.imag), ...
     */
    void save(const std::string& data_name = "epsilon") const;
};

/**
 * @brief Data class for holding AC conductivity.
 *
 */
class ACCondData : public BaseData {
public:
    /// @brief energy grid corresponding to e_min, e_max and e_step in eV
    /// dimension: (num_eng, 1)
    Eigen::VectorXd energies;

    /// @brief AC conductivity in e**2/(h_bar*nm) in 3d case and e**2/h_bar in 2d case
    /// dimension: (num_eng, 1)
    Eigen::VectorXcd ac_cond;

    /**
     * @brief Construct a new ACCondData object
     *
     * @param[in] prefix prefix of data files
     */
    explicit ACCondData(const std::string& prefix = "sample")
        : BaseData(prefix)
    {
    }

    /**
     * @brief Save to file
     *
     * @param[in] data_name name of data
     *
     * @note
     * 1. Output files:
     *      ${prefix}_${data_name}.dat (no HDF5)
     *      ${prefix}_diag.h5, with datasets ${data_name}.energies and ${data_name}.ac_cond
     *
     * 2. Data format of dat output:
     *      dimension (num_eng, 3)
     *      columns: energies, (ac_cond.real, ac_cond.imag)
     */
    void save(const std::string& data_name = "ac_cond") const;
};
} // namespace tbplas::diag
#endif