#include <ctype.h>
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gsl/gsl_spline.h>
#include "cctk.h"
#include "cctk_Arguments.h"
#include "cctk_Parameters.h"

#include "mdefs.h"

#define MAX(x,y) (x) > (y) ? (x) : (y)
#define SQR(x) ((x)*(x))

#ifndef M_SQRT2
#define M_SQRT2 1.4142135623730951
#endif
#ifndef M_SQRT1_2
#define M_SQRT1_2 0.7071067811865476
#endif

#define MASS 1
/*
 * the constant C from e.g. 0908.1063v4
 */
#define TRUMPET_CONST (3*sqrt(3)*SQR(MASS)/4)
/*
 * small number to avoid r=0 singularities
 */
#define EPS 1E-08
/*
 * isotropic/coordinate radius
 */
#define ISO_R(x, y, z) (sqrt(SQR(x) + SQR(y) + SQR(z) + EPS))

#define TRUMPET_ALPHA(R) (sqrt(1 - 2*MASS/R + SQR(TRUMPET_CONST)/SQR(SQR(R))))

static inline double r_areal_to_isotropic(double r, double m)
{
    double par   = sqrt(4*SQR(r) + 4*m*r + 3*SQR(m));
    double term1 = (2*r + m + par)/4;
    double term2 = (4 + 3*M_SQRT2)*(2*r - 3*m)/(8*r + 6*m + 3*M_SQRT2*par);
    return term1*pow(term2, M_SQRT1_2);
}

/**
 * get maximal isotropic r
 */
static CCTK_REAL get_max_r(CCTK_INT dim)
{
    CCTK_REAL phys_min[dim], phys_max[dim], ext_min[dim], ext_max[dim], int_min[dim], int_max[dim], step[dim];
    CCTK_REAL max = 0;

    GetDomainSpecification(dim, phys_min, phys_max, int_min, int_max, ext_min, ext_max, step);

    for (int i = 0; i < dim; i++) {
        max = MAX(max, abs(ext_min[i]));
        max = MAX(max, abs(ext_max[i]));
    }
    return sqrt(dim)*2*max;
}

static void get_r_tables(double **pr_iso, double **pr_areal, CCTK_INT *size, CCTK_INT dim)
{
#define STEP 0.0001
    double *r_iso, *r_areal, max = get_max_r(dim);
    int count = max*2/STEP, i = 0;

    r_iso   = malloc(count*sizeof(*r_iso));
    r_areal = malloc(count*sizeof(*r_areal));

    for (i = 0;;i++) {
        if (i >= count) {
            count *= 2;
            r_iso = realloc(r_iso, count*sizeof(*r_iso));
            r_areal = realloc(r_areal, count*sizeof(*r_areal));
        }

        r_areal[i] = i*STEP + 1.5;
        r_iso[i]   = r_areal_to_isotropic(r_areal[i], 1);
        if (r_iso[i] > max)
            break;
    }
    *pr_iso   = r_iso;
    *pr_areal = r_areal;
    *size     = i;
}

/* those equations come from trumpet.nb */
static long double sqrt_factor(long double R, long double r, long double x, long double alpha, long double beta)
{
    long double R2 = SQR(R);
    return sqrtl((R2*(R + alpha*beta*r) - beta*TRUMPET_CONST*x)*(R2*(R - alpha*beta*r) - beta*TRUMPET_CONST*x));
}

static inline CCTK_REAL K_11(CCTK_REAL x, CCTK_REAL y, CCTK_REAL z,
                             CCTK_REAL R, CCTK_REAL r,
                             CCTK_REAL beta, CCTK_REAL M,
                             CCTK_REAL alpha)
{
    long double b2 = SQR(beta);
    long double r2 = SQR(r);
    long double R2 = SQR(R);

    long double n1 = TRUMPET_CONST*(1.0 - 3*SQR(x/r))*(beta*x*TRUMPET_CONST/R - R2)/r2;
    long double n2 = b2*beta*M*(-Power(TRUMPET_CONST,2)/(R2*R2) + Power(alpha,2))*x;
    long double n3 = b2*TRUMPET_CONST*M*(1.0 + 2*SQR(x/r))/R;
    long double n4 = beta*x*(SQR(R/r)*(-2*M + (-1 + alpha)*alpha*R));

    long double den = (b2 - 1)*sqrt_factor(R, r, x, alpha, beta);

    return (n1 + n2 + n3 + n4)/den;
}

