module;

#include <concepts>
#include <cstddef>
#include <string_view>
#include <type_traits>

export module mean_field:field.base;

export namespace mean_field::field {
    template <typename... Ts> struct TypeList { };

    template <typename T, typename ListT> struct TypeListContains;

    template <typename T, typename... Ts>
    struct TypeListContains<T, TypeList<Ts...>>
        : std::bool_constant<(std::same_as<T, Ts> || ...)> { };

    template <typename T, typename ListT>
    inline constexpr bool typeListContains = TypeListContains<T, ListT>::value;

    enum class StorageKind { finite_element, global_scalar };

    inline constexpr int dynamicBlockSize = -1;

    // -------------------------------------------------------------------------
    // Function-space tags
    // -------------------------------------------------------------------------

    struct L2 {
        static constexpr std::string_view name = "L2";
    };

    struct H1 {
        static constexpr std::string_view name = "H1";
    };

    struct RT {
        static constexpr std::string_view name = "RT";
    };

    struct ND {
        static constexpr std::string_view name = "ND";
    };

    template <typename SpaceT>
    concept SpaceTag = std::same_as<SpaceT, L2> || std::same_as<SpaceT, H1> ||
                       std::same_as<SpaceT, RT> || std::same_as<SpaceT, ND>;

    template <SpaceTag SpaceT, int RankV>
    inline constexpr bool spaceSupportsRank =
        (std::same_as<SpaceT, H1> && (RankV == 0 || RankV == 1)) ||
        (std::same_as<SpaceT, L2> && (RankV == 0 || RankV == 1)) ||
        (std::same_as<SpaceT, RT> && RankV == 1) ||
        (std::same_as<SpaceT, ND> && RankV == 1);

    // -------------------------------------------------------------------------
    // Discretization descriptors
    //
    // familyOrder is the order passed to the backend's FE collection
    // constructor. It is deliberately not called polynomialOrder because those
    // values differ for some spaces, notably Raviart-Thomas elements in MFEM.
    // -------------------------------------------------------------------------

    template <SpaceTag SpaceT, int FamilyOrderV> struct Disc {
        using Space                      = SpaceT;

        static constexpr int familyOrder = FamilyOrderV;

        static_assert(
            FamilyOrderV >= 0,
            "Finite-element family order must be non-negative."
        );
    };

    template <typename T>
    concept DiscretizationTag = requires {
        typename T::Space;
        { T::familyOrder } -> std::convertible_to<int>;
    } && SpaceTag<typename T::Space>;

    // -------------------------------------------------------------------------
    // Physical relations between quantities
    // -------------------------------------------------------------------------

    struct FieldRelation {
        struct Independent { };

        template <typename SourceT> struct Gradient {
            using Source = SourceT;
        };

        template <typename SourceT> struct Divergence {
            using Source = SourceT;
        };

        template <typename SourceT> struct Curl {
            using Source = SourceT;
        };
    };

    template <typename T> struct IsGradient : std::false_type { };

    template <typename T> struct IsDivergence : std::false_type { };

    template <typename T> struct IsCurl : std::false_type { };

    template <typename SourceT>
    struct IsGradient<FieldRelation::Gradient<SourceT>> : std::true_type { };

    template <typename SourceT>
    struct IsDivergence<FieldRelation::Divergence<SourceT>> : std::true_type {
    };

    template <typename SourceT>
    struct IsCurl<FieldRelation::Curl<SourceT>> : std::true_type { };

    template <typename RelationT>
    concept ValidRelation =
        std::same_as<RelationT, FieldRelation::Independent> ||
        IsGradient<RelationT>::value || IsDivergence<RelationT>::value ||
        IsCurl<RelationT>::value;

    template <typename RelationT> struct RelationTarget {
        using Type = void;
    };

    template <typename SourceT>
    struct RelationTarget<FieldRelation::Gradient<SourceT>> {
        using Type = SourceT;
    };

    template <typename SourceT>
    struct RelationTarget<FieldRelation::Divergence<SourceT>> {
        using Type = SourceT;
    };

    template <typename SourceT>
    struct RelationTarget<FieldRelation::Curl<SourceT>> {
        using Type = SourceT;
    };

    template <typename QuantityT>
    using RelationTargetT =
        typename RelationTarget<typename QuantityT::Relation>::Type;

    // -------------------------------------------------------------------------
    // Field quantities
    // -------------------------------------------------------------------------

    template <int RankV, ValidRelation RelationT, DiscretizationTag DiscT>
    struct Quantity {
        using Relation                           = RelationT;
        using Discretization                     = DiscT;
        using Space                              = typename DiscT::Space;

        static constexpr int rankValue           = RankV;
        static constexpr int familyOrder         = DiscT::familyOrder;

        static constexpr StorageKind storageKind = StorageKind::finite_element;
        static constexpr int staticBlockSize     = dynamicBlockSize;

        static_assert(
            RankV >= 0,
            "A field quantity cannot have a negative tensor rank."
        );

        static_assert(
            spaceSupportsRank<
                Space,
                RankV>,
            "This function space cannot represent a quantity of this rank."
        );
    };

    template <ValidRelation RelationT, DiscretizationTag DiscT>
    using ScalarQ = Quantity<0, RelationT, DiscT>;

    template <ValidRelation RelationT, DiscretizationTag DiscT>
    using VectorQ = Quantity<1, RelationT, DiscT>;

    struct GlobalScalarQ {
        using Relation                           = FieldRelation::Independent;

