module;

#include <cmath>
#include <concepts>
#include <memory>
#include <type_traits>
#include <utility>

#include <mfem.hpp>

export module mean_field:operators.prepared_surface_constraint;

export import :field.mfem;
export import :surface.compiled;

namespace mean_field::operators::detail {
    template <dimensions::ThermodynamicQuantityType Quantity> struct SingleQuantitySurfaceState final {
        dimensions::QuantityValue<Quantity> quantityValue;

        [[nodiscard]] dimensions::QuantityValue<Quantity> value(Quantity) const noexcept {
            return quantityValue;
        }
    };
} // namespace mean_field::operators::detail

export namespace mean_field::operators {
    /*
     * Runtime enforcement currently supports a pointwise pressure constraint
     * whose row field is also its sole state field. The concept is expressed
     * entirely in compiled-constraint metadata: no thermodynamic carrier or
     * concrete field is selected by this prepared layer.
     */
    template <typename Candidate>
    concept SingleFieldPressureSurfaceConstraint =
        requires {
            typename std::remove_cvref_t<Candidate>::PhysicalQuantity;
            typename std::remove_cvref_t<Candidate>::CarrierQuantity;
            typename std::remove_cvref_t<Candidate>::CarrierField;
            typename std::remove_cvref_t<Candidate>::SurfaceDependencies;
        } && std::same_as<typename std::remove_cvref_t<Candidate>::PhysicalQuantity, dimensions::quantity::Pressure> &&
        std::same_as<
            typename std::remove_cvref_t<Candidate>::SurfaceDependencies::RowField,
            typename std::remove_cvref_t<Candidate>::CarrierField> &&
        std::same_as<
            typename std::remove_cvref_t<Candidate>::SurfaceDependencies::StateFieldTypes,
            field::TypeList<typename std::remove_cvref_t<Candidate>::CarrierField>>;

    template <typename Candidate, typename Field>
    concept SingleFieldPressureSurfaceConstraintFor =
        SingleFieldPressureSurfaceConstraint<Candidate> &&
        std::same_as<typename std::remove_cvref_t<Candidate>::SurfaceDependencies::RowField, Field>;

    /*
     * Non-owning runtime bridge for a statically compiled pressure constraint.
     * There is one function-pointer dispatch per complete row application;
     * the concrete loop remains templated so EOS operations can be inlined.
     */
    class PressureSurfaceConstraintView final {
    public:
        template <SingleFieldPressureSurfaceConstraint Constraint>
        explicit PressureSurfaceConstraintView(const Constraint &constraint) noexcept
            : m_constraint(std::addressof(constraint)),
              m_applyResidualRows(&applyResidualRows<Constraint>),
              m_applyJacobianRows(&applyJacobianRows<Constraint>),
              m_descriptor(constraint.descriptor()) {
        }

        void ApplyResidualRows(
            const mfem::Vector &surfaceState,
            const field::FieldBoundaryDofMap &surfaceRows,
            mfem::Vector &rowResidual
        ) const {
            m_applyResidualRows(m_constraint, surfaceState, surfaceRows, rowResidual);
        }

        void ApplyJacobianRows(
            const mfem::Vector &surfaceState,
            const field::FieldBoundaryDofMap &surfaceRows,
            const mfem::Vector &stateVariation,
            mfem::Vector &rowAction
        ) const {
            m_applyJacobianRows(m_constraint, surfaceState, surfaceRows, stateVariation, rowAction);
        }

        [[nodiscard]] surface::PressureSurfaceDescriptor descriptor() const noexcept {
            return m_descriptor;
        }

    private:
        using ApplyResidualRowsFunction = void (*)(
            const void *,
            const mfem::Vector &,
            const field::FieldBoundaryDofMap &,
            mfem::Vector &
        );
        using ApplyJacobianRowsFunction = void (*)(
            const void *,
            const mfem::Vector &,
            const field::FieldBoundaryDofMap &,
            const mfem::Vector &,
            mfem::Vector &
        );

        template <SingleFieldPressureSurfaceConstraint Constraint>
        static void applyResidualRows(
            const void *constraint,
            const mfem::Vector &surfaceState,
            const field::FieldBoundaryDofMap &surfaceRows,
            mfem::Vector &rowResidual
        ) {
            using CarrierQuantity = typename Constraint::CarrierQuantity;

            for (int surfaceIndex = 0; surfaceIndex < surfaceRows.size(); ++surfaceIndex) {
                const detail::SingleQuantitySurfaceState<CarrierQuantity> state{
                    dimensions::QuantityValue<CarrierQuantity>{surfaceState(surfaceIndex)}
                };
                rowResidual(surfaceRows.reduced_dofs()[surfaceIndex]) =
                    static_cast<const Constraint *>(constraint)->residual(state);
            }
        }

