#include <ctype.h>
#include <errno.h>
#include <float.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 <gsl/gsl_spline.h>

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

#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.0f : -1.0f)

/*
 * 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*/
    CCTK_REAL (*eval)      (CCTK_REAL coord, int idx);
    /* evaluate the first derivative of the idx-th basis function at the specified point*/
    CCTK_REAL (*eval_diff1)(CCTK_REAL coord, int idx);
    /* evaluate the second derivative of the idx-th basis function at the specified point*/
    CCTK_REAL (*eval_diff2)(CCTK_REAL coord, int idx);
    /**
     * Get the idx-th collocation point for the specified order.
     * idx runs from 0 to order - 1 (inclusive)
     */
    CCTK_REAL (*colloc_point)(int order, int idx);
} BasisSet;

/*
 * 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 beginning and decay as 1/x in infinity.
 */

#define SCALE_FACTOR 1

static CCTK_REAL sb_even_eval(CCTK_REAL coord, int idx)
{
    CCTK_REAL val = (coord == 0.0) ? M_PI_2 : atan(SCALE_FACTOR / fabs(coord));

    idx *= 2;   // even only

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

static CCTK_REAL sb_even_eval_diff1(CCTK_REAL coord, int idx)
{
    CCTK_REAL val = (coord == 0.0) ? M_PI_2 : atan(SCALE_FACTOR / coord );

    idx *= 2;   // even only

    return - SCALE_FACTOR * (idx + 1) * SGN(coord) * cos((idx + 1) * val) / (SQR(SCALE_FACTOR) + SQR(coord));
}

static CCTK_REAL sb_even_eval_diff2(CCTK_REAL coord, int idx)
{
    CCTK_REAL val = (coord == 0.0) ? M_PI_2 : atan(SCALE_FACTOR / coord);

    idx *= 2;   // even only

    return SCALE_FACTOR * (idx + 1) * SGN(coord) * (2 * coord * cos((idx + 1) * val) - SCALE_FACTOR * (idx + 1) * sin((idx + 1) * val)) / SQR(SQR(SCALE_FACTOR) + SQR(coord));
}

static CCTK_REAL sb_even_colloc_point(int order, int idx)
{
    CCTK_REAL t;
    order *= 2;

    t = (idx + 2) * M_PI / (order + 4);
    return SCALE_FACTOR / tan(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,
};

typedef struct BrillDataContext {
    /* options */
    int cheb_order;
    float amplitude;

    double (*q_func)(struct BrillDataContext *bd, double rho, double z);
    double (*laplace_q_func)(struct BrillDataContext *bd, double rho, double z);

    /* state */
    const BasisSet *basis;

    int nb_coeffs_x;
    int nb_coeffs_z;
    int nb_coeffs;

    gsl_vector *psi_coeffs;
} BrillDataContext;

static int brill_construct_matrix(BrillDataContext *bd, gsl_matrix *mat,
                                  gsl_vector *rhs)
{
    const int nb_gridpoints_z = bd->cheb_order;
    const int nb_gridpoints_x = bd->cheb_order;
    const int nb_gridpoints   = nb_gridpoints_x * nb_gridpoints_z;

    CCTK_REAL z_physical, x_physical;

    int idx_coeff_x, idx_coeff_z, idx_grid_x, idx_grid_z;
    int idx_coeff, idx_grid;

    /* interior */
    for (idx_grid_z = 0; idx_grid_z < nb_gridpoints_z; idx_grid_z++) {
        CCTK_REAL z_physical = bd->basis->colloc_point(nb_gridpoints_z, idx_grid_z);

        fprintf(stdout, "coord %d %g\n", idx_grid_z, z_physical);

        for (idx_grid_x = 0; idx_grid_x < nb_gridpoints_x; idx_grid_x++) {
            CCTK_REAL x_physical = bd->basis->colloc_point(nb_gridpoints_x, idx_grid_x);
            CCTK_REAL d2q        = bd->laplace_q_func(bd, x_physical, z_physical);
            int idx_grid         = idx_grid_z * nb_gridpoints_z + idx_grid_x;

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

                    mat->data[idx_grid * mat->tda + idx_coeff] = bd->basis->eval_diff2(x_physical, idx_coeff_x) * bd->basis->eval(z_physical, idx_coeff_z) * x_physical +
                                                                 bd->basis->eval_diff1(x_physical, idx_coeff_x) * bd->basis->eval(z_physical, idx_coeff_z) +
                                                                 bd->basis->eval_diff2(z_physical, idx_coeff_z) * bd->basis->eval(x_physical, idx_coeff_x) * x_physical +
                                                                 bd->basis->eval      (x_physical, idx_coeff_x) * bd->basis->eval(z_physical, idx_coeff_z) * 0.25 * d2q * x_physical;
                }
            rhs->data[idx_grid] = -0.25 * d2q * x_physical;
        }
    }

    return 0;
}

/* solve the equation
 * Δψ + ¼ ψ (∂²q/∂r² + ∂²q/∂z²) = 0
 * for the coefficients of spectral approximation of ψ:
 * ψ(r, z) = ΣaᵢⱼTᵢ(r)Tⱼ(z)
 * The cofficients are exported in bd->psi_coeffs
 */
static int brill_solve(BrillDataContext *bd)
{
    gsl_matrix *mat;
    gsl_permutation *perm;
    gsl_vector *coeffs, *rhs;

    int signum;
    int ret = 0;

    /* allocate and fill the pseudospectral matrix */
    mat    = gsl_matrix_alloc(bd->nb_coeffs, bd->nb_coeffs);

    perm   = gsl_permutation_alloc(bd->nb_coeffs);
    coeffs = gsl_vector_alloc     (bd->nb_coeffs);
    rhs    = gsl_vector_calloc    (bd->nb_coeffs);
    if (!mat || !perm || !coeffs || !rhs) {
        ret = -ENOMEM;
        goto fail;
    }

    /* fill the matrix */
    brill_construct_matrix(bd, mat, rhs);

    /* solve for the coeffs */
    gsl_linalg_LU_decomp(mat, perm, &signum);
    gsl_linalg_LU_solve(mat, perm, rhs, coeffs);

    bd->psi_coeffs = coeffs;

fail:
    if (ret < 0)
        gsl_vector_free(coeffs);

    gsl_vector_free(rhs);
    gsl_matrix_free(mat);
    gsl_permutation_free(perm);

    return ret;
}

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

static double laplace_q_gundlach(BrillDataContext *bd, double rho, double z)
{
    double r2 = SQR(rho);
    double z2 = SQR(z);
    double e  = exp(-r2 - z2);

    return 2 * bd->amplitude * e * (1.0 + 2 * r2 * (r2 + z2 - 3));
}

static BrillDataContext *init_bd(cGH *cctkGH, float amplitude, int cheb_order)
{
    BrillDataContext *bd;

    int i, j;

    bd = malloc(sizeof(*bd));

    bd->q_func         = q_gundlach;
    bd->laplace_q_func = laplace_q_gundlach;
    bd->amplitude      = amplitude;

    bd->basis      = &sb_even_basis;
    bd->cheb_order = cheb_order;

    bd->nb_coeffs_x = cheb_order;
    bd->nb_coeffs_z = cheb_order;

    bd->nb_coeffs = bd->nb_coeffs_x * bd->nb_coeffs_z;

    return bd;
}

void brill_data(CCTK_ARGUMENTS)
{
    static BrillDataContext *prev_bd;

    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    BrillDataContext *bd;

    CCTK_REAL *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, cheb_order);

        brill_solve(bd);

        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_x * 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);
    }

    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_x; k++)
                basis_val_r[(j * cctk_lsh[0] + i) * bd->nb_coeffs_x + k] = bd->basis->eval(r, k);
        }

    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);
                CCTK_REAL xx = x[index], yy = y[index], zz = z[index];

                CCTK_REAL r2 = SQR(xx) + SQR(yy);
                CCTK_REAL r  = sqrt(r2);

                CCTK_REAL q   = bd->q_func(bd, r, zz);
                CCTK_REAL e2q = exp(2 * q);

                CCTK_REAL psi = 1.0, psi4;

                for (int n = 0; n < bd->nb_coeffs_z; n++)
                    for (int m = 0; m < bd->nb_coeffs_x; m++)
                        psi += bd->psi_coeffs->data[n * bd->nb_coeffs_z + m] * basis_val_r[(j * cctk_lsh[0] + i) * bd->nb_coeffs_x + 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);
            }
}