//
// Created by Emily Boudreaux on 4/10/25.
//

#ifndef MFEM_SMOUT_H
#define MFEM_SMOUT_H

#include "mfem.hpp"
#include <iostream>
#include <fstream>
#include <vector>
#include <array>
#include <iomanip>
#include <tuple>
#include <ranges>

/**
 * @brief Saves an mfem::SparseMatrix to a custom compact binary file (.csrbin).
 *
 * @param mat The mfem::SparseMatrix to save (assumed to be in CSR format).
 * @param filename The path to the output file.
 * @return true if saving was successful, false otherwise.
 *
 * File Format (.csrbin):
 * - Magic (4 bytes): 'C','S','R','B'
 * - Version (1 byte): 1
 * - IntSize (1 byte): 8 (using int64_t for indices/dims)
 * - FltSize (1 byte): 8 (using double for data)
 * - Reserved (1 byte): 0
 * - Height (uint64_t): Number of rows
 * - Width (uint64_t): Number of columns
 * - NNZ (uint64_t): Number of non-zeros
 * - I array (int64_t * (Height + 1)): CSR Row Pointers
 * - J array (int64_t * NNZ): CSR Column Indices
 * - Data array (double * NNZ): CSR Non-zero values
 */
inline void write_sparse_matrix(mfem::Operator &op, std::ostream &outfile, bool neg) {
    mfem::SparseMatrix* mat;
    try {
        mat = dynamic_cast<mfem::SparseMatrix*>(&op);
        if (!mat) {
            throw std::runtime_error("The operator is not a SparseMatrix.");
        }
    } catch (const std::runtime_error&) {
        try {
            const auto& blf = dynamic_cast<mfem::BilinearForm*>(&op);
            if (!blf) {
                throw std::runtime_error("The operator is not a SparseMatrix or BilinearForm.");
            }
            mat = &blf->SpMat();
        } catch (const std::runtime_error&) {
            auto mblf = dynamic_cast<mfem::MixedBilinearForm*>(&op);
            if (!mblf) {
                throw std::runtime_error("The operator is not a SparseMatrix or BilinearForm or MixedBilinear Form.");
            }
            mat = &mblf->SpMat();
        }
    }
    if (neg) {
        *mat *= -1.0;
    }
    // --- Get Data Pointers and Dimensions from MFEM Matrix ---
    const int* mfem_I = mat->GetI();
    const int* mfem_J = mat->GetJ();
    const double* mfem_data = mat->GetData();

    auto height = static_cast<uint64_t>(mat->Height());
    auto width = static_cast<uint64_t>(mat->Width());
    auto nnz = static_cast<uint64_t>(mat->NumNonZeroElems());
    uint64_t i_count = height + 1;
    uint64_t j_count = nnz;
    uint64_t data_count = nnz;


    // --- Write Header ---
    const char magic[4] = {'C', 'S', 'R', 'B'};
    const uint8_t version = 1;
    const uint8_t int_size = 8;
    const uint8_t flt_size = 8;
    const uint8_t reserved = 0;

    outfile.write(magic, 4);
    outfile.write(reinterpret_cast<const char*>(&version), 1);
    outfile.write(reinterpret_cast<const char*>(&int_size), 1);
    outfile.write(reinterpret_cast<const char*>(&flt_size), 1);
    outfile.write(reinterpret_cast<const char*>(&reserved), 1);

    outfile.write(reinterpret_cast<const char*>(&height), sizeof(height));
    outfile.write(reinterpret_cast<const char*>(&width), sizeof(width));
    outfile.write(reinterpret_cast<const char*>(&nnz), sizeof(nnz));

    if (!outfile) throw std::runtime_error("Error writing header.");

    // --- Write Arrays (Converting int to int64_t for I and J) ---
    std::vector<int64_t> i_buffer(i_count);
    for (uint64_t idx = 0; idx < i_count; ++idx) {
        i_buffer[idx] = static_cast<int64_t>(mfem_I[idx]);
    }
    outfile.write(reinterpret_cast<const char*>(i_buffer.data()), i_count * sizeof(int64_t));
    if (!outfile) throw std::runtime_error("Error writing I array.");

    std::vector<int64_t> j_buffer(j_count);
    for (uint64_t idx = 0; idx < j_count; ++idx) {
        j_buffer[idx] = static_cast<int64_t>(mfem_J[idx]);
    }
    outfile.write(reinterpret_cast<const char*>(j_buffer.data()), j_count * sizeof(int64_t));
    if (!outfile) throw std::runtime_error("Error writing J array.");

    outfile.write(reinterpret_cast<const char*>(mfem_data), data_count * sizeof(double));
    if (!outfile) throw std::runtime_error("Error writing Data array.");
}

