/*
 * Copyright 2014-2015 Anton Khirnov <anton@khirnov.net>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

/**
 * @file
 * evaluation of the solution on a grid
 */

#include <errno.h>
#include <math.h>
#include <stdlib.h>
#include <string.h>

#include "brill_data.h"
#include "internal.h"

int bd_eval_psi(const BDContext *bd,
                const double *rho, int nb_coords_rho,
                const double *z,   int nb_coords_z,
                const unsigned int diff_order[2],
                double *psi, unsigned int psi_stride)
{
    BDPriv *s = bd->priv;
    const double *sf = bd->basis_scale_factor;
    double (*eval)(double coord, int idx, double sf);

    double *basis_val_rho = NULL, *basis_val_z = NULL;

    double add = 0.0;

    int ret = 0;

    if (diff_order[0] > 2 || diff_order[1] > 2) {
        bdi_log(bd, 0, "Derivatives of higher order than 2 are not supported.\n");
        return -ENOSYS;
    }

    if (diff_order[0] == 0 && diff_order[1] == 0)
        add = 1.0;

    /* precompute the basis values on the grid points */
    basis_val_rho = malloc(sizeof(*basis_val_rho) * s->nb_coeffs[0] * nb_coords_rho);
    basis_val_z   = malloc(sizeof(*basis_val_z)   * s->nb_coeffs[1] * nb_coords_z);
    if (!basis_val_rho || !basis_val_z) {
        ret = -ENOMEM;
        goto fail;
    }

    switch (diff_order[0]) {
    case 0: eval = s->basis[0]->eval;       break;
    case 1: eval = s->basis[0]->eval_diff1; break;
    case 2: eval = s->basis[0]->eval_diff2; break;
    }

    for (int i = 0; i < nb_coords_rho; i++) {
        double rrho = rho[i];
        for (int j = 0; j < s->nb_coeffs[0]; j++)
            basis_val_rho[i * s->nb_coeffs[0] + j] = eval(rrho, j, sf[0]);
    }

    switch (diff_order[1]) {
    case 0: eval = s->basis[1]->eval;       break;
    case 1: eval = s->basis[1]->eval_diff1; break;
    case 2: eval = s->basis[1]->eval_diff2; break;
    }
    for (int i = 0; i < nb_coords_z; i++) {
        double zz   = z[i];
        for (int j = 0; j < s->nb_coeffs[1]; j++)
            basis_val_z[i * s->nb_coeffs[1] + j] = eval(zz, j, sf[1]);
    }

    for (int i = 0; i < nb_coords_z; i++) {
        double *dst = psi + i * psi_stride;

        for (int j = 0; j < nb_coords_rho; j++) {
            double ppsi = add;

            for (int m = 0; m < s->nb_coeffs[1]; m++)
                for (int n = 0; n < s->nb_coeffs[0]; n++)
                    ppsi += s->coeffs[m * s->nb_coeffs[0] + n] * basis_val_rho[j * s->nb_coeffs[0] + n] * basis_val_z[i * s->nb_coeffs[1] + m];

            *dst++ = ppsi;
        }
    }

fail:
    free(basis_val_rho);
    free(basis_val_z);

    return ret;

}

