/*
 * Pseudospectral solver for 2nd order 2D linear PDE systems
 * Copyright (C) 2014-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_PSSOLVE_H
#define TEUKOLSKY_PSSOLVE_H

/**
 * The problem being solved is a sequence of N linear partial differential
 * equations
 *
 *    /                                                                   \
 *  ∑ |   ∑ C_{ab}^{ik} ∂_a　∂_b u_k  + ∑ C_{a}^{ik} ∂_a u_k + C^{ik} u_k | = S^i
 *  k \  a,b                            a                                 /
 *
 *  where
 *      * i numbers the equations and runs from 0 to N-1
 *      * k numbers the unknown functions and also runs from 0 to N-1
 *      * a and b identify spatial directions and run from 0 to 1
 *      * u_k = u_k(x_a) is the k-th unknown function
 *      * C_{ab}^{ik}, C_{a}^{ik} and C^{ik} are the coefficients in front of
 *        the corresponding derivative of k-th unknown function in the i-th
 *        equation
 *      * S^i is the right-hand side of the i-th equation
 * C_{*}^{ik} and S^i are all (known) functions of space and define the
 * equations to be solved.
 */

#include "config.h"

#if HAVE_OPENCL
#include <cl.h>
#else
typedef void* cl_context;
typedef void* cl_command_queue;
#endif

#include <stdint.h>

#include "basis.h"
#include "log.h"
#include "threadpool.h"

enum PSSolveDiffOrder {
    PSSOLVE_DIFF_ORDER_00,
    PSSOLVE_DIFF_ORDER_10,
    PSSOLVE_DIFF_ORDER_01,
    PSSOLVE_DIFF_ORDER_11,
    PSSOLVE_DIFF_ORDER_20,
    PSSOLVE_DIFF_ORDER_02,
    PSSOLVE_DIFF_ORDER_NB,
};

typedef struct PSSolvePriv PSSolvePriv;

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

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

    /**
     * Number of equations/unknown functions in the set.
     * Set by tdi_pssolve_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_pssolve_context_alloc(), must be filled by
     * by the caller before tdi_pssolve_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_pssolve_context_alloc(), must be filled by the caller
     * before tdi_pssolve_context_init().
     */
    unsigned int (*solve_order)[2];

    /**
     * Locations of the collocation points. The equation coefficients passed to
     * tdi_pssolve_solve() 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_pssolve_context_init().
     */
    double *(*colloc_grid)[2];

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

    cl_context       ocl_ctx;
    cl_command_queue ocl_queue;

    uint64_t lu_solves_count;
    uint64_t lu_solves_time;

    uint64_t cg_solve_count;
    uint64_t cg_iter_count;
    uint64_t cg_time_total;

    uint64_t construct_matrix_count;
    uint64_t construct_matrix_time;
} PSSolveContext;

/**
 * Allocate a new solver.
 *
 * @param ctx The newly allocated solver context will be written here.
 * @param nb_equations number of equations to solve (equal to the number of
 *                     unknown functions to solve for)
 *
 * @return 0 on success a negative error code on failure.
 */
int tdi_pssolve_context_alloc(PSSolveContext **ctx, unsigned int nb_equations);

/**
 * Initialize the solver for use after all the context options have been set.
 * This function must be called exactly once before any calls to
 * tdi_pssolve_solve().
 *
 * @return 0 on success, a negative error code on failure.
 */
int tdi_pssolve_context_init(PSSolveContext *ctx);

/**
 * Free the solver and write NULL to the supplied pointer.
 */
void tdi_pssolve_context_free(PSSolveContext **ctx);

/**
 * Solve a PDE. This function may be called multiple times in succession to
 * solve multiple related PDEs (it will be efficient if the equation
 * coefficients do not change too much).
 *
 * @param ctx the solver context
 * @param eq_coeffs the equation coefficients at the collocation points.
 *                  eq_coeffs[i][j][k] is the array of coefficients for the k-th
 *                  derivative (as per enum PSSolveDiffOrder) of the j-th
 *                  unknown function in the i-th equation.
 * @param rhs the right-hand side of the equation at the collocation points.
 * @param coeffs the spectral coefficients of the solution will be written here.
 *
 * @return 0 on success, a negative error code on failure. The contents of
 *         coeffs are undefined on failure.
 */
int tdi_pssolve_solve(PSSolveContext *ctx,
                      const double *(**eq_coeffs)[PSSOLVE_DIFF_ORDER_NB],
                      const double *rhs, double *coeffs);

int tdi_pssolve_diff_order(enum PSSolveDiffOrder order, unsigned int dir);

#endif /* TEUKOLSKY_PSSOLVE_H */