/*
 * Copyright 2014-2015 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 BRILL_DATA_H
#define BRILL_DATA_H

#include <stdarg.h>
#include <stddef.h>
#include <stdint.h>

/**
 * API usage:
 *
 * First, allocate the solver context with bd_context_alloc(). All interaction
 * with the solver is done through this context.
 *
 * Fill any fields in the context that are described as settable by the caller.
 * Call bd_solve() to solve the equation determined by the option values.
 *
 * The Brill data is defined by the line element in cylindrical coordinates
 * { ρ, z, φ } as
 * dl^2 = ψ^4 (exp(-2 q(ρ, z)) (dρ^2 + dz^2) + dφ^2)
 * The conformal factor ψ and its derivatives may be evaluated with
 * bd_eval_psi() and the q function in the exponential can be evaluated with
 * bd_eval_q(). The full metric can be evaluated with bd_eval_metric(). *
 *
 * Finally, free the solver context with bd_context_free().
 */

enum BDQFuncType {
    /**
     * q(ρ, z) = A ρ^2 exp(-(ρ-ρ_0)^2 - z^2)
     */
    BD_Q_FUNC_GUNDLACH,
    /**
     * q(ρ, z) = A ρ^2 / (1 + r^n); where r^2 = ρ^2 + z^2
     */
    BD_Q_FUNC_EPPLEY,
};

enum BDMetricComponent {
    BD_METRIC_COMPONENT_RHORHO,
    BD_METRIC_COMPONENT_RHOZ,
    BD_METRIC_COMPONENT_RHOPHI,
    BD_METRIC_COMPONENT_ZRHO,
    BD_METRIC_COMPONENT_ZZ,
    BD_METRIC_COMPONENT_ZPHI,
    BD_METRIC_COMPONENT_PHIRHO,
    BD_METRIC_COMPONENT_PHIZ,
    BD_METRIC_COMPONENT_PHIPHI,
};

typedef struct BDContext {
    /**
     *  Solver internals, not to be accessed by the caller
     */
    void *priv;

    /********************************
     *  options, set by the caller  *
     ********************************/

    /**
     * A callback that will be used to print diagnostic messages.
     *
     * Defaults to fprintf(stderr, ...)
     */
    void (*log_callback)(const struct BDContext *bd, int level,
                         const char *fmt, va_list);

    /**
     * Arbitrary user data, e.g. to be used by the log callback.
     */
    void *opaque;

    /********************************
     *  initial data parameters     *
     ********************************/

    /**
     * The choice of the q function in the exponent.
     * Defaults to BD_Q_FUNC_GUNDLACH.
     */
    enum BDQFuncType q_func_type;
    /**
     * The amplitude in the q function.
     * Defaults to 1.
     */
    double amplitude;
    /**
     * For BD_Q_FUNC_EPPLEY, the power in the denominator.
     * Must be >= 4.
     * Defaults to 5.
     */
    unsigned int eppley_n;

    /**
     * Center of the wave for BD_Q_FUNC_GUNDLACH.
     * Defaults to 0.
     */
    double rho0;

    /********************************
     *        solver options        *
     ********************************/

    /**
     * The number of basis functions in each direction.
     * [0] - radial, [1] - angular
     */
    unsigned int nb_coeffs[2];

    /**
     * The scaling factor used in the basis functions.
     * [0] - radial, [1] - angular
     */
    double basis_scale_factor[2];

    /**********
     * output *
     **********/
    /**
     * The coefficients of the solution.
     * The adjacent elements in one row correspond to the increasing order of
     * the first direction (radial), going between rows corresponds to the
     * second direction (angular).
     *
     * The data is owned and managed by this library and is read-only for
     * the caller.
     */
    const double *psi_minus1_coeffs;
    /**
     * The number of array elements between two rows.
     */
    unsigned int stride;

    /**
     * Number of threads to use for parallelization. Set to 0 to automatically
     * detect according to the number of CPU cores.
     */
    unsigned int nb_threads;
} BDContext;

/**
 * Allocate and initialize the solver.
 */
BDContext *bd_context_alloc(void);

/**
 * Free the solver and everything associated with it.
 */
void bd_context_free(BDContext **bd);

/**
 * Solve the equation for the conformal factor ψ and export the expansion coefficients in the
 * context.
 *
 * @return >= 0 on success, a negative error code on failure
 */
int bd_solve(BDContext *bd);

