/*
 * Copyright 2014 Anton Khirnov <anton@khirnov.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <errno.h>
#include <limits.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>

#include <lapacke.h>

#include "brill_data.h"

#define ACC_TEST 1

#define MAX(x, y) ((x) > (y) ? (x) : (y))
#define MIN(x, y) ((x) > (y) ? (y) : (x))
#define SQR(x) ((x) * (x))
#define SGN(x) ((x) >= 0.0 ? 1.0 : -1.0)

/*
 * small number to avoid r=0 singularities
 */
#define EPS 1E-08

/* a set of basis functions */
typedef struct BasisSet {
    /* evaluate the idx-th basis function at the specified point*/
    long double (*eval)      (long double coord, int idx, long double sf);
    /* evaluate the first derivative of the idx-th basis function at the specified point*/
    long double (*eval_diff1)(long double coord, int idx, long double sf);
    /* evaluate the second derivative of the idx-th basis function at the specified point*/
    long double (*eval_diff2)(long double coord, int idx, long double sf);
    /**
     * Get the idx-th collocation point for the specified order.
     * idx runs from 0 to order - 1 (inclusive)
     */
    long double (*colloc_point)(int order, int idx, long double sf);
} BasisSet;

typedef struct QFunc {
    double         (*q)(struct BDContext *bd, double rho, double z);
    double    (*dq_rho)(struct BDContext *bd, double rho, double z);
    double      (*dq_z)(struct BDContext *bd, double rho, double z);
    double   (*d2q_rho)(struct BDContext *bd, double rho, double z);
    double     (*d2q_z)(struct BDContext *bd, double rho, double z);
    double (*d2q_rho_z)(struct BDContext *bd, double rho, double z);
} QFunc;

typedef struct BDPriv {
    const BasisSet *basis;
    const QFunc    *q_func;

    int nb_colloc_points;
    int nb_colloc_points_rho;
    int nb_colloc_points_z;

    int colloc_grid_order_rho;
    int colloc_grid_order_z;
    int nb_coeffs;

    gsl_vector *coeffs;
} BDPriv;

/*
 * The basis of even (n = 2 * idx) SB functions (Boyd 2000, Ch 17.9)
 * SB(x, n) = sin((n + 1) arccot(|x| / L))
 * They are symmetric wrt origin and decay as 1/x in infinity.
 */

static long double sb_even_eval(long double coord, int idx, long double sf)
{
    long double val = (coord == 0.0) ? M_PI_2 : atanl(sf / fabsl(coord));

    idx *= 2;   // even only

    return sinl((idx + 1) * val);
}

static long double sb_even_eval_diff1(long double coord, int idx, long double sf)
{
    long double val = (coord == 0.0) ? M_PI_2 : atanl(sf / fabsl(coord));

    idx *= 2;   // even only

    return - sf * (idx + 1) * SGN(coord) * cosl((idx + 1) * val) / (SQR(sf) + SQR(coord));
}

static long double sb_even_eval_diff2(long double coord, int idx, long double sf)
{
    long double val = (coord == 0.0) ? M_PI_2 : atanl(sf / fabsl(coord));

    idx *= 2;   // even only

    return sf * (idx + 1) * SGN(coord) * (2 * fabsl(coord) * cosl((idx + 1) * val) - sf * (idx + 1) * sinl((idx + 1) * val)) / SQR(SQR(sf) + SQR(coord));
}

static long double sb_even_colloc_point(int order, int idx, long double sf)
{
    long double t;

    idx = order - idx - 1;
    order *= 2;

    t = (idx + 2) * M_PI / (order + 4);
    return sf / tanl(t);
}

static const BasisSet sb_even_basis = {
    .eval         = sb_even_eval,
    .eval_diff1   = sb_even_eval_diff1,
    .eval_diff2   = sb_even_eval_diff2,
    .colloc_point = sb_even_colloc_point,
};

