module;

#include <compare>
#include <cstdint>

#include <mfem.hpp>

export module mean_field:operators.prepared_displacement_residual;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.context.gravity_field;
export import :operators.prepared_gravity_displacement_force;
export import :operators.prepared_pressure_force;
export import :operators.prepared_rotational_displacement_force;
export import :eos.polytrope;
export import :physics.rigid_rotation;
export import :utils.blocks;

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

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

    struct DisplacementResidualDependencies final {
        DisplacementResidualDependencyStamp discretization;
        DisplacementResidualDependencyStamp density;
        DisplacementResidualDependencyStamp displacement;
        DisplacementResidualDependencyStamp gravityGradient;
        DisplacementResidualDependencyStamp enthalpy;
        DisplacementResidualDependencyStamp rotation;

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

    /*
     * Density, displacement, and gravity gradient deliberately do not appear
     * here. They are obtained from the shared, already-prepared gravity-field
     * linearization context so every mechanical-force contribution consumes
     * the same frozen density and geometry as the gravity equations.
     */
    struct DisplacementResidualStateView final {
        const mfem::Vector &enthalpy;
    };

    struct PreparedDisplacementResidualReport final {
        PreparedPressureForceReport pressure;
        PreparedGravityDisplacementForceReport gravity;
        PreparedRotationalDisplacementForceReport rotation;

        bool assembledResidual{false};

        [[nodiscard]] bool DidAnyChildWork() const noexcept {
            return pressure.DidAnyWork() || gravity.DidAnyWork() || rotation.DidAnyWork();
        }

        [[nodiscard]] bool DidAnyWork() const noexcept {
            return DidAnyChildWork() || assembledResidual;
        }
    };

    struct PreparedDisplacementResidualActionStatistics final {
        std::uint64_t densityApplications{0};
        std::uint64_t displacementApplications{0};
        std::uint64_t gravityGradientApplications{0};
        std::uint64_t enthalpyApplications{0};
        std::uint64_t completeApplications{0};

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

    /*
     * Row-level composer for
     *
     *   R_d = R_d^pressure + R_d^gravity + R_d^rotation.
     *
     * This class owns the three prepared contributors and only orchestrates
     * their existing residual and Jacobian APIs. It contains no force kernel
     * and no independent copy of the gravity linearization context.
     */
    class PreparedDisplacementResidualOperator final {
    public:
        PreparedDisplacementResidualOperator(
            const fem::FEM &f,
            const mapping::DomainMapperStateless &domainMapper,
            const eos::Polytrope &barotrope,
            const context::gravity_field::GravityFieldLinearizationContext &gravityContext
        );

        PreparedDisplacementResidualOperator(const PreparedDisplacementResidualOperator &)            = delete;

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

        PreparedDisplacementResidualOperator(PreparedDisplacementResidualOperator &&)                 = delete;

        PreparedDisplacementResidualOperator &operator=(PreparedDisplacementResidualOperator &&)      = delete;

        PreparedDisplacementResidualReport Prepare(
            const DisplacementResidualStateView &state,
            const DisplacementResidualDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

        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 ApplyGravityGradientJacobianAction(
            const mfem::Vector &gravityGradientVariation,
            mfem::Vector &action
        ) const;

        void ApplyEnthalpyJacobianAction(
            const mfem::Vector &enthalpyVariation,
            mfem::Vector &action
        ) const;

        void ApplyCompleteJacobianAction(
            const mfem::Vector &densityVariation,
            const mfem::Vector &displacementVariation,
            const mfem::Vector &gravityGradientVariation,
            const mfem::Vector &enthalpyVariation,
            mfem::Vector &action
        ) const;

        [[nodiscard]] bool IsPrepared() const noexcept;

        [[nodiscard]] std::uint64_t GetResidualPreparationCount() const noexcept;

        [[nodiscard]] std::uint64_t GetResidualApplicationCount() const noexcept;

        [[nodiscard]] const PreparedDisplacementResidualActionStatistics &GetActionStatistics() const noexcept;

        [[nodiscard]] const PreparedPressureForceOperator &GetPressureOperator() const noexcept;

        [[nodiscard]] const PreparedGravityDisplacementForceOperator &GetGravityOperator() const noexcept;

        [[nodiscard]] const PreparedRotationalDisplacementForceOperator &GetRotationalOperator() const noexcept;

        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;

        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;

    private:
        void AssembleResidual();
        void VerifyPrepared() const;

        const fem::FEM &m_fem;
        const mapping::DomainMapperStateless &m_domainMapper;
        const context::gravity_field::GravityFieldLinearizationContext &m_gravityContext;

        PreparedPressureForceOperator m_pressureOperator;
        PreparedGravityDisplacementForceOperator m_gravityOperator;
        PreparedRotationalDisplacementForceOperator m_rotationalOperator;

        DisplacementResidualDependencies m_preparedDependencies;
        mfem::Vector m_cachedResidual;

        std::uint64_t m_residualPreparationCount{0};
        mutable std::uint64_t m_residualApplicationCount{0};
        mutable PreparedDisplacementResidualActionStatistics m_actionStatistics;

        bool m_isPrepared{false};
    };

    using DisplacementResidualLayout = utils::blocks::form_layout<utils::blocks::barotropic_equilibrium_form>;

    class PreparedDisplacementResidualJacobianOperator final : public mfem::Operator {
    public:
        PreparedDisplacementResidualJacobianOperator(
            const DisplacementResidualLayout &layout,
            const PreparedDisplacementResidualOperator &preparedOperator
        );

        void Mult(
            const mfem::Vector &direction,
            mfem::Vector &action
        ) const override;

        [[nodiscard]] const DisplacementResidualLayout &GetLayout() const noexcept;

    private:
        DisplacementResidualLayout m_layout;
        const PreparedDisplacementResidualOperator &m_preparedOperator;
    };
} // namespace mean_field::operators