module;

#include <concepts>
#include <string_view>

export module mean_field:field.registry;

export import :field.base;
export import :quadrature.policy;

export namespace mean_field::field {
    // =========================================================================
    // Density
    // =========================================================================

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

        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;
        static constexpr int fluxOrder         = 2;

        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;

        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>>;

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

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

    struct BarotropicConstant {
        static constexpr std::string_view name = "barotropic_constant";

        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);
    };

    // =========================================================================
    // 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;

        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;

            { T::name } -> std::convertible_to<std::string_view>;
        } && 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>);
} // namespace mean_field::field