static int brill_construct_matrix(BDContext *bd, gsl_matrix *mat,
                                  gsl_vector *rhs)
{
    BDPriv *s = bd->priv;
    const long double sfr = bd->basis_scale_factor_rho, sfz = bd->basis_scale_factor_z;

    double *basis_val, *basis_dval, *basis_d2val;
    int idx_coeff_rho, idx_coeff_z, idx_grid_rho, idx_grid_z;

    /* precompute the basis values we will need */
    // FIXME: assumes same number of points in ρ and z directions
    basis_val   = malloc(sizeof(*basis_val)   * s->nb_colloc_points_rho * bd->nb_coeffs_rho);
    basis_dval  = malloc(sizeof(*basis_dval)  * s->nb_colloc_points_rho * bd->nb_coeffs_rho);
    basis_d2val = malloc(sizeof(*basis_d2val) * s->nb_colloc_points_rho * bd->nb_coeffs_rho);
    for (int i = 0; i < s->nb_colloc_points_rho; i++) {
        double coord = s->basis->colloc_point(s->colloc_grid_order_rho, i, sfr);
        for (int j = 0; j < bd->nb_coeffs_rho; j++) {
            basis_val  [i * bd->nb_coeffs_rho + j] = s->basis->eval(coord, j, sfr);
            basis_dval [i * bd->nb_coeffs_rho + j] = s->basis->eval_diff1(coord, j, sfr);
            basis_d2val[i * bd->nb_coeffs_rho + j] = s->basis->eval_diff2(coord, j, sfr);
        }
    }

    for (idx_grid_z = 0; idx_grid_z < s->nb_colloc_points_z; idx_grid_z++) {
        double z_physical = s->basis->colloc_point(s->colloc_grid_order_z, idx_grid_z, sfz);

        for (idx_grid_rho = 0; idx_grid_rho < s->nb_colloc_points_rho; idx_grid_rho++) {
            double x_physical = s->basis->colloc_point(s->colloc_grid_order_rho, idx_grid_rho, sfr);
            double d2q        = s->q_func->d2q_rho(bd, x_physical, z_physical) + s->q_func->d2q_z(bd, x_physical, z_physical);
            int idx_grid         = idx_grid_z * s->nb_colloc_points_z + idx_grid_rho;

            for (idx_coeff_z = 0; idx_coeff_z < bd->nb_coeffs_z; idx_coeff_z++)
                for (idx_coeff_rho = 0; idx_coeff_rho < bd->nb_coeffs_rho; idx_coeff_rho++) {
                    int idx_coeff = idx_coeff_z * bd->nb_coeffs_z + idx_coeff_rho;

                    mat->data[idx_grid + mat->tda * idx_coeff] = basis_d2val[idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho] * basis_val[idx_grid_z   * bd->nb_coeffs_z + idx_coeff_z] * x_physical +
                                                                 basis_dval [idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho] * basis_val[idx_grid_z   * bd->nb_coeffs_z + idx_coeff_z] +
                                                                 basis_d2val[idx_grid_z   * bd->nb_coeffs_z   + idx_coeff_z]   * basis_val[idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho] * x_physical +
                                                                 basis_val  [idx_grid_rho * bd->nb_coeffs_rho + idx_coeff_rho] * basis_val[idx_grid_z   * bd->nb_coeffs_z + idx_coeff_z] * 0.25 * d2q * x_physical;
                }
            rhs->data[idx_grid] = -0.25 * d2q * x_physical;
        }
    }

    free(basis_val);
    free(basis_dval);
    free(basis_d2val);

    return 0;
}

static int brill_solve_linear(gsl_matrix *mat, gsl_vector **px, gsl_vector **prhs)
{
    int *ipiv;
    gsl_matrix *mat_f;
    double cond, ferr, berr, rpivot;

    gsl_vector   *x = *px;
    gsl_vector *rhs = *prhs;
    char      equed = 'N';
    int         ret = 0;

    ipiv  = malloc(mat->size1 * sizeof(*ipiv));
    mat_f = gsl_matrix_alloc(mat->size1, mat->size2);
    if (!ipiv || !mat_f) {
        ret = -ENOMEM;
        goto fail;
    }

    LAPACKE_dgesvx(LAPACK_COL_MAJOR, 'N', 'N', mat->size1, 1,
                   mat->data, mat->tda, mat_f->data, mat_f->tda, ipiv, &equed,
                   NULL, NULL, rhs->data, rhs->size, x->data, x->size,
                   &cond, &ferr, &berr, &rpivot);

    fprintf(stderr, "LU factorization solution to a %zdx%zd matrix: "
            "condition number %16.16g; forward error %16.16g backward error %16.16g\n",
            mat->size1, mat->size2, cond, ferr, berr);
fail:
    gsl_matrix_free(mat_f);
    free(ipiv);

    return ret;
}

