module;
#include "xad_promote_polyfill.h"
#include <XAD/XAD.hpp>
#include <mfem.hpp>
export module mean_field:integrators.pressure_gradient;
import :mapping.domain_mapper;
import :utils.misc;

export namespace mean_field::integrators {
    template <utils::is_xad EOS_T> class PressureGradientIntegrator : public mfem::BlockNonlinearFormIntegrator {
    public:
        PressureGradientIntegrator(
            const mapping::DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate,
            utils::EOS_P<EOS_T> eos
        );

        void AssembleElementVector(
            const mfem::Array<const mfem::FiniteElement *> &el,
            mfem::ElementTransformation &Tr,
            const mfem::Array<const mfem::Vector *> &elfun,
            const mfem::Array<mfem::Vector *> &elvec
        ) override;
        void AssembleElementGrad(
            const mfem::Array<const mfem::FiniteElement *> &el,
            mfem::ElementTransformation &Tr,
            const mfem::Array<const mfem::Vector *> &elfun,
            const mfem::Array2D<mfem::DenseMatrix *> &elmats
        ) override;

    private:
        mapping::GridFunctionMappingEvaluator m_mapping;
        utils::EOS_P<EOS_T> m_eos;
    };

    template <utils::is_xad EOS_T>
    PressureGradientIntegrator<EOS_T>::PressureGradientIntegrator(
        const mapping::DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        utils::EOS_P<EOS_T> eos
    )
        : m_mapping(mapper, displacement, compactification_coordinate),
          m_eos(std::move(eos)) {
    }

    template <utils::is_xad EOS_T>
    void PressureGradientIntegrator<EOS_T>::AssembleElementVector(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array<mfem::Vector *> &elvec
    ) {
        m_mapping.InvalidateCache();

        if (utils::is_vacuum(Tr, elvec)) {
            return;
        }

        const mfem::FiniteElement *fe_v   = el[0];
        const mfem::FiniteElement *fe_rho = el[1];

        const int dof_v                   = fe_v->GetDof();
        const int dof_rho                 = fe_rho->GetDof();
        const int dim                     = Tr.GetSpaceDim();

        const mfem::Vector &rho_dofs      = *elfun[1];

        mfem::Vector &r_v                 = *elvec[0];
        r_v.SetSize(dof_v * dim);
        r_v = 0.0;
        if (elvec[1]) {
            elvec[1]->SetSize(dof_rho);
            *elvec[1] = 0.0;
        }

        mfem::DenseMatrix dshape_v_ref(dof_v, dim), dshape_v_phys(dof_v, dim);
        mfem::Vector shape_rho(dof_rho);

        const mfem::IntegrationRule *ir = &mfem::IntRules.Get(fe_v->GetGeomType(), 2 * fe_v->GetOrder());

        for (int q = 0; q < ir->GetNPoints(); ++q) {
            const mfem::IntegrationPoint &ip = ir->IntPoint(q);
            Tr.SetIntPoint(&ip);

            auto [J_inv, detJ, weight] = m_mapping.GetQuadratureContext(Tr, ip);

            fe_v->CalcDShape(ip, dshape_v_ref);
            mfem::Mult(dshape_v_ref, J_inv, dshape_v_phys);
            fe_rho->CalcShape(ip, shape_rho);

            double rho_val = 0.0;
            for (int i = 0; i < dof_rho; ++i)
                rho_val += rho_dofs(i) * shape_rho(i);

            // Guard against negative density from Newton solver overshoots
            if (rho_val < 1e-15)
                rho_val = 1e-15;

            // Evaluate the exact Equation of State Pressure
            EOS_T x_rho  = rho_val;
            double P_val = m_eos(x_rho, EOS_T(0.0)).value();

            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    r_v(i + c * dof_v) -= dshape_v_phys(i, c) * P_val * weight;
                }
            }
        }
    }

    template <utils::is_xad EOS_T>
    void PressureGradientIntegrator<EOS_T>::AssembleElementGrad(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array2D<mfem::DenseMatrix *> &elmats
    ) {
        m_mapping.InvalidateCache();

        const mfem::FiniteElement *fe_v   = el[0];
        const mfem::FiniteElement *fe_rho = el[1];

        const int dof_v                   = fe_v->GetDof();
        const int dof_rho                 = fe_rho->GetDof();
        const int dim                     = Tr.GetSpaceDim();

        const mfem::Vector &rho_dofs      = *elfun[1];

        mfem::DenseMatrix *dv_dv          = elmats(0, 0);
        mfem::DenseMatrix *dv_drho        = elmats(0, 1);

        if (dv_dv)
            *dv_dv = 0.0;
        if (dv_drho)
            *dv_drho = 0.0;
        if (!dv_drho)
            return;

        mfem::DenseMatrix dshape_v_ref(dof_v, dim), dshape_v_phys(dof_v, dim);
        mfem::Vector shape_rho(dof_rho);

        const mfem::IntegrationRule *ir = &mfem::IntRules.Get(fe_v->GetGeomType(), 2 * fe_v->GetOrder());

        for (int q = 0; q < ir->GetNPoints(); ++q) {
            using Scalar = EOS_T::value_type;
            xad::Tape<Scalar> tape;
            const mfem::IntegrationPoint &ip = ir->IntPoint(q);
            Tr.SetIntPoint(&ip);

            auto [J_inv, detJ, weight] = m_mapping.GetQuadratureContext(Tr, ip);

            fe_v->CalcDShape(ip, dshape_v_ref);
            mfem::Mult(dshape_v_ref, J_inv, dshape_v_phys);
            fe_rho->CalcShape(ip, shape_rho);

            EOS_T x_rho(0.0);
            tape.registerInput(x_rho);
            tape.newRecording();

            for (int i = 0; i < dof_rho; ++i) {
                x_rho += rho_dofs(i) * shape_rho(i);
            }
            if (x_rho < 1e-15)
                x_rho = EOS_T(1e-15);
            EOS_T x_P = m_eos(x_rho, EOS_T(0.0));
            tape.registerOutput(x_P);
            x_P.setAdjoint(1.0);
            tape.computeAdjoints();
            double dP_drho     = x_rho.getAdjoint();

            double debug_K     = 1.5;
            double debug_n     = 3.0;
            double analytic_dp = debug_K * (1.0 + 1.0 / debug_n) * std::pow(xad::value(x_rho), 1.0 / debug_n);

            double ad_err      = std::abs(dP_drho - analytic_dp);

            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    int row = i + c * dof_v;
                    for (int j = 0; j < dof_rho; ++j) {
                        int col     = j;
                        double term = dshape_v_phys(i, c) * dP_drho * shape_rho(j);
                        (*dv_drho)(row, col) -= term * weight;
                    }
                }
            }
        }
    }
} // namespace mean_field::integrators