/*
 * Copyright 2014-2016 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"
#include "qfunc.h"

static int eval_psi_cylindric(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;
    enum BSEvalType type;

    double *basis_val_rho = NULL, *basis_val_z = NULL;

    double add = (diff_order[0] || diff_order[1]) ? 0.0 : 1.0;

    int ret = 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: type = BS_EVAL_TYPE_VALUE; break;
    case 1: type = BS_EVAL_TYPE_DIFF1; break;
    case 2: type = BS_EVAL_TYPE_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] = bdi_basis_eval(s->basis[0], type, rrho, j);
    }

    switch (diff_order[1]) {
    case 0: type = BS_EVAL_TYPE_VALUE; break;
    case 1: type = BS_EVAL_TYPE_DIFF1; break;
    case 2: type = BS_EVAL_TYPE_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] = bdi_basis_eval(s->basis[1], type, zz, j);
    }

    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;

}

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

    switch (s->coord_system) {
    case COORD_SYSTEM_CYLINDRICAL: ret = eval_psi_cylindric(bd, rho, nb_coords_rho, z, nb_coords_z,
                                                            diff_order, psi, psi_stride);
                                   break;
    case COORD_SYSTEM_RADIAL:      ret = bdi_eval_psi_radial(bd, rho, nb_coords_rho, z, nb_coords_z,
                                                             diff_order, psi, psi_stride);
                                   break;
    }

    return ret;

}

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[9] = { NULL }, *q[9] = { 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 */
    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;

        for (int j = 0; j < ARRAY_ELEMS(psi); j++) {
            if (psi[j])
                continue;
            if (j / 3 + j % 3 > 2)
                continue;
            if (!(j % 3 <= diff_order[i][0] ||
                  j / 3 <= diff_order[i][1]))
                continue;

            psi[j] = malloc(nb_coords * sizeof(*psi[j]));
            if (!psi[j]) {
                ret = -ENOMEM;
                goto fail;
            }

            ret = bd_eval_psi(bd, rho, nb_coords_rho, z, nb_coords_z,
                              (unsigned int[2]){ j % 3, j / 3 }, psi[j],
                              nb_coords_rho);
            if (ret < 0)
                goto fail;

            q[j] = malloc(nb_coords * sizeof(*q[j]));
            if (!q[j]) {
                ret = -ENOMEM;
                goto fail;
            }

            ret = bd_eval_q(bd, rho, nb_coords_rho, z, nb_coords_z,
                            (unsigned int[2]){ j % 3, j / 3 }, q[j],
                            nb_coords_rho);
            if (ret < 0)
                goto fail;
        }
    }

    /* evaluate the requested metric components */
    for (int i = 0; i < nb_comp; i++) {
        bdi_eval_metric(bd, rho, nb_coords_rho, z, nb_coords_z,
                        comp[i], diff_order[i],
                        psi, q, out[i], out_strides[i]);
    }

fail:
    for (int i = 0; i < ARRAY_ELEMS(psi); i++)
        free(psi[i]);
    for (int i = 0; i < ARRAY_ELEMS(q); i++)
        free(q[i]);

    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;
    int diff_idx;

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

    diff_idx = diff_order[1] * 3 + diff_order[0];

    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++ = bdi_qfunc_eval(s->qfunc, rho[j], z[i], diff_idx);
    }

    return 0;
}