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

#include <lapacke.h>

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

#define SQR(x) ((x) * (x))

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

void brill_data(CCTK_ARGUMENTS)
{
    static BDContext *prev_bd;

    DECLARE_CCTK_ARGUMENTS;
    DECLARE_CCTK_PARAMETERS;

    BDContext *bd;
    double *r_val, *z_val, *psi_val, *q_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 equation for ψ */
    if (!prev_bd) {

        bd = bd_context_alloc();
        if (!bd)
            CCTK_WARN(0, "Memory allocation failed\n");

        bd->q_func_type  = BD_Q_FUNC_GUNDLACH;
        bd->amplitude    = amplitude;
        bd->eppley_n     = eppley_n;
        bd->nb_coeffs[0] = basis_order_0;
        bd->nb_coeffs[1] = basis_order_1;
        bd->basis_scale_factor[0] = scale_factor;
        bd->basis_scale_factor[1] = scale_factor;

        ret = bd_solve(bd);
        if (ret < 0)
            CCTK_WARN(0, "Error solving the Brill wave initial data equation\n");

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

    /* construct the coordinate vectors to be passed to the library */
    r_val = malloc(sizeof(*r_val) * cctk_lsh[1] * cctk_lsh[0]);
    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));

            r_val[j * cctk_lsh[0] + i] = r;
        }
    z_val = malloc(sizeof(*z_val) * cctk_lsh[2]);
    for (int i = 0; i < cctk_lsh[2]; i++)
        z_val[i] = z[CCTK_GFINDEX3D(cctkGH, 0, 0, i)];

    psi_val = malloc(sizeof(*psi_val) * grid_size);
    q_val   = malloc(sizeof(*q_val)   * grid_size);

#pragma omp parallel for
    for (int j = 0; j < cctkGH->cctk_lsh[1]; j++) {
        double *r_y   = r_val   + j * cctkGH->cctk_lsh[0];
        double *psi_y = psi_val + j * cctkGH->cctk_lsh[0] * cctkGH->cctk_lsh[2];
        double *q_y   = q_val   + j * cctkGH->cctk_lsh[0] * cctkGH->cctk_lsh[2];

        int err;

        err = bd_eval_psi(bd, r_y, cctk_lsh[0], z_val, cctk_lsh[2],
                          (unsigned int[2]){ 0, 0}, psi_y, cctk_lsh[0]);
        if (err < 0)
            CCTK_WARN(0, "Error evaluating the conformal factor");

        err = bd_eval_q(bd, r_y, cctk_lsh[0], z_val, cctk_lsh[2],
                        (unsigned int[2]){ 0, 0}, q_y, cctk_lsh[0]);
        if (err < 0)
            CCTK_WARN(0, "Error evaluating the q function");

        for (int k = 0; k < cctk_lsh[2]; k++)
            for (int i = 0; i < cctk_lsh[0]; i++) {
                int    index = CCTK_GFINDEX3D(cctkGH, i, j, k);
                double xx = x[index], yy = y[index];
                double r2 = SQR(xx) + SQR(yy);

                double e2q = exp(2 * q_y[k * cctkGH->cctk_lsh[0] + i]);
                double psi = psi_y[k * cctk_lsh[0] + i];

                double 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);

            }
    }

    free(r_val);
    free(z_val);
    free(psi_val);
    free(q_val);
}