module;

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

#include <mfem.hpp>

export module mean_field:model.compiled_fixed_angular_momentum;

export import :field.registry;
export import :model.compiled_fixed_mass;
export import :physics.rigid_rotation;
export import :utils.blocks;

export namespace mean_field::models {
    using FixedAngularMomentumLayoutRequest = ConstraintLayoutRequest<
        FixedAngularMomentum,
        PhysicalCoordinateFor<FixedAngularMomentum>,
        ResidualFor<FixedAngularMomentum>,
        utils::blocks::fixed_angular_momentum::angular_velocity::value,
        utils::blocks::fixed_angular_momentum::angular_velocity::residual,
        utils::blocks::fixed_angular_momentum::angular_velocity,
        utils::blocks::density::mass::value,
        utils::blocks::surface_deformation::parameters::value,
        utils::blocks::fixed_angular_momentum::angular_velocity::value>;

    class CompiledFixedAngularMomentum final {
    public:
        using SpecificationType    = FixedAngularMomentum;
        using LayoutRequest        = FixedAngularMomentumLayoutRequest;
        using AngularVelocityType  = typename LayoutRequest::GeneratedValueType;
        using ResidualType         = typename LayoutRequest::GeneratedResidualType;
        using AngularVelocityField = field::AngularVelocity;

        explicit CompiledFixedAngularMomentum(const FixedAngularMomentum specification) noexcept
            : m_specification(specification) {
        }

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

        [[nodiscard]] dimensions::AngularMomentumValue targetAngularMomentum() const noexcept {
            return m_specification.targetAngularMomentum();
        }

        [[nodiscard]] physics::RigidRotation makeRotation(const double angularVelocity) const {
            mfem::Vector velocity(3);
            mfem::Vector center(3);
            for (int component = 0; component < 3; ++component) {
                velocity(component) = angularVelocity * m_specification.axis()[static_cast<std::size_t>(component)];
                center(component)   = m_specification.center()[static_cast<std::size_t>(component)];
            }
            return {velocity, center};
        }

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

    private:
        FixedAngularMomentum m_specification;
    };

    [[nodiscard]] inline CompiledFixedAngularMomentum
    compileConstraint(const FixedAngularMomentum specification) noexcept {
        return CompiledFixedAngularMomentum{specification};
    }

    static_assert(ConstraintLayoutRequestType<FixedAngularMomentumLayoutRequest>);
    static_assert(CompiledConstraint<CompiledFixedAngularMomentum>);
} // namespace mean_field::models