MeanField_SERiF/tests/include/serif/tests/equations/volume_equation_fixture.hpp

#pragma once

#include "serif/utils/error/errors.hpp"
#include <optional>
#include <stdexcept>
#include <utility>

#include "serif/backend/mfem/operators/volume_computations.hpp"
#include "serif/backend/mfem/operators/volume_application.hpp"

namespace serif::tests::equations {
    namespace backend = serif::operators::backend::mfem;

    /// Derive the coefficient context used by an equation's unique input terms.
    template <typename ColumnsT, typename SchemaT>
    struct EquationCoefficients;

    template <typename... Operands, typename SchemaT>
    struct EquationCoefficients<utils::types::TypeList<Operands...>, SchemaT> {
        using Type = decltype(backend::make_term_coefficients(
            backend::bind(
                std::declval<const discretization::backend::mfem::TermBinding<typename Operands::term_type, SchemaT>&>(),
                std::declval<const ::mfem::Vector&>()
            )...
        ));
    };

    /// Keep legacy numerical comparisons concise while production uses dependency caches directly.
    template <operators::VolumeEquation EquationT, discretization::domain::schema::IsSchema SchemaT>
    class VolumeEquationFixture {
    public:
        using Plan = operators::EquationPlan<EquationT>;
        using Discretization = backend::VolumeDiscretization<typename Plan::Test, SchemaT>;
        using RowBinding = discretization::backend::mfem::TermBinding<typename Plan::Test, SchemaT>;
        using DisplacementBinding = discretization::backend::mfem::TermBinding<discretization::blocks::displacement::vector, SchemaT>;
        using Coefficients = EquationCoefficients<typename Plan::Columns, SchemaT>::Type;
        using Inputs = decltype(backend::make_volume_inputs<EquationT>(
            std::declval<const Discretization&>(), std::declval<const Coefficients&>(), std::declval<const ::mfem::Vector&>()
        ));
        using Cache = decltype(backend::make_volume_equation_cache(
            std::declval<EquationT>(), std::declval<Discretization>(), std::declval<const Inputs&>(), backend::VolumeOptions{}
        ));

        /// Retain the fixed test configuration until the first numerical input is supplied.
        VolumeEquationFixture(
            EquationT equation,
            const RowBinding& residual_binding,
            const DisplacementBinding& displacement_binding,
            const ::mfem::GridFunction& compactification,
            const mapping::DomainMapper& mapper,
            discretization::quadrature::Policy quadrature_policy = discretization::quadrature::Policy{{}},
            backend::ContractionMode contraction_mode = backend::ContractionMode::automatic
        ) :
            m_equation(std::move(equation)),
            m_discretization{
                .residual_binding = residual_binding,
                .displacement_binding = displacement_binding,
                .compactification = compactification,
                .mapper = mapper
            },
            m_options{
                .quadrature_policy = std::move(quadrature_policy),
                .contraction_mode = contraction_mode
            } {}

        /// Exercise the production validation and graph execution with the old test call shape.
        template <typename CoefficientsT>
        std::expected<void, backend::PreparationRejection> try_prepare(
            const CoefficientsT& coefficients,
            const ::mfem::Vector& displacement,
            backend::Workspace&,
            backend::PreparationMode mode = backend::PreparationMode::rebuild
        ) {
            try {
                if (!m_cache) {
                    const auto inputs = backend::make_volume_inputs<EquationT>(m_discretization, coefficients, displacement);
                    m_cache.emplace(backend::make_volume_equation_cache(std::move(m_equation), m_discretization, inputs, m_options));
                }
            } catch (const std::domain_error& error) {
                return std::unexpected(backend::PreparationRejection{
                    .reason = backend::PreparationRejectionReason::NON_FINITE,
                    .message = error.what()
                });
            }
            return backend::try_prepare_volume(*m_cache, m_discretization, coefficients, displacement, mode);
        }

        /// Turn a numerical rejection into the exception expected by the numerical fixtures.
        template <typename CoefficientsT>
        void prepare(
            const CoefficientsT& coefficients,
            const ::mfem::Vector& displacement,
            backend::Workspace& workspace,
            backend::PreparationMode mode = backend::PreparationMode::rebuild
        ) {
            const auto result = try_prepare(coefficients, displacement, workspace, mode);
            if (!result) {
                serif::utils::errors::serif_error<std::domain_error, serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT>("{}", result.error().message);
            }
        }

        /// Access the real cache for MFEM-adapter tests.
        Cache& cache() {
            if (!m_cache) {
                serif::utils::errors::serif_error<std::logic_error, serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT>("The numerical fixture has not received its inputs.");
            }
            return *m_cache;
        }

        /// Read the same accepted result consumed by production functions.
        const auto& data() const {
            if (!m_cache) {
                serif::utils::errors::serif_error<std::logic_error, serif::utils::errors::SERiFErrorCode::INVALID_ARGUMENT>("The numerical fixture has not received its inputs.");
            }
            return m_cache->template get<backend::volume_keys::Linearization>();
        }

        /// Report readiness directly from the dependency cache.
        bool is_prepared() const {
            return m_cache && m_cache->ready();
        }

        /// Count successful preparations, including explicitly requested rebuilds.
        std::size_t preparation_count() const {
            return m_cache ? m_cache->preparation_count() : 0;
        }

        /// Read derived numerical work without owning a second preparation report.
        auto preparation_report() const {
            return backend::volume_work_report(*m_cache);
        }

        /// Discard cached results through the dependency module.
        void invalidate() {
            if (m_cache) {
                m_cache->invalidate();
            }
        }

        /// Return the supported row size before preparation is needed.
        int height() const {
            return m_discretization.residual_binding.size();
        }

        /// Return the accepted Jacobian column count.
        int width() const {
            return backend::volume_width(data());
        }

        /// Inspect which numerical path the production preparation selected.
        std::size_t tensor_element_count() const {
            return backend::tensor_element_count(data());
        }

        /// Inspect the selected quadrature order for a supported element.
        int integration_order(int element) const {
            return backend::integration_order(data(), element);
        }

        /// Fetch the residual already assembled by the graph.
        void residual(backend::Workspace&, ::mfem::Vector& result) const {
            backend::volume_residual(data(), result);
        }

        /// Delegate every Jacobian call to the production numerical kernel.
        template <typename... Args>
        void apply_jacobian(Args&&... args) const {
            backend::apply_jacobian(data(), std::forward<Args>(args)...);
        }

        /// Delegate transpose application to the production numerical kernel.
        void apply_jacobian_transpose(const ::mfem::Vector& dual, backend::Workspace& workspace, ::mfem::Vector& result) const {
            backend::apply_jacobian_transpose(data(), dual, workspace, result);
        }

    private:
        EquationT m_equation;
        Discretization m_discretization;
        backend::VolumeOptions m_options;
        std::optional<Cache> m_cache;
    };

    template <operators::VolumeEquation EquationT, discretization::domain::schema::IsSchema SchemaT>
    VolumeEquationFixture(
        EquationT,
        const discretization::backend::mfem::TermBinding<typename EquationT::Test, SchemaT>&,
        const discretization::backend::mfem::TermBinding<discretization::blocks::displacement::vector, SchemaT>&,
        const ::mfem::GridFunction&,
        const mapping::DomainMapper&,
        discretization::quadrature::Policy = discretization::quadrature::Policy{{}},
        backend::ContractionMode = backend::ContractionMode::automatic
    ) -> VolumeEquationFixture<EquationT, SchemaT>;
} // namespace serif::tests::equations