/*
 * 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
 * definitions of the q functions in the exponential in the Brill data
 */

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

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

typedef double (*QFunc)(const struct QFuncContext *s, double rho, double z);

struct QFuncContext {
    const QFunc *qfunc;

    double amplitude;
    double rho0;
    int n;
};

static double q_gundlach(const QFuncContext *s, double rho, double z)
{
    return s->amplitude * SQR(rho) * exp(-(SQR(rho - s->rho0) + SQR(z)));
}

static double dq_rho_gundlach(const QFuncContext *s, double rho, double z)
{
    return s->amplitude * 2 * rho * exp(-SQR(rho - s->rho0) - SQR(z)) * (1.0 - rho * (rho - s->rho0));
}

static double d2q_rho_gundlach(const QFuncContext *s, double rho, double z)
{
    double rho2 = SQR(rho);
    return s->amplitude * 2 * exp(-SQR(rho - s->rho0) - SQR(z)) * (1.0 - 4.0 * rho * (rho - s->rho0) - rho2 + 2.0 * rho2 * SQR(rho - s->rho0));
}

static double d2q_rho_z_gundlach(const QFuncContext *s, double rho, double z)
{
    return s->amplitude * 4 * z * rho * exp(-SQR(rho - s->rho0) - SQR(z)) * (rho * (rho - s->rho0) - 1.0);
}

static double dq_z_gundlach(const QFuncContext *s, double rho, double z)
{
    return -s->amplitude * 2 * z * SQR(rho) * exp(-SQR(rho - s->rho0) - SQR(z));
}

static double d2q_z_gundlach(const QFuncContext *s, double rho, double z)
{
    return s->amplitude * 2 * SQR(rho) * exp(-SQR(rho - s->rho0) - SQR(z)) * (2 * SQR(z) - 1.0);
}

// q function from from PHYSICAL REVIEW D 88, 103009 (2013)
// with σ and ρ_0 hardcoded to 1 for now
static const QFunc q_func_gundlach[9] = {
    [0] = q_gundlach,
    [1] = dq_rho_gundlach,
    [3] = dq_z_gundlach,
    [2] = d2q_rho_gundlach,
    [6] = d2q_z_gundlach,
    [4] = d2q_rho_z_gundlach,
};

static double q_eppley(const QFuncContext *s, double rho, double z)
{
    double r2 = SQR(rho) + SQR(z);
    return s->amplitude * SQR(rho) / (1.0 + pow(r2, s->n / 2.0));
}

static double dq_rho_eppley(const QFuncContext *s, double rho, double z)
{
    double A = s->amplitude;
    double n = s->n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm2 = pow(r, n - 2);

    return A * rho * (2 * (1 + rnm2 * r2) - n * rho2 * rnm2) / SQR(1 + rnm2 * r2);
}

static double dq_z_eppley(const QFuncContext *s, double rho, double z)
{
    double A = s->amplitude;
    double n = s->n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm2 = pow(r, n - 2);

    return - A * n * rho2 * z * rnm2 / SQR(1 + rnm2 * r2);
}

static double d2q_rho_eppley(const QFuncContext *s, double rho, double z)
{
    double A = s->amplitude;
    double n = s->n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm4 = pow(r, n - 4);
    double rn   = rnm4 * SQR(r2);
    double rnp1 = 1 + rn;

    return A * (SQR(n) * SQR(rho2) * rnm4 * (rn - 1) - n * rho2 * rnm4 * (3 * rho2 + 5 * z2) * rnp1 + 2 * SQR(rnp1)) / (rnp1 * SQR(rnp1));
}

static double d2q_z_eppley(const QFuncContext *s, double rho, double z)
{
    double A = s->amplitude;
    double n = s->n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm4 = pow(r, n - 4);
    double rn   = rnm4 * SQR(r2);
    double rnp1 = 1 + rn;

    return A * n * rho2 * rnm4 * (z2 - rho2 - n * z2 + rn * ((1 + n) * z2 - rho2)) / (rnp1 * SQR(rnp1));
}

static double d2q_rho_z_eppley(const QFuncContext *s, double rho, double z)
{
    double A = s->amplitude;
    double n = s->n;
    double rho2 = SQR(rho);
    double z2   = SQR(z);

    double r2 = rho2 + z2;
    double r  = sqrt(r2);
    double rnm4 = pow(r, n - 4);
    double rn   = rnm4 * SQR(r2);
    double rnp1 = 1 + rn;

    return A * n * rho * z * rnm4 * (rn * (n * rho2 - 2 * z2) - 2 * z2 - n * rho2) / (rnp1 * SQR(rnp1));
}

// q function from from PHYSICAL REVIEW D 16, 1609 (1977)
// with λ hardcoded to 1 for now
static const QFunc q_func_eppley[9] = {
    [0] = q_eppley,
    [1] = dq_rho_eppley,
    [3] = dq_z_eppley,
    [2] = d2q_rho_eppley,
    [6] = d2q_z_eppley,
    [4] = d2q_rho_z_eppley,
};

double bdi_qfunc_eval(QFuncContext *s, double rho, double z, int diff_idx)
{
    return s->qfunc[diff_idx](s, rho, z);
}

void bdi_qfunc_free(QFuncContext **ps)
{
    QFuncContext *s = *ps;
    free(s);
    *ps = NULL;
}

int bdi_qfunc_init(const BDContext *bd, QFuncContext **out, enum BDQFuncType type,
                   double amplitude, int n, double rho0)
{
    QFuncContext *s = calloc(1, sizeof(*s));
    int ret = 0;

    if (!s)
        return -ENOMEM;

    switch (type) {
    case BD_Q_FUNC_GUNDLACH:
        s->qfunc = q_func_gundlach;
        break;
    case BD_Q_FUNC_EPPLEY:
        if (n < 4) {
            bdi_log(bd, 0, "Invalid n: %d < 4\n", n);
            ret = -EINVAL;
            goto fail;
        }
        s->qfunc = q_func_eppley;
        break;
    default:
        bdi_log(bd, 0, "Unknown q function type: %d\n", type);
        ret = -EINVAL;
        goto fail;
    }

    s->amplitude = amplitude;
    s->n         = n;
    s->rho0      = rho0;

    *out = s;
    return 0;
fail:
    free(s);
    return ret;
}