#include <errno.h>
#include <float.h>
#include <limits.h>
#include <math.h>
#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

#include "cctk.h"
#include "cctk_Arguments.h"
#include "cctk_Parameters.h"
#include "util_Table.h"
#include "Symmetry.h"

#define SQR(x) ((x) * (x))
#define SGN(x) ((x) >= 0.0 ? 1.0 : -1.0)
#define MAX(x, y) ((x) > (y) ? (x) : (y))
#define MIN(x, y) ((x) > (y) ? (y) : (x))
#define ABS(x) ((x >= 0) ? (x) : -(x))
#define ARRAY_ELEMS(arr) (sizeof(arr) / sizeof(*arr))

/* indices (in our code, not cactus structs) of the grid functions which we'll need to
 * interpolate on the photon positions */
enum MetricVars {
    GTXX = 0,
    GTYY,
    GTZZ,
    GTXY,
    GTXZ,
    GTYZ,
    ATXX,
    ATYY,
    ATZZ,
    ATXY,
    ATXZ,
    ATYZ,
    PHI,
    ALPHA,
    TRK,
    NB_METRIC_VARS,
};

/* indices of the interpolated values of the above grid functions and their derivatives */
enum InterpMetricVars {
    I_GTXX = 0,
    I_GTYY,
    I_GTZZ,
    I_GTXY,
    I_GTXZ,
    I_GTYZ,
    I_GTXX_DX,
    I_GTYY_DX,
    I_GTZZ_DX,
    I_GTXZ_DX,
    I_GTXX_DZ,
    I_GTYY_DZ,
    I_GTZZ_DZ,
    I_GTXZ_DZ,
    I_ATXX,
    I_ATYY,
    I_ATZZ,
    I_ATXY,
    I_ATXZ,
    I_ATYZ,
    I_PHI,
    I_PHI_DX,
    I_PHI_DZ,
    I_ALPHA,
    I_ALPHA_DX,
    I_ALPHA_DZ,
    I_TRK,
    NB_INTERP_VARS,
};

typedef struct NSContext {
    cGH *gh;

    int reflevel;
    double dt;

    int nb_surfaces;
    int photons_per_surface;

    // interpolation parameters
    int coord_system;
    int interp_operator;
    int interp_params;

    int        interp_vars_indices[NB_METRIC_VARS];
    double    *interp_values[NB_INTERP_VARS];
    int        interp_value_codes[NB_INTERP_VARS];

    double    *rk_scratch[4];
} NSContext;

static NSContext ns_ctx;

/* mapping between our indices and thorn names */
static const char *metric_vars[] = {
    [GTXX]  = "ML_BSSN::gt11",
    [GTYY]  = "ML_BSSN::gt22",
    [GTZZ]  = "ML_BSSN::gt33",
    [GTXY]  = "ML_BSSN::gt12",
    [GTXZ]  = "ML_BSSN::gt13",
    [GTYZ]  = "ML_BSSN::gt23",
    [ATXX]  = "ML_BSSN::At11",
    [ATYY]  = "ML_BSSN::At22",
    [ATZZ]  = "ML_BSSN::At33",
    [ATXY]  = "ML_BSSN::At12",
    [ATXZ]  = "ML_BSSN::At13",
    [ATYZ]  = "ML_BSSN::At23",
    [PHI]   = "ML_BSSN::phi",
    [ALPHA] = "ML_BSSN::alpha",
    [TRK]   = "ML_BSSN::trK",
};

/* mapping between the cactus grid values and interpolated values */
static const int interp_operation_indices[] = {
    [I_GTXX]     = GTXX,
    [I_GTYY]     = GTYY,
    [I_GTZZ]     = GTZZ,
    [I_GTXY]     = GTXY,
    [I_GTXZ]     = GTXZ,
    [I_GTYZ]     = GTYZ,
    [I_GTXX_DX]  = GTXX,
    [I_GTYY_DX]  = GTYY,
    [I_GTZZ_DX]  = GTZZ,
    [I_GTXZ_DX]  = GTXZ,
    [I_GTXX_DZ]  = GTXX,
    [I_GTYY_DZ]  = GTYY,
    [I_GTZZ_DZ]  = GTZZ,
    [I_GTXZ_DZ]  = GTXZ,
    [I_ATXX]     = ATXX,
    [I_ATYY]     = ATYY,
    [I_ATZZ]     = ATZZ,
    [I_ATXY]     = ATXY,
    [I_ATXZ]     = ATXZ,
    [I_ATYZ]     = ATYZ,
    [I_PHI]      = PHI,
    [I_PHI_DX]   = PHI,
    [I_PHI_DZ]   = PHI,
    [I_ALPHA]    = ALPHA,
    [I_ALPHA_DX] = ALPHA,
    [I_ALPHA_DZ] = ALPHA,
    [I_TRK]      = TRK,
};

