module;

#include <algorithm>
#include <array>
#include <limits>
#include <cmath>
#include <functional>
#include <map>
#include <memory>
#include <mutex>
#include <stdexcept>
#include <utility>
#include <vector>

#include <mfem.hpp>
#include <mpi.h>

module mean_field;
import :fem.reference_tables;

namespace mean_field::fem {
    namespace {
        // Three one-dimensional contractions, with x fastest throughout.
        // Input/output must not alias either scratch array.
        void tensor_values(
            int dx, int dy, int dz, int q, const double *bx, const double *by, const double *bz,
            const double *input, double *output, double *sx, double *sxy, bool transpose
        ) {
            if (!transpose) {
                for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int x = 0; x < q; ++x) {
                    double v = 0.0;
                    for (int i = 0; i < dx; ++i) v += bx[x + q*i] * input[i + dx*(j + dy*k)];
                    sx[x + q*(j + dy*k)] = v;
                }
                for (int k = 0; k < dz; ++k) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
                    double v = 0.0;
                    for (int j = 0; j < dy; ++j) v += by[y + q*j] * sx[x + q*(j + dy*k)];
                    sxy[x + q*(y + q*k)] = v;
                }
                for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
                    double v = 0.0;
                    for (int k = 0; k < dz; ++k) v += bz[z + q*k] * sxy[x + q*(y + q*k)];
                    output[x + q*(y + q*z)] = v;
                }
            } else {
                for (int k = 0; k < dz; ++k) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
                    double v = 0.0;
                    for (int z = 0; z < q; ++z) v += bz[z + q*k] * input[x + q*(y + q*z)];
                    sxy[x + q*(y + q*k)] = v;
                }
                for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int x = 0; x < q; ++x) {
                    double v = 0.0;
                    for (int y = 0; y < q; ++y) v += by[y + q*j] * sxy[x + q*(y + q*k)];
                    sx[x + q*(j + dy*k)] = v;
                }
                for (int k = 0; k < dz; ++k) for (int j = 0; j < dy; ++j) for (int i = 0; i < dx; ++i) {
                    double v = 0.0;
                    for (int x = 0; x < q; ++x) v += bx[x + q*i] * sx[x + q*(j + dy*k)];
                    output[i + dx*(j + dy*k)] = v;
                }
            }
        }

        struct ReferenceTableKey {
            const mfem::FiniteElement *element;
            std::vector<std::array<double, 4>> points;

            bool operator<(const ReferenceTableKey &other) const {
                if (element != other.element)
                    return std::less<const mfem::FiniteElement *>{}(element, other.element);
                return points < other.points;
            }
        };

        ReferenceTableKey make_key(
            const mfem::FiniteElement &element,
            const mfem::IntegrationRule &rule
        ) {
            ReferenceTableKey key{.element = &element, .points = {}};
            key.points.reserve(rule.GetNPoints());
            for (int q = 0; q < rule.GetNPoints(); ++q) {
                const auto &point = rule.IntPoint(q);
                const std::array<double, 4> values{
                    point.x, element.GetDim() > 1 ? point.y : 0.0, element.GetDim() > 2 ? point.z : 0.0, point.weight
                };
                for (const double value : values) {
                    if (!std::isfinite(value))
                        throw std::invalid_argument("Reference table quadrature entries must be finite.");
                }
                key.points.push_back(values);
            }
            return key;
        }
    } // namespace

    struct ReferenceTableCache::Storage {
        std::mutex mutex;
        std::map<ReferenceTableKey, std::shared_ptr<const ScalarReferenceTable>> scalar_tables;
        std::map<ReferenceTableKey, std::shared_ptr<const VectorReferenceTable>> vector_tables;
    };

    ReferenceTableCache::ReferenceTableCache() : m_storage(std::make_unique<Storage>()) {
    }
    ReferenceTableCache::~ReferenceTableCache() = default;

    std::shared_ptr<const ScalarReferenceTable> ReferenceTableCache::GetScalarTable(
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    ) const {
        if (element.GetRangeType() != mfem::FiniteElement::SCALAR)
            throw std::invalid_argument("A scalar reference table requires a scalar finite element.");
        auto key = make_key(element, rule);
        const std::lock_guard lock(m_storage->mutex);
        if (const auto found = m_storage->scalar_tables.find(key); found != m_storage->scalar_tables.end())
            return found->second;
        auto table = std::shared_ptr<const ScalarReferenceTable>(new ScalarReferenceTable(element, rule));
        m_storage->scalar_tables.emplace(std::move(key), table);
        return table;
    }

    std::shared_ptr<const VectorReferenceTable> ReferenceTableCache::GetVectorTable(
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    ) const {
        if (element.GetRangeType() != mfem::FiniteElement::VECTOR)
            throw std::invalid_argument("A vector reference table requires a vector finite element.");
        auto key = make_key(element, rule);
        const std::lock_guard lock(m_storage->mutex);
        if (const auto found = m_storage->vector_tables.find(key); found != m_storage->vector_tables.end())
            return found->second;
        auto table = std::shared_ptr<const VectorReferenceTable>(new VectorReferenceTable(element, rule));
        m_storage->vector_tables.emplace(std::move(key), table);
        return table;
    }

    ScalarReferenceTable::ScalarReferenceTable(
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    )
        : m_values(
              rule.GetNPoints(),
              element.GetDof()
          ),
          m_dimension(element.GetDim()) {
        mfem::Vector values(element.GetDof());
        if (element.GetDerivType() == mfem::FiniteElement::GRAD)
            m_gradients.resize(rule.GetNPoints());
        for (int q = 0; q < rule.GetNPoints(); ++q) {
            const auto &point = rule.IntPoint(q);
            element.CalcShape(point, values);
            for (int dof = 0; dof < element.GetDof(); ++dof)
                m_values(q, dof) = values(dof);
            if (!m_gradients.empty()) {
                auto &gradient = m_gradients[q];
                gradient.SetSize(element.GetDof(), m_dimension);
                element.CalcDShape(point, gradient);
            }
        }
        PrepareTensorEvaluation3D(element, rule);
    }

    const mfem::DenseMatrix &ScalarReferenceTable::GetValues() const {
        return m_values;
    }
    const mfem::DenseMatrix &ScalarReferenceTable::GetGradients(const int point) const {
        return m_gradients.at(point);
    }
    int ScalarReferenceTable::GetPointCount() const {
        return m_values.Height();
    }
    int ScalarReferenceTable::GetDofCount() const {
        return m_values.Width();
    }
    int ScalarReferenceTable::GetDimension() const {
        return m_dimension;
    }

    void ScalarReferenceTable::PrepareTensorEvaluation3D(
        const mfem::FiniteElement &element, const mfem::IntegrationRule &rule
    ) {
        // Nodal H1 and L2 hexes share the same verified tensor contract.
        // Other elements and non-Cartesian rules keep the existing dense path.
        const auto *hex = dynamic_cast<const mfem::NodalTensorFiniteElement *>(&element);
        if (dynamic_cast<const mfem::H1_HexahedronElement *>(&element) == nullptr &&
            dynamic_cast<const mfem::L2_HexahedronElement *>(&element) == nullptr) return;
        if (hex == nullptr || element.GetMapType() != mfem::FiniteElement::VALUE ||
            element.GetDerivType() != mfem::FiniteElement::GRAD || rule.GetNPoints() == 0) return;
        const int d = element.GetOrder() + 1;
        const int q = static_cast<int>(std::lround(std::cbrt(rule.GetNPoints())));
        if (d * d * d != element.GetDof() || q * q * q != rule.GetNPoints()) return;
        for (int z = 0; z < q; ++z) {
            for (int y = 0; y < q; ++y) {
                for (int x = 0; x < q; ++x) {
                    const auto &point = rule.IntPoint(x + q * (y + q * z));
                    if (point.x != rule.IntPoint(x).x ||
                        point.y != rule.IntPoint(q * y).y ||
                        point.z != rule.IntPoint(q * q * z).z) return;
                }
            }
        }
        const auto &map = hex->GetDofMap();
        if (map.Size() != 0 && map.Size() != element.GetDof()) return;
        m_tensorDofMap.SetSize(element.GetDof());
        std::vector<bool> seen(element.GetDof(), false);
        for (int i = 0; i < element.GetDof(); ++i) {
            const int native = map.Size() == 0 ? i : map[i];
            if (native < 0 || native >= element.GetDof() || seen[native]) return;
            seen[native] = true;
            m_tensorDofMap[i] = native;
        }
        mfem::Vector b(d), g(d);
        for (int axis = 0; axis < 3; ++axis) {
            m_tensorValues[axis].SetSize(q, d);
            m_tensorGradients[axis].SetSize(q, d);
            for (int i = 0; i < q; ++i) {
                const double coordinate = axis == 0 ? rule.IntPoint(i).x :
                    axis == 1 ? rule.IntPoint(q * i).y : rule.IntPoint(q * q * i).z;
                hex->GetBasis1D().Eval(coordinate, b, g);
                for (int j = 0; j < d; ++j) {
                    m_tensorValues[axis](i, j) = b(j);
                    m_tensorGradients[axis](i, j) = g(j);
                }
            }
        }
        // Verify the native DOF map and factorization once per shared table.
        // This also guards against an unexpected MFEM basis convention.
        const auto agrees = [](double actual, double expected) {
            return std::isfinite(actual) && std::abs(actual - expected) <=
                64.0 * std::numeric_limits<double>::epsilon() * std::max(1.0, std::abs(expected));
        };
        for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
            const int point = x + q * (y + q * z);
            for (int k = 0; k < d; ++k) for (int j = 0; j < d; ++j) for (int i = 0; i < d; ++i) {
                const int native = m_tensorDofMap[i + d * (j + d * k)];
                const double bx = m_tensorValues[0](x, i), by = m_tensorValues[1](y, j),
                             bz = m_tensorValues[2](z, k);
                if (!agrees(bx * by * bz, m_values(point, native)) ||
                    !agrees(m_tensorGradients[0](x, i) * by * bz, m_gradients[point](native, 0)) ||
                    !agrees(bx * m_tensorGradients[1](y, j) * bz, m_gradients[point](native, 1)) ||
                    !agrees(bx * by * m_tensorGradients[2](z, k), m_gradients[point](native, 2))) return;
            }
        }
        m_tensorDofsPerAxis = d;
        m_tensorPointsPerAxis = q;
    }

    bool ScalarReferenceTable::HasTensorEvaluation3D() const noexcept {
        return m_tensorPointsPerAxis != 0;
    }

    void ScalarReferenceTable::EvaluateTensor3D(
        const mfem::DenseMatrix &elementDofs, mfem::DenseMatrix &valuesAndGradients,
        mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && elementDofs.Height() == GetDofCount(),
                    "Tensor evaluation requires a compatible verified hexahedral reference table.");
        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis;
        const int nx = q * d * d, nxy = q * q * d, nq = q * q * q;
        workspace.SetSize(2 * nx + 3 * nxy);
        double *xValue = workspace.HostWrite(), *xGradient = xValue + nx;
        double *xyValue = xGradient + nx, *xyDx = xyValue + nxy, *xyDy = xyDx + nxy;
        valuesAndGradients.SetSize(nq, 4 * elementDofs.Width());
        const double *u = elementDofs.Data();
        double *out = valuesAndGradients.Data();
        const double *bx = m_tensorValues[0].Data(), *by = m_tensorValues[1].Data(),
                     *bz = m_tensorValues[2].Data(), *gx = m_tensorGradients[0].Data(),
                     *gy = m_tensorGradients[1].Data(), *gz = m_tensorGradients[2].Data();
        const int ndof = GetDofCount();
        for (int c = 0; c < elementDofs.Width(); ++c) {
            // Contract x, retaining the value and x derivative.
            for (int k = 0; k < d; ++k) for (int j = 0; j < d; ++j) for (int x = 0; x < q; ++x) {
                double value = 0.0, dx = 0.0;
                for (int i = 0; i < d; ++i) {
                    const double coefficient = u[m_tensorDofMap[i + d * (j + d * k)] + c * ndof];
                    value += bx[x + q * i] * coefficient;
                    dx += gx[x + q * i] * coefficient;
                }
                const int index = x + q * (j + d * k);
                xValue[index] = value; xGradient[index] = dx;
            }
            // Contract y; propagate dx and form dy.
            for (int k = 0; k < d; ++k) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
                double value = 0.0, dx = 0.0, dy = 0.0;
                for (int j = 0; j < d; ++j) {
                    const int index = x + q * (j + d * k);
                    value += by[y + q * j] * xValue[index];
                    dx += by[y + q * j] * xGradient[index];
                    dy += gy[y + q * j] * xValue[index];
                }
                const int index = x + q * (y + q * k);
                xyValue[index] = value; xyDx[index] = dx; xyDy[index] = dy;
            }
            // Contract z and write all four fields in quadrature ordering.
            for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
                double value = 0.0, dx = 0.0, dy = 0.0, dz = 0.0;
                for (int k = 0; k < d; ++k) {
                    const int index = x + q * (y + q * k);
                    value += bz[z + q * k] * xyValue[index];
                    dx += bz[z + q * k] * xyDx[index];
                    dy += bz[z + q * k] * xyDy[index];
                    dz += gz[z + q * k] * xyValue[index];
                }
                const int point = x + q * (y + q * z);
                out[point + nq * (4 * c)] = value;
                out[point + nq * (4 * c + 1)] = dx;
                out[point + nq * (4 * c + 2)] = dy;
                out[point + nq * (4 * c + 3)] = dz;
            }
        }
    }

    void ScalarReferenceTable::ApplyGradientTransposeTensor3D(
        const mfem::DenseMatrix &quadratureDuals, mfem::DenseMatrix &elementDuals,
        mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Height() == GetPointCount() &&
                    quadratureDuals.Width() % 3 == 0, "Tensor gradient transpose has incompatible duals.");
        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
        workspace.SetSize(n + q*d*d + q*q*d);
        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
        elementDuals.SetSize(n, quadratureDuals.Width()/3);
        elementDuals = 0.0;
        for (int c = 0; c < elementDuals.Width(); ++c) for (int a = 0; a < 3; ++a) {
            tensor_values(d, d, d, q,
                (a == 0 ? m_tensorGradients[0] : m_tensorValues[0]).Data(),
                (a == 1 ? m_tensorGradients[1] : m_tensorValues[1]).Data(),
                (a == 2 ? m_tensorGradients[2] : m_tensorValues[2]).Data(),
                quadratureDuals.GetColumn(3*c+a), lex, sx, sxy, true);
            for (int i = 0; i < n; ++i) elementDuals(m_tensorDofMap[i], c) += lex[i];
        }
    }

    void ScalarReferenceTable::EvaluateValuesTensor3D(
        const mfem::Vector &elementDofs, mfem::Vector &values, mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && elementDofs.Size() == GetDofCount(),
                    "Tensor values require compatible element DOFs.");
        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
        workspace.SetSize(n + q*d*d + q*q*d);
        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
        const double *u = elementDofs.HostRead();
        for (int i = 0; i < n; ++i) lex[i] = u[m_tensorDofMap[i]];
        values.SetSize(GetPointCount());
        tensor_values(d, d, d, q, m_tensorValues[0].Data(), m_tensorValues[1].Data(),
                      m_tensorValues[2].Data(), lex, values.HostWrite(), sx, sxy, false);
    }

    void ScalarReferenceTable::ApplyTransposeTensor3D(
        const mfem::Vector &quadratureDuals, mfem::Vector &elementDuals, mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Size() == GetPointCount(),
                    "Tensor transpose requires compatible quadrature duals.");
        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
        workspace.SetSize(n + q*d*d + q*q*d);
        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
        tensor_values(d, d, d, q, m_tensorValues[0].Data(), m_tensorValues[1].Data(),
                      m_tensorValues[2].Data(), quadratureDuals.HostRead(), lex, sx, sxy, true);
        elementDuals.SetSize(n);
        double *out = elementDuals.HostWrite();
        for (int i = 0; i < n; ++i) out[m_tensorDofMap[i]] = lex[i];
    }

    void ScalarReferenceTable::ApplyTransposeTensor3D(
        const mfem::DenseMatrix &quadratureDuals, mfem::DenseMatrix &elementDuals, mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Height() == GetPointCount(),
                    "Tensor transpose requires compatible quadrature duals.");
        const int d = m_tensorDofsPerAxis, q = m_tensorPointsPerAxis, n = GetDofCount();
        workspace.SetSize(n + q*d*d + q*q*d);
        double *lex = workspace.HostWrite(), *sx = lex + n, *sxy = sx + q*d*d;
        elementDuals.SetSize(n, quadratureDuals.Width());
        for (int c = 0; c < quadratureDuals.Width(); ++c) {
            tensor_values(d, d, d, q, m_tensorValues[0].Data(), m_tensorValues[1].Data(),
                          m_tensorValues[2].Data(), quadratureDuals.GetColumn(c), lex, sx, sxy, true);
            for (int i = 0; i < n; ++i) elementDuals(m_tensorDofMap[i], c) = lex[i];
        }
    }

    VectorReferenceTable::VectorReferenceTable(
        const mfem::FiniteElement &element,
        const mfem::IntegrationRule &rule
    )
        : m_dof_count(element.GetDof()),
          m_dimension(element.GetRangeDim()) {
        m_values.resize(rule.GetNPoints());
        for (int q = 0; q < rule.GetNPoints(); ++q) {
            auto &values = m_values[q];
            values.SetSize(m_dof_count, m_dimension);
            element.CalcVShape(rule.IntPoint(q), values);
        }
        PrepareTensorEvaluation3D(element, rule);
    }

    void VectorReferenceTable::PrepareTensorEvaluation3D(
        const mfem::FiniteElement &element, const mfem::IntegrationRule &rule
    ) {
        const auto *rt = dynamic_cast<const mfem::RT_HexahedronElement *>(&element);
        if (rt == nullptr || element.GetMapType() != mfem::FiniteElement::H_DIV || rule.GetNPoints() == 0) return;
        const int q = static_cast<int>(std::lround(std::cbrt(rule.GetNPoints())));
        if (q*q*q != rule.GetNPoints()) return;
        for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
            const auto &point = rule.IntPoint(x + q*(y + q*z));
            if (point.x != rule.IntPoint(x).x || point.y != rule.IntPoint(q*y).y ||
                point.z != rule.IntPoint(q*q*z).z) return;
        }
        int d = 2;
        while (3*d*(d-1)*(d-1) < m_dof_count) ++d;
        if (3*d*(d-1)*(d-1) != m_dof_count) return;
        const auto &map = rt->GetDofMap();
        if (map.Size() != m_dof_count) return;
        m_tensorDofMap = map;
        std::vector<bool> seen(m_dof_count, false);
        for (int i = 0; i < m_dof_count; ++i) {
            const int native = map[i] < 0 ? -1-map[i] : map[i];
            if (native < 0 || native >= m_dof_count || seen[native]) return;
            seen[native] = true;
        }
        const int block = m_dof_count / 3;
        const auto basis = [&](int point, int lex, int component) {
            const int encoded = map[lex], native = encoded < 0 ? -1-encoded : encoded;
            return (encoded < 0 ? -1.0 : 1.0) * m_values[point](native, component);
        };
        // Extract owned 1D factors from the full basis. This supports nodal and
        // integrated open bases without retaining DofToQuad pointers to rules.
        for (int c = 0; c < 3; ++c) {
            const int dims[3]{c == 0 ? d : d-1, c == 1 ? d : d-1, c == 2 ? d : d-1};
            int anchor = 0;
            for (int point = 1; point < q*q*q; ++point)
                if (std::abs(basis(point, c*block, c)) > std::abs(basis(anchor, c*block, c))) anchor = point;
            const double base = basis(anchor, c*block, c);
            if (!std::isfinite(base) || base == 0.0) return;
            const int coordinates[3]{anchor % q, (anchor/q) % q, anchor/(q*q)};
            const int strides[3]{1, dims[0], dims[0]*dims[1]};
            for (int axis = 0; axis < 3; ++axis) {
                auto &b = m_tensorValues[c][axis];
                b.SetSize(q, dims[axis]);
                for (int i = 0; i < dims[axis]; ++i) for (int t = 0; t < q; ++t) {
                    int xyz[3]{coordinates[0], coordinates[1], coordinates[2]};
                    xyz[axis] = t;
                    b(t, i) = basis(xyz[0] + q*(xyz[1] + q*xyz[2]), c*block + strides[axis]*i, c) /
                              (axis == 0 ? 1.0 : base);
                }
            }
            // Check every native basis entry, including zero components.
            for (int k = 0; k < dims[2]; ++k) for (int j = 0; j < dims[1]; ++j)
                for (int i = 0; i < dims[0]; ++i) {
                    const int lex = c*block + i + dims[0]*(j + dims[1]*k);
                    for (int z = 0; z < q; ++z) for (int y = 0; y < q; ++y) for (int x = 0; x < q; ++x) {
                        const int point = x + q*(y + q*z);
                        const double product = m_tensorValues[c][0](x,i) * m_tensorValues[c][1](y,j) *
                                               m_tensorValues[c][2](z,k);
                        for (int component = 0; component < 3; ++component) {
                            const double expected = basis(point, lex, component);
                            const double actual = component == c ? product : 0.0;
                            if (!std::isfinite(actual) || std::abs(actual - expected) >
                                128.0 * std::numeric_limits<double>::epsilon() * std::max(1.0, std::abs(expected))) return;
                        }
                    }
                }
        }
        m_tensorClosedDofs = d;
        m_tensorPointsPerAxis = q;
    }

    bool VectorReferenceTable::HasTensorEvaluation3D() const noexcept {
        return m_tensorPointsPerAxis != 0;
    }

    void VectorReferenceTable::EvaluateTensor3D(
        const mfem::Vector &elementDofs, mfem::DenseMatrix &values, mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && elementDofs.Size() == m_dof_count,
                    "RT tensor evaluation requires compatible element DOFs.");
        const int d = m_tensorClosedDofs, q = m_tensorPointsPerAxis, block = m_dof_count/3;
        workspace.SetSize(block + q*d*d + q*q*d);
        double *lex = workspace.HostWrite(), *sx = lex + block, *sxy = sx + q*d*d;
        values.SetSize(q*q*q, 3);
        for (int c = 0; c < 3; ++c) {
            for (int i = 0; i < block; ++i) {
                const int encoded = m_tensorDofMap[c*block+i];
                lex[i] = encoded < 0 ? -elementDofs(-1-encoded) : elementDofs(encoded);
            }
            tensor_values(c == 0 ? d : d-1, c == 1 ? d : d-1, c == 2 ? d : d-1, q,
                          m_tensorValues[c][0].Data(), m_tensorValues[c][1].Data(), m_tensorValues[c][2].Data(),
                          lex, values.GetColumn(c), sx, sxy, false);
        }
    }

    void VectorReferenceTable::ApplyTransposeTensor3D(
        const mfem::DenseMatrix &quadratureDuals, mfem::Vector &elementDuals, mfem::Vector &workspace
    ) const {
        MFEM_VERIFY(HasTensorEvaluation3D() && quadratureDuals.Height() == GetPointCount() && quadratureDuals.Width() == 3,
                    "RT tensor transpose requires three quadrature dual components.");
        const int d = m_tensorClosedDofs, q = m_tensorPointsPerAxis, block = m_dof_count/3;
        workspace.SetSize(block + q*d*d + q*q*d);
        double *lex = workspace.HostWrite(), *sx = lex + block, *sxy = sx + q*d*d;
        elementDuals.SetSize(m_dof_count);
        for (int c = 0; c < 3; ++c) {
            tensor_values(c == 0 ? d : d-1, c == 1 ? d : d-1, c == 2 ? d : d-1, q,
                          m_tensorValues[c][0].Data(), m_tensorValues[c][1].Data(), m_tensorValues[c][2].Data(),
                          quadratureDuals.GetColumn(c), lex, sx, sxy, true);
            for (int i = 0; i < block; ++i) {
                const int encoded = m_tensorDofMap[c*block+i];
                if (encoded < 0) elementDuals(-1-encoded) = -lex[i]; else elementDuals(encoded) = lex[i];
            }
        }
    }

    const mfem::DenseMatrix &VectorReferenceTable::GetValues(const int point) const {
        return m_values.at(point);
    }
    int VectorReferenceTable::GetPointCount() const {
        return static_cast<int>(m_values.size());
    }
    int VectorReferenceTable::GetDofCount() const {
        return m_dof_count;
    }
    int VectorReferenceTable::GetDimension() const {
        return m_dimension;
    }

    struct PreparedGeometryInspector::Storage {
        struct Element {
            std::shared_ptr<const ScalarReferenceTable> basis;
            mfem::DenseMatrix inverseJacobians; // rows: points; columns: row-major entries
        };
        const mfem::ParFiniteElementSpace *space{nullptr};
        const mfem::Mesh *mesh{nullptr};
        long spaceSequence{-1}, meshSequence{-1}, nodesSequence{-1};
        std::unique_ptr<ReferenceTableCache> tables;
        std::vector<Element> elements;
        std::unique_ptr<mfem::ParGridFunction> field;
        mfem::Array<int> dofs;
        mfem::Vector elementDofs, workspace;
        mfem::DenseMatrix dofMatrix, samples, gradient;
        std::uint64_t builds{0};

        void Clear() {
            elements.clear();
            tables.reset();
            field.reset();
            space = nullptr;
            mesh = nullptr;
        }
    };

    PreparedGeometryInspector::PreparedGeometryInspector() : m_storage(std::make_unique<Storage>()) {}
    PreparedGeometryInspector::~PreparedGeometryInspector() = default;
    PreparedGeometryInspector::PreparedGeometryInspector(PreparedGeometryInspector &&) noexcept = default;
    PreparedGeometryInspector &PreparedGeometryInspector::operator=(PreparedGeometryInspector &&) noexcept = default;
    void PreparedGeometryInspector::ClearCache() const { m_storage->Clear(); }
    int PreparedGeometryInspector::GetTensorElementCount() const noexcept {
        int count = 0;
        for (const auto &element : m_storage->elements)
            count += element.basis && element.basis->HasTensorEvaluation3D();
        return count;
    }
    std::uint64_t PreparedGeometryInspector::GetGeometryBuildCount() const noexcept { return m_storage->builds; }

    double PreparedGeometryInspector::Inspect(
        mfem::ParFiniteElementSpace &space, const mfem::Vector &displacement, bool reference
    ) const {
        MFEM_VERIFY(displacement.Size() == space.GetTrueVSize(), "Geometry inspection received incompatible DOFs.");
        auto &s = *m_storage;
        auto *mesh = space.GetMesh();
        const int dimension = mesh->SpaceDimension();
        MFEM_VERIFY(space.GetVDim() == dimension, "Geometry inspection requires a spatial displacement vector.");
        if (s.space != &space || s.mesh != mesh || s.spaceSequence != space.GetSequence() ||
            s.meshSequence != mesh->GetSequence() || s.nodesSequence != mesh->GetNodesSequence()) {
            s.Clear();
            s.space = &space; s.mesh = mesh;
            s.spaceSequence = space.GetSequence(); s.meshSequence = mesh->GetSequence();
            s.nodesSequence = mesh->GetNodesSequence();
            s.field = std::make_unique<mfem::ParGridFunction>(&space);
            s.tables = std::make_unique<ReferenceTableCache>();
            s.elements.resize(mesh->GetNE());
        }
        s.field->SetFromTrueDofs(displacement);
        double minimum = std::numeric_limits<double>::infinity();
        int finite = 1;
        auto accumulate = [&](double determinant) {
            if (!std::isfinite(determinant)) finite = 0;
            else minimum = std::min(minimum, determinant);
        };
        for (int e = 0; e < mesh->GetNE(); ++e) {
            auto &entry = s.elements[e];
            auto &transformation = *mesh->GetElementTransformation(e);
            const auto &element = *space.GetFE(e);
            // Preserve the existing d*p geometry inspection order, including
            // vacuum elements and points beyond the physics quadrature.
            const int order = std::max(element.GetOrder() + 2, 2 * dimension * element.GetOrder());
            const auto &rule = mfem::IntRules.Get(transformation.GetGeometryType(), order);
            if (!reference && dimension == 3 && element.GetDim() == 3 &&
                element.GetMapType() == mfem::FiniteElement::VALUE &&
                dynamic_cast<const mfem::H1_HexahedronElement *>(&element)) {
                if (!entry.basis) {
                    entry.basis = s.tables->GetScalarTable(element, rule);
                    if (entry.basis->HasTensorEvaluation3D()) {
                        entry.inverseJacobians.SetSize(rule.GetNPoints(), 9);
                        for (int q = 0; q < rule.GetNPoints(); ++q) {
                            transformation.SetIntPoint(&rule.IntPoint(q));
                            const auto &inverse = transformation.InverseJacobian();
                            for (int a = 0; a < 3; ++a) for (int b = 0; b < 3; ++b)
                                entry.inverseJacobians(q, a*3+b) = inverse(a,b);
                        }
                        ++s.builds;
                    }
                }
            }
            if (!reference && entry.basis && entry.basis->HasTensorEvaluation3D()) {
                auto *dofTransformation = space.GetElementVDofs(e, s.dofs);
                s.field->GetSubVector(s.dofs, s.elementDofs);
                if (dofTransformation) dofTransformation->InvTransformPrimal(s.elementDofs);
                const int nd = element.GetDof();
                s.dofMatrix.SetSize(nd, 3);
                // GetElementVDofs groups local entries by component for either
                // global ordering, as used by GridFunction::GetVectorGradientHat.
                for (int c = 0; c < 3; ++c) for (int i = 0; i < nd; ++i)
                    s.dofMatrix(i,c) = s.elementDofs(i+c*nd);
                entry.basis->EvaluateTensor3D(s.dofMatrix, s.samples, s.workspace);
                s.gradient.SetSize(3);
                for (int q = 0; q < rule.GetNPoints(); ++q) {
                    for (int c = 0; c < 3; ++c) for (int b = 0; b < 3; ++b) {
                        double value = c == b ? 1.0 : 0.0;
                        for (int a = 0; a < 3; ++a)
                            value += s.samples(q, 4*c+a+1) * entry.inverseJacobians(q, a*3+b);
                        s.gradient(c,b) = value;
                    }
                    accumulate(s.gradient.Det());
                }
            } else {
                for (int q = 0; q < rule.GetNPoints(); ++q) {
                    transformation.SetIntPoint(&rule.IntPoint(q));
                    s.field->GetVectorGradient(transformation, s.gradient);
                    for (int c = 0; c < dimension; ++c) s.gradient(c,c) += 1.0;
                    accumulate(s.gradient.Det());
                }
            }
        }
        double globalMinimum = 0.0;
        int globalFinite = 0;
        const int minimumStatus = MPI_Allreduce(&minimum, &globalMinimum, 1, MPI_DOUBLE, MPI_MIN, space.GetComm());
        const int finiteStatus = MPI_Allreduce(&finite, &globalFinite, 1, MPI_INT, MPI_MIN, space.GetComm());
        if (minimumStatus != MPI_SUCCESS || finiteStatus != MPI_SUCCESS)
            throw std::runtime_error("Could not reduce the mapped-geometry inspection.");
        return globalFinite ? globalMinimum : std::numeric_limits<double>::quiet_NaN();
    }

} // namespace mean_field::fem