/*
Copyright 2016 Ian Jauslin

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*
  Base functions for integration

  see integral_*.h for the values taken by INTEGRAL_FUNC, etc...
*/


// compute the integral of a real function of 1 variable using the Gauss-Legendre scheme
int INTEGRAL_FUNC(integrate_gauss) (INTEGRAL_FLOAT_TYPE* out, int (*func)(INTEGRAL_FLOAT_TYPE*, INTEGRAL_FLOAT_TYPE, void*), INTEGRAL_FLOAT_TYPE lower, INTEGRAL_FLOAT_TYPE upper, INTEGRAL_FLOATARRAY_TYPE abcissa, INTEGRAL_FLOATARRAY_TYPE weights, void* extra_args){
  unsigned int i;
  INTEGRAL_FLOAT_TYPE valf, delta, avg, x;
  int ret;

  // check arguments
  if(abcissa.length != weights.length){
    return(LIBINUM_ERROR_SIZE_MISMATCH);
  }

  #ifdef INTEGRAL_FLOAT_INIT
    INTEGRAL_FLOAT_INIT(valf);
    INTEGRAL_FLOAT_INIT(delta);
    INTEGRAL_FLOAT_INIT(avg);
    INTEGRAL_FLOAT_INIT(x);
  #endif

  // init to 0
  INTEGRAL_FLOAT_SET_UI(*out, 0);

  // half length of interval
  INTEGRAL_FLOAT_SUB(delta, upper, lower);
  INTEGRAL_FLOAT_DIV_UI(delta, delta, 2);
  // average of interval
  INTEGRAL_FLOAT_ADD(avg, upper, lower);
  INTEGRAL_FLOAT_DIV_UI(avg, avg, 2);

  for(i=0;i<abcissa.length;i++){
    // evaluate at x
    INTEGRAL_FLOAT_MUL(x, delta, abcissa.values[i]);
    INTEGRAL_FLOAT_ADD(x, x, avg);

    ret=(*func)(&valf, x, extra_args);
    if(ret<0){
      #ifdef INTEGRAL_FLOAT_FREE
	INTEGRAL_FLOAT_FREE(valf);
	INTEGRAL_FLOAT_FREE(delta);
	INTEGRAL_FLOAT_FREE(avg);
	INTEGRAL_FLOAT_FREE(x);
      #endif
      return(ret);
    }
    // check whether valf is a number
    if(! INTEGRAL_FLOAT_ISNUMBER(valf)){
      #ifdef INTEGRAL_FLOAT_FREE
	INTEGRAL_FLOAT_FREE(valf);
	INTEGRAL_FLOAT_FREE(delta);
	INTEGRAL_FLOAT_FREE(avg);
	INTEGRAL_FLOAT_FREE(x);
      #endif
      return(LIBINUM_ERROR_NAN);
    }
      
    INTEGRAL_FLOAT_MUL(valf, valf, weights.values[i]);
    INTEGRAL_FLOAT_ADD(*out, *out, valf);
  }

  INTEGRAL_FLOAT_MUL(*out, *out, delta);

  #ifdef INTEGRAL_FLOAT_FREE
    INTEGRAL_FLOAT_FREE(valf);
    INTEGRAL_FLOAT_FREE(delta);
    INTEGRAL_FLOAT_FREE(avg);
    INTEGRAL_FLOAT_FREE(x);
  #endif

  return(0);
}

