module;
#include <mfem.hpp>
module mean_field;
import :solver.fields;

namespace {
    using namespace mean_field;

    constexpr int velocity_block =
        solver::block_index(solver::FieldBlock::velocity);
    constexpr int density_block =
        solver::block_index(solver::FieldBlock::density);
    constexpr int gravity_gradient_block =
        solver::block_index(solver::FieldBlock::gravity_gradient);
    constexpr int displacement_block =
        solver::block_index(solver::FieldBlock::displacement);
} // namespace

namespace mean_field::integrators {
    GravityMomentumIntegrator::GravityMomentumIntegrator(
        const mapping::DomainMapper &map,
        const GravityForceJacobianMode jacobian_mode
    )
        : m_map(map),
          m_jacobian_mode(jacobian_mode) {
    }

    void GravityMomentumIntegrator::SetJacobianMode(
        const GravityForceJacobianMode jacobian_mode
    ) {
        m_jacobian_mode = jacobian_mode;
    }

    void GravityMomentumIntegrator::SetIntegrationRule(
        const mfem::IntegrationRule &integration_rule
    ) {
        m_integration_rule = &integration_rule;
    }

    GravityForceJacobianMode
    GravityMomentumIntegrator::GetJacobianMode() const {
        return m_jacobian_mode;
    }

