#pragma once

#include "mfem/linalg/densemat.hpp"
#include "mfem/linalg/vector.hpp"
#include "serif/backend/mfem/mapping/domain_mapper.hpp"

namespace serif::tests {
    /**
     * @brief Differentiate the volume map along one direction through the solver's direction-variation path.
     * @note Samples the direction's value and canonical-element gradient at the point and forms the inverse mesh
     * Jacobian from the prepared context, as the geometry derivatives do.
     */
    [[nodiscard]] inline serif::mapping::MappingStatus evaluate_volume_variation(
        const serif::mapping::DomainMapper& mapper,
        const serif::mapping::ElementMappingData& element_data,
        const serif::mapping::ElementDisplacementData& direction,
        ::mfem::ElementTransformation& transformation,
        const ::mfem::IntegrationPoint& ip,
        const serif::mapping::VolumeMappingContext& context,
        serif::mapping::DMWorkspace& workspace,
        serif::mapping::VolumeMappingVariation& variation,
        const serif::mapping::CompactificationPointData* compactification_sample = nullptr
    ) {
        const ::mfem::FiniteElement& element = direction.get_element();
        ::mfem::Vector shape(element.GetDof());
        ::mfem::DenseMatrix shape_gradient(element.GetDof(), element.GetDim());
        element.CalcShape(ip, shape);
        element.CalcDShape(ip, shape_gradient);
        ::mfem::Vector value(direction.get_dimension());
        direction.get_dof_matrix().MultTranspose(shape, value);
        ::mfem::DenseMatrix gradient(direction.get_dimension(), element.GetDim());
        ::mfem::MultAtB(direction.get_dof_matrix(), shape_gradient, gradient);
        ::mfem::DenseMatrix inverse_mesh_jacobian(context.mapping.mapping_jacobian.Height());
        ::mfem::Mult(context.quadrature.J_element_to_physical_inverse, context.mapping.mapping_jacobian, inverse_mesh_jacobian);
        return mapper.evaluate_volume_direction_variation(element_data, direction, transformation, ip, context, workspace, variation,
            serif::mapping::DMReferenceDisplacementSample{value, gradient}, compactification_sample, inverse_mesh_jacobian);
    }
}