static int brill_solve_svd(gsl_matrix *mat, gsl_vector **px, gsl_vector **prhs)
{
    double *sv;
    int rank;

    gsl_vector   *x = *px;
    gsl_vector *rhs = *prhs;

    sv = malloc(sizeof(*sv) * x->size);
    if (!sv)
        return -ENOMEM;

    LAPACKE_dgelsd(LAPACK_COL_MAJOR, rhs->size, x->size, 1, mat->data, mat->tda,
                   rhs->data, rhs->size, sv, -1, &rank);

    fprintf(stderr, "Least squares SVD solution to a %zdx%zd matrix: "
            "rank %d, condition number %16.16g\n", mat->size1, mat->size2,
            rank, sv[x->size - 1] / sv[0]);

    gsl_vector_free(x);

    *px         = rhs;
    (*px)->size = mat->size1;
    *prhs       = NULL;

    free(sv);

    return 0;
}

/*
 * Solve the equation
 * Δψ + ¼ ψ (∂²q/∂r² + ∂²q/∂z²) = 0
 * for the coefficients of spectral approximation of ψ:
 * ψ(r, z) = 1 + ΣaᵢⱼTᵢ(r)Tⱼ(z)
 * where i =  { 0, ... , bd->nb_coeffs_rho };
 *       j =  { 0, ... , bd->nb_coeffs_z };
 * q(r, z) and Tᵢ(x) are defined by bd->q_func, bd->laplace_q_func and
 * bd->basis.
 *
 * The cofficients are exported in bd->priv->coeffs
 */
static int brill_solve(BDContext *bd)
{
    BDPriv *s = bd->priv;
    gsl_matrix *mat = NULL;
    gsl_vector *coeffs = NULL, *rhs = NULL;

#if ACC_TEST
    gsl_vector *rhs_bkp = NULL;
    gsl_matrix *mat_bkp = NULL;
#endif

    int ret = 0;

    /* allocate and fill the pseudospectral matrix */
    mat    = gsl_matrix_alloc (s->nb_coeffs, s->nb_colloc_points); // inverted order for lapack
    coeffs = gsl_vector_alloc (s->nb_coeffs);
    rhs    = gsl_vector_calloc(s->nb_colloc_points);
    if (!mat || !coeffs || !rhs) {
        ret = -ENOMEM;
        goto fail;
    }

    /* fill the matrix */
    ret = brill_construct_matrix(bd, mat, rhs);
    if (ret < 0)
        goto fail;

#if ACC_TEST
    /* make backups of the matrix and RHS, since they might be destroyed later */
    mat_bkp = gsl_matrix_alloc(mat->size2, mat->size1);
    rhs_bkp = gsl_vector_alloc(rhs->size);
    if (!mat_bkp || !rhs_bkp) {
        ret = -ENOMEM;
        goto fail;
    }

    gsl_vector_memcpy(rhs_bkp, rhs);
    gsl_matrix_transpose_memcpy(mat_bkp, mat);
#endif

    /* solve for the coeffs */
    if (s->nb_colloc_points > s->nb_coeffs)
        ret = brill_solve_svd(mat, &coeffs, &rhs);
    else
        ret = brill_solve_linear(mat, &coeffs, &rhs);

    /* export the result */
    s->coeffs = coeffs;

#if ACC_TEST
    {
        double min, max, rmin, rmax;

        gsl_vector_minmax(rhs_bkp, &rmin, &rmax);
        gsl_blas_dgemv(CblasNoTrans, 1.0, mat_bkp, coeffs, -1.0, rhs_bkp);
        gsl_vector_minmax(rhs_bkp, &min, &max);

        fprintf(stderr, "max(|A·x - rhs|) / max(|rhs|): %16.16g\n",
                MAX(fabs(min), fabs(max)) / MAX(fabs(rmin), fabs(rmax)));
    }
#endif

fail:

#if ACC_TEST
    gsl_matrix_free(mat_bkp);
    gsl_vector_free(rhs_bkp);
#endif

    if (ret < 0)
        gsl_vector_free(coeffs);

    gsl_vector_free(rhs);
    gsl_matrix_free(mat);

    return ret;
}

// q function form from PHYSICAL REVIEW D 88, 103009 (2013)
// with σ and ρ_0 hardcoded to 1 for now
static double q_gundlach(BDContext *bd, double rho, double z)
{
    return bd->amplitude * SQR(rho) * exp(- (SQR(rho) + SQR(z)));
}

static double dq_rho_gundlach(BDContext *bd, double rho, double z)
{
    return bd->amplitude * 2 * rho * exp(-SQR(rho) - SQR(z)) * (1 - SQR(rho));
}

