module;

#include <concepts>
#include <string_view>

#ifndef MEAN_FIELD_UNIFORM_POLYNOMIAL_ORDER_INCREMENT
#define MEAN_FIELD_UNIFORM_POLYNOMIAL_ORDER_INCREMENT 0
#endif

export module mean_field:field.registry;

export import :dimensions.quantities;
export import :field.base;
export import :quadrature.policy;
export import :utils.domain;

export namespace mean_field::field {
    inline constexpr int uniformPolynomialOrderIncrement = MEAN_FIELD_UNIFORM_POLYNOMIAL_ORDER_INCREMENT;
    static_assert(uniformPolynomialOrderIncrement >= 0);

    // =========================================================================
    // Density
    // =========================================================================

    struct Density {
        static constexpr std::string_view name = "density";
        static constexpr int scalarOrder       = 2 + uniformPolynomialOrderIncrement;

        using PhysicalQuantity                 = dimensions::quantity::Density;
        using Support                          = DomainSupport<utils::domain::Stellar>;

        struct Scalar final : ScalarQ<FieldRelation::Independent, Disc<L2, scalarOrder>> {
            static constexpr std::string_view symbol = "ρ";
        };

        using Quantities                          = TypeList<Scalar>;
        using Constraints                         = TypeList<>;

        static constexpr bool constraintsAreValid = validate_constraints(Constraints{});

        static_assert(constraintsAreValid);

        struct Form {
            // Density-space mass matrix: (rho, q).
            using ProjectionMass = FormSpec<quadrature::Term::density_projection, 0, Operand<Scalar>, Operand<Scalar>>;

            // Projection RHS with one runtime coefficient order.
            using ProjectionSource  = FormSpec<quadrature::Term::density_projection, 1, Operand<Scalar>>;

            // Density-space contribution to the barotropic EOS closure:
            // (rho, q_rho).
            using EosClosureMass    = FormSpec<quadrature::Term::eos_closure, 0, Operand<Scalar>, Operand<Scalar>>;

            // Integral of density over the physical volume.
            using MassConservation  = FormSpec<quadrature::Term::mass_conservation, 0, Operand<Scalar>>;

            // The same physical integral used as a nonlinear normalization
            // constraint. It has a distinct policy key so solver assembly and
            // diagnostics can be overintegrated independently.
            using MassNormalization = FormSpec<quadrature::Term::mass_normalization, 0, Operand<Scalar>>;

            // Integral of rho * x. The combined position-coefficient order is
            // supplied as one dynamic order.
            using CenterOfMass      = FormSpec<quadrature::Term::center_of_mass, 1, Operand<Scalar>>;

            // Integral of rho times the quadratic position tensor. The
            // combined tensor-coefficient order is supplied dynamically.
            using Quadrupole        = FormSpec<quadrature::Term::quadrupole, 1, Operand<Scalar>>;

            using ErrorNorm         = FormSpec<quadrature::Term::error_norm, 0, Operand<Scalar>, Operand<Scalar>>;
        };

        using FormList = TypeList<
            Form::ProjectionMass,
            Form::ProjectionSource,
            Form::EosClosureMass,
            Form::MassConservation,
            Form::MassNormalization,
            Form::CenterOfMass,
            Form::Quadrupole,
            Form::ErrorNorm>;
    };

    // =========================================================================
    // Gravity
    // =========================================================================

    struct Gravity {
        static constexpr std::string_view name = "gravity";

        static constexpr int potentialOrder    = 2 + uniformPolynomialOrderIncrement;
        static constexpr int fluxOrder         = 2 + uniformPolynomialOrderIncrement;

        using Support                          = DomainSupport<utils::domain::All>;

        struct Potential final : ScalarQ<FieldRelation::Independent, Disc<L2, potentialOrder>> {
            static constexpr std::string_view symbol = "φ";
        };

        struct Flux final : VectorQ<FieldRelation::Gradient<Potential>, Disc<RT, fluxOrder>> {
            static constexpr std::string_view symbol = "∇φ";
        };

        using Quantities                          = TypeList<Potential, Flux>;

        using Constraints                         = TypeList<RtL2StablePair<Flux, Potential>>;

        static constexpr bool constraintsAreValid = validate_constraints(Constraints{});

        static_assert(constraintsAreValid);

        struct Form {
            using HDivMass           = FormSpec<quadrature::Term::gravity_hdiv_mass, 0, Operand<Flux>, Operand<Flux>>;

            using DivergenceCoupling = FormSpec<
                quadrature::Term::gravity_divergence,
                0,
                Operand<Flux, FieldOperation::Divergence>,
                Operand<Potential>>;

            using Boundary = FormSpec<
                quadrature::Term::gravity_boundary,
                0,
                Operand<Flux, FieldOperation::NormalTrace>,
                Operand<Flux, FieldOperation::NormalTrace>>;

