/* -*- mode: C; mode: fold; -*- */
/*
  Copyright (c) 2007 Massachusetts Institute of Technology

  This software was developed by the MIT Center for Space Research
  under contract SV1-61010 from the Smithsonian Institution.

  Permission to use, copy, modify, distribute, and sell this software
  and its documentation for any purpose is hereby granted without fee,
  provided that the above copyright notice appear in all copies and
  that both that copyright notice and this permission notice appear in
  the supporting documentation, and that the name of the Massachusetts
  Institute of Technology not be used in advertising or publicity
  pertaining to distribution of the software without specific, written
  prior permission.  The Massachusetts Institute of Technology makes
  no representations about the suitability of this software for any
  purpose.  It is provided "as is" without express or implied warranty.

  THE MASSACHUSETTS INSTITUTE OF TECHNOLOGY DISCLAIMS ALL WARRANTIES
  WITH REGARD TO THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF
  MERCHANTABILITY AND FITNESS.  IN NO EVENT SHALL THE MASSACHUSETTS
  INSTITUTE OF TECHNOLOGY BE LIABLE FOR ANY SPECIAL, INDIRECT OR
  CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM LOSS
  OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT,
  NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION
  WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
*/

/* Author: John E. Davis (jed@jedsoft.org) */

#include <stdio.h>
#include <string.h>
#include <slang.h>

#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_eigen.h>

#include "config.h"
#include "slgsl.h"
#include "version.h"

#ifdef __cplusplus
extern "C"
{
#endif
/* SLANG_MODULE(gslmatrix); */
#ifdef __cplusplus
}
#endif

static int check_for_complex_args (int nargs, SLtype *tp)
{
   unsigned int i, n;

   *tp = SLANG_DOUBLE_TYPE;
   if (nargs <= 0)
     return 0;

   n = (unsigned int) nargs;
   for (i = 0; i < n; i++)
     {
	int type = SLang_peek_at_stack1_n (i);
	if (type == -1)
	  return -1;
	if (type == SLANG_COMPLEX_TYPE)
	  {
	     *tp = SLANG_COMPLEX_TYPE;
	     return 0;
	  }
     }
   return 0;
}

static int pop_permutation (gsl_permutation **pp)
{
   gsl_permutation *p;
   SLang_Array_Type *at;
   unsigned int i, n;
   unsigned int *data;
   size_t *pdata;

   *pp = NULL;

   if (-1 == SLang_pop_array_of_type (&at, SLANG_UINT_TYPE))
     return -1;

   data = (unsigned int *) at->data;
   n = at->num_elements;

   if (n == 0)
     {
	SLang_verror (SL_INVALID_PARM, "Empty permutation array");
	SLang_free_array (at);
	return -1;
     }

   if (NULL == (p = gsl_permutation_alloc (n)))
     {
	SLang_free_array (at);
	return -1;
     }
   pdata = p->data;

   for (i = 0; i < n; i++)
     {
	if (data[i] >= n)
	  {
	     SLang_verror (SL_INVALID_PARM, "Invalid permutation array");
	     SLang_free_array (at);
	     gsl_permutation_free (p);
	     return -1;
	  }
	pdata[i] = data[i];
     }
   SLang_free_array (at);
   *pp = p;
   return 0;
}

static int push_permutation (gsl_permutation *p)
{
   SLang_Array_Type *at;
   unsigned int *data;
   size_t *pdata;
   SLindex_Type i, n;

   n = p->size;

   if (NULL == (at = SLang_create_array (SLANG_UINT_TYPE, 0, NULL, &n, 1)))
     return -1;

   pdata = p->data;
   data = (unsigned int *) at->data;

   for (i = 0; i < n; i++)
     data[i] = pdata[i];

   return SLang_push_array (at, 1);
}