/**
 * Evaluate the conformal factor ψ or its derivatives at the specified rectangular grid.
 *
 * @param bd the solver context
 * @param rho the array of ρ coordinates.
 * @param nb_coords_rho the number of elements in rho.
 * @param z   the array of z coordinates.
 * @param nb_coords_z the number of elements in z.
 * @param diff_order the order of the derivatives of ψ to evaluate. The first element specifies the
 *                   derivative wrt ρ, the second wrt z. I.e. diff_order = { 0, 0 } evaluates ψ
 *                   itself, diff_order = { 0, 1 } evaluates ∂ψ/∂z etc.
 * @param psi the array into which the values of ψ will be written. ψ is evaluated on the grid
 *            formed by the outer product of the rho and z vectors.
 *            I.e. psi[j * psi_stride + i] = ψ(rho[i], z[j]). The length of psi must be
 *            at least psi_stride * nb_coords_z.
 * @param psi_stride the distance (in double-sized elements) in psi between two elements corresponding
 *                   to the same ρ but one step in z. Must be at least nb_coords_rho.
 *
 * @return >= 0 on success, a negative error code on failure.
 */
int bd_eval_psi(const BDContext *bd, const double *rho, int nb_coords_rho,
                const double *z, int nb_coords_z,
                const unsigned int diff_order[2],
                double *psi, unsigned int psi_stride);

/**
 * Evaluate the 3-metric γ_ij at the specified rectangular grid (in cylindrical
 * coordinates { ρ, z, φ }).
 *
 * @param bd the solver context
 * @param rho the array of ρ coordinates.
 * @param nb_coords_rho the number of elements in rho.
 * @param z   the array of z coordinates.
 * @param nb_coords_z the number of elements in z.
 * @param nb_comp number of needed components of the metric
 * @param comp a nb_comp-sized array specifying the components of the metric to evaluate
 * @param diff_order a nb_comp-sized array specifying the order of the derivatives of the
 *                   metric to evaluate. The first element specifies the derivative wrt ρ,
 *                   the second wrt z. I.e. diff_order[i] = { 0, 0 } evaluates the metric
 *                   itself, diff_order[i] = { 0, 1 } evaluates ∂γ/∂z etc.
 * @param out a nb_comp-sized array of pointers to the arrays into which the values of the
 *            metric will be written. Each requested component is evaluated on the grid
 *            formed by the outer product of the rho and z vectors. I.e.
 *            out[l][j * out_strides[l] + i] = γ_comp[l](rho[i], z[j]). The length of each
 *            array in out must be nb_coords_rho * nb_coords_z.
 * @param out_strides a nb_comp-sized array of distances (in double-sized elements), for each
 *                    array in out, between two elements corresponding to the same ρ but one
 *                    step in z. Each element in out_strides must be at least nb_coords_rho.
 *
 * @return >= 0 on success, a negative error code on failure.
 */
int bd_eval_metric(const BDContext *bd, const double *rho, int nb_coords_rho,
                   const double *z, int nb_coords_z,
                   int nb_comp, const enum BDMetricComponent *comp,
                   const unsigned int (*diff_order)[2],
                   double **out, unsigned int *out_strides);

/**
 * Evaluate the q function at the specified rectangular grid (in cylindrical
 * coordinates { ρ, z, φ }).
 *
 * @param bd the solver context
 * @param rho the array of ρ coordinates.
 * @param nb_coords_rho the number of elements in rho.
 * @param z   the array of z coordinates.
 * @param nb_coords_z the number of elements in z.
 * @param diff_order the order of the derivatives of the q function to evaluate. The first element
 *                   specifies the derivative wrt ρ, the second wrt z. I.e. diff_order = { 0, 0 }
 *                   evaluates the q function itself, diff_order = { 0, 1 } evaluates ∂q/∂z etc.
 * @param out the array into which the values of the q function will be written. The q function
 *            is evaluated on the grid formed by the outer product of the rho
 *            and z vectors. I.e. out[j * out_stride + i] = γ_comp[0]comp[1](rho[i], z[j]).
 *            The length of out must be nb_coords_rho * nb_coords_z.
 * @param out_stride the distance (in double-sized elements) in out between two elements corresponding
 *                   to the same ρ but one step in z. Must be at least nb_coords_rho.
 *
 * @return >= 0 on success, a negative error code on failure.
 */
int bd_eval_q(const BDContext *bd, const double *rho, int nb_coords_rho,
              const double *z, int nb_coords_z, const unsigned int diff_order[2],
              double *out, unsigned int out_stride);

#endif /* BRILL_DATA_H */