#include <inttypes.h>

#include <cl.h>

#include "cctk.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 ARRAY_ELEMS(arr) (sizeof(arr) / sizeof(*arr))

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

#define SCALE_FACTOR scale_factor

/* indices (in our code, not cactus structs) of the grid functions which we'll need to
 * interpolate on the pseudospectral grid */
enum MetricVars {
    GTXX = 0,
    GTYY,
    GTZZ,
    GTXY,
    GTXZ,
    GTYZ,
    PHI,
    ATXX,
    ATYY,
    ATZZ,
    ATXY,
    ATXZ,
    ATYZ,
    K,
    XTX,
    XTY,
    XTZ,
    BETAX,
    BETAY,
    BETAZ,
    NB_METRIC_VARS,
};

/* indices of the interpolated values of the above grid functions and their derivatives */
enum InterpMetricVars {
    I_GTXX = 0,
    I_GTXX_DX,
    I_GTXX_DY,
    I_GTXX_DZ,
    I_GTYY,
    I_GTYY_DX,
    I_GTYY_DY,
    I_GTYY_DZ,
    I_GTZZ,
    I_GTZZ_DX,
    I_GTZZ_DY,
    I_GTZZ_DZ,
    I_GTXY,
    I_GTXY_DX,
    I_GTXY_DY,
    I_GTXY_DZ,
    I_GTXZ,
    I_GTXZ_DX,
    I_GTXZ_DY,
    I_GTXZ_DZ,
    I_GTYZ,
    I_GTYZ_DX,
    I_GTYZ_DY,
    I_GTYZ_DZ,
    I_PHI,
    I_PHI_DX,
    I_PHI_DY,
    I_PHI_DZ,
    I_ATXX,
    I_ATYY,
    I_ATZZ,
    I_ATXY,
    I_ATXZ,
    I_ATYZ,
    I_K,
    I_K_DX,
    I_K_DY,
    I_K_DZ,
    I_XTX,
    I_XTY,
    I_XTZ,
    I_BETAX,
    I_BETAY,
    I_BETAZ,
    NB_INTERP_VARS,
};

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

extern const BasisSet msa_cheb_basis;
extern const BasisSet msa_full_basis;
extern const BasisSet msa_tb_even_basis;
extern const BasisSet msa_sb_even_basis;

extern double scale_factor;

/* precomputed values for a given refined grid */
typedef struct CoordPatch {
    CCTK_REAL origin[3];
    CCTK_REAL delta[3];
    CCTK_INT size[3];

    // basis values on the grid
    double *basis_val_r;
    double *basis_val_z;

    double *transform_z;
    double *one;
} CoordPatch;

/* state and scratch storage for the BiCGSTAB solver */
typedef struct BiCGSTABContext {
    double *p, *v, *y, *z, *t;
    double *res, *res0;
    double *k;

    cl_mem cl_p, cl_v, cl_y, cl_z, cl_t;
    cl_mem cl_res, cl_res0;
    cl_mem cl_k, cl_mat;
    cl_mem cl_rho, cl_alpha, cl_beta, cl_omega, cl_omega1;
    cl_mem cl_tmp, cl_tmp1;

    int64_t solve_total;
    int64_t iter_total;
    int64_t time_total;
} BiCGSTABContext;

typedef struct MaximalSlicingContext {
    cGH *gh;
    const BasisSet *basis;

    BiCGSTABContext bicgstab;
    int steps_since_inverse;

    int64_t lu_solves_total;
    int64_t lu_solves_time;

    // the grid of collocation points
    CCTK_REAL *grid_x;
    CCTK_REAL *grid_z;

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

    CCTK_REAL *interp_coords[3];

    int interp_vars_indices[NB_METRIC_VARS];

    CCTK_REAL *interp_values[NB_INTERP_VARS];
    CCTK_INT   interp_value_codes[NB_INTERP_VARS];

    CCTK_REAL *metric_u[6];

    CCTK_REAL *kij_kij;
    CCTK_REAL *trk;

    int nb_coeffs_x;
    int nb_coeffs_z;
    int nb_coeffs;

    int nb_colloc_points_x;
    int nb_colloc_points_z;
    int nb_colloc_points;

    int colloc_grid_order_x;
    int colloc_grid_order_z;

    double *mat;
    double *mat_f;
    double *rhs;
    double *coeffs;
    int *ipiv;
    double *basis_x_val;
    double *basis_x_dval;
    double *basis_x_d2val;

    double *basis_z_val;
    double *basis_z_dval;
    double *basis_z_d2val;

    double *basis_val_00;
    double *basis_val_20;
    double *basis_val_02;
    double *basis_val_11;
    double *basis_val_10;
    double *basis_val_01;

    CoordPatch *patches;
    int nb_patches;

    // OpenCL / CLBLAS stuff
    cl_context cl_ctx;
    cl_command_queue cl_queue;

    cl_mem ocl_coeffs;
} MaximalSlicingContext;

int msa_maximal_solve(MaximalSlicingContext *ms);

#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;
}