static void linalg_LU_decomp (void)
{
   SLang_Ref_Type *signum_ref = NULL;
   SLGSL_Matrix_Type *matrix;
   gsl_permutation *p;
   int signum;
   SLtype type;
   int nargs = SLang_Num_Function_Args;

   if (-1 == check_for_complex_args (nargs, &type))
     return;

   switch (nargs)
     {
      case 2:
	if (-1 == SLang_pop_ref (&signum_ref))
	  return;
	/* fall through */
      case 1:
	if (-1 == slgsl_pop_square_matrix (&matrix, type, 1))
	  {
	     if (signum_ref != NULL)
	       SLang_free_ref (signum_ref);
	     return;
	  }
	break;

      default:
	SLang_verror (SL_USAGE_ERROR, "Usage: (LU, p) = linalg_LU_decomp(A [,&signum])");
	return;
     }

   if (NULL == (p = gsl_permutation_alloc (matrix->size1)))
     {
	slgsl_free_matrix (matrix);
	if (signum_ref != NULL)
	  SLang_free_ref (signum_ref);
	return;
     }

   slgsl_reset_errors ();
   if (type == SLANG_COMPLEX_TYPE)
     gsl_linalg_complex_LU_decomp (&matrix->m.c, p, &signum);
   else
     gsl_linalg_LU_decomp (&matrix->m.d, p, &signum);

   slgsl_check_errors ("linalg_LU_decomp");

   if ((0 == slgsl_push_matrix (matrix))
       && (0 == push_permutation (p))
       && (signum_ref != NULL))
     (void) SLang_assign_to_ref (signum_ref, SLANG_INT_TYPE, (VOID_STAR)&signum);

   if (signum_ref != NULL)
     SLang_free_ref (signum_ref);

   gsl_permutation_free (p);
   slgsl_free_matrix (matrix);
}

static void linalg_LU_solve (void)
{
   SLGSL_Matrix_Type *lu = NULL;
   SLGSL_Vector_Type *b = NULL;
   SLGSL_Vector_Type *x = NULL;
   gsl_permutation *p = NULL;
   SLtype type;
   int nargs = SLang_Num_Function_Args;

   if (-1 == check_for_complex_args (nargs, &type))
     return;

   switch (nargs)
     {
      case 3:
	if ((-1 == slgsl_pop_vector (&b, type, 0))
	    || (-1 == pop_permutation (&p))
	    || (-1 == slgsl_pop_square_matrix (&lu, type, 0)))
	  goto return_error;

	if ((lu->size2 != b->size)
	    || (p->size != b->size))
	  {
	     SLang_verror (SL_INVALID_PARM, "matrices have incompatible dimensions");
	     goto return_error;
	  }
	break;

      default:
	SLang_verror (SL_USAGE_ERROR, "Usage: x = linalg_LU_solve(LU, p, b);");
	return;
     }

   if (NULL == (x = slgsl_new_vector (type, b->size, 0, NULL)))
     goto return_error;

   slgsl_reset_errors ();
   if (type == SLANG_COMPLEX_TYPE)
     gsl_linalg_complex_LU_solve (&lu->m.c, p, &b->v.c, &x->v.c);
   else
     gsl_linalg_LU_solve (&lu->m.d, p, &b->v.d, &x->v.d);
   slgsl_check_errors ("linalg_LU_solve");

   if (0 == SLang_get_error ())
     (void) slgsl_push_vector (x);

   /* drop */

   return_error:

   slgsl_free_vector (x);
   slgsl_free_matrix (lu);
   gsl_permutation_free (p);
   slgsl_free_vector (b);
}

static void do_linalg_LU_det (int nargs, int do_log)
{
   SLGSL_Matrix_Type *matrix;
   SLtype type;
   int signum;
   const char *func;

   if (do_log)
     {
	func = "linalg_LU_lndet";
	if (nargs != 1)
	  {
	     SLang_verror (SL_USAGE_ERROR, "Usage: det = linalg_LU_lndet (LU)");
	     return;
	  }
	signum = 0;
     }
   else
     {
	func = "linalg_LU_det";
	if (nargs != 2)
	  {
	     SLang_verror (SL_USAGE_ERROR, "Usage: det = linalg_LU_det (LU, signum)");
	     return;
	  }
	if (-1 == SLang_pop_int (&signum))
	  return;
     }

   if (-1 == check_for_complex_args (1, &type))
     return;

   if (-1 == slgsl_pop_square_matrix (&matrix, type, 0))
     return;

   slgsl_reset_errors ();
   if (type == SLANG_COMPLEX_TYPE)
     {
	if (do_log)
	  {
	     double d = gsl_linalg_complex_LU_lndet (&matrix->m.c);
	     (void) SLang_push_double (d);
	  }
	else
	  {
	     gsl_complex c = gsl_linalg_complex_LU_det (&matrix->m.c, signum);
	     (void) SLang_push_complex (c.dat[0], c.dat[1]);
	  }
     }
   else
     {
	double d;
	if (do_log)
	  d = gsl_linalg_LU_lndet (&matrix->m.d);
	else
	  d = gsl_linalg_LU_det (&matrix->m.d, signum);
	(void) SLang_push_double (d);
     }
   slgsl_check_errors (func);

   slgsl_free_matrix (matrix);
}

