#include "operator.h"
#include "mfem.hpp"
#include "linalg/vector.hpp"
#include <memory>

#include "debug.h"

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

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

    m_isFinalized = true;
}


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

    // Caucluate 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_tofs.Size(); i++) {
        int idx = m_theta_ess_tofs[i];
        if (idx >= 0 && idx < y_R0.Size()) {
            y_block.GetBlock(0)[idx] = 0.0; // Zero out the essential theta dofs in the bilinear form
        }
    }

    for (int i = 0; i < m_phi_ess_tofs.Size(); i++) {
        int idx = m_phi_ess_tofs[i];
        if (idx >= 0 && idx < y_R1.Size()) {
            y_block.GetBlock(1)[idx] = 0.0; // Zero out the essential phi dofs in the bilinear form
        }
    }

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

    mfem::Operator& J00 = m_f->GetGradient(x_theta);
    if (m_jacobian == nullptr) {
        m_jacobian = std::make_unique<mfem::BlockOperator>(m_blockOffsets);
        m_jacobian->SetBlock(0, 0, &J00); // df/dθ (state-dependent)
        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)
    } else {
        // The Jacobian already exists, we only need to update the first block
        // since the other blocks have a constant derivitive (they are linear)
        m_jacobian->SetBlock(0, 0, &J00);
    }

    return *m_jacobian;
}

void PolytropeOperator::SetEssentialTrueDofs(const mfem::Array<int> &theta_ess_tofs,
                                             const mfem::Array<int> &phi_ess_tofs) {
    m_isFinalized = false;
    m_theta_ess_tofs = theta_ess_tofs;
    m_phi_ess_tofs = phi_ess_tofs;

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