static void eval_metric(const BDContext *bd,
                        const double *rho, int nb_coords_rho,
                        const double *z,   int nb_coords_z,
                        enum BDMetricComponent comp,
                        const unsigned int diff_order[2],
                        const double *psi,
                        const double *dpsi_rho, const double *dpsi_z,
                        const double *d2psi_rho, const double *d2psi_z, const double *d2psi_rho_z,
                        double *out, unsigned int out_stride)
{
    const BDPriv *s = bd->priv;

    if (comp != BD_METRIC_COMPONENT_RHORHO &&
        comp != BD_METRIC_COMPONENT_ZZ     &&
        comp != BD_METRIC_COMPONENT_PHIPHI) {
        memset(out, 0, out_stride * nb_coords_z * sizeof(*out));
        return;
    }

/* the template for the loop over the grid points */
#define EVAL_LOOP(DO_EVAL)                                          \
do {                                                                \
    for (int i = 0; i < nb_coords_z; i++) {                         \
        const double zz = z[i];                                     \
        double *dst = out + i * out_stride;                         \
                                                                    \
        for (int j = 0; j < nb_coords_rho; j++) {                   \
            const double rrho = rho[j];                             \
            const double ppsi = psi[i * nb_coords_rho + j];         \
            const double psi2 = SQR(ppsi);                          \
            const double psi3 = psi2 * ppsi;                        \
            const double psi4 = SQR(psi2);                          \
            double val;                                             \
                                                                    \
            DO_EVAL;                                                \
                                                                    \
            *dst++ = val;                                           \
        }                                                           \
    }                                                               \
} while (0)

#define GRID(arr) (arr[i * nb_coords_rho + j])

    if (comp == BD_METRIC_COMPONENT_PHIPHI) {
        switch (diff_order[0]) {
        case 0:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(val = SQR(rrho) * psi4);                                                                                                                    break;
            case 1: EVAL_LOOP(val = 4 * SQR(rrho) * psi3 * GRID(dpsi_z));                                                                                                 break; // ∂_z
            case 2: EVAL_LOOP(val = 4 * SQR(rrho) * (GRID(d2psi_z) * psi3 + 3 * psi2 * SQR(GRID(dpsi_z))));                                                               break; // ∂_z ∂_z
            }
            break;
        case 1:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(val = 2 * rrho * psi4 + 4 * SQR(rrho) * psi3 * GRID(dpsi_rho));                                                                             break; // ∂_ρ
            case 1: EVAL_LOOP(val = 8 * rrho * psi3 * GRID(dpsi_z) + 12 * SQR(rrho) * psi2 * GRID(dpsi_z) * GRID(dpsi_rho) + 4 * SQR(rrho) * psi3 * GRID(d2psi_rho_z));   break; // ∂_ρ ∂_z
            }
            break;
        case 2:     EVAL_LOOP(val = 4 * SQR(rrho) * psi3 * GRID(d2psi_rho) + 12 * SQR(rrho * ppsi * GRID(dpsi_rho)) + 16 * rrho * psi3 * GRID(dpsi_rho) + 2 * psi4);      break; // ∂_ρ ∂_ρ
        }
    } else {
        // γ_ρρ / γ_zz

/* a wrapper around the actual evaluation expression that provides the q function and its
 * derivatives as needed */
#define DO_EVAL_Q_WRAPPER(DO_EVAL, eval_drho, eval_dz, eval_d2rho, eval_d2z, eval_d2rhoz)                                                    \
do {                                                                                                                                         \
    const double base = psi4 * exp(2 * s->q_func->q(bd, rrho, zz));                                                                          \
    double drho, dz, d2rho, d2z, d2rho_z;                                                                                                    \
                                                                                                                                             \
    if (eval_drho)   drho    = s->q_func->dq_rho(bd, rrho, zz)    + 2 * GRID(dpsi_rho)    / ppsi;                                            \
    if (eval_dz)     dz      = s->q_func->dq_z(bd, rrho, zz)      + 2 * GRID(dpsi_z)      / ppsi;                                            \
    if (eval_d2rho)  d2rho   = s->q_func->d2q_rho(bd, rrho, zz)   + 2 * GRID(d2psi_rho)   / ppsi - 2 * SQR(GRID(dpsi_rho) / ppsi);           \
    if (eval_d2z)    d2z     = s->q_func->d2q_z(bd, rrho, zz)     + 2 * GRID(d2psi_z)     / ppsi - 2 * SQR(GRID(dpsi_z)   / ppsi);           \
    if (eval_d2rhoz) d2rho_z = s->q_func->d2q_rho_z(bd, rrho, zz) + 2 * GRID(d2psi_rho_z) / ppsi - 2 * GRID(dpsi_rho) * GRID(dpsi_z) / psi2; \
                                                                                                                                             \
    DO_EVAL;                                                                                                                                 \
} while (0)

        switch (diff_order[0]) {
        case 0:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = base,                                      0, 0, 0, 0, 0)); break;
            case 1: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 2 * base * dz,                             0, 1, 0, 0, 0)); break;        // ∂_z
            case 2: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * SQR(dz) * base + 2 * base * d2z,       0, 1, 0, 1, 0)); break;        // ∂_z ∂_z
            }
            break;
        case 1:
            switch (diff_order[1]) {
            case 0: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 2 * base * drho,                           1, 0, 0, 0, 0)); break;        // ∂_ρ
            case 1: EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * base * drho * dz + 2 * base * d2rho_z, 1, 1, 0, 0, 1)); break;        // ∂_ρ ∂_z
            }
            break;
        case 2:     EVAL_LOOP(DO_EVAL_Q_WRAPPER(val = 4 * SQR(drho) * base + 2 * base * d2rho,   1, 0, 1, 0, 0)); break;        // ∂_ρ ∂_ρ
        }
    }
}

