module;

#include <array>
#include <cstdint>
#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;

        // Available only for a verified nodal H1/L2 hexahedron and Cartesian quadrature
        // ordering (x fastest). Results have columns [value, dx, dy, dz] for
        // each input component. Workspace is caller-owned; tables stay immutable.
        [[nodiscard]] bool HasTensorEvaluation3D() const noexcept;
        void EvaluateTensor3D(
            const mfem::DenseMatrix &elementDofs,
            mfem::DenseMatrix &valuesAndGradients,
            mfem::Vector &workspace
        ) const;

        // Transpose of the reference gradient. Input columns are [dx,dy,dz]
        // per output component. Callers supply all physical factors and weights.
        void ApplyGradientTransposeTensor3D(const mfem::DenseMatrix &quadratureDuals,
                                           mfem::DenseMatrix &elementDuals,
                                           mfem::Vector &workspace) const;

        // Value-only scalar contractions avoid forming unused gradients.
        void EvaluateValuesTensor3D(const mfem::Vector &elementDofs, mfem::Vector &values,
                                    mfem::Vector &workspace) const;
        void ApplyTransposeTensor3D(const mfem::Vector &quadratureDuals, mfem::Vector &elementDuals,
                                    mfem::Vector &workspace) const;

        // Apply the transpose of the value basis to quadrature duals. No weights
        // are introduced here; callers supply their already weighted duals.
        void ApplyTransposeTensor3D(
            const mfem::DenseMatrix &quadratureDuals, mfem::DenseMatrix &elementDuals,
            mfem::Vector &workspace
        ) const;

    private:
        void PrepareTensorEvaluation3D(const mfem::FiniteElement &element, const mfem::IntegrationRule &rule);
        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;
        int m_tensorDofsPerAxis{0};
        int m_tensorPointsPerAxis{0};
        mfem::Array<int> m_tensorDofMap;
        std::array<mfem::DenseMatrix, 3> m_tensorValues;
        std::array<mfem::DenseMatrix, 3> m_tensorGradients;
    };

    class VectorReferenceTable {
    public:
        [[nodiscard]] const mfem::DenseMatrix &GetValues(int point) const;
        [[nodiscard]] int GetPointCount() const;
        [[nodiscard]] int GetDofCount() const;
        [[nodiscard]] int GetDimension() const;
        [[nodiscard]] bool HasTensorEvaluation3D() const noexcept;
        void EvaluateTensor3D(
            const mfem::Vector &elementDofs, mfem::DenseMatrix &values, mfem::Vector &workspace
        ) const;
        void ApplyTransposeTensor3D(
            const mfem::DenseMatrix &quadratureDuals, mfem::Vector &elementDuals, mfem::Vector &workspace
        ) const;

    private:
        void PrepareTensorEvaluation3D(const mfem::FiniteElement &element, const mfem::IntegrationRule &rule);
        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;
        int m_tensorClosedDofs{0}, m_tensorPointsPerAxis{0};
        mfem::Array<int> m_tensorDofMap;
        std::array<std::array<mfem::DenseMatrix, 3>, 3> m_tensorValues;
    };

    // 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;
    };
    // The production mapped-geometry inspection, independently benchmarkable.
    // Keeps the original conservative rule and checks every volume element.
    // Mesh/FES/node sequences invalidate the fixed inverse-Jacobian cache;
    // raw node edits require Mesh::NodesUpdated() or ClearCache(). Displacement
    // is always re-evaluated. Like prepared operators, scratch is not thread-safe.
    class PreparedGeometryInspector {
    public:
        PreparedGeometryInspector();
        ~PreparedGeometryInspector();
        PreparedGeometryInspector(const PreparedGeometryInspector &) = delete;
        PreparedGeometryInspector &operator=(const PreparedGeometryInspector &) = delete;
        PreparedGeometryInspector(PreparedGeometryInspector &&) noexcept;
        PreparedGeometryInspector &operator=(PreparedGeometryInspector &&) noexcept;

        // Collective minimum of det(I + grad(u)); NaN if any sampled value is
        // non-finite. reference=true retains GridFunction::GetVectorGradient.
        [[nodiscard]] double Inspect(mfem::ParFiniteElementSpace &space,
                                     const mfem::Vector &displacement, bool reference = false) const;
        void ClearCache() const;
        [[nodiscard]] int GetTensorElementCount() const noexcept;
        [[nodiscard]] std::uint64_t GetGeometryBuildCount() const noexcept;
    private:
        struct Storage;
        std::unique_ptr<Storage> m_storage;
    };
} // namespace mean_field::fem