    void GravityMomentumIntegrator::AssembleElementVector(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array<mfem::Vector *> &elvec
    ) {
        if (utils::is_vacuum(Tr, elvec)) {
            return;
        }

        MFEM_VERIFY(
            m_integration_rule,
            "GravityForceIntegrator must be configured with an "
            "integration rule before assembly."
        );
        MFEM_VERIFY(
            el.Size() > gravity_gradient_block,
            "GravityForceIntegrator requires velocity, density, and "
            "gravity-gradient finite elements."
        );
        MFEM_VERIFY(
            elfun.Size() > gravity_gradient_block,
            "GravityForceIntegrator requires velocity, density, and "
            "gravity-gradient element states."
        );
        MFEM_VERIFY(
            elvec.Size() > velocity_block && elvec[velocity_block],
            "GravityForceIntegrator requires a velocity residual block."
        );
        MFEM_VERIFY(
            el[velocity_block] && el[density_block] &&
                el[gravity_gradient_block],
            "GravityForceIntegrator received a null finite element."
        );
        MFEM_VERIFY(
            elfun[density_block] && elfun[gravity_gradient_block],
            "GravityForceIntegrator received a null element state."
        );

        const mfem::FiniteElement *velocity_element = el[velocity_block];
        const mfem::FiniteElement *density_element  = el[density_block];
        const mfem::FiniteElement *gravity_gradient_element =
            el[gravity_gradient_block];

        const int velocity_dofs_count = velocity_element->GetDof();
        const int density_dofs_count  = density_element->GetDof();
        const int gravity_gradient_dofs_count =
            gravity_gradient_element->GetDof();
        const int dim                    = Tr.GetSpaceDim();

        const mfem::Vector &density_dofs = *elfun[density_block];
        const mfem::Vector &gravity_gradient_dofs =
            *elfun[gravity_gradient_block];

        MFEM_VERIFY(
            density_dofs.Size() == density_dofs_count,
            "GravityForceIntegrator received an incorrectly sized density "
            "state."
        );
        MFEM_VERIFY(
            gravity_gradient_dofs.Size() == gravity_gradient_dofs_count,
            "GravityForceIntegrator received an incorrectly sized "
            "gravity-gradient "
            "state."
        );

        for (int block = 0; block < elvec.Size(); ++block) {
            if (elvec[block]) {
                *elvec[block] = 0.0;
            }
        }

        mfem::Vector &velocity_residual = *elvec[velocity_block];
        velocity_residual.SetSize(dim * velocity_dofs_count);
        velocity_residual = 0.0;

        if (elvec.Size() > density_block && elvec[density_block]) {
            elvec[density_block]->SetSize(density_dofs_count);
            *elvec[density_block] = 0.0;
        }

        if (elvec.Size() > gravity_gradient_block &&
            elvec[gravity_gradient_block]) {
            elvec[gravity_gradient_block]->SetSize(gravity_gradient_dofs_count);
            *elvec[gravity_gradient_block] = 0.0;
        }

        mfem::Vector velocity_shape(velocity_dofs_count);
        mfem::Vector density_shape(density_dofs_count);
        mfem::DenseMatrix gravity_gradient_shape(
            gravity_gradient_dofs_count, dim
        );
        mfem::Vector gravity_gradient_element_value(dim);
        mfem::Vector gravity_gradient_physical_value(dim);

        const mfem::IntegrationRule &integration_rule = *m_integration_rule;

        for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
            const mfem::IntegrationPoint &integration_point =
                integration_rule.IntPoint(q);
            Tr.SetIntPoint(&integration_point);

            const mapping::VolumeQuadratureContext context =
                m_map.GetQuadratureContext(Tr, integration_point);

            velocity_element->CalcShape(integration_point, velocity_shape);
            density_element->CalcShape(integration_point, density_shape);
            gravity_gradient_element->CalcVShape(Tr, gravity_gradient_shape);

            gravity_gradient_shape.MultTranspose(
                gravity_gradient_dofs, gravity_gradient_element_value
            );
            context.J_inv.MultTranspose(
                gravity_gradient_element_value, gravity_gradient_physical_value
            );

            double density_value = 0.0;
            for (int i = 0; i < density_dofs_count; ++i) {
                density_value += density_dofs(i) * density_shape(i);
            }

            for (int i = 0; i < velocity_dofs_count; ++i) {
                for (int component = 0; component < dim; ++component) {
                    velocity_residual(i + component * velocity_dofs_count) +=
                        velocity_shape(i) * density_value *
                        gravity_gradient_physical_value(component) *
                        context.weight;
                }
            }
        }
    }

    void GravityMomentumIntegrator::AssembleElementGrad(
        const mfem::Array<const mfem::FiniteElement *> &el,
        mfem::ElementTransformation &Tr,
        const mfem::Array<const mfem::Vector *> &elfun,
        const mfem::Array2D<mfem::DenseMatrix *> &elmats
    ) {
        if (utils::is_vacuum(Tr, elmats)) {
            return;
        }

        MFEM_VERIFY(
            m_integration_rule,
            "GravityForceIntegrator must be configured with an "
            "integration rule before assembly."
        );
        MFEM_VERIFY(
            el.Size() > gravity_gradient_block,
            "GravityForceIntegrator requires velocity, density, and "
            "gravity-gradient finite elements."
        );
        MFEM_VERIFY(
            elfun.Size() > gravity_gradient_block,
            "GravityForceIntegrator requires velocity, density, and "
            "gravity-gradient element states."
        );
        MFEM_VERIFY(
            el[velocity_block] && el[density_block] &&
                el[gravity_gradient_block],
            "GravityForceIntegrator received a null finite element."
        );
        MFEM_VERIFY(
            elfun[density_block] && elfun[gravity_gradient_block],
            "GravityForceIntegrator received a null element state."
        );

        for (int row = 0; row < elmats.NumRows(); ++row) {
            for (int column = 0; column < elmats.NumCols(); ++column) {
                if (elmats(row, column)) {
                    *elmats(row, column) = 0.0;
                }
            }
        }

        if (m_jacobian_mode == GravityForceJacobianMode::exact) {
            MFEM_ABORT(
                "Exact GravityForceIntegrator geometry Jacobian is unavailable "
                "until "
                "DomainMapper linearization is "
                "implemented."
            );
        }

        const mfem::FiniteElement *velocity_element = el[velocity_block];
        const mfem::FiniteElement *density_element  = el[density_block];
        const mfem::FiniteElement *gravity_gradient_element =
            el[gravity_gradient_block];

        const int velocity_dofs_count = velocity_element->GetDof();
        const int density_dofs_count  = density_element->GetDof();
        const int gravity_gradient_dofs_count =
            gravity_gradient_element->GetDof();
        const int dim                    = Tr.GetSpaceDim();

        const mfem::Vector &density_dofs = *elfun[density_block];
        const mfem::Vector &gravity_gradient_dofs =
            *elfun[gravity_gradient_block];

        MFEM_VERIFY(
            density_dofs.Size() == density_dofs_count,
            "GravityForceIntegrator received an incorrectly sized density "
            "state."
        );
        MFEM_VERIFY(
            gravity_gradient_dofs.Size() == gravity_gradient_dofs_count,
            "GravityForceIntegrator received an incorrectly sized "
            "gravity-gradient "
            "state."
        );

        mfem::DenseMatrix *dv_drho = elmats(velocity_block, density_block);
        mfem::DenseMatrix *dv_dgrad_phi =
            m_jacobian_mode == GravityForceJacobianMode::field_coupled
                ? elmats(velocity_block, gravity_gradient_block)
                : nullptr;

        if (!dv_drho && !dv_dgrad_phi) {
            return;
        }

        mfem::Vector velocity_shape(velocity_dofs_count);
        mfem::Vector density_shape(density_dofs_count);
        mfem::DenseMatrix gravity_gradient_shape(
            gravity_gradient_dofs_count, dim
        );
        mfem::Vector gravity_gradient_element_value(dim);
        mfem::Vector gravity_gradient_physical_value(dim);
        mfem::Vector gravity_basis_element(dim);
        mfem::Vector gravity_basis_physical(dim);

        const mfem::IntegrationRule &integration_rule = *m_integration_rule;

        for (int q = 0; q < integration_rule.GetNPoints(); ++q) {
            const mfem::IntegrationPoint &integration_point =
                integration_rule.IntPoint(q);
            Tr.SetIntPoint(&integration_point);

            const mapping::VolumeQuadratureContext context =
                m_map.GetQuadratureContext(Tr, integration_point);

            velocity_element->CalcShape(integration_point, velocity_shape);
            density_element->CalcShape(integration_point, density_shape);
            gravity_gradient_element->CalcVShape(Tr, gravity_gradient_shape);

            gravity_gradient_shape.MultTranspose(
                gravity_gradient_dofs, gravity_gradient_element_value
            );
            context.J_inv.MultTranspose(
                gravity_gradient_element_value, gravity_gradient_physical_value
            );

            double density_value = 0.0;
            for (int i = 0; i < density_dofs_count; ++i) {
                density_value += density_dofs(i) * density_shape(i);
            }

            if (dv_drho) {
                for (int i = 0; i < velocity_dofs_count; ++i) {
                    for (int component = 0; component < dim; ++component) {
                        const int row = i + component * velocity_dofs_count;

                        for (int j = 0; j < density_dofs_count; ++j) {
                            (*dv_drho)(row, j) +=
                                velocity_shape(i) * density_shape(j) *
                                gravity_gradient_physical_value(component) *
                                context.weight;
                        }
                    }
                }
            }

            if (dv_dgrad_phi) {
                for (int j = 0; j < gravity_gradient_dofs_count; ++j) {
                    for (int component = 0; component < dim; ++component) {
                        gravity_basis_element(component) =
                            gravity_gradient_shape(j, component);
                    }

                    context.J_inv.MultTranspose(
                        gravity_basis_element, gravity_basis_physical
                    );

                    for (int i = 0; i < velocity_dofs_count; ++i) {
                        for (int component = 0; component < dim; ++component) {
                            const int row = i + component * velocity_dofs_count;
                            (*dv_dgrad_phi)(row, j) +=
                                velocity_shape(i) * density_value *
                                gravity_basis_physical(component) *
                                context.weight;
                        }
                    }
                }
            }
        }
    }
} // namespace mean_field::integrators