module;

#include <memory>
#include <vector>

#include <mfem.hpp>

export module mean_field:fem.reference_tables;

export namespace mean_field::fem {
    // These tables contain only reference-element data: no mesh coordinates,
    // orientation, element DOF transforms, or state-dependent mapping factors.
    class ScalarReferenceTable {
    public:
        [[nodiscard]] const mfem::DenseMatrix &GetValues() const;
        // Available for GRAD elements; otherwise throws std::out_of_range.
        [[nodiscard]] const mfem::DenseMatrix &GetGradients(int point) const;
        [[nodiscard]] int GetPointCount() const;
        [[nodiscard]] int GetDofCount() const;
        [[nodiscard]] int GetDimension() const;

    private:
        friend class ReferenceTableCache;
        ScalarReferenceTable(
            const mfem::FiniteElement &element,
            const mfem::IntegrationRule &rule
        );

        mfem::DenseMatrix m_values;
        std::vector<mfem::DenseMatrix> m_gradients;
        int m_dimension;
    };

    class VectorReferenceTable {
    public:
        [[nodiscard]] const mfem::DenseMatrix &GetValues(int point) const;
        [[nodiscard]] int GetPointCount() const;
        [[nodiscard]] int GetDofCount() const;
        [[nodiscard]] int GetDimension() const;

    private:
        friend class ReferenceTableCache;
        VectorReferenceTable(
            const mfem::FiniteElement &element,
            const mfem::IntegrationRule &rule
        );

        std::vector<mfem::DenseMatrix> m_values;
        int m_dof_count;
        int m_dimension;
    };

    // Owned by one discretization, never process-global. Finite elements must
    // remain alive and immutable while this cache is used; rebuild the cache if
    // their collections are replaced. Rules are keyed by their actual points
    // and weights, so temporary, copied, or modified rules are safe to use.
    // Returned immutable handles also keep tables alive after cache destruction.
    class ReferenceTableCache {
    public:
        ReferenceTableCache();
        ~ReferenceTableCache();
        ReferenceTableCache(const ReferenceTableCache &)            = delete;
        ReferenceTableCache &operator=(const ReferenceTableCache &) = delete;

        [[nodiscard]] std::shared_ptr<const ScalarReferenceTable> GetScalarTable(
            const mfem::FiniteElement &element,
            const mfem::IntegrationRule &rule
        ) const;
        [[nodiscard]] std::shared_ptr<const VectorReferenceTable> GetVectorTable(
            const mfem::FiniteElement &element,
            const mfem::IntegrationRule &rule
        ) const;

    private:
        struct Storage;
        std::unique_ptr<Storage> m_storage;
    };
} // namespace mean_field::fem