/* ***********************************************************************
//
//   Copyright (C) 2025 -- The 4D-STAR Collaboration
//   File Author: Emily Boudreaux
//   Last Modified: April 21, 2025
//
//   4DSSE is free software; you can use it and/or modify
//   it under the terms and restrictions the GNU General Library Public
//   License version 3 (GPLv3) as published by the Free Software Foundation.
//
//   4DSSE is distributed in the hope that it will be useful,
//   but WITHOUT ANY WARRANTY; without even the implied warranty of
//   MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
//   See the GNU Library General Public License for more details.
//
//   You should have received a copy of the GNU Library General Public License
//   along with this software; if not, write to the Free Software
//   Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
//
// *********************************************************************** */
#include "operator.h"
#include "4DSTARTypes.h"
#include "mfem.hpp"
#include <memory>

static int s_PolyOperatorMultCalledCount = 0;

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.
 */
 void writeOperatorToCSV(const mfem::Operator &op,
    const std::string &filename,
    const 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);
}

void approxJacobiInvert(const mfem::SparseMatrix& mat, std::unique_ptr<mfem::SparseMatrix>& invMat, const std::string& name="matrix") {
    // PERF: This likely can be made much more efficient and will probably be called in tight loops, a good
    // PERF:   place for some easy optimization might be here.

    // Confirm that mat is a square matrix
    MFEM_ASSERT(mat.Height() == mat.Width(), "Matrix " + name + " is not square, cannot invert.");

    mfem::Vector diag;
    mat.GetDiag(diag);

    // Invert the diagonal
    for (int i = 0; i < diag.Size(); i++) {
        MFEM_ASSERT(diag(i) != 0, "Diagonal element (" + std::to_string(i) +") in " + name + " is zero, cannot invert.");
        diag(i) = 1.0 / diag(i);
    }

    // If the matrix is already inverted, just set the diagonal to avoid reallocation
    if (invMat != nullptr) {
        MFEM_ASSERT(invMat->Height() == invMat->Width(), "invMat (result matrix) is not square, cannot invert " + name + " into it.");
        MFEM_ASSERT(invMat->Height() == mat.Height(), "Incompatible matrix sizes for inversion of " + name + ", expected " + std::to_string(mat.Height()) + " but got " + std::to_string(invMat->Height()));
        for (int i = 0; i < diag.Size(); i++) {
            MFEM_ASSERT(diag(i) != 0, "Diagonal element (" + std::to_string(i) +") in " + name + " is zero, resulting matrix would be singular.");
            invMat->Elem(i, i) = diag(i);
        }
    } else { // The matrix has not been allocated yet so that needs to be done. Sparse Matrix has a constructor that can build from the diagonals
        invMat = std::make_unique<mfem::SparseMatrix>(diag);
    }
}

PolytropeOperator::PolytropeOperator(

  std::unique_ptr<mfem::MixedBilinearForm> M,
  std::unique_ptr<mfem::MixedBilinearForm> Q,
  std::unique_ptr<mfem::BilinearForm> D,
  std::unique_ptr<mfem::NonlinearForm> f,
  const mfem::Array<int> &blockOffsets) : 

  mfem::Operator(blockOffsets.Last()), // Initialize the base class with the total size of the block offset vector
  m_blockOffsets(blockOffsets),
  m_jacobian(nullptr) {

    m_M = std::move(M);
    m_Q = std::move(Q);
    m_D = std::move(D);
    m_f = std::move(f);
}

void PolytropeOperator::finalize(const mfem::Vector &initTheta) {
    if (m_isFinalized) {
        return; // do nothing if already finalized
    }

    m_negM_op = std::make_unique<mfem::ScaledOperator>(m_M.get(), -1.0);
    m_negQ_op = std::make_unique<mfem::ScaledOperator>(m_Q.get(), -1.0);

    // Set up the constant parts of the jacobian now
    m_jacobian = std::make_unique<mfem::BlockOperator>(m_blockOffsets);
    m_jacobian->SetBlock(0, 1, m_negM_op.get()); //<- -M (constant)
    m_jacobian->SetBlock(1, 0, m_negQ_op.get()); //<- -Q (constant)
    m_jacobian->SetBlock(1, 1, m_D.get());       //<-  D (constant)

    // Build the initial preconditioner based on some initial guess
    const auto &grad = m_f->GetGradient(initTheta);
    updatePreconditioner(grad);

    m_isFinalized = true;
}

