/*
  Equality of self and other in view of numerical array.
  i.e., both arrays have same shape and corresponding elements are equal.
  @overload == other
  @param [Object] other
  @return [Boolean] true if self and other is equal.
*/
VALUE
na_equal(VALUE self, volatile VALUE other)
{
    volatile VALUE vbool;
    narray_t *na1, *na2;
    int i;

    GetNArray(self,na1);

    if (!rb_obj_is_kind_of(other,cNArray)) {
        other = rb_funcall(CLASS_OF(self), rb_intern("cast"), 1, other);
    }

    GetNArray(other,na2);
    if (na1->ndim != na2->ndim) {
        return Qfalse;
    }
    for (i=0; i<na1->ndim; i++) {
        if (na1->shape[i] != na2->shape[i]) {
            return Qfalse;
        }
    }
    vbool = rb_funcall(self, rb_intern("eq"), 1, other);
    return (rb_funcall(vbool, rb_intern("count_false"), 0)==INT2FIX(0)) ? Qtrue : Qfalse;
}

/* @overload initialize( inputs, targets )
 * Creates a new dataset from example data.
 * @param [NArray<sfloat>] inputs the input examples that a nn_model will process
 * @param [NArray<sfloat>] targets known outputs that can be used to train or assess a nn_model
 * @return [RuNeNe::DataSet] new dataset
 */
VALUE dataset_class_initialize( VALUE self, VALUE rv_inputs, VALUE rv_targets ) {
  volatile VALUE val_inputs;
  volatile VALUE val_targets;
  struct NARRAY *na_inputs;
  struct NARRAY *na_targets;
  DataSet *dataset = get_dataset_struct( self );

  val_inputs = na_cast_object( rv_inputs, NA_SFLOAT );
  GetNArray( val_inputs, na_inputs );

  val_targets = na_cast_object( rv_targets, NA_SFLOAT );
  GetNArray( val_targets, na_targets );

  if ( na_inputs->rank < 2 ) {
    rb_raise( rb_eArgError, "Inputs rank should be at least 2, but got %d", na_inputs->rank );
  }

  if ( na_targets->rank < 2 ) {
    rb_raise( rb_eArgError, "Targets rank should be at least 2, but got %d", na_targets->rank );
  }

  if ( na_inputs->shape[ na_inputs->rank - 1 ] != na_targets->shape[ na_targets->rank - 1 ] ) {
    rb_raise( rb_eArgError, "Number of input items %d not same as target items %d",
        na_inputs->shape[ na_inputs->rank - 1 ], na_targets->shape[ na_targets->rank - 1 ] );
  }

  dataset__init_from_narray( dataset, val_inputs, val_targets );

  return self;
}

/* @overload run( input )
 * Runs nn_model forward and generates a result
 * @param [NArray<sfloat>] input single input vector
 * @return [NArray<sfloat>] output of nn_model
 */
VALUE nn_model_rbobject__run( VALUE self, VALUE rv_input ) {
  NNModel *nn_model = get_nn_model_struct( self );
  int out_shape[1] = { nn_model->num_outputs };

  struct NARRAY *na_input;
  volatile VALUE val_input = na_cast_object(rv_input, NA_SFLOAT);
  GetNArray( val_input, na_input );

  // Shouldn't happen, but we don't want a segfault
  if ( nn_model->num_layers < 1 ) {
    return Qnil;
  }

  if ( na_input->total != nn_model->num_inputs ) {
    rb_raise( rb_eArgError, "Input array must be size %d, but it was size %d", nn_model->num_inputs, na_input->total );
  }

  struct NARRAY *na_output;

  volatile VALUE val_output = na_make_object( NA_SFLOAT, 1, out_shape, cNArray );
  GetNArray( val_output, na_output );

  nn_model__run( nn_model, (float*) na_input->ptr );

  memcpy( (float*) na_output->ptr, nn_model->activations[nn_model->num_layers-1], nn_model->num_outputs * sizeof(float) );

  return val_output;
}

