/*****************************************************************************
 *
 * MODULE:       Grass numerical math interface
 * AUTHOR(S):    Soeren Gebbert, Berlin (GER) Dec 2006
 *                 soerengebbert <at> googlemail <dot> com
 *
 * PURPOSE:      linear equation system solvers
 *                 part of the gmath library
 *
 * COPYRIGHT:    (C) 2010 by the GRASS Development Team
 *
 *               This program is free software under the GNU General Public
 *               License (>=v2). Read the file COPYING that comes with GRASS
 *               for details.
 *
 *****************************************************************************/

#include <assert.h>
#include <math.h>
#include <unistd.h>
#include <stdio.h>
#include <string.h>
#include <grass/gis.h>
#include <grass/gmath.h>
#include <grass/glocale.h>

static G_math_spvector **create_diag_precond_matrix(double **A,
                                                    G_math_spvector **Asp,
                                                    int rows, int prec);
static int solver_pcg(double **A, G_math_spvector **Asp, double *x, double *b,
                      int rows, int maxit, double err, int prec, int has_band,
                      int bandwidth);
static int solver_cg(double **A, G_math_spvector **Asp, double *x, double *b,
                     int rows, int maxit, double err, int has_band,
                     int bandwidth);
static int solver_bicgstab(double **A, G_math_spvector **Asp, double *x,
                           double *b, int rows, int maxit, double err);

/*!
 * \brief The iterative preconditioned conjugate gradients solver for symmetric
 * positive definite matrices
 *
 * This iterative solver works with symmetric positive definite  regular
 * quadratic matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param A (double **) -- the matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \param prec (int) -- the preconditioner which should be used 1,2 or 3
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_pcg(double **A, double *x, double *b, int rows, int maxit,
                      double err, int prec)
{

    return solver_pcg(A, NULL, x, b, rows, maxit, err, prec, 0, 0);
}

/*!
 * \brief The iterative preconditioned conjugate gradients solver for symmetric
 * positive definite band matrices
 *
 * WARNING: The preconditioning of symmetric band matrices is not implemented
 * yet
 *
 * This iterative solver works with symmetric positive definite band matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param A (double **) -- the positive definite band matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param bandwidth (int) -- bandwidth of matrix A
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \param prec (int) -- the preconditioner which should be used 1,2 or 3
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_pcg_sband(double **A, double *x, double *b, int rows,
                            int bandwidth, int maxit, double err, int prec)
{
    G_fatal_error("Preconditioning of band matrics is not implemented yet");
    return solver_pcg(A, NULL, x, b, rows, maxit, err, prec, 1, bandwidth);
}

/*!
 * \brief The iterative preconditioned conjugate gradients solver for sparse
 * symmetric positive definite matrices
 *
 * This iterative solver works with symmetric positive definite sparse matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param Asp (G_math_spvector **) -- the sparse matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \param prec (int) -- the preconditioner which should be used 1,2 or 3
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_sparse_pcg(G_math_spvector **Asp, double *x, double *b,
                             int rows, int maxit, double err, int prec)
{

    return solver_pcg(NULL, Asp, x, b, rows, maxit, err, prec, 0, 0);
}

int solver_pcg(double **A, G_math_spvector **Asp, double *x, double *b,
               int rows, int maxit, double err, int prec, int has_band,
               int bandwidth)
{
    double *r, *z;

    double *p;

    double *v;

    double s = 0.0;

    double a0 = 0, a1 = 0, mygamma, tmp = 0;

    int m, i;

    int finished = 2;

    int error_break;

    G_math_spvector **M;

    r = G_alloc_vector(rows);
    p = G_alloc_vector(rows);
    v = G_alloc_vector(rows);
    z = G_alloc_vector(rows);

    error_break = 0;

    /*compute the preconditioning matrix, this is a sparse matrix */
    M = create_diag_precond_matrix(A, Asp, rows, prec);

    /*
     * residual calculation
     */
