module;

#include <compare>
#include <cstdint>
#include <expected>
#include <mfem.hpp>
#include <optional>
#include <vector>

export module mean_field:operators.prepared_mass_normalization;

export import :fem;
export import :mapping.domain_mapper;
export import :model.compiled_fixed_mass;
export import :operators.context.gravity_field;
export import :operators.prepared_constraint;
export import :utils.blocks;

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

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

    struct MassNormalizationDependencies final {
        MassNormalizationDependencyStamp discretization;
        MassNormalizationDependencyStamp density;
        MassNormalizationDependencyStamp displacement;
        MassNormalizationDependencyStamp targetMass;

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

    struct MassNormalizationStateView final {
        double targetMass{0.0};
    };

    struct FixedMassJacobianInput final {
        const mfem::Vector &densityVariation;
        const mfem::Vector &displacementVariation;
    };

    struct FixedMassJacobianTransposeOutput final {
        mfem::Vector &densityDual;
        mfem::Vector &displacementDual;
    };

    struct PreparedMassNormalizationReport final {
        bool rebuiltStaticPlan{false};
        bool refreshedGeometry{false};
        bool refreshedDensity{false};
        bool updatedTargetMass{false};
        bool assembledResidual{false};

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

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

    enum class MassNormalizationPreparationRejectionReason : std::uint8_t {
        mapping_failure,
        non_finite_density_interpolation,
        non_finite_assembled_mass
    };

    /*
     * Candidate rejection is deliberately represented without text or owned
     * storage.  That makes the result cheap to propagate through a line search
     * and gives the MPI implementation a deterministic, allocation-free value
     * to select on every rank.
     */
    struct MassNormalizationPreparationRejection final {
        MassNormalizationPreparationRejectionReason reason{
            MassNormalizationPreparationRejectionReason::mapping_failure
        };
        mapping::MappingStatus mappingStatus{mapping::MappingStatus::non_finite_result};
    };

    using MassNormalizationPreparationResult =
        std::expected<PreparedMassNormalizationReport, MassNormalizationPreparationRejection>;

    struct PreparedMassNormalizationActionStatistics final {
        std::uint64_t densityApplications{0};
        std::uint64_t displacementApplications{0};
        std::uint64_t completeApplications{0};
        std::uint64_t transposeApplications{0};

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

    /*
     * Prepared scalar row
     *
     *   R_M(rho, d) = integral_{Omega_star(d)} rho dV - M_target.
     *
     * Density and displacement are borrowed from the shared gravity-field
     * linearization context. This keeps the mass row on exactly the same
     * frozen state and geometry as the gravity and mechanical rows.
     * Jacobian directions use the corresponding registered FieldDof
     * coordinates; expansion to MFEM true-DOF vectors is private.
     */
    class PreparedMassNormalizationOperator final {
    public:
        using SpecificationType       = models::FixedTotalMass;
        using CompiledConstraintType  = models::CompiledFixedMass;
        using Dependencies            = MassNormalizationDependencies;
        using Report                  = PreparedMassNormalizationReport;
        using JacobianInput           = FixedMassJacobianInput;
        using JacobianTransposeOutput = FixedMassJacobianTransposeOutput;

        PreparedMassNormalizationOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const context::gravity_field::GravityFieldLinearizationContext &gravityContext
        );

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

        PreparedMassNormalizationReport Prepare(
            const MassNormalizationStateView &state,
            const MassNormalizationDependencies &dependencies
        );

        PreparedMassNormalizationReport Prepare(
            const models::CompiledFixedMass &constraint,
            const MassNormalizationDependencies &dependencies
        );

        [[nodiscard]] MassNormalizationPreparationResult TryPrepare(
            const MassNormalizationStateView &state,
            const MassNormalizationDependencies &dependencies
        );

        [[nodiscard]] MassNormalizationPreparationResult TryPrepare(
            const models::CompiledFixedMass &constraint,
            const MassNormalizationDependencies &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 ApplyCompleteJacobianAction(
            const mfem::Vector &densityVariation,
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

        void ApplyJacobian(
            const FixedMassJacobianInput &input,
            mfem::Vector &action
        ) const;

        void ApplyCompleteJacobianTransposeAction(
            double residualDual,
            mfem::Vector &densityDual,
            mfem::Vector &displacementDual
        ) const;

        void ApplyJacobianTranspose(
            const mfem::Vector &residualDual,
            FixedMassJacobianTransposeOutput output
        ) const;

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] double GetCurrentMass() const;
        [[nodiscard]] double GetTargetMass() const;
        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept;
        [[nodiscard]] std::uint64_t GetResidualApplicationCount() const noexcept;
        [[nodiscard]] const PreparedMassNormalizationActionStatistics &GetActionStatistics() const noexcept;
        [[nodiscard]] const fem::FEM &GetFEM() const noexcept;
        [[nodiscard]] const context::gravity_field::GravityFieldLinearizationContext &
        GetGravityContext() const noexcept;

    private:
        struct QuadraturePointData final {
            mfem::IntegrationPoint integrationPoint;
            mfem::Vector densityShape;
            mapping::VolumeMappingContext mappingContext;
            double density{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();
        [[nodiscard]] std::optional<MassNormalizationPreparationRejection>
        RefreshGeometry(const mfem::Vector &displacement);
        [[nodiscard]] std::optional<MassNormalizationPreparationRejection> RefreshDensity(const mfem::Vector &density);
        [[nodiscard]] std::optional<MassNormalizationPreparationRejection> AssembleResidual();
        [[nodiscard]] std::optional<MassNormalizationPreparationRejection> UpdateResidualForTargetMass();
        void VerifyPrepared() const;

        [[nodiscard]] double EvaluateDensityActionLocal(const mfem::Vector &densityVariation) const;

        [[nodiscard]] double EvaluateDisplacementActionLocal(const mfem::Vector &displacementVariation) const;

        void AssembleDensityTransposeAction(
            double residualDual,
            mfem::Vector &densityDual
        ) const;

        void AssembleDisplacementTransposeAction(
            double residualDual,
            mfem::Vector &displacementDual
        ) const;

        [[nodiscard]] double GlobalSum(double localValue) const;

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

        std::vector<ElementPAData> m_elements;
        MassNormalizationDependencies m_preparedDependencies;
        mfem::Vector m_cachedResidual;

        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;

        double m_currentMass{0.0};
        double m_targetMass{0.0};

        std::uint64_t m_preparationCount{0};
        mutable std::uint64_t m_residualApplicationCount{0};
        mutable PreparedMassNormalizationActionStatistics m_actionStatistics;
        bool m_isPrepared{false};
    };

    using PreparedFixedMass = PreparedMassNormalizationOperator;

    static_assert(PreparedConstraint<PreparedFixedMass>);

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

    class PreparedMassNormalizationJacobianOperator final : public mfem::Operator {
    public:
        PreparedMassNormalizationJacobianOperator(
            const MassNormalizationLayout &layout,
            const PreparedMassNormalizationOperator &preparedOperator
        );

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

        void MultTranspose(
            const mfem::Vector &residualDual,
            mfem::Vector &stateDual
        ) const override;

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

    private:
        MassNormalizationLayout m_layout;
        const PreparedMassNormalizationOperator &m_preparedOperator;
    };
} // namespace mean_field::operators