/*
 * 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 <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "basis.h"
#include "common.h"
#include "log.h"
#include "nlsolve.h"
#include "pssolve.h"

#define N_X 8
#define N_Z 4
#define FUNC_X 4
#define FUNC_Z 2
#define TOL 5e-15
#define ITER 10
static int
scalar_linear_func(void *opaque, unsigned int eq_idx,
                   const unsigned int *colloc_grid_order,
                   const double * const *colloc_grid,
                   const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
                   double *dst)
{
    BasisSetContext *basis = opaque;

    for (int j = 0; j < colloc_grid_order[1]; j++) {
        const double z = colloc_grid[1][j];

        for (int i = 0; i < colloc_grid_order[0]; i++) {
            const double x = colloc_grid[0][i];

            const int idx = j * colloc_grid_order[0] + i;

            const double d2x = vars[0][PSSOLVE_DIFF_ORDER_20][idx];
            const double d2z = vars[0][PSSOLVE_DIFF_ORDER_02][idx];
            const double val = vars[0][PSSOLVE_DIFF_ORDER_00][idx];

            const double d2x_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, x, FUNC_X) *
                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);
            const double d2z_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
                                  tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, z, FUNC_Z);
            const double val_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);

            dst[idx] = (d2x + d2z + val) - (d2x_ex + d2z_ex + val_ex);
        }
    }

    return 0;
}

static int scalar_linear(void)
{
    BasisSetContext *basis = NULL;
    NLSolveContext    *ctx = NULL;

    double coeffs[N_X * N_Z];
    int ret;

    ret = tdi_basis_init(&basis, BASIS_FAMILY_SB_EVEN, 1.0);
    if (ret < 0)
        goto finish;

    ret = tdi_nlsolve_context_alloc(&ctx, 1);
    if (ret < 0)
        return ret;

    ctx->basis[0][0] = basis;
    ctx->basis[0][1] = basis;

    ctx->solve_order[0][0] = N_X;
    ctx->solve_order[0][1] = N_Z;

    ctx->maxiter = ITER;
    ctx->atol    = TOL;

    ctx->logger.log = tdi_log_default_callback;

    ret = tdi_nlsolve_context_init(ctx);
    if (ret < 0)
        goto finish;

    // seed the initial guess with garbage, but not too large, numbers
    for (int i = 0; i < N_X; i++)
        for (int j = 0; j < N_Z; j++)
            coeffs[j * N_X + i] = 0.2 * j + 0.1 * sin(138.0 * i);

    ret = tdi_nlsolve_solve(ctx, scalar_linear_func, NULL, basis, coeffs);
    if (ret < 0)
        goto finish;

    for (int j = 0; j < N_Z; j++)
        for (int i = 0; i < N_X; i++) {
            const int idx = j * N_X + i;
            const double val = (i == FUNC_X && j == FUNC_Z) ? 1.0 : 0.0;

            if (fabs(coeffs[idx] - val) > TOL) {
                fprintf(stderr, "unexpected value %g at %d/%d\n",
                        coeffs[idx], i, j);
                ret = -1;
                goto finish;
            }
        }

finish:
    tdi_nlsolve_context_free(&ctx);
    tdi_basis_free(&basis);

    return ret;
}
#undef N_X
#undef N_Z
#undef FUNC_X
#undef FUNC_Z
#undef TOL
#undef ITER

#define N_X 16
#define N_Z 16
#define L   0.1
#define FUNC_X 4
#define FUNC_Z 2
#define ATOL 1e-14
#define TOL 1e-14
#define ITER 10
static int
scalar_quadratic_func(void *opaque, unsigned int eq_idx,
                      const unsigned int *colloc_grid_order,
                      const double * const *colloc_grid,
                      const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
                      double *dst)
{
    BasisSetContext *basis = opaque;

    for (int j = 0; j < colloc_grid_order[1]; j++) {
        const double z = colloc_grid[1][j];

        for (int i = 0; i < colloc_grid_order[0]; i++) {
            const double x = colloc_grid[0][i];

            const int idx = j * colloc_grid_order[0] + i;

            const double d2x = vars[0][PSSOLVE_DIFF_ORDER_20][idx];
            const double d2z = vars[0][PSSOLVE_DIFF_ORDER_02][idx];
            const double val = vars[0][PSSOLVE_DIFF_ORDER_00][idx];

            const double d2x_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, x, FUNC_X) *
                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);
            const double d2z_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
                                  tdi_basis_eval(basis, BS_EVAL_TYPE_DIFF2, z, FUNC_Z);
            const double val_ex = tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, x, FUNC_X) *
                                  tdi_basis_eval(basis, BS_EVAL_TYPE_VALUE, z, FUNC_Z);

            dst[idx] = (d2x - d2x_ex) + (d2z - d2z_ex) + (SQR(val) - SQR(val_ex));
        }
    }

    return 0;
}

static int
scalar_quadratic_jac(void *opaque, unsigned int eq_idx,
                     unsigned int var_idx, enum PSSolveDiffOrder diff_order,
                     const unsigned int *colloc_grid_order,
                     const double * const *colloc_grid,
                     const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
                     double *dst)
{
    BasisSetContext *basis = opaque;

    for (int j = 0; j < colloc_grid_order[1]; j++) {
        const double z = colloc_grid[1][j];

        for (int i = 0; i < colloc_grid_order[0]; i++) {
            const double x = colloc_grid[0][i];

            const int idx = j * colloc_grid_order[0] + i;

            const double val = vars[0][PSSOLVE_DIFF_ORDER_00][idx];

            if (diff_order == PSSOLVE_DIFF_ORDER_20 ||
                diff_order == PSSOLVE_DIFF_ORDER_02) {
                dst[idx] = 1.0;
            } else if (diff_order == PSSOLVE_DIFF_ORDER_00) {
                dst[idx] = 2.0 * val;
            } else {
                dst[idx] = 0.0;
            }
        }
    }

    return 0;
}

static int scalar_quadratic(void)
{
    BasisSetContext *basis = NULL;
    NLSolveContext    *ctx = NULL;

    double coeffs[N_X * N_Z];
    int ret;

    ret = tdi_basis_init(&basis, BASIS_FAMILY_SB_EVEN, L);
    if (ret < 0)
        goto finish;

    ret = tdi_nlsolve_context_alloc(&ctx, 1);
    if (ret < 0)
        return ret;

    ctx->basis[0][0] = basis;
    ctx->basis[0][1] = basis;

    ctx->solve_order[0][0] = N_X;
    ctx->solve_order[0][1] = N_Z;

    ctx->maxiter = ITER;
    ctx->atol    = ATOL;

    ctx->logger.log = tdi_log_default_callback;

    ret = tdi_nlsolve_context_init(ctx);
    if (ret < 0)
        goto finish;

    memset(coeffs, 0, sizeof(coeffs));

    ret = tdi_nlsolve_solve(ctx, scalar_quadratic_func, scalar_quadratic_jac, basis, coeffs);
    if (ret < 0)
        goto finish;

    for (int j = 0; j < N_Z; j++)
        for (int i = 0; i < N_X; i++) {
            const int idx = j * N_X + i;
            const double val = (i == FUNC_X && j == FUNC_Z) ? 1.0 : 0.0;

            if (fabs(coeffs[idx] - val) > TOL) {
                fprintf(stderr, "unexpected value %g at %d/%d\n",
                        coeffs[idx], i, j);
                ret = -1;
                goto finish;
            }
        }

finish:
    tdi_nlsolve_context_free(&ctx);
    tdi_basis_free(&basis);

    return ret;
}
#undef N_X
#undef N_Z
#undef L
#undef FUNC_X
#undef FUNC_Z
#undef ATOL
#undef TOL
#undef ITER

#define N_R   16
#define N_PHI 8
#define L   0.5
#define FUNC0_R   4
#define FUNC0_PHI 2
#define FUNC1_R   2
#define FUNC1_PHI 3
#define ATOL 1e-14
#define TOL 1e-14
#define ITER 128
static int
vector_quadratic_func(void *opaque, unsigned int eq_idx,
                      const unsigned int *colloc_grid_order,
                      const double * const *colloc_grid,
                      const double * const (*vars)[PSSOLVE_DIFF_ORDER_NB],
                      double *dst)
{
    BasisSetContext **basis = opaque;
    unsigned int exact_r   = eq_idx == 0 ? FUNC0_R : FUNC1_R;
    unsigned int exact_phi = eq_idx == 0 ? FUNC0_PHI : FUNC1_PHI;

    for (int j = 0; j < colloc_grid_order[1]; j++) {
        const double phi = colloc_grid[1][j];

        for (int i = 0; i < colloc_grid_order[0]; i++) {
            const double r = colloc_grid[0][i];

            const int idx = j * colloc_grid_order[0] + i;

            const double  dr   = vars[eq_idx][PSSOLVE_DIFF_ORDER_10][idx];
            const double d2r   = vars[eq_idx][PSSOLVE_DIFF_ORDER_20][idx];
            const double d2phi = vars[eq_idx][PSSOLVE_DIFF_ORDER_02][idx];
            const double val0 = vars[0][PSSOLVE_DIFF_ORDER_00][idx];
            const double val1 = vars[1][PSSOLVE_DIFF_ORDER_00][idx];

            const double  dr_ex  = tdi_basis_eval(basis[0], BS_EVAL_TYPE_DIFF1, r,    exact_r) *
                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi,  exact_phi);
            const double d2r_ex  = tdi_basis_eval(basis[0], BS_EVAL_TYPE_DIFF2, r,    exact_r) *
                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi,  exact_phi);
            const double d2phi_ex = tdi_basis_eval(basis[0], BS_EVAL_TYPE_VALUE, r,   exact_r) *
                                    tdi_basis_eval(basis[1], BS_EVAL_TYPE_DIFF2, phi, exact_phi);
            const double val0_ex = tdi_basis_eval(basis[0], BS_EVAL_TYPE_VALUE, r,   FUNC0_R) *
                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi, FUNC0_PHI);
            const double val1_ex = tdi_basis_eval(basis[0], BS_EVAL_TYPE_VALUE, r,   FUNC1_R) *
                                   tdi_basis_eval(basis[1], BS_EVAL_TYPE_VALUE, phi, FUNC1_PHI);

            dst[idx] = (d2r - d2r_ex) + (dr - dr_ex) / r + (d2phi - d2phi_ex) / SQR(r) + (val0 * val1 - val0_ex * val1_ex);
        }
    }

    return 0;
}

static int vector_quadratic(void)
{
    BasisSetContext *basis[2] = { NULL };
    NLSolveContext    *ctx = NULL;

    double coeffs[2 * N_R * N_PHI];
    int ret;

    ret = tdi_basis_init(&basis[0], BASIS_FAMILY_SB_EVEN,  L);
    ret = tdi_basis_init(&basis[1], BASIS_FAMILY_COS_EVEN, L);
    if (ret < 0)
        goto finish;

    ret = tdi_nlsolve_context_alloc(&ctx, 2);
    if (ret < 0)
        return ret;

    ctx->basis[0][0] = basis[0];
    ctx->basis[1][0] = basis[0];
    ctx->basis[0][1] = basis[1];
    ctx->basis[1][1] = basis[1];

    ctx->solve_order[0][0] = N_R;
    ctx->solve_order[0][1] = N_PHI;
    ctx->solve_order[1][0] = N_R;
    ctx->solve_order[1][1] = N_PHI;

    ctx->maxiter = ITER;
    ctx->atol    = ATOL;

    ctx->logger.log = tdi_log_default_callback;

    ret = tdi_nlsolve_context_init(ctx);
    if (ret < 0)
        goto finish;

    memset(coeffs, 0, sizeof(coeffs));

    ret = tdi_nlsolve_solve(ctx, vector_quadratic_func, NULL, basis, coeffs);
    if (ret < 0)
        goto finish;

    for (int j = 0; j < N_PHI; j++)
        for (int i = 0; i < N_R; i++) {
            const int idx = j * N_R + i;
            const double val = (i == FUNC0_R && j == FUNC0_PHI) ? 1.0 : 0.0;

            if (fabs(coeffs[idx] - val) > TOL) {
                fprintf(stderr, "unexpected value %g at %d/%d\n",
                        coeffs[idx], i, j);
                ret = -1;
                goto finish;
            }
        }
    for (int j = 0; j < N_PHI; j++)
        for (int i = 0; i < N_R; i++) {
            const int idx = N_PHI * N_R + j * N_R + i;
            const double val = (i == FUNC1_R && j == FUNC1_PHI) ? 1.0 : 0.0;

            if (fabs(coeffs[idx] - val) > TOL) {
                fprintf(stderr, "unexpected value %g at %d/%d\n",
                        coeffs[idx], i, j);
                ret = -1;
                goto finish;
            }
        }

finish:
    tdi_nlsolve_context_free(&ctx);
    tdi_basis_free(&basis[0]);
    tdi_basis_free(&basis[1]);

    return ret;
}
#undef N_R
#undef N_PHI
#undef L
#undef FUNC0_R
#undef FUNC0_PHI
#undef FUNC1_R
#undef FUNC1_PHI
#undef ATOL
#undef TOL
#undef ITER

static const struct {
    const char *name;
    int (*test)(void);
} tests[] = {
    { "scalar_linear",    scalar_linear },
    { "scalar_quadratic", scalar_quadratic },
    { "vector_quadratic", vector_quadratic },
};

int main(void)
{
    int ret = 0;

    for (int i = 0; i < ARRAY_ELEMS(tests); i++) {
        fprintf(stderr, "executing test '%s'\n", tests[i].name);
        ret = tests[i].test();
        if (ret < 0) {
            fprintf(stderr, "test '%s' failed\n", tests[i].name);
            return -1;
        }
        fprintf(stderr, "test '%s' succeeded\n", tests[i].name);
    }

    return 0;
}