// multithreaded version
// arguments to pass to each thread
struct INTEGRAL_FUNC(pthread_integrate_gauss_args) {
  // partial sum of the values assigned to the thread
  INTEGRAL_FLOAT_TYPE out;
  // abcissa
  INTEGRAL_FLOATARRAY_TYPE x;
  // weights
  INTEGRAL_FLOATARRAY_TYPE weights;
  // pointer to the function to evaluate
  int (*func)(INTEGRAL_FLOAT_TYPE*, INTEGRAL_FLOAT_TYPE, void*);
  // extra arguments passed to func
  void* extra_args;
  // return value
  int ret;
};
int INTEGRAL_FUNC(integrate_gauss_multithread) (INTEGRAL_FLOAT_TYPE* out, int (*func)(INTEGRAL_FLOAT_TYPE*, INTEGRAL_FLOAT_TYPE, void*), INTEGRAL_FLOAT_TYPE lower, INTEGRAL_FLOAT_TYPE upper, INTEGRAL_FLOATARRAY_TYPE abcissa, INTEGRAL_FLOATARRAY_TYPE weights, void* extra_args, unsigned int threads, array_pthread_t* thread_ids){
  unsigned int i;
  INTEGRAL_FLOAT_TYPE x, delta, avg;
  struct INTEGRAL_FUNC(pthread_integrate_gauss_args) args[threads];
  int ret=0;
  unsigned int thread_nr;

  array_pthread_t_init(thread_ids, threads);
  thread_ids->length=threads;

  // check arguments
  if(abcissa.length != weights.length){
    return(LIBINUM_ERROR_SIZE_MISMATCH);
  }

  #ifdef INTEGRAL_FLOAT_INIT
    INTEGRAL_FLOAT_INIT(delta);
    INTEGRAL_FLOAT_INIT(avg);
    INTEGRAL_FLOAT_INIT(x);
  #endif

  // init to 0
  INTEGRAL_FLOAT_SET_UI(*out, 0);

  // half length of interval
  INTEGRAL_FLOAT_SUB(delta, upper, lower);
  INTEGRAL_FLOAT_DIV_UI(delta, delta, 2);
  // average of interval
  INTEGRAL_FLOAT_ADD(avg, upper, lower);
  INTEGRAL_FLOAT_DIV_UI(avg, avg, 2);

  // inits
  for(thread_nr=0;thread_nr<threads;thread_nr++){
    INTEGRAL_FLOATARRAY_FUNC(init) (&(args[thread_nr].x),abcissa.length/threads+1);
    INTEGRAL_FLOATARRAY_FUNC(init) (&(args[thread_nr].weights),abcissa.length/threads+1);
    #ifdef INTEGRAL_FLOAT_INIT
      INTEGRAL_FLOAT_INIT(args[thread_nr].out);
    #endif
  }

  // set abcissa and weights
  for(i=0,thread_nr=0;i<abcissa.length;i++,thread_nr=(thread_nr+1)%threads){
    INTEGRAL_FLOAT_MUL(x, delta, abcissa.values[i]);
    INTEGRAL_FLOAT_ADD(x, x, avg);

    INTEGRAL_FLOATARRAY_FUNC(append) (x, &(args[thread_nr].x));
    INTEGRAL_FLOATARRAY_FUNC(append) (weights.values[i], &(args[thread_nr].weights));
  }

  for(thread_nr=0;thread_nr<threads;thread_nr++){
    // set func
    args[thread_nr].func=func;
    // set extra_args
    args[thread_nr].extra_args=extra_args;
    // init ret
    args[thread_nr].ret=0;
    // run threads
    pthread_create(thread_ids->values+thread_nr, NULL, INTEGRAL_FUNC(integrate_gauss_thread), (void*)(args+thread_nr));
  }

  // wait for completion and join
  for(thread_nr=0;thread_nr<threads;thread_nr++){
    pthread_join(thread_ids->values[thread_nr], NULL);

    if(args[thread_nr].ret<0){
      ret=args[thread_nr].ret;
    }
    else{
      INTEGRAL_FLOAT_ADD(*out, *out, args[thread_nr].out);
    }
  }

  // multiply by size of interval
  INTEGRAL_FLOAT_MUL(*out, *out, delta);

  for(thread_nr=0;thread_nr<threads;thread_nr++){
    INTEGRAL_FLOATARRAY_FUNC(free) (args[thread_nr].x);
    INTEGRAL_FLOATARRAY_FUNC(free) (args[thread_nr].weights);
    #ifdef INTEGRAL_FLOAT_FREE
      INTEGRAL_FLOAT_FREE(args[thread_nr].out);
    #endif
  }

  #ifdef INTEGRAL_FLOAT_FREE
    INTEGRAL_FLOAT_FREE(delta);
    INTEGRAL_FLOAT_FREE(avg);
    INTEGRAL_FLOAT_FREE(x);
  #endif


  return(ret);
}
// per-thread function
void* INTEGRAL_FUNC(integrate_gauss_thread) (void* args){
  unsigned int i;
  INTEGRAL_FLOAT_TYPE valf;
  int ret;
  struct INTEGRAL_FUNC(pthread_integrate_gauss_args)* argument=((struct INTEGRAL_FUNC(pthread_integrate_gauss_args)*)args);
  #ifdef INTEGRAL_FLOAT_INIT
    INTEGRAL_FLOAT_INIT(valf);
  #endif

  INTEGRAL_FLOAT_SET_UI(argument->out, 0);

  for(i=0;i<argument->x.length;i++){
    // evaluate
    ret=(*(argument->func))(&valf, argument->x.values[i], argument->extra_args);

    if(ret<0){
      #ifdef INTEGRAL_FLOAT_FREE
	INTEGRAL_FLOAT_FREE(valf);
      #endif
      argument->ret=ret;
      return(NULL);
    }
    // check whether valf is a number
    if(! INTEGRAL_FLOAT_ISNUMBER(valf)){
      #ifdef INTEGRAL_FLOAT_FREE
	INTEGRAL_FLOAT_FREE(valf);
      #endif
      argument->ret=LIBINUM_ERROR_NAN;
      return(NULL);
    }

    INTEGRAL_FLOAT_MUL(valf, valf, argument->weights.values[i]);
    INTEGRAL_FLOAT_ADD(argument->out, argument->out, valf);
  }
  #ifdef INTEGRAL_FLOAT_FREE
    INTEGRAL_FLOAT_FREE(valf);
  #endif

  return(NULL);
}



