module;

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

#include <mfem.hpp>

export module mean_field:operators.prepared_rotational_displacement_force;

export import :fem;
export import :mapping.domain_mapper;
export import :operators.context.rotational_displacement_force;
export import :operators.kernels.rotational_displacement_force;
export import :physics.rigid_rotation;
export import :utils.blocks;

export namespace mean_field::operators {
    struct PreparedRotationalDisplacementForceReport final {
        context::rotational_displacement_force::RotationalDisplacementForcePreparationReport contextReport;

        bool updatedRotation{false};
        bool preparedResidual{false};

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

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

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

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

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

    class PreparedRotationalDisplacementForceOperator final {
    public:
        PreparedRotationalDisplacementForceOperator(
            const fem::FEM &f,
            const mapping::DomainMapper &domainMapper
        );

        PreparedRotationalDisplacementForceOperator(const PreparedRotationalDisplacementForceOperator &) = delete;

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

        PreparedRotationalDisplacementForceOperator(PreparedRotationalDisplacementForceOperator &&)            = delete;

        PreparedRotationalDisplacementForceOperator &operator=(PreparedRotationalDisplacementForceOperator &&) = delete;

        PreparedRotationalDisplacementForceReport Prepare(
            const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
            const context::rotational_displacement_force::RotationalDisplacementForceDependencies &dependencies,
            const physics::RigidRotation &rotation
        );

        [[nodiscard]] std::expected<
            PreparedRotationalDisplacementForceReport,
            kernels::RotationalDisplacementForceRejection>
        TryPrepare(
            const context::rotational_displacement_force::RotationalDisplacementForceStateView &state,
            const context::rotational_displacement_force::RotationalDisplacementForceDependencies &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 ApplyCompleteJacobianAction(
            const mfem::Vector &densityVariation,
            const mfem::Vector &displacementVariation,
            mfem::Vector &action
        ) const;

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

        [[nodiscard]] const context::rotational_displacement_force::RotationalDisplacementForcePreparationStatistics &
        GetContextPreparationStatistics() const noexcept;

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

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

        [[nodiscard]] const PreparedRotationalDisplacementForceColumnStatistics &
        GetDensityJacobianStatistics() const noexcept;

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

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

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

        [[nodiscard]] const context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext &
        GetContext() const noexcept;

    private:
        void VerifyPrepared() const;
        [[nodiscard]] std::expected<
            void,
            kernels::RotationalDisplacementForceRejection>
        TryPrepareElementData();
        void ApplyPreparedCompleteJacobianActionTrue(
            const mfem::Vector &densityVariationTrue,
            const mfem::Vector &displacementVariationTrue,
            mfem::Vector &actionTrue
        ) const;

        struct ElementPAData {
            int elementId{-1};
            mfem::Array<int> densityDofs;
            mfem::Array<int> displacementDofs;
            mfem::DofTransformation *densityDofTransformation{nullptr};
            mfem::DofTransformation *displacementDofTransformation{nullptr};
            const mfem::IntegrationRule *integrationRule{nullptr};
            mfem::DenseMatrix inverseElementJacobians;
            mfem::DenseMatrix centrifugalAccelerations;
            mfem::Vector baseDensityValues;
            mfem::Vector quadratureWeights;
        };

        const fem::FEM &m_fem;
        const mapping::DomainMapper &m_domainMapper;

        context::rotational_displacement_force::RotationalDisplacementForceLinearizationContext m_context;

        std::optional<physics::RigidRotation> m_rotation;
        mfem::Vector m_cachedResidual;
        std::vector<ElementPAData> m_elements;
        mutable mfem::Vector m_densityVariationTrue;
        mutable mfem::Vector m_displacementVariationTrue;
        mutable mfem::Vector m_actionTrue;
        mutable mfem::Vector m_densityVariationLocal;
        mutable mfem::Vector m_displacementVariationLocal;
        mutable mfem::Vector m_localAction;
        mutable mfem::Vector m_elementDensityVariation;
        mutable mfem::Vector m_elementDisplacementVariation;
        mutable mfem::Vector m_elementAction;
        mutable mfem::Vector m_densityShape;
        mutable mfem::Vector m_displacementShape;
        mutable mfem::Vector m_physicalPositionVariation;
        mutable mfem::Vector m_centrifugalAcceleration;
        mutable mfem::Vector m_centrifugalAccelerationVariation;
        mutable mfem::Vector m_weightedForce;
        mutable mfem::DenseMatrix m_referenceDisplacementDShape;
        mutable mfem::DenseMatrix m_referenceDisplacementJacobian;

        context::rotational_displacement_force::RotationalDisplacementForceDependencies m_preparedDependencies;

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

        mutable PreparedRotationalDisplacementForceColumnStatistics m_densityJacobianStatistics;

        mutable PreparedRotationalDisplacementForceColumnStatistics m_displacementJacobianStatistics;

        mutable PreparedRotationalDisplacementForceCompleteStatistics m_completeJacobianStatistics;

        bool m_isPrepared{false};
    };

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

    class PreparedRotationalDisplacementForceJacobianOperator final : public mfem::Operator {
    public:
        PreparedRotationalDisplacementForceJacobianOperator(
            const RotationalDisplacementForceLayout &layout,
            const PreparedRotationalDisplacementForceOperator &preparedOperator
        );

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

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

    private:
        RotationalDisplacementForceLayout m_layout;
        const PreparedRotationalDisplacementForceOperator &m_preparedOperator;
    };
} // namespace mean_field::operators