module;

#include <compare>
#include <cstdint>
#include <vector>

#include <mfem.hpp>

export module mean_field:operators.prepared_angular_momentum;

export import :fem;
export import :mapping.domain_mapper;
export import :model.compiled_fixed_angular_momentum;
export import :operators.context.gravity_field;

export namespace mean_field::operators {
    struct AngularMomentumDependencyStamp final {
        std::uint64_t identity{0};
        std::uint64_t revision{0};

        constexpr auto operator<=>(const AngularMomentumDependencyStamp &) const = default;
    };

    struct AngularMomentumDependencies final {
        AngularMomentumDependencyStamp discretization;
        AngularMomentumDependencyStamp density;
        AngularMomentumDependencyStamp displacement;
        AngularMomentumDependencyStamp rotation;

        constexpr auto operator<=>(const AngularMomentumDependencies &) const = default;
    };

    struct PreparedAngularMomentumReport final {
        bool rebuiltStaticPlan{false};
        bool refreshedGeometry{false};
        bool refreshedDensity{false};
        bool updatedAngularVelocity{false};
        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return rebuiltStaticPlan || refreshedGeometry || refreshedDensity || updatedAngularVelocity ||
                   assembledResidual;
        }

        constexpr auto operator<=>(const PreparedAngularMomentumReport &) const = default;
    };

    struct AngularMomentumConstraintReport final {
        double targetAngularMomentum;
        double achievedAngularMomentum;
        double momentOfInertia;
        double angularVelocity;
        double dimensionalResidual;
        double scaledResidual;
    };

    struct PreparedAngularMomentumActionStatistics final {
        std::uint64_t densityApplications{0};
        std::uint64_t displacementApplications{0};
        std::uint64_t angularVelocityApplications{0};
        std::uint64_t completeApplications{0};

        constexpr auto operator<=>(const PreparedAngularMomentumActionStatistics &) const = default;
    };

    /*
     * Prepared scalar invariant
     *
     *   R_J(rho, d, Omega) = Omega I_axis(rho, d) - J_target,
     *   I_axis             = integral rho |(x-x_0)_perp|^2 dV.
     *
     * The axis is normalized by CompiledFixedAngularMomentum. Density and
     * geometry are borrowed from the shared gravity context, so this row is
     * linearized at exactly the same mapped state as every physical equation.
     */
    class PreparedAngularMomentumOperator final {
    public:
        using SpecificationType      = models::FixedAngularMomentum;
        using CompiledConstraintType = models::CompiledFixedAngularMomentum;
        using Dependencies           = AngularMomentumDependencies;
        using Report                 = PreparedAngularMomentumReport;

        PreparedAngularMomentumOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const context::gravity_field::GravityFieldLinearizationContext &gravityContext,
            models::CompiledFixedAngularMomentum constraint
        );

        PreparedAngularMomentumOperator(const PreparedAngularMomentumOperator &)            = delete;
        PreparedAngularMomentumOperator &operator=(const PreparedAngularMomentumOperator &) = delete;
        PreparedAngularMomentumOperator(PreparedAngularMomentumOperator &&)                 = delete;
        PreparedAngularMomentumOperator &operator=(PreparedAngularMomentumOperator &&)      = delete;

        PreparedAngularMomentumReport Prepare(
            double angularVelocity,
            const AngularMomentumDependencies &dependencies
        );

        void BuildResidual(mfem::Vector &residual) const;

        void ApplyDensityJacobianAction(
            const mfem::Vector &densityVariation,
            mfem::Vector &action
        ) const;

        void ApplyDisplacementJacobianAction(
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        void ApplyAngularVelocityJacobianAction(
            double angularVelocityVariation,
            mfem::Vector &action
        ) const;

        void ApplyCompleteJacobianAction(
            const mfem::Vector &densityVariation,
            const mfem::Vector &displacementVariation,
            double angularVelocityVariation,
            mfem::Vector &action
        ) const;

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] double GetMomentOfInertia() const;
        [[nodiscard]] double GetAngularVelocity() const;
        [[nodiscard]] double GetCurrentAngularMomentum() const;
        [[nodiscard]] double GetTargetAngularMomentum() const noexcept;
        [[nodiscard]] physics::RigidRotation GetRotation() const;
        [[nodiscard]] AngularMomentumConstraintReport GetConstraintReport() const;
        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept;
        [[nodiscard]] std::uint64_t GetResidualApplicationCount() const noexcept;
        [[nodiscard]] const PreparedAngularMomentumActionStatistics &GetActionStatistics() const noexcept;
        [[nodiscard]] const models::CompiledFixedAngularMomentum &GetCompiledConstraint() const noexcept;

    private:
        struct QuadraturePointData final {
            mfem::IntegrationPoint integrationPoint;
            mfem::Vector densityShape;
            mapping::VolumeMappingContext mappingContext;
            double density{0.0};
            double cylindricalRadiusSquared{0.0};
        };

        struct ElementPAData final {
            int elementId{-1};
            mfem::Array<int> densityDofs;
            mfem::Array<int> displacementDofs;
            mfem::Array<int> compactificationDofs;
            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *displacementDofTransformation{nullptr};
            mfem::DofTransformation *compactificationDofTransformation{nullptr};
            mfem::Vector baseDisplacement;
            mfem::Vector compactification;
            std::vector<QuadraturePointData> quadraturePoints;
        };

        void BuildStaticPlan();
        void RefreshGeometry(const mfem::Vector &displacement);
        void RefreshDensity(const mfem::Vector &density);
        void AssembleResidual();
        void VerifyPrepared() const;

        [[nodiscard]] double EvaluateDensityMomentActionLocal(const mfem::Vector &densityVariation) const;
        [[nodiscard]] double EvaluateDisplacementMomentActionLocal(const mfem::Vector &displacementVariation) const;
        [[nodiscard]] double CylindricalRadiusSquared(const mfem::Vector &physicalPosition) const noexcept;
        [[nodiscard]] double CylindricalRadiusSquaredVariation(
            const mfem::Vector &physicalPosition,
            const mfem::Vector &physicalPositionVariation
        ) const noexcept;
        [[nodiscard]] double GlobalSum(double localValue) const;

        const fem::FEM &m_fem;
        const mapping::DomainMapper &m_domainMapper;
        const context::gravity_field::GravityFieldLinearizationContext &m_gravityContext;
        models::CompiledFixedAngularMomentum m_constraint;

        std::vector<ElementPAData> m_elements;
        AngularMomentumDependencies m_preparedDependencies;
        mfem::Vector m_cachedResidual;
        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;
        mutable mfem::Vector m_densityVariationLocal;
        mutable mfem::Vector m_displacementVariationLocal;
        mutable mfem::Vector m_elementDensityVariation;
        mutable mfem::Vector m_elementDisplacementVariation;

        double m_momentOfInertia{0.0};
        double m_angularVelocity{0.0};
        double m_currentAngularMomentum{0.0};
        std::uint64_t m_preparationCount{0};
        mutable std::uint64_t m_residualApplicationCount{0};
        mutable PreparedAngularMomentumActionStatistics m_actionStatistics;
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators