/*
 * Multigrid 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_H
#define MG2D_H

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

#include <mpi.h>

#include "mg2d_boundary.h"
#include "mg2d_constants.h"

typedef struct MG2DInternal MG2DInternal;

/**
 * Flags for diff coeffs boundary specifications.
 */
enum MG2DDCFlag {
    /**
     * The coefficients have a discontinuity at the boundary, i.e.
     * C(x, y) = C0(x, y) + D(x, y),
     * where C is the resulting value used for the solve, C0 is continuous
     * everywhere, and D is zero everywhere except along the boundary. D is
     * assumed to be continuous along the boundary.
     *
     * The values of D should be stored in MG2DDCBoundarySpec.val
     */
    MG2D_DC_FLAG_DISCONT = (1 << 0),
    /**
     * The coefficients have a first order pole at the boundary, i.e.
     * C(x, y) · (x_i - x_i^0)                  (1)
     * is continuous, where C is the actual value of the coefficient, x_i is the
     * coordinate which is constant along the boundary and x_i^0 is the value of
     * that coordinate along the boundary.
     *
     * The value of (1) along the boundary should be stored in
     * MG2DDCBoundarySpec.val.
     */
    MG2D_DC_FLAG_POLE    = (1 << 1),
};

/**
 * Specification of the behaviour of differential equation coefficients on a
 * boundary.
 */
typedef struct MG2DDCBoundarySpec {
    /**
     * A combination of MG2DDCFlag
     */
    int flags;
    /**
     * Value determined by flags. Need not be filled if the flags are zero.
     */
    double *val;
} MG2DDCBoundarySpec;

/**
 * Specification of a single coefficient of the PDE.
 */
typedef struct MG2DDiffCoeffs {
    /**
     * Values of the coefficient at the grid points. Values corresponding to
     * successive x locations are contiguous in memory.
     */
    double         *data;
    /**
     * Number of elements between successive y locations in data.
     */
    ptrdiff_t       stride;

    /**
     * Behaviour of the coefficient on the boundaries.
     */
    MG2DDCBoundarySpec boundaries[4];
} MG2DDiffCoeffs;

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

    /**
     * A callback that will be used to print diagnostic messages.
     *
     * The default implementation prints to stderr.
     *
     * May be set to NULL to disable logging.
     */
    void (*log_callback)(const struct MG2DContext *ctx, int level,
                         const char *fmt, va_list);

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

    /**
     * Number of grid points in each direction, set by mg2d_solver_alloc().
     * Read-only for the caller.
     */
    size_t domain_size;

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

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

    /**
     * Boundary specification, indexed by MG2DBoundaryLoc.
     * To be filled by the caller.
     */
    MG2DBoundary *boundaries[4];

    /**
     * Maximum number of solver iterations.
     */
    unsigned int maxiter;

    /**
     * Value of the residual at which to consider the equation solved.
     */
    double tol;

    /**
     * Number of cycles on each subgrid.
     */
    unsigned int nb_cycles;

    /**
     * Number of relaxation steps before and after each coarser-grid solve.
     */
    unsigned int nb_relax_pre;
    unsigned int nb_relax_post;

    /**
     * Values of the unknown function.
     *
     * Allocated and initialized to zero by the solver in mg2d_solver_alloc(),
     * owned by the solver.
     * May be filled by the caller before solving to set the initial guess.
     * Afterwards updated by mg2d_solve() if it returns 0,
     * MG2D_ERR_MAXITER_REACHED or MG2D_ERR_DIVERGE. If mg2d_solve() returnes
     * another error code, the contents of u are unspecified.
     */
    double   *u;
    /**
     * Distance between neighbouring gridpoints along coord1.
     */
    ptrdiff_t u_stride;

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

    /**
     * Coefficients C_{*} that define the differential equation.
     *
     * Allocated by the solver in mg2d_solver_alloc(), owned by the solver.
     * Must be filled by the caller before solving.
     */
    MG2DDiffCoeffs *diff_coeffs[MG2D_DIFF_COEFF_NB];

    /**
     * Number of threads to use for processing.
     *
     * Set to 0 to autodetect. Defaults to 1.
     */
    unsigned int nb_threads;

    /**
     * Time-stepping factor to use for relaxation.
     */
    double cfl_factor;

    /**
     * Maximum level of messages printed by the default logging callback. Has no
     * effect when log_callback is overridden.
     */
    enum MG2DLogLevel log_level;

    /**
     * Maximum number of refinement levels. May only be modified by the caller
     * before the first call to mg2d_solve().
     */
    unsigned int max_levels;

    /**
     * Maximum size (along one dimensions) of a refinement level on which an
     * exact solve is performed.
     */
    size_t max_exact_size;

    /**
     * Indices of the lower left corner of the local domain in the full
     * computational grid, for distributed solvers. Always { 0, 0 } for
     * single-component solvers.
     *
     * Set by mg2d_solver_alloc[_mpi]().
     */
    ptrdiff_t local_start[2];

    /**
     * Number of points in the local component in each direction. Equal to
     * { domain_size, domain_size } for single-component solvers.
     *
     * Set by mg2d_solver_alloc[_mpi]().
     */
    size_t    local_size[2];

    /**
     * Maximum of the absolute value of the residual.
     *
     * Set by mg2d_solve() if it returns 0 or MG2D_ERR_MAXITER_REACHED.
     */
    double residual_max;

    /**
     * Unused, do not set.
     */
    int adaptive_step;
} MG2DContext;

/**
 * Allocate the solver for the given domain size.
 *
 * @return The solver context on success, NULL on failure.
 */
MG2DContext *mg2d_solver_alloc(size_t domain_size);

/**
 * Allocate a solver component in a multi-component MPI-based solve.
 *
 * @param comm The MPI communicator used to communicate with the other
 *             components.
 * @param local_start Indices of this component's lower left corner in the full
 *                    computational domain.
 * @param local_size Size of this component in each direction.
 *
 * @return The solver context on success, NULL on failure.
 */
MG2DContext *mg2d_solver_alloc_mpi(MPI_Comm comm, const size_t local_start[2],
                                   const size_t local_size[2]);
/**
 * Solve the equation, after all the required fields have been filled by the
 * caller.
 *
 * This function may be called more than once.
 *
 * @return
 *   - 0 on success
 *   - MG2D_ERR_MAXITER_REACHED if desired tolerance was not reached after
 *     maximum allowed number of iterations were performed. The final
 *     value of the solution is still exported in ctx->u, ctx->residual_max is
 *     set to the final residual norm.
 *   - MG2D_ERR_DIVERGE if the iteration process has diverged.
 *   - another negative error code on other types of failure
 */
int mg2d_solve(MG2DContext *ctx);
/**
 * Free the solver and write NULL to the provided pointer.
 */
void mg2d_solver_free(MG2DContext **ctx);

/**
 * Print the solver timing statistics using the logging callback.
 *
 * @param prefix the string prepended to each output line, may be NULL to mean
 *               none
 */
void mg2d_print_stats(MG2DContext *ctx, const char *prefix);

/**
 * @return maximum supported value of fd_stencil.
 */
unsigned int mg2d_max_fd_stencil(void);

int mg2d_init_guess(MG2DContext *ctx, const double *src,
                    ptrdiff_t src_stride,
                    const ptrdiff_t src_start[2],
                    const size_t src_size[2],
                    const double src_step[2]);

#endif /* MG2D_H */