#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <grass/gis.h>
#include <grass/gmath.h>
#include <grass/glocale.h>

/**
 * \brief Cholesky decomposition of a symmetric band matrix
 *
 * \param A (double**) the input symmetric band matrix
 * \param T (double**) the resulting lower tringular symmetric band matrix
 * \param rows (int) number of rows
 * \param bandwidth (int) the bandwidth of the symmetric band matrix
 *
 * */

void G_math_cholesky_sband_decomposition(double **A, double **T, int rows,
                                         int bandwidth)
{
    int i, j, k, end;
    double sum;

    G_debug(2, "G_math_cholesky_sband_decomposition(): n=%d  bandwidth=%d",
            rows, bandwidth);

    for (i = 0; i < rows; i++) {
        G_percent(i, rows, 9);
        /* For j = 0 */
        sum = A[i][0];
        end = ((bandwidth - 0) < (i + 1) ? (bandwidth - 0) : (i + 1));
        for (k = 1; k < end; k++)
            sum -= T[i - k][k] * T[i - k][0 + k];
        if (sum <= 0.0)
            G_fatal_error(_("Decomposition failed at row %i and col %i"), i, 0);
        T[i][0] = sqrt(sum);

#pragma omp parallel for schedule(static) private(j, k, end, sum) \
    shared(A, T, i, bandwidth)
        for (j = 1; j < bandwidth; j++) {
            sum = A[i][j];
            end = ((bandwidth - j) < (i + 1) ? (bandwidth - j) : (i + 1));
            for (k = 1; k < end; k++)
                sum -= T[i - k][k] * T[i - k][j + k];
            T[i][j] = sum / T[i][0];
        }
    }

    G_percent(i, rows, 2);
    return;
}

/**
 * \brief Cholesky symmetric band matrix solver for linear equation systems of
 * type Ax = b
 *
 * \param A (double**) the input symmetric band matrix
 * \param x (double*) the resulting vector, result is written in here
 * \param b (double*) the right hand side of Ax = b
 * \param rows (int) number of rows
 * \param bandwidth (int) the bandwidth of the symmetric band matrix
 *
 * */

void G_math_solver_cholesky_sband(double **A, double *x, double *b, int rows,
                                  int bandwidth)
{

    double **T;

    T = G_alloc_matrix(rows, bandwidth);

    G_math_cholesky_sband_decomposition(A, T, rows,
                                        bandwidth); /* T computation */
    G_math_cholesky_sband_substitution(T, x, b, rows, bandwidth);

    G_free_matrix(T);

    return;
}

/**
 * \brief Forward and backward substitution of a lower tringular symmetric band
 * matrix of A from system Ax = b
 *
 * \param T (double**) the lower triangle symmetric band matrix
 * \param x (double*) the resulting vector
 * \param b (double*) the right hand side of Ax = b
 * \param rows (int) number of rows
 * \param bandwidth (int) the bandwidth of the symmetric band matrix
 *
 * */

void G_math_cholesky_sband_substitution(double **T, double *x, double *b,
                                        int rows, int bandwidth)
{

    int i, j, start, end;

    /* Forward substitution */
    x[0] = b[0] / T[0][0];
    for (i = 1; i < rows; i++) {
        x[i] = b[i];
        /* start = 0 or i - bandwidth + 1 */
        start = ((i - bandwidth + 1) < 0 ? 0 : (i - bandwidth + 1));
        /* end = i */
        for (j = start; j < i; j++)
            x[i] -= T[j][i - j] * x[j];
        x[i] = x[i] / T[i][0];
    }

    /* Backward substitution */
    x[rows - 1] = x[rows - 1] / T[rows - 1][0];
    for (i = rows - 2; i >= 0; i--) {
        /* start = i + 1 */
        /* end = rows or bandwidth + i */
        end = (rows < (bandwidth + i) ? rows : (bandwidth + i));
        for (j = i + 1; j < end; j++)
            x[i] -= T[i][j - i] * x[j];
        x[i] = x[i] / T[i][0];
    }

    return;
}

/*--------------------------------------------------------------------------------------*/
/* Tcholetsky matrix invertion */

void G_math_cholesky_sband_invert(double **A, double *invAdiag, int rows,
                                  int bandwidth)
{
    double **T = NULL;
    double *vect = NULL;
    int i, j, k, start;
    double sum;

    T = G_alloc_matrix(rows, bandwidth);
    vect = G_alloc_vector(rows);

    /* T computation                */
    G_math_cholesky_sband_decomposition(A, T, rows, bandwidth);

    /* T Diagonal invertion */
    for (i = 0; i < rows; i++) {
        T[i][0] = 1.0 / T[i][0];
    }

    /* A Diagonal invertion */
    for (i = 0; i < rows; i++) {
        vect[0] = T[i][0];
        invAdiag[i] = vect[0] * vect[0];
        for (j = i + 1; j < rows; j++) {
            sum = 0.0;
            /* start = i or j - bandwidth + 1 */
            start = ((j - bandwidth + 1) < i ? i : (j - bandwidth + 1));
            /* end = j */
            for (k = start; k < j; k++) {
                sum -= vect[k - i] * T[k][j - k];
            }
            vect[j - i] = sum * T[j][0];
            invAdiag[i] += vect[j - i] * vect[j - i];
        }
    }

    G_free_matrix(T);
    G_free_vector(vect);

    return;
}

/*--------------------------------------------------------------------------------------*/
/* Tcholetsky matrix solution and invertion */

void G_math_solver_cholesky_sband_invert(double **A, double *x, double *b,
                                         double *invAdiag, int rows,
                                         int bandwidth)
{

    double **T = NULL;
    double *vect = NULL;
    int i, j, k, start;
    double sum;

    T = G_alloc_matrix(rows, bandwidth);
    vect = G_alloc_vector(rows);

    /* T computation                */
    G_math_cholesky_sband_decomposition(A, T, rows, bandwidth);
    G_math_cholesky_sband_substitution(T, x, b, rows, bandwidth);

    /* T Diagonal invertion */
    for (i = 0; i < rows; i++) {
        T[i][0] = 1.0 / T[i][0];
    }

    /* A Diagonal invertion */
    for (i = 0; i < rows; i++) {
        vect[0] = T[i][0];
        invAdiag[i] = vect[0] * vect[0];
        for (j = i + 1; j < rows; j++) {
            sum = 0.0;
            /* start = i or j - bandwidth + 1 */
            start = ((j - bandwidth + 1) < i ? i : (j - bandwidth + 1));
            /* end = j */
            for (k = start; k < j; k++) {
                sum -= vect[k - i] * T[k][j - k];
            }
            vect[j - i] = sum * T[j][0];
            invAdiag[i] += vect[j - i] * vect[j - i];
        }
    }

    G_free_matrix(T);
    G_free_vector(vect);

    return;
}