#pragma omp parallel
    {
        if (Asp)
            G_math_Ax_sparse(Asp, x, v, rows);
        else if (has_band)
            G_math_Ax_sband(A, x, v, rows, bandwidth);
        else
            G_math_d_Ax(A, x, v, rows, rows);

        G_math_d_ax_by(b, v, r, 1.0, -1.0, rows);
        /*perform the preconditioning */
        G_math_Ax_sparse(M, r, p, rows);

        /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s)
        for (i = 0; i < rows; i++) {
            s += p[i] * r[i];
        }
    }

    a0 = s;
    s = 0.0;

    /* ******************* */
    /* start the iteration */
    /* ******************* */
    for (m = 0; m < maxit; m++) {
#pragma omp parallel default(shared)
        {
            if (Asp)
                G_math_Ax_sparse(Asp, p, v, rows);
            else if (has_band)
                G_math_Ax_sband(A, p, v, rows, bandwidth);
            else
                G_math_d_Ax(A, p, v, rows, rows);

                /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s)
            for (i = 0; i < rows; i++) {
                s += v[i] * p[i];
            }

            /* barrier */
#pragma omp single
            {
                tmp = s;
                mygamma = a0 / tmp;
                s = 0.0;
            }

            G_math_d_ax_by(p, x, x, mygamma, 1.0, rows);

            if (m % 50 == 1) {
                if (Asp)
                    G_math_Ax_sparse(Asp, x, v, rows);
                else if (has_band)
                    G_math_Ax_sband(A, x, v, rows, bandwidth);
                else
                    G_math_d_Ax(A, x, v, rows, rows);

                G_math_d_ax_by(b, v, r, 1.0, -1.0, rows);
            }
            else {
                G_math_d_ax_by(r, v, r, 1.0, -1.0 * mygamma, rows);
            }

            /*perform the preconditioning */
            G_math_Ax_sparse(M, r, z, rows);

            /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s)
            for (i = 0; i < rows; i++) {
                s += z[i] * r[i];
            }

            /* barrier */
#pragma omp single
            {
                a1 = s;
                tmp = a1 / a0;
                a0 = a1;
                s = 0.0;

                if (a1 < 0 || a1 == 0 || a1 > 0) {
                    ;
                }
                else {
                    G_warning(_("Unable to solve the linear equation system"));
                    error_break = 1;
                }
            }
            G_math_d_ax_by(p, z, p, tmp, 1.0, rows);
        }

        if (Asp != NULL)
            G_message(_("Sparse PCG -- iteration %i error  %g\n"), m, a0);
        else
            G_message(_("PCG -- iteration %i error  %g\n"), m, a0);

        if (error_break == 1) {
            finished = -1;
            break;
        }

        if (a0 < err) {
            finished = 1;
            break;
        }
    }

    G_free(r);
    G_free(p);
    G_free(v);
    G_free(z);
    G_math_free_spmatrix(M, rows);

    return finished;
}

/*!
 * \brief The iterative conjugate gradients solver for symmetric positive
 * definite matrices
 *
 * This iterative solver works with symmetric positive definite  regular
 * quadratic matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param A (double **) -- the matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_cg(double **A, double *x, double *b, int rows, int maxit,
                     double err)
{
    return solver_cg(A, NULL, x, b, rows, maxit, err, 0, 0);
}

/*!
 * \brief The iterative conjugate gradients solver for symmetric positive
 * definite band matrices
 *
 * This iterative solver works with symmetric positive definite band matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param A (double **) -- the symmetric positive definite band matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param bandwidth (int) -- the bandwidth of matrix A
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_cg_sband(double **A, double *x, double *b, int rows,
                           int bandwidth, int maxit, double err)
{
    return solver_cg(A, NULL, x, b, rows, maxit, err, 1, bandwidth);
}

/*!
 * \brief The iterative conjugate gradients solver for sparse symmetric positive
 * definite matrices
 *
 * This iterative solver works with symmetric positive definite sparse matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param Asp (G_math_spvector **) -- the sparse matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_sparse_cg(G_math_spvector **Asp, double *x, double *b,
                            int rows, int maxit, double err)
{
    return solver_cg(NULL, Asp, x, b, rows, maxit, err, 0, 0);
}

int solver_cg(double **A, G_math_spvector **Asp, double *x, double *b, int rows,
              int maxit, double err, int has_band, int bandwidth)
{
    double *r;

    double *p;

    double *v;

    double s = 0.0;

    double a0 = 0, a1 = 0, mygamma, tmp = 0;

    int m, i;

    int finished = 2;

    int error_break;

    r = G_alloc_vector(rows);
    p = G_alloc_vector(rows);
    v = G_alloc_vector(rows);

    error_break = 0;
    /*
     * residual calculation
     */
