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

#ifndef TBPLAS_BASE_RESULTS_H
#define TBPLAS_BASE_RESULTS_H

#include <fstream>
#include <iomanip>

#include <eigen3/Eigen/Dense>

namespace tbplas::base {

/**
 * @brief Data class controlling output format of matrices.
 *
 */
class MatrixFormat {
public:
    /// @brief whether to output in transposed form
    bool transpose = false;

    /// @brief whether to use scientific notation
    bool scientific = true;

    /// @brief number of digits of output
    int precision = 15;

    /// @brief width of output
    int width = 25;
};

/**
 * @brief Save real Eigen matrix to file
 *
 * @tparam derived_t type of Eigen matrix
 * @param[in] matrix matrix to save
 * @param[in] file_name file name of data
 * @param[in] fmt output format of matrix
 *
 * @note The compiler cannot distinguish MatrixXd and MatrixXcd, so we have to
 *   use two separate functions.
 */
template <typename derived_t>
void save_matrix_real(
    const Eigen::MatrixBase<derived_t>& matrix,
    const std::string& file_name = "data.dat",
    const MatrixFormat& fmt = MatrixFormat())
{
    std::ofstream outfile;
    outfile.open(file_name, std::ios::out);
    if (fmt.scientific) {
        outfile.setf(std::ios::scientific);
    } else {
        outfile.setf(std::ios::fixed);
    }

    // Meta info.
    outfile.precision(fmt.precision);
        outfile << "# datatype real\n";
    if (fmt.transpose) {
        outfile << "# transposed true\n";
    } else {
        outfile << "# transposed false\n";
    }

    // Data block
    if (fmt.transpose) {
        for (auto i_col = 0; i_col < matrix.cols(); ++i_col) {
            for (auto i_row = 0; i_row < matrix.rows(); ++i_row) {
                outfile << std::setw(fmt.width) << matrix(i_row, i_col);
            }
            outfile << "\n";
        }
    } else {
        for (auto i_row = 0; i_row < matrix.rows(); ++i_row) {
            for (auto i_col = 0; i_col < matrix.cols(); ++i_col) {
                outfile << std::setw(fmt.width) << matrix(i_row, i_col);
            }
            outfile << "\n";
        }
    }
    outfile.close();
}

/**
 * @brief Save complex Eigen matrix to file
 *
 * @tparam derived_t type of Eigen matrix
 * @param[in] matrix matrix to save
 * @param[in] file_name file name of data
 * @param[in] fmt output format of matrix
 *
 * @note The compiler cannot distinguish MatrixXd and MatrixXcd, so we have to
 *   use two separate functions.
 */
template <typename derived_t>
void save_matrix_complex(
    const Eigen::MatrixBase<derived_t>& matrix,
    const std::string& file_name = "data.dat",
    const MatrixFormat& fmt = MatrixFormat())
{
    std::ofstream outfile;
    outfile.open(file_name, std::ios::out);
    if (fmt.scientific) {
        outfile.setf(std::ios::scientific);
    } else {
        outfile.setf(std::ios::fixed);
    }
    outfile.precision(fmt.precision);

    // Meta info.
    outfile << "# datatype complex\n";
    if (fmt.transpose) {
        outfile << "# transposed true\n";
    } else {
        outfile << "# transposed false\n";
    }

    // Data block
    if (fmt.transpose) {
        for (auto i_col = 0; i_col < matrix.cols(); ++i_col) {
            for (auto i_row = 0; i_row < matrix.rows(); ++i_row) {
                outfile << std::setw(fmt.width) << matrix(i_row, i_col).real()
                        << std::setw(fmt.width) << matrix(i_row, i_col).imag();
            }
            outfile << "\n";
        }
    } else {
        for (auto i_row = 0; i_row < matrix.rows(); ++i_row) {
            for (auto i_col = 0; i_col < matrix.cols(); ++i_col) {
                outfile << std::setw(fmt.width) << matrix(i_row, i_col).real()
                        << std::setw(fmt.width) << matrix(i_row, i_col).imag();
            }
            outfile << "\n";
        }
    }
    outfile.close();
}

} // namespace tbplas::base
#endif