#include "operator.h"
#include "4DSTARTypes.h"
#include "linalg/densemat.hpp"
#include "linalg/sparsemat.hpp"
#include "mfem.hpp"
#include "linalg/vector.hpp"
#include <memory>

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,
    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);
}

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() {
    if (m_isFinalized) {
        return;
    }
    m_Mmat = std::make_unique<mfem::SparseMatrix>(m_M->SpMat());
    m_Qmat = std::make_unique<mfem::SparseMatrix>(m_Q->SpMat());
    m_Dmat = std::make_unique<mfem::SparseMatrix>(m_D->SpMat());

    for (const int thetaDof : m_theta_ess_tdofs.first) {
        m_Mmat->EliminateRow(thetaDof);
        m_Qmat->EliminateCol(thetaDof);
    }
    // These are commented out because they theoretically are wrong (need to think more about how to apply essential dofs to a vector div field)
    // for (const int phiDof : m_phi_ess_tdofs.first) {
    //     if (phiDof >=0 && phiDof < m_Dmat->Height()) {
    //         m_Dmat->EliminateRowCol(phiDof);
    //         m_Qmat->EliminateRow(phiDof);
    //         m_Mmat->EliminateCol(phiDof);
    //     }
    // }

    m_negM_op = std::make_unique<mfem::ScaledOperator>(m_Mmat.get(), -1.0);
    m_negQ_op = std::make_unique<mfem::ScaledOperator>(m_Qmat.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_Dmat.get()); // D (constant)

    m_isFinalized = true;
}

const mfem::BlockOperator &PolytropeOperator::GetJacobianOperator() const {
     if (m_jacobian == nullptr) {
         MFEM_ABORT("Jacobian has not been initialized before GetJacobianOperator() call.");
     }
     return *m_jacobian;
}

mfem::Vector PolytropeOperator::GetJ00Diag() const {
}

mfem::Vector PolytropeOperator::GetJ01Diag() const {
     if (!m_isFinalized) {
         MFEM_ABORT("PolytropeOperator::Get01Diag -> GetJ01Diag called before finalization of PolytropeOperator");
     }
     if (m_Mmat == nullptr) {
         MFEM_ABORT("PolytropeOperator::Get01Diag -> M sparse matrix has not been initialized before GetJ01Diag() call.");
     }

     mfem::Vector J01diag;
     m_Mmat->GetDiag(J01diag);
     J01diag *= -1;
     return J01diag;
}

mfem::Vector PolytropeOperator::GetJ10diag() const {
     if (!m_isFinalized) {
         MFEM_ABORT("PolytropeOperator::Get10Diag -> GetJ10Diag called before finalization of PolytropeOperator");
     }
     if (m_Qmat == nullptr) {
         MFEM_ABORT("PolytropeOperator::Get10Diag -> Q sparse matrix has not been initialized before GetJ10Diag() call.");
     }
     mfem::Vector J10diag;
     m_Qmat->GetDiag(J10diag);
     J10diag *= -1;
     return J10diag;
}