static VALUE rb_gsl_spline_init(VALUE obj, VALUE xxa, VALUE yya)
{
  rb_gsl_spline *sp = NULL;
  gsl_spline *p = NULL;
  gsl_vector *xa = NULL, *ya = NULL;
  size_t i, size;
  int flagx = 0, flagy = 0;
  double *ptr1 = NULL, *ptr2 = NULL;
#ifdef HAVE_NARRAY_H
  struct NARRAY *nax = NULL, *nay = NULL;
#endif
  Data_Get_Struct(obj, rb_gsl_spline, sp);
  p = sp->s;
  if (TYPE(xxa) == T_ARRAY) {
    size = RARRAY_LEN(xxa);
    xa = gsl_vector_alloc(size);
    for (i = 0; i < size; i++) gsl_vector_set(xa, i, NUM2DBL(rb_ary_entry(xxa, i)));
    ptr1 = xa->data;
    flagx = 1;
  } else if (VECTOR_P(xxa)) {
    Data_Get_Struct(xxa, gsl_vector, xa);
    size = xa->size;
    ptr1 = xa->data;
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(xxa)) {
      GetNArray(xxa, nax);
      size = nax->total;
      ptr1 = (double *) nax->ptr;
#endif
  } else {
    rb_raise(rb_eTypeError, "not a vector");
  }
  if (TYPE(yya) == T_ARRAY) {
    ya = gsl_vector_alloc(size);
    for (i = 0; i < size; i++) gsl_vector_set(ya, i, NUM2DBL(rb_ary_entry(yya, i)));
    ptr2 = ya->data;
    flagy = 1;
#ifdef HAVE_NARRAY_H
  } else if (NA_IsNArray(yya)) {
      GetNArray(yya, nay);
      ptr2 = (double *) nay->ptr;
#endif
  } else if (VECTOR_P(yya)) {
    Data_Get_Struct(yya, gsl_vector, ya);
    ptr2 = ya->data;
  } else {
    rb_raise(rb_eTypeError, "not a vector");
  }
  gsl_spline_init(p, ptr1, ptr2, size);
  if (flagx == 1) gsl_vector_free(xa);
  if (flagy == 1) gsl_vector_free(ya);
  return obj;
}

void
na_copy_flags(VALUE src, VALUE dst)
{
    narray_t *na1, *na2;

    GetNArray(src,na1);
    GetNArray(dst,na2);

    na2->flag[0] = na1->flag[0];
    na2->flag[1] = na1->flag[1];

    RBASIC(dst)->flags |= (RBASIC(src)->flags) &
        (FL_USER1|FL_USER2|FL_USER3|FL_USER4|FL_USER5|FL_USER6|FL_USER7);
}

void mbgd_layer__deep_copy( MBGDLayer *mbgd_layer_copy, MBGDLayer *mbgd_layer_orig ) {
    struct NARRAY *narr;
    GradientDescent_SGD * gd_sgd;
    GradientDescent_NAG * gd_nag;
    GradientDescent_RMSProp * gd_rmsprop;

    mbgd_layer_copy->num_inputs = mbgd_layer_orig->num_inputs;
    mbgd_layer_copy->num_outputs = mbgd_layer_orig->num_outputs;
    mbgd_layer_copy->learning_rate = mbgd_layer_orig->learning_rate;
    mbgd_layer_copy->gradient_descent_type = mbgd_layer_orig->gradient_descent_type;

    switch ( mbgd_layer_copy->gradient_descent_type ) {
    case GD_TYPE_SGD:
        Data_Get_Struct( mbgd_layer_orig->gradient_descent, GradientDescent_SGD, gd_sgd );
        mbgd_layer_copy->gradient_descent = Data_Wrap_Struct( RuNeNe_GradientDescent_SGD,
                                            gd_sgd__gc_mark, gd_sgd__destroy, gd_sgd__clone( gd_sgd ) );
        break;

    case GD_TYPE_NAG:
        Data_Get_Struct( mbgd_layer_orig->gradient_descent, GradientDescent_NAG, gd_nag );
        mbgd_layer_copy->gradient_descent = Data_Wrap_Struct( RuNeNe_GradientDescent_NAG,
                                            gd_nag__gc_mark, gd_nag__destroy, gd_nag__clone( gd_nag ) );
        break;

    case GD_TYPE_RMSPROP:
        Data_Get_Struct( mbgd_layer_orig->gradient_descent, GradientDescent_RMSProp, gd_rmsprop );
        mbgd_layer_copy->gradient_descent = Data_Wrap_Struct( RuNeNe_GradientDescent_RMSProp,
                                            gd_rmsprop__gc_mark, gd_rmsprop__destroy, gd_rmsprop__clone( gd_rmsprop ) );
        break;
    }

    mbgd_layer_copy->max_norm = mbgd_layer_orig->max_norm;
    mbgd_layer_copy->weight_decay = mbgd_layer_orig->weight_decay;

    mbgd_layer_copy->narr_de_dz = na_clone( mbgd_layer_orig->narr_de_dz );
    GetNArray( mbgd_layer_copy->narr_de_dz, narr );
    mbgd_layer_copy->de_dz = (float *) narr->ptr;

    mbgd_layer_copy->narr_de_da = na_clone( mbgd_layer_orig->narr_de_da );
    GetNArray( mbgd_layer_copy->narr_de_da, narr );
    mbgd_layer_copy->de_da = (float *) narr->ptr;

    mbgd_layer_copy->narr_de_dw = na_clone( mbgd_layer_orig->narr_de_dw );
    GetNArray( mbgd_layer_copy->narr_de_dw, narr );
    mbgd_layer_copy->de_dw = (float *) narr->ptr;

    return;
}

