module;

#include "mean_field.h"

export module mean_field:mapping.domain_mapper;

export namespace mean_field::mapping {
    class DomainMapper {
    public:
        struct VolumeQuadratureContext {
            mfem::DenseMatrix J_inv;
            double detJ;
            double weight;
        };

        struct FaceQuadratureContext {
            mfem::Vector normal;
            double ds;
            double v_dot_n_scale;
        };

    public:
        explicit DomainMapper(const double r_star_ref, const double r_inf_ref);

        explicit DomainMapper(const mfem::GridFunction &d, const double r_star_ref, const double r_inf_ref);

        [[nodiscard]] bool is_vacuum(const mfem::ElementTransformation &T) const;

        void SetDisplacement(const mfem::GridFunction &d);

        [[nodiscard]] bool IsIdentity() const;

        void ResetDisplacement();

        void ComputeJacobian(mfem::ElementTransformation &T, mfem::DenseMatrix &J) const;

        double ComputeDetJ(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) const;

        void ComputeMappedDiffusionTensor(mfem::ElementTransformation &T, mfem::DenseMatrix &D) const;

        void ComputeInverseJacobian(mfem::ElementTransformation &T, mfem::DenseMatrix &JInv) const;

        VolumeQuadratureContext GetQuadratureContext(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip) const;

        FaceQuadratureContext GetFaceQuadratureContext(mfem::FaceElementTransformations &T, const mfem::IntegrationPoint &ip) const;

        void GetPhysicalPoint(mfem::ElementTransformation &T, const mfem::IntegrationPoint &ip, mfem::Vector &x_phys) const;

        void GetVectorValue(const int i, const mfem::IntegrationPoint &ip, mfem::Vector &val) const;

        [[nodiscard]] const mfem::GridFunction *GetDisplacement() const;

        [[nodiscard]] double GetPhysInfRadius() const;

        [[nodiscard]] size_t GetCacheHits() const;

        [[nodiscard]] size_t GetCacheMisses() const;

        [[nodiscard]] double GetCacheHitRate() const;

        void ResetCacheStats() const;

    private:
        void InitAllScratchSpaces() const;

        void ApplyKelvinMapping(const mfem::Vector &x_ref, mfem::Vector &x_phys) const;

        void ComputeKelvinJacobian(const mfem::Vector &x_ref, const mfem::Vector &x_disp, const mfem::DenseMatrix &J_D, mfem::DenseMatrix &J) const;

        void InvalidateCache() const;

        void UpdateElementCache(const mfem::ElementTransformation &T) const;

    private:
        const mfem::GridFunction *m_d;
        std::unique_ptr<mfem::GridFunction> m_internal_d;
        const int m_dim{3};
        const int m_vacuum_attr{3};
        const double m_r_star_ref{1.0};
        const double m_r_inf_ref{2.0};
        const double m_xi_clamp{0.9999};

        mutable int m_cached_elem_id{-1};
        mutable int m_cached_elem_type{mfem::ElementTransformation::ELEMENT};
        mutable const mfem::FiniteElement *m_fe{nullptr};

        mutable mfem::Vector m_elem_dofs;
        mutable mfem::DenseMatrix m_dof_mat;
        mutable mfem::DenseMatrix m_dshape;
        mutable mfem::Vector m_shape;

        mutable size_t m_cache_hits{0};
        mutable size_t m_cache_misses{0};

        mutable mfem::DenseMatrix m_J_D;
        mutable mfem::DenseMatrix m_J_temp;
        mutable mfem::DenseMatrix m_JInv_temp;
        mutable mfem::Vector m_x_ref;
        mutable mfem::Vector m_x_disp;
        mutable mfem::Vector m_d_val;
    };

}