            // Density is a registered coefficient field and potential is the
            // test field, so the full polynomial order is compile-time data.
            using SourceLinear =
                FormSpec<quadrature::Term::gravity_source, 0, Operand<Density::Scalar>, Operand<Potential>>;

            // Mixed density-to-potential projection. Both trial and test
            // orders are registered quantities.
            using SourceProjection =
                FormSpec<quadrature::Term::gravity_source, 0, Operand<Density::Scalar>, Operand<Potential>>;

            using PotentialErrorNorm =
                FormSpec<quadrature::Term::error_norm, 0, Operand<Potential>, Operand<Potential>>;

            using FluxErrorNorm = FormSpec<quadrature::Term::error_norm, 0, Operand<Flux>, Operand<Flux>>;
        };

        using FormList = TypeList<
            Form::HDivMass,
            Form::DivergenceCoupling,
            Form::Boundary,
            Form::SourceLinear,
            Form::SourceProjection,
            Form::PotentialErrorNorm,
            Form::FluxErrorNorm>;
    };

    // =========================================================================
    // Displacement
    // =========================================================================

    struct Displacement {
        static constexpr std::string_view name = "displacement";
        static constexpr int vectorOrder       = 3 + uniformPolynomialOrderIncrement;

        using Support                          = DomainSupport<utils::domain::All>;

        struct Vector final : VectorQ<FieldRelation::Independent, Disc<H1, vectorOrder>> {
            static constexpr std::string_view symbol = "d";
        };

        using Quantities                          = TypeList<Vector>;
        using Constraints                         = TypeList<>;

        static constexpr bool constraintsAreValid = validate_constraints(Constraints{});

        static_assert(constraintsAreValid);

        struct Form {
            // Harmonic or pseudoelastic interior mesh extension. For the
            // initial Laplacian model this is (grad d, grad w).
            using MeshExtension = FormSpec<
                quadrature::Term::mesh_extension,
                0,
                Operand<Vector, FieldOperation::Gradient>,
                Operand<Vector, FieldOperation::Gradient>>;

            // Positive gravitational contribution to the displacement row:
            //
            //     int rho grad(phi) . w dV.
            //
            // Both the base geometry Jacobian and the displacement test
            // function contribute to the polynomial order. The RT flux is
            // mapped to physical space by the contravariant Piola map.
            using GravityForce = FormSpec<
                quadrature::Term::gravity_force,
                0,
                Operand<Density::Scalar>,
                Operand<Gravity::Flux>,
                Operand<Vector, FieldOperation::Gradient>,
                Operand<Vector>>;

            // Rigid-rotation contribution to the displacement row:
            //
            //   -int rho grad(Psi_rotation) . w dV.
            //
            // grad(Psi_rotation) is linear in physical position, so its
            // polynomial order is supplied as one runtime contribution.
            using CentrifugalForce =
                FormSpec<quadrature::Term::centrifugal, 1, Operand<Density::Scalar>, Operand<Vector>>;

            using ErrorNorm = FormSpec<quadrature::Term::error_norm, 0, Operand<Vector>, Operand<Vector>>;
        };

        using FormList = TypeList<Form::MeshExtension, Form::GravityForce, Form::CentrifugalForce, Form::ErrorNorm>;
    };

    // Current realization of MultiplierFor<FixedTotalMass>. This remains a
    // barotrope-specific field representation: the specification compiler,
    // rather than the universal state registry, decides when it is present.
    struct BarotropicConstant {
        static constexpr std::string_view name = "barotropic_constant";

        using Support                          = NonSpatialSupport;

        struct Scalar final : GlobalScalarQ {
            static constexpr std::string_view symbol = "C";
        };

        using Quantities                          = TypeList<Scalar>;
        using Constraints                         = TypeList<>;
        using FormList                            = TypeList<>;

        static constexpr bool constraintsAreValid = validate_constraints(Constraints{});

        static_assert(constraintsAreValid);
    };

    // Solver border generated by FixedCentralDensity. This is deliberately a
    // non-spatial numerical coordinate rather than a physical stellar field.
    struct CentralDensityBorder {
        static constexpr std::string_view name = "central_density_border";

        using Support                          = NonSpatialSupport;

        struct Scalar final : GlobalScalarQ {
            static constexpr std::string_view symbol = "lambda_rho_c";
        };

        using Quantities                          = TypeList<Scalar>;
        using Constraints                         = TypeList<>;
        using FormList                            = TypeList<>;

        static constexpr bool constraintsAreValid = validate_constraints(Constraints{});

        static_assert(constraintsAreValid);
    };

    // =========================================================================
    // Specific enthalpy
    //
    // Pressure is deliberately not registered as an independent field. For a
    // barotrope it is derived from h through the EOS, while h supplies the
    // continuous H1 trace used to define the isobaric stellar surface.
    // =========================================================================

    struct Enthalpy {
        static constexpr std::string_view name = "specific_enthalpy";
        static constexpr int scalarOrder       = 3 + uniformPolynomialOrderIncrement;