/* the operation (plain value or x/y/z-derivative) to apply during interpolation */
static const CCTK_INT interp_operation_codes[] = {
    [I_GTXX]     = 0,
    [I_GTYY]     = 0,
    [I_GTZZ]     = 0,
    [I_GTXY]     = 0,
    [I_GTXZ]     = 0,
    [I_GTYZ]     = 0,
    [I_GTXX_DX]  = 1,
    [I_GTYY_DX]  = 1,
    [I_GTZZ_DX]  = 1,
    [I_GTXZ_DX]  = 1,
    [I_GTXX_DZ]  = 3,
    [I_GTYY_DZ]  = 3,
    [I_GTZZ_DZ]  = 3,
    [I_GTXZ_DZ]  = 3,
    [I_ATXX]     = 0,
    [I_ATYY]     = 0,
    [I_ATZZ]     = 0,
    [I_ATXY]     = 0,
    [I_ATXZ]     = 0,
    [I_ATYZ]     = 0,
    [I_PHI]      = 0,
    [I_PHI_DX]   = 1,
    [I_PHI_DZ]   = 3,
    [I_ALPHA]    = 0,
    [I_ALPHA_DX] = 1,
    [I_ALPHA_DZ] = 3,
    [I_TRK]      = 0,
};

static int ctz(int a)
{
    int ret = 0;

    if (!a)
        return INT_MAX;

    while (!(a & 1)) {
        a >>= 1;
        ret++;
    }

    return ret;
}

static void surface_launch(NSContext *ns, double *pos_x, double *pos_z, double *mom_x, double *mom_z)
{
    const int origin_idx = CCTK_GFINDEX3D(ns->gh, ns->gh->cctk_nghostzones[0], 0, ns->gh->cctk_nghostzones[2]);

    double gtxx  = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[GTXX]))[origin_idx];
    double gtyy  = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[GTYY]))[origin_idx];
    double gtzz  = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[GTZZ]))[origin_idx];
    double gtxy  = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[GTXY]))[origin_idx];
    double gtxz  = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[GTXZ]))[origin_idx];
    double gtyz  = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[GTYZ]))[origin_idx];
    double phi   = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[PHI]))[origin_idx];
    double alpha = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[ALPHA]))[origin_idx];

    const double g[3][3] = { { gtxx / SQR(phi), gtxy / SQR(phi), gtxz / SQR(phi) },
                             { gtxy / SQR(phi), gtyy / SQR(phi), gtyz / SQR(phi) },
                             { gtxz / SQR(phi), gtyz / SQR(phi), gtzz / SQR(phi) }};

    /**
     * send out photons along geodesics evenly spaced in angles θ between x- and z-axis
     * p^μ = { 1, p^x, 0, p^z }
     * g_μν p^μ p^ν = 0
     * p^x = p^r sinθ
     * p^z = p^r cosθ
     */
#pragma omp parallel for
    for (int i = 0; i < ns->photons_per_surface; i++) {
        const double theta = (double)i / (ns->photons_per_surface - 1) * M_PI / 2;
        const double st = sin(theta);
        const double ct = cos(theta);

        const double a = g[0][0] * SQR(st) + g[2][2] * SQR(ct) + 2.0 * g[0][2] * st * ct;
        const double b = 0.0; // assume shift is zero
        const double c = -SQR(alpha);

        const double p_r = (-b + sqrt(SQR(b) - 4 * a * c)) / (2 * a);

        pos_x[i] = 0.0;
        pos_z[i] = 0.0;

        mom_x[i] = st * p_r;
        mom_z[i] = ct * p_r;
    }
}

