#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 <teukolsky_data.h>

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

#define SQR(x) ((x) * (x))
#define SIGN(x) ((x) > 0 ? 1.0 : -1.0)

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

void teukolsky_data(CCTK_ARGUMENTS)
{
    static TDContext *prev_td;

    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    TDContext *td;
    double *r_val, *theta_val, *psi_val, *krr_val, *kpp_val, *krt_val;
    int ret;

    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 constraints */
    if (!prev_td) {
        const char *omp_threads = getenv("OMP_NUM_THREADS");

        td = td_context_alloc();
        if (!td)
            CCTK_WARN(0, "Memory allocation failed\n");

        td->amplitude    = amplitude;
        td->nb_coeffs[0] = basis_order_0;
        td->nb_coeffs[1] = basis_order_1;
        td->basis_scale_factor[0] = scale_factor;
        td->basis_scale_factor[1] = scale_factor;
        td->solution_branch = solution_branch;
        td->max_iter        = max_iter;
        td->family          = family;

        if (omp_threads)
            td->nb_threads = strtol(omp_threads, NULL, 0);
        if (td->nb_threads <= 0)
            td->nb_threads = 1;


        ret = td_solve(td, NULL);
        if (ret < 0)
            CCTK_WARN(0, "Error solving the Teukolsky wave initial data equations\n");

        prev_td = td;
    } else
        td = prev_td;

    memset(gxz, 0, sizeof(*gxz) * grid_size);
    memset(gxy, 0, sizeof(*gxy) * grid_size);
    memset(gyz, 0, sizeof(*gyz) * grid_size);

    memset(kxy, 0, sizeof(*kxy) * grid_size);
    memset(kyz, 0, sizeof(*kyz) * grid_size);

    /* construct the coordinate vectors to be passed to the library */
    r_val     = malloc(sizeof(*r_val)     * (cctk_lsh[2] * cctk_lsh[0] + 1));
    theta_val = malloc(sizeof(*theta_val) * (cctk_lsh[2] * cctk_lsh[0] + 1));
    psi_val   = malloc(sizeof(*psi_val)   * (cctk_lsh[2] * cctk_lsh[0] + 1));
    krr_val   = malloc(sizeof(*krr_val)   * (cctk_lsh[2] * cctk_lsh[0] + 1));
    kpp_val   = malloc(sizeof(*kpp_val)   * (cctk_lsh[2] * cctk_lsh[0] + 1));
    krt_val   = malloc(sizeof(*krt_val)   * (cctk_lsh[2] * cctk_lsh[0] + 1));
    for (int j = 0; j < cctk_lsh[2]; j++)
        for (int i = 0; i < cctk_lsh[0]; i++) {
            double xx = x[CCTK_GFINDEX3D(cctkGH, i, 0, j)];
            double zz = z[CCTK_GFINDEX3D(cctkGH, i, 0, j)];
            double r     = sqrt(SQR(xx) + SQR(zz));
            double theta = (fabs(r) > 1e-15) ? acos(zz / r) : 0.0;

            r_val[j * cctk_lsh[0] + i]     = r;
            theta_val[j * cctk_lsh[0] + i] = theta;
        }
    r_val[cctk_lsh[2] * cctk_lsh[0]] = 0.0;
    theta_val[cctk_lsh[2] * cctk_lsh[0]] = M_PI / 2;

    ret  = td_eval_psi(td, cctk_lsh[2] * cctk_lsh[0] + 1, r_val, theta_val,
                       (const unsigned int [2]){ 0, 0},
                       psi_val);
    ret |= td_eval_krr(td, cctk_lsh[2] * cctk_lsh[0] + 1, r_val, theta_val,
                       (const unsigned int [2]){ 0, 0},
                       krr_val);
    ret |= td_eval_kpp(td, cctk_lsh[2] * cctk_lsh[0] + 1, r_val, theta_val,
                       (const unsigned int [2]){ 0, 0},
                       kpp_val);
    ret |= td_eval_krt(td, cctk_lsh[2] * cctk_lsh[0] + 1, r_val, theta_val,
                       (const unsigned int [2]){ 0, 0},
                       krt_val);
    if (ret)
        CCTK_WARN(0, "Error evaluating the variables\n");

    for (int j = 0; j < cctkGH->cctk_lsh[2]; j++) {
        for (int i = 0; i < cctk_lsh[0]; i++) {
            int    idx_dst = CCTK_GFINDEX3D(cctkGH, i, 0, j);
            int    idx_src = j * cctk_lsh[0] + i;
            double xx = x[idx_dst], zz = z[idx_dst];
            double     r = r_val[idx_src];
            double theta = theta_val[idx_src];
            double s2t    = sin(2.0 * theta);

            double psi = psi_val[idx_src];
            double psi4 = SQR(SQR(psi));

            double krr = krr_val[idx_src];
            double kpp = kpp_val[idx_src];
            double ktt = -(krr + kpp);
            double krt = krt_val[idx_src];

            gxx[idx_dst] = psi4;
            gyy[idx_dst] = psi4;
            gzz[idx_dst] = psi4;

            kyy[idx_dst] = psi4 * kpp;

            if (fabs(r) > 1e-15) {
                kxx[idx_dst] = psi4 * (SQR(xx / r)      * krr + SQR(zz / r)      * ktt + 2.0 * zz * fabs(xx) / (r * SQR(r)) * krt);
                kzz[idx_dst] = psi4 * (SQR(zz / r)      * krr + SQR(xx / r)      * ktt - 2.0 * zz * fabs(xx) / (r * SQR(r)) * krt);
                kxz[idx_dst] = psi4 * (xx * zz / SQR(r) * krr - xx * zz / SQR(r) * ktt + (-xx * fabs(xx) + SIGN(xx) * SQR(zz)) * krt / (r * SQR(r)));
            } else {
                kxx[idx_dst] = psi4 * krr_val[cctk_lsh[0] * cctk_lsh[2]];
                kzz[idx_dst] = psi4 * krr;
                kxz[idx_dst] = 0.0;
            }
        }
    }

    if (!strcmp(initial_lapse, "teukolsky_max")) {
        ret = td_eval_lapse(td, cctk_lsh[2] * cctk_lsh[0], r_val, theta_val,
                            (const unsigned int [2]){ 0, 0}, alp);
        if (ret)
            CCTK_WARN(0, "Error evaluating the lapse\n");
    }

    free(r_val);
    free(theta_val);
    free(psi_val);
    free(krr_val);
    free(kpp_val);
    free(krt_val);
}