module;
#include <mfem.hpp>

module mean_field;
import :operators.gravity_field_jacobian;
import :operators.kernels.gravity_field;
import :utils.blocks;

namespace {
    template <int index>
    mfem::Vector make_read_only_value_view(
        const mfem::Vector &vector,
        const mfem::Array<int> &offsets,
        const mean_field::utils::blocks::value_block<index>
    ) {
        const int offset = offsets[index];
        const int size   = offsets[index + 1] - offset;
        return mfem::Vector(const_cast<mfem::real_t *>(vector.GetData()) + offset, size);
    }

    template <int index>
    mfem::Vector make_residual_view(
        mfem::Vector &vector,
        const mfem::Array<int> &offsets,
        const mean_field::utils::blocks::residual_block<index>
    ) {
        const int offset = offsets[index];
        const int size   = offsets[index + 1] - offset;
        return mfem::Vector(vector.GetData() + offset, size);
    }

    template <int index>
    int get_block_size(
        const mfem::Array<int> &offsets,
        const mean_field::utils::blocks::value_block<index>
    ) {
        return offsets[index + 1] - offsets[index];
    }

    template <int index>
    int get_block_size(
        const mfem::Array<int> &offsets,
        const mean_field::utils::blocks::residual_block<index>
    ) {
        return offsets[index + 1] - offsets[index];
    }

    void validate_offsets(
        const mfem::Array<int> &offsets,
        const int block_count,
        const char *message
    ) {
        MFEM_VERIFY(offsets.Size() == block_count + 1, message);
        MFEM_VERIFY(offsets[0] == 0, "Block offsets must begin at zero.");

        for (int i = 0; i < block_count; ++i)
            MFEM_VERIFY(offsets[i + 1] >= offsets[i], "Block offsets must be nondecreasing.");
    }

    void validate_layout(
        const mean_field::fem::FEM &f,
        const mfem::Array<int> &state_offsets,
        const mfem::Array<int> &residual_offsets
    ) {
        using form = mean_field::utils::blocks::gravity_field_form;

        constexpr auto density_block =
            mean_field::utils::blocks::get_value_block<form>(mean_field::utils::blocks::density_field.mass_term);
        constexpr auto displacement_block = mean_field::utils::blocks::get_value_block<form>(
            mean_field::utils::blocks::displacement_field.geometry_term
        );
        constexpr auto gravity_gradient_block =
            mean_field::utils::blocks::get_value_block<form>(mean_field::utils::blocks::gravity_field.gradient_term);
        constexpr auto gravity_potential_block =
            mean_field::utils::blocks::get_value_block<form>(mean_field::utils::blocks::gravity_field.poisson_term);
        constexpr auto gravity_gradient_residual_block =
            mean_field::utils::blocks::get_residual_block<form>(mean_field::utils::blocks::gravity_field.gradient_term);
        constexpr auto gravity_poisson_residual_block =
            mean_field::utils::blocks::get_residual_block<form>(mean_field::utils::blocks::gravity_field.poisson_term);

        validate_offsets(
            state_offsets, form::value_block_count,
            "Gravity Jacobian state offsets do not match the gravity field "
            "form."
        );
        validate_offsets(
            residual_offsets, form::residual_block_count,
            "Gravity Jacobian residual offsets do not match the gravity field "
            "form."
        );

        MFEM_VERIFY(
            get_block_size(state_offsets, density_block) == f.densityFes->GetTrueVSize(),
            "The Jacobian density block has the wrong size."
        );
        MFEM_VERIFY(
            get_block_size(state_offsets, displacement_block) == f.displacementFes->GetTrueVSize(),
            "The Jacobian displacement block has the wrong size."
        );
        MFEM_VERIFY(
            get_block_size(state_offsets, gravity_gradient_block) == f.gravityFluxFes->GetTrueVSize(),
            "The Jacobian gravity-gradient block has the wrong size."
        );
        MFEM_VERIFY(
            get_block_size(state_offsets, gravity_potential_block) == f.gravityPotentialFes->GetTrueVSize(),
            "The Jacobian gravity-potential block has the wrong size."
        );
        MFEM_VERIFY(
            get_block_size(residual_offsets, gravity_gradient_residual_block) == f.gravityFluxFes->GetTrueVSize(),
            "The Jacobian gradient-residual block has the wrong size."
        );
        MFEM_VERIFY(
            get_block_size(residual_offsets, gravity_poisson_residual_block) == f.gravityPotentialFes->GetTrueVSize(),
            "The Jacobian Poisson-residual block has the wrong size."
        );
    }
} // namespace