static void null_geodesic_rhs(size_t len,
                              double *dpos_dt_x, double *dpos_dt_z,
                              double *dmom_dt_x, double *dmom_dt_z,
                              const double *pos_x, const double *pos_z,
                              const double *mom_x, const double *mom_z,
                              const double * const metric_vars[NB_INTERP_VARS])
{
#pragma omp parallel for
    for (int idx = 0; idx < len; idx++) {
        const double x = pos_x[idx];

        const int zaxis = x <= 1e-12;

        const double gtxx = metric_vars[I_GTXX][idx];
        const double gtyy = metric_vars[I_GTYY][idx];
        const double gtzz = metric_vars[I_GTZZ][idx];
        const double gtxy = metric_vars[I_GTXY][idx];
        const double gtxz = metric_vars[I_GTXZ][idx];
        const double gtyz = metric_vars[I_GTYZ][idx];

        const double gt[3][3] = {{ gtxx, gtxy, gtxz },
                                 { gtxy, gtyy, gtyz },
                                 { gtxz, gtyz, gtzz }};

        const double dx_gt11 = metric_vars[I_GTXX_DX][idx];
        const double dx_gt22 = metric_vars[I_GTYY_DX][idx];
        const double dx_gt33 = metric_vars[I_GTZZ_DX][idx];
        const double dx_gt13 = metric_vars[I_GTXZ_DX][idx];

        const double dz_gt11 = metric_vars[I_GTXX_DZ][idx];
        const double dz_gt22 = metric_vars[I_GTYY_DZ][idx];
        const double dz_gt33 = metric_vars[I_GTZZ_DZ][idx];
        const double dz_gt13 = metric_vars[I_GTXZ_DZ][idx];

        const double dgt[3][3][3] = {
            {
                { dx_gt11,     0.0, dx_gt13 },
                {     0.0, dx_gt22,     0.0 },
                { dx_gt13,     0.0, dx_gt33 },
            },
            {
                {     0.0, zaxis ? dx_gt11 - dx_gt22 : (gtxx - gtyy) / x, 0.0 },
                { zaxis ? dx_gt11 - dx_gt22 : (gtxx - gtyy) / x, 0.0, zaxis ? dx_gt13 : gtxz / x },
                { 0.0, zaxis ? dx_gt13 : gtxz / x, 0.0 },
            },
            {
                { dz_gt11,     0.0, dz_gt13 },
                {     0.0, dz_gt22,     0.0 },
                { dz_gt13,     0.0, dz_gt33 },
            },
        };

        const double Atxx = metric_vars[I_ATXX][idx];
        const double Atyy = metric_vars[I_ATYY][idx];
        const double Atzz = metric_vars[I_ATZZ][idx];
        const double Atxy = metric_vars[I_ATXY][idx];
        const double Atxz = metric_vars[I_ATXZ][idx];
        const double Atyz = metric_vars[I_ATYZ][idx];

        const double At[3][3] = {{ Atxx, Atxy, Atxz },
                                 { Atxy, Atyy, Atyz },
                                 { Atxz, Atyz, Atzz }};

        const double trK    = metric_vars[I_TRK][idx];

        const double alpha     = metric_vars[I_ALPHA][idx];
        const double dx_alpha  = metric_vars[I_ALPHA_DX][idx];
        const double dz_alpha  = metric_vars[I_ALPHA_DZ][idx];

        const double phi    = metric_vars[I_PHI][idx];
        const double dx_phi = metric_vars[I_PHI_DX][idx];
        const double dz_phi = metric_vars[I_PHI_DZ][idx];

        const double dphi[3] = { dx_phi, 0.0, dz_phi };

        double g[3][3], gu[3][3], gtu[3][3], dg[3][3][3];
        double K[3][3];

        double g4[4][4], gu4[4][4], dg4[4][4][4];
        double G4_x[4][4], G4_z[4][4];
        double mom4[4];

        double tmp, shift_term;

        const double det = gtxx * gtyy * gtzz + 2 * gtxy * gtyz * gtxz - gtzz * SQR(gtxy) - SQR(gtxz) * gtyy - gtxx * SQR(gtyz);

        // \tilde{γ}^{ij}
        gtu[0][0] =  (gtyy * gtzz - SQR(gtyz)) / det;
        gtu[1][1] =  (gtxx * gtzz - SQR(gtxz)) / det;
        gtu[2][2] =  (gtxx * gtyy - SQR(gtxy)) / det;
        gtu[0][1] = -(gtxy * gtzz - gtyz * gtxz) / det;
        gtu[0][2] =  (gtxy * gtyz - gtyy * gtxz) / det;
        gtu[1][2] = -(gtxx * gtyz - gtxy * gtxz) / det;
        gtu[1][0] = gtu[0][1];
        gtu[2][0] = gtu[0][2];
        gtu[2][1] = gtu[1][2];

        // γ_{jk}/^{jk}
        // K_{ij}
        for (int j = 0; j < 3; j++)
            for (int k = 0; k < 3; k++) {
                gu[j][k] = SQR(phi) * gtu[j][k];
                g[j][k]  = gt[j][k] / SQR(phi);
                K[j][k]  = At[j][k] / SQR(phi) + (1.0 / 3.0) * g[j][k] * trK;
            }

        // ∂_j γ_{kl}
        for (int j = 0; j < 3; j++)
            for (int k = 0; k < 3; k++)
                for (int l = 0; l < 3; l++) {
                    dg[j][k][l] = -2.0 * dphi[j] * gt[k][l] / (phi * SQR(phi)) + dgt[j][k][l] / SQR(phi);
                }

        // 4-metric
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                g4[1 + i][1 + j] = g[i][j];
        for (int i = 0; i < 3; i++) {
            g4[0][1 + i] = 0.0; // assume zero shift
            g4[1 + i][0] = 0.0;
        }
        g4[0][0] = -SQR(alpha);

        /* derivative of the 4-metric */
        /* spatial part */
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++) {
                for (int k = 0; k < 3; k++)
                    dg4[1 + k][1 + i][1 + j] = dg[k][i][j];
                dg4[0][1 + i][1 + j] = -2.0 * alpha * K[i][j]; // + Dbeta[i][j] + Dbeta[j][i]
            }

        /* temporal part, assume shift is zero */
        for (int i = 0; i < 4; i++)
            for (int j = 0; j < 3; j++) {
                dg4[i][1 + j][0] = 0.0;
                dg4[i][0][1 + j] = 0.0;
            }
        dg4[1][0][0] = -2 * alpha * dx_alpha;
        dg4[2][0][0] = 0.0;
        dg4[3][0][0] = -2 * alpha * dz_alpha;
        dg4[0][0][0] = 1e6; // should only be used multiplied by zero

        /* inverse 4-metric */
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                gu4[1 + i][1 + j] = gu[i][j];

        for (int i = 0; i < 3; i++) {
            gu4[1 + i][0] = 0.0;
            gu4[0][1 + i] = 0.0;
        }
        gu4[0][0] = -1. / SQR(alpha);

        for (int i = 0; i < 4; i++)
            for (int j = 0; j < 4; j++) {
                double val_x = 0.0;
                double val_z = 0.0;
                for (int k = 0; k < 4; k++) {
                    val_x += gu4[1][k] * (dg4[i][k][j] + dg4[j][k][i] - dg4[k][i][j]);
                    val_z += gu4[3][k] * (dg4[i][k][j] + dg4[j][k][i] - dg4[k][i][j]);
                }
                G4_x[i][j] = 0.5 * val_x;
                G4_z[i][j] = 0.5 * val_z;
            }

        mom4[1] = mom_x[idx];
        mom4[2] = 0.0;
        mom4[3] = mom_z[idx];

        shift_term = 0.0;
        for (int i = 0; i < 3; i++)
            shift_term += g4[0][1 + i] * mom4[1 + i];

        tmp = 0.0;
        for (int i = 0; i < 3; i++)
            for (int j = 0; j < 3; j++)
                tmp += g4[1 + i][1 + j] * mom4[1 + i] * mom4[1 + j];
        mom4[0] = 1.0 / g4[0][0] * (-shift_term - sqrt(SQR(shift_term) - g4[0][0] * tmp));

        tmp = 0.0;
        for (int i = 0; i < 4; i++)
            for (int j = 0; j < 4; j++)
                tmp += g4[i][j] * mom4[i] * mom4[j];
        if (fabs(tmp) > 1e-10)
            abort();

        dpos_dt_x[idx] = mom_x[idx] / mom4[0];
        dpos_dt_z[idx] = mom_z[idx] / mom4[0];

        tmp = 0.0;
        for (int i = 0; i < 4; i++)
            for (int j = 0; j < 4; j++)
                tmp += G4_x[i][j] * mom4[i] * mom4[j];
        dmom_dt_x[idx] = -tmp / mom4[0];

        tmp = 0.0;
        for (int i = 0; i < 4; i++)
            for (int j = 0; j < 4; j++)
                tmp += G4_z[i][j] * mom4[i] * mom4[j];
        dmom_dt_z[idx] = -tmp / mom4[0];
    }
}

