/* spgetset.c

 * 

 * Copyright (C) 2012 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_math.h>

#include <gsl/gsl_errno.h>

#include <gsl/gsl_spmatrix.h>

#include "avl.c"

static void *tree_find(const gsl_spmatrix *m, const size_t i, const size_t j);

double

gsl_spmatrix_get(const gsl_spmatrix *m, const size_t i, const size_t j)

{

  if (i >= m->size1)

    {

      GSL_ERROR_VAL("first index out of range", GSL_EINVAL, 0.0);

    }

  else if (j >= m->size2)

    {

      GSL_ERROR_VAL("second index out of range", GSL_EINVAL, 0.0);

    }

  else if (m->nz == 0)

    {

      /* no non-zero elements added to matrix */

      return (0.0);

    }

  else

    {

      if (GSL_SPMATRIX_ISTRIPLET(m))

        {

          /* traverse binary tree to search for (i,j) element */

          void *ptr = tree_find(m, i, j);

          double x = ptr ? *(double *) ptr : 0.0;

          return x;

        }

      else if (GSL_SPMATRIX_ISCCS(m))

        {

          const size_t *mi = m->i;

          const size_t *mp = m->p;

          size_t p;

          /* loop over column j and search for row index i */

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

            {

              if (mi[p] == i)

                return m->data[p];

            }

        }

      else if (GSL_SPMATRIX_ISCRS(m))

        {

          const size_t *mj = m->i;

          const size_t *mp = m->p;

          size_t p;

          /* loop over row i and search for column index j */

          for (p = mp[i]; p < mp[i + 1]; ++p)

            {

              if (mj[p] == j)

                return m->data[p];

            }

        }

      else

        {

          GSL_ERROR_VAL("unknown sparse matrix type", GSL_EINVAL, 0.0);

        }

      /* element not found; return 0 */

      return 0.0;

    }

} /* gsl_spmatrix_get() */

/*

gsl_spmatrix_set()

  Add an element to a matrix in triplet form

Inputs: m - spmatrix

        i - row index

        j - column index

        x - matrix value

*/

int

gsl_spmatrix_set(gsl_spmatrix *m, const size_t i, const size_t j,

                 const double x)

{

  if (!GSL_SPMATRIX_ISTRIPLET(m))

    {

      GSL_ERROR("matrix not in triplet representation", GSL_EINVAL);

    }

  else if (x == 0.0)

    {

      /* traverse binary tree to search for (i,j) element */

      void *ptr = tree_find(m, i, j);

      /*

       * just set the data element to 0; it would be easy to

       * delete the node from the tree with avl_delete(), but

       * we'd also have to delete it from the data arrays which

       * is less simple

       */

      if (ptr != NULL)

        *(double *) ptr = 0.0;

      return GSL_SUCCESS;

    }

  else

    {

      int s = GSL_SUCCESS;

      void *ptr;

      /* check if matrix needs to be realloced */

      if (m->nz >= m->nzmax)

        {

          s = gsl_spmatrix_realloc(2 * m->nzmax, m);

          if (s)

            return s;

        }

      /* store the triplet (i, j, x) */

      m->i[m->nz] = i;

      m->p[m->nz] = j;

      m->data[m->nz] = x;

      ptr = avl_insert(m->tree_data->tree, &m->data[m->nz]);

      if (ptr != NULL)

        {

          /* found duplicate entry (i,j), replace with new x */

          *((double *) ptr) = x;

        }

      else

        {

          /* no duplicate (i,j) found, update indices as needed */

          /* increase matrix dimensions if needed */

          m->size1 = GSL_MAX(m->size1, i + 1);

          m->size2 = GSL_MAX(m->size2, j + 1);

          ++(m->nz);

        }

      return s;

    }

} /* gsl_spmatrix_set() */

double *

gsl_spmatrix_ptr(gsl_spmatrix *m, const size_t i, const size_t j)

{

  if (i >= m->size1)

    {

      GSL_ERROR_NULL("first index out of range", GSL_EINVAL);

    }

  else if (j >= m->size2)

    {

      GSL_ERROR_NULL("second index out of range", GSL_EINVAL);

    }

  else

    {

      if (GSL_SPMATRIX_ISTRIPLET(m))

        {

          /* traverse binary tree to search for (i,j) element */

          void *ptr = tree_find(m, i, j);

          return (double *) ptr;

        }

      else if (GSL_SPMATRIX_ISCCS(m))

        {

          const size_t *mi = m->i;

          const size_t *mp = m->p;

          size_t p;

          /* loop over column j and search for row index i */

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

            {

              if (mi[p] == i)

                return &(m->data[p]);

            }

        }

      else if (GSL_SPMATRIX_ISCRS(m))

        {

          const size_t *mj = m->i;

          const size_t *mp = m->p;

          size_t p;

          /* loop over row i and search for column index j */

          for (p = mp[i]; p < mp[i + 1]; ++p)

            {

              if (mj[p] == j)

                return &(m->data[p]);

            }

        }

      else

        {

          GSL_ERROR_NULL("unknown sparse matrix type", GSL_EINVAL);

        }

      /* element not found; return 0 */

      return NULL;

    }

} /* gsl_spmatrix_ptr() */

/*

tree_find()

  Find node in tree corresponding to matrix entry (i,j). Adapted

from avl_find()

Inputs: m - spmatrix

        i - row index

        j - column index

Return: pointer to tree node data if found, NULL if not found

*/

static void *

tree_find(const gsl_spmatrix *m, const size_t i, const size_t j)

{

  const struct avl_table *tree = (struct avl_table *) m->tree_data->tree;

  const struct avl_node *p;

  for (p = tree->avl_root; p != NULL; )

    {

      size_t n = (double *) p->avl_data - m->data;

      size_t pi = m->i[n];

      size_t pj = m->p[n];

      int cmp = gsl_spmatrix_compare_idx(i, j, pi, pj);

      if (cmp < 0)

        p = p->avl_link[0];

      else if (cmp > 0)

        p = p->avl_link[1];

      else /* |cmp == 0| */

        return p->avl_data;

    }

  return NULL;

} /* tree_find() */