inline bool saveSparseMatrixBinary(mfem::SparseMatrix& mat, const std::string& filename) {
    std::ofstream outfile(filename, std::ios::binary | std::ios::trunc);
    if (!outfile) {
        std::cerr << "Error: Cannot open file for writing: " << filename << std::endl;
        return false;
    }

    try {
        write_sparse_matrix(mat, outfile, false);


    } catch (const std::exception& e) {
        std::cerr << "Error during binary matrix save: " << e.what() << std::endl;
        outfile.close();
        return false;
    }

    outfile.close();
    if (!outfile) {
        std::cerr << "Error closing file after writing: " << filename << std::endl;
        return false;
    }
    return true;
}

inline void writeDenseMatrixToCSV(const std::string &filename, int precision, const mfem::DenseMatrix *mat) {
    if (!mat) {
        throw std::runtime_error("The operator is not a SparseMatrix.");
    }

    std::ofstream outfile(filename);
    if (!outfile.is_open()) {
        throw std::runtime_error("Failed to open file: " + filename);
    }


    int height = mat->Height();
    int width = mat->Width();

    // Set precision for floating-point output
    outfile << std::fixed << std::setprecision(precision);

    for (int i = 0; i < width; i++) {
        outfile << i;
        if (i < width - 1) {
            outfile << ",";
        }
        else {
            outfile << "\n";
        }
    }

    // Iterate through rows
    for (int i = 0; i < height; ++i) {
        for (int j = 0; j < width; ++j) {
            outfile << mat->Elem(i, j);
            if (j < width - 1) {
                outfile << ",";
            }
        }
        outfile << std::endl;
    }

    outfile.close();
}

/**
 * @brief Writes the dense representation of an MFEM Operator (if it's a SparseMatrix) to a CSV file.
 *
 * @param op The MFEM Operator to write.
 * @param filename The name of the output CSV file.
 * @param precision Number of decimal places for floating-point values.
 */
inline void writeOperatorToCSV(const mfem::Operator &op,
                               const std::string &filename,
                               int precision = 6) // Add precision argument
{
    // Attempt to cast the Operator to a SparseMatrix
    const auto *sparse_mat = dynamic_cast<const mfem::SparseMatrix*>(&op);
    if (!sparse_mat) {
        throw std::runtime_error("The operator is not a SparseMatrix.");
    }
    const mfem::DenseMatrix *mat = sparse_mat->ToDenseMatrix();
    writeDenseMatrixToCSV(filename, precision, mat);
}

inline void saveBlockFormToBinary(const std::vector<mfem::Operator*> &block_diags, const std::vector<std::array<int, 2>>& block, const std::vector<bool>& neg, const std::string &filename) {
    // First write a magic number and version

    // --- Open the file ---
    std::ofstream outfile(filename, std::ios::binary | std::ios::trunc);
    if (!outfile) {
        std::cerr << "Error: Cannot open file for writing: " << filename << std::endl;
        return;
    }

    // --- Write Header ---
    const char magic[4] = {'B', 'L', 'C', 'K'};
    const char datastart[9] = {'D', 'A', 'T', 'A', 'S', 'T', 'A', 'R', 'T'};
    const char dataend[7] = {'D', 'A', 'T', 'A', 'E', 'N', 'D'};
    const uint8_t size = block_diags.size();

    outfile.write(reinterpret_cast<const char*>(&magic), 4);
    outfile.write(reinterpret_cast<const char*>(&size), sizeof(size));

    for (const auto&& [block_diag, blockIDs, isNeg] : std::views::zip(block_diags, block, neg)) {
        // Write the sparse matrix data
        outfile.write(reinterpret_cast<const char*>(&datastart), 9);
        outfile.write(reinterpret_cast<const char*>(&blockIDs), sizeof(blockIDs));
        write_sparse_matrix(*block_diag, outfile, isNeg);
        outfile.write(reinterpret_cast<const char*>(&dataend), 7);
    }

}

#endif //MFEM_SMOUT_H