        using PhysicalQuantity                 = dimensions::quantity::SpecificEnthalpy;
        using Support                          = DomainSupport<utils::domain::Stellar>;

        struct Scalar final : ScalarQ<FieldRelation::Independent, Disc<H1, scalarOrder>> {
            static constexpr std::string_view symbol = "h";
        };

        using Quantities                          = TypeList<Scalar>;
        using Constraints                         = TypeList<>;

        static constexpr bool constraintsAreValid = validate_constraints(Constraints{});

        static_assert(constraintsAreValid);

        struct Form {
            // EOS source contribution (rho(h), q_rho). The dynamic order is
            // the extra polynomial order introduced by the nonlinear EOS
            // beyond the registered order of h. For an n=3 polytrope this is
            // 2 * hOrder, making rho(h) cubic in h.
            using EosClosureSource =
                FormSpec<quadrature::Term::eos_closure, 1, Operand<Scalar>, Operand<Density::Scalar>>;

            // (h, q_h) contribution to
            // h + phi - Psi_rotation - C = 0.
            using EquilibriumEnthalpy =
                FormSpec<quadrature::Term::hydrostatic_equilibrium, 0, Operand<Scalar>, Operand<Scalar>>;

            // (phi, q_h) contribution to hydrostatic equilibrium.
            using EquilibriumGravity =
                FormSpec<quadrature::Term::hydrostatic_equilibrium, 0, Operand<Gravity::Potential>, Operand<Scalar>>;

            // (Psi_rotation, q_h). The rotation-potential order is supplied
            // dynamically because it belongs to runtime rotation data.
            using EquilibriumRotation = FormSpec<quadrature::Term::hydrostatic_equilibrium, 1, Operand<Scalar>>;

            // (C, q_h), where C is spatially constant.
            using EquilibriumConstant = FormSpec<
                quadrature::Term::hydrostatic_equilibrium,
                0,
                Operand<BarotropicConstant::Scalar>,
                Operand<Scalar>>;

            // Boundary trace form available for weak enforcement, testing, or
            // a future multiplier formulation of h|Gamma_star = 0.
            using IsobaricSurface  = FormSpec<quadrature::Term::isobaric_surface, 0, Operand<Scalar>, Operand<Scalar>>;

            // Integral of P(h). The dynamic order is the extra EOS order
            // beyond the registered order of h.
            using PressureIntegral = FormSpec<quadrature::Term::pressure_integral, 1, Operand<Scalar>>;

            // Weak pressure force in the displacement test space:
            //
            //     -int P(h) I : grad(w) dV
            //
            // which is equivalent to -int P(h) div(w) dV. The dynamic order
            // is the extra EOS order beyond the registered order of h. For an
            // n=3 polytrope this is 3 * hOrder, making P(h) quartic in h.
            using PressureForce    = FormSpec<
                quadrature::Term::pressure_force,
                1,
                Operand<Scalar>,
                Operand<Displacement::Vector, FieldOperation::Gradient>>;

            using ErrorNorm = FormSpec<quadrature::Term::error_norm, 0, Operand<Scalar>, Operand<Scalar>>;
        };

        using FormList = TypeList<
            Form::EosClosureSource,
            Form::EquilibriumEnthalpy,
            Form::EquilibriumGravity,
            Form::EquilibriumRotation,
            Form::EquilibriumConstant,
            Form::IsobaricSurface,
            Form::PressureIntegral,
            Form::PressureForce,
            Form::ErrorNorm>;
    };

    // =========================================================================
    // Field definition concept
    // =========================================================================

    template <typename T>
    concept FieldTag =
        requires {
            typename T::Quantities;
            typename T::Constraints;
            typename T::FormList;
            typename T::Support;

            { T::name } -> std::convertible_to<std::string_view>;
        } && IsFieldSupport<typename T::Support> && isRegisteredQuantityList<typename T::Quantities> &&
        isFieldFormList<typename T::FormList>;

    static_assert(FieldTag<Gravity>);
    static_assert(FieldTag<Displacement>);
    static_assert(FieldTag<Density>);
    static_assert(FieldTag<Enthalpy>);
    static_assert(FieldTag<BarotropicConstant>);

    static_assert(DerivedQuantity<Gravity::Flux>);

    static_assert(std::same_as<
                  RelationTargetT<Gravity::Flux>,
                  Gravity::Potential>);

    static_assert(std::same_as<
                  FieldDomainT<Density>,
                  utils::domain::Stellar>);

    static_assert(std::same_as<
                  FieldDomainT<Enthalpy>,
                  utils::domain::Stellar>);

    static_assert(std::same_as<
                  FieldDomainT<Gravity>,
                  utils::domain::All>);

    static_assert(std::same_as<
                  FieldDomainT<Displacement>,
                  utils::domain::All>);

    static_assert(NonSpatialField<BarotropicConstant>);
} // namespace mean_field::field