#include "results.h"

#ifdef WITH_HDF5
#include "base/io_base.h"
#include "base/io_hdf5.h"
#else
#include "base/io_text.h"
#endif

namespace tbplas::diag {
// Aliases
#ifdef WITH_HDF5
using tbplas::base::HDFFile;
using tbplas::base::OptBinMode;
#else
using tbplas::base::TextFile;
#endif

void BandData::save(const std::string& data_name) const
{
#ifdef WITH_HDF5
    HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
    data_file.save_matrix_real(data_name + ".k_len", k_len);
    data_file.save_matrix_real(data_name + ".bands", bands);
    data_file.close();
#else
    size_t num_kpt = bands.cols();
    size_t num_states = bands.rows();
    Eigen::MatrixXd data(num_states + 1, bands.cols());
    data.row(0) = k_len.transpose();
    data.block(1, 0, num_states, num_kpt) = bands;
    TextFile data_file(prefix_, true);
    data_file.save_matrix_real(data_name, data);
#endif
}

void DOSData::save(const std::string& data_name) const
{
#ifdef WITH_HDF5
    HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
    data_file.save_matrix_real(data_name + ".energies", energies);
    data_file.save_matrix_real(data_name + ".dos", dos);
    data_file.close();
#else
    Eigen::MatrixXd data(dos.size(), 2);
    data.col(0) = energies;
    data.col(1) = dos;
    TextFile data_file(prefix_, false);
    data_file.save_matrix_real(data_name, data);
#endif
}

void EigenStateData::save_eigen_values(const std::string& data_name) const
{
    BandData::save(data_name);
}

void EigenStateData::save_eigen_states(
    const int& ik,
    const Eigen::MatrixXcd& eigen_states,
    const std::string& data_name) const
{
#ifdef WITH_HDF5
    // Must use separate h5 files for each k-point since all MPI processes have
    // data to write.
    HDFFile data_file(prefix_ + "_" + data_name + "_k" + std::to_string(ik), OptBinMode::UPDATE);
    data_file.save_matrix_complex("eigen_states", eigen_states);
    data_file.close();
#else
    TextFile data_file(prefix_, false);
    data_file.save_matrix_complex(data_name + "_k" + std::to_string(ik), eigen_states);
#endif
}

void SpinData::save(const std::string& data_name) const
{
#ifdef WITH_HDF5
    HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
    data_file.save_matrix_real(data_name + ".k_cart", k_cart);
    data_file.save_matrix_real(data_name + ".bands", bands);
    data_file.save_matrix_real(data_name + ".sigma_x", sigma_x);
    data_file.save_matrix_real(data_name + ".sigma_y", sigma_y);
    data_file.save_matrix_real(data_name + ".sigma_z", sigma_z);
    data_file.close();
#else
    TextFile data_file(prefix_, true);
    data_file.save_matrix_real(data_name + "_k_cart", k_cart);
    data_file.save_matrix_real(data_name + "_bands", bands);
    data_file.save_matrix_real(data_name + "_sigma_x", sigma_x);
    data_file.save_matrix_real(data_name + "_sigma_y", sigma_y);
    data_file.save_matrix_real(data_name + "_sigma_z", sigma_z);
#endif
}

void BerryData::save(const std::string& data_name) const
{
#ifdef WITH_HDF5
    HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
    data_file.save_matrix_real(data_name + ".k_cart", k_cart);
    data_file.save_matrix_real(data_name + ".bands", bands);
    data_file.save_matrix_real(data_name + ".omega_xy", omega_xy);
    data_file.close();
#else
    TextFile data_file(prefix_, true);
    data_file.save_matrix_real(data_name + "_k_cart", k_cart);
    data_file.save_matrix_real(data_name + "_bands", bands);
    data_file.save_matrix_real(data_name + "_omega_xy", omega_xy);
#endif
}

void Z2Data::save(const std::string& data_name) const
{
#ifdef WITH_HDF5
    HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
    data_file.save_matrix_real(data_name + ".k_b", k_b);
    data_file.save_matrix_real(data_name + ".phases", phases);
    data_file.close();
#else
    size_t num_kpt = phases.cols();
    size_t num_states = phases.rows();
    Eigen::MatrixXd data(num_states + 1, phases.cols());
    data.row(0) = k_b.transpose();
    data.block(1, 0, num_states, num_kpt) = phases;
    TextFile data_file(prefix_, true);
    data_file.save_matrix_real(data_name, data);
#endif
}

void DynPolData::save(const std::string& data_name) const
{
    size_t num_eng = energies.size();
    size_t num_qpt = dyn_pol.cols();
    if ((num_eng > 0) && (num_qpt > 0)) {
#ifdef WITH_HDF5
        HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
        data_file.save_matrix_real(data_name + ".energies", energies);
        data_file.save_matrix_complex(data_name + ".dyn_pol", dyn_pol);
        data_file.close();
#else
        Eigen::MatrixXd data(num_eng, 1 + 2 * num_qpt);
        data.col(0) = energies;
        for (ptrdiff_t iq = 0; iq < num_qpt; ++iq) {
            data.col(iq * 2 + 1) = dyn_pol.col(iq).real();
            data.col(iq * 2 + 2) = dyn_pol.col(iq).imag();
        }
        TextFile data_file(prefix_, false);
        data_file.save_matrix_real(data_name, data);
#endif
    }
}

void EpsilonData::save(const std::string& data_name) const
{
    size_t num_eng = energies.size();
    size_t num_qpt = epsilon.cols();
    if ((num_eng > 0) && (epsilon.cols() > 0)) {
#ifdef WITH_HDF5
        HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
        data_file.save_matrix_real(data_name + ".energies", energies);
        data_file.save_matrix_complex(data_name + ".epsilon", epsilon);
        data_file.close();
#else
        Eigen::MatrixXd data(num_eng, 1 + 2 * num_qpt);
        data.col(0) = energies;
        for (ptrdiff_t iq = 0; iq < num_qpt; ++iq) {
            data.col(iq * 2 + 1) = epsilon.col(iq).real();
            data.col(iq * 2 + 2) = epsilon.col(iq).imag();
        }
        TextFile data_file(prefix_, false);
        data_file.save_matrix_real(data_name, data);
#endif
    }
}

void ACCondData::save(const std::string& data_name) const
{
    size_t num_eng = energies.size();
    if (ac_cond.size() > 0) {
#ifdef WITH_HDF5
        HDFFile data_file(prefix_ + "_diag", OptBinMode::UPDATE);
        data_file.save_matrix_real(data_name + ".energies", energies);
        data_file.save_matrix_complex(data_name + ".ac_cond", ac_cond);
        data_file.close();
#else
        Eigen::MatrixXd data(num_eng, 3);
        data.col(0) = energies;
        data.col(1) = ac_cond.real();
        data.col(2) = ac_cond.imag();
        TextFile data_file(prefix_, false);
        data_file.save_matrix_real(data_name, data);
#endif
    }
}
} // namespace tbplas::diag