/*
  array.c
  Ruby/GSL: Ruby extension library for GSL (GNU Scientific Library)
    (C) Copyright 2001-2004 by Yoshiki Tsunesada

  Ruby/GSL is free software: you can redistribute it and/or modify it
  under the terms of the GNU General Public License.
  This library is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY.
*/

#include "rb_gsl_common.h"
#include "rb_gsl_array.h"
#include "rb_gsl_complex.h"
#ifdef HAVE_NARRAY_H
#include "rb_gsl_with_narray.h"
#endif

/* global variables */
VALUE cgsl_block;
VALUE cgsl_vector, cgsl_vector_view;
VALUE cgsl_vector_view_ro;
VALUE cgsl_matrix_view_ro;
VALUE cgsl_vector_complex;
VALUE cgsl_vector_complex_view;
VALUE cgsl_matrix;
VALUE cgsl_matrix_view;
VALUE cgsl_matrix_complex;
VALUE cgsl_matrix_complex_view;
VALUE cgsl_vector_complex_view_ro;
VALUE cgsl_matrix_complex_view_ro;
VALUE cgsl_permutation;
VALUE cgsl_range, cgsl_linspace, cgsl_logspace;
VALUE cgsl_vector_col, cgsl_vector_col_view;
VALUE cgsl_vector_complex_col, cgsl_vector_complex_col_view;