static double d2q_rho_gundlach(BDContext *bd, double rho, double z)
{
    double rho2 = SQR(rho);
    return bd->amplitude * 2 * exp(-rho2 - SQR(z)) * ((1 - rho2) * (1 - 2 * rho2) - 2 * rho2);
}

static double d2q_rho_z_gundlach(BDContext *bd, double rho, double z)
{
    return -bd->amplitude * 4 * z * rho * exp(-SQR(rho) - SQR(z)) * (1 - SQR(rho));
}

static double dq_z_gundlach(BDContext *bd, double rho, double z)
{
    return - bd->amplitude * 2 * z * SQR(rho) * exp(-SQR(rho) - SQR(z));
}

static double d2q_z_gundlach(BDContext *bd, double rho, double z)
{
    return bd->amplitude * 2 * SQR(rho) * exp(-SQR(rho) - SQR(z)) * (2 * SQR(z) - 1);
}

static const QFunc q_func_gundlach = {
    .q         = q_gundlach,
    .dq_rho    = dq_rho_gundlach,
    .dq_z      = dq_z_gundlach,
    .d2q_rho   = d2q_rho_gundlach,
    .d2q_z     = d2q_z_gundlach,
    .d2q_rho_z = d2q_rho_z_gundlach,
};

static double q_eppley(BDContext *bd, double rho, double z)
{
    double r2 = SQR(rho) + SQR(z);
    return bd->amplitude * SQR(rho) / (1.0 + pow(r2, bd->eppley_n / 2.0));
}

static double dq_rho_eppley(BDContext *bd, double rho, double z)
{
    double A = bd->amplitude;
    double n = bd->eppley_n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm2 = pow(r, n - 2);

    return A * rho * (2 * (1 + rnm2 * r2) - n * rho2 * rnm2) / SQR(1 + rnm2 * r2);
}

static double dq_z_eppley(BDContext *bd, double rho, double z)
{
    double A = bd->amplitude;
    double n = bd->eppley_n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm2 = pow(r, n - 2);

    return - A * n * rho2 * z * rnm2 / SQR(1 + rnm2 * r2);
}

static double d2q_rho_eppley(BDContext *bd, double rho, double z)
{
    double A = bd->amplitude;
    double n = bd->eppley_n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm4 = pow(r, n - 4);
    double rn   = rnm4 * SQR(r2);
    double rnp1 = 1 + rn;

    return A * (SQR(n) * SQR(rho2) * rnm4 * (rn - 1) - n * rho2 * rnm4 * (3 * rho2 + 5 * z2) * rnp1 + 2 * SQR(rnp1)) / (rnp1 * SQR(rnp1));
}

static double d2q_z_eppley(BDContext *bd, double rho, double z)
{
    double A = bd->amplitude;
    double n = bd->eppley_n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm4 = pow(r, n - 4);
    double rn   = rnm4 * SQR(r2);
    double rnp1 = 1 + rn;

    return A * n * rho2 * rnm4 * (z2 - rho2 - n * z2 + rn * ((1 + n) * z2 - rho2)) / (rnp1 * SQR(rnp1));
}

static double d2q_rho_z_eppley(BDContext *bd, double rho, double z)
{
    double A = bd->amplitude;
    double n = bd->eppley_n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm4 = pow(r, n - 4);
    double rn   = rnm4 * SQR(r2);
    double rnp1 = 1 + rn;

    return A * n * rho * z * rnm4 * (rn * (n * rho2 - 2 * z2) - 2 * z2 - n * rho2) / (rnp1 * SQR(rnp1));
}

static const QFunc q_func_eppley = {
    .q         = q_eppley,
    .dq_rho    = dq_rho_eppley,
    .dq_z      = dq_z_eppley,
    .d2q_rho   = d2q_rho_eppley,
    .d2q_z     = d2q_z_eppley,
    .d2q_rho_z = d2q_rho_z_eppley,
};

static double d2q_eppley(BDContext *bd, double rho, double z)
{
    double rho2 = SQR(rho);
    double z2   = SQR(z);
    double rho4 = SQR(rho2);
    double z4 = SQR(z2);
    double z6 = z4 * z2;

    double r2 = rho2 + z2;
    double r = sqrt(r2);

    return bd->amplitude * (2.0 + r2 * (7 * rho4 * rho4 + 23 * rho4 * rho2 * z2 + 27 * rho4 * z4 + rho2 * (13 * z6 - 41 * r) + 2 * z2 * (z6 + 2 * r))) / pow(1.0 + r2 * r2 * r, 3);
}