//
// NArray#load_from_shmemutex(shmemutex,timeout=-1)
//
VALUE rb_NArray_load_from_shmemutex(int argc,VALUE *argv,VALUE self)
{
    VALUE shm,timeout;
    struct NARRAY *n_na;
    ShMemutex *ptr;
    long to=-1;
    size_t sz;

    if(rb_scan_args(argc,argv,"11",&shm,&timeout)==2)
        to=NUM2LONG(timeout);

    if(rb_obj_is_kind_of(shm,
                         rb_const_get(rb_cObject,
                                      rb_intern("ShMemutex")))!=Qtrue)
        rb_raise(rb_eTypeError,"1st. argument must be ShMemutex object.");
    Data_Get_Struct(shm,ShMemutex,ptr);
    

    GetNArray(self,n_na);
    sz=n_na->total*na_sizeof[n_na->type];
    if(!sz)
        rb_raise(rb_eTypeError,"NArray size is too small");

    ptr->read(n_na->ptr,sz,to);

    return self;
}

//
// ShMemutex::write(src,timeout=-1)
//
VALUE rb_ShMemutex_write(int argc,VALUE *argv,VALUE self)
{
    VALUE src,timeout;
    struct NARRAY *n_na;
    ShMemutex *ptr;
    long to=-1;
    void *buf;
    size_t sz;

    if(rb_scan_args(argc,argv,"11",&src,&timeout)==2)
        to=NUM2LONG(timeout);

    Data_Get_Struct(self,ShMemutex,ptr);

    if(IsNArray(src)){
        GetNArray(src,n_na);

        sz=n_na->total*na_sizeof[n_na->type];
        buf=(void *)n_na->ptr;
    }
    else{
        StringValue(src);
        sz=RSTRING(src)->len+1;
        buf=RSTRING(src)->ptr;
    }

    if((!sz)||(!buf))
        return INT2FIX(0);
    
    return INT2NUM(ptr->write(buf,sz,to));
}

/* @overload init_weights( mult = 1.0 )
 * Initialises weights in all layers.
 * @param [Float] mult optional size factor
 * @return [RuNeNe::NNModel] self
 */
VALUE nn_model_rbobject__init_weights( int argc, VALUE* argv, VALUE self ) {
  NNModel *nn_model = get_nn_model_struct( self );
  VALUE rv_mult;
  Layer_FF *layer_ff;
  float m = 1.0;
  int i, j, t;
  struct NARRAY *narr;

  rb_scan_args( argc, argv, "01", &rv_mult );
  if ( ! NIL_P( rv_mult ) ) {
    m = NUM2FLT( rv_mult );
  }

  for ( i = 0; i < nn_model->num_layers; i++ ) {
    // TODO: This only works for Layer_FF layers, we need a more flexible system
    Data_Get_Struct( nn_model->layers[i], Layer_FF, layer_ff );

    layer_ff__init_weights( layer_ff );

    if ( m != 0 ) {
      GetNArray( layer_ff->narr_weights, narr );
      t = narr->total;
      for ( j = 0; j < t; j++ ) {
        layer_ff->weights[j] *= m;
      }
    }
  }

  return self;
}

