/*
 * Relaxation solver for a 2nd order 2D linear PDE.
 * Copyright 2018 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 MG2D_ELL_RELAX_H
#define MG2D_ELL_RELAX_H

/**
 * The problem being solved is a linear partial differential
 * equation
 *
 *   ∑ C_{ab} ∂_a　∂_b u  + ∑ C_{a} ∂_a u + C u = rhs
 *  a,b                     a
 *
 *  where
 *      * a and b identify spatial directions and run from 0 to 1
 *      * u = u(x_a) is the unknown function
 *      * C_{ab}, C_{a} and C are the coefficients in front of
 *        the corresponding derivative of unknown function
 *      * rhs is the right-hand side of the equation
 * C_{*} and rhs are all (known) functions of space and define the
 * equation to be solved.
 */

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

#include "log.h"

/**
 * Type of the boundary condition on a given boundary.
 */
enum EllRelaxBCType {
    /**
     * The value of the unknown function is fixed on the boundary.
     */
    ELL_RELAX_BC_TYPE_FIXVAL,
    /**
     * The normal derivative of the unkown function is fixed on the boundary.
     *
     * TODO: the only supported value of the derivative is currently zero, i.e.
     * periodic boundary conditions.
     */
    ELL_RELAX_BC_TYPE_FIXDIFF,
};

/**
 * Location of the boundary.
 */
enum EllRelaxBoundaryLoc {
    /**
     * coord0 lower
     */
    ELL_RELAX_BOUNDARY_0L,
    /**
     * coord0 upper
     */
    ELL_RELAX_BOUNDARY_0U,
    /**
     * coord1 lower
     */
    ELL_RELAX_BOUNDARY_1L,
    /**
     * coord1 upper
     */
    ELL_RELAX_BOUNDARY_1U,
};

/**
 * Derivative specifier.
 */
enum EllRelaxDiffCoeff {
    /**
     * Zeroth derivative, i.e. function value.
     */
    ELL_RELAX_DIFF_COEFF_00,
    /**
     * First derivative wrt coord1.
     */
    ELL_RELAX_DIFF_COEFF_01,
    /**
     * First derivative wrt coord0.
     */
    ELL_RELAX_DIFF_COEFF_10,
    /**
     * Second derivative, once wrt coord0 and once wrt coord1
     */
    ELL_RELAX_DIFF_COEFF_11,
    /**
     * Second derivative wrt coord1
     */
    ELL_RELAX_DIFF_COEFF_02,
    /**
     * Second derivative wrt coord0
     */
    ELL_RELAX_DIFF_COEFF_20,
    /**
     * Total number of allowed derivatives.
     */
    ELL_RELAX_DIFF_COEFF_NB,
};

typedef struct EllRelaxInternal EllRelaxInternal;

/**
 * Boundary condition definition.
 */
typedef struct EllRelaxBoundary {
    /**
     * Type of the boundary condition.
     */
    enum EllRelaxBCType type;
    /**
     * For type = ELL_RELAX_BC_TYPE_FIXVAL:
     *     Values of the unknown function on the boundary.
     *     The number of boundary layers is equal to fd_stencil.
     *     The first boundary layer has the number of points equal to the
     *     corresponding domain_size. Each subsequent boundary layer has one
     *     more boundary point at each end of the domain.
     *
     * Ignored otherwise.
     */
    double *val;
    /**
     * Number of elements between rows in val. I.e. if val[0] is the first
     * boundary point, then val[val_stride - 1] is the first boundary point in
     * the second row and so on.
     */
    size_t  val_stride;
} EllRelaxBoundary;

typedef struct EllRelaxContext {
    /**
     * Solver private data, not to be accessed in any way by the caller.
     */
    EllRelaxInternal *priv;

    /**
     * The logging context, to be filled by the caller before mg2di_ell_relax_init().
     */
    MG2DLogger logger;

    /**
     * The thread pool used for execution. May be set by the caller before
     * mg2di_ell_relax_init().
     */
    TPContext *tp;

    /**
     * Flags indicating supported CPU features.
     */
    int cpuflags;

    /**
     * Size of the solver grid, set by mg2di_ell_relax_alloc().
     * Read-only for the caller.
     */
    size_t domain_size[2];

    /**
     * Distance between the neighbouring grid points.
     * Must be set by the caller before mg2di_ell_relax_init().
     */
    double step[2];

    /**
     * Order of the finite difference operators used for approximating derivatives.
     * Must be set by the caller before mg2di_ell_relax_init().
     */
    size_t fd_stencil;

    /**
     * Boundary specification, indexed by EllRelaxBoundaryLoc.
     * To be filled by the caller before mg2di_ell_relax_init().
     */
    EllRelaxBoundary boundaries[4];

    /**
     * The time stepping factor in relaxation.
     */
    double relax_factor;

    /**
     * Values of the unknown function.
     *
     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
     * Must be filled by the caller before mg2di_ell_relax_init() to set the
     * initial guess.
     * Afterwards updated in mg2di_ell_relax_step().
     */
    double   *u;
    /**
     * Distance between neighbouring gridpoints along coord1.
     */
    ptrdiff_t u_stride;

    /**
     * Values of the right-hand side.
     *
     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
     * Must be filled by the caller before mg2di_ell_relax_init().
     */
    double    *rhs;
    /**
     * Distance between neighbouring gridpoints along coord1.
     */
    ptrdiff_t  rhs_stride;

    /**
     * Values of the right-hand side.
     *
     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
     * Read-only for the caller. Initialized after mg2di_ell_relax_init(),
     * afterwards updated in mg2di_ell_relax_step().
     */
    double   *residual;
    /**
     * Distance between neighbouring gridpoints along coord1.
     */
    ptrdiff_t residual_stride;

    /**
     * Coefficients C_{*} that define the differential equation.
     *
     * Allocated by the solver in mg2di_ell_relax_alloc(), owned by the solver.
     * Must be filled by the caller before mg2di_ell_relax_init().
     */
    double   *diff_coeffs[ELL_RELAX_DIFF_COEFF_NB];
    /**
     * Distance between neighbouring gridpoints along coord1.
     */
    ptrdiff_t diff_coeffs_stride;

    /* timings */
    int64_t count_correct;
    int64_t time_correct;

    int64_t time_boundaries;
    int64_t count_boundaries;
    int64_t time_res_calc;
    int64_t count_res;
} EllRelaxContext;

/**
 * Allocate the solver for the given domain size.
 *
 * @param domain_size number of grid points in each direction.
 *
 * @return The solver context on success, NULL on failure.
 */
EllRelaxContext *mg2di_ell_relax_alloc(size_t domain_size[2]);
/**
 * Initialize the solver for use, after all the required fields are filled by
 * the caller.
 *
 * This function may be called multiple times to re-initialize the solver after
 * certain parameters (e.g. the right-hand side or the equation coefficients)
 * change.
 *
 * @return 0 on success, a negative error code on failure.
 */
int              mg2di_ell_relax_init(EllRelaxContext *ctx);
/**
 * Free the solver and write NULL to the provided pointer.
 */
void             mg2di_ell_relax_free(EllRelaxContext **ctx);

/**
 * Perform a single relaxation step.
 *
 * @return 0 on success, a negative error code on failure.
 */
int mg2di_ell_relax_step(EllRelaxContext *ctx);

#endif /* MG2D_ELL_RELAX_H */