/*
 * Basis sets for pseudospectral methods
 * Copyright (C) 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/>.
 */

#include "config.h"

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

#include "basis.h"
#include "common.h"

typedef struct BasisSet {
    /* evaluate the idx-th basis function at the specified point*/
    double (*eval)      (const BasisSetContext *s, double coord, unsigned int idx);
    /* evaluate the first derivative of the idx-th basis function at the specified point*/
    double (*eval_diff1)(const BasisSetContext *s, double coord, unsigned int idx);
    /* evaluate the second derivative of the idx-th basis function at the specified point*/
    double (*eval_diff2)(const BasisSetContext *s, double coord, unsigned int idx);
    /**
     * Get the idx-th collocation point for the specified order.
     * idx runs from 0 to order - 1 (inclusive)
     */
    double (*colloc_point)(const BasisSetContext *s, unsigned int order, unsigned int idx);
} BasisSet;

struct BasisSetContext {
    const BasisSet *bs;
    double sf;
};

/*
 * The basis of even (n = 2 * idx) SB functions (Boyd 2000, Ch 17.9)
 * SB(x, n) = sin((n + 1) arccot(|x| / L))
 * They are symmetric wrt origin and decay as 1/x in infinity.
 */
static double sb_even_eval(const BasisSetContext *s, double coord, unsigned int idx)
{
    double val = atan2(s->sf, coord);

    idx *= 2;   // even only

    return sin((idx + 1) * val);
}

static double sb_even_eval_diff1(const BasisSetContext *s, double coord, unsigned int idx)
{
    double val = atan2(s->sf, coord);

    idx *= 2;   // even only

    return -s->sf * (idx + 1) * cos((idx + 1) * val) / (SQR(s->sf) + SQR(coord));
}

static double sb_even_eval_diff2(const BasisSetContext *s, double coord, unsigned int idx)
{
    const double sf = s->sf;
    double val = atan2(sf, coord);

    idx *= 2;   // even only

    return sf * (idx + 1) * (2 * coord * cos((idx + 1) * val) - sf * (idx + 1) * sin((idx + 1) * val)) / SQR(SQR(sf) + SQR(coord));
}

static double sb_even_colloc_point(const BasisSetContext *s, unsigned int order, unsigned int idx)
{
    double t;

    idx = order - idx - 1;
    //order *= 2;

    //t = (idx + 2) * M_PI / (order + 4);
#if TD_POLAR
    t = (idx + 2) * M_PI / (2 * order + 3);
#else
    t = (idx + 2) * M_PI / (2 * order + 2);
#endif
    return s->sf / tan(t);
}

static const BasisSet sb_even_basis = {
    .eval         = sb_even_eval,
    .eval_diff1   = sb_even_eval_diff1,
    .eval_diff2   = sb_even_eval_diff2,
    .colloc_point = sb_even_colloc_point,
};

static double sb_odd_eval(const BasisSetContext *s, double coord, unsigned int idx)
{
    double val = atan2(s->sf, coord);

    idx = 2 * idx + 2;   // odd only

    return sin((idx) * val);
}

static double sb_odd_eval_diff1(const BasisSetContext *s, double coord, unsigned int idx)
{
    double val = atan2(s->sf, coord);

    idx = 2 * idx + 2;   // odd only

    return -s->sf * (idx) * cos((idx) * val) / (SQR(s->sf) + SQR(coord));
}

static double sb_odd_eval_diff2(const BasisSetContext *s, double coord, unsigned int idx)
{
    const double sf = s->sf;
    double val = atan2(sf, coord);

    idx = 2 * idx + 2;   // odd only

    return sf * (idx) * (2 * coord * cos((idx) * val) - sf * (idx) * sin((idx) * val)) / SQR(SQR(sf) + SQR(coord));
}

static double sb_odd_colloc_point(const BasisSetContext *s, unsigned int order, unsigned int idx)
{
    double t;

    idx = order - idx - 1;
    //order *= 2;

    //t = (idx + 2) * M_PI / (order + 4);
#if TD_POLAR
    t = (idx + 2) * M_PI / (2 * order + 3);
#else
    t = (idx + 2) * M_PI / (2 * order + 3);
#endif
    return s->sf / tan(t);
}

static const BasisSet sb_odd_basis = {
    .eval         = sb_odd_eval,
    .eval_diff1   = sb_odd_eval_diff1,
    .eval_diff2   = sb_odd_eval_diff2,
    .colloc_point = sb_odd_colloc_point,
};

static double tb_even_eval(const BasisSetContext *s, double coord, unsigned int idx)
{
    double val = (coord == 0.0) ? M_PI_2 : atan(s->sf / fabs(coord));

    idx++;
    idx *= 2;   // even only

    return cos(idx * val) - 1.0;
}

static double tb_even_eval_diff1(const BasisSetContext *s, double coord, unsigned int idx)
{
    double val = (coord == 0.0) ? M_PI_2 : atan(s->sf / fabs(coord));

    idx++;
    idx *= 2;   // even only

    return s->sf * idx * SGN(coord) * sin(idx * val) / (SQR(s->sf) + SQR(coord));
}

static double tb_even_eval_diff2(const BasisSetContext *s, double coord, unsigned int idx)
{
    const double sf = s->sf;
    double val = (coord == 0.0) ? M_PI_2 : atan(sf / fabs(coord));

    idx++;
    idx *= 2;   // even only

    return -sf * idx * SGN(coord) * (2 * fabs(coord) * sin(idx * val) + sf * idx * cos(idx * val)) / SQR(SQR(sf) + SQR(coord));
}

