module;

#include <concepts>
#include <type_traits>

export module mean_field:model.compiled_fixed_mass;

export import :field.registry;
export import :model.specifications;
export import :utils.blocks;

export namespace mean_field::models {
    enum class ConstraintRowInjection { append, solver_border };

    template <
        ModelSpecification Specification,
        typename GeneratedValue,
        typename GeneratedResidual,
        typename ValueBlock,
        typename ResidualBlock,
        typename Term,
        typename... StateValueBlocks>
    struct ConstraintLayoutRequest final {
        using SpecificationType                              = Specification;
        using GeneratedValueType                             = GeneratedValue;
        using GeneratedResidualType                          = GeneratedResidual;
        using ValueBlockType                                 = ValueBlock;
        using ResidualBlockType                              = ResidualBlock;
        using TermType                                       = Term;
        using StateValueBlockTypes                           = ModelTypeList<StateValueBlocks...>;

        static constexpr ConstraintRowInjection rowInjection = ConstraintRowInjection::append;
        static constexpr std::size_t valueArity              = GeneratedValue::scalarArity;
        static constexpr std::size_t residualArity           = GeneratedResidual::scalarArity;

        template <typename Form> [[nodiscard]] static consteval auto valueBlock() {
            return utils::blocks::get_value_block<Form>(Term{});
        }

        template <typename Form> [[nodiscard]] static consteval auto residualBlock() {
            return utils::blocks::get_residual_block<Form>(Term{});
        }
    };

    template <typename Candidate>
    concept ConstraintLayoutRequestType = requires {
        typename std::remove_cvref_t<Candidate>::SpecificationType;
        typename std::remove_cvref_t<Candidate>::GeneratedValueType;
        typename std::remove_cvref_t<Candidate>::GeneratedResidualType;
        typename std::remove_cvref_t<Candidate>::ValueBlockType;
        typename std::remove_cvref_t<Candidate>::ResidualBlockType;
        typename std::remove_cvref_t<Candidate>::StateValueBlockTypes;
        requires ModelSpecification<typename std::remove_cvref_t<Candidate>::SpecificationType>;
        requires std::remove_cvref_t<Candidate>::valueArity == std::remove_cvref_t<Candidate>::residualArity;
    };

    using FixedMassLayoutRequest = ConstraintLayoutRequest<
        FixedTotalMass,
        MultiplierFor<FixedTotalMass>,
        ResidualFor<FixedTotalMass>,
        utils::blocks::fixed_total_mass::mass_normalization::value,
        utils::blocks::fixed_total_mass::mass_normalization::residual,
        utils::blocks::fixed_total_mass::mass_normalization,
        utils::blocks::density::mass::value,
        utils::blocks::displacement::geometry::value>;

    class CompiledFixedMass final {
    public:
        using SpecificationType = FixedTotalMass;
        using LayoutRequest     = FixedMassLayoutRequest;
        using MultiplierType    = typename LayoutRequest::GeneratedValueType;
        using ResidualType      = typename LayoutRequest::GeneratedResidualType;

        // In the current barotropic formulation, the multiplier generated by
        // FixedTotalMass is realized by the historical scalar C field.
        using MultiplierField   = field::BarotropicConstant;

        explicit CompiledFixedMass(const FixedTotalMass specification) noexcept : m_specification(specification) {
        }

        [[nodiscard]] const FixedTotalMass &specification() const noexcept {
            return m_specification;
        }

        [[nodiscard]] dimensions::MassValue targetMass() const noexcept {
            return m_specification.targetMass();
        }

        [[nodiscard]] static consteval LayoutRequest layoutRequest() noexcept {
            return {};
        }

    private:
        FixedTotalMass m_specification;
    };

    template <typename Candidate>
    concept CompiledConstraint = requires(const std::remove_cvref_t<Candidate> &constraint) {
        typename std::remove_cvref_t<Candidate>::SpecificationType;
        typename std::remove_cvref_t<Candidate>::LayoutRequest;
        requires ModelSpecification<typename std::remove_cvref_t<Candidate>::SpecificationType>;
        requires ConstraintLayoutRequestType<typename std::remove_cvref_t<Candidate>::LayoutRequest>;
        {
            constraint.specification()
        } noexcept -> std::same_as<const typename std::remove_cvref_t<Candidate>::SpecificationType &>;
        { constraint.layoutRequest() } noexcept -> std::same_as<typename std::remove_cvref_t<Candidate>::LayoutRequest>;
    };

    [[nodiscard]] inline CompiledFixedMass compileConstraint(const FixedTotalMass specification) noexcept {
        return CompiledFixedMass{specification};
    }

    static_assert(ConstraintLayoutRequestType<FixedMassLayoutRequest>);
    static_assert(CompiledConstraint<CompiledFixedMass>);
} // namespace mean_field::models