static void interp_geometry(NSContext *ns, int surfaces_active,
                            const double *pos_x, const double *pos_y, const double *pos_z)
{
    const void * const coords[] = { pos_x, pos_y, pos_z };
    int ret;

    ret = CCTK_InterpGridArrays(ns->gh, 3, ns->interp_operator, ns->interp_params,
                                ns->coord_system, surfaces_active * ns->photons_per_surface,
                                CCTK_VARIABLE_REAL, coords,
                                ARRAY_ELEMS(ns->interp_vars_indices), ns->interp_vars_indices,
                                ARRAY_ELEMS(ns->interp_values), ns->interp_value_codes,
                                (void * const *)ns->interp_values);
    if (ret < 0)
        CCTK_WARN(0, "Error interpolating");
}

static void
ns_evol_rk_substep(NSContext *ns, int substep, int surfaces_active)
{
    const ptrdiff_t rk_rhs_stride = ns->nb_surfaces * ns->photons_per_surface;

    double *pos_x_rhs = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_x_rhs") + rk_rhs_stride * substep;
    double *pos_z_rhs = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_z_rhs") + rk_rhs_stride * substep;
    double *mom_x_rhs = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_x_rhs") + rk_rhs_stride * substep;
    double *mom_z_rhs = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_z_rhs") + rk_rhs_stride * substep;

    double *pos_x = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_x");
    double *pos_y = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_y");
    double *pos_z = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_z");
    double *mom_x = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_x");
    double *mom_z = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_z");

    const int origin_idx = CCTK_GFINDEX3D(ns->gh, ns->gh->cctk_nghostzones[0],
                                          0,      ns->gh->cctk_nghostzones[2]);
    double *tau_rhs = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::origin_proper_time_rhs") + substep;
    double  alpha   = ((double*)CCTK_VarDataPtr(ns->gh, 0, metric_vars[ALPHA]))[origin_idx];

    double *step_pos_x, *step_pos_z, *step_mom_x, *step_mom_z;

    /* proper time at origin */
    *tau_rhs = alpha;

    /* compute the coordinate/momenta to evaluate the RHS for */
    switch (substep) {
        case 0:
            step_pos_x = pos_x;
            step_pos_z = pos_z;
            step_mom_x = mom_x;
            step_mom_z = mom_z;
            break;
        case 1:
        case 2:
        case 3: {
            double fact = ns->dt * (substep == 3 ? 1.0 : 0.5);

            step_pos_x = ns->rk_scratch[0];
            step_pos_z = ns->rk_scratch[1];
            step_mom_x = ns->rk_scratch[2];
            step_mom_z = ns->rk_scratch[3];

#pragma omp parallel for
            for (int i = 0; i < surfaces_active * ns->photons_per_surface; i++) {
                step_pos_x[i] = pos_x[i] + fact * pos_x_rhs[i - rk_rhs_stride];
                step_pos_z[i] = pos_z[i] + fact * pos_z_rhs[i - rk_rhs_stride];
                step_mom_x[i] = mom_x[i] + fact * mom_x_rhs[i - rk_rhs_stride];
                step_mom_z[i] = mom_z[i] + fact * mom_z_rhs[i - rk_rhs_stride];
            }
            break;
        }
    }

    // interpolate the metric variables at this step's coordinates
    interp_geometry(ns, surfaces_active, step_pos_x, pos_y, step_pos_z);

    // eval RHS
    null_geodesic_rhs(ns->photons_per_surface * surfaces_active,
                      pos_x_rhs, pos_z_rhs, mom_x_rhs, mom_z_rhs,
                      step_pos_x, step_pos_z, step_mom_x, step_mom_z,
                      ns->interp_values);
}

