/*
  signal.c
  Ruby/GSL: Ruby extension library for GSL (GNU Scientific Library)
    (C) Copyright 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_config.h"
#include "rb_gsl_fft.h"

enum FFT_CONV_CORR {
  RB_GSL_FFT_CONVOLVE = 0,
  RB_GSL_FFT_CORRELATE = 1,
  RB_GSL_FFT_REAL = 2,
  RB_GSL_FFT_HALFCOMPLEX = 3,
};

#ifndef WAVETABLE_P
#define WAVETABLE_P(x) (rb_obj_is_kind_of(x,cgsl_fft_halfcomplex_wavetable)?1:0)
#endif

#ifndef CHECK_WAVETABLE
#define CHECK_WAVETABLE(x) if(!rb_obj_is_kind_of(x,cgsl_fft_halfcomplex_wavetable))\
    rb_raise(rb_eTypeError, "wrong argument type (FFT::HalfComplex::Wavetable expected)");
#endif

#ifndef WORKSPACE_P
#define WORKSPACE_P(x) (rb_obj_is_kind_of(x,cgsl_fft_real_workspace)?1:0)
#endif

#ifndef CHECK_WORKSPACE
#define CHECK_WORKSPACE(x) if(!rb_obj_is_kind_of(x,cgsl_fft_real_workspace))\
    rb_raise(rb_eTypeError, "wrong argument type (FFT::Real::Workspace expected)");
#endif

/* data1, data2: FFTed data */
static void rbgsl_convolve_c(const double *data1, const double *data2, 
			     double *data3, size_t size,
			     gsl_fft_halfcomplex_wavetable *table,
			     gsl_fft_real_workspace *space)
{
  size_t i;
  double re1, re2, im1, im2;
  data3[0] = data1[0]*data2[0];
  data3[size-1] = data1[size-1]*data2[size-1];
  for (i = 1; i < size-1; i+=2) {
    re1 = data1[i];    im1 = data1[i+1];
    re2 = data2[i];    im2 = data2[i+1];
    /* complex multiplication */
    data3[i]   = re1*re2 - im1*im2;
    data3[i+1] = re1*im2 + im1*re2;
  }
  gsl_fft_halfcomplex_inverse(data3, 1, size, table, space);
}

static void rbgsl_correlate_c(const double *data1, const double *data2,
			      double *data3, size_t size,
			      gsl_fft_halfcomplex_wavetable *table,
			      gsl_fft_real_workspace *space)
{
  size_t i;
  double re1, re2, im1, im2;
  data3[0] = data1[0]*data2[0];
  data3[size-1] = data1[size-1]*data2[size-1];
  for (i = 1; i < size-1; i+=2) {
    re1 = data1[i];    im1 = data1[i+1];
    re2 = data2[i];    im2 = data2[i+1];
    /* complex conjugate multiplication */
    data3[i]   =  re1*re2 + im1*im2;
    data3[i+1] = -re1*im2 + im1*re2;
  }
  gsl_fft_halfcomplex_inverse(data3, 1, size, table, space);
}

/* Singleton method, GSL::FFT::HalfComplex */
static VALUE rb_gsl_fft_conv_corr(int argc, VALUE *argv, VALUE obj,
				  enum FFT_CONV_CORR flag1,
				  enum FFT_CONV_CORR flag2)
{
  double *data1, *data2, *data3;
  size_t stride1, stride2, stride3 = 1, size1, size2;
  int naflag1, naflag2, shape;
  gsl_vector *v = NULL;
  gsl_fft_halfcomplex_wavetable *table = NULL;
  gsl_fft_real_wavetable *rtable = NULL;
  gsl_fft_real_workspace *space = NULL, *space2 = NULL;
  int flagt = 0, flagw = 0;
  gsl_vector *vtmp1 = NULL, *vtmp2 = NULL;
  VALUE ary;
  switch (argc) {
  case 4:
    data1 = get_ptr_double3(argv[0], &size1, &stride1, &naflag1);
    data2 = get_ptr_double3(argv[1], &size2, &stride2, &naflag2);
    CHECK_WAVETABLE(argv[2]);
    Data_Get_Struct(argv[2], gsl_fft_halfcomplex_wavetable, table);
    CHECK_WORKSPACE(argv[3]);
    Data_Get_Struct(argv[3], gsl_fft_real_workspace, space);
    break;
  case 3:
    data1 = get_ptr_double3(argv[0], &size1, &stride1, &naflag1);
    data2 = get_ptr_double3(argv[1], &size2, &stride2, &naflag2);
    if (WAVETABLE_P(argv[2])) {
      Data_Get_Struct(argv[2], gsl_fft_halfcomplex_wavetable, table);
      space = gsl_fft_real_workspace_alloc(size1);
      flagw = 1;
    } else if (WORKSPACE_P(argv[2])) {
      Data_Get_Struct(argv[2], gsl_fft_real_workspace, space);
      table = gsl_fft_halfcomplex_wavetable_alloc(size1);
      flagt = 1;
    } else {
      rb_raise(rb_eTypeError, 
	       "wrong argument type %s "
	       "(FFT::HalfComplex::Wavetable or FFT::Real::Workspace expected)", 
	       rb_class2name(CLASS_OF(argv[2])));
    }
    break;
  case 2:
    data1 = get_ptr_double3(argv[0], &size1, &stride1, &naflag1);
    data2 = get_ptr_double3(argv[1], &size2, &stride2, &naflag2);
    table = gsl_fft_halfcomplex_wavetable_alloc(size1);
    space = gsl_fft_real_workspace_alloc(size1);
    flagt = 1;
    flagw = 1;
    break;
  default:
    rb_raise(rb_eArgError, "wrong number of arguments (%d for 2-4)", argc);
  }