namespace mean_field::operators {
    GravityFieldJacobianOperator::GravityFieldJacobianOperator(
        fem::FEM &f,
        const mapping::DomainMapperStateless &domain_mapper,
        const context::gravity_field::GravityFieldLinearizationContext &linearization_context,
        const mfem::Array<int> &state_true_offsets,
        const mfem::Array<int> &residual_true_offsets
    )
        : Operator(
              residual_true_offsets.Last(),
              state_true_offsets.Last()
          ),
          m_fem(f),
          m_domain_mapper(domain_mapper),
          m_linearization_context(linearization_context),
          m_state_true_offsets(state_true_offsets),
          m_residual_true_offsets(residual_true_offsets) {
        MFEM_VERIFY(
            f.densityFes != nullptr, "GravityFieldJacobianOperator requires the density finite-element "
                                     "space."
        );
        MFEM_VERIFY(
            f.gravityPotentialFes != nullptr, "GravityFieldJacobianOperator requires the gravity-potential "
                                              "finite-element space."
        );
        MFEM_VERIFY(
            f.gravityFluxFes != nullptr, "GravityFieldJacobianOperator requires the "
                                         "gravity-gradient finite-element space."
        );
        MFEM_VERIFY(
            f.displacementFes != nullptr, "GravityFieldJacobianOperator requires the "
                                          "displacement finite-element space."
        );
        MFEM_VERIFY(
            f.gravityContext.b_form != nullptr, "GravityFieldJacobianOperator requires the divergence operator."
        );
        MFEM_VERIFY(
            f.gravityContext.BT != nullptr, "GravityFieldJacobianOperator requires the transpose divergence "
                                            "operator."
        );
        MFEM_VERIFY(
            f.quadratureFactory != nullptr, "GravityFieldJacobianOperator requires the quadrature-rule factory."
        );
        MFEM_VERIFY(
            domain_mapper.GetDimension() == f.mesh->Dimension(),
            "GravityFieldJacobianOperator received a domain mapper with the "
            "wrong "
            "dimension."
        );

        validate_layout(f, m_state_true_offsets, m_residual_true_offsets);
    }

    void GravityFieldJacobianOperator::Mult(
        const mfem::Vector &direction,
        mfem::Vector &action
    ) const {
        MFEM_VERIFY(
            m_linearization_context.IsPrepared(), "GravityFieldJacobianOperator requires a prepared linearization "
                                                  "context."
        );
        MFEM_VERIFY(
            direction.Size() == Width(), "GravityFieldJacobianOperator received a direction with the wrong "
                                         "size."
        );

        using form                   = utils::blocks::gravity_field_form;

        constexpr auto density_block = utils::blocks::get_value_block<form>(utils::blocks::density_field.mass_term);
        constexpr auto displacement_block =
            utils::blocks::get_value_block<form>(utils::blocks::displacement_field.geometry_term);
        constexpr auto gravity_gradient_block =
            utils::blocks::get_value_block<form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravity_potential_block =
            utils::blocks::get_value_block<form>(utils::blocks::gravity_field.poisson_term);
        constexpr auto gravity_gradient_residual_block =
            utils::blocks::get_residual_block<form>(utils::blocks::gravity_field.gradient_term);
        constexpr auto gravity_poisson_residual_block =
            utils::blocks::get_residual_block<form>(utils::blocks::gravity_field.poisson_term);

        const context::gravity_field::GravityFieldGeometryContext &geometry_context =
            m_linearization_context.GetGeometryContext();
        const mfem::Vector &density          = m_linearization_context.GetDensity();
        const mfem::Vector &displacement     = geometry_context.GetDisplacement();
        const mfem::Vector &gravity_gradient = m_linearization_context.GetGravityGradient();

        const mfem::Vector density_direction =
            make_read_only_value_view(direction, m_state_true_offsets, density_block);
        const mfem::Vector displacement_direction =
            make_read_only_value_view(direction, m_state_true_offsets, displacement_block);
        const mfem::Vector gravity_gradient_direction =
            make_read_only_value_view(direction, m_state_true_offsets, gravity_gradient_block);
        const mfem::Vector gravity_potential_direction =
            make_read_only_value_view(direction, m_state_true_offsets, gravity_potential_block);

        action.SetSize(Height());
        action = 0.0;

        mfem::Vector gravity_gradient_action =
            make_residual_view(action, m_residual_true_offsets, gravity_gradient_residual_block);
        mfem::Vector gravity_poisson_action =
            make_residual_view(action, m_residual_true_offsets, gravity_poisson_residual_block);

        mfem::Vector transpose_divergence_action;
        mfem::Vector source_action;
        mfem::Vector mass_variation_action;
        mfem::Vector source_variation_action;

        geometry_context.GetMassOperator().Mult(gravity_gradient_direction, gravity_gradient_action);
        geometry_context.GetSourceOperator().Mult(density_direction, source_action);

        kernels::apply_mapped_hdiv_mass_variation(
            m_fem, m_domain_mapper, gravity_gradient, displacement, displacement_direction, mass_variation_action
        );

        kernels::apply_mapped_source_variation(
            m_fem, m_domain_mapper, density, displacement, displacement_direction, source_variation_action
        );

        transpose_divergence_action.SetSize(gravity_gradient_action.Size());
        m_fem.gravityContext.BT->Mult(gravity_potential_direction, transpose_divergence_action);

        gravity_gradient_action += transpose_divergence_action;
        gravity_gradient_action += mass_variation_action;

        m_fem.gravityContext.b_form->Mult(gravity_gradient_direction, gravity_poisson_action);

        gravity_poisson_action -= source_action;
        gravity_poisson_action -= source_variation_action;
    }

    const context::gravity_field::GravityFieldLinearizationContext &
    GravityFieldJacobianOperator::GetLinearizationContext() const noexcept {
        return m_linearization_context;
    }
} // namespace mean_field::operators