static void ns_evol_rk_finish(NSContext *ns, int surfaces_active)
{
    const ptrdiff_t rk_rhs_stride = ns->nb_surfaces * ns->photons_per_surface;

    double *pos_x = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_x");
    double *pos_z = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_z");
    double *mom_x = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_x");
    double *mom_z = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_z");

    const double *pos_x_rhs = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_x_rhs");
    const double *pos_z_rhs = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::pos_z_rhs");
    const double *mom_x_rhs = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_x_rhs");
    const double *mom_z_rhs = CCTK_VarDataPtr(ns->gh, 0, "NullSurf::mom_z_rhs");

    const double *pos_x_rhs0 = pos_x_rhs + 0 * rk_rhs_stride;
    const double *pos_x_rhs1 = pos_x_rhs + 1 * rk_rhs_stride;
    const double *pos_x_rhs2 = pos_x_rhs + 2 * rk_rhs_stride;
    const double *pos_x_rhs3 = pos_x_rhs + 3 * rk_rhs_stride;

    const double *pos_z_rhs0 = pos_z_rhs + 0 * rk_rhs_stride;
    const double *pos_z_rhs1 = pos_z_rhs + 1 * rk_rhs_stride;
    const double *pos_z_rhs2 = pos_z_rhs + 2 * rk_rhs_stride;
    const double *pos_z_rhs3 = pos_z_rhs + 3 * rk_rhs_stride;

    const double *mom_x_rhs0 = mom_x_rhs + 0 * rk_rhs_stride;
    const double *mom_x_rhs1 = mom_x_rhs + 1 * rk_rhs_stride;
    const double *mom_x_rhs2 = mom_x_rhs + 2 * rk_rhs_stride;
    const double *mom_x_rhs3 = mom_x_rhs + 3 * rk_rhs_stride;

    const double *mom_z_rhs0 = mom_z_rhs + 0 * rk_rhs_stride;
    const double *mom_z_rhs1 = mom_z_rhs + 1 * rk_rhs_stride;
    const double *mom_z_rhs2 = mom_z_rhs + 2 * rk_rhs_stride;
    const double *mom_z_rhs3 = mom_z_rhs + 3 * rk_rhs_stride;

    double *tau_rhs = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::origin_proper_time_rhs");
    double *tau     = (double*)CCTK_VarDataPtr(ns->gh, 0, "NullSurf::origin_proper_time");

    *tau += ns->dt * (tau_rhs[0] / 6. + tau_rhs[1] / 3. + tau_rhs[2] / 3. + tau_rhs[3] / 6.);

#pragma omp parallel for
    for (int i = 0; i < surfaces_active * ns->photons_per_surface; i++) {
        pos_x[i] += ns->dt * (pos_x_rhs0[i] / 6. + pos_x_rhs1[i] / 3. + pos_x_rhs2[i] / 3. + pos_x_rhs3[i] / 6.);
        pos_z[i] += ns->dt * (pos_z_rhs0[i] / 6. + pos_z_rhs1[i] / 3. + pos_z_rhs2[i] / 3. + pos_z_rhs3[i] / 6.);
        mom_x[i] += ns->dt * (mom_x_rhs0[i] / 6. + mom_x_rhs1[i] / 3. + mom_x_rhs2[i] / 3. + mom_x_rhs3[i] / 6.);
        mom_z[i] += ns->dt * (mom_z_rhs0[i] / 6. + mom_z_rhs1[i] / 3. + mom_z_rhs2[i] / 3. + mom_z_rhs3[i] / 6.);
    }
}

