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 DomainMapper {
    public:
        class Workspace {
        public:
            explicit Workspace(int dimension = 3);

            void SetDimension(int dimension);

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

        private:
            friend class DomainMapper;

            int m_dimension;

            mfem::Vector m_shape;
            mfem::DenseMatrix m_reference_dshape;
            mfem::DenseMatrix m_mesh_dshape;
            mfem::DenseMatrix m_reference_field_jacobian;
            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:
        DomainMapper(
            utils::DomainMapperOptions options,
            std::unique_ptr<const compactification::ExteriorDomainMap> exterior_map
        );

        DomainMapper(const DomainMapper &)            = delete;
        DomainMapper &operator=(const DomainMapper &) = delete;
        DomainMapper(DomainMapper &&)                 = default;
        DomainMapper &operator=(DomainMapper &&)      = 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]] 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 mfem::DenseMatrix *inverse_mesh_jacobian
        ) const;

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

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

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

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

        utils::DomainMapperOptions m_options;
        std::unique_ptr<const compactification::ExteriorDomainMap> m_exterior_map;
    };

    class GridFunctionMappingEvaluator {
    public:
        /*
         * The evaluator references the supplied grid functions and caches copies of
         * their element-local DOFs. Call InvalidateCache() or Refresh() after either
         * grid function's values are modified. Finite-element-space sequence changes
         * are detected automatically.
         *
         * This object owns mutable workspace and cache state and is not thread-safe.
         */
        GridFunctionMappingEvaluator(
            const DomainMapper &mapper,
            const mfem::GridFunction &displacement,
            const mfem::GridFunction &compactification_coordinate
        );

        /*
         * Discard all element-local field data. The next evaluation reloads its
         * requested element lazily. This operation is idempotent.
         */
        void InvalidateCache() noexcept;

        /*
         * Reload the currently cached element immediately. If no element has been
         * evaluated yet, Refresh() is a validated no-op. If either finite-element
         * space changed sequence, the old element ID is discarded and the next
         * evaluation reloads lazily against the updated spaces.
         */
        void Refresh();

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

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

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

        [[nodiscard]] VolumeQuadratureContext GetQuadratureContext(
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point
        );

        [[nodiscard]] FaceQuadratureContext GetFaceQuadratureContext(
            mfem::FaceElementTransformations &transformation,
            const mfem::IntegrationPoint &integration_point,
            FaceElementSide side = FaceElementSide::element_1
        );

        void GetPhysicalPoint(
            mfem::ElementTransformation &transformation,
            const mfem::IntegrationPoint &integration_point,
            mfem::Vector &physical_position
        );

    private:
        void ValidateFieldBindings() const;
        [[nodiscard]] bool InvalidateForChangedSpaces();
        void LoadElement(int element_id);

        const DomainMapper &m_mapper;
        const mfem::GridFunction &m_displacement;
        const mfem::GridFunction &m_compactification_coordinate;
        const mfem::FiniteElementSpace *m_displacement_space;
        const mfem::FiniteElementSpace *m_compactification_space;
        long m_displacement_space_sequence;
        long m_compactification_space_sequence;
        DomainMapper::Workspace m_workspace;

        mfem::Array<int> m_displacement_dofs;
        mfem::Array<int> m_compactification_dofs;
        mfem::Vector m_element_displacement;
        mfem::Vector m_element_compactification;
        std::unique_ptr<ElementDisplacementData> m_displacement_data;
        std::unique_ptr<ElementCompactificationData> m_compactification_data;
        int m_cached_element_id{-1};
    };

} // namespace mean_field::mapping