static double tb_even_colloc_point(const BasisSetContext *s, unsigned int order, unsigned int idx)
{
    double t;

    idx = order - idx - 1;
    //order *= 2;

    //t = (idx + 2) * M_PI / (order + 4);
    t = (idx + 2) * M_PI / (2 * order + 4);
    return s->sf / tan(t);
}

static const BasisSet tb_even_basis = {
    .eval         = tb_even_eval,
    .eval_diff1   = tb_even_eval_diff1,
    .eval_diff2   = tb_even_eval_diff2,
    .colloc_point = tb_even_colloc_point,
};

static double cos_eval(const BasisSetContext *s, double coord, unsigned int idx)
{
    return cos(idx * coord);
}

static double cos_eval_diff1(const BasisSetContext *s, double coord, unsigned int idx)
{
    return -1.0 * idx * sin(idx * coord);
}

static double cos_eval_diff2(const BasisSetContext *s, double coord, unsigned int idx)
{
    return -1.0 * SQR(idx) * cos(idx * coord);
}

static double cos_colloc_point(const BasisSetContext *s, unsigned int order, unsigned int idx)
{
    return M_PI * (idx + 1) / (order + 2);
}

static const BasisSet cos_basis = {
    .eval         = cos_eval,
    .eval_diff1   = cos_eval_diff1,
    .eval_diff2   = cos_eval_diff2,
    .colloc_point = cos_colloc_point,
};

static double cos_even_eval(const BasisSetContext *s, double coord, unsigned int idx)
{
    return cos(2 * idx * coord);
}

static double cos_even_eval_diff1(const BasisSetContext *s, double coord, unsigned int idx)
{
    return -2.0 * idx * sin(2.0 * idx * coord);
}

static double cos_even_eval_diff2(const BasisSetContext *s, double coord, unsigned int idx)
{
    return -4.0 * SQR(idx) * cos(2.0 * idx * coord);
}

static double cos_even_colloc_point(const BasisSetContext *s, unsigned int order, unsigned int idx)
{
    return M_PI * (idx + 1) / (2 * order + 2);
}

static const BasisSet cos_even_basis = {
    .eval         = cos_even_eval,
    .eval_diff1   = cos_even_eval_diff1,
    .eval_diff2   = cos_even_eval_diff2,
    .colloc_point = cos_even_colloc_point,
};

static double cos_4_eval(const BasisSetContext *s, double coord, unsigned int idx)
{
    return cos(4 * idx * coord);
}

static double cos_4_eval_diff1(const BasisSetContext *s, double coord, unsigned int idx)
{
    return -4.0 * idx * sin(4.0 * idx * coord);
}

static double cos_4_eval_diff2(const BasisSetContext *s, double coord, unsigned int idx)
{
    return -16.0 * SQR(idx) * cos(4.0 * idx * coord);
}

static double cos_4_colloc_point(const BasisSetContext *s, unsigned int order, unsigned int idx)
{
    return M_PI * (idx + 1) / (4 * order + 4);
}

static const BasisSet cos_4_basis = {
    .eval         = cos_4_eval,
    .eval_diff1   = cos_4_eval_diff1,
    .eval_diff2   = cos_4_eval_diff2,
    .colloc_point = cos_4_colloc_point,
};

double tdi_basis_eval(const BasisSetContext *s, enum BSEvalType type,
                      double coord, unsigned int order)
{
    double (*eval)(const BasisSetContext *, double, unsigned int) = NULL;

    switch (type) {
    case BS_EVAL_TYPE_VALUE: eval = s->bs->eval;       break;
    case BS_EVAL_TYPE_DIFF1: eval = s->bs->eval_diff1; break;
    case BS_EVAL_TYPE_DIFF2: eval = s->bs->eval_diff2; break;
    }

    return eval(s, coord, order);
}

double tdi_basis_colloc_point(const BasisSetContext *s, unsigned int order,
                              unsigned int idx)
{
    return s->bs->colloc_point(s, order, idx);
}

void tdi_basis_free(BasisSetContext **pctx)
{
    BasisSetContext *ctx = *pctx;

    if (!ctx)
        return;

    free(ctx);
    *pctx = NULL;
}

int tdi_basis_init(BasisSetContext **pctx, enum BasisFamily family, double sf)
{
    BasisSetContext *ctx;

    ctx = calloc(1, sizeof(*ctx));
    if (!ctx)
        return -ENOMEM;

    switch (family) {
    case BASIS_FAMILY_TB_EVEN:   ctx->bs = &tb_even_basis;   break;
    case BASIS_FAMILY_SB_EVEN:   ctx->bs = &sb_even_basis;   break;
    case BASIS_FAMILY_SB_ODD:    ctx->bs = &sb_odd_basis;    break;
    case BASIS_FAMILY_COS:       ctx->bs = &cos_basis;       break;
    case BASIS_FAMILY_COS_EVEN:  ctx->bs = &cos_even_basis;  break;
    case BASIS_FAMILY_COS_4:     ctx->bs = &cos_4_basis;     break;
    default:
        free(ctx);
        return -EINVAL;
    }

    ctx->sf = sf;

    *pctx = ctx;
    return 0;
}