// smart management of temporary variables: initialize as many as needed but allow them to be re-used instead of freeing them
#ifdef INTEGRAL_FLOAT_INIT
int INTEGRAL_FUNC(integrate_gauss_smarttmp) (INTEGRAL_FLOAT_TYPE* out, int (*func)(INTEGRAL_FLOAT_TYPE*, INTEGRAL_FLOAT_TYPE, void*), INTEGRAL_FLOAT_TYPE lower, INTEGRAL_FLOAT_TYPE upper, INTEGRAL_FLOATARRAY_TYPE abcissa, INTEGRAL_FLOATARRAY_TYPE weights, INTEGRAL_FLOATARRAY_TYPE* tmps, void* extra_args){
  unsigned int i;
  int ret;

  // check arguments
  if(abcissa.length != weights.length){
    return(LIBINUM_ERROR_SIZE_MISMATCH);
  }

  // allocate tmp values if needed
  INTEGRAL_FLOATARRAY_FUNC(alloc_tmps)(4, tmps);

  // init to 0
  INTEGRAL_FLOAT_SET_UI(*out, 0);

  // half length of interval
  INTEGRAL_FLOAT_SUB(tmps->values[0], upper, lower);
  INTEGRAL_FLOAT_DIV_UI(tmps->values[0], tmps->values[0], 2);
  // average of interval
  INTEGRAL_FLOAT_ADD(tmps->values[1], upper, lower);
  INTEGRAL_FLOAT_DIV_UI(tmps->values[1], tmps->values[1], 2);

  for(i=0;i<abcissa.length;i++){
    // evaluation point
    INTEGRAL_FLOAT_MUL(tmps->values[2], tmps->values[0], abcissa.values[i]);
    INTEGRAL_FLOAT_ADD(tmps->values[2], tmps->values[2], tmps->values[1]);

    ret=(*func)(tmps->values+3, tmps->values[2], extra_args);
    if(ret<0){
      return(ret);
    }
    // check whether tmps->values[3] is a number
    if(! INTEGRAL_FLOAT_ISNUMBER(tmps->values[3])){
      return(LIBINUM_ERROR_NAN);
    }
      
    INTEGRAL_FLOAT_MUL(tmps->values[3], tmps->values[3], weights.values[i]);
    INTEGRAL_FLOAT_ADD(*out, *out, tmps->values[3]);
  }

  INTEGRAL_FLOAT_MUL(*out, *out, tmps->values[0]);

  return(0);
}
#endif