static int
na_get_result_dimension_for_slice(VALUE self, int argc, VALUE *argv)
{
    int i;
    int count_new=0;
    int count_rest=0;
    narray_t *na;
    VALUE a;

    GetNArray(self,na);
    if (na->size == 0) {
        rb_raise(nary_eShapeError, "cannot get element of empty array");
    }
    for (i=0; i<argc; i++) {
        a = argv[i];
        switch(TYPE(a)) {
        case T_FALSE:
        case T_SYMBOL:
            if (a==sym_rest || a==sym_tilde || a==Qfalse) {
                argv[i] = Qfalse;
                count_rest++;
            } else if (a==sym_new || a==sym_minus) {
                argv[i] = sym_new;
                count_new++;
            }
        }
    }

    return check_index_count(argc, na->ndim, count_new, count_rest);
}

VALUE
rb_narray_debug_info(VALUE self)
{
    int i;
    narray_t *na;
    GetNArray(self,na);

    printf("%s:\n",rb_class2name(CLASS_OF(self)));
    printf("  id     = 0x%"SZF"x\n", self);
    printf("  type   = %d\n", na->type);
    printf("  flag   = [%d,%d]\n", na->flag[0], na->flag[1]);
    printf("  size   = %"SZF"d\n", na->size);
    printf("  ndim   = %d\n", na->ndim);
    printf("  shape  = 0x%"SZF"x\n", (size_t)na->shape);
    if (na->shape) {
        printf("  shape  = [");
        for (i=0;i<na->ndim;i++)
            printf(" %"SZF"d", na->shape[i]);
        printf(" ]\n");
    }

    switch(na->type) {
    case NARRAY_DATA_T:
    case NARRAY_FILEMAP_T:
        rb_narray_debug_info_nadata(self);
        break;
    case NARRAY_VIEW_T:
        rb_narray_debug_info_naview(self);
        break;
    }
    return Qnil;
}

void
na_setup(VALUE self, int ndim, size_t *shape)
{
    narray_t *na;
    GetNArray(self,na);
    na_setup_shape(na, ndim, shape);
}

static VALUE
nary_struct_cast_array(VALUE klass, VALUE rary)
{
    volatile VALUE vnc, nary;
    narray_t *na;
    na_compose_t *nc;
    VALUE opt;
    ndfunc_arg_in_t ain[3] = {{rb_cArray,0},{Qnil,0},{sym_option}};
    ndfunc_t ndf = { iter_nstruct_from_a, NO_LOOP, 3, 0, ain, 0 };

    //fprintf(stderr,"rary:");rb_p(rary);
    //fprintf(stderr,"class_of(rary):");rb_p(CLASS_OF(rary));

    vnc = na_ary_composition_for_struct(klass, rary);
    Data_Get_Struct(vnc, na_compose_t, nc);
    nary = rb_narray_new(klass, nc->ndim, nc->shape);
    GetNArray(nary,na);
    //fprintf(stderr,"na->size=%lu\n",na->size);
    //fprintf(stderr,"na->ndim=%d\n",na->ndim);
    if (na->size>0) {
        opt = nst_create_member_views(nary);
        rb_funcall(nary, rb_intern("allocate"), 0);
        na_ndloop_cast_rarray_to_narray2(&ndf, rary, nary, opt);
    }
    return nary;
}

