/*
 * Copyright 2017 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 <errno.h>
#include <math.h>

#include "common.h"
#include "pssolve.h"
#include "td_constraints.h"

struct TDConstraintEvalPriv {
    double *k_rtheta;
    double *dk_rtheta[2];
};

typedef struct TDFamilyDef {
    double (*eval_krt)   (TDConstraintEvalContext *ctx, double r, double theta);
    double (*eval_krt_dr)(TDConstraintEvalContext *ctx, double r, double theta);
    double (*eval_krt_dt)(TDConstraintEvalContext *ctx, double r, double theta);
    double a_converge;
    double a_diverge;
} TDFamilyDef;

static double k_rtheta_eval_ae_time_antisym(TDConstraintEvalContext *ctx,
                                            double r, double theta)
{
    const double r2 = SQR(r);
    const double r3 = r * r2;
    return -60.0 * sqrt(2) * (3.0 * r - 2.0 * r3) * exp(-r2) * sin(2.0 * theta) * ctx->amplitude / sqrt(M_PI);
}

static double k_rtheta_eval_dr_ae_time_antisym(TDConstraintEvalContext *ctx,
                                               double r, double theta)
{
    const double r2 = SQR(r);
    const double r4 = SQR(r2);
    return -60.0 * sqrt(2) * (3.0 - 12.0 * r2 + 4.0 * r4) * exp(-r2) * sin(2.0 * theta) * ctx->amplitude / sqrt(M_PI);
}

static double k_rtheta_eval_dt_ae_time_antisym(TDConstraintEvalContext *ctx,
                                               double r, double theta)
{
    const double r2 = SQR(r);
    const double r3 = r * r2;
    return -120.0 * sqrt(2) * (3.0 * r - 2.0 * r3) * exp(-r2) * cos(2.0 * theta) * ctx->amplitude / sqrt(M_PI);
}

static const TDFamilyDef ae_time_antisym = {
    .eval_krt    = k_rtheta_eval_ae_time_antisym,
    .eval_krt_dr = k_rtheta_eval_dr_ae_time_antisym,
    .eval_krt_dt = k_rtheta_eval_dt_ae_time_antisym,
    .a_converge  = 0.0092078837420,
    .a_diverge   = 0.0092078837427,
};

#if 0
static const TDFamilyDef simple_time_antisym = {
    .eval_krt    = k_rtheta_eval_simple_time_antisym,
    .eval_krt_dr = k_rtheta_eval_dr_simple_time_antisym,
    .eval_krt_dt = k_rtheta_eval_dt_simple_time_antisym,
    .a_converge  = 1.018918628476058,
    .a_diverge   = 1.018918628476968,
};
#endif

static const TDFamilyDef *td_families[] = {
    [TD_FAMILY_AE_TIME_ANTISYM] = &ae_time_antisym,
};

static void
constraint_eval_ham_confflat(TDConstraintEvalContext *ctx,
                             double *vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                             double *out)
{
    TDConstraintEvalPriv *priv = ctx->priv;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];
        const double tt = tan(theta);

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            const double K_rt = priv->k_rtheta[idx];

            const double K_rr     = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_00][idx];
            const double K_pp     = vars[TD_CONSTRAINT_VAR_CURV_PHIPHI][PSSOLVE_DIFF_ORDER_00][idx];
            const double psi      = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_00][idx] + 1.0;
            const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
            const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];
            const double d2psi_rr = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_20][idx];
            const double d2psi_tt = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_02][idx];

            const double Km[3][3] = {{ K_rr,                    K_rt,    0 },
                                     { K_rt / SQR(r), -(K_rr + K_pp),    0 },
                                     { 0,                          0, K_pp }};

            const double r2   = SQR(r);
            const double psi2 = SQR(psi);
            const double psi4 = SQR(psi2);
            const double Rscal = -8.0 * (d2psi_rr + 2.0 * dpsi_r / r + d2psi_tt / r2 + dpsi_t / (r2 * tt)) / (psi4 * psi);

            double K2 = 0.0;
            for (int i = 0; i < 3; i++)
                for (int j = 0; j < 3; j++)
                    K2 += Km[i][j] * Km[j][i];

            out[idx] = Rscal - K2;
        }
    }
}

static void
constraint_eval_mom_r_confflat(TDConstraintEvalContext *ctx,
                               double *vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                               double *out)
{
    TDConstraintEvalPriv *priv = ctx->priv;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];
        const double tt = tan(theta);

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            const double K_rt     = priv->k_rtheta[idx];
            const double dK_rt_t  = priv->dk_rtheta[1][idx];

            const double K_rr     = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_00][idx];
            const double dK_rr_r  = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_10][idx];
            const double psi      = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_00][idx] + 1.0;
            const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
            const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];

            const double r2   = SQR(r);

            out[idx] = 6.0 * (K_rr * dpsi_r + K_rt * dpsi_t / r2) / psi +
                       dK_rr_r + 3.0 * K_rr / r + K_rt / (r2 * tt) + dK_rt_t / r2;
        }
    }
}

static void
constraint_eval_mom_t_confflat(TDConstraintEvalContext *ctx,
                               double *vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                               double *out)
{
    TDConstraintEvalPriv *priv = ctx->priv;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];
        const double tt = tan(theta);

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            const double K_rt     = priv->k_rtheta[idx];
            const double dK_rt_r  = priv->dk_rtheta[0][idx];

            const double K_rr     = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_00][idx];
            const double dK_rr_t  = vars[TD_CONSTRAINT_VAR_CURV_RR][PSSOLVE_DIFF_ORDER_01][idx];
            const double K_pp     = vars[TD_CONSTRAINT_VAR_CURV_PHIPHI][PSSOLVE_DIFF_ORDER_00][idx];
            const double dK_pp_t  = vars[TD_CONSTRAINT_VAR_CURV_PHIPHI][PSSOLVE_DIFF_ORDER_01][idx];
            const double psi      = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_00][idx] + 1.0;
            const double dpsi_r   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_10][idx];
            const double dpsi_t   = vars[TD_CONSTRAINT_VAR_CONFFAC][PSSOLVE_DIFF_ORDER_01][idx];

            out[idx] = -(K_rr + 2.0 * K_pp) / tt - 6.0 * (K_rr + K_pp) * dpsi_t / psi +
                       6.0 * K_rt * dpsi_r / psi - dK_pp_t - dK_rr_t + dK_rt_r + 2.0 * K_rt / r;
        }
    }
}

int tdi_constraint_eval_init(TDConstraintEvalContext *ctx)
{
    TDConstraintEvalPriv *priv = ctx->priv;
    const TDFamilyDef *fd = td_families[ctx->family];
    int ret;

    ret = posix_memalign((void**)&priv->k_rtheta, 32,
                         sizeof(*priv->k_rtheta) * ctx->nb_coords[0] * ctx->nb_coords[1]);
    if (ret)
        return -ENOMEM;

    for (int i = 0; i < ARRAY_ELEMS(priv->dk_rtheta); i++) {
        ret = posix_memalign((void**)&priv->dk_rtheta[i], 32,
                             sizeof(*priv->dk_rtheta[i]) * ctx->nb_coords[0] * ctx->nb_coords[1]);
        if (ret)
            return -ENOMEM;
    }

    ctx->a_converge = fd->a_converge;
    ctx->a_diverge  = fd->a_diverge;

    for (int idx1 = 0; idx1 < ctx->nb_coords[1]; idx1++) {
        const double theta = ctx->coords[1][idx1];

        for (int idx0 = 0; idx0 < ctx->nb_coords[0]; idx0++) {
            const double r = ctx->coords[0][idx0];

            const int idx = idx1 * ctx->nb_coords[0] + idx0;

            priv->k_rtheta[idx]     = fd->eval_krt   (ctx, r, theta);
            priv->dk_rtheta[0][idx] = fd->eval_krt_dr(ctx, r, theta);
            priv->dk_rtheta[1][idx] = fd->eval_krt_dt(ctx, r, theta);
        }
    }

    return 0;
}

int tdi_constraint_eval_alloc(TDConstraintEvalContext **pctx, enum TDFamily family)
{
    TDConstraintEvalContext *ctx;
    int ret;

    if (family < 0 || family >= ARRAY_ELEMS(td_families))
        return -EINVAL;

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

    ctx->priv = calloc(1, sizeof(*ctx->priv));
    if (!ctx->priv) {
        tdi_constraint_eval_free(&ctx);
        return -ENOMEM;
    }

    ctx->family = family;

    *pctx = ctx;
    return 0;
}

void tdi_constraint_eval_free(TDConstraintEvalContext **pctx)
{
    TDConstraintEvalContext *ctx = *pctx;

    if (!ctx)
        return;

    if (ctx->priv) {
        free(ctx->priv->k_rtheta);
        free(ctx->priv->dk_rtheta[0]);
        free(ctx->priv->dk_rtheta[1]);
    }
    free(ctx->priv);

    free(ctx);
    *pctx = NULL;
}

static const void
(*constraint_funcs[TD_CONSTRAINT_EQ_NB])(TDConstraintEvalContext *ctx,
                                         double *vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                                         double *out) = {
    [TD_CONSTRAINT_EQ_HAM]   = constraint_eval_ham_confflat,
    [TD_CONSTRAINT_EQ_MOM_0] = constraint_eval_mom_r_confflat,
    [TD_CONSTRAINT_EQ_MOM_1] = constraint_eval_mom_t_confflat,
};

int tdi_constraint_eval(TDConstraintEvalContext *ctx, enum TDConstraintEq eq,
                        const double *vars[TD_CONSTRAINT_VAR_NB][PSSOLVE_DIFF_ORDER_NB],
                        double *out)
{
    constraint_funcs[eq](ctx, vars, out);
}