double* get_vector_ptr(VALUE ary, size_t *stride, size_t *n)
{
  gsl_vector *v = NULL;
#ifdef HAVE_NARRAY_H
  VALUE ary2;
#endif
  if (rb_obj_is_kind_of(ary, cgsl_vector)) {
    Data_Get_Struct(ary, gsl_vector, v);
    *stride = v->stride;
    *n = v->size;
    return v->data;
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(ary)) {
    *n = NA_TOTAL(ary);
    *stride = 1;
    ary2 = na_change_type(ary, NA_DFLOAT);
    return NA_PTR_TYPE(ary2,double*);
#endif
  } else {
    rb_raise(rb_eTypeError,
	     "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  }
}

gsl_vector* get_cvector(VALUE obj)
{
  gsl_vector *v = NULL;
  if (rb_obj_is_kind_of(obj, cgsl_vector)) {
    Data_Get_Struct(obj, gsl_vector,  v);
#ifdef HAVE_NARRAY_H
  } else if (NA_IsArray(obj)) {
    v = make_cvector_from_rarrays(obj);
#endif
  } else {
    rb_raise(rb_eTypeError,
             "wrong argument type %s", rb_class2name(CLASS_OF(obj)));
  }
  return v; 
}

VALUE make_rarray_from_cvector(gsl_vector *v)
{
  int i;
  VALUE ary;
  ary = rb_ary_new2(v->size);
  for (i = 0; i < v->size; i++) {
    rb_ary_store(ary, i, rb_float_new(gsl_vector_get(v, i)));
  }
  return ary;
}

VALUE make_rarray_from_cpermutation(gsl_permutation *v)
{
  int i;
  VALUE ary;
  ary = rb_ary_new2(v->size);
  for (i = 0; i < v->size; i++) {
    rb_ary_store(ary, i, rb_float_new(gsl_permutation_get(v, i)));
  }
  return ary;
}

gsl_vector* get_vector(VALUE ary)
{
  gsl_vector *v = NULL;
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  if (TYPE(ary) == T_ARRAY) {
    return make_cvector_from_rarray(ary);
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(ary)) {
    return make_cvector_from_narray(ary);
#endif
  } else if (VECTOR_P(ary)) {
    Data_Get_Struct(ary, gsl_vector, v);
    return v;
  } else {
    rb_raise(rb_eTypeError,
	     "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  }
}

gsl_vector* make_cvector_from_rarrays(VALUE ary)
{
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  if (TYPE(ary) == T_ARRAY) {
    return make_cvector_from_rarray(ary);
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(ary)) {
    return make_cvector_from_narray(ary);
#endif
  } else {
    rb_raise(rb_eTypeError,
	     "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  }
}

gsl_vector* make_cvector_from_rarray(VALUE ary)
{
  gsl_vector *v = NULL;
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  Check_Type(ary, T_ARRAY);
  v = gsl_vector_alloc(RARRAY(ary)->len);
  if (v == NULL) rb_raise(rb_eNoMemError, "gsl_vector_alloc failed");
  cvector_set_from_rarray(v, ary);
  return v;
}

void cvector_set_from_carray(gsl_vector *v, const double *a)
{
  int i;
  for (i = 0; i < v->size; i++) gsl_vector_set(v, i, a[i]);
}

void cvector_set_from_rarrays(gsl_vector *v, VALUE ary)
{
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  if (TYPE(ary) == T_ARRAY) {
    return cvector_set_from_rarray(v, ary);
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(ary)) {
    return cvector_set_from_narray(v, ary);
#endif
  } else {
    rb_raise(rb_eTypeError,
	     "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  }
}

void cvector_set_from_rarray(gsl_vector *v, VALUE ary)
{
  int i;
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  Check_Type(ary, T_ARRAY);
  if (RARRAY(ary)->len == 0) return;
  for (i = 0; i < v->size; i++) gsl_vector_set(v, i, NUM2DBL(rb_ary_entry(ary, i)));
}

void carray_set_from_cvector(double *a, gsl_vector *v)
{
  int i;
  for (i = 0; i < v->size; i++) a[i] = gsl_vector_get(v, i);
}

void carray_set_from_rarrays(double *a, VALUE ary)
{
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  if (TYPE(ary) == T_ARRAY) {
    return carray_set_from_rarray(a, ary);
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(ary)) {
    return carray_set_from_narray(a, ary);
#endif
  } else {
    rb_raise(rb_eTypeError,
	     "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  }
}

void carray_set_from_rarray(double *a, VALUE ary)
{
  int i, size;
  VALUE val;
  if (CLASS_OF(ary) == rb_cRange) ary = rb_gsl_range2ary(ary);
  Check_Type(ary, T_ARRAY);
  size = RARRAY(ary)->len;
  if (size == 0) return;
  for (i = 0; i < size; i++) {
    val = rb_ary_entry(ary, i);
    Need_Float(val);
    a[i] = RFLOAT(val)->value;
  }
}

#ifdef HAVE_NARRAY_H
/* NArray -> CArray */
void carray_set_from_narray(double *a, VALUE ary)
{
  int size;
  VALUE ary2;
  if (!NA_IsNArray(ary))
    rb_raise(rb_eTypeError,
             "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  size = NA_TOTAL(ary);
  if (size == 0) return;
  ary2 = na_change_type(ary, NA_DFLOAT);
  memcpy(a, NA_PTR_TYPE(ary2,double*), size*sizeof(double));
}

/* NArray -> GSL::Vector */
gsl_vector* make_cvector_from_narray(VALUE ary)
{
  gsl_vector *v = NULL;
  size_t size;
  VALUE ary2;
  if (!NA_IsNArray(ary))
    rb_raise(rb_eTypeError,
             "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  size = NA_TOTAL(ary);
  v = gsl_vector_alloc(size);
  if (v == NULL) rb_raise(rb_eNoMemError, "gsl_vector_alloc failed");
  ary2 = na_change_type(ary, NA_DFLOAT);
  memcpy(v->data, NA_PTR_TYPE(ary2,double*), size*sizeof(double));
  /*  cvector_set_from_narray(v, ary);*/
  return v;
}

void cvector_set_from_narray(gsl_vector *v, VALUE ary)
{
  int size;
  if (!NA_IsNArray(ary))
    rb_raise(rb_eTypeError,
             "wrong argument type %s", rb_class2name(CLASS_OF(ary)));
  size = NA_TOTAL(ary);
  carray_set_from_narray(v->data, ary);
}
#endif

gsl_vector* make_vector_clone(gsl_vector *v)
{
  gsl_vector *vnew = NULL;
  vnew = gsl_vector_alloc(v->size);
  if (vnew == NULL) rb_raise(rb_eNoMemError, "gsl_vector_alloc failed");
  gsl_vector_memcpy(vnew, v);
  return vnew;
}

gsl_vector_complex* make_vector_complex_clone(gsl_vector_complex *v)
{
  gsl_vector_complex *vnew = NULL;
  vnew = gsl_vector_complex_alloc(v->size);
  if (vnew == NULL) rb_raise(rb_eNoMemError, "gsl_vector_complex_alloc failed");
  gsl_vector_complex_memcpy(vnew, v);
  return vnew;
}

gsl_matrix* make_matrix_clone(gsl_matrix *m)
{
  gsl_matrix *mnew = NULL;
  mnew = gsl_matrix_alloc(m->size1, m->size2);
  if (mnew == NULL) rb_raise(rb_eNoMemError, "gsl_matrix_alloc failed");
  gsl_matrix_memcpy(mnew, m);
  return mnew;
}

VALUE make_matrix_clone2(VALUE vm)
{
  gsl_matrix *m = NULL, *mnew = NULL;
  Data_Get_Struct(vm, gsl_matrix, m);
  mnew = gsl_matrix_alloc(m->size1, m->size2);
  if (mnew == NULL) rb_raise(rb_eNoMemError, "gsl_matrix_alloc failed");
  gsl_matrix_memcpy(mnew, m);
  return Data_Wrap_Struct(cgsl_matrix, 0, gsl_matrix_free, mnew);;
}

gsl_matrix_complex* make_matrix_complex_clone(gsl_matrix_complex *m)
{
  gsl_matrix_complex *mnew = NULL;
  mnew = gsl_matrix_complex_alloc(m->size1, m->size2);
  if (mnew == NULL) rb_raise(rb_eNoMemError, "gsl_matrix_complex_alloc failed");
  gsl_matrix_complex_memcpy(mnew, m);
  return mnew;
}

gsl_vector_complex* vector_to_complex(gsl_vector *v)
{
  gsl_vector_complex *cv = NULL;
  gsl_complex z;
  int i;
  cv = gsl_vector_complex_alloc(v->size);
  if (cv == NULL) rb_raise(rb_eNoMemError, "gsl_vector_complex_alloc failed");
  for (i = 0; i < v->size; i++) {
    z = gsl_complex_rect(gsl_vector_get(v, i), 0.0);
    gsl_vector_complex_set(cv, i, z);
  }
  return cv;
}

gsl_matrix_complex* matrix_to_complex(gsl_matrix *m)
{
  gsl_matrix_complex *cm = NULL;
  gsl_complex z;
  size_t i, j;
  cm = gsl_matrix_complex_alloc(m->size1, m->size2);
  if (cm == NULL) rb_raise(rb_eNoMemError, "gsl_matrix_complex_alloc failed");
  for (i = 0; i < m->size1; i++) {
    for (j = 0; j < m->size2; j++) {
      z = gsl_complex_rect(gsl_matrix_get(m, i, j), 0.0);
      gsl_matrix_complex_set(cm, i, j, z);
    }
  }
  return cm;
}

void gsl_matrix_complex_mul(gsl_matrix_complex *mnew, gsl_matrix_complex *m, 
			    gsl_matrix_complex *mb)
{
  gsl_complex a, b, c, sum;
  int i, j, k;
  for (i = 0; i < m->size1; i++) {
    for (j = 0; j < m->size2; j++) {
      sum = gsl_complex_rect(0.0, 0.0);
      for (k = 0; k < m->size2; k++) {
	a = gsl_matrix_complex_get(m, j, k);
	b = gsl_matrix_complex_get(mb, k, i);
	c = gsl_complex_mul(a, b);
	sum = gsl_complex_add(sum, c);
      }
      gsl_matrix_complex_set(mnew, j, i, sum);
    }
  }
}
/*
void gsl_matrix_mul(gsl_matrix *mnew, gsl_matrix *m, gsl_matrix *b)
{
  double val;
  int i, j, k;
  for (i = 0; i < m->size1; i++) {
    for (j = 0; j < m->size2; j++) {
      val = 0;
      for (k = 0; k < m->size2; k++) {
	val += gsl_matrix_get(m, j, k)*gsl_matrix_get(b, k, i);
      }
      gsl_matrix_set(mnew, j, i, val);
    }
  }
}
*/
void gsl_matrix_mul_vector(gsl_vector *vnew, gsl_matrix *m, gsl_vector *v)
{
  int i, j;
  double val;
  for (i = 0; i < m->size1; i++) {
    val = 0;
    for (j = 0; j < m->size2; j++) 
      val += gsl_matrix_get(m, i, j)*gsl_vector_get(v, j);
    gsl_vector_set(vnew, i, val);
  }
}

void gsl_matrix_complex_mul_vector(gsl_vector_complex *vnew, gsl_matrix_complex *m, 
				   gsl_vector_complex *v)
{
  gsl_complex a, b, c, sum;
  int i, j;
  for (i = 0; i < m->size1; i++) {
    sum = gsl_complex_rect(0.0, 0.0);
    for (j = 0; j < m->size2; j++) {
      a = gsl_matrix_complex_get(m, i, j);
      b = gsl_vector_complex_get(v, j);
      c = gsl_complex_mul(a, b);
      sum = gsl_complex_add(sum, c);
    }
    gsl_vector_complex_set(vnew, i, sum);
  }
  
}

/*****/

int gsl_vector_complex_add(gsl_vector_complex *cv, const gsl_vector_complex *cv2)
{
  int i;
  gsl_complex a, b, c;
  for (i = 0; i < cv->size; i++) {
    a = gsl_vector_complex_get(cv, i);
    b = gsl_vector_complex_get(cv2, i);
    c = gsl_complex_add(a, b);
    gsl_vector_complex_set(cv, i, c);
  }
  return 0;
}

int gsl_vector_complex_add_constant(gsl_vector_complex *cv, gsl_complex b)
{
  int i;
  gsl_complex a, c;
  for (i = 0; i < cv->size; i++) {
    a = gsl_vector_complex_get(cv, i);
    c = gsl_complex_add(a, b);
    gsl_vector_complex_set(cv, i, c);
  }
  return 0;
}

int gsl_vector_complex_scale(gsl_vector_complex *cv, gsl_complex b)
{
  int i;
  gsl_complex a, c;
  for (i = 0; i < cv->size; i++) {
    a = gsl_vector_complex_get(cv, i);
    c = gsl_complex_mul(a, b);
    gsl_vector_complex_set(cv, i, c);
  }
  return 0;
}


int gsl_vector_complex_sub(gsl_vector_complex *cv, const gsl_vector_complex *cv2)
{
  int i;
  gsl_complex a, b, c;
  for (i = 0; i < cv->size; i++) {
    a = gsl_vector_complex_get(cv, i);
    b = gsl_vector_complex_get(cv2, i);
    c = gsl_complex_sub(a, b);
    gsl_vector_complex_set(cv, i, c);
  }
  return 0;
}

int gsl_vector_complex_mul(gsl_vector_complex *cv, const gsl_vector_complex *cv2)
{
  int i;
  gsl_complex a, b, c;
  for (i = 0; i < cv->size; i++) {
    a = gsl_vector_complex_get(cv, i);
    b = gsl_vector_complex_get(cv2, i);
    c = gsl_complex_mul(a, b);
    gsl_vector_complex_set(cv, i, c);
  }
  return 0;
}

int gsl_vector_complex_div(gsl_vector_complex *cv, const gsl_vector_complex *cv2)
{
  int i;
  gsl_complex a, b, c;
  for (i = 0; i < cv->size; i++) {
    a = gsl_vector_complex_get(cv, i);
    b = gsl_vector_complex_get(cv2, i);
    c = gsl_complex_div(a, b);
    gsl_vector_complex_set(cv, i, c);
  }
  return 0;
}

gsl_vector_view* rb_gsl_make_vector_view(double *data, size_t size, size_t stride)
{
  gsl_vector_view *v = NULL;
  v = ALLOC(gsl_vector_view);
  v->vector.size = size;
  v->vector.stride = stride;
  v->vector.owner = 0;
  v->vector.data = data;
  return v;
}

VALUE rb_gsl_range2ary(VALUE obj)
{
  int beg, en, n, i;
  VALUE ary;
  if (CLASS_OF(obj) != rb_cRange) 
    rb_raise(rb_eTypeError, "wrong argument type %s (Range expected)",
	     rb_class2name(CLASS_OF(obj)));
  beg = NUM2INT(rb_ivar_get(obj, rb_gsl_id_beg));
  en = NUM2INT(rb_ivar_get(obj, rb_gsl_id_end));
  if (RTEST(rb_ivar_get(obj, rb_gsl_id_excl)))
    n = en - beg;
  else
    n = en - beg + 1;
  ary = rb_ary_new2(n);
  for (i = 0; i < n; i++) rb_ary_store(ary, i, INT2FIX(beg+i));
  return ary;
}

void Init_gsl_array()
{
  cgsl_block = rb_define_class_under(mgsl, "Block", 
					 rb_cObject);
  cgsl_vector = rb_define_class_under(mgsl, "Vector", 
				      rb_cObject);
  cgsl_vector_col = rb_define_class_under(cgsl_vector, "Col", 
					 cgsl_vector);
  cgsl_vector_complex = rb_define_class_under(cgsl_vector, "Complex", 
					      rb_cObject);
  cgsl_vector_complex_col = rb_define_class_under(cgsl_vector_complex, "Col",
					      cgsl_vector_complex);
  cgsl_matrix = rb_define_class_under(mgsl, "Matrix", rb_cObject);
  cgsl_matrix_complex = rb_define_class_under(cgsl_matrix, "Complex", rb_cObject);

  cgsl_vector_view = rb_define_class_under(cgsl_vector, "View", cgsl_vector);
  cgsl_vector_col_view = rb_define_class_under(cgsl_vector_col, "View", cgsl_vector_col);

  cgsl_vector_complex_view = rb_define_class_under(cgsl_vector_complex, "View", 
						   cgsl_vector_complex);
  cgsl_vector_complex_col_view = rb_define_class_under(cgsl_vector_complex_col, "View", 
						   cgsl_vector_complex_col);

  cgsl_matrix_view = rb_define_class_under(cgsl_matrix, "View", 
					   cgsl_matrix);
  cgsl_matrix_complex_view = rb_define_class_under(cgsl_matrix_complex, "View", 
						   cgsl_matrix_complex);
  cgsl_permutation = rb_define_class_under(mgsl, "Permutation", rb_cObject);  

  cgsl_range = rb_define_class_under(mgsl, "Range", cgsl_vector);
  cgsl_linspace = rb_define_class_under(mgsl, "Linspace", cgsl_range);
  cgsl_logspace = rb_define_class_under(mgsl, "Logspace", cgsl_range);

  cgsl_vector_view_ro = rb_define_class_under(cgsl_vector_view, "ReadOnly",
					      cgsl_vector_view);
  cgsl_matrix_view_ro = rb_define_class_under(cgsl_matrix_view, "ReadOnly",
					      cgsl_matrix_view);

  cgsl_vector_complex_view_ro = rb_define_class_under(cgsl_vector_complex_view, 
						      "ReadOnly",
						      cgsl_vector_complex_view);
  cgsl_matrix_complex_view_ro = rb_define_class_under(cgsl_matrix_complex_view, 
						      "ReadOnly",
						      cgsl_matrix_complex_view);


  Init_gsl_vector();
  Init_gsl_range();
  Init_gsl_vector_complex();
  Init_gsl_matrix();
  Init_gsl_matrix_complex();
  Init_gsl_permutation();
  Init_gsl_combination();

  Init_gsl_array_complex();

  rb_define_method(cgsl_vector_view_ro, "set", rb_gsl_obj_read_only, -1);
  rb_define_method(cgsl_matrix_view_ro, "set", rb_gsl_obj_read_only, -1);
  rb_define_method(cgsl_vector_complex_view_ro, "set", rb_gsl_obj_read_only, -1);
  rb_define_method(cgsl_matrix_complex_view_ro, "set", rb_gsl_obj_read_only, -1);
}