        static constexpr int rankValue           = 0;
        static constexpr StorageKind storageKind = StorageKind::global_scalar;
        static constexpr int staticBlockSize     = 1;
    };

    template <typename T>
    concept FieldQuantity =
        requires {
            typename T::Relation;
            typename T::Discretization;
            typename T::Space;

            { T::rankValue } -> std::convertible_to<int>;
            { T::familyOrder } -> std::convertible_to<int>;
            { T::storageKind } -> std::convertible_to<StorageKind>;
            { T::staticBlockSize } -> std::convertible_to<int>;
        } && SpaceTag<typename T::Space> &&
        T::storageKind == StorageKind::finite_element;

    template <typename T>
    concept GlobalScalarQuantity =
        requires {
            typename T::Relation;

            { T::rankValue } -> std::convertible_to<int>;
            { T::storageKind } -> std::convertible_to<StorageKind>;
            { T::staticBlockSize } -> std::convertible_to<int>;
        } && T::rankValue == 0 &&
        T::storageKind == StorageKind::global_scalar && T::staticBlockSize == 1;

    template <typename T>
    concept RegisteredQuantity = FieldQuantity<T> || GlobalScalarQuantity<T>;

    template <typename QuantityT>
    concept DerivedQuantity = FieldQuantity<QuantityT> &&
                              (!std::same_as<RelationTargetT<QuantityT>, void>);

    // -------------------------------------------------------------------------
    // Compile-time discretization constraints
    // -------------------------------------------------------------------------

    template <FieldQuantity FluxT, FieldQuantity PotentialT>
    struct RtL2StablePair {
        static consteval void validate() {
            static_assert(
                std::same_as<typename FluxT::Space, RT>,
                "The flux in an RT/L2 pair must use Raviart-Thomas elements."
            );

            static_assert(
                std::same_as<typename PotentialT::Space, L2>,
                "The potential in an RT/L2 pair must use L2 elements."
            );

            static_assert(
                FluxT::rankValue == 1,
                "The flux in an RT/L2 pair must be vector-valued."
            );

            static_assert(
                PotentialT::rankValue == 0,
                "The potential in an RT/L2 pair must be scalar-valued."
            );

            static_assert(
                FluxT::familyOrder == PotentialT::familyOrder,
                "The MFEM RT and L2 family orders must match."
            );
        }
    };

    template <typename... ConstraintTs>
    consteval bool validate_constraints(TypeList<ConstraintTs...>) {
        (ConstraintTs::validate(), ...);
        return true;
    }

    // -------------------------------------------------------------------------
    // Operations applied to quantities inside weak forms
    //
    // These describe the mathematics. Backend-specific polynomial-order rules
    // are provided by field.mfem.
    // -------------------------------------------------------------------------

    struct FieldOperation {
        struct Value { };
        struct Gradient { };
        struct Divergence { };
        struct Curl { };
        struct NormalTrace { };
    };

    template <typename OperationT>
    concept FieldOperationTag =
        std::same_as<OperationT, FieldOperation::Value> ||
        std::same_as<OperationT, FieldOperation::Gradient> ||
        std::same_as<OperationT, FieldOperation::Divergence> ||
        std::same_as<OperationT, FieldOperation::Curl> ||
        std::same_as<OperationT, FieldOperation::NormalTrace>;

    template <
        RegisteredQuantity QuantityT,
        FieldOperationTag OperationT = FieldOperation::Value>
    struct Operand {
        using Quantity  = QuantityT;
        using Operation = OperationT;

        static_assert(
            FieldQuantity<QuantityT> || std::same_as<
                                            OperationT,
                                            FieldOperation::Value>,
            "Global scalar quantities support only the value operation."
        );
    };

    template <typename T>
    concept FieldOperand =
        requires {
            typename T::Quantity;
            typename T::Operation;
        } && RegisteredQuantity<typename T::Quantity> &&
        FieldOperationTag<typename T::Operation>;

    // -------------------------------------------------------------------------
    // Weak-form descriptions
    //
    // PolicyKeyV associates the form with a runtime quadrature-policy key.
    //
    // DynamicOrderCountV is the number of polynomial-order contributions that
    // cannot yet be derived from registered quantities. For example, a source
    // coefficient supplied at runtime contributes one dynamic order.
    // -------------------------------------------------------------------------

    template <
        auto PolicyKeyV,
        std::size_t DynamicOrderCountV,
        FieldOperand... OperandTs>
    struct FormSpec {
        static constexpr auto policyKey                = PolicyKeyV;
        static constexpr std::size_t dynamicOrderCount = DynamicOrderCountV;

        using Operands                                 = TypeList<OperandTs...>;
    };

    template <typename T>
    concept FieldForm = requires {
        typename T::Operands;

        T::policyKey;

        { T::dynamicOrderCount } -> std::convertible_to<std::size_t>;
    };

    template <typename ListT>
    struct IsRegisteredQuantityList : std::false_type { };

    template <RegisteredQuantity... QuantityTs>
    struct IsRegisteredQuantityList<TypeList<QuantityTs...>> : std::true_type {
    };

    template <typename ListT>
    inline constexpr bool isRegisteredQuantityList =
        IsRegisteredQuantityList<ListT>::value;

    template <typename ListT> struct IsFieldFormList : std::false_type { };

    template <FieldForm... FormTs>
    struct IsFieldFormList<TypeList<FormTs...>> : std::true_type { };

    template <typename ListT>
    inline constexpr bool isFieldFormList = IsFieldFormList<ListT>::value;
} // namespace mean_field::field