static void linalg_LU_det (void)
{
   do_linalg_LU_det (SLang_Num_Function_Args, 0);
}

static void linalg_LU_lndet (void)
{
   do_linalg_LU_det (SLang_Num_Function_Args, 1);
}

static void linalg_LU_invert (void)
{
   SLGSL_Matrix_Type *lu = NULL;
   SLGSL_Matrix_Type *inv = NULL;
   gsl_permutation *p = NULL;
   SLtype type;
   int nargs = SLang_Num_Function_Args;

   if (-1 == check_for_complex_args (nargs, &type))
     return;

   if (nargs != 2)
     {
	SLang_verror (SL_USAGE_ERROR, "Usage: inv = linalg_LU_invert(LU, p);");
	return;
     }
   if ((-1 == pop_permutation (&p))
       || (-1 == slgsl_pop_square_matrix (&lu, type, 1)))
     goto return_error;

   if (NULL == (inv = slgsl_new_matrix (type, lu->size1, lu->size2, 0, NULL)))
     goto return_error;

   slgsl_reset_errors ();
   if (type == SLANG_COMPLEX_TYPE)
     gsl_linalg_complex_LU_invert (&lu->m.c, p, &inv->m.c);
   else
     gsl_linalg_LU_invert (&lu->m.d, p, &inv->m.d);

   slgsl_check_errors ("linalg_LU_solve");

   if (0 == SLang_get_error ())
     (void) slgsl_push_matrix (inv);

   /* drop */

   return_error:

   slgsl_free_matrix (inv);
   slgsl_free_matrix (lu);
   gsl_permutation_free (p);
}

static void linalg_QR_decomp (void)
{
   SLGSL_Matrix_Type *matrix;
   SLGSL_Vector_Type *tau;
   unsigned int n;
   SLtype type = SLANG_DOUBLE_TYPE;
   int nargs = SLang_Num_Function_Args;

   if (nargs != 1)
     {
	SLang_verror (SL_USAGE_ERROR, "Usage: (QR, tau) = linalg_QR_decomp(A)");
	return;
     }

   if (-1 == slgsl_pop_matrix (&matrix, type, 1))
     return;

   n = matrix->size1;
   if (matrix->size2 < n)
     n = matrix->size2;

   if (NULL == (tau = slgsl_new_vector (SLANG_DOUBLE_TYPE, n, 0, NULL)))
     {
	slgsl_free_matrix (matrix);
	return;
     }

   slgsl_reset_errors ();
   gsl_linalg_QR_decomp (&matrix->m.d, &tau->v.d);
   slgsl_check_errors ("linalg_LU_decomp");

   (void) slgsl_push_matrix (matrix);
   (void) slgsl_push_vector (tau);
   slgsl_free_vector (tau);
   slgsl_free_matrix (matrix);
}

static void linalg_QR_solve (void)
{
   SLGSL_Matrix_Type *qr = NULL;
   SLGSL_Vector_Type *b = NULL;
   SLGSL_Vector_Type *x = NULL;
   SLGSL_Vector_Type *tau = NULL;
   SLGSL_Vector_Type *residual = NULL;
   SLang_Ref_Type *ref = NULL;
   SLtype type;
   int nargs = SLang_Num_Function_Args;

   type = SLANG_DOUBLE_TYPE;

   switch (nargs)
     {
      case 4:
	if (-1 == SLang_pop_ref (&ref))
	  return;
	/* fall through */
      case 3:
	if ((-1 == slgsl_pop_vector (&b, type, 0))
	    || (-1 == slgsl_pop_vector (&tau, type, 0))
	    || (-1 == slgsl_pop_matrix (&qr, type, 0)))
	  goto return_error;
	break;

      default:
	SLang_verror (SL_USAGE_ERROR, "Usage: x = linalg_QR_solve(QR, tau, b [,&residual]);");
	return;
     }

   if (qr->size2 != b->size)
     {
	SLang_verror (SL_INVALID_PARM, "matrices have incompatible dimensions");
	goto return_error;
     }

   if (NULL == (x = slgsl_new_vector (type, b->size, 0, NULL)))
     goto return_error;

   if ((ref != NULL)
       || (qr->size1 != qr->size2))
     {
	if (NULL == (residual = slgsl_new_vector (type, b->size, 0, NULL)))
	  goto return_error;
     }

   slgsl_reset_errors ();
   if (residual == NULL)
     gsl_linalg_QR_solve (&qr->m.d, &tau->v.d, &b->v.d, &x->v.d);
   else
     gsl_linalg_QR_lssolve (&qr->m.d, &tau->v.d, &b->v.d, &x->v.d, &residual->v.d);
   slgsl_check_errors ("linalg_LU_solve");

   if (0 == SLang_get_error ())
     {
	(void) slgsl_push_vector (x);
	if (ref != NULL)
	  (void) slgsl_assign_vector_to_ref (residual, ref);
     }

   /* drop */
   return_error:

   slgsl_free_vector (x);
   slgsl_free_matrix (qr);
   slgsl_free_vector (tau);
   slgsl_free_vector (b);
   if (ref != NULL)
     SLang_free_ref (ref);
   if (residual != NULL)
     slgsl_free_vector (residual);
}

static void linalg_SV_decomp (void)
{
   SLGSL_Matrix_Type *a = NULL, *v = NULL;
   SLGSL_Vector_Type *s = NULL;
   gsl_vector *work = NULL;
   size_t N,M;
   SLtype type;
   int nargs = SLang_Num_Function_Args;

   if (nargs != 1)
     {
	SLang_verror (SL_USAGE_ERROR, "Usage: (U,S,V) = linalg_SV_decomp(A); %% ==> A=U#S#transpose(V)");
	return;
     }
   if (-1 == check_for_complex_args (nargs, &type))
     return;

   if (type == SLANG_COMPLEX_TYPE)
     {
	SLang_verror (SL_NOT_IMPLEMENTED, "GSL does not support the SVD of complex arrays");
	return;
     }

   if (-1 == slgsl_pop_matrix (&a, type, 1))
     return;

   M = a->size1;
   N = a->size2;

   if (M < N)
     {
	SLang_verror (SL_INVALID_PARM, "Expecting a matrix with nrows>=ncols");
	slgsl_free_matrix (a);
	return;
     }

   if ((NULL == (s = slgsl_new_vector (type, N, 0, NULL)))
       || (NULL == (v = slgsl_new_matrix (type, N, N, 0, NULL)))
       || (NULL == (work = gsl_vector_alloc (N))))
     goto return_error;

   slgsl_reset_errors ();
   (void) gsl_linalg_SV_decomp (&a->m.d, &v->m.d, &s->v.d, work);
   slgsl_check_errors ("linalg_SV_decomp");

   if (0 == SLang_get_error ())
     {
	(void) slgsl_push_matrix (a);
	(void) slgsl_push_vector (s);
	(void) slgsl_push_matrix (v);
     }

   /* drop */
   return_error:
   if (work != NULL)
     gsl_vector_free (work);
   slgsl_free_matrix (v);
   slgsl_free_vector (s);
   slgsl_free_matrix (a);
}