// multithreaded version with smart management of temporary variables (only for datatypes that need to be initialized)
#ifdef INTEGRAL_FLOAT_INIT
// arguments to pass to each thread
struct INTEGRAL_FUNC(pthread_integrate_gauss_smarttmp_args) {
  // partial sum of the values assigned to the thread
  INTEGRAL_FLOAT_TYPE* out;
  // abcissa
  INTEGRAL_FLOATARRAY_TYPE x;
  // weights
  INTEGRAL_FLOATARRAY_TYPE weights;
  // tmp
  INTEGRAL_FLOAT_TYPE* tmp;
  // pointer to the function to evaluate
  int (*func)(INTEGRAL_FLOAT_TYPE*, INTEGRAL_FLOAT_TYPE, void*);
  // extra arguments passed to func
  void* extra_args;
  // return value
  int ret;
};
int INTEGRAL_FUNC(integrate_gauss_smarttmp_multithread) (INTEGRAL_FLOAT_TYPE* out, int (*func)(INTEGRAL_FLOAT_TYPE*, INTEGRAL_FLOAT_TYPE, void*), INTEGRAL_FLOAT_TYPE lower, INTEGRAL_FLOAT_TYPE upper, INTEGRAL_FLOATARRAY_TYPE abcissa, INTEGRAL_FLOATARRAY_TYPE weights, INTEGRAL_FLOATARRAY_TYPE* tmps, void* extra_args, unsigned int threads, array_pthread_t* thread_ids){
  unsigned int i;
  struct INTEGRAL_FUNC(pthread_integrate_gauss_smarttmp_args) args[threads];
  int ret=0;
  unsigned int thread_nr;

  array_pthread_t_init(thread_ids, threads);
  thread_ids->length=threads;

  // check arguments
  if(abcissa.length != weights.length){
    return(LIBINUM_ERROR_SIZE_MISMATCH);
  }

  // allocate tmps if needed
  if(tmps->memory<2+2*threads+abcissa.length){
    // no need to resize since the values should not be kept
    INTEGRAL_FLOATARRAY_FUNC(free)(*tmps);
    INTEGRAL_FLOATARRAY_FUNC(init)(tmps, 2+2*threads+abcissa.length);
  }
  for (i=tmps->length;i<2+2*threads+abcissa.length;i++){
    INTEGRAL_FLOAT_INIT(tmps->values[i]);
    (tmps->length)++;
  }


  // init to 0
  INTEGRAL_FLOAT_SET_UI(*out, 0);

  // half length of interval
  INTEGRAL_FLOAT_SUB(tmps->values[0], upper, lower);
  INTEGRAL_FLOAT_DIV_UI(tmps->values[0], tmps->values[0], 2);
  // average of interval
  INTEGRAL_FLOAT_ADD(tmps->values[1], upper, lower);
  INTEGRAL_FLOAT_DIV_UI(tmps->values[1], tmps->values[1], 2);

  // inits
  for(thread_nr=0;thread_nr<threads;thread_nr++){
    INTEGRAL_FLOATARRAY_FUNC(init) (&(args[thread_nr].x),abcissa.length/threads+1);
    INTEGRAL_FLOATARRAY_FUNC(init) (&(args[thread_nr].weights),abcissa.length/threads+1);
    args[thread_nr].out=tmps->values+2+thread_nr;
    args[thread_nr].tmp=tmps->values+2+threads+thread_nr;
  }

  // set abcissa and weights
  for(i=0,thread_nr=0;i<abcissa.length;i++,thread_nr=(thread_nr+1)%threads){
    INTEGRAL_FLOAT_MUL(tmps->values[2+2*threads+i], tmps->values[0], abcissa.values[i]);
    INTEGRAL_FLOAT_ADD(tmps->values[2+2*threads+i], tmps->values[2+2*threads+i], tmps->values[1]);

    INTEGRAL_FLOATARRAY_FUNC(append_noinit) (tmps->values[2+2*threads+i], &(args[thread_nr].x));
    INTEGRAL_FLOATARRAY_FUNC(append_noinit) (weights.values[i], &(args[thread_nr].weights));
  }

  for(thread_nr=0;thread_nr<threads;thread_nr++){
    // set func
    args[thread_nr].func=func;
    // set extra_args
    args[thread_nr].extra_args=extra_args;
    // init ret
    args[thread_nr].ret=0;
    // run threads
    pthread_create(thread_ids->values+thread_nr, NULL, INTEGRAL_FUNC(integrate_gauss_smarttmp_thread), (void*)(args+thread_nr));
  }

  // wait for completion and join
  for(thread_nr=0;thread_nr<threads;thread_nr++){
    pthread_join(thread_ids->values[thread_nr], NULL);

    if(args[thread_nr].ret<0){
      ret=args[thread_nr].ret;
    }
    else{
      INTEGRAL_FLOAT_ADD(*out, *out, *(args[thread_nr].out));
    }
  }

  // multiply by size of interval
  INTEGRAL_FLOAT_MUL(*out, *out, tmps->values[0]);

  // free x and weights
  for(thread_nr=0;thread_nr<threads;thread_nr++){
    INTEGRAL_FLOATARRAY_FUNC(free_vects)(args[thread_nr].x);
    INTEGRAL_FLOATARRAY_FUNC(free_vects)(args[thread_nr].weights);
  }
  return(ret);
}
// per-thread function
void* INTEGRAL_FUNC(integrate_gauss_smarttmp_thread) (void* args){
  unsigned int i;
  int ret;
  struct INTEGRAL_FUNC(pthread_integrate_gauss_smarttmp_args)* argument=((struct INTEGRAL_FUNC(pthread_integrate_gauss_smarttmp_args)*)args);

  INTEGRAL_FLOAT_SET_UI(*(argument->out), 0);

  for(i=0;i<argument->x.length;i++){
    // evaluate
    ret=(*(argument->func))(argument->tmp, argument->x.values[i], argument->extra_args);

    if(ret<0){
      argument->ret=ret;
      return(NULL);
    }
    // check whether argument->tmp is a number
    if(! INTEGRAL_FLOAT_ISNUMBER(*(argument->tmp))){
      argument->ret=LIBINUM_ERROR_NAN;
      return(NULL);
    }

    INTEGRAL_FLOAT_MUL(*(argument->tmp), *(argument->tmp), argument->weights.values[i]);
    INTEGRAL_FLOAT_ADD(*(argument->out), *(argument->out), *(argument->tmp));
  }

  return(NULL);
}
#endif


