module;
#include <mfem.hpp>
module mean_field;

namespace mean_field::integrators {
    CentrifugalForceIntegrator::CentrifugalForceIntegrator(
        const mapping::DomainMapper &mapper,
        const mfem::GridFunction &displacement,
        const mfem::GridFunction &compactification_coordinate,
        const mfem::Vector &omega
    )
        : m_mapping(
              mapper,
              displacement,
              compactification_coordinate
          ),
          m_omega(3) {
        MFEM_ASSERT(omega.Size() == 3, "Omega vector must be 3D");
        m_omega = omega;
    }

    void CentrifugalForceIntegrator::SetOmega(const mfem::Vector &omega) {
        MFEM_ASSERT(omega.Size() == 3, "Omega vector must be 3D");
        m_omega = omega;
    }

    void CentrifugalForceIntegrator::SetIntegrationRule(const mfem::IntegrationRule &ir) {
        m_ir = &ir;
    }

    void CentrifugalForceIntegrator::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::Vector shape_v(dof_v), shape_rho(dof_rho);
        mfem::Vector a(dim), b(dim);
        mapping::VolumeMappingContext mapping_context;

        MFEM_VERIFY(
            m_ir, "CentrifugalForceIntegrator must be configured with an "
                  "integration rule before assembly. Call "
                  "SetIntegrationRule first."
        );
        const mfem::IntegrationRule *ir = m_ir;

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

            const mapping::MappingStatus mapping_status = m_mapping.EvaluateVolume(Tr, ip, mapping_context);
            MFEM_VERIFY(
                mapping_status == mapping::MappingStatus::valid,
                "Centrifugal-force assembly encountered an invalid volume mapping."
            );
            const double weight = mapping_context.quadrature.weight;

            fe_v->CalcShape(ip, shape_v);
            fe_rho->CalcShape(ip, shape_rho);

            const mfem::Vector &x_phys = mapping_context.mapping.physical_position;

            // ω x r
            a(0)                       = m_omega(1) * x_phys(2) - m_omega(2) * x_phys(1);
            a(1)                       = m_omega(2) * x_phys(0) - m_omega(0) * x_phys(2);
            a(2)                       = m_omega(0) * x_phys(1) - m_omega(1) * x_phys(0);

            // ω x (ω x r) [centrifugal acceleration]
            b(0)                       = m_omega(1) * a(2) - m_omega(2) * a(1);
            b(1)                       = m_omega(2) * a(0) - m_omega(0) * a(2);
            b(2)                       = m_omega(0) * a(1) - m_omega(1) * a(0);

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

            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    r_v(i + c * dof_v) += shape_v(i) * rho_val * b(c) * weight;
                }
            }
        }
    }

    void CentrifugalForceIntegrator::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();

        if (utils::is_vacuum(Tr, elmats)) {
            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();

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

        if (dv_dv)
            *dv_dv = 0.0;
        if (elmats(1, 0))
            *elmats(1, 0) = 0.0;
        if (elmats(1, 1))
            *elmats(1, 1) = 0.0;
        if (dv_drho)
            *dv_drho = 0.0;
        if (!dv_drho)
            return;

        mfem::Vector shape_v(dof_v), shape_rho(dof_rho);
        mfem::Vector a(dim), b(dim);
        mapping::VolumeMappingContext mapping_context;

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

            const mapping::MappingStatus mapping_status = m_mapping.EvaluateVolume(Tr, ip, mapping_context);
            MFEM_VERIFY(
                mapping_status == mapping::MappingStatus::valid,
                "Centrifugal-force Jacobian assembly encountered an invalid volume mapping."
            );
            const double weight = mapping_context.quadrature.weight;

            fe_v->CalcShape(ip, shape_v);
            fe_rho->CalcShape(ip, shape_rho);

            const mfem::Vector &x_phys = mapping_context.mapping.physical_position;

            // ω x r
            a(0)                       = m_omega(1) * x_phys(2) - m_omega(2) * x_phys(1);
            a(1)                       = m_omega(2) * x_phys(0) - m_omega(0) * x_phys(2);
            a(2)                       = m_omega(0) * x_phys(1) - m_omega(1) * x_phys(0);

            // ω x (ω x r) [centrifugal acceleration]
            b(0)                       = m_omega(1) * a(2) - m_omega(2) * a(1);
            b(1)                       = m_omega(2) * a(0) - m_omega(0) * a(2);
            b(2)                       = m_omega(0) * a(1) - m_omega(1) * a(0);

            // dR_dv_i_c / drho_j = φ_i * φ_j * b_c
            for (int i = 0; i < dof_v; ++i) {
                for (int c = 0; c < dim; ++c) {
                    const int row = i + c * dof_v;
                    for (int j = 0; j < dof_rho; ++j) {
                        (*dv_drho)(row, j) += shape_v(i) * shape_rho(j) * b(c) * weight;
                    }
                }
            }
        }
    }
} // namespace mean_field::integrators