module;

#include <array>
#include <concepts>
#include <cstddef>
#include <memory>
#include <stdexcept>

#include <mfem.hpp>

export module mean_field:field.mfem;

export import :field.registry;

namespace mean_field::field::detail {
    template <typename T> inline constexpr bool alwaysFalse = false;

    // -------------------------------------------------------------------------
    // MFEM finite-element collection construction
    // -------------------------------------------------------------------------

    template <typename SpaceT> struct FecFor;

    template <> struct FecFor<L2> {
        static std::unique_ptr<mfem::FiniteElementCollection> make(
            int familyOrder,
            int dimension
        ) {
            return std::make_unique<mfem::L2_FECollection>(
                familyOrder, dimension
            );
        }
    };

    template <> struct FecFor<H1> {
        static std::unique_ptr<mfem::FiniteElementCollection> make(
            int familyOrder,
            int dimension
        ) {
            return std::make_unique<mfem::H1_FECollection>(
                familyOrder, dimension
            );
        }
    };

    template <> struct FecFor<RT> {
        static std::unique_ptr<mfem::FiniteElementCollection> make(
            int familyOrder,
            int dimension
        ) {
            return std::make_unique<mfem::RT_FECollection>(
                familyOrder, dimension
            );
        }
    };

    template <> struct FecFor<ND> {
        static std::unique_ptr<mfem::FiniteElementCollection> make(
            int familyOrder,
            int dimension
        ) {
            return std::make_unique<mfem::ND_FECollection>(
                familyOrder, dimension
            );
        }
    };

    // -------------------------------------------------------------------------
    // MFEM polynomial-order interpretation
    //
    // familyOrder is the collection constructor argument.
    //
    // For RT_p:
    //   value order       = p + 1
    //   divergence order  = p
    //   normal-trace order = p
    //
    // This distinction is what allows Disc<RT, p> and Disc<L2, p> to form a
    // compatible pair while still giving different value-shape orders.
    // -------------------------------------------------------------------------

    template <typename OperandT> struct MfemOperandOrder;

    template <RegisteredQuantity QuantityT, FieldOperationTag OperationT>
    struct MfemOperandOrder<Operand<QuantityT, OperationT>> {
        static constexpr int orderValue = []() consteval {
            if constexpr (GlobalScalarQuantity<QuantityT>) {
                static_assert(
                    std::same_as<OperationT, FieldOperation::Value>,
                    "Global scalars support only the value operation."
                );

                return 0;
            } else {
                using Space               = typename QuantityT::Space;
                constexpr int familyOrder = QuantityT::familyOrder;

                if constexpr (std::same_as<OperationT, FieldOperation::Value>) {
                    if constexpr (std::same_as<Space, RT>) {
                        return familyOrder + 1;
                    } else {
                        return familyOrder;
                    }
                } else if constexpr (
                    std::same_as<OperationT, FieldOperation::Divergence>
                ) {
                    static_assert(
                        std::same_as<Space, RT>,
                        "Only RT quantities currently support the divergence "
                        "polynomial-order rule."
                    );

                    return familyOrder;
                } else if constexpr (
                    std::same_as<OperationT, FieldOperation::Gradient>
                ) {
                    static_assert(
                        std::same_as<Space, H1>,
                        "Only H1 quantities currently support the gradient "
                        "polynomial-order rule."
                    );

                    return familyOrder > 0 ? familyOrder - 1 : 0;
                } else if constexpr (
                    std::same_as<OperationT, FieldOperation::Curl>
                ) {
                    static_assert(
                        std::same_as<Space, ND>,
                        "Only ND quantities currently support the curl "
                        "polynomial-order rule."
                    );

                    return familyOrder > 0 ? familyOrder - 1 : 0;
                } else if constexpr (
                    std::same_as<OperationT, FieldOperation::NormalTrace>
                ) {
                    static_assert(
                        std::same_as<Space, RT>,
                        "Only RT quantities currently support the normal-trace "
                        "polynomial-order rule."
                    );

                    return familyOrder;
                } else {
                    static_assert(
                        alwaysFalse<OperationT>,
                        "Unsupported MFEM field operation."
                    );
                }
            }
        }();
    };

    // -------------------------------------------------------------------------
    // Static polynomial-order contribution of an entire form
    // -------------------------------------------------------------------------

    template <typename FormT> struct MfemFormOrder;

    template <
        auto PolicyKeyV,
        std::size_t DynamicOrderCountV,
        FieldOperand... OperandTs>
    struct MfemFormOrder<
        FormSpec<PolicyKeyV, DynamicOrderCountV, OperandTs...>> {
        static constexpr int staticOrder =
            (MfemOperandOrder<OperandTs>::orderValue + ... + 0);
    };

    // -------------------------------------------------------------------------
    // MFEM vector-dimension and ordering rules
    //
    // Vector H1/L2 fields are represented using multiple copies of a scalar
    // finite-element space. RT and ND elements are intrinsically vector-valued
    // and therefore use vdim = 1.
    // -------------------------------------------------------------------------

    template <FieldQuantity QuantityT> int get_vdim(int spaceDimension) {
        if (spaceDimension <= 0) {
            throw std::invalid_argument("Space dimension must be positive.");
        }

        if constexpr (QuantityT::rankValue == 0) {
            return 1;
        } else if constexpr (
            std::same_as<typename QuantityT::Space, H1> ||
            std::same_as<typename QuantityT::Space, L2>
        ) {
            return spaceDimension;
        } else {
            return 1;
        }
    }