const mfem::BlockOperator &PolytropeOperator::GetJacobianOperator() const {
    if (m_jacobian == nullptr) {
        MFEM_ABORT("Jacobian has not been initialized before GetJacobianOperator() call.");
    }
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator not finalized prior to call to GetJacobianOperator().");
    }

    return *m_jacobian;
}

mfem::BlockDiagonalPreconditioner& PolytropeOperator::GetPreconditioner() const {
    if (m_schurPreconditioner == nullptr) {
        MFEM_ABORT("Schur preconditioner has not been initialized before GetPreconditioner() call.");
    }
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator not finalized prior to call to GetPreconditioner().");
    }
    return *m_schurPreconditioner;
}

void PolytropeOperator::Mult(const mfem::Vector &x, mfem::Vector &y) const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator::Mult called before finalize");
    }
    // -- Create BlockVector views for input x and output y
    mfem::BlockVector x_block(const_cast<mfem::Vector&>(x), m_blockOffsets);
    mfem::BlockVector y_block(y, m_blockOffsets);

    // -- Get Vector views for individual blocks
    const mfem::Vector &x_theta = x_block.GetBlock(0);
    const mfem::Vector &x_phi = x_block.GetBlock(1);
    mfem::Vector &y_R0 = y_block.GetBlock(0); // Residual Block 0 (theta)
    mfem::Vector &y_R1 = y_block.GetBlock(1); // Residual Block 1 (phi)

    int theta_size = m_blockOffsets[1] - m_blockOffsets[0];
    int phi_size = m_blockOffsets[2] - m_blockOffsets[1];

    mfem::Vector f_term(theta_size);
    mfem::Vector Mphi_term(theta_size);
    mfem::Vector Dphi_term(phi_size);
    mfem::Vector Qtheta_term(phi_size);

    // Calculate R0 and R1 terms
    // R0 = f(θ) - Mɸ
    // R1 = Dɸ - Qθ

    MFEM_ASSERT(m_f.get() != nullptr, "NonlinearForm m_f is null in PolytropeOperator::Mult");

    m_f->Mult(x_theta, f_term); // fixme: this may be the wrong way to assemble m_f?
    m_M->Mult(x_phi, Mphi_term);
    m_D->Mult(x_phi, Dphi_term);
    m_Q->Mult(x_theta, Qtheta_term);

    subtract(f_term, Mphi_term, y_R0);
    subtract(Dphi_term, Qtheta_term, y_R1);

    // -- Apply essential boundary conditions --
    for (int i = 0; i < m_theta_ess_tdofs.first.Size(); i++) {
        if (int idx = m_theta_ess_tdofs.first[i]; idx >= 0 && idx < y_R0.Size()) {
            const double &targetValue = m_theta_ess_tdofs.second[i];
            // y_block.GetBlock(0)[idx] = targetValue - x_theta(idx); // inhomogenous essential bc. This is commented out since seems it results in dramatic instabilies arrising
            y_block.GetBlock(0)[idx] = 0; // Zero out the essential theta dofs in the bi-linear form
        }
    }

    // TODO: Are the residuals for φ being calculated correctly?

    std::ofstream outfileθ("PolyOperatorMultCallTheta_" + std::to_string(s_PolyOperatorMultCalledCount) + ".csv", std::ios::out | std::ios::trunc);
    outfileθ << "dof,R_θ\n";
    for (int i = 0; i < y_R0.Size(); i++) {
        outfileθ << i << "," << y_R0(i) << "\n";
    }
    outfileθ.close();
    std::ofstream outfileφ("PolyOperatorMultCallPhi_" + std::to_string(s_PolyOperatorMultCalledCount) + ".csv", std::ios::out | std::ios::trunc);
    outfileφ << "dof,R_ɸ\n";
    for (int i = 0; i < y_R1.Size(); i++) {
        outfileφ << i << "," << y_R1(i) << "\n";
    }
    outfileφ.close();
    ++s_PolyOperatorMultCalledCount;

}


void PolytropeOperator::updateInverseNonlinearJacobian(const mfem::Operator &grad) const {
    if (const auto *sparse_mat = dynamic_cast<const mfem::SparseMatrix*>(&grad); sparse_mat != nullptr) {
        approxJacobiInvert(*sparse_mat, m_invNonlinearJacobian, "Nonlinear Jacobian");
    } else {
        MFEM_ABORT("PolytropeOperator::GetGradient called on nonlinear jacobian where nonlinear jacobian is not dynamically castable to a sparse matrix");
    }
}