static int brill_init_check_options(BDContext *bd)
{
    BDPriv *s = bd->priv;

    switch (bd->q_func_type) {
    case BD_Q_FUNC_GUNDLACH:
        s->q_func   = &q_func_gundlach;
        break;
    case BD_Q_FUNC_EPPLEY:
        if (bd->eppley_n < 4) {
            fprintf(stderr, "Invalid n: %d < 4\n", bd->eppley_n);
            return -EINVAL;
        }

        s->q_func   = &q_func_eppley;
        break;
    default:
        fprintf(stderr, "Unknown q function type: %d\n", bd->q_func_type);
        return -EINVAL;
    }

    if (!bd->nb_coeffs_z)
        bd->nb_coeffs_z = bd->nb_coeffs_rho;

    if (bd->nb_coeffs_rho != bd->nb_coeffs_z) {
        fprintf(stderr, "Different ρ and z basis orders are not supported.\n");
        return -EINVAL;
    }

    s->basis = &sb_even_basis;

    s->nb_coeffs = bd->nb_coeffs_rho * bd->nb_coeffs_z;

    s->nb_colloc_points_rho = bd->nb_coeffs_rho + bd->overdet_rho;
    s->nb_colloc_points_z   = bd->nb_coeffs_z   + bd->overdet_z;

    s->nb_colloc_points     = s->nb_colloc_points_rho * s->nb_colloc_points_z;

    s->colloc_grid_order_rho = bd->nb_coeffs_rho + bd->colloc_grid_offset_rho;
    s->colloc_grid_order_z   = bd->nb_coeffs_z   + bd->colloc_grid_offset_z;

    return 0;
}

int bd_solve(BDContext *bd)
{
    BDPriv *s = bd->priv;
    int ret;

    if (bd->psi_minus1_coeffs) {
        fprintf(stderr, "Solve called multiple times on the same context.\n");
        return -EINVAL;
    }

    ret = brill_init_check_options(bd);
    if (ret < 0)
        return ret;

    ret = brill_solve(bd);
    if (ret < 0)
        return ret;

    bd->psi_minus1_coeffs = s->coeffs->data;
    bd->stride            = bd->nb_coeffs_rho;

    return 0;
}

BDContext *bd_context_alloc(void)
{
    BDContext *bd = calloc(1, sizeof(*bd));

    if (!bd)
        return NULL;

    bd->q_func_type = BD_Q_FUNC_GUNDLACH;
    bd->amplitude   = 1.0;
    bd->eppley_n    = 5;

    bd->nb_coeffs_rho = 80;
    bd->nb_coeffs_z   = 0;

    bd->colloc_grid_offset_rho = 3;
    bd->colloc_grid_offset_z   = 3;

    bd->basis_scale_factor_rho = 3.0;
    bd->basis_scale_factor_z   = 3.0;

    bd->priv = calloc(1, sizeof(BDPriv));
    if (!bd->priv) {
        free(bd);
        return NULL;
    }

    return bd;
}

void bd_context_free(BDContext **pbd)
{
    BDContext *bd = *pbd;
    BDPriv *s;

    if (!bd)
        return;

    s = bd->priv;

    gsl_vector_free(s->coeffs);

    free(bd->priv);
    free(bd);
    *pbd = NULL;
}