void ns_evol(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    NSContext *ns = &ns_ctx;

    const int reflevel = ctz(cctk_levfac[0]);

    /* check if we are on the right reflevel */
    if (reflevel != solve_level)
        return;

    fprintf(stderr, "ns evol step %d at time %g; pos[zaxis] = %g; pos[xaxis] = %g\n",
            *evol_next_step, cctk_time, pos_z[0], pos_x[photons_per_surface - 1]);

    /* do the last Runge-Kutta sub-step */
    if (*evol_next_step == 3) {
        ns_evol_rk_substep(ns, *evol_next_step, *surfaces_active);
        ns_evol_rk_finish(ns, *surfaces_active);
        *evol_next_step = 0;
    }

    /* launch a new surface if its time has come */
    if (*surfaces_active < nb_surfaces &&
        *origin_proper_time >= *surfaces_active * launch_delta) {
        const int offset = *surfaces_active * photons_per_surface;
        fprintf(stderr, "launching surface %d at t=%g/tau=%g\n",
                *surfaces_active, cctk_time, *origin_proper_time);
        surface_launch(ns, pos_x + offset, pos_z + offset, mom_x + offset, mom_z + offset);
        *surfaces_active += 1;
    }

    /* do the Runge-Kutta evolution step(s): either 0->1 or 1->2 */
    ns_evol_rk_substep(ns, *evol_next_step, *surfaces_active);
    *evol_next_step += 1;

    if (*evol_next_step == 2) {
        ns_evol_rk_substep(ns, *evol_next_step, *surfaces_active);
        *evol_next_step += 1;
    }
}