mfem::Vector PolytropeOperator::GetJ11diag() const {
     if (!m_isFinalized) {
         MFEM_ABORT("PolytropeOperator::Get11Diag -> GetJ11Diag called before finalization of PolytropeOperator");
     }
     if (m_Dmat == nullptr) {
         MFEM_ABORT("PolytropeOperator::Get11Diag -> D sparse matrix has not been initialized before GetJ11Diag() call.");
     }
     mfem::Vector J11diag;
     m_Dmat->GetDiag(J11diag);
     J11diag *= -1;
     return J11diag;
}


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");
    MFEM_ASSERT(m_Mmat.get() != nullptr, "SparseMatrix m_Mmat is null in PolytropeOperator::Mult");
    MFEM_ASSERT(m_Dmat.get() != nullptr, "SparseMatrix m_Dmat is null in PolytropeOperator::Mult");
    MFEM_ASSERT(m_Qmat.get() != nullptr, "SparseMatrix m_Qmat is null in PolytropeOperator::Mult");

    m_f->Mult(x_theta, f_term);
    m_Mmat->Mult(x_phi, Mphi_term);
    m_Dmat->Mult(x_phi, Dphi_term);
    m_Qmat->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++) {
        int idx = m_theta_ess_tdofs.first[i];
        if (idx >= 0 && idx < y_R0.Size()) {
            const double &targetValue = m_theta_ess_tdofs.second[i];
            y_block.GetBlock(0)[idx] = targetValue - x_theta(idx); // Zero out the essential theta dofs in the bi-linear form
            // y_block.GetBlock(0)[idx] = 0; // Zero out the essential theta dofs in the bi-linear form
        }
    }

    // TODO look into how the true dof -> vector component works
    // for (int i = 0; i < m_phi_ess_tdofs.first.Size(); i++) {
    //     int idx = m_phi_ess_tdofs.first[i];
    //     if (idx >= 0 && idx < y_R1.Size()) {
    //         // const double &targetValue = m_phi_ess_tdofs.second[i];
    //         // y_block.GetBlock(1)[idx] = targetValue - x_phi(idx); // Zero out the essential phi dofs in the bi-linear form
    //         y_block.GetBlock(1)[idx] = 0; // Zero out the essential phi dofs in the bi-linear form
    //     }
    // }

}

void PolytropeOperator::updateInverseNonlinearJacobian(const mfem::Operator &grad) const {
    if (const auto *sparse_mat = dynamic_cast<const mfem::SparseMatrix*>(&grad); sparse_mat != nullptr) {
        mfem::Vector gradDiag;
        sparse_mat->GetDiag(gradDiag);
        for (int i = 0; i < gradDiag.Size(); i++) {
            gradDiag(i) = 1.0/gradDiag(i); // Invert the diagonals of the jacobian
        }
        m_invNonlinearJacobian = std::make_unique<mfem::SparseMatrix>(gradDiag);
    } else {
        MFEM_ABORT("PolytropeOperator::GetGradient called on nonlinear Jacobian");
    }
}

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_Dmat.get()); // This is now a copy of D

    // Calculate S = D - Q df^{-1} M
    mfem::SparseMatrix* temp_QxdfInv = mfem::Mult(*m_Qmat, *m_invNonlinearJacobian); // Q * df^{-1}
    mfem::SparseMatrix* temp_QxdfInvxM = mfem::Mult(*temp_QxdfInv, *m_Mmat); // Q * df^{-1} * M

    schurCompliment->Add(-1, *temp_QxdfInvxM); // D - Q * df^{-1} * M

    writeOperatorToCSV(*schurCompliment, "/Users/tboudreaux/Desktop/SchursCompliment.csv");
    std::cout << "Wrote" << std::endl;

    // Note: EMB (April 9, 2025) Where I left for the day, so I don't pull my hair out tomorrow:
    /*
     * Need to calculate the inverse of Schur's compliment to precondition the jacobian
     * I think I have this close; however, when running there is a seg fault when converting to a DenseMatrix for
     *   writing out to CSV. This bodes poorly for the underlying memory structure.
     *
     * Also, there is a lot of math happening in these tight loops, probably need to think about that and
     *   make it more efficient.
     */

    // TODO Also I need to actually invert the schur compliment since so far I have just found its non inverted state
    // Should probably add an energy check to make sure it is roughly diagonal

    // I did this manually in python for the non linear jacobian and the energy along the diagonal was ~99.999 % of the
    // total energy indicating that it is very strongly diagonal.
}

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);
    updateInverseNonlinearJacobian(grad);
    updateInverseSchurCompliment();

    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);

        // This should be zeroing out the row; however, I am getting a segfault
    } else {
        MFEM_ABORT("m_f 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);
}