int bd_eval_psi(BDContext *bd, const double *rho, int nb_coords_rho,
                const double *z, int nb_coords_z,
                const unsigned int diff_order[2],
                double *psi)
{
    BDPriv *s = bd->priv;
    const long double sfr = bd->basis_scale_factor_rho, sfz = bd->basis_scale_factor_z;
    long double (*eval)(long double coord, int idx, long double sf);

    long double *basis_val_rho, *basis_val_z;

    long double add = 0.0;

    if (diff_order[0] > 2 || diff_order[1] > 2) {
        fprintf(stderr, "Derivatives of higher order than 2 are not supported.\n");
        return -ENOSYS;
    }

    if (diff_order[0] == 0 && diff_order[1] == 0)
        add = 1.0;

    /* precompute the basis values on the grid points */
    basis_val_rho = malloc(sizeof(*basis_val_rho) * bd->nb_coeffs_rho * nb_coords_rho);
    basis_val_z   = malloc(sizeof(*basis_val_z)   * bd->nb_coeffs_z   * nb_coords_z);

    switch (diff_order[0]) {
    case 0: eval = s->basis->eval;       break;
    case 1: eval = s->basis->eval_diff1; break;
    case 2: eval = s->basis->eval_diff2; break;
    }

    for (int i = 0; i < nb_coords_rho; i++) {
        double rrho = rho[i];
        for (int j = 0; j < bd->nb_coeffs_rho; j++)
            basis_val_rho[i * bd->nb_coeffs_rho + j] = eval(rrho, j, sfr);
    }

    switch (diff_order[1]) {
    case 0: eval = s->basis->eval;       break;
    case 1: eval = s->basis->eval_diff1; break;
    case 2: eval = s->basis->eval_diff2; break;
    }
    for (int i = 0; i < nb_coords_z; i++) {
        double zz   = z[i];
        for (int j = 0; j < bd->nb_coeffs_z; j++)
            basis_val_z[i * bd->nb_coeffs_z + j] = eval(zz, j, sfz);
    }

    for (int i = 0; i < nb_coords_z; i++)
        for (int j = 0; j < nb_coords_rho; j++) {
            long double ppsi = add;

            for (int m = 0; m < bd->nb_coeffs_z; m++)
                for (int n = 0; n < bd->nb_coeffs_rho; n++)
                    ppsi += s->coeffs->data[m * bd->nb_coeffs_rho + n] * basis_val_rho[j * bd->nb_coeffs_rho + n] * basis_val_z[i * bd->nb_coeffs_z + m];

            psi[i * nb_coords_rho + j] = ppsi;
        }

    free(basis_val_rho);
    free(basis_val_z);

    return 0;

}

