/*
 * Newton-Kantorovich iterative solver for nonlinear PDE systems
 * Copyright (C) 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/>.
 */

#ifndef TEUKOLSKY_DATA_NLSOLVE_H
#define TEUKOLSKY_DATA_NLSOLVE_H

#include <threadpool.h>

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

typedef struct NLSolvePriv NLSolvePriv;

typedef struct NLSolveContext {
    /**
     * Solver private data, not to be touched by the caller.
     */
    NLSolvePriv *priv;

    /**
     * The logging context.
     * Set by the caller before tdi_nlsolve_context_init().
     */
    TDLogger logger;

    /**
     * The thread pool used for multithreaded execution. May be set by the
     * caller before tdi_nlsolve_context_init(), otherwise a single thread will
     * be used.
     */
    TPContext *tp;

    /**
     * Number of equations/unknown functions in the set.
     * Set by tdi_nlsolve_context_alloc().
     */
    unsigned int nb_equations;

    /**
     * The basis sets.
     *
     * basis[i][j] is the basis set used for i-th variable in j-th direction.
     *
     * The array is allocated by tdi_nlsolve_context_alloc(), must be filled
     * by the caller before tdi_nlsolve_context_init().
     */
    const BasisSetContext *(*basis)[2];

    /**
     * Order of the solver.
     *
     * solve_order[i][j] is the order of the solver (i.e. the number of the
     * basis functions used) for i-th variable in j-th direction.
     *
     * Allocated by tdi_nlsolve_context_alloc(), must be filled by the caller
     * before tdi_nlsolve_context_init().
     */
    unsigned int (*solve_order)[2];

    /**
     * Locations of the collocation points. The equation coefficients in NLEqCallback and
     * NLEqJacobianCallback should be evaluated at those grid positions.
     *
     * colloc_grid[i][j] is an array of length solve_order[i][j] and contains
     * the collocation points for the i-th equation in the j-th direction.
     *
     * Set by the solver after tdi_nlsolve_context_init().
     */
    const double * const (*colloc_grid)[2];

    // solver parameters
    unsigned int maxiter;
    double atol;
} NLSolveContext;

typedef int (*NLEqCallback)(
    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);

typedef int (*NLEqJacobianCallback)(
    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);

int tdi_nlsolve_context_alloc(NLSolveContext **ctx, unsigned int nb_equations);

int tdi_nlsolve_context_init(NLSolveContext *ctx);

void tdi_nlsolve_context_free(NLSolveContext **ctx);

int tdi_nlsolve_solve(NLSolveContext *ctx, NLEqCallback eq_eval,
                      NLEqJacobianCallback eq_jac_eval,
                      void *opaque, double *coeffs, int fast_abort);

void tdi_nlsolve_print_stats(NLSolveContext *ctx);

#endif // TEUKOLSKY_DATA_NLSOLVE_H