module;

#include <compare>
#include <cstddef>
#include <cstdint>
#include <expected>
#include <memory>
#include <optional>
#include <vector>

#include <mfem.hpp>

export module mean_field:operators.prepared_pressure_force;

export import :eos.polytrope;
export import :fem;
export import :fem.reference_tables;
export import :field.mfem;
export import :mapping.domain_mapper;
export import :mapping.prepared_cache;
export import :operators.context.pressure_force;
export import :utils.blocks;

export namespace mean_field::operators {
    enum class PressureForcePreparationRejectionReason : std::uint8_t {
        equation_of_state,
        invalid_mapping,
        non_finite_arithmetic
    };

    struct PressureForcePreparationRejection final {
        PressureForcePreparationRejectionReason reason{PressureForcePreparationRejectionReason::equation_of_state};
        eos::EvaluationErrorCode equationOfStateCode{eos::EvaluationErrorCode::nonfinite_result};
        mapping::MappingStatus mappingStatus{mapping::MappingStatus::valid};
    };

    struct PreparedPressureForceReport final {
        context::pressure_force::PressureForcePreparationReport contextReport;

        bool preparedEnthalpyJacobianData{false};
        bool preparedDisplacementJacobianData{false};
        bool preparedResidual{false};

        [[nodiscard]]
        bool DidAnyWork() const noexcept {
            return contextReport.DidAnyWork() || preparedEnthalpyJacobianData || preparedDisplacementJacobianData ||
                   preparedResidual;
        }
    };

    struct PreparedPressureForceEnthalpyJacobianStatistics final {
        std::uint64_t preparations{0};
        std::uint64_t applications{0};

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

    struct PreparedPressureForceDisplacementJacobianStatistics final {
        std::uint64_t preparations{0};
        std::uint64_t applications{0};

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

    struct PreparedPressureForceCompleteJacobianStatistics final {
        std::uint64_t applications{0};

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

    /*
     * Prepared pressure contribution
     *
     *   R_d^P(w)
     *       =
     *       - integral_{Omega_star(d)}
     *           P(h) div(w) dV.
     *
     * Public state and Jacobian directions are expressed in FieldDof
     * coordinates.
     *
     * Current registry:
     *
     *   h   -> Stellar -> reduced
     *   d   -> All     -> identity/full
     *   R_d -> All     -> identity/full
     *
     * Full MFEM true/local vectors are private implementation details.
     */
    class PreparedPressureForceOperator final {
    public:
        PreparedPressureForceOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState
        );

        PreparedPressureForceOperator(const PreparedPressureForceOperator &)            = delete;

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

        PreparedPressureForceOperator(PreparedPressureForceOperator &&)                 = delete;

        PreparedPressureForceOperator &operator=(PreparedPressureForceOperator &&)      = delete;

        PreparedPressureForceReport Prepare(
            const context::pressure_force::PressureForceStateView &state,
            const context::pressure_force::PressureForceDependencies &dependencies
        );

        [[nodiscard]] std::expected<
            PreparedPressureForceReport,
            PressureForcePreparationRejection>
        TryPrepare(
            const context::pressure_force::PressureForceStateView &state,
            const context::pressure_force::PressureForceDependencies &dependencies
        );

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

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

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

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

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

        [[nodiscard]]
        int GetEnthalpySize() const noexcept;

        [[nodiscard]]
        int GetDisplacementSize() const noexcept;

        [[nodiscard]]
        const context::pressure_force::PressureForceLinearizationContext &GetContext() const noexcept;

        [[nodiscard]]
        const context::pressure_force::PressureForcePreparationStatistics &
        GetContextPreparationStatistics() const noexcept;

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

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

        [[nodiscard]]
        const PreparedPressureForceEnthalpyJacobianStatistics &GetEnthalpyJacobianStatistics() const noexcept;

        [[nodiscard]]
        const PreparedPressureForceDisplacementJacobianStatistics &GetDisplacementJacobianStatistics() const noexcept;

        [[nodiscard]]
        const PreparedPressureForceCompleteJacobianStatistics &GetCompleteJacobianStatistics() const noexcept;

        [[nodiscard]]
        std::size_t GetStellarElementCount() const noexcept;

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

        template <typename Visitor> void VisitMappedGeometryRules(Visitor &&visitor) const {
            for (const ElementPAData &data : m_elements) {
                visitor(data.elementId, *data.integrationRule);
            }
        }

