module;

#include "mean_field.h"

export module mean_field:mapping.domain_mapper;
export import :mapping.types;
import :mapping.compactification;
import :utils.user;

export namespace mean_field::mapping {
    enum class FaceElementSide : uint8_t { element_1, element_2 };

    class ElementDisplacementData {
    public:
        ElementDisplacementData(
            const mfem::FiniteElement &element,
            const mfem::Vector &displacement_dofs,
            mfem::Ordering::Type ordering = mfem::Ordering::byNODES
        );

        [[nodiscard]] const mfem::FiniteElement &GetElement() const noexcept;
        [[nodiscard]] const mfem::DenseMatrix &GetDofMatrix() const noexcept;
        [[nodiscard]] int GetDimension() const noexcept;
        [[nodiscard]] int GetDofCount() const noexcept;
        [[nodiscard]] mfem::Ordering::Type GetOrdering() const noexcept;

    private:
        const mfem::FiniteElement *m_element;
        mfem::DenseMatrix m_dof_matrix;
        int m_dimension;
        mfem::Ordering::Type m_ordering;
    };

    struct CompactificationPointData {
        double coordinate{0.0};
        mfem::Vector coordinate_gradient;
    };

    [[nodiscard]] ElementDisplacementData ElementDisplacementDataFromElementVDofs(
        const mfem::FiniteElement &element,
        const mfem::Vector &displacement_dofs
    );

    class ElementCompactificationData {
    public:
        ElementCompactificationData(
            const mfem::FiniteElement &element,
            const mfem::Vector &dofs
        );

        [[nodiscard]] const mfem::FiniteElement &GetElement() const noexcept;
        [[nodiscard]] const mfem::Vector &GetDofs() const noexcept;
        [[nodiscard]] int GetDofCount() const noexcept;

    private:
        const mfem::FiniteElement *m_element;
        mfem::Vector m_dofs;
    };

    struct ElementMappingData {
        const ElementDisplacementData &displacement;
        const ElementCompactificationData &compactification;
    };

    class DomainMapperStateless {
    public:
        class Workspace {
        public:
            explicit Workspace(int dimension = 3);

            void SetDimension(int dimension);

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

        private:
            friend class DomainMapperStateless;

            int m_dimension;

            mfem::Vector m_shape;
            mfem::DenseMatrix m_mesh_dshape;
            mfem::Vector m_field_value;
            mfem::DenseMatrix m_field_jacobian;

            mfem::Vector m_compactification_shape;
            mfem::DenseMatrix m_compactification_dshape;
            CompactificationPointData m_compactification_point;

            mfem::Vector m_reference_normal;
            mfem::Vector m_mapped_normal;
            mfem::DenseMatrix m_full_element_jacobian;

            mfem::Vector m_vector_temp;
            mfem::DenseMatrix m_matrix_temp_1;
            mfem::DenseMatrix m_matrix_temp_2;

            compactification::ExteriorMapResult m_exterior_result;
            compactification::ExteriorMapVariation m_exterior_variation;
        };

    public:
        DomainMapperStateless(
            utils::DomainMapperStatelessOptions options,
            std::unique_ptr<const compactification::ExteriorDomainMap> exterior_map
        );

        DomainMapperStateless(const DomainMapperStateless &)            = delete;
        DomainMapperStateless &operator=(const DomainMapperStateless &) = delete;
        DomainMapperStateless(DomainMapperStateless &&)                 = default;
        DomainMapperStateless &operator=(DomainMapperStateless &&)      = default;

        [[nodiscard]] MappingStatus EvaluatePoint(
            const ElementMappingData &element_data,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            Workspace &workspace,
            MappingPointContext &context
        ) const;

        [[nodiscard]] MappingStatus EvaluateVolume(
            const ElementMappingData &element_data,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            Workspace &workspace,
            VolumeMappingContext &context
        ) const;

        [[nodiscard]] MappingStatus EvaluateFace(
            const ElementMappingData &element_data,
            mfem::FaceElementTransformations &transformation,
            FaceElementSide side,
            const mfem::IntegrationPoint &integration_point,
            Workspace &workspace,
            FaceMappingContext &context
        ) const;

        [[nodiscard]] MappingStatus EvaluatePointVariation(
            const ElementMappingData &element_data,
            const ElementDisplacementData &direction,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            const MappingPointContext &base_context,
            Workspace &workspace,
            MappingPointVariation &variation
        ) const;

        [[nodiscard]] MappingStatus EvaluateVolumeVariation(
            const ElementMappingData &element_data,
            const ElementDisplacementData &direction,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            const VolumeMappingContext &base_context,
            Workspace &workspace,
            VolumeMappingVariation &variation
        ) const;

        [[nodiscard]] MappingStatus EvaluateFaceVariation(
            const ElementMappingData &element_data,
            const ElementDisplacementData &direction,
            mfem::FaceElementTransformations &transformation,
            FaceElementSide side,
            const mfem::IntegrationPoint &integration_point,
            const FaceMappingContext &base_context,
            Workspace &workspace,
            FaceMappingVariation &variation
        ) const;

        [[nodiscard]] bool IsCompactifiedElement(const mfem::ElementTransformation &transformation) const noexcept;
        [[nodiscard]] int GetDimension() const noexcept;
        [[nodiscard]] int GetVacuumElementAttribute() const noexcept;
        [[nodiscard]] const compactification::ExteriorDomainMap &GetExteriorMap() const noexcept;

    private:
        void ValidateElementData(const ElementMappingData &element_data) const;

        void EvaluateField(
            const ElementDisplacementData &field,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            Workspace &workspace,
            mfem::Vector &value,
            mfem::DenseMatrix &jacobian
        ) const;

        [[nodiscard]] MappingStatus EvaluateCompactificationCoordinate(
            const ElementCompactificationData &compactification,
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            Workspace &workspace,
            CompactificationPointData &point_data
        ) const;

        [[nodiscard]] static mfem::ElementTransformation &SelectFaceElementTransformation(
            mfem::FaceElementTransformations &transformation,
            FaceElementSide side
        );

        [[nodiscard]] static const mfem::IntegrationPoint &SelectFaceElementIntegrationPoint(
            mfem::FaceElementTransformations &transformation,
            FaceElementSide side
        );

        utils::DomainMapperStatelessOptions m_options;
        std::unique_ptr<const compactification::ExteriorDomainMap> m_exterior_map;
    };
    class DomainMapper {

    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 HasCompactification() const noexcept;
        [[nodiscard]] bool HasDisplacementField() const noexcept;
        [[nodiscard]] bool CalcIsIdentity() 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;

        void MapHDivFluxToPhysical(
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            const mfem::Vector &reference_flux,
            mfem::Vector &physical_flux
        ) const;

        void MapPhysicalFluxToHDivReference(
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            const mfem::Vector &physical_flux,
            mfem::Vector &reference_flux
        ) const;

        void MapReferenceGradientToPhysical(
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            const mfem::Vector &reference_gradient,
            mfem::Vector &physical_gradient
        ) 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;

        bool m_displacement_is_identity{true};
    };

} // namespace mean_field::mapping