module;

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

export module mean_field:operators.prepared_barotropic_closure;

export import :eos.polytrope;
export import :fem;
export import :field.mfem;
export import :mapping.domain_mapper;
export import :operators.context.barotropic_closure_linearization;

export namespace mean_field::operators {
    struct PreparedBarotropicClosureReport final {
        context::barotropic::BarotropicClosurePreparationReport contextReport;
        bool preparedElementData{false};

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

    enum class BarotropicClosurePreparationRejectionReason : std::uint8_t {
        mapping_failure,
        invalid_quadrature_data,
        equation_of_state
    };

    /*
     * A rejected candidate is part of the nonlinear-solver control flow, not
     * an exceptional API failure.  Keep the payload fixed-size so it can be
     * selected deterministically across ranks without allocating.  Only the
     * detail associated with `reason` is meaningful.
     */
    struct BarotropicClosurePreparationRejection final {
        BarotropicClosurePreparationRejectionReason reason{
            BarotropicClosurePreparationRejectionReason::mapping_failure
        };
        mapping::MappingStatus mappingStatus{mapping::MappingStatus::non_finite_result};
        eos::EvaluationErrorCode equationOfStateError{eos::EvaluationErrorCode::nonfinite_result};
    };

    using BarotropicClosurePreparationResult =
        std::expected<PreparedBarotropicClosureReport, BarotropicClosurePreparationRejection>;

    class PreparedBarotropicClosureOperator final : public mfem::Operator {
    public:
        PreparedBarotropicClosureOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper,
            const eos::Polytrope &equationOfState
        );

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

        PreparedBarotropicClosureReport Prepare(
            const context::barotropic::BarotropicClosureStateView &state,
            const context::barotropic::BarotropicClosureDependencies &dependencies
        );

        [[nodiscard]] BarotropicClosurePreparationResult TryPrepare(
            const context::barotropic::BarotropicClosureStateView &state,
            const context::barotropic::BarotropicClosureDependencies &dependencies
        );

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

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

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

        // Exact diagonal of the prepared density-to-closure block, expressed
        // in the reduced density coordinates used by the root operator.
        void AssembleDensityJacobianDiagonal(mfem::Vector &diagonal) const;

        [[nodiscard]] bool IsPrepared() const noexcept;
        [[nodiscard]] std::uint64_t GetPreparationCount() const noexcept;
        [[nodiscard]] int GetDensitySize() const noexcept;
        [[nodiscard]] int GetEnthalpySize() const noexcept;
        [[nodiscard]] int GetDisplacementSize() const noexcept;

        [[nodiscard]] const context::barotropic::BarotropicClosureLinearizationContext &GetContext() const noexcept;
        [[nodiscard]] const context::barotropic::BarotropicClosurePreparationStatistics &
        GetContextPreparationStatistics() const noexcept;

    private:
        struct ConstructionData;

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

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

        void VerifyPrepared() const;

        void ApplyThermodynamicActionFull(
            const mfem::Vector &densityVariationTrue,
            const mfem::Vector &enthalpyVariationTrue,
            mfem::Vector &actionTrue
        ) const;

        void ApplyDisplacementActionFull(
            const mfem::Vector &displacementVariationTrue,
            mfem::Vector &actionTrue
        ) const;

        struct ElementPAData {
            int elementId{-1};
            mfem::Array<int> densityDofs;
            mfem::Array<int> enthalpyDofs;
            mfem::Array<int> displacementDofs;

            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *enthalpyDofTransformation{nullptr};
            mfem::DofTransformation *displacementDofTransformation{nullptr};

            mfem::DenseMatrix densityBasis;
            mfem::DenseMatrix enthalpyBasis;
            mfem::DenseMatrix inverseElementJacobians;

            mfem::Vector weightedResidual;
            mfem::Vector quadratureWeights;
            mfem::Vector weightedEnthalpyDerivative;
        };

        const fem::FEM &m_fem;
        const mapping::DomainMapper &m_domainMapper;
        const eos::Polytrope &m_equationOfState;

        field::FieldDofMap m_densityMap;
        field::FieldDofMap m_enthalpyMap;
        field::FieldDofMap m_displacementMap;

        context::barotropic::BarotropicClosureLinearizationContext m_context;

        std::vector<ElementPAData> m_elements;

        mfem::Vector m_baseDensityTrue;
        mfem::Vector m_baseEnthalpyTrue;
        mfem::Vector m_baseDisplacementTrue;

        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_enthalpyVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;
        mutable mfem::Vector m_fullThermodynamicAction;
        mutable mfem::Vector m_fullDisplacementAction;
        mutable mfem::Vector m_fullResidual;

        mutable mfem::Vector m_displacementVariationLocal;
        mutable mfem::Vector m_localDisplacementAction;
        mutable mfem::Vector m_elementDisplacementVariation;
        mutable mfem::Vector m_quadratureDisplacementAction;
        mutable mfem::Vector m_elementDisplacementAction;
        mutable mfem::DenseMatrix m_referenceDShape;
        mutable mfem::DenseMatrix m_referenceDisplacementJacobian;

        std::uint64_t m_preparationCount{0};
        bool m_isPrepared{false};
    };
} // namespace mean_field::operators