#pragma omp parallel
    {
        if (Asp)
            G_math_Ax_sparse(Asp, x, v, rows);
        else if (has_band)
            G_math_Ax_sband(A, x, v, rows, bandwidth);
        else
            G_math_d_Ax(A, x, v, rows, rows);

        G_math_d_ax_by(b, v, r, 1.0, -1.0, rows);
        G_math_d_copy(r, p, rows);

        /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s)
        for (i = 0; i < rows; i++) {
            s += r[i] * r[i];
        }
    }

    a0 = s;
    s = 0.0;

    /* ******************* */
    /* start the iteration */
    /* ******************* */
    for (m = 0; m < maxit; m++) {
#pragma omp parallel default(shared)
        {
            if (Asp)
                G_math_Ax_sparse(Asp, p, v, rows);
            else if (has_band)
                G_math_Ax_sband(A, p, v, rows, bandwidth);
            else
                G_math_d_Ax(A, p, v, rows, rows);

                /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s)
            for (i = 0; i < rows; i++) {
                s += v[i] * p[i];
            }

            /* barrier */
#pragma omp single
            {
                tmp = s;
                mygamma = a0 / tmp;
                s = 0.0;
            }

            G_math_d_ax_by(p, x, x, mygamma, 1.0, rows);

            if (m % 50 == 1) {
                if (Asp)
                    G_math_Ax_sparse(Asp, x, v, rows);
                else if (has_band)
                    G_math_Ax_sband(A, x, v, rows, bandwidth);
                else
                    G_math_d_Ax(A, x, v, rows, rows);

                G_math_d_ax_by(b, v, r, 1.0, -1.0, rows);
            }
            else {
                G_math_d_ax_by(r, v, r, 1.0, -1.0 * mygamma, rows);
            }

            /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s)
            for (i = 0; i < rows; i++) {
                s += r[i] * r[i];
            }

            /* barrier */
#pragma omp single
            {
                a1 = s;
                tmp = a1 / a0;
                a0 = a1;
                s = 0.0;

                if (a1 < 0 || a1 == 0 || a1 > 0) {
                    ;
                }
                else {
                    G_warning(_("Unable to solve the linear equation system"));
                    error_break = 1;
                }
            }
            G_math_d_ax_by(p, r, p, tmp, 1.0, rows);
        }

        if (Asp != NULL)
            G_message(_("Sparse CG -- iteration %i error  %g\n"), m, a0);
        else
            G_message(_("CG -- iteration %i error  %g\n"), m, a0);

        if (error_break == 1) {
            finished = -1;
            break;
        }

        if (a0 < err) {
            finished = 1;
            break;
        }
    }

    G_free(r);
    G_free(p);
    G_free(v);

    return finished;
}