// compute the abcissa and weights for the Gauss-Legendre numerical integration scheme
int INTEGRAL_FUNC(gauss_legendre_weights) (unsigned int order, INTEGRAL_FLOAT_TYPE tolerance, unsigned int maxiter, INTEGRAL_FLOATARRAY_TYPE* abcissa, INTEGRAL_FLOATARRAY_TYPE* weights){
  unsigned int i;
  INTEGRAL_FLOAT_TYPE x, tmp;
  INTEGRAL_POLYNOMIALARRAY_TYPE legendre;
  int ret;

  INTEGRAL_FLOATARRAY_FUNC(init) (abcissa, order);
  INTEGRAL_FLOATARRAY_FUNC(init) (weights, order);

  #ifdef INTEGRAL_FLOAT_INIT
    INTEGRAL_FLOAT_INIT(x);
    INTEGRAL_FLOAT_INIT(tmp);
  #endif

  INTEGRAL_POLYNOMIALARRAY_FUNC(init) (&legendre, 2);
  // compute the roots of the 'order'-th Legendre polynomial
  INTEGRAL_POLYNOMIAL_FUNC(legendre) (order, legendre.values);
  INTEGRAL_POLYNOMIAL_FUNC(derive) (legendre.values[0], legendre.values+1);
  legendre.length=2;

  for(i=0;i<order;i++){
    // initial guess
    INTEGRAL_FLOAT_SET_D(x, cos(3.1415926*(i+0.75)/(order+0.5)));
    ret=INTEGRAL_FUNC(root_newton_inplace) (&x, &INTEGRAL_FUNC(legendre_wrapper), &INTEGRAL_FUNC(deriv_legendre_wrapper), tolerance, maxiter, &legendre);

    if(ret<0){
      #ifdef INTEGRAL_FLOAT_FREE
	INTEGRAL_FLOAT_FREE(x);
	INTEGRAL_FLOAT_FREE(tmp);
      #endif
      INTEGRAL_POLYNOMIALARRAY_FUNC(free) (legendre);
      INTEGRAL_FLOATARRAY_FUNC(free) (*abcissa);
      INTEGRAL_FLOATARRAY_FUNC(free) (*weights);
      return(ret);
    }

    INTEGRAL_FLOATARRAY_FUNC(append) (x, abcissa);

    // weight: 2/((1-x^2)*(deriv_Legendre(x)))^2
    INTEGRAL_POLYNOMIAL_FUNC(evaluate) (&tmp, x, legendre.values[1]);
    INTEGRAL_FLOAT_POW_UI(tmp, tmp, 2);
    INTEGRAL_FLOAT_POW_UI(x, x, 2);
    INTEGRAL_FLOAT_UI_SUB(x, 1, x);
    INTEGRAL_FLOAT_MUL(x, x, tmp);
    INTEGRAL_FLOAT_SET_UI(tmp, 2);
    INTEGRAL_FLOAT_DIV(x, tmp, x);

    INTEGRAL_FLOATARRAY_FUNC(append) (x, weights);
  }

  #ifdef INTEGRAL_FLOAT_FREE
    INTEGRAL_FLOAT_FREE(x);
    INTEGRAL_FLOAT_FREE(tmp);
  #endif
  INTEGRAL_POLYNOMIALARRAY_FUNC(free) (legendre);

  return(0);
}

// wrapper functions to evaluate the legendre polynomial and its derivative
int INTEGRAL_FUNC(legendre_wrapper) (INTEGRAL_FLOAT_TYPE* out, INTEGRAL_FLOAT_TYPE in, void* legendre){
  INTEGRAL_POLYNOMIAL_FUNC(evaluate) (out, in, ((INTEGRAL_POLYNOMIALARRAY_TYPE*)legendre)->values[0]);
  return(0);
}
int INTEGRAL_FUNC(deriv_legendre_wrapper) (INTEGRAL_FLOAT_TYPE* out, INTEGRAL_FLOAT_TYPE in, void* legendre){
  INTEGRAL_POLYNOMIAL_FUNC(evaluate) (out, in, ((INTEGRAL_POLYNOMIALARRAY_TYPE*)legendre)->values[1]);
  return(0);
}