  switch (naflag1*naflag2) {
  case 0:
    v = gsl_vector_alloc(size1);
    ary = Data_Wrap_Struct(cgsl_vector, 0, gsl_vector_free, v);
    data3 = v->data;
    stride3 = 1;
    break;
  case 1:
#ifdef HAVE_NARRAY_H
    shape = (int) size1;
    ary = na_make_object(NA_DFLOAT, 1, &shape, cNArray);
    data3 = NA_PTR_TYPE(ary, double*);
    stride3 = 1;
#endif
    break;
  default:
    break;
  }
  switch (flag1) {
  case RB_GSL_FFT_REAL:
    vtmp1 = gsl_vector_alloc(size1);
    vtmp2 = gsl_vector_alloc(size2);
    memcpy(vtmp1->data, data1, sizeof(double)*size1);
    memcpy(vtmp2->data, data2, sizeof(double)*size2);
    data1 = vtmp1->data;
    data2 = vtmp2->data;
    rtable = gsl_fft_real_wavetable_alloc(size1);
    if (size1 == space->n) {
      gsl_fft_real_transform(data1, stride1, size1, rtable, space); 
    } else {
      space2 = gsl_fft_real_workspace_alloc(size1);
      gsl_fft_real_transform(data1, stride1, size1, rtable, space2); 
      /* no freeing space2 here */
    }
    if (size1 != size2) {
      if (rtable) gsl_fft_real_wavetable_free(rtable);
      rtable = gsl_fft_real_wavetable_alloc(size2);
    }
    if (size2 == space->n) {
      gsl_fft_real_transform(data2, stride2, size2, rtable, space); 
    } else if (size2 == size1) {
      gsl_fft_real_transform(data2, stride2, size2, rtable, space2); 
      gsl_fft_real_workspace_free(space2);
    } else {
      if (space2) gsl_fft_real_workspace_free(space2);
      space2 = gsl_fft_real_workspace_alloc(size2);
      gsl_fft_real_transform(data2, stride2, size2, rtable, space2); 
      gsl_fft_real_workspace_free(space2);
    }
    gsl_fft_real_wavetable_free(rtable);
    space2 = NULL;
    rtable = NULL;
    break;
  case RB_GSL_FFT_HALFCOMPLEX:
    /* do nothing */
    break;
  default:
    /* not occur */
    break;
  }
  switch (flag2) {
  case RB_GSL_FFT_CONVOLVE:
    rbgsl_convolve_c(data1, data2, data3, size1, table, space);
    break;
  case RB_GSL_FFT_CORRELATE:
    rbgsl_correlate_c(data1, data2, data3, size1, table, space);
    break;
  default:
    break;
  }
  if (flagt == 1) gsl_fft_halfcomplex_wavetable_free(table);
  if (flagw == 1) gsl_fft_real_workspace_free(space);
  if (vtmp1) gsl_vector_free(vtmp1);
  if (vtmp2) gsl_vector_free(vtmp2);
  return ary;
}

/* Singleton method, GSL::FFT::Real */
static VALUE rb_gsl_fft_real_convolve(int argc, VALUE *argv, VALUE obj)
{
  return rb_gsl_fft_conv_corr(argc, argv, obj, 
			      RB_GSL_FFT_REAL,
			      RB_GSL_FFT_CONVOLVE);
}

/* Singleton method, GSL::FFT::Real */
static VALUE rb_gsl_fft_real_correlate(int argc, VALUE *argv, VALUE obj)
{
  return rb_gsl_fft_conv_corr(argc, argv, obj, 
			      RB_GSL_FFT_REAL,
			      RB_GSL_FFT_CORRELATE);
}

/* Singleton method, GSL::FFT::HalfComplex */
static VALUE rb_gsl_fft_halfcomplex_convolve(int argc, VALUE *argv, VALUE obj)
{
  return rb_gsl_fft_conv_corr(argc, argv, obj, 
			      RB_GSL_FFT_HALFCOMPLEX,
			      RB_GSL_FFT_CONVOLVE);
}

/* Singleton method, GSL::FFT::HalfComplex */
static VALUE rb_gsl_fft_halfcomplex_correlate(int argc, VALUE *argv, VALUE obj)
{
  return rb_gsl_fft_conv_corr(argc, argv, obj, 
			      RB_GSL_FFT_HALFCOMPLEX,
			      RB_GSL_FFT_CORRELATE);
}

void Init_gsl_signal(VALUE module)
{
  rb_define_singleton_method(mgsl_fft_real, "convolve",
			     rb_gsl_fft_real_convolve, -1);
  rb_define_singleton_method(mgsl_fft_real, "correlate",
			     rb_gsl_fft_real_correlate, -1);	
  rb_define_singleton_method(mgsl_fft_halfcomplex, "convolve",
			     rb_gsl_fft_halfcomplex_convolve, -1);
  rb_define_singleton_method(mgsl_fft_halfcomplex, "correlate",
			     rb_gsl_fft_halfcomplex_correlate, -1);			     
}

#undef WAVETABLE_P
#undef CHECK_WAVETABLE
#undef WORKSPACE_P
#undef CHECK_WORKSPACE