/*!
 * \brief The iterative biconjugate gradients solver with stabilization for
 * unsymmetric non-definite matrices
 *
 * This iterative solver works with regular quadratic matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param A (double **) -- the matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_bicgstab(double **A, double *x, double *b, int rows,
                           int maxit, double err)
{
    return solver_bicgstab(A, NULL, x, b, rows, maxit, err);
}

/*!
 * \brief The iterative biconjugate gradients solver with stabilization for
 * unsymmetric non-definite matrices
 *
 * This iterative solver works with sparse matrices.
 *
 * This solver solves the linear equation system:
 *  A x = b
 *
 * The parameter <i>maxit</i> specifies the maximum number of iterations. If the
 * maximum is reached, the solver will abort the calculation and writes the
 * current result into the vector x. The parameter <i>err</i> defines the error
 * break criteria for the solver.
 *
 * \param Asp (G_math_spvector **) -- the sparse matrix
 * \param x (double *) -- the value vector
 * \param b (double *) -- the right hand side
 * \param rows (int)
 * \param maxit (int) -- the maximum number of iterations
 * \param err (double) -- defines the error break criteria
 * \return (int) -- 1 - success, 2 - not finished but success, 0 - matrix
 * singular, -1 - could not solve the les
 *
 * */
int G_math_solver_sparse_bicgstab(G_math_spvector **Asp, double *x, double *b,
                                  int rows, int maxit, double err)
{
    return solver_bicgstab(NULL, Asp, x, b, rows, maxit, err);
}

int solver_bicgstab(double **A, G_math_spvector **Asp, double *x, double *b,
                    int rows, int maxit, double err)
{
    double *r;

    double *r0;

    double *p;

    double *v;

    double *s;

    double *t;

    double s1 = 0.0, s2 = 0.0, s3 = 0.0;

    double alpha = 0, beta = 0, omega, rr0 = 0, error;

    int m, i;

    int finished = 2;

    int error_break;

    r = G_alloc_vector(rows);
    r0 = G_alloc_vector(rows);
    p = G_alloc_vector(rows);
    v = G_alloc_vector(rows);
    s = G_alloc_vector(rows);
    t = G_alloc_vector(rows);

    error_break = 0;

#pragma omp parallel
    {
        if (Asp)
            G_math_Ax_sparse(Asp, x, v, rows);
        else
            G_math_d_Ax(A, x, v, rows, rows);

        G_math_d_ax_by(b, v, r, 1.0, -1.0, rows);
        G_math_d_copy(r, r0, rows);
        G_math_d_copy(r, p, rows);
    }

    s1 = s2 = s3 = 0.0;

    /* ******************* */
    /* start the iteration */
    /* ******************* */
    for (m = 0; m < maxit; m++) {

#pragma omp parallel default(shared)
        {
            if (Asp)
                G_math_Ax_sparse(Asp, p, v, rows);
            else
                G_math_d_Ax(A, p, v, rows, rows);

                /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s1, s2, s3)
            for (i = 0; i < rows; i++) {
                s1 += r[i] * r[i];
                s2 += r[i] * r0[i];
                s3 += v[i] * r0[i];
            }

#pragma omp single
            {
                error = s1;

                if (error < 0 || error == 0 || error > 0) {
                    ;
                }
                else {
                    G_warning(_("Unable to solve the linear equation system"));
                    error_break = 1;
                }

                rr0 = s2;
                alpha = rr0 / s3;
                s1 = s2 = s3 = 0.0;
            }

            G_math_d_ax_by(r, v, s, 1.0, -1.0 * alpha, rows);
            if (Asp)
                G_math_Ax_sparse(Asp, s, t, rows);
            else
                G_math_d_Ax(A, s, t, rows, rows);

                /* scalar product */
#pragma omp for schedule(static) private(i) reduction(+ : s1, s2)
            for (i = 0; i < rows; i++) {
                s1 += t[i] * s[i];
                s2 += t[i] * t[i];
            }

#pragma omp single
            {
                omega = s1 / s2;
                s1 = s2 = 0.0;
            }

            G_math_d_ax_by(p, s, r, alpha, omega, rows);
            G_math_d_ax_by(x, r, x, 1.0, 1.0, rows);
            G_math_d_ax_by(s, t, r, 1.0, -1.0 * omega, rows);

#pragma omp for schedule(static) private(i) reduction(+ : s1)
            for (i = 0; i < rows; i++) {
                s1 += r[i] * r0[i];
            }

#pragma omp single
            {
                beta = alpha / omega * s1 / rr0;
                s1 = s2 = s3 = 0.0;
            }

            G_math_d_ax_by(p, v, p, 1.0, -1.0 * omega, rows);
            G_math_d_ax_by(p, r, p, beta, 1.0, rows);
        }

        if (Asp != NULL)
            G_message(_("Sparse BiCGStab -- iteration %i error  %g\n"), m,
                      error);
        else
            G_message(_("BiCGStab -- iteration %i error  %g\n"), m, error);

        if (error_break == 1) {
            finished = -1;
            break;
        }

        if (error < err) {
            finished = 1;
            break;
        }
    }

    G_free(r);
    G_free(r0);
    G_free(p);
    G_free(v);
    G_free(s);
    G_free(t);

    return finished;
}