int bd_eval_metric(BDContext *bd, const double *rho, int nb_coords_rho,
                   const double *z, int nb_coords_z,
                   const unsigned int comp[2], const unsigned int diff_order[2],
                   double *out)
{
    BDPriv *s = bd->priv;
    const int nb_coords = nb_coords_rho * nb_coords_z;
    double *psi = NULL, *dpsi_rho = NULL, *dpsi_z = NULL, *d2psi_rho = NULL, *d2psi_rho_z = NULL, *d2psi_z = NULL;
    int ret = 0;

    /* check the parameters for validity */
    if (comp[0] > 2 || comp[1] > 2) {
        fprintf(stderr, "Invalid component indices: %d%d\n", comp[0], comp[1]);
        return -EINVAL;
    }

    if (comp[0] != comp[1]) {
        memset(out, 0, nb_coords * sizeof(*out));
        return 0;
    }

    if (diff_order[0] + diff_order[1] > 2) {
        fprintf(stderr, "At most second order derivatives are supported.\n");
        return -ENOSYS;
    }

    /* evaluate the conformal factor and its derivatives if necessary */
#define EVAL_PSI(arr, order)                                    \
do {                                                            \
    arr = malloc(nb_coords * sizeof(*arr));                     \
    if (!arr) {                                                 \
        ret = -ENOMEM;                                          \
        goto fail;                                              \
    }                                                           \
                                                                \
    ret = bd_eval_psi(bd, rho, nb_coords_rho, z, nb_coords_z,   \
                      order, arr);                              \
    if (ret < 0)                                                \
        goto fail;                                              \
} while (0)

    EVAL_PSI(psi,             ((unsigned int[2]){ 0, 0 }));
    if (diff_order[0])
        EVAL_PSI(dpsi_rho,    ((unsigned int[2]){ 1, 0 }));
    if (diff_order[0] > 1)
        EVAL_PSI(d2psi_rho,   ((unsigned int[2]){ 2, 0 }));
    if (diff_order[1])
        EVAL_PSI(dpsi_z,      ((unsigned int[2]){ 0, 1 }));
    if (diff_order[1] > 1)
        EVAL_PSI(d2psi_z,     ((unsigned int[2]){ 0, 2 }));
    if (diff_order[0] && diff_order[1])
        EVAL_PSI(d2psi_rho_z, ((unsigned int[2]){ 1, 1 }));

/* the template for the loop over the grid points */
#define EVAL_LOOP(DO_EVAL)                                          \
do {                                                                \
    for (int i = 0; i < nb_coords_z; i++) {                         \
        const double zz = z[i];                                     \
                                                                    \
        for (int j = 0; j < nb_coords_rho; j++) {                   \
            const double rrho = rho[j];                             \
            const double ppsi = psi[i * nb_coords_rho + j];         \
            const double psi2 = SQR(ppsi);                          \
            const double psi3 = psi2 * ppsi;                        \
            const double psi4 = SQR(psi2);                          \
            double val;                                             \
                                                                    \
            DO_EVAL;                                                \
                                                                    \
            out[i * nb_coords_rho + j] = val;                       \
        }                                                           \
    }                                                               \
} while (0)

#define GRID(arr) (arr[i * nb_coords_rho + j])

    if (comp[0] == 2) {
        // γ_φφ
        switch (diff_order[0]) {
        case 0:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(val = SQR(rrho) * psi4);                                                                                                                    break;
            case 1: EVAL_LOOP(val = 4 * SQR(rrho) * psi3 * GRID(dpsi_z));                                                                                                 break; // ∂_z
            case 2: EVAL_LOOP(val = 4 * SQR(rrho) * (GRID(d2psi_z) * psi3 + 3 * psi2 * SQR(GRID(dpsi_z))));                                                               break; // ∂_z ∂_z
            }
            break;
        case 1:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(val = 2 * rrho * psi4 + 4 * SQR(rrho) * psi3 * GRID(dpsi_rho));                                                                             break; // ∂_ρ
            case 1: EVAL_LOOP(val = 8 * rrho * psi3 * GRID(dpsi_z) + 12 * SQR(rrho) * psi2 * GRID(dpsi_z) * GRID(dpsi_rho) + 4 * SQR(rrho) * psi3 * GRID(d2psi_rho_z));   break; // ∂_ρ ∂_z
            }
            break;
        case 2:     EVAL_LOOP(val = 4 * SQR(rrho) * psi3 * GRID(d2psi_rho) + 12 * SQR(rrho * ppsi * GRID(dpsi_rho)) + 16 * rrho * psi3 * GRID(dpsi_rho) + 2 * psi4);      break; // ∂_ρ ∂_ρ
        }
    } else {
        // γ_ρρ / γ_zz

/* a wrapper around the actual evaluation expression that provides the q function and its
 * derivatives as needed */
#define DO_EVAL_Q_WRAPPER(DO_EVAL, eval_drho, eval_dz, eval_d2rho, eval_d2z, eval_d2rhoz)                                                    \
do {                                                                                                                                         \
    const double base = psi4 * exp(2 * s->q_func->q(bd, rrho, zz));                                                                          \
    double drho, dz, d2rho, d2z, d2rho_z;                                                                                                    \
                                                                                                                                             \
    if (eval_drho)   drho    = s->q_func->dq_rho(bd, rrho, zz)    + 2 * GRID(dpsi_rho)    / ppsi;                                            \
    if (eval_dz)     dz      = s->q_func->dq_z(bd, rrho, zz)      + 2 * GRID(dpsi_z)      / ppsi;                                            \
    if (eval_d2rho)  d2rho   = s->q_func->d2q_rho(bd, rrho, zz)   + 2 * GRID(d2psi_rho)   / ppsi - 2 * SQR(GRID(dpsi_rho) / ppsi);           \
    if (eval_d2z)    d2z     = s->q_func->d2q_z(bd, rrho, zz)     + 2 * GRID(d2psi_z)     / ppsi - 2 * SQR(GRID(dpsi_z)   / ppsi);           \
    if (eval_d2rhoz) d2rho_z = s->q_func->d2q_rho_z(bd, rrho, zz) + 2 * GRID(d2psi_rho_z) / ppsi - 2 * GRID(dpsi_rho) * GRID(dpsi_z) / psi2; \
                                                                                                                                             \
    DO_EVAL;                                                                                                                                 \
} while (0)

        switch (diff_order[0]) {
        case 0:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = base,                                      0, 0, 0, 0, 0)); break;
            case 1: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 2 * base * dz,                             0, 1, 0, 0, 0)); break;        // ∂_z
            case 2: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * SQR(dz) * base + 2 * base * d2z,       0, 1, 0, 1, 0)); break;        // ∂_z ∂_z
            }
            break;
        case 1:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 2 * base * drho,                           1, 0, 0, 0, 0)); break;        // ∂_ρ
            case 1: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * base * drho * dz + 2 * base * d2rho_z, 1, 1, 0, 0, 1)); break;        // ∂_ρ ∂_z
            }
            break;
        case 2:     EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * SQR(drho) * base + 2 * base * d2rho,   1, 0, 1, 0, 0)); break;        // ∂_ρ ∂_ρ
        }
    }

fail:
    free(psi);
    free(dpsi_rho);
    free(dpsi_z);
    free(d2psi_rho);
    free(d2psi_rho_z);
    free(d2psi_z);

    return ret;
}