        template <SingleFieldPressureSurfaceConstraint Constraint>
        static void applyJacobianRows(
            const void *constraint,
            const mfem::Vector &surfaceState,
            const field::FieldBoundaryDofMap &surfaceRows,
            const mfem::Vector &stateVariation,
            mfem::Vector &rowAction
        ) {
            using CarrierQuantity = typename Constraint::CarrierQuantity;

            for (int surfaceIndex = 0; surfaceIndex < surfaceRows.size(); ++surfaceIndex) {
                const int reducedDof = surfaceRows.reduced_dofs()[surfaceIndex];
                const detail::SingleQuantitySurfaceState<CarrierQuantity> state{
                    dimensions::QuantityValue<CarrierQuantity>{surfaceState(surfaceIndex)}
                };
                const detail::SingleQuantitySurfaceState<CarrierQuantity> variation{
                    dimensions::QuantityValue<CarrierQuantity>{stateVariation(reducedDof)}
                };
                rowAction(reducedDof) = static_cast<const Constraint *>(constraint)->jacobianAction(state, variation);
            }
        }

        const void *m_constraint;
        ApplyResidualRowsFunction m_applyResidualRows;
        ApplyJacobianRowsFunction m_applyJacobianRows;
        surface::PressureSurfaceDescriptor m_descriptor;
    };

    struct PreparedSurfaceConstraintReport final {
        bool cachedSurfaceState{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return cachedSurfaceState;
        }
    };

    class PreparedPressureSurfaceConstraint final {
    public:
        PreparedPressureSurfaceConstraint(
            field::FieldBoundaryDofMap surfaceRows,
            const PressureSurfaceConstraintView constraint
        )
            : m_surfaceRows(std::move(surfaceRows)),
              m_constraint(constraint),
              m_surfaceState(m_surfaceRows.size()) {
        }

        [[nodiscard]] PreparedSurfaceConstraintReport Prepare(
            const mfem::Vector &reducedState,
            const bool stateChanged
        ) {
            MFEM_VERIFY(
                reducedState.Size() == m_surfaceRows.field_size(),
                "The pressure surface constraint received a state vector with the wrong size."
            );

            PreparedSurfaceConstraintReport report;
            if (!m_isPrepared || stateChanged) {
                for (int surfaceIndex = 0; surfaceIndex < m_surfaceRows.size(); ++surfaceIndex) {
                    const double value = reducedState(m_surfaceRows.reduced_dofs()[surfaceIndex]);
                    MFEM_VERIFY(std::isfinite(value), "The pressure surface constraint received non-finite state.");
                    m_surfaceState(surfaceIndex) = value;
                }
                report.cachedSurfaceState = true;
            }

            m_isPrepared = true;
            return report;
        }

        void ApplyResidualRows(mfem::Vector &rowResidual) const {
            VerifyPrepared();
            MFEM_VERIFY(
                rowResidual.Size() == m_surfaceRows.field_size(),
                "The pressure surface constraint received a residual vector with the wrong size."
            );
            m_constraint.ApplyResidualRows(m_surfaceState, m_surfaceRows, rowResidual);
        }

        void ApplyJacobianRows(
            const mfem::Vector &stateVariation,
            mfem::Vector &rowAction
        ) const {
            VerifyPrepared();
            MFEM_VERIFY(
                stateVariation.Size() == m_surfaceRows.field_size() && rowAction.Size() == m_surfaceRows.field_size(),
                "The pressure surface constraint received a Jacobian vector with the wrong size."
            );
            m_constraint.ApplyJacobianRows(m_surfaceState, m_surfaceRows, stateVariation, rowAction);
        }

        [[nodiscard]] bool IsPrepared() const noexcept {
            return m_isPrepared;
        }

        [[nodiscard]] const field::FieldBoundaryDofMap &GetSurfaceRows() const noexcept {
            return m_surfaceRows;
        }

        [[nodiscard]] surface::PressureSurfaceDescriptor GetPhysicalCondition() const noexcept {
            return m_constraint.descriptor();
        }

    private:
        void VerifyPrepared() const {
            MFEM_VERIFY(m_isPrepared, "The pressure surface constraint must be prepared before row application.");
        }

        field::FieldBoundaryDofMap m_surfaceRows;
        PressureSurfaceConstraintView m_constraint;
        mfem::Vector m_surfaceState;
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators