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

#include <mg2d.h>

#include <cblas.h>

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

#define SQR(x) ((x) * (x))
#define ARRAY_ELEMS(x) (sizeof(x) / sizeof(*x))

#define CPINDEX(cp, i, j, k) ((k * cp->grid_size[1] + j) * cp->grid_size[0] + i)

#if 0
#define LOGDEBUG(...) fprintf(stderr, __VA_ARGS)
#else
#define LOGDEBUG
#endif

static const struct {
    const char *str;
    enum MG2DLogLevel level;
} log_levels[] = {
    { "fatal",   MG2D_LOG_FATAL   },
    { "error",   MG2D_LOG_ERROR   },
    { "warning", MG2D_LOG_WARNING },
    { "info",    MG2D_LOG_INFO    },
    { "verbose", MG2D_LOG_VERBOSE },
    { "debug",   MG2D_LOG_DEBUG   },
};

#include <sys/time.h>
static inline int64_t gettime(void)
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return (int64_t)tv.tv_sec * 1000000 + tv.tv_usec;
}

typedef struct MaximalSlicingContext MaximalSlicingContext;

/* precomputed values for a given refined grid */
typedef struct CoordPatch {
    int    level;
    ptrdiff_t grid_size[3];

    MG2DContext *solver;

    ptrdiff_t y_idx;

    /* number of x/z grid points by which the elliptic solver domain is offset
     * from the cactus grid */
    ptrdiff_t offset_left[2];
    ptrdiff_t offset_right[2];

    /**
     * pointers into boundary_interp_vals
     * the first row are the boundary points for the solver
     * [0] is the upper-z patch
     *     rows run from x=0 to x[lsh[0] - 1]
     * [1] is the upper-x patch
     *     rows run from z=0 to z=<upper physical boundary>
     */
    ptrdiff_t boundary_val_stride[2];
} CoordPatch;

typedef struct MSMGContext {
    cGH *gh;

    int fd_stencil;

    CoordPatch *patches;
    int nb_patches;

    MaximalSlicingContext *ps_solver;

    int log_level;

    /* timings */
    int64_t time_solve;
    int64_t count_solve;

    int64_t time_solve_ps;
    int64_t count_solve_ps;

    int64_t time_solve_mg;
    int64_t count_solve_mg;
    int64_t time_solve_fill_coeffs;
    int64_t time_solve_boundaries;
    int64_t time_solve_mg2d;
    int64_t time_solve_export;
    int64_t time_solve_history;

    int64_t time_eval;
    int64_t count_eval;

    int64_t time_eval_ps;
    int64_t count_eval_ps;

    int64_t time_extrapolate;
    int64_t count_extrapolate;

    int64_t time_fine_solve;
    int64_t count_fine_solve;
    int64_t time_fine_fill_coeffs;
    int64_t time_fine_boundaries;
    int64_t time_fine_mg2d;
    int64_t time_fine_export;
} MSMGContext;

int msa_context_init(cGH *gh, MaximalSlicingContext **ctx, int basis_order0, int basis_order1, double scale_factor, double filter_power);
void msa_context_free(MaximalSlicingContext **ctx);
void msa_lapse_solve(MaximalSlicingContext *ctx);
void msa_lapse_eval(MaximalSlicingContext *ctx, double *dst);

static MSMGContext *ms;

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

    if (!a)
        return INT_MAX;

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

    return ret;
}

static void coord_patch_free(CoordPatch *cp)
{
    mg2d_solver_free(&cp->solver);
}

static void log_callback(const MG2DContext *ctx, int level,
                         const char *fmt, va_list vl)
{
    MSMGContext *ms = ctx->opaque;
    int target_level = ms->log_level;
    if (level > target_level)
        return;
    fprintf(stderr, "[%d] t=%g ", ctz(ms->gh->cctk_levfac[0]), ms->gh->cctk_time);
    vfprintf(stderr, fmt, vl);
}

static CoordPatch *get_coord_patch(MSMGContext *ms, int level)
{
    cGH *gh = ms->gh;

    const size_t grid_size = gh->cctk_lsh[2] * gh->cctk_lsh[1] * gh->cctk_lsh[0];
    const double *a_x = CCTK_VarDataPtr(gh, 0, "grid::x");
    const double *a_y = CCTK_VarDataPtr(gh, 0, "grid::y");
    const double *a_z = CCTK_VarDataPtr(gh, 0, "grid::z");

    CoordPatch *cp;
    size_t domain_size[2];
    int integrator_substeps = MoLNumIntegratorSubsteps();
    int i, ret;

    for (int i = 0; i < ms->nb_patches; i++) {
        cp = &ms->patches[i];

        if (cp->level == level)
            return cp;
    }

    /* create a new patch */
    ms->patches = realloc(ms->patches, sizeof(*ms->patches) * (ms->nb_patches + 1));
    cp = &ms->patches[ms->nb_patches];

    memset(cp, 0, sizeof(*cp));

    cp->level = ctz(ms->gh->cctk_levfac[0]);
    cp->grid_size[0] = gh->cctk_lsh[0];
    cp->grid_size[1] = gh->cctk_lsh[1];
    cp->grid_size[2] = gh->cctk_lsh[2];

    for (i = 0; i < gh->cctk_lsh[1]; i++)
        if (fabs(a_y[CCTK_GFINDEX3D(gh, 0, i, 0)]) < 1e-8) {
            cp->y_idx = i;
            break;
        }
    if (i == gh->cctk_lsh[1])
        CCTK_WARN(0, "The grid does not include y==0");

    for (int i = 0; i < 2; i++) {
        ptrdiff_t offset_left = ms->gh->cctk_nghostzones[2 * i];
        ptrdiff_t offset_right = offset_left;

        /* account for grid tapering on refined levels */
        if (cp->level > 0)
            offset_right *= integrator_substeps * 2;
        ///* the outer boundary layer overlaps with the first ghost zone */
        //if (cp->level > 0)
        //    offset_right++;
        if (cp->level > 0)
            offset_right--;

        domain_size[i] = ms->gh->cctk_lsh[2 * i] - offset_left - offset_right;
        /* the outer boundary layer overlaps with the first ghost zone */
        if (cp->level > 0)
            domain_size[i]++;

        cp->offset_left[i]  = offset_left;
        cp->offset_right[i] = offset_right;
    }

    if (domain_size[0] != domain_size[1])
        CCTK_WARN(0, "The grid is non-square, only square grids are supported");

    cp->solver = mg2d_solver_alloc(domain_size[0]);
    if (!cp->solver)
        CCTK_WARN(0, "Error allocating the solver");

    cp->solver->step[0] = a_x[1] - a_x[0];
    cp->solver->step[1] = cp->solver->step[0];

    cp->solver->fd_stencil = ms->fd_stencil;

    cp->solver->boundaries[MG2D_BOUNDARY_0L]->type = MG2D_BC_TYPE_FIXDIFF;
    cp->solver->boundaries[MG2D_BOUNDARY_1L]->type = MG2D_BC_TYPE_FIXDIFF;
    cp->solver->boundaries[MG2D_BOUNDARY_0U]->type = MG2D_BC_TYPE_FIXVAL;
    cp->solver->boundaries[MG2D_BOUNDARY_1U]->type = MG2D_BC_TYPE_FIXVAL;

    cp->solver->maxiter   = 32;
    cp->solver->tol       = 1e-10;
    cp->solver->nb_cycles = 2;

    cp->solver->opaque       = ms;
    cp->solver->log_callback = log_callback;


    /* initialize boundary values to zero,
     * non-zero values on the outer boundaries of refined levels are filled in elsewhere */
    for (int i = 0; i < 4; i++) {
        for (size_t j = 0; j < cp->solver->fd_stencil; j++)
            memset(cp->solver->boundaries[i]->val + j * cp->solver->boundaries[i]->val_stride - j,
                   0, sizeof(*cp->solver->boundaries[i]->val) * (domain_size[0] + 2 * j));
    }

    if (cp->level > 0) {
        cp->boundary_val_stride[0] = gh->cctk_lsh[0] - cp->offset_left[0];
        cp->boundary_val_stride[1] = cp->solver->domain_size;
    }

    ms->nb_patches++;
    return cp;
}

static void print_stats(MSMGContext *ms)
{
    int orig_log_level;
    int64_t total = ms->time_solve + ms->time_eval;

    fprintf(stderr, "%2.2f%% eval: %ld runs; %g s total; %g ms avg per run\n",
            (double)ms->time_eval * 100.0 / total, ms->count_eval,
            ms->time_eval / 1e6, ms->time_eval / 1e3 / ms->count_eval);
    fprintf(stderr, "  %2.2f%% eval ps: %ld runs; %g s total; %g ms avg per run\n",
            (double)ms->time_eval_ps * 100.0 / total, ms->count_eval_ps,
            ms->time_eval_ps / 1e6, ms->time_eval_ps / 1e3 / ms->count_eval_ps);
    fprintf(stderr, "  %2.2f%% eval extrapolate: %ld runs; %g s total; %g ms avg per run\n",
            (double)ms->time_extrapolate * 100.0 / total, ms->count_extrapolate,
            ms->time_extrapolate / 1e6, ms->time_extrapolate / 1e3 / ms->count_extrapolate);
    fprintf(stderr, "  %2.2f%% fine solve: %ld runs; %g s total; %g ms avg per run || "
            "%2.2f%% fill coeffs %2.2f%% boundaries %2.2f%% mg2d %2.2f%% export\n",
            (double)ms->time_fine_solve * 100.0 / total, ms->count_fine_solve,
            ms->time_fine_solve / 1e6, ms->time_fine_solve / 1e3 / ms->count_fine_solve,
            (double)ms->time_fine_fill_coeffs * 100.0 / ms->time_fine_solve,
            (double)ms->time_fine_boundaries  * 100.0 / ms->time_fine_solve,
            (double)ms->time_fine_mg2d        * 100.0 / ms->time_fine_solve,
            (double)ms->time_fine_export      * 100.0 / ms->time_fine_solve);

    fprintf(stderr, "%2.2f%% solve: %ld runs; %g s total; %g ms avg per run\n",
            (double)ms->time_solve * 100.0 / total, ms->count_solve,
            ms->time_solve / 1e6, ms->time_solve / 1e3 / ms->count_solve);
    fprintf(stderr, "  %2.2f%% solve ps: %ld runs; %g s total; %g ms avg per run\n",
            (double)ms->time_solve_ps * 100.0 / total, ms->count_solve_ps,
            ms->time_solve_ps / 1e6, ms->time_solve_ps / 1e3 / ms->count_solve_ps);
    fprintf(stderr, "  %2.2f%% solve mg2d: %ld runs; %g s total; %g ms avg per run || "
            "%2.2f%% fill coeffs %2.2f%% boundaries %2.2f%% mg2d %2.2f%% export %2.2f%% history\n",
            (double)ms->time_solve_mg * 100.0 / total, ms->count_solve_mg,
            ms->time_solve_mg / 1e6, ms->time_solve_mg / 1e3 / ms->count_solve_mg,
            (double)ms->time_solve_fill_coeffs * 100.0 / ms->time_solve_mg,
            (double)ms->time_solve_boundaries  * 100.0 / ms->time_solve_mg,
            (double)ms->time_solve_mg2d        * 100.0 / ms->time_solve_mg,
            (double)ms->time_solve_export      * 100.0 / ms->time_solve_mg,
            (double)ms->time_solve_history     * 100.0 / ms->time_solve_mg);

    orig_log_level = ms->log_level;
    ms->log_level = MG2D_LOG_VERBOSE;

    for (int i = 0; i < ms->nb_patches; i++) {
        CoordPatch *cp = get_coord_patch(ms, i);
        char indent_buf[64];

        snprintf(indent_buf, sizeof(indent_buf), " [%d] ", i);

        mg2d_print_stats(cp->solver, indent_buf);
    }
    ms->log_level = orig_log_level;
}

static int context_init(cGH *gh, int fd_stencil, const char *loglevel_str,
                        MSMGContext **ctx)
{
    MSMGContext *ms;
    int loglevel = -1;

    ms = calloc(1, sizeof(*ms));
    if (!ms)
        return -ENOMEM;

    ms->gh = gh;
    ms->fd_stencil = fd_stencil;

    for (int i = 0; i < ARRAY_ELEMS(log_levels); i++) {
        if (!strcmp(loglevel_str, log_levels[i].str)) {
            loglevel = log_levels[i].level;
            break;
        }
    }
    if (loglevel < 0) {
        fprintf(stderr, "Invalid loglevel: %s\n", loglevel_str);
        return -EINVAL;
    }

    ms->log_level = loglevel;

    *ctx = ms;

    return 0;
}

static void context_free(MSMGContext **pms)
{
    MSMGContext *ms = *pms;

    if (!ms)
        return;

    for (int i = 0; i < ms->nb_patches; i++)
        coord_patch_free(&ms->patches[i]);
    free(ms->patches);

    free(ms);
    *pms = NULL;
}

static void fill_eq_coeffs(cGH *gh, CoordPatch *cp)
{
    int ret;

    const ptrdiff_t stride_z = CCTK_GFINDEX3D(gh, 0, 0, 1);

    double *a_x = CCTK_VarDataPtr(gh, 0, "grid::x");
    double *a_z = CCTK_VarDataPtr(gh, 0, "grid::z");

    const double dx = a_x[1] - a_x[0];
    const double dz = a_z[stride_z] - a_z[0];

    double *a_gtxx = CCTK_VarDataPtr(gh, 0, "ML_BSSN::gt11");
    double *a_gtyy = CCTK_VarDataPtr(gh, 0, "ML_BSSN::gt22");
    double *a_gtzz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::gt33");
    double *a_gtxy = CCTK_VarDataPtr(gh, 0, "ML_BSSN::gt12");
    double *a_gtxz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::gt13");
    double *a_gtyz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::gt23");
    double *a_Atxx = CCTK_VarDataPtr(gh, 0, "ML_BSSN::At11");
    double *a_Atyy = CCTK_VarDataPtr(gh, 0, "ML_BSSN::At22");
    double *a_Atzz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::At33");
    double *a_Atxy = CCTK_VarDataPtr(gh, 0, "ML_BSSN::At12");
    double *a_Atxz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::At13");
    double *a_Atyz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::At23");
    double  *a_phi = CCTK_VarDataPtr(gh, 0, "ML_BSSN::phi");
    double  *a_Xtx = CCTK_VarDataPtr(gh, 0, "ML_BSSN::Xt1");
    double  *a_Xtz = CCTK_VarDataPtr(gh, 0, "ML_BSSN::Xt3");

    for (int idx_z = 0; idx_z < cp->solver->domain_size; idx_z++)
        for (int idx_x = 0; idx_x < cp->solver->domain_size; idx_x++) {
            const int idx_src = CCTK_GFINDEX3D(gh, idx_x + gh->cctk_nghostzones[0], cp->y_idx, idx_z + gh->cctk_nghostzones[2]);
            const int idx_dc  = idx_z * cp->solver->diff_coeffs_stride + idx_x;
            const int idx_rhs = idx_z * cp->solver->rhs_stride + idx_x;

            const double x = a_x[idx_src];

            const double gtxx = a_gtxx[idx_src];
            const double gtyy = a_gtyy[idx_src];
            const double gtzz = a_gtzz[idx_src];
            const double gtxy = a_gtxy[idx_src];
            const double gtxz = a_gtxz[idx_src];
            const double gtyz = a_gtyz[idx_src];
            const double Atxx = a_Atxx[idx_src];
            const double Atyy = a_Atyy[idx_src];
            const double Atzz = a_Atzz[idx_src];
            const double Atxy = a_Atxy[idx_src];
            const double Atxz = a_Atxz[idx_src];
            const double Atyz = a_Atyz[idx_src];
            const double  phi =  a_phi[idx_src];
            const double  Xtx =  a_Xtx[idx_src];
            const double  Xtz =  a_Xtz[idx_src];

#if 1
            const double phi_dx = (a_phi[idx_src + 1] - a_phi[idx_src - 1]) / (2.0 * dx);
            const double phi_dz = (a_phi[idx_src + stride_z] - a_phi[idx_src - stride_z]) / (2.0 * dz);
#else
            const double phi_dx = (-1.0 * a_phi[idx_src + 2] + 8.0 * a_phi[idx_src + 1] - 8.0 * a_phi[idx_src - 1] + a_phi[idx_src - 1]) / (12.0 * dx);
            const double phi_dz = (-1.0 * a_phi[idx_src + 2 * stride_z] + 8.0 * a_phi[idx_src + 1 * stride_z] - 8.0 * a_phi[idx_src - 1 * stride_z] + a_phi[idx_src - 1 * stride_z]) / (12.0 * dz);
#endif

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

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


            double Xx, Xz, k2;

            double Am[3][3], gtu[3][3];

			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];

			for (int j = 0; j < 3; j++)
				for (int k = 0; k < 3; k++) {
					double val = 0.0;
					for (int l = 0; l < 3; l++)
						val += gtu[j][l] * At[l][k];
					Am[j][k] = val;
				}

			// K_{ij} K^{ij}
			k2 = 0.0;
			for (int j = 0; j < 3; j++)
				for (int k = 0; k < 3; k++)
					k2 += Am[j][k] * Am[k][j];

            Xx = SQR(phi) * (Xtx + (phi_dx * gtu[0][0] + phi_dz * gtu[0][2]) / phi);
            Xz = SQR(phi) * (Xtz + (phi_dx * gtu[0][2] + phi_dz * gtu[2][2]) / phi);

            cp->solver->diff_coeffs[MG2D_DIFF_COEFF_20][idx_dc] = SQR(phi) * (gtu[0][0] + ((idx_x == 0) ? gtu[1][1] : 0.0));
            cp->solver->diff_coeffs[MG2D_DIFF_COEFF_02][idx_dc] = SQR(phi) * gtu[2][2];
            cp->solver->diff_coeffs[MG2D_DIFF_COEFF_11][idx_dc] = SQR(phi) * gtu[0][2] * 2;
            cp->solver->diff_coeffs[MG2D_DIFF_COEFF_10][idx_dc] = -Xx + (idx_x ? SQR(phi) * gtu[1][1] / x : 0.0);
            cp->solver->diff_coeffs[MG2D_DIFF_COEFF_01][idx_dc] = -Xz;
            cp->solver->diff_coeffs[MG2D_DIFF_COEFF_00][idx_dc] = -k2;
            cp->solver->rhs[idx_rhs]                            = k2;
        }
}

static void solution_to_grid(CoordPatch *cp, double *dst)
{
    for (int j = 0; j < cp->solver->domain_size; j++)
        for (int i = 0; i < cp->solver->domain_size; i++) {
            const ptrdiff_t idx_dst = CPINDEX(cp, i + cp->offset_left[0], cp->y_idx, j + cp->offset_left[1]);
            const int idx_src = j * cp->solver->u_stride + i;
            dst[idx_dst] = 1.0 + cp->solver->u[idx_src];
        }

    /* fill in the axial symmetry ghostpoints by mirroring */
    for (int idx_z = cp->offset_left[1]; idx_z < cp->grid_size[2]; idx_z++)
        for (int idx_x = 0; idx_x < cp->offset_left[0]; idx_x++) {
            const ptrdiff_t idx_dst = CPINDEX(cp, idx_x, cp->y_idx, idx_z);
            const ptrdiff_t idx_src = CPINDEX(cp, 2 * cp->offset_left[0] - idx_x, cp->y_idx, idx_z);
            dst[idx_dst] = dst[idx_src];
        }
    for (int idx_z = 0; idx_z < cp->offset_left[1]; idx_z++)
        for (int idx_x = 0; idx_x < cp->grid_size[0]; idx_x++) {
            const ptrdiff_t idx_dst = CPINDEX(cp, idx_x, cp->y_idx, idx_z);
            const ptrdiff_t idx_src = CPINDEX(cp, idx_x, cp->y_idx, 2 * cp->offset_left[1] - idx_z);
            dst[idx_dst] = dst[idx_src];
        }
}

static void solve_ps(MSMGContext *ctx, double *dst)
{
    if (!ctx->ps_solver)
        msa_context_init(ctx->gh, &ctx->ps_solver, 60, 60, 8.0, 64.0);

    msa_lapse_solve(ctx->ps_solver);
    msa_lapse_eval(ctx->ps_solver, dst);
}

static int history_initialized(double *lapse_prev0_time, double *lapse_prev1_time, int nb_levels)
{
    for (int i = 0; i < nb_levels; i++) {
        if (lapse_prev0_time[i] == DBL_MAX || lapse_prev1_time[i] == DBL_MAX)
            return 0;
    }
    return 1;
}

void msa_mg_eval(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    int64_t total_start;

    int nb_levels;
    int nb_levels_type, ret;

    const size_t grid_size = cctkGH->cctk_lsh[2] * cctkGH->cctk_lsh[1] * cctkGH->cctk_lsh[0];
    const double t = cctkGH->cctk_time;
    const int reflevel = ctz(cctkGH->cctk_levfac[0]);
    double time_interp_step, fact0, fact1;

    total_start = gettime();

    LOGDEBUG( "evaluating lapse at rl=%d, t=%g\n", reflevel, t);

    nb_levels = *(int*)CCTK_ParameterGet("num_levels_1", "CarpetRegrid2", &nb_levels_type);

    if (!history_initialized(lapse_prev0_time, lapse_prev1_time, nb_levels)) {
        int64_t eval_ps_start = gettime();

        solve_ps(ms, lapse_mg_eval);
        memcpy(alpha, lapse_mg_eval, sizeof(*alpha) * grid_size);

        ms->time_eval_ps += gettime() - eval_ps_start;
        ms->count_eval_ps++;

        LOGDEBUG( "pseudospectral solve\n");
        goto finish;
    }

    if (lapse_prev1_time[reflevel] <= lapse_prev0_time[reflevel] ||
        fabs(lapse_prev1_time[reflevel] - lapse_prev0_time[reflevel]) < 1e-14)
        CCTK_WARN(0, "Invalid past value times");

    time_interp_step = lapse_prev1_time[reflevel] - lapse_prev0_time[reflevel];
    fact0 = (lapse_prev1_time[reflevel] - t) / time_interp_step;
    fact1 = (t - lapse_prev0_time[reflevel]) / time_interp_step;

    if (reflevel < nb_levels - 1) {
        /* on coarse levels use extrapolated past solutions */
        int64_t extrap_start = gettime();

        LOGDEBUG( "extrapolating from t1=%g (%g) and t0=%g (%g)\n", lapse_prev1_time[reflevel], fact1, lapse_prev0_time[reflevel], fact0);

        for (size_t i = 0; i < grid_size; i++)
            lapse_mg_eval[i] = fact0 * lapse_prev0[i] + fact1 * lapse_prev1[i];

        ms->time_extrapolate += gettime() - extrap_start;
        ms->count_extrapolate++;
    } else {
        /* on the finest level, use the extrapolation only for the boundary
         * values and solve for the interior */
        int64_t fine_solve_start = gettime();
        int64_t start;
        CoordPatch *cp = get_coord_patch(ms, reflevel);

        start = gettime();
        fill_eq_coeffs(ms->gh, cp);
        ms->time_fine_fill_coeffs += gettime() - start;

        LOGDEBUG( "extrapolating BCs from t1=%g (%g) and t0=%g (%g)\n", lapse_prev1_time[reflevel], fact1, lapse_prev0_time[reflevel], fact0);

        start = gettime();
        for (ptrdiff_t idx_z = 0; idx_z < cp->offset_right[1]; idx_z++)
            for (ptrdiff_t idx_x = 0; idx_x < cp->boundary_val_stride[0]; idx_x++) {
                const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, idx_x + cp->offset_left[0], cp->y_idx, cctkGH->cctk_lsh[2] - cp->offset_right[1] + idx_z);
                lapse_mg_eval[idx] = fact0 * lapse_prev0[idx] + fact1 * lapse_prev1[idx];
            }
        for (int idx_z = 0; idx_z < cp->boundary_val_stride[1]; idx_z++)
            for (int idx_x = 0; idx_x < cp->offset_right[0]; idx_x++) {
                const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, cctkGH->cctk_lsh[0] - cp->offset_right[0] + idx_x, cp->y_idx,
                                                     cp->offset_left[1] + idx_z);
                lapse_mg_eval[idx] = fact0 * lapse_prev0[idx] + fact1 * lapse_prev1[idx];
            }

        for (int j = 0; j < cp->solver->fd_stencil; j++) {
            MG2DBoundary *bnd = cp->solver->boundaries[MG2D_BOUNDARY_0U];
            double *dst = bnd->val + j * bnd->val_stride;
            for (ptrdiff_t i = -j; i < (ptrdiff_t)cp->solver->domain_size + j; i++) {
                const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, i + cp->offset_left[0], cp->y_idx, cctkGH->cctk_lsh[2] - cp->offset_right[1] + j);
                dst[i] = lapse_mg_eval[idx] - 1.0;
            }
        }
        for (int j = 0; j < cp->solver->fd_stencil; j++) {
            MG2DBoundary *bnd = cp->solver->boundaries[MG2D_BOUNDARY_1U];
            double *dst = bnd->val + j * bnd->val_stride;
            for (ptrdiff_t i = -j; i < (ptrdiff_t)cp->solver->domain_size + j; i++) {
                const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, cctkGH->cctk_lsh[0] - cp->offset_right[0] + j, cp->y_idx, cp->offset_left[1] + i);
                dst[i] = lapse_mg_eval[idx] - 1.0;
            }
        }
        ms->time_fine_boundaries += gettime() - start;

        LOGDEBUG( "mg solve\n");

        start = gettime();
        ret = mg2d_solve(cp->solver);
        if (ret < 0)
            CCTK_WARN(0, "Error solving the maximal slicing equation");
        ms->time_fine_mg2d += gettime() - start;

        start = gettime();
        solution_to_grid(cp, lapse_mg_eval);
        ms->time_fine_export = gettime() - start;

        ms->time_fine_solve += gettime() - fine_solve_start;
        ms->count_fine_solve++;
    }
    memcpy(alpha, lapse_mg_eval, sizeof(*alpha) * grid_size);

finish:
    ms->time_eval += gettime() - total_start;
    ms->count_eval++;
}

void msa_mg_solve(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    CoordPatch *cp;

    int64_t total_start, mg_start, start;

    int nb_levels, reflevel_top, ts_tmp;
    int nb_levels_type, ret;

    const size_t grid_size = cctkGH->cctk_lsh[2] * cctkGH->cctk_lsh[1] * cctkGH->cctk_lsh[0];
    const double t = cctkGH->cctk_time;
    const int timestep = lrint(t * cctkGH->cctk_levfac[0] / cctkGH->cctk_delta_time);
    const int reflevel = ctz(cctkGH->cctk_levfac[0]);

    double *lapse_mg1;

    total_start = gettime();

    LOGDEBUG( "solve lapse at rl=%d, t=%g; step %d\n", reflevel, t, timestep);

    nb_levels = *(int*)CCTK_ParameterGet("num_levels_1", "CarpetRegrid2", &nb_levels_type);

    reflevel_top = reflevel;
    ts_tmp = timestep;
    while (reflevel_top > 0 && !(ts_tmp % 2)) {
        ts_tmp /= 2;
        reflevel_top--;
    }

    if (!history_initialized(lapse_prev0_time, lapse_prev1_time, nb_levels)) {
        if (reflevel == 0 || !(timestep % 2)) {
            int64_t solve_ps_start = gettime();

            if (lapse_prev1_time[reflevel] != DBL_MAX &&
                reflevel_top == 0) {
                LOGDEBUG( "copy %g\n", lapse_prev1_time[reflevel]);
                memcpy(lapse_prev0, lapse_prev1, grid_size * sizeof(*lapse_prev1));
                lapse_prev0_time[reflevel] = lapse_prev1_time[reflevel];
            }

            LOGDEBUG( "solve ps\n");
            solve_ps(ms, alpha);
            memcpy(lapse_prev1, alpha, grid_size * sizeof(*lapse_prev0));
            lapse_prev1_time[reflevel] = t;

            ms->time_solve_ps += gettime() - solve_ps_start;
            ms->count_solve_ps++;
        }
        goto finish;
    }

    mg_start = gettime();
    LOGDEBUG( "mg solve cur %d top %d\n", reflevel, reflevel_top);

    /* fill in the equation coefficients */
    cp = get_coord_patch(ms, reflevel);
    start = gettime();
    fill_eq_coeffs(cctkGH, cp);
    ms->time_solve_fill_coeffs += gettime() - start;

    start = gettime();
    if (reflevel == 0 || timestep % 2) {
        /* outer-most level for this solve, use extrapolated lapse as the
         * boundary condition */
        double time_interp_step = lapse_prev1_time[reflevel] - lapse_prev0_time[reflevel];
        double fact0 = (lapse_prev1_time[reflevel] - t) / time_interp_step;
        double fact1 = (t - lapse_prev0_time[reflevel]) / time_interp_step;

        LOGDEBUG( "extrapolating BCs from t1=%g (%g) and t0=%g (%g)\n", lapse_prev1_time[reflevel], fact1, lapse_prev0_time[reflevel], fact0);

        for (int j = 0; j < cp->solver->fd_stencil; j++) {
            for (ptrdiff_t i = -j; i < (ptrdiff_t)cp->solver->domain_size + j; i++) {
                const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, i + cp->offset_left[0], cp->y_idx, cctkGH->cctk_lsh[2] - cp->offset_right[1] + j);
                lapse_mg[idx] = fact0 * lapse_prev0[idx] + fact1 * lapse_prev1[idx];
            }
        }
        for (int j = 0; j < cp->solver->fd_stencil; j++) {
            for (ptrdiff_t i = -j; i < (ptrdiff_t)cp->solver->domain_size + j; i++) {
                const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, cctkGH->cctk_lsh[0] - cp->offset_right[0] + j, cp->y_idx, cp->offset_left[1] + i);
                lapse_mg[idx] = fact0 * lapse_prev0[idx] + fact1 * lapse_prev1[idx];
            }
        }
    }
    /* if the condition above was false, then lapse_mg should be filled by
     * prolongation from the coarser level */

    /* fill the solver boundary conditions */
    for (int j = 0; j < cp->solver->fd_stencil; j++) {
        MG2DBoundary *bnd = cp->solver->boundaries[MG2D_BOUNDARY_0U];
        double *dst = bnd->val + j * bnd->val_stride;
        for (ptrdiff_t i = -j; i < (ptrdiff_t)cp->solver->domain_size + j; i++) {
            const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, i + cp->offset_left[0], cp->y_idx, cctkGH->cctk_lsh[2] - cp->offset_right[1] + j);
            dst[i] = lapse_mg[idx] - 1.0;
        }
    }
    for (int j = 0; j < cp->solver->fd_stencil; j++) {
        MG2DBoundary *bnd = cp->solver->boundaries[MG2D_BOUNDARY_1U];
        double *dst = bnd->val + j * bnd->val_stride;
        for (ptrdiff_t i = -j; i < (ptrdiff_t)cp->solver->domain_size + j; i++) {
            const ptrdiff_t idx = CCTK_GFINDEX3D(cctkGH, cctkGH->cctk_lsh[0] - cp->offset_right[0] + j, cp->y_idx, cp->offset_left[1] + i);
            dst[i] = lapse_mg[idx] - 1.0;
        }
    }
    ms->time_solve_boundaries += gettime() - start;

    /* do the elliptic solve */
    start = gettime();
    ret = mg2d_solve(cp->solver);
    if (ret < 0)
        CCTK_WARN(0, "Error solving the maximal slicing equation");
    ms->time_solve_mg2d += gettime() - start;

    /* export the result */
    start = gettime();
    solution_to_grid(cp, lapse_mg);

    lapse_mg1 = CCTK_VarDataPtr(cctkGH, 1, "MaximalSlicingAxiMG::lapse_mg");
    memcpy(lapse_mg1, lapse_mg, grid_size * sizeof(*lapse_mg1));
    ms->time_solve_export += gettime() - start;

    /* update lapse history for extrapolation */
    if (reflevel == 0 || !(timestep % 2)) {
        const double vel_fact = 1.0 / (1 << (reflevel - reflevel_top));

        start = gettime();
        for (size_t j = 0; j < grid_size; j++) {
            const double sol_new = lapse_mg[j];
            const double delta = sol_new - lapse_mg_eval[j];
            lapse_prev0[j] = lapse_prev1[j] + delta - delta * vel_fact;
            lapse_prev1[j] = sol_new;
            alpha[j]       = sol_new;
        }
        lapse_prev0_time[reflevel] = lapse_prev1_time[reflevel];
        lapse_prev1_time[reflevel] = cctkGH->cctk_time;

        ms->time_solve_history += gettime() - start;
    }

    ms->time_solve_mg += gettime() - mg_start;
    ms->count_solve_mg++;