#if 0
void brill_data(CCTK_ARGUMENTS)
{
    static BrillDataContext *prev_bd;

    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    BrillDataContext *bd;

    long double *basis_val_r, *basis_val_z;

    int64_t grid_size = CCTK_GFINDEX3D(cctkGH,
                                       cctk_lsh[0] - 1,
                                       cctk_lsh[1] - 1,
                                       cctk_lsh[2] - 1) + 1;

    /* on the first run, solve the equation for ψ */
    if (!prev_bd) {
        bd = init_bd(cctkGH, amplitude, eppley_n, basis_order_r, basis_order_z,
                     colloc_grid_offset_r, colloc_grid_offset_z,
                     scale_factor, overdet);

        brill_solve(bd);

        if (export_coeffs)
            memcpy(brill_coeffs, bd->psi_coeffs->data, sizeof(*brill_coeffs) * bd->nb_coeffs);

        prev_bd = bd;
    } else
        bd = prev_bd;

    memset(kxx, 0, sizeof(*kxx) * grid_size);
    memset(kyy, 0, sizeof(*kyy) * grid_size);
    memset(kzz, 0, sizeof(*kzz) * grid_size);
    memset(kxy, 0, sizeof(*kxy) * grid_size);
    memset(kxz, 0, sizeof(*kxz) * grid_size);
    memset(kyz, 0, sizeof(*kyz) * grid_size);
    memset(gxz, 0, sizeof(*kxz) * grid_size);
    memset(gyz, 0, sizeof(*kyz) * grid_size);

    /* precompute the basis values on the grid points */
    basis_val_r = malloc(sizeof(*basis_val_r) * bd->nb_coeffs_rho * cctk_lsh[1] * cctk_lsh[0]);
    basis_val_z = malloc(sizeof(*basis_val_z) * bd->nb_coeffs_z * cctk_lsh[2]);

    for (int i = 0; i < cctk_lsh[2]; i++) {
        CCTK_REAL zz = z[CCTK_GFINDEX3D(cctkGH, 0, 0, i)];
        for (int j = 0; j < bd->nb_coeffs_z; j++)
            basis_val_z[i * bd->nb_coeffs_z + j] = bd->basis->eval(zz, j, bd->basis_scale_factor_rho);
    }

    for (int j = 0; j < cctk_lsh[1]; j++)
        for (int i = 0; i < cctk_lsh[0]; i++) {
            CCTK_REAL xx = x[CCTK_GFINDEX3D(cctkGH, i, j, 0)];
            CCTK_REAL yy = y[CCTK_GFINDEX3D(cctkGH, i, j, 0)];
            CCTK_REAL r = sqrt(SQR(xx) + SQR(yy));

            for (int k = 0; k < bd->nb_coeffs_rho; k++)
                basis_val_r[(j * cctk_lsh[0] + i) * bd->nb_coeffs_rho + k] = bd->basis->eval(r, k, bd->basis_scale_factor_rho);
        }

#pragma omp parallel for
    for (int k = 0; k < cctk_lsh[2]; k++)
        for (int j = 0; j < cctk_lsh[1]; j++)
            for (int i = 0; i < cctk_lsh[0]; i++) {
                int    index = CCTK_GFINDEX3D(cctkGH, i, j, k);
                long double xx = x[index], yy = y[index], zz = z[index];

                long double r2 = SQR(xx) + SQR(yy);
                long double r  = sqrtl(r2);

                long double q   = bd->q_func(bd, r, zz);
                long double e2q = expl(2 * q);

                long double psi = 1.0, psi4;

                for (int n = 0; n < bd->nb_coeffs_z; n++)
                    for (int m = 0; m < bd->nb_coeffs_rho; m++)
                        psi += bd->psi_coeffs->data[n * bd->nb_coeffs_z + m] * basis_val_r[(j * cctk_lsh[0] + i) * bd->nb_coeffs_rho + m] * basis_val_z[k * bd->nb_coeffs_z + n];

                psi4 = SQR(SQR(psi));

                if (r2 < EPS)
                    r2 = EPS;

                gxx[index] = psi4 * (e2q + (1 - e2q) * SQR(yy) / r2);
                gyy[index] = psi4 * (e2q + (1 - e2q) * SQR(xx) / r2);
                gzz[index] = psi4 * e2q;
                gxy[index] = psi4 * (-(1.0 - e2q) * xx * yy / r2);
            }
}
#endif