/**
 * @brief Functions and classes for I/O.
 *
 * @author Yunhai Li (liyunhai1016@hotmail.com)
 *
 * @copyright Copyright (c) 2025, TBPLaS develop team.
 */

#ifndef TBPLAS_BASE_IO_TEXT_H
#define TBPLAS_BASE_IO_TEXT_H

#include <fstream>
#include <iomanip>
#include <string>

#include <eigen3/Eigen/Dense>

namespace tbplas::base {

/**
 * @brief Plain text file.
 *
 */
class TextFile {
public:
    /**
     * @brief Construct a new Text File object
     *
     * @param[in] prefix prefix of data files
     * @param[in] transpose whether to transpose the matrix
     */
    explicit TextFile(
        const std::string& prefix = "sample",
        bool transpose = false)
        : prefix_(prefix)
        , transpose_(transpose)
    {
    }

    /// @brief Setters for attributes
    void set_prefix(const std::string& prefix);
    void set_transpose(bool transpose);
    void set_scientific(bool scientific);
    void set_precision(int precision);
    void set_width(int width);

    /**
     * @brief Save real Eigen matrix to file
     *
     * @tparam derived_t type of Eigen matrix
     * @param[in] data_name name of data
     * @param[in] matrix matrix to save
     *
     * @note
     *   1. The compiler cannot distinguish MatrixXd and MatrixXcd, so we have to
     *      use two separate functions.
     *   2. The actual file name will be $prefix_$data_name.dat.
     */
    template <typename derived_t>
    void save_matrix_real(
        const std::string& data_name,
        const Eigen::MatrixBase<derived_t>& matrix) const
    {
        std::ofstream outfile;
        outfile.open(prefix_ + "_" + data_name + ".dat", std::ios::out);

        // Set precision
        if (scientific_) {
            outfile.setf(std::ios::scientific);
        } else {
            outfile.setf(std::ios::fixed);
        }

        // Meta info.
        outfile.precision(precision_);
        outfile << "# datatype real\n";
        if (transpose_) {
            outfile << "# transposed true\n";
            outfile << "# num_rows " << matrix.cols() << "\n";
            outfile << "# num_cols " << matrix.rows() << "\n";
        } else {
            outfile << "# transposed false\n";
            outfile << "# num_rows " << matrix.rows() << "\n";
            outfile << "# num_cols " << matrix.cols() << "\n";
        }

        // Data block
        if (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(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(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] data_name name of data
     * @param[in] matrix matrix to save
     *
     * @note
     *   1. The compiler cannot distinguish MatrixXd and MatrixXcd, so we have to
     *      use two separate functions.
     *   2. The actual file name will be $prefix_$data_name.dat.
     */
    template <typename derived_t>
    void save_matrix_complex(
        const std::string& data_name,
        const Eigen::MatrixBase<derived_t>& matrix) const
    {
        std::ofstream outfile;
        outfile.open(prefix_ + "_" + data_name + ".dat", std::ios::out);

        // Set precision
        if (scientific_) {
            outfile.setf(std::ios::scientific);
        } else {
            outfile.setf(std::ios::fixed);
        }
        outfile.precision(precision_);

        // Meta info.
        outfile << "# datatype complex\n";
        if (transpose_) {
            outfile << "# transposed true\n";
            outfile << "# num_rows " << matrix.cols() << "\n";
            outfile << "# num_cols " << matrix.rows() << "\n";
        } else {
            outfile << "# transposed false\n";
            outfile << "# num_rows " << matrix.rows() << "\n";
            outfile << "# num_cols " << matrix.cols() << "\n";
        }

        // Data block
        if (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(width_) << matrix(i_row, i_col).real()
                            << std::setw(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(width_) << matrix(i_row, i_col).real()
                            << std::setw(width_) << matrix(i_row, i_col).imag();
                }
                outfile << "\n";
            }
        }
        outfile.close();
    }

private:
    /// @brief prefix of data files
    std::string prefix_ = "sample";

    /// @brief whether to transpose the matrix
    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 for each number
    int width_ = 25;
};
} // namespace tbplas::base

#endif