/*!
 * \brief Compute a diagonal preconditioning matrix for krylov space solver
 *
 * \param A (double **) -- the matrix for which the precondition should be
 * computed (if the sparse matrix is used, set it to NULL) \param Asp
 * (G_math_spvector **) -- the matrix for which the precondition should be
 * computed \param rows (int) \param prec (int) -- which preconditioner should
 * be used 1, 2 or 3
 *
 * */
G_math_spvector **create_diag_precond_matrix(double **A, G_math_spvector **Asp,
                                             int rows, int prec)
{
    G_math_spvector **Msp;

    unsigned int i, j, cols = (unsigned int)rows;

    double sum;

    assert(rows >= 0);

    Msp = G_math_alloc_spmatrix(rows);

    if (A != NULL) {
#pragma omp parallel for schedule(static) private(i, j, sum) \
    shared(A, Msp, rows, cols, prec)
        for (i = 0; i < (unsigned int)rows; i++) {
            G_math_spvector *spvect = G_math_alloc_spvector(1);

            switch (prec) {
            case G_MATH_ROWSCALE_EUKLIDNORM_PRECONDITION:
                sum = 0;
                for (j = 0; j < cols; j++)
                    sum += A[i][j] * A[i][j];
                spvect->values[0] = 1.0 / sqrt(sum);
                break;
            case G_MATH_ROWSCALE_ABSSUMNORM_PRECONDITION:
                sum = 0;
                for (j = 0; j < cols; j++)
                    sum += fabs(A[i][j]);
                spvect->values[0] = 1.0 / (sum);
                break;
            case G_MATH_DIAGONAL_PRECONDITION:
            default:
                spvect->values[0] = 1.0 / A[i][i];
                break;
            }

            spvect->index[0] = i;
            spvect->cols = 1;
            ;
            G_math_add_spvector(Msp, spvect, i);
        }
    }
    else {
#pragma omp parallel for schedule(static) private(i, j, sum) \
    shared(Asp, Msp, rows, cols, prec)
        for (i = 0; i < (unsigned int)rows; i++) {
            G_math_spvector *spvect = G_math_alloc_spvector(1);

            switch (prec) {
            case G_MATH_ROWSCALE_EUKLIDNORM_PRECONDITION:
                sum = 0;
                for (j = 0; j < Asp[i]->cols; j++)
                    sum += Asp[i]->values[j] * Asp[i]->values[j];
                spvect->values[0] = 1.0 / sqrt(sum);
                break;
            case G_MATH_ROWSCALE_ABSSUMNORM_PRECONDITION:
                sum = 0;
                for (j = 0; j < Asp[i]->cols; j++)
                    sum += fabs(Asp[i]->values[j]);
                spvect->values[0] = 1.0 / (sum);
                break;
            case G_MATH_DIAGONAL_PRECONDITION:
            default:
                for (j = 0; j < Asp[i]->cols; j++)
                    if (i == Asp[i]->index[j])
                        spvect->values[0] = 1.0 / Asp[i]->values[j];
                break;
            }

            spvect->index[0] = i;
            spvect->cols = 1;
            ;
            G_math_add_spvector(Msp, spvect, i);
        }
    }
    return Msp;
}