static inline CCTK_REAL K_22(CCTK_REAL x, CCTK_REAL y, CCTK_REAL z,
                             CCTK_REAL R, CCTK_REAL r,
                             CCTK_REAL beta, CCTK_REAL M,
                             CCTK_REAL alpha)
{
    long double r2 = SQR(r);
    long double R2 = SQR(R);

    long double n1 = TRUMPET_CONST*(R2 - beta*TRUMPET_CONST*(x/R))*(1.0 - 3*SQR(y/r))/r2;
    long double n2 = alpha*(alpha - 1)*beta*SQR(R/r)*R*x;

    long double sq = sqrt_factor(R, r, x, alpha, beta);

    return (n1 + n2)/sq;
}

static inline CCTK_REAL K_33(CCTK_REAL x, CCTK_REAL y, CCTK_REAL z,
                             CCTK_REAL R, CCTK_REAL r,
                             CCTK_REAL beta, CCTK_REAL M,
                             CCTK_REAL alpha)
{
    return K_22(x, z, y, R, r, beta, M, alpha);
}

static inline CCTK_REAL K_12(CCTK_REAL x, CCTK_REAL y, CCTK_REAL z,
                             CCTK_REAL R, CCTK_REAL r,
                             CCTK_REAL beta, CCTK_REAL M,
                             CCTK_REAL alpha)
{
    long double b2 = SQR(beta);
    long double r2 = SQR(r);
    long double R2 = SQR(R);

    long double n1 = TRUMPET_CONST*(b2*M/R + 3*SQR(R/r))*x - 3*beta*SQR(TRUMPET_CONST)*SQR(x/r)/R;
    long double n2 = beta*SQR(R/r)*(-M + (-1 + alpha)*alpha*R);
    long double den = -sqrt_factor(R, r, x, alpha, beta)*sqrtl(1 - b2);     // really minus?

    return y*(n1/r2 + n2)/den;
}

static inline CCTK_REAL K_13(CCTK_REAL x, CCTK_REAL y, CCTK_REAL z,
                             CCTK_REAL R, CCTK_REAL r,
                             CCTK_REAL beta, CCTK_REAL M,
                             CCTK_REAL alpha)
{
    return K_12(x, z, y, R, r, beta, M, alpha);
}

static inline CCTK_REAL K_23(CCTK_REAL x, CCTK_REAL y, CCTK_REAL z,
                             CCTK_REAL R, CCTK_REAL r,
                             CCTK_REAL beta, CCTK_REAL M,
                             CCTK_REAL alpha)
{
    long double r2 = SQR(r);
    long double R2 = SQR(R);

    long double num = -3*TRUMPET_CONST*(R2 - beta*TRUMPET_CONST*x/R)*(y/r)*(z/r);
    long double den = r2*sqrt_factor(R, r, x, alpha, beta);

    return num/den;
}

void trumpet_data(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    double gamma = 1.0/sqrt(1.0 - SQR(boost_velocity));

    gsl_interp_accel *acc;
    gsl_spline    *spline;

    double *r_iso, *r_areal;
    CCTK_INT size;

    get_r_tables(&r_iso, &r_areal, &size, cctk_dim);

    spline = gsl_spline_alloc(gsl_interp_cspline, size);
    gsl_spline_init(spline, r_iso, r_areal, size);
    acc = gsl_interp_accel_alloc();

    memset(gyz, 0, sizeof(*gxy)*(CCTK_GFINDEX3D(cctkGH, cctk_lsh[0]-1, cctk_lsh[1]-1, cctk_lsh[2]-1) + 1));

#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);
                CCTK_REAL xx = gamma*x[index], yy = y[index], zz = z[index];
                CCTK_REAL r = ISO_R(xx, yy, zz);
                CCTK_REAL R = gsl_spline_eval(spline, r, acc);
                CCTK_REAL alpha = TRUMPET_ALPHA(R);

                long double r2 = SQR(r);
                long double R2 = SQR(R);
                long double b2 = SQR(boost_velocity);

                kxx[index] = K_11(xx, yy, zz, R, r, boost_velocity, MASS, alpha);
                kyy[index] = K_22(xx, yy, zz, R, r, boost_velocity, MASS, alpha);
                kzz[index] = K_33(xx, yy, zz, R, r, boost_velocity, MASS, alpha);
                kxy[index] = K_12(xx, yy, zz, R, r, boost_velocity, MASS, alpha);
                kxz[index] = K_13(xx, yy, zz, R, r, boost_velocity, MASS, alpha);
                kyz[index] = K_23(xx, yy, zz, R, r, boost_velocity, MASS, alpha);

                gxx[index] = (-SQR(R/r) + b2*(-SQR(TRUMPET_CONST/R2) + Power(alpha,2)) + 2*boost_velocity*TRUMPET_CONST*xx/(R*r2)) / (b2 - 1);
                gyy[index] = SQR(R/r);
                gzz[index] = SQR(R/r);
                gxy[index] = -((boost_velocity*TRUMPET_CONST*yy)/(Sqrt(1 - b2)*r2*R));
                gxz[index] = -((boost_velocity*TRUMPET_CONST*zz)/(Sqrt(1 - b2)*r2*R));
            }
    free(r_iso);
    free(r_areal);
    gsl_interp_accel_free(acc);
    gsl_spline_free(spline);
}

void trumpet_lapse(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    double gamma = 1.0/sqrt(1.0 - SQR(boost_velocity));

    gsl_interp_accel *acc;
    gsl_spline    *spline;

    double *r_iso, *r_areal;
    CCTK_INT size;

    get_r_tables(&r_iso, &r_areal, &size, cctk_dim);

    spline = gsl_spline_alloc(gsl_interp_cspline, size);
    gsl_spline_init(spline, r_iso, r_areal, size);
    acc = gsl_interp_accel_alloc();

#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);
                CCTK_REAL xx = gamma*x[index], yy = y[index], zz = z[index];
                CCTK_REAL r = ISO_R(xx, yy, zz);
                CCTK_REAL R = gsl_spline_eval(spline, r, acc);
                CCTK_REAL alpha = TRUMPET_ALPHA(R);

                long double r2 = SQR(r);
                long double R2 = SQR(R);
                long double b2 = SQR(boost_velocity);

                alp[index] = alpha*R2*R*sqrt((b2 - 1) / (Power(alpha,2)*b2*r2*R2*R2 - Power(R2*R - boost_velocity*TRUMPET_CONST*xx,2)));
            }
    free(r_iso);
    free(r_areal);
    gsl_interp_accel_free(acc);
    gsl_spline_free(spline);
}

void trumpet_shift(CCTK_ARGUMENTS)
{
    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    double gamma = 1.0/sqrt(1.0 - SQR(boost_velocity));

    gsl_interp_accel *acc;
    gsl_spline    *spline;

    double *r_iso, *r_areal;
    CCTK_INT size;

    get_r_tables(&r_iso, &r_areal, &size, cctk_dim);

    spline = gsl_spline_alloc(gsl_interp_cspline, size);
    gsl_spline_init(spline, r_iso, r_areal, size);
    acc = gsl_interp_accel_alloc();

#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);
                CCTK_REAL xx = gamma*x[index], yy = y[index], zz = z[index];
                CCTK_REAL r = ISO_R(xx, yy, zz);
                CCTK_REAL R = gsl_spline_eval(spline, r, acc);
                CCTK_REAL alpha = TRUMPET_ALPHA(R);

                long double r2 = SQR(r);
                long double R2 = SQR(R);
                long double b2 = SQR(boost_velocity);
                long double a2 = SQR(alpha);

                betax[index] = -((a2*boost_velocity*r2*R2*R2 + TRUMPET_CONST*Power(R,3)*xx + b2*TRUMPET_CONST*Power(R,3)*xx - boost_velocity*(Power(R,6) + Power(TRUMPET_CONST,2)*Power(xx,2)))/
     (a2*b2*r2*R2*R2 - Power(Power(R,3) - boost_velocity*TRUMPET_CONST*xx,2)));
                betay[index] = (Sqrt(1 - b2)*TRUMPET_CONST*(-Power(R,3) + boost_velocity*TRUMPET_CONST*xx)*yy)/(a2*b2*r2*R2*R2 - Power(Power(R,3) - boost_velocity*TRUMPET_CONST*xx,2));
                betaz[index] = (Sqrt(1 - b2)*TRUMPET_CONST*(-Power(R,3) + boost_velocity*TRUMPET_CONST*xx)*zz)/(a2*b2*r2*R2*R2 - Power(Power(R,3) - boost_velocity*TRUMPET_CONST*xx,2));

            }
    free(r_iso);
    free(r_areal);
    gsl_interp_accel_free(acc);
    gsl_spline_free(spline);
}