    private:
        struct ConstructionData;

        [[nodiscard]]
        static ConstructionData MakeConstructionData(const fem::FEM &f);

        PreparedPressureForceOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState,
            ConstructionData constructionData
        );

        struct ElementPAData final {
            int elementId{-1};

            mfem::Array<int> enthalpyDofs;
            mfem::Array<int> displacementDofs;
            mfem::Array<int> compactificationDofs;

            mfem::DofTransformation *enthalpyDofTransformation{nullptr};

            mfem::DofTransformation *displacementDofTransformation{nullptr};

            mfem::DofTransformation *compactificationDofTransformation{nullptr};

            const mfem::IntegrationRule *integrationRule{nullptr};

            // Immutable reference values and gradients are shared by FE/rule.
            std::shared_ptr<const fem::ScalarReferenceTable> enthalpyReferenceTable;
            std::shared_ptr<const fem::ScalarReferenceTable> displacementReferenceTable;

            [[nodiscard]] const mfem::DenseMatrix &GetEnthalpyBasis() const {
                return enthalpyReferenceTable->GetValues();
            }

            /*
             * Each entry is:
             *
             *   scalar displacement DOF
             *       x
             *   physical dimension.
             */
            std::vector<mfem::DenseMatrix> physicalTestGradients;

            mapping::VolumeMappingCache baseMappingContexts;

            std::optional<mapping::ElementDisplacementData> baseDisplacementData;

            std::optional<mapping::ElementCompactificationData> compactificationData;

            mfem::Vector quadratureWeights;
            mfem::Vector pressure;
            mfem::Vector pressureDerivative;

            mfem::Vector elementResidual;
            mfem::DenseMatrix enthalpyJacobian;
        };

        void PrepareStaticPlan();
        [[nodiscard]] std::optional<PressureForcePreparationRejection> PrepareGeometry();
        [[nodiscard]] std::optional<PressureForcePreparationRejection> PrepareMaterialState();

        void FinalizeDisplacementJacobianPreparation();

        [[nodiscard]] std::optional<PressureForcePreparationRejection> AssembleCachedResidual();

        void VerifyPrepared() const;

        const fem::FEM &m_fem;

        const mapping::DomainMapper &m_domainMapper;

        const eos::Polytrope &m_equationOfState;

        field::FieldDofMap m_enthalpyMap;

        field::FieldDofMap m_displacementMap;

        context::pressure_force::PressureForceLinearizationContext m_context;

        std::vector<ElementPAData> m_elements;

        /*
         * Canonical full-MFEM expansion of the frozen FieldDof state.
         */
        mfem::Vector m_baseEnthalpyTrue;
        mfem::Vector m_baseDisplacementTrue;

        /*
         * Reusable Krylov work storage.
         */
        mutable mfem::Vector m_enthalpyVariationTrue;

        mutable mfem::Vector m_displacementVariationTrue;

        mutable mfem::Vector m_fullDisplacementAction;

        /*
         * Solver-facing cached residual in Displacement FieldDof
         * coordinates.
         */
        mfem::Vector m_cachedResidual;

        std::uint64_t m_residualPreparationCount{0};

        mutable std::uint64_t m_residualApplicationCount{0};

        mutable PreparedPressureForceEnthalpyJacobianStatistics m_enthalpyJacobianStatistics;

        mutable PreparedPressureForceDisplacementJacobianStatistics m_displacementJacobianStatistics;

        mutable PreparedPressureForceCompleteJacobianStatistics m_completeJacobianStatistics;

        bool m_isPrepared{false};
    };

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

    /*
     * Coupled-layout adapter around the modern prepared pressure-force
     * Jacobian.
     *
     * Reads:
     *
     *   delta d
     *   delta h
     *
     * Writes:
     *
     *   R_d
     *
     * It deliberately imposes no raw-FES-size assumptions on unrelated
     * coupled blocks.
     */
    class PreparedPressureForceJacobianOperator final : public mfem::Operator {
    public:
        PreparedPressureForceJacobianOperator(
            const BarotropicEquilibriumLayout &layout,
            const PreparedPressureForceOperator &preparedOperator
        );

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

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

    private:
        BarotropicEquilibriumLayout m_layout;

        const PreparedPressureForceOperator &m_preparedOperator;
    };
} // namespace mean_field::operators