void PolytropeOperator::updateInverseSchurCompliment() const {
    // TODO Add a flag in to make sure this tracks in parallel (i.e. every time the non linear jacobian inverse is updated set the flag to true and then check if the flag is true here and if so do work (if not throw error). then at the end of this function set it to false.
    if (m_invNonlinearJacobian == nullptr) {
        MFEM_ABORT("PolytropeOperator::updateInverseSchurCompliment called before updateInverseNonlinearJacobian");
    }
    mfem::SparseMatrix* schurCompliment(&m_D->SpMat()); // This is now a copy of D

    // Calculate S = D - Q df^{-1} M
    mfem::SparseMatrix* temp_QxdfInv = mfem::Mult(m_Q->SpMat(), *m_invNonlinearJacobian); // Q * df^{-1}
    const mfem::SparseMatrix* temp_QxdfInvxM = mfem::Mult(*temp_QxdfInv, m_M->SpMat()); // Q * df^{-1} * M

    // PERF: Could potentially add some caching in here to only update the preconditioner when some condition has been met
    schurCompliment->Add(-1, *temp_QxdfInvxM); // D - Q * df^{-1} * M
    approxJacobiInvert(*schurCompliment, m_invSchurCompliment, "Schur Compliment");

    if (m_schurPreconditioner == nullptr) {
        m_schurPreconditioner = std::make_unique<mfem::BlockDiagonalPreconditioner>(m_blockOffsets);
    }

    // ⎡ḟ(θ)^-1 0⎤
    // ⎣0     S^-1⎦
    m_schurPreconditioner->SetDiagonalBlock(0, m_invNonlinearJacobian.get());
    m_schurPreconditioner->SetDiagonalBlock(1, m_invSchurCompliment.get());
}

void PolytropeOperator::updatePreconditioner(const mfem::Operator &grad) const {
    updateInverseNonlinearJacobian(grad);
    updateInverseSchurCompliment();
}


mfem::Operator& PolytropeOperator::GetGradient(const mfem::Vector &x) const {
    if (!m_isFinalized) {
        MFEM_ABORT("PolytropeOperator::GetGradient called before finalize");
    }
    // --- Get the gradient of f ---
    mfem::BlockVector x_block(const_cast<mfem::Vector&>(x), m_blockOffsets);
    const mfem::Vector& x_theta = x_block.GetBlock(0);

    auto &grad = m_f->GetGradient(x_theta);
    updatePreconditioner(grad);

    m_jacobian->SetBlock(0, 0, &grad);
    // The other blocks are set up in finalize since they are constant. Only J00 depends on the current state of theta

    return *m_jacobian;
}

void PolytropeOperator::SetEssentialTrueDofs(const SSE::MFEMArrayPair& theta_ess_tdofs, const SSE::MFEMArrayPair& phi_ess_tdofs) {
    m_isFinalized = false;
    m_theta_ess_tdofs = theta_ess_tdofs;
    m_phi_ess_tdofs = phi_ess_tdofs;

    if (m_f) {
        m_f->SetEssentialTrueDofs(theta_ess_tdofs.first);
    } else {
        MFEM_ABORT("m_f is null in PolytropeOperator::SetEssentialTrueDofs");
    }

    if (m_M) {
        m_M->EliminateTestEssentialBC(theta_ess_tdofs.first);
        m_M->EliminateTrialEssentialBC(phi_ess_tdofs.first);
    } else {
        MFEM_ABORT("m_M is null in PolytropeOperator::SetEssentialTrueDofs");
    }

    if (m_Q) {
        m_Q->EliminateTrialEssentialBC(theta_ess_tdofs.first);
        m_Q->EliminateTestEssentialBC(phi_ess_tdofs.first);
    } else {
        MFEM_ABORT("m_Q is null in PolytropeOperator::SetEssentialTrueDofs");
    }

    if (m_D) {
        m_D->EliminateEssentialBC(phi_ess_tdofs.first);
    } else {
        MFEM_ABORT("m_D is null in PolytropeOperator::SetEssentialTrueDofs");
    }
}

void PolytropeOperator::SetEssentialTrueDofs(const SSE::MFEMArrayPairSet& ess_tdof_pair_set) {
    SetEssentialTrueDofs(ess_tdof_pair_set.first, ess_tdof_pair_set.second);
}

SSE::MFEMArrayPairSet PolytropeOperator::GetEssentialTrueDofs() const {
     return std::make_pair(m_theta_ess_tdofs, m_phi_ess_tdofs);
}