static void linalg_SV_solve (void)
{
   SLGSL_Matrix_Type *u = NULL, *v = NULL;
   SLGSL_Vector_Type *b = NULL, *x = NULL, *s = NULL;
   size_t M, N;
   SLtype type;
   int nargs = SLang_Num_Function_Args;

   if (nargs != 4)
     {
	SLang_verror (SL_USAGE_ERROR, "Usage: x = linalg_SV_solve (U,V,S,b);");
	return;
     }

   if (-1 == check_for_complex_args (nargs, &type))
     return;

   if (type == SLANG_COMPLEX_TYPE)
     {
	SLang_verror (SL_NOT_IMPLEMENTED, "GSL does not support the SVD of complex arrays");
	return;
     }

   if ((-1 == slgsl_pop_vector (&b, type, 0))      /* N */
       || (-1 == slgsl_pop_vector (&s, type, 0))      /* N */
       || (-1 == slgsl_pop_square_matrix (&v, type, 0))   /* N */
       || (-1 == slgsl_pop_matrix (&u, type, 0)))   /* MxN */
     goto return_error;

   N = b->size;
   if ((s->size != N)
       || (v->size1 != N)
       || (u->size2 != N))
     {
	SLang_verror (SL_INVALID_PARM, "matrices have incompatible dimensions");
	goto return_error;
     }
   M = u->size1;
   if (M < N)
     {
	SLang_verror (SL_INVALID_PARM, "Context requires a matrix with nrows>=ncols");
	goto return_error;
     }

   if (NULL == (x = slgsl_new_vector (type, N, 0, NULL)))
     goto return_error;

   slgsl_reset_errors ();
   gsl_linalg_SV_solve (&u->m.d, &v->m.d, &s->v.d, &b->v.d, &x->v.d);
   slgsl_check_errors ("linalg_SV_solve");

   if (0 == SLang_get_error ())
     (void) slgsl_push_vector (x);

   /* drop */

   return_error:

   slgsl_free_vector (x);
   slgsl_free_vector (b);
   slgsl_free_vector (s);
   slgsl_free_matrix (v);
   slgsl_free_matrix (u);
}

/* Eigenvalue Routines */
static void eigen_symmv (void)
{
   SLGSL_Matrix_Type *matrix;
   SLtype type = SLANG_DOUBLE_TYPE;
   SLGSL_Vector_Type *eigvals = NULL;
   SLGSL_Matrix_Type *eigvecs = NULL;
   unsigned int n;

   if (SLang_Num_Function_Args != 1)
     {
	SLang_verror (SL_USAGE_ERROR, "Usage: (eigvecs, eigvals)=eigen_symmv(A)");
	return;
     }

   if (-1 == check_for_complex_args (1, &type))
     return;

   if (-1 == slgsl_pop_square_matrix (&matrix, type, 1))
     return;

   n = matrix->size1;

   if ((NULL == (eigvals = slgsl_new_vector (SLANG_DOUBLE_TYPE, n, 0, NULL)))
       || (NULL == (eigvecs = slgsl_new_matrix (type, n, n, 0, NULL))))
     goto return_error;

   slgsl_reset_errors ();
   if (type == SLANG_COMPLEX_TYPE)
     {
	gsl_eigen_hermv_workspace *w = gsl_eigen_hermv_alloc (n);
	if (w == NULL)
	  goto return_error;
	(void) gsl_eigen_hermv (&matrix->m.c, &eigvals->v.d, &eigvecs->m.c, w);
	gsl_eigen_hermv_free (w);
     }
   else
     {
	gsl_eigen_symmv_workspace *w = gsl_eigen_symmv_alloc (n);
	if (w == NULL)
	  goto return_error;
	(void) gsl_eigen_symmv (&matrix->m.d, &eigvals->v.d, &eigvecs->m.d, w);
	gsl_eigen_symmv_free (w);
     }
   slgsl_check_errors ("eigen_symmv");

   if (0 == SLang_get_error ())
     {
	if (type == SLANG_COMPLEX_TYPE)
	  gsl_eigen_hermv_sort (&eigvals->v.d, &eigvecs->m.c, GSL_EIGEN_SORT_ABS_DESC);
	else
	  gsl_eigen_symmv_sort (&eigvals->v.d, &eigvecs->m.d, GSL_EIGEN_SORT_ABS_DESC);
	(void) slgsl_push_matrix (eigvecs);
	(void) slgsl_push_vector (eigvals);
     }
   /* drop */
   return_error:

   slgsl_free_matrix (eigvecs);
   slgsl_free_vector (eigvals);
   slgsl_free_matrix (matrix);
}

