#pragma once

#include <mfem.hpp>

namespace {
    // U and G are n x 3, component/derivative columns respectively. Keep the
    // nine reductions together instead of invoking a general GEMM per point.
    // The returned 3 x 3 matrix is column-major: A(c,j) = sum_i U(i,c) G(i,j).
    inline void displacement_gradient_3d(
        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &gradients, double *a
    ) {
        const int n = direction.Height();
        const double *u = direction.Data();
        const double *g = gradients.Data();
        for (int j = 0; j < 3; ++j) {
            double x = 0.0, y = 0.0, z = 0.0;
            const double *gj = g + j * n;
            for (int i = 0; i < n; ++i) {
                const double v = gj[i];
                x += u[i] * v;
                y += u[i + n] * v;
                z += u[i + 2 * n] * v;
            }
            a[3 * j] = x; a[3 * j + 1] = y; a[3 * j + 2] = z;
        }
    }

    inline void vector_basis_value_3d(
        const mfem::DenseMatrix &basis, const mfem::Vector &coefficients, double *value
    ) {
        const int n = basis.Height();
        const double *b = basis.Data();
        const double *v = coefficients.HostRead();
        double x = 0.0, y = 0.0, z = 0.0;
        for (int i = 0; i < n; ++i) {
            x += b[i] * v[i]; y += b[i + n] * v[i]; z += b[i + 2 * n] * v[i];
        }
        value[0] = x; value[1] = y; value[2] = z;
    }

    // Matrix cache layout: nq x 9, with row-major matrix entries in columns.
    inline void cached_matrix_vector_3d(
        const mfem::DenseMatrix &cache, const int point, const double *v, double *out
    ) {
        const int nq = cache.Height();
        const double *m = cache.Data() + point;
        for (int r = 0; r < 3; ++r)
            out[r] = m[(3 * r) * nq] * v[0] + m[(3 * r + 1) * nq] * v[1] +
                     m[(3 * r + 2) * nq] * v[2];
    }

    inline double displacement_trace_3d(
        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &gradients,
        const mfem::DenseMatrix &inverseJacobians, const int point
    ) {
        double a[9];
        displacement_gradient_3d(direction, gradients, a);
        const double *k = inverseJacobians.Data() + point;
        const int nq = inverseJacobians.Height();
        double trace = 0.0;
        for (int j = 0; j < 3; ++j)
            for (int c = 0; c < 3; ++c) trace += k[(3 * j + c) * nq] * a[c + 3 * j];
        return trace;
    }

    // (du/dxi) v, without constructing du/dxi. v is fixed at this point.
    inline void displacement_directional_value_3d(
        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &gradients,
        const double *v, double *out
    ) {
        const int n = direction.Height();
        const double *u = direction.Data();
        const double *g = gradients.Data();
        double x = 0.0, y = 0.0, z = 0.0;
        for (int i = 0; i < n; ++i) {
            const double d = g[i] * v[0] + g[i + n] * v[1] + g[i + 2 * n] * v[2];
            x += u[i] * d; y += u[i + n] * d; z += u[i + 2 * n] * d;
        }
        out[0] = x; out[1] = y; out[2] = z;
    }

    inline void pressure_displacement_action_3d(
        const mfem::DenseMatrix &direction, const mfem::DenseMatrix &physicalGradients,
        const double pressureWeight, mfem::Vector &action
    ) {
        // H = du/dx; d(G dx) = (tr(H) G - G H) dx.
        double h[9];
        displacement_gradient_3d(direction, physicalGradients, h);
        const double trace = h[0] + h[4] + h[8];
        const int n = direction.Height();
        const double *g = physicalGradients.Data();
        double *r = action.HostReadWrite();
        for (int c = 0; c < 3; ++c) {
            const double a = pressureWeight * (h[3 * c] - (c == 0 ? trace : 0.0));
            const double b = pressureWeight * (h[3 * c + 1] - (c == 1 ? trace : 0.0));
            const double d = pressureWeight * (h[3 * c + 2] - (c == 2 ? trace : 0.0));
            for (int i = 0; i < n; ++i) r[i + c * n] += a * g[i] + b * g[i + n] + d * g[i + 2 * n];
        }
    }

    // Element VDofs are component-blocked, including for a byVDIM global space;
    // this is the ordering consumed by ElementDisplacementDataFromElementVDofs.
    // For an ordinary interior VALUE field, dx = du and
    // d(log det J_element) = tr(J_element^-1 * du/dxi).
    inline void add_interior_functional_row(
        const mfem::DenseMatrix &gradients,
        const mfem::DenseMatrix &values,
        const int point,
        const mfem::DenseMatrix &inverseElementJacobian,
        const double determinantCoefficient,
        const double *positionCoefficient,
        mfem::Vector &row
    ) {
        const int n = gradients.Height();
        const int dimension = gradients.Width();
        const int nq = values.Height();
        const double *g = gradients.Data();
        const double *b = values.Data() + point;
        double *r = row.HostReadWrite();
        for (int component = 0; component < dimension; ++component) {
            for (int i = 0; i < n; ++i) {
                double coefficient = positionCoefficient[component] * b[i * nq];
                for (int derivative = 0; derivative < dimension; ++derivative) {
                    coefficient += determinantCoefficient * inverseElementJacobian(derivative, component) *
                                   g[i + derivative * n];
                }
                r[i + component * n] += coefficient;
            }
        }
    }

    inline double apply_interior_functional(
        const mfem::DenseMatrix &direction,
        const mfem::DenseMatrix &gradients,
        const mfem::DenseMatrix &values,
        const int point,
        const double *inverseElementJacobian,
        const double determinantCoefficient,
        const double *positionCoefficient
    ) {
        const int n = direction.Height();
        const int dimension = direction.Width();
        const int nq = values.Height();
        const double *u = direction.Data();
        const double *g = gradients.Data();
        const double *b = values.Data() + point;
        double result = 0.0;
        for (int component = 0; component < dimension; ++component) {
            double value = 0.0;
            for (int i = 0; i < n; ++i) value += u[i + component * n] * b[i * nq];
            result += positionCoefficient[component] * value;
            for (int derivative = 0; derivative < dimension; ++derivative) {
                double gradient = 0.0;
                for (int i = 0; i < n; ++i)
                    gradient += u[i + component * n] * g[i + derivative * n];
                result += determinantCoefficient * inverseElementJacobian[derivative + component * dimension] *
                          gradient;
            }
        }
        return result;
    }
}