int bd_eval_metric(const BDContext *bd, const double *rho, int nb_coords_rho,
                   const double *z, int nb_coords_z,
                   int nb_comp, const enum BDMetricComponent *comp,
                   const unsigned int (*diff_order)[2],
                   double **out, unsigned int *out_strides)
{
    const int nb_coords = nb_coords_rho * nb_coords_z;
    double *psi = NULL, *dpsi_rho = NULL, *dpsi_z = NULL, *d2psi_rho = NULL, *d2psi_rho_z = NULL, *d2psi_z = NULL;
    int ret = 0;

    /* check the parameters for validity */
    for (int i = 0; i < nb_comp; i++) {
        if (comp[i] >= 9 || comp[i] < 0) {
            bdi_log(bd, 0, "Invalid component %d: %d\n", i, comp[i]);
            return -EINVAL;
        }

        if (diff_order[i][0] + diff_order[i][1] > 2) {
            bdi_log(bd, 0, "At most second order derivatives are supported.\n");
            return -ENOSYS;
        }
    }

    /* evaluate the conformal factor and its derivatives as necessary */
#define EVAL_PSI(arr, order)                                    \
do {                                                            \
    if (arr)                                                    \
        break;                                                  \
                                                                \
    arr = malloc(nb_coords * sizeof(*arr));                     \
    if (!arr) {                                                 \
        ret = -ENOMEM;                                          \
        goto fail;                                              \
    }                                                           \
                                                                \
    ret = bd_eval_psi(bd, rho, nb_coords_rho, z, nb_coords_z,   \
                      order, arr, nb_coords_rho);               \
    if (ret < 0)                                                \
        goto fail;                                              \
} while (0)

    for (int i = 0; i < nb_comp; i++) {
        if (comp[i] != BD_METRIC_COMPONENT_RHORHO &&
            comp[i] != BD_METRIC_COMPONENT_ZZ     &&
            comp[i] != BD_METRIC_COMPONENT_PHIPHI)
            continue;

        EVAL_PSI(psi,             ((unsigned int[2]){ 0, 0 }));
        if (diff_order[i][0])
            EVAL_PSI(dpsi_rho,    ((unsigned int[2]){ 1, 0 }));
        if (diff_order[i][0] > 1)
            EVAL_PSI(d2psi_rho,   ((unsigned int[2]){ 2, 0 }));
        if (diff_order[i][1])
            EVAL_PSI(dpsi_z,      ((unsigned int[2]){ 0, 1 }));
        if (diff_order[i][1] > 1)
            EVAL_PSI(d2psi_z,     ((unsigned int[2]){ 0, 2 }));
        if (diff_order[i][0] && diff_order[i][1])
            EVAL_PSI(d2psi_rho_z, ((unsigned int[2]){ 1, 1 }));
    }

    /* evaluate the requested metric components */
    for (int i = 0; i < nb_comp; i++)
        eval_metric(bd, rho, nb_coords_rho, z, nb_coords_z,
                    comp[i], diff_order[i],
                    psi, dpsi_rho, dpsi_z, d2psi_rho, d2psi_z, d2psi_rho_z,
                    out[i], out_strides[i]);

fail:
    free(psi);
    free(dpsi_rho);
    free(dpsi_z);
    free(d2psi_rho);
    free(d2psi_rho_z);
    free(d2psi_z);

    return ret;
}

int bd_eval_q(const BDContext *bd, const double *rho, int nb_coords_rho,
              const double *z, int nb_coords_z, const unsigned int diff_order[2],
              double *out, unsigned int out_stride)
{
    BDPriv *s = bd->priv;
    double (*eval)(const struct BDContext *bd, double rho, double z);

    if (diff_order[0] > 2 || diff_order[1] > 2) {
        bdi_log(bd, 0, "Derivatives of higher order than 2 are not supported.\n");
        return -ENOSYS;
    }

    switch (diff_order[0]) {
    case 0:
        switch (diff_order[1]) {
        case 0: eval = s->q_func->q;     break;
        case 1: eval = s->q_func->dq_z;  break;
        case 2: eval = s->q_func->d2q_z; break;
        }
        break;
    case 1:
        switch (diff_order[1]) {
        case 0: eval = s->q_func->dq_rho;    break;
        case 1: eval = s->q_func->d2q_rho_z; break;
        }
        break;
    case 2:     eval = s->q_func->d2q_rho;   break;
    }

    for (int i = 0; i < nb_coords_z; i++) {
        double *dst = out + i * out_stride;
        for (int j = 0; j < nb_coords_rho; j++)
            *dst++ = eval(bd, rho[j], z[i]);
    }

    return 0;
}