/* spdgemm.c

 * 

 * Copyright (C) 2014 Patrick Alken

 * 

 * 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, write to the Free Software

 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.

 */

#include <config.h>

#include <stdlib.h>

#include <math.h>

#include <gsl/gsl_spmatrix.h>

#include <gsl/gsl_spblas.h>

#include <gsl/gsl_errno.h>

/*

gsl_spblas_dgemm()

  Multiply two sparse matrices

Inputs: alpha - scalar factor

        A     - sparse matrix

        B     - sparse matrix

        C     - (output) C = alpha * A * B

Return: success or error

Notes:

1) based on CSparse routine cs_multiply

*/

int

gsl_spblas_dgemm(const double alpha, const gsl_spmatrix *A,

                 const gsl_spmatrix *B, gsl_spmatrix *C)

{

  if (A->size2 != B->size1 || A->size1 != C->size1 || B->size2 != C->size2)

    {

      GSL_ERROR("matrix dimensions do not match", GSL_EBADLEN);

    }

  else if (A->sptype != B->sptype || A->sptype != C->sptype)

    {

      GSL_ERROR("matrix storage formats do not match", GSL_EINVAL);

    }

  else if (!GSL_SPMATRIX_ISCCS(A))

    {

      GSL_ERROR("compressed column format required", GSL_EINVAL);

    }

  else

    {

      int status = GSL_SUCCESS;

      const size_t M = A->size1;

      const size_t N = B->size2;

      size_t *Bi = B->i;

      size_t *Bp = B->p;

      double *Bd = B->data;

      size_t *w = (size_t *) A->work; /* workspace of length M */

      double *x = (double *) C->work; /* workspace of length M */

      size_t *Cp, *Ci;

      double *Cd;

      size_t j, p;

      size_t nz = 0;

      if (C->nzmax < A->nz + B->nz)

        {

          status = gsl_spmatrix_realloc(A->nz + B->nz, C);

          if (status)

            {

              GSL_ERROR("unable to realloc matrix C", status);

            }

        }

      /* initialize workspace to 0 */

      for (j = 0; j < M; ++j)

        w[j] = 0;

      Cp = C->p;

      Ci = C->i;

      Cd = C->data;

      for (j = 0; j < N; ++j)

        {

          if (nz + M > C->nzmax)

            {

              status = gsl_spmatrix_realloc(2 * C->nzmax + M, C);

              if (status)

                {

                  GSL_ERROR("unable to realloc matrix C", status);

                }

              /* these pointers could have changed due to reallocation */

              Ci = C->i;

              Cd = C->data;

            }

          Cp[j] = nz; /* column j of C starts here */

          for (p = Bp[j]; p < Bp[j + 1]; ++p)

            {

              nz = gsl_spblas_scatter(A, Bi[p], Bd[p], w, x, j + 1, C, nz);

            }

          for (p = Cp[j]; p < nz; ++p)

            Cd[p] = x[Ci[p]];

        }

      Cp[N] = nz;

      C->nz = nz;

      /* scale by alpha */

      gsl_spmatrix_scale(C, alpha);

      return status;

    }

} /* gsl_spblas_dgemm() */

/*

gsl_spblas_scatter()

  Keep a running total x -> x + alpha*A(:,j) for adding matrices together in CCS,

which will eventually be stored in C(:,j)

  When a new non-zero element with row index i is found, update C->i with

the row index. C->data is updated only by the calling function after all

matrices have been added via this function.

Inputs: A     - sparse matrix m-by-n

        j     - column index

        alpha - scalar factor

        w     - keeps track which rows of column j have been added to C;

                initialize to 0 prior to first call

        x     - column vector of length m

        mark  -

        C     - output matrix whose jth column will be added to A(:,j)

        nz    - (input/output) number of non-zeros in matrix C

Notes:

1) This function is designed to be called successively when adding multiple

matrices together. Column j of C is stored contiguously as per CCS but not

necessarily in order - ie: the row indices C->i may not be in ascending order.

2) based on CSparse routine cs_scatter

*/

size_t

gsl_spblas_scatter(const gsl_spmatrix *A, const size_t j, const double alpha,

                   size_t *w, double *x, const size_t mark, gsl_spmatrix *C,

                   size_t nz)

{

  size_t p;

  size_t *Ai = A->i;

  size_t *Ap = A->p;

  double *Ad = A->data;

  size_t *Ci = C->i;

  for (p = Ap[j]; p < Ap[j + 1]; ++p)

    {

      size_t i = Ai[p];          /* A(i,j) is nonzero */

      if (w[i] < mark)           /* check if row i has been stored in column j yet */

        {

          w[i] = mark;           /* i is new entry in column j */

          Ci[nz++] = i;          /* add i to pattern of C(:,j) */

          x[i] = alpha * Ad[p];  /* x(i) = alpha * A(i,j) */

        }

      else                       /* this (i,j) exists in C from a previous call */

        {

          x[i] += alpha * Ad[p]; /* add alpha*A(i,j) to C(i,j) */

        }

    }

  return (nz) ;

} /* gsl_spblas_scatter() */