finish:
    ms->time_solve += gettime() - total_start;
    ms->count_solve++;

    if (stats_every > 0 && reflevel == 0 && ms->count_solve > 0 && ms->count_eval > 0 &&
        !(ms->count_solve % stats_every))
        print_stats(ms);
}

void msa_mg_sync(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;
    const int reflevel = ctz(cctkGH->cctk_levfac[0]);
    LOGDEBUG( "\nsync %g %d\n\n", cctkGH->cctk_time, reflevel);

    return;
}

#if 0
static void maximal_slicing_axi_mg_old(CCTK_ARGUMENTS)
{
    CoordPatch *cp;

    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    double *coeffs = NULL;
    int i, ret;

    fprintf(stderr, "ms mg solve: level %d time %g\n", ctz(ms->gh->cctk_levfac[0]), ms->gh->cctk_time);

    cp = get_coord_patch(ms);

    fill_eq_coeffs(ms->gh, cp);

    CCTK_TimerStart("MaximalSlicingAxiMG_Solve");
    ret = mg2d_solve(cp->solver);
    if (ret < 0)
        CCTK_WARN(0, "Error solving the maximal slicing equation");
    CCTK_TimerStop("MaximalSlicingAxiMG_Solve");


    CCTK_TimerStart("MaximalSlicingAxiMG_Copy");
//#pragma omp parallel for
    for (int j = 0; j < cp->solver->domain_size; j++)
        for (int i = 0; i < cp->solver->domain_size; i++) {
            const int idx_dst = CCTK_GFINDEX3D(ms->gh, i + cp->offset_left[0], cp->y_idx, j + cp->offset_left[1]);
            const int idx_src = j * cp->solver->u_stride + i;
            alpha[idx_dst] = 1.0 + cp->solver->u[idx_src];
        }
    if (cp->level > 0) {
        /* fill in the interpolated outer buffer points */
        for (int idx_z = 0; idx_z < cp->offset_right[1]; idx_z++)
            for (int idx_x = 0; idx_x < cp->boundary_val_stride[0]; idx_x++) {
                const ptrdiff_t idx_src = idx_z * cp->boundary_val_stride[0] + idx_x;
                const ptrdiff_t idx_dst = CCTK_GFINDEX3D(ms->gh, cp->offset_left[0] + idx_x, cp->y_idx, ms->gh->cctk_lsh[2] - cp->offset_right[1] + idx_z);
                alpha[idx_dst] = cp->boundary_val[0][idx_src];
            }
        for (int idx_z = 0; idx_z < cp->boundary_val_stride[1]; idx_z++)
            for (int idx_x = 0; idx_x < cp->offset_right[0]; idx_x++) {
                const ptrdiff_t idx_src = idx_x * cp->boundary_val_stride[1] + idx_z;
                const ptrdiff_t idx_dst = CCTK_GFINDEX3D(ms->gh, ms->gh->cctk_lsh[0] - cp->offset_right[0] + idx_x, cp->y_idx,
                                                         cp->offset_left[1] + idx_z);
                alpha[idx_dst] = cp->boundary_val[1][idx_src];
            }
    }

    /* fill in the axial symmetry ghostpoints by mirroring */
    for (int idx_z = cp->offset_left[1]; idx_z < ms->gh->cctk_lsh[2]; idx_z++)
        for (int idx_x = 0; idx_x < cp->offset_left[0]; idx_x++) {
            const ptrdiff_t idx_dst = CCTK_GFINDEX3D(ms->gh, idx_x, cp->y_idx, idx_z);
            const ptrdiff_t idx_src = CCTK_GFINDEX3D(ms->gh, 2 * cp->offset_left[0] - idx_x, cp->y_idx, idx_z);
            alpha[idx_dst] = alpha[idx_src];
        }
    for (int idx_z = 0; idx_z < cp->offset_left[1]; idx_z++)
        for (int idx_x = 0; idx_x < ms->gh->cctk_lsh[0]; idx_x++) {
            const ptrdiff_t idx_dst = CCTK_GFINDEX3D(ms->gh, idx_x, cp->y_idx, idx_z);
            const ptrdiff_t idx_src = CCTK_GFINDEX3D(ms->gh, idx_x, cp->y_idx, 2 * cp->offset_left[1] - idx_z);
            alpha[idx_dst] = alpha[idx_src];
        }
    CCTK_TimerStop("MaximalSlicingAxiMG_Copy");
}
#endif

void msa_mg_init(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;
    int ret;

    ret = context_init(cctkGH, fd_stencil, loglevel, &ms);
    if (ret < 0)
        CCTK_WARN(0, "Error initializing the solver context");
}

void msa_mg_inithist(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;
    const size_t grid_size = cctkGH->cctk_lsh[2] * cctkGH->cctk_lsh[1] * cctkGH->cctk_lsh[0];

    for (size_t i = 0; i < grid_size; i++)
        lapse_mg[i] = 1.0;

    for (int i = 0; i < 32; i++) {
        lapse_prev0_time[i] = DBL_MAX;
        lapse_prev1_time[i] = DBL_MAX;
    }
}