static int context_init(NSContext *ns, cGH *gh, int reflevel,
                        int nb_surfaces, int photons_per_surface)
{
    int ret;

    ns->gh       = gh;
    ns->reflevel = reflevel;

    /* we do one null surface evolution step per two cactus time steps */
    ns->dt = 2.0 * gh->cctk_delta_time / (1 << reflevel);

    ns->coord_system = CCTK_CoordSystemHandle("cart3d");
    if (ns->coord_system < 0)
        CCTK_WARN(0, "Error getting the coordinate system");

    ns->interp_operator = CCTK_InterpHandle("Lagrange polynomial interpolation (tensor product)");
    if (ns->interp_operator < 0)
        CCTK_WARN(0, "Error getting the interpolation operator");

    ns->interp_params = Util_TableCreateFromString("order=4 want_global_mode=1");
    if (ns->interp_params < 0)
        CCTK_WARN(0, "Error creating interpolation parameters table");

    ret = Util_TableSetIntArray(ns->interp_params, NB_INTERP_VARS,
                                interp_operation_codes, "operation_codes");
    if (ret < 0)
        CCTK_WARN(0, "Error setting operation codes");

    ret = Util_TableSetIntArray(ns->interp_params, NB_INTERP_VARS,
                                interp_operation_indices, "operand_indices");
    if (ret < 0)
        CCTK_WARN(0, "Error setting operand indices");

    for (int i = 0; i < ARRAY_ELEMS(metric_vars); i++) {
        ns->interp_vars_indices[i] = CCTK_VarIndex(metric_vars[i]);
        if (ns->interp_vars_indices[i] < 0)
            CCTK_VWarn(0, __LINE__, __FILE__, CCTK_THORNSTRING, "Error getting the index of variable: %s\n", metric_vars[i]);
    }

    for (int i = 0; i < ARRAY_ELEMS(ns->interp_values); i++) {
        ret = posix_memalign((void**)&ns->interp_values[i], 32,
                             nb_surfaces * photons_per_surface * sizeof(*ns->interp_values[i]));
        if (ret)
            goto fail;
        ns->interp_value_codes[i] = CCTK_VARIABLE_REAL;
    }

    for (int i = 0; i < ARRAY_ELEMS(ns->rk_scratch); i++) {
        ret = posix_memalign((void**)&ns->rk_scratch[i], 32,
                             nb_surfaces * photons_per_surface * sizeof(*ns->rk_scratch[i]));
        if (ret)
            goto fail;
    }

    ns->nb_surfaces         = nb_surfaces;
    ns->photons_per_surface = photons_per_surface;

    return 0;
fail:
    return -ENOMEM;
}

void ns_evol_init(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    const int reflevel = ctz(cctk_levfac[0]);

    if (reflevel != solve_level)
        return;

    memset(&ns_ctx, 0, sizeof(ns_ctx));

    context_init(&ns_ctx, cctkGH, reflevel, nb_surfaces, photons_per_surface);

    /* initialize y-values of the photon positions */
    memset(pos_y, 0, sizeof(*pos_y) * nb_surfaces * photons_per_surface);
}

void ns_init(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    const int reflevel = ctz(cctk_levfac[0]);

    if (reflevel != solve_level)
        return;

    *surfaces_active = 0;
    *evol_next_step  = 0;
}