    template <FieldQuantity QuantityT>
    constexpr mfem::Ordering::Type get_ordering() {
        if constexpr (
            QuantityT::rankValue == 1 &&
            (std::same_as<typename QuantityT::Space, H1> ||
             std::same_as<typename QuantityT::Space, L2>)
        ) {
            return mfem::Ordering::byVDIM;
        } else {
            return mfem::Ordering::byNODES;
        }
    }

    // -------------------------------------------------------------------------
    // Quantity-specific MFEM realization
    //
    // Backend choices that are part of a field definition live here rather
    // than leaking into FEM setup or call sites.
    // -------------------------------------------------------------------------

    template <FieldQuantity QuantityT> struct MfemQuantityTraits {
        static std::unique_ptr<mfem::FiniteElementCollection>
        make_fec(int dimension) {
            return FecFor<typename QuantityT::Space>::make(
                QuantityT::familyOrder, dimension
            );
        }

        static constexpr mfem::Ordering::Type ordering =
            get_ordering<QuantityT>();
    };

    template <> struct MfemQuantityTraits<Gravity::Flux> {
        static std::unique_ptr<mfem::FiniteElementCollection>
        make_fec(int dimension) {
            return std::make_unique<mfem::RT_FECollection>(
                Gravity::Flux::familyOrder, dimension,
                mfem::BasisType::GaussLobatto, mfem::BasisType::IntegratedGLL
            );
        }

        static constexpr mfem::Ordering::Type ordering =
            mfem::Ordering::byNODES;
    };

    template <> struct MfemQuantityTraits<Displacement::Vector> {
        static std::unique_ptr<mfem::FiniteElementCollection>
        make_fec(int dimension) {
            return FecFor<H1>::make(
                Displacement::Vector::familyOrder, dimension
            );
        }

        static constexpr mfem::Ordering::Type ordering =
            mfem::Ordering::byNODES;
    };
} // namespace mean_field::field::detail

export namespace mean_field::field {
    // -------------------------------------------------------------------------
    // User-facing field type
    //
    // The object itself is currently a zero-cost compile-time descriptor:
    //
    //     Field<Gravity> gravityField;
    //
    // MFEM construction and typed quadrature-query generation are provided as
    // static operations. Runtime ownership can later be added without changing
    // Gravity, Displacement, or their form definitions.
    // -------------------------------------------------------------------------

    template <FieldTag TagT> class Field {
    public:
        using Tag = TagT;

        // ---------------------------------------------------------------------
        // MFEM finite-element collection construction
        // ---------------------------------------------------------------------

        template <FieldQuantity QuantityT>
            requires typeListContains<
                QuantityT,
                typename TagT::Quantities>
        static std::unique_ptr<mfem::FiniteElementCollection>
        make_fec(int dimension) {
            if (dimension <= 0) {
                throw std::invalid_argument("Mesh dimension must be positive.");
            }

            return detail::MfemQuantityTraits<QuantityT>::make_fec(dimension);
        }

        // ---------------------------------------------------------------------
        // MFEM parallel finite-element space construction
        //
        // The finite-element collection must outlive the returned space.
        // ---------------------------------------------------------------------

        template <FieldQuantity QuantityT>
            requires typeListContains<
                QuantityT,
                typename TagT::Quantities>
        static std::unique_ptr<mfem::ParFiniteElementSpace> make_fespace(
            mfem::ParMesh &mesh,
            mfem::FiniteElementCollection &finiteElementCollection
        ) {
            return std::make_unique<mfem::ParFiniteElementSpace>(
                &mesh, &finiteElementCollection,
                detail::get_vdim<QuantityT>(mesh.SpaceDimension()),
                detail::MfemQuantityTraits<QuantityT>::ordering
            );
        }

        // ---------------------------------------------------------------------
        // Typed quadrature-query construction
        //
        // geometryWeightOrder is supplied at runtime because it depends on the
        // actual element transformation.
        //
        // dynamicOrders contains the form-specific polynomial orders that are
        // not represented by registered compile-time quantities.
        //
        // Examples:
        //
        // Density::Form::CenterOfMass:
        //     { positionOrder }
        //
        // Gravity source forms need no dynamic orders because density and
        // potential are both registered quantities.
        //
        // The completed base order is stored in Query::base_order, so Policy
        // does not need to understand divergence, RT conventions, or individual
        // field layouts.
        // ---------------------------------------------------------------------

        template <FieldForm FormT>
            requires typeListContains<
                FormT,
                typename TagT::FormList>
        static constexpr quadrature::Query make_query(
            quadrature::QuadratureRole role,
            int geometryWeightOrder,
            std::array<
                int,
                FormT::dynamicOrderCount> dynamicOrders = {},
            utils::DOMAINS domain                       = utils::DOMAINS::ALL,
            quadrature::MappingKind mapping = quadrature::MappingKind::none
        ) {
            if (geometryWeightOrder < 0) {
                throw std::invalid_argument(
                    "Geometry weight order cannot be negative."
                );
            }

            int baseOrder =
                detail::MfemFormOrder<FormT>::staticOrder + geometryWeightOrder;

            for (const int dynamicOrder : dynamicOrders) {
                if (dynamicOrder < 0) {
                    throw std::invalid_argument(
                        "Dynamic polynomial orders cannot be negative."
                    );
                }

                baseOrder += dynamicOrder;
            }

            return {
                .term                  = FormT::policyKey,
                .role                  = role,
                .domain                = domain,
                .mapping               = mapping,
                .trial_order           = 0,
                .test_order            = 0,
                .coefficient_order     = 0,
                .geometry_weight_order = geometryWeightOrder,
                .base_order            = baseOrder
            };
        }
    };

    static_assert(FieldTag<Gravity>);
    static_assert(FieldTag<Displacement>);
    static_assert(FieldTag<Density>);
    static_assert(FieldTag<BarotropicConstant>);
} // namespace mean_field::field