static VALUE
nst_allocate(VALUE self)
{
    narray_t *na;
    char *ptr;
    VALUE velmsz;

    GetNArray(self,na);

    switch(NA_TYPE(na)) {
    case NARRAY_DATA_T:
        ptr = NA_DATA_PTR(na);
        if (na->size > 0 && ptr == NULL) {
            velmsz = rb_const_get(CLASS_OF(self), rb_intern("element_byte_size"));
            ptr = xmalloc(NUM2SIZET(velmsz) * na->size);
            NA_DATA_PTR(na) = ptr;
        }
        break;
    case NARRAY_VIEW_T:
        rb_funcall(NA_VIEW_DATA(na), rb_intern("allocate"), 0);
        break;
    case NARRAY_FILEMAP_T:
        //ptr = ((narray_filemap_t*)na)->ptr;
        // to be implemented
    default:
        rb_bug("invalid narray type : %d",NA_TYPE(na));
    }
    return self;
}

static size_t
check_array_1d(VALUE item, size_t size) {
    narray_t *na;
    size_t i, len;

    if (TYPE(item) == T_ARRAY) {
        len = RARRAY_LEN(item);
        if (size != len) {
            return 0;
        }
        for (i=0; i<len; i++) {
            if (!check_array(RARRAY_AREF(item,i))) {
                return 0;
            }
        }
        return 1;
    }
    if (RTEST(rb_obj_is_kind_of(item, cNArray))) {
        GetNArray(item,na);
        if (na->ndim == 1 && na->size == size) {
            return 1;
        } else {
            return 0;
        }
    }
    return 0;
}

/* method: size() -- returns the total number of typeents */
static VALUE
na_ndim(VALUE self)
{
    narray_t *na;
    GetNArray(self,na);
    return INT2NUM(na->ndim);
}

void
na_array_to_internal_shape(VALUE self, VALUE ary, size_t *shape)
{
    size_t    i, n, c, s;
    VALUE     v;
    narray_t *na;
    int       flag = 0;

    n = RARRAY_LEN(ary);

    if (RTEST(self)) {
        GetNArray(self, na);
        flag = TEST_COLUMN_MAJOR(na);
    }
    if (flag) {
        c = n-1;
        s = -1;
    } else {
        c = 0;
        s = 1;
    }
    for (i=0; i<n; i++) {
        v = RARRAY_AREF(ary,i);
        if (!FIXNUM_P(v) && !rb_obj_is_kind_of(v, rb_cInteger)) {
            rb_raise(rb_eTypeError, "array size must be Integer");
        }
        if (RTEST(rb_funcall(v, rb_intern("<"), 1, INT2FIX(0)))) {
            rb_raise(rb_eArgError,"size must be non-negative");
        }
        shape[c] = NUM2SIZE(v);
        c += s;
    }
}

char *
na_get_pointer_for_read(VALUE self)
{
    char  *ptr;
    narray_t *na;
    GetNArray(self,na);

    //if (NA_TEST_LOCK(na)) {
    //    rb_raise(rb_eRuntimeError, "cannot read locked NArray.");
    //}

    if (NA_TYPE(na) == NARRAY_DATA_T) {
        ptr = NA_DATA_PTR(na);
    } else {
        ptr = na_get_pointer(self);
    }

    if (NA_SIZE(na) > 0 && ptr == NULL) {
        rb_raise(rb_eRuntimeError,"cannot read unallocated NArray");
    }

    //NA_SET_LOCK(na);

    return ptr;
}

/* method: size() -- returns the total number of typeents */
static VALUE
na_size(VALUE self)
{
    narray_t *na;
    GetNArray(self,na);
    return SIZE2NUM(na->size);
}

char *
na_get_pointer_for_write(VALUE self)
{
    char *ptr;
    narray_t *na;
    GetNArray(self,na);

    if (OBJ_FROZEN(self)) {
        rb_raise(rb_eRuntimeError, "cannot write to frozen NArray.");
    }

    if (NA_TYPE(na) == NARRAY_DATA_T) {
        ptr = NA_DATA_PTR(na);
        if (na->size > 0 && ptr == NULL) {
            rb_funcall(self, id_allocate, 0);
            ptr = NA_DATA_PTR(na);
        }
    } else {
        ptr = na_get_pointer(self);
        if (NA_SIZE(na) > 0 && ptr == NULL) {
            rb_raise(rb_eRuntimeError,"cannot write to unallocated NArray");
        }
    }

    //NA_SET_LOCK(na);

    return ptr;
}