#if GSL_VERSION_INT >= 10900
static void eigen_nonsymmv (void)
{
   SLGSL_Matrix_Type *matrix;
   SLGSL_Vector_Type *eigvals = NULL;
   SLGSL_Matrix_Type *eigvecs = NULL;
   gsl_eigen_nonsymmv_workspace *w = NULL;
   unsigned int n;

   if (SLang_Num_Function_Args != 1)
     {
	SLang_verror (SL_USAGE_ERROR, "Usage: (eigvecs, eigvals)=eigen_nonsymmv(A)");
	return;
     }

   if (-1 == slgsl_pop_square_matrix (&matrix, SLANG_DOUBLE_TYPE, 1))
     return;

   n = matrix->size1;

   if ((NULL == (eigvals = slgsl_new_vector (SLANG_COMPLEX_TYPE, n, 0, NULL)))
       || (NULL == (eigvecs = slgsl_new_matrix (SLANG_COMPLEX_TYPE, n, n, 0, NULL)))
       || (NULL == (w = gsl_eigen_nonsymmv_alloc (n))))
     goto return_error;

   slgsl_reset_errors ();
   (void) gsl_eigen_nonsymmv (&matrix->m.d, &eigvals->v.c, &eigvecs->m.c, w);
   slgsl_check_errors ("eigen_nonsymmv");

   if (0 == SLang_get_error ())
     {
	gsl_eigen_nonsymmv_sort (&eigvals->v.c, &eigvecs->m.c, GSL_EIGEN_SORT_ABS_DESC);
	(void) slgsl_push_matrix (eigvecs);
	(void) slgsl_push_vector (eigvals);
     }
   /* drop */

   return_error:

   gsl_eigen_nonsymmv_free (w);
   slgsl_free_matrix (eigvecs);
   slgsl_free_vector (eigvals);
   slgsl_free_matrix (matrix);
}
#endif

#define V SLANG_VOID_TYPE
static SLang_Intrin_Fun_Type Module_Intrinsics [] =
{
   MAKE_INTRINSIC_0("linalg_LU_decomp", linalg_LU_decomp, V),
   MAKE_INTRINSIC_0("linalg_LU_det", linalg_LU_det, V),
   MAKE_INTRINSIC_0("linalg_LU_lndet", linalg_LU_lndet, V),
   MAKE_INTRINSIC_0("linalg_LU_invert", linalg_LU_invert, V),
   MAKE_INTRINSIC_0("linalg_LU_solve", linalg_LU_solve, V),
   MAKE_INTRINSIC_0("linalg_QR_decomp", linalg_QR_decomp, V),
   MAKE_INTRINSIC_0("linalg_QR_solve", linalg_QR_solve, V),
   MAKE_INTRINSIC_0("linalg_SV_decomp", linalg_SV_decomp, V),
   MAKE_INTRINSIC_0("linalg_SV_solve", linalg_SV_solve, V),

   MAKE_INTRINSIC_0("eigen_symmv", eigen_symmv, V),
#if GSL_VERSION_INT >= 10900
   MAKE_INTRINSIC_0("eigen_nonsymmv", eigen_nonsymmv, V),
#endif
   SLANG_END_INTRIN_FUN_TABLE
};
#undef V

static SLang_Intrin_Var_Type Module_Variables [] =
{
   MAKE_VARIABLE("_gslmatrix_module_version_string", &Module_Version_String, SLANG_STRING_TYPE, 1),
   SLANG_END_INTRIN_VAR_TABLE
};

static SLang_IConstant_Type Module_IConstants [] =
{
   MAKE_ICONSTANT("_gslmatrix_module_version", MODULE_VERSION_NUMBER),
   SLANG_END_ICONST_TABLE
};

int init_gslmatrix_module_ns (char *ns_name)
{
   SLang_NameSpace_Type *ns = SLns_create_namespace (ns_name);
   if (ns == NULL)
     return -1;

   if (
       (-1 == SLns_add_intrin_fun_table (ns, Module_Intrinsics, NULL))
       || (-1 == SLns_add_intrin_var_table (ns, Module_Variables, NULL))
       || (-1 == SLns_add_iconstant_table (ns, Module_IConstants, NULL))
      )
     return -1;

   return 0;
}

/* This function is optional */
void deinit_gslmatrix_module (void)
{
}
