/**
 * Subroutines to correct for the intra-pixel sensitivity
 * variation in e.g. NICMOS pixels. There are routines
 * to corect both, an extracted SPC's and a list
 * of PET pixels.
 *
 */
#include "fitsio.h"
#include <gsl/gsl_interp.h>
#include <gsl/gsl_roots.h>
#include <gsl/gsl_multifit_nlin.h>
#include "aXe_grism.h"
#include "aXe_utils.h"
#include "aXe_errors.h"
#include "spc_spc.h"
#include "spce_pathlength.h"
#include "fringe_conf.h"
#include "spc_wl_calib.h"
#include "trfit_utils.h"
#include "ipixcorr_utils.h"

#define MAX(x,y) (((x)>(y))?(x):(y))
#define MIN(x,y) (((x)<(y))?(x):(y))

/**
 * Function: get_intpix_corr
 * The function determines and returns a correction factor with
 * values stored in an interpolator at a certain wavelength.
 * To derive the result the dispersion solution is inverted,
 * then the trace position for the wavelength is determined.
 * The the fractional pixel value in y is computed and the
 * corresponding correction factor is returned.
 *
 * Parameters:
 * @param actbeam - beam to correct
 * @param cdisp   - the dispersion coefficients for the beam
 * @param ipcorr  - the correction values stored in an interpolator
 * @param lambda  - the wavelength to correct
 *
 * Returns:
 * @return cfactor - the correction factor
 */
double
get_intpix_corr(beam *actbeam, gsl_vector *cdisp, interpolator *ipcorr,
		float lambda)
{
  double tlength;

  double xvalue;
  double yfract;

  double cfactor=0.0;

  // get the trace length for the dispersion
  // solution and wavelength
  tlength = get_tlength_from_dispersion(cdisp, lambda);

  // compute the x-offset from the reference point
  // for this trace length value
  xvalue = get_xvalue_from_tlength(actbeam, tlength);

  // compute the fractional y-value
  yfract = get_yfract_for_xvalue(actbeam, xvalue);

  // get the correction factor
  cfactor = eval_interp(ipcorr, yfract);

  fprintf(stdout, "xdiff: %e, yfrac: %e, factor: %e ", xvalue, yfract, cfactor);

  //fprintf(stdout, "lambda: %f, tlength: %f, xval: %f, yval: %f, yfract: %f, cfactor: %f\n",
  // 	  lambda, tlength, actbeam->refpoint.x+xvalue, actbeam->refpoint.y+actbeam->spec_trace->func (xvalue, actbeam->spec_trace->data),yfract, cfactor);

  // return the correction factor
  return cfactor;
}

/**
 * Function: get_yfract_for_xvalue
 * The function computes the fractional y-value for a trace
 * position in a beam. The trace position is given as the
 * x-offset position with respect to the reference position
 *
 * Parameters:
 * @param actbeam - beam to correct
 * @param xvalue  - the x-value
 *
 * Returns:
 * @return yfract - the fractional y-value
 */
double
get_yfract_for_xvalue(const beam *actbeam, double xvalue)
{
  double dy;
  double yabs;
  double yfract;

  // compute the y-value relative to the reference point
  dy = actbeam->spec_trace->func (xvalue, actbeam->spec_trace->data);

  // compute the absolute y-value on the chip
  yabs = actbeam->refpoint.y + dy;

  // compute the fractional pixel of this y-value
  yfract = yabs - floor(yabs);

  // return the fractional y-value
  return yfract;
}

/**
 * Function: get_xvalue_from_tlength
 * The function computes the x-offset position from the reference point
 * for a given tracelength in a given beam.
 * The solution is numerically derived and therefore
 * does work for any reasonable tracefunction.
 *
 * Parameters:
 * @param actbeam - the beam to compute the x-offset
 * @param tlength - tracelength
 *
 * Returns:
 * @return xdiff - the x-offset
 */
double
get_xvalue_from_tlength(beam *actbeam, double tlength)
{
  int iter=0;
  int status;

  double xdiff=0.0;
  d_point x_interv;

  // define and initialize the solver
  const gsl_root_fsolver_type *T = gsl_root_fsolver_brent;
  gsl_root_fsolver            *s = gsl_root_fsolver_alloc (T);

  gsl_function F;
  tlength_pars *tpars;

  // derive the x-interval from the beam boundaries
  x_interv = get_xinterv_from_beam(actbeam);
  //  fprintf(stdout, "xpos: %f, ypos: %f\n", x_interv.x, x_interv.y);

  // allocate and fill the parameters
  tpars = (tlength_pars *) malloc(sizeof(tlength_pars));
  tpars->actbeam = actbeam;
  tpars->tlength = tlength;

  // fille the GSL-function
  F.function = &zero_tlength;
  F.params =   tpars;

  // set the boundaries for the solver
  gsl_root_fsolver_set (s, &F, x_interv.x, x_interv.y);

  // iterate to find the zeropoint
  do
    {
      // increment the iteration counter
      iter++;

      // iterate on the solver
      status = gsl_root_fsolver_iterate (s);

      // get a new guess from the solver
      xdiff = gsl_root_fsolver_root (s);

      // derive and set new boundaries
      x_interv.x = gsl_root_fsolver_x_lower (s);
      x_interv.y = gsl_root_fsolver_x_upper (s);

      // check the accuracy
      status = gsl_root_test_interval (x_interv.x, x_interv.y,
				       0, 0.0001);
      //--------------CAVEAT---------------------------------------
      // until March 28th 08 the code, wriongly was:
      //status = gsl_root_test_interval (x_interv.x, x_interv.x,
      //                                  0, 0.0001);
      // somehow this made no difference.....
      //-----------------------------------------------------------
    }
  // check for the break condition
  while (status == GSL_CONTINUE && iter < MAX_ITER);

  // free the memory
  free(tpars);

  // free the memory
  gsl_root_fsolver_free (s);

  // return the result
  return xdiff;
}

/**
 * Function: zero_tlength
 * Helper function for the 1D-root finding routine.
 * Evaluates a functional value at the position given
 * in the first parameter using the values given
 * as second parameter.
 *
 * Parameters:
 * @param x      - the independent value
 * @param params - ingredients to evaluate the function value
 *
 * Returns:
 * @return tzero - the function value
 */
double
zero_tlength (double x, void *params)
{
  double tzero;

  // extract the components from the parameter-structure
  beam  *actbeam = ((tlength_pars *) params)->actbeam;
  double tlength = ((tlength_pars *) params)->tlength;

  // compute the function value
  tzero = actbeam->spec_trace->path_len (x, actbeam->spec_trace->data) - tlength;

  // return the function value
  return tzero;
}

/**
 * Function: get_xinterv_from_beam
 * Determines the maximum and minimum extension of a given
 * beam along the x-axis. For security the interval is extended
 * at both ends by a fixed amount (quantified in the header-file).
 *
 * Parameters:
 * @param actbeam - the beam
 *
 * Returns:
 * @return x_interv - xmin, xmax covered by the beam
 */
d_point
get_xinterv_from_beam(const beam *actbeam)
{
  d_point x_interv;

  // compute the beam extension
  // towards the negative x-axis
  x_interv.x = (double)MIN(actbeam->corners[0].x,
			   MIN(actbeam->corners[1].x,
			       MIN(actbeam->corners[2].x,
				   actbeam->corners[3].x)))-actbeam->refpoint.x - INTERVEXT;

  // compute the beam extension
  // towards the positive x-axis
  x_interv.y = (double)MAX(actbeam->corners[0].x,
			   MAX(actbeam->corners[1].x,
			       MAX(actbeam->corners[2].x,
				   actbeam->corners[3].x)))-actbeam->refpoint.x + INTERVEXT;

  // return the interval
  return x_interv;
}

/**
 * Function: get_tlength_from_dispersion
 * The function computes and returnes the trace length for
 * a given diseprsion solution and wavelength. The tracelength
 * value is returned. Currently only linear and quadratic
 * dispersio solution are considered.
 *
 * Parameters:
 * @param cdisp  - the dispersion coefficient polynomial
 * @param lambda - the wavelength
 *
 * Returns:
 * @return tlength - the trace length
 */
double
get_tlength_from_dispersion(gsl_vector *cdisp, float lambda)
{
  double tlength=0.0;
  //double tlength_a=0.0;
  double det;

  // check whether the polynomial is linear
  if (cdisp->size == 2)
    {
      // compute the tracelength for linear dispersion
      tlength = (lambda-gsl_vector_get(cdisp, 0))/gsl_vector_get(cdisp, 1);
    }
  // check whether the polynomial is quadratic
  else if (cdisp->size == 3)
    {
      // compute the determinante
      det = gsl_vector_get(cdisp, 1) * gsl_vector_get(cdisp, 1)
		 - 4.0 * gsl_vector_get(cdisp, 0) * gsl_vector_get(cdisp, 2);

      // gove error if det < 0.0
      if (det < 0.0)
	aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		     "get_tlength_from_dispersion: Can not determine the tracelength since det < 0.0!\n");

      // compute the tracelength
      tlength   = (-gsl_vector_get(cdisp, 1)+sqrt(det))/(2.0*gsl_vector_get(cdisp, 2));
      //    tlength_a = (-gsl_vector_get(cdisp, 1)-sqrt(det))/(2.0*gsl_vector_get(cdisp, 2));
      //    fprintf(stdout, "plus: %f, minus: f\n", tlength, tlength_a);
    }
  else
    {
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_tlength_from_dispersion: The dispersion solution has %i coefficients and can not be inverted!\n", cdisp->size);
    }

  // return the tracelength
  return tlength;
}

/**
 * Function: condense_dispersion
 * The function strips off leading zeroes in the high order terms
 * of the dispersion solution derived from the configuration file.
 * These usually become a problem in the inversion to derive the
 * tracelength.
 *
 * Parameters:
 * @param disp__input - the configuration file dispersion
 *
 * Returns:
 * @return cdisp - the effective dispersion relation
 */
gsl_vector *
condense_dispersion(gsl_vector *disp_input)
{
  int index;
  int i;

  gsl_vector *cdisp;

  // get the size of the incomingg dispersion
  // vector
  index = disp_input->size;

  // determine the true size by subtracting
  // high order zeros
  while (index > -1 && !gsl_vector_get(disp_input, index-1))
    index--;

  // allocate space for the new vector
  cdisp = gsl_vector_alloc(index);

  // tranasfer the relevant values
  // from the old to the new vector
  for (i=0; i < index; i++)
    gsl_vector_set(cdisp, i, gsl_vector_get(disp_input, i));

  // return the new dispersion vector
  return cdisp;
}

/**
 * Function: fitting_ipc_corr
 *
 * Parameters:
 * @param actbeam        - the beam to examine
 * @param conf_file_path - the full pathname too the configuration file
 * @param ipcorr         - the interpolator with the correction factor
 * @param SPC            - the full spectrum to correct
 * @param obs            - the grism image
 * @param data_matrix    - background subtracted grism image
 * @param ipc_file_path  - file name for the phase data
 *
 * Returns:
 * @return icorr - marker whether the correction was applied or not (1/0)
 */
int
fitting_ipc_corr(beam act_beam, char conf_file_path[], interpolator *ipcorr,
		 full_spectr *SPC, observation *obs, gsl_matrix *data_matrix,
		 char ipc_file_path[], beam *beam_ptr)
{
  aperture_conf *conf;

  d_point        fit_qual;

  int            icorr     = 0;

  trace_func    *tracefun  = act_beam.spec_trace;
  double        *tracedata =  (double *)(tracefun->data);

  // load the configuration file
  conf = get_aperture_descriptor(conf_file_path);

  // determine the phase data and fix the phase
  // by fitting a function to it
  fit_qual = kappa_sigma_klipp_ipc(conf, obs, data_matrix, act_beam, ipc_file_path);

  // report on the shift and the quality
  fprintf(stdout, "shift-difference: %e, quality: %e\n", fit_qual.x, fit_qual.y);

  // check whether the quality is good enough
  if (fit_qual.y < MAXQUALITY)
    {
      // apply a correction to the
      // reference point to achive the
      // right corrections
      act_beam.refpoint.y += fit_qual.x;
      beam_ptr->refpoint.y += fit_qual.x;
      beam_ptr->ignore = -100;

      // report the new y-shift
      fprintf(stdout, "new y-shift: %e\n", act_beam.refpoint.y - (double)(int)(act_beam.refpoint.y + tracedata[1]));

      // apply the sensitivity correction to the spectrum
      intpix_corr_beam(act_beam, conf_file_path, ipcorr, SPC);

      // set the correction marker
      icorr=1;
    }

  // free allocated memory
  free_aperture_conf(conf);

  // return the marker
  return icorr;
}

/**
 * Function: intpix_corr_beam
 * The functio corrects a full spectrum for the intrapixel sensitivity
 * variations.
 * For every spectral element its fractional y-value is determined,
 * and then the correction factor for this y-value is computed
 * and applied to the appropriate elements of the spectral bin.
 *
 * Parameters:
 * @param actbeam        - the beam to examine
 * @param conf_file_path - the full pathname too the configuration file
 * @param ipcorr         - the interpolator with the correction factor
 * @param SPC            - the full spectrum to correct
 *
 * Returns:
 * @return -
 */
void
intpix_corr_beam(beam actbeam, char conf_file_path[], interpolator *ipcorr,
		 full_spectr *SPC)
{
  int            index=0;
  int            for_grism=1;

  double         cfactor;

  double         lambda;

  gsl_vector     *cdisp;

  d_point        pixel;

  dispstruct     *beam_disp;

  aperture_conf  *conf;

  // load the configuration file
  conf = get_aperture_descriptor(conf_file_path);

  // check whether it is grism (for_grism=1)
  // or prism (for_grism=0) data
  // give an error if there is a prism solution
  for_grism = check_for_grism (conf_file_path, actbeam.ID);
  if (!for_grism)
    aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		 "intpix_corr_beam: Only grism dispersion solution can be corrected.\n");

  // get the wavelength dispersion relation at
  // position "refpoint". conf->refx and conf->refy
  // are used at this point to allow for a non (0,0) centered
  // 2D field dependence.
  pixel.x = actbeam.refpoint.x - conf->refx;
  pixel.y = actbeam.refpoint.y - conf->refy;

  // derive the dispersion at the object position
  beam_disp = get_dispstruct_at_pos(conf_file_path, for_grism,
				    actbeam.ID, pixel);

  // skipp high order zeroes in the dispersion solution
  cdisp = condense_dispersion(beam_disp->pol);

  for (index=0; index < SPC->nelems; index++)
    {
      lambda = SPC->fgr_spec->spec[index].lambda_mean;

      if (!gsl_isnan (lambda) && lambda)
	{
	  cfactor = get_intpix_corr(&actbeam, cdisp, ipcorr, lambda);

    	  fprintf(stdout, "lambda: %e, factor: %e\n", lambda, cfactor);
	  /*
	  SPC->fgr_spec->spec[index].count  = SPC->fgr_spec->spec[index].count / cfactor;
	  SPC->fgr_spec->spec[index].error  = SPC->fgr_spec->spec[index].error / cfactor;
	  SPC->fgr_spec->spec[index].flux   = SPC->fgr_spec->spec[index].flux / cfactor;
	  SPC->fgr_spec->spec[index].ferror = SPC->fgr_spec->spec[index].ferror / cfactor;

	  SPC->bck_spec->spec[index].count  = SPC->bck_spec->spec[index].count / cfactor;
	  SPC->bck_spec->spec[index].error  = SPC->bck_spec->spec[index].error / cfactor;
	  SPC->bck_spec->spec[index].flux   = SPC->bck_spec->spec[index].flux / cfactor;
	  SPC->bck_spec->spec[index].ferror = SPC->bck_spec->spec[index].ferror / cfactor;

	  SPC->obj_spec->spec[index].count  = SPC->obj_spec->spec[index].count / cfactor;
	  SPC->obj_spec->spec[index].error  = SPC->obj_spec->spec[index].error / cfactor;
	  SPC->obj_spec->spec[index].flux   = SPC->obj_spec->spec[index].flux / cfactor;
	  SPC->obj_spec->spec[index].ferror = SPC->obj_spec->spec[index].ferror / cfactor;

	  this is the wqrong version!!
	  SPC->fgr_spec->spec[index].count  = SPC->fgr_spec->spec[index].count * cfactor;
	  SPC->fgr_spec->spec[index].error  = SPC->fgr_spec->spec[index].error * cfactor;
	  SPC->fgr_spec->spec[index].flux   = SPC->fgr_spec->spec[index].flux * cfactor;
	  SPC->fgr_spec->spec[index].ferror = SPC->fgr_spec->spec[index].ferror * cfactor;

	  SPC->bck_spec->spec[index].count  = SPC->bck_spec->spec[index].count * cfactor;
	  SPC->bck_spec->spec[index].error  = SPC->bck_spec->spec[index].error * cfactor;
	  SPC->bck_spec->spec[index].flux   = SPC->bck_spec->spec[index].flux * cfactor;
	  SPC->bck_spec->spec[index].ferror = SPC->bck_spec->spec[index].ferror * cfactor;

	  SPC->obj_spec->spec[index].count  = SPC->obj_spec->spec[index].count * cfactor;
	  SPC->obj_spec->spec[index].error  = SPC->obj_spec->spec[index].error * cfactor;
	  SPC->obj_spec->spec[index].flux   = SPC->obj_spec->spec[index].flux * cfactor;
	  SPC->obj_spec->spec[index].ferror = SPC->obj_spec->spec[index].ferror * cfactor;
	  */
	}
    }

  // free the configuration structure
  free_aperture_conf(conf);

  // free the dispersion struct
  free_dispstruct(beam_disp);

  // free the memory for the dispersion
  gsl_vector_free(cdisp);
}

/**
 * Function: get_ipc_lambdas
 * The function computes the wavelengths for a list of x-offsets from the
 * reference point of a certain beam. The wavelength values are
 * returned as a gsl-vector.
 *
 * Parameters:
 * @param actbeam        - the beam
 * @param conf_file_path - the full path to the aXe configuration file
 * @param xvalues        - the list of x-offsets
 *
 * Returns:
 * @return lambdas - the list of wavelength values
 */
gsl_vector *
get_ipc_lambdas(const beam actbeam, char conf_file_path[], gsl_vector *xvalues)
{
  aperture_conf  *conf;

  dispstruct     *disp;
  calib_function *wl_calib;

  gsl_vector *lambdas;

  d_point pixel;
  int for_grism;

  int i;

  // load the configuration file
  conf = get_aperture_descriptor(conf_file_path);

  // check whether it is grism (for_grism=1)
  // or prism (for_grism=0) data
  for_grism = check_for_grism (conf_file_path, actbeam.ID);

  // determine the referencee point position
  pixel.x = actbeam.refpoint.x - conf->refx;
  pixel.y = actbeam.refpoint.y - conf->refy;

  // determine the dispersion at the reference point
  disp = get_dispstruct_at_pos(conf_file_path, for_grism,
			       actbeam.ID,pixel);

  // make a calibration structure from the dispersion
  wl_calib = create_calib_from_gsl_vector(for_grism, disp->pol);

  // convert the x-values to tracelength-values
  abscissa_to_pathlength (actbeam.spec_trace, xvalues);

  // allocate memory for the wavelengths
  lambdas = gsl_vector_alloc(xvalues->size);

  // go over all tracelength values
  for (i=0; i < (int)xvalues->size; i++)
    {
      // determine and store the wavelength for each tracelength
      gsl_vector_set(lambdas, i,
		     wl_calib->func(gsl_vector_get(xvalues, i), wl_calib->order,
				    wl_calib->coeffs));
    }

  // free the configuration structure
  free_aperture_conf(conf);

  // free the memory in the calibration structure
  free_calib(wl_calib);

  // free the dispersion structure
  free_dispstruct(disp);

  // return the wavelengths
  return lambdas;
}

/**
 * Function: get_ipc_cvalues
 * The function computes the intra-pixel correction factors for a certain
 * beam on a set of trace positions specified by their x-offset from
 * the reference point.
 * For every x-offset position the trace positon and its fractional y-pixel
 * (in absolute coordinates) is evaluated. Then the correction factor is
 * determined using the input interpolator.
 *
 * Parameters:
 * @param actbeam - the beam to correct
 * @param ipcorr  - the correction values depending on fractional y-pixel
 * @param xvalues - the list x-offsets from the reference point
 *
 * Returns:
 * @return cvalues - the list of correction values
 */
gsl_vector *
get_ipc_cvalues(const beam actbeam, interpolator *ipcorr, gsl_vector *xvalues)
{
  gsl_vector *cvalues;
  double yfract=0.0;
  int i;

  // allocate memory
  cvalues = gsl_vector_alloc(xvalues->size);

  // go over all x-values
  for (i=0; i < (int)xvalues->size; i++)
    {
      // determine the y-fraction for the x-value
      yfract = get_yfract_for_xvalue(&actbeam, gsl_vector_get(xvalues, i));

      // detyermine and store the correction factor in the array
      gsl_vector_set(cvalues, i, eval_interp(ipcorr, yfract));
    }

  // return the array
  return cvalues;
}

/**
 * Function: get_ipc_xvalues
 * The function computes a list of x-offsets from the reference
 * point of a beam. the regularly space offsets span the range
 * of x-values covered by the pixels of the beam.
 *
 * Parameters:
 * @param actbeam  - the beam to compute the x-offsets for
 *
 * Returns:
 * @return xvalues - the list of x-offsets
 */
gsl_vector *
get_ipc_xvalues(const beam actbeam)
{
  gsl_vector *xvalues;
  d_point xrange;

  int npoints;
  int i;

  // determine the x-range covered by the beam
  xrange =  get_xinterv_from_beam(&actbeam);

  // determine the number of points within the x-range
  npoints = (xrange.y - xrange.x) / XSTEPSIZE + 1;

  // allocate and fill a proper vector
  // with the x-values
  xvalues = gsl_vector_alloc(npoints);
  for (i=0; i < npoints; i++)
    {
      gsl_vector_set(xvalues, i, xrange.x + (double)i * XSTEPSIZE);
    }

  // return the vector with the x-values
  return xvalues;
}

/**
 * Function: get_ipclambda
 * The function creates an interpolator for the intra-pixel correction
 * as a function of wavelength based on this correction as function
 * of fractional y-pixel and the full calibration information
 * on a beam.
 * The interpolator is created after stepping along the beam trace
 * and combining the wavelength values with the correction values
 * at the trace positions.
 *
 * Parameters:
 * @param actbeam        - the beam to find the correction values for
 * @param conf_file_path - the full path to the aXe cofiguration file
 * @param ipcorr         - the correction values as function of y-fraction
 *
 * Returns:
 * @return ipclambda     - the correction values as function of wavelength
 */
interpolator *
get_ipclambda(beam actbeam, char conf_file_path[],
	      interpolator *ipcorr)
{
  gsl_vector *xvalues;
  gsl_vector *cvalues;
  gsl_vector *lambdas;

  double *cv;
  double *lv;

  interpolator *ipclambda;

  int i=0;
  int j=0;

  // get the x values around the reference point
  xvalues = get_ipc_xvalues(actbeam);

  // get the correction factors for these x-values
  cvalues = get_ipc_cvalues(actbeam, ipcorr, xvalues);

  // get the wavelength at the correction factors
  lambdas = get_ipc_lambdas(actbeam, conf_file_path, xvalues);

  // allocate memory for the dependent values
  cv = (double *) malloc(cvalues->size * sizeof(double));
  if (!cv) {
    aXe_message (aXe_M_ERROR, __FILE__, __LINE__,
		 "Memory allocation failed");
  }

  // allocate memory for the independent values
  lv = (double *) malloc(cvalues->size * sizeof(double));
  if (!lv) {
    aXe_message (aXe_M_ERROR, __FILE__, __LINE__,
		 "Memory allocation failed");
  }

  if (gsl_vector_get(lambdas, lambdas->size-1) > gsl_vector_get(lambdas, 0))
    {
      // transfer the values from the
      // gsl vectors to the c-vectors
      for (i = 0; i < (int)cvalues->size; i++)
	{
	  cv[i] = gsl_vector_get(cvalues, i);
	  lv[i] = gsl_vector_get(lambdas, i);
	}
    }
  else
    {
      // transfer the values from the
      // gsl vectors to the c-vectors
      // invert the order from the
      // gsl-vectors
      j = cvalues->size - 1;
      for (i = 0; i < (int)cvalues->size; i++)
	{
	  cv[j] = gsl_vector_get(cvalues, i);
	  lv[j] = gsl_vector_get(lambdas, i);
	  j--;
	}
    }

  // create the interpolator
  ipclambda = create_interp(cvalues->size, FILTER_INTERP_TYPE, lv, cv);

  // free the memory allocated to
  // the vectors
  gsl_vector_free(xvalues);
  gsl_vector_free(cvalues);
  gsl_vector_free(lambdas);

  // return the interpolator
  return ipclambda;
}

/**
 * Function: apply_corr_pet
 * The function applies an intra-pixel sensitivity correction to
 * a list of PET pixels. The correction values are given depending
 * on the wavelength of the pixel, and are applied to the PET pixels
 * in place.
 *
 * Parameters:
 * @param ipclambda - the ipc correction factor as function of wavelength
 * @param PET       - the list of PET pixels to correct
 *
 * Returns:
 * -
 */
void
apply_corr_pet(interpolator *ipclambda, ap_pixel *PET)
{
  int j = 0;

  double cvalue=0.0;

  // go along the PET pixels until you meet
  // the last
  while (PET[j].p_x != -1)
    {
      // get the correction factor
      cvalue = eval_interp(ipclambda, PET[j].lambda);

      // apply the corection factor
      // to the counts and the error
      PET[j].count = PET[j].count / cvalue;
      PET[j].error = PET[j].error / cvalue;

      //      fprintf(stdout, "lambda: %f, correction: %f\n", PET[j].lambda, cvalue);
      j++;
    }

}

/**
 * Function: intpix_corr_pet
 * The function applies the intra-pixel sensitivity correction to
 * a list of PET pixels. The values in the pixels are corrected in
 * place using the correction function, gemoetrical parameters and
 * dispersion relation in the various parameters.
 *
 * Parameters:
 * @param actbeam        - the beam to correct
 * @param conf_file_path - full path to the axe configuration file
 * @param ipcorr         - the ipc correction function
 * @param PET            - the list of PET pixels to correct
 *
 * Returns:
 * -
 */
void
intpix_corr_pet(beam actbeam, char conf_file_path[],
		interpolator *ipcorr, ap_pixel *PET)
{
  interpolator *ipclambda;
  aperture_conf  *conf;

  int for_grism=0;

  // load the configuration file
  conf = get_aperture_descriptor(conf_file_path);

  // check whether it is grism (for_grism=1)
  // or prism (for_grism=0) data
  // give an error if there is a prism solution
  for_grism = check_for_grism (conf_file_path, actbeam.ID);
  if (!for_grism)
    aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		 "intpix_corr_beam: Only grism dispersion solution can be corrected.\n");

  // determine the correction factor versus wavelength
  ipclambda = get_ipclambda(actbeam, conf_file_path, ipcorr);

  print_interp(ipclambda);

  // apply the correction to the PET
  apply_corr_pet(ipclambda, PET);

  // free the configuration structure
  free_aperture_conf(conf);

  // free the memory in the interpolator
  free_interp(ipclambda);
}

/**
 * Function: is_pointlike
 * The function evaluates all criteria to decide whether an object
 * is pointlike or not. In the current implementation an object
 * is considered pointlike if a special OAF is given and the
 * the beam is NOT excluded OR if the spatial extension of the
 * object is smaller than the maximal extension.
 *
 * Parameters:
 * @param actbeam   - the beam to examine
 * @param spec_OAF  - indicates a special OAF file
 * @param max_ext   - the maximal allowed extension
 *
 * Returns:
 * @return point_like - pointer to the opened fits file
 */
int
is_pointlike(beam actbeam, int spec_OAF, double max_ext)
{
  int point_like=0;

  int OAF_crit=0;
  int EXT_crit=0;

  // check whether a special OAF is given
  // and the beam should NOT be ignored
  if (spec_OAF && !actbeam.ignore)
    OAF_crit=1;

  // check whether the maximal extension
  // is given and the object extension
  // is smaller;
  // also check whether non-default values
  // for the object width are set
  if (actbeam.awidth > -1.0 && actbeam.bwidth > -1.0 && max_ext && actbeam.awidth < max_ext && actbeam.bwidth < max_ext)
    EXT_crit=1;

  // set to pointlike if at least
  // one of the criteria is set
  if (OAF_crit || EXT_crit)
    point_like=1;

  // return the result
  return point_like;
}

/**
 * Function: get_SPC_opened
 * The function opens an existing SPC file and returns the pointer
 * to it.
 * As of now, the mode of opening it is automatically READWRITE.
 * Later on a differentiation using the free parameter "mode"
 * might be added  to generalize the function.
 *
 * Parameters:
 * @param SPCname - name of the SPC file
 * @param mode    - mode to open it (not yet used)
 *
 * Returns:
 * @return SPC_ptr - pointer to the opened fits file
 *
 */
fitsfile *
get_SPC_opened(char SPCname[], int mode)
{
    fitsfile *SPC_ptr;
    int f_status=0;

  // Open the OPET file for reading/writing
  fits_open_file (&SPC_ptr, SPCname, READWRITE, &f_status);
  if (f_status)
    {
      ffrprt (stdout, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "aXe_INTPIXCORR: Could not open file: %s\n",
		   SPCname);
    }

  // return the pointer to the fits file
  return SPC_ptr;
}

/**
 * Function: get_ALL_from_next_in_SPC
 * The function creates and fills a full spectrum structure with
 * the content of a SPC table extension. This is only done when,
 * according to the beam ID, this beam should be corrected.
 * For extensions with higher order beams which are not corrected
 * an emply structure is returned.
 *
 * Parameters:
 * @param SPCn_ptr - pointer to the opened SPC file
 * @param aperID   - pointer to aperture identification number
 * @param beamID   - pointer to beam identification number
 *
 * Returns:
 * @return SPC - the full spectrum structure
 */
full_spectr *
get_ALL_from_next_in_SPC(fitsfile *SPC_ptr, int *aperID, int *beamID)
{
  int f_status=0, hdutype;

  long tmp;
  //long nrows=0;
  char comment[FLEN_COMMENT];

  full_spectr *SPC;

  fits_movrel_hdu (SPC_ptr, 1, &hdutype, &f_status);

  if (f_status)
    {
      *aperID = -1;
      *beamID = -1;
      SPC = NULL;
      return SPC;
    }

  // read the beam ID number
  fits_read_key_lng (SPC_ptr, "BEAMID", &tmp, comment, &f_status);
  if (f_status)
    {
      ffrprt (stderr, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_ALL_from_next_in_SPC: Error getting index keyword OBJECTID");
    }
  *beamID = (int)tmp;

  // check whether this beam should be correct
  if (*beamID > CORRMAX)
    {
      // set it to NULL and return
      SPC = NULL;
      //      fprintf (stdout, "aXe_PETFF: Skipping beam: %c.\n", BEAM(*beamID));
      return SPC;
    }

  // the beam shall be corrected and
  // first must be read in
  SPC = (full_spectr *) malloc (sizeof (full_spectr ));

  // transfer the beam ID
  SPC->beamID = (int)tmp;

  // read the aperture number
  fits_read_key_lng (SPC_ptr, "OBJECTID", &tmp, comment, &f_status);
  if (f_status)
    {
      ffrprt (stderr, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_ALL_from_next_in_SPC: Error getting index keyword OBJECTID");
    }
  // transfer the aperture ID
  *aperID = (int)tmp;
  SPC->aperID = (int)tmp;

  // Get the number of rows
  fits_get_num_rows (SPC_ptr, &tmp, &f_status);
  if (f_status) {
    ffrprt (stderr, f_status);
    aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		 "get_ALL_from_next_in_SPC: "
		 "Could not determine the number of rows in"
		 " correction function table!");
  }
  SPC->nelems = (int)tmp;

  // load the background subtracted object spectrum
  SPC->obj_spec = get_spectrum_from_SPC(SPC_ptr, "COUNT", "ERROR", SPC->nelems);

  // load the total object spectrum
  SPC->fgr_spec = get_spectrum_from_SPC(SPC_ptr, "TCOUNT", "TERROR", SPC->nelems);

  // load the background spectrum
  SPC->bck_spec = get_spectrum_from_SPC(SPC_ptr, "BCOUNT", "BERROR", SPC->nelems);

  // return the filled structure
  return SPC;
}

/**
 * Function: free_full_spectr
 * The function frees the memory allocated in a full
 * spectrum structure.
 *
 * Parameters:
 * @param SPC - the full spectrum structure
 *
 * Returns:
 * @return -
 */
void
free_full_spectr(full_spectr *SPC)
{
  // free the three spectra in the
  // full spectrum structure
  free_spectrum(SPC->obj_spec);
  free_spectrum(SPC->fgr_spec);
  free_spectrum(SPC->bck_spec);

  // free the full spectrum
  free(SPC);

  // set the structure to NULL
  SPC = NULL;
}

/**
 * Function: get_spectrum_from_SPC
 * The function fills a spectrum structure with data in an SPC extension.
 * There is more data in an SPC extension than can be stored in a
 * spectrum structure. Two parameters in this function specify partly which
 * data should be loaded.
 *
 * Parameters:
 * @param SPC_ptr   - pointer to the SPC extension
 * @param count_col - column to load for 'count'
 * @param error_col - column to load for 'error'
 * @param nelems    - the number ol elements in the spectrum
 *
 * Returns:
 * @return act_spec - the filled spectrum structure
 */
spectrum *
get_spectrum_from_SPC(fitsfile *SPC_ptr, char count_col[],
		      char error_col[], int nelems)
{
  int index=0;

  spectrum *act_spec;

  double *count;
  double *lambda;
  double *error;
  double *flux;
  double *ferror;
  double *weight;
  double *contam;
  long   *dq;

  // allocate the spectrum
  act_spec = allocate_spectrum (nelems);

  // transfer the column entries into a vector
  lambda = get_dcolumn_from_SPC_opened(SPC_ptr, "LAMBDA", nelems);
  count  = get_dcolumn_from_SPC_opened(SPC_ptr, count_col, nelems);
  error  = get_dcolumn_from_SPC_opened(SPC_ptr, error_col, nelems);
  flux   = get_dcolumn_from_SPC_opened(SPC_ptr, "FLUX", nelems);
  ferror = get_dcolumn_from_SPC_opened(SPC_ptr, "FERROR", nelems);
  weight = get_dcolumn_from_SPC_opened(SPC_ptr, "WEIGHT", nelems);
  contam = get_dcolumn_from_SPC_opened(SPC_ptr, "CONTAM", nelems);
  dq     = get_lcolumn_from_SPC_opened(SPC_ptr, "DQ", nelems);

  // transfer the data from the vector into the
  // spectrum structure
  for (index = 0; index < nelems; index++)
    {
      act_spec->spec[index].lambda_mean = lambda[index];
      act_spec->spec[index].count       = count[index];
      act_spec->spec[index].error       = error[index];
      act_spec->spec[index].flux        = flux[index];
      act_spec->spec[index].ferror      = ferror[index];
      act_spec->spec[index].weight      = weight[index];
      act_spec->spec[index].contam      = contam[index];
      act_spec->spec[index].dq          = dq[index];
    }

  // free the arrays
  free(lambda);
  free(count);
  free(error);
  free(flux);
  free(ferror);
  free(weight);
  free(contam);
  free(dq);

  // return the spectrum
  return act_spec;
}

/**
 * Function: get_dcolumn_from_SPC_opened
 * The function reads the values from a double column specified
 * by its name into a array of doubles. The array is returned.
 *
 * Parameters:
 * @param SPC_ptr   - pointer to the SPC extension
 * @param count_col - name of the column to load
 * @param nelems    - number of elements in the column
 *
 * Returns:
 * @return values - pointer to a filled array
 */
double *
get_dcolumn_from_SPC_opened(fitsfile *SPC_ptr, char colname[], int nelems)
{
  int colnum=0;
  int anynul;
  int f_status=0;
  double *values;

  // allocate the return array;
  // give an error if allocation fails
  values = (double *) malloc(nelems * sizeof(double));
  if (!values) {
    aXe_message (aXe_M_ERROR, __FILE__, __LINE__,
		 "Memory allocation failed");
  }

  // get the desired column number;
  // give an error if the column name
  // can not be read
  fits_get_colnum (SPC_ptr, CASEINSEN, colname, &colnum, &f_status);
  if (f_status)
    {
      ffrprt (stderr, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_dcolumn_from_SPC_opened: "
		   "Could not determine column %s in "
		   " table!\n", colname);
    }

  // read all data in the column
  fits_read_col (SPC_ptr, TDOUBLE, colnum, 1, 1, nelems, NULL, values,
                    &anynul, &f_status);
  if (f_status)
    {
      ffrprt (stderr, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_dcolumn_from_SPC_opened: "
		   "Could not read column %s"
		   " from BINARY table!",colname);
    }

  // return the filled vector
  return values;
}

/**
 * Function: get_lcolumn_from_SPC_opened
 * The function reads the values from a column with long specified
 * by its name into a array of type long. The array is returned.
 *
 * Parameters:
 * @param SPC_ptr   - pointer to the SPC extension
 * @param count_col - name of the column to load
 * @param nelems    - number of elements in the column
 *
 * Returns:
 * @return values - pointer to a filled array
 */
long *
get_lcolumn_from_SPC_opened(fitsfile *SPC_ptr, char colname[], int nelems)
{
  int colnum=0;
  int anynul;
  int f_status=0;
  long *values;

  // allocate the return array;
  // give an error if allocation fails
  values = (long *) malloc(nelems * sizeof(long));
  if (!values) {
    aXe_message (aXe_M_ERROR, __FILE__, __LINE__,
		 "Memory allocation failed");
  }

  // get the desired column number;
  // give an error if the column name
  // can not be read
  fits_get_colnum (SPC_ptr, CASEINSEN, colname, &colnum, &f_status);
  if (f_status)
    {
      ffrprt (stderr, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_dcolumn_from_SPC_opened: "
		   "Could not determine column %s in "
		   " table!\n", colname);
    }

  // read all data in the column
  fits_read_col (SPC_ptr, TLONG, colnum, 1, 1, nelems, NULL, values,
                    &anynul, &f_status);
  if (f_status)
    {
      ffrprt (stderr, f_status);
      aXe_message (aXe_M_FATAL, __FILE__, __LINE__,
		   "get_dcolumn_from_SPC_opened: "
		   "Could not read column %s"
		   " from BINARY table!",colname);
    }

  // return the filled array
  return values;
}

/**
 * Function: create_nlincor
 * The function creates an interpolator for the nonlinearity
 * correction applied to NICMOS data.
 *
 * Parameters:
 *
 * Returns:
 * @return nlincorr - the interpolator created
 */
interpolator *
create_nlincor()
{
  interpolator *nlincorr;
  /*
  double x[14] = {8250.0, 8750.0, 9250.0, 9750.0, 11000.0, 12000.0, 13000.0, 14000.0, 15000.0, 16000.0, 17000.0, 18000.0, 19000.0, 20000.0};
  double y[14] = {  .069,   .057,   .052,   .050,    .049,    .048,      .041,    .023,    .013,    .008,    .004,    .0,      .0,      .0};
  */
  double *xx;
  double *yy;

  xx = (double *) malloc(14 * sizeof(double));
  yy = (double *) malloc(14 * sizeof(double));

  xx[0]  = 8250.0;
  xx[1]  = 8750.0;
  xx[2]  = 9250.0;
  xx[3]  = 9750.0;
  xx[4]  = 11000.0;
  xx[5]  = 12000.0;
  xx[6]  = 13000.0;
  xx[7]  = 14000.0;
  xx[8]  = 15000.0;
  xx[9]  = 16000.0;
  xx[10] = 17000.0;
  xx[11] = 18000.0;
  xx[12] = 19000.0;
  xx[13] = 20000.0;

  yy[0]  = .069;
  yy[1]  = .057;
  yy[2]  = .052;
  yy[3]  = .050;
  yy[4]  = .049;
  yy[5]  = .048;
  yy[6]  = .041;
  yy[7]  = .023;
  yy[8]  = .013;
  yy[9]  = .008;
  yy[10] = .004;
  yy[11] = .0;
  yy[12] = .0;
  yy[13] = .0;

  // create the interpolator
  nlincorr = create_interp(14, NLINCORR_INTERP_TYPE, xx, yy);

  // return the interpolator
  return nlincorr;
}

/**
 * Function: nlin_corr_beam
 *
 * Parameters:
 * @param nlincorr       - the interpolator with the correction factor
 * @param SPC            - the full spectrum to correct
 *
 * Returns:
 * @return -
 */
void
nlin_corr_beam(interpolator *nlincorr, double adcgain, full_spectr *SPC)
{
  int            index=0;
  //  int            for_grism=1;

  double         cfactor;
  double         cps;

  double         lambda;

  for (index=0; index < SPC->nelems; index++)
    {
      // get the independent spectral values,
      // the wavelenth and the cps value
      lambda = SPC->obj_spec->spec[index].lambda_mean;
      cps    = SPC->obj_spec->spec[index].count;

      // check whether the spectral element
      // is not corrupt
      if (!gsl_isnan (lambda) && lambda)
	{
	  if (cps > 0.0)
	    // compute the correction factor
	    cfactor = get_nlin_corr(nlincorr, lambda, cps/adcgain);
	  else
	    // make a dummy factor
	    cfactor = 1.0;
	  // correct what should be corrected
	  SPC->fgr_spec->spec[index].count  = SPC->fgr_spec->spec[index].count / cfactor;
	  SPC->fgr_spec->spec[index].error  = SPC->fgr_spec->spec[index].error / cfactor;
	  SPC->fgr_spec->spec[index].flux   = SPC->fgr_spec->spec[index].flux / cfactor;
	  SPC->fgr_spec->spec[index].ferror = SPC->fgr_spec->spec[index].ferror / cfactor;

	  SPC->bck_spec->spec[index].count  = SPC->bck_spec->spec[index].count / cfactor;
	  SPC->bck_spec->spec[index].error  = SPC->bck_spec->spec[index].error / cfactor;
	  SPC->bck_spec->spec[index].flux   = SPC->bck_spec->spec[index].flux / cfactor;
	  SPC->bck_spec->spec[index].ferror = SPC->bck_spec->spec[index].ferror / cfactor;

	  SPC->obj_spec->spec[index].count  = SPC->obj_spec->spec[index].count / cfactor;
	  SPC->obj_spec->spec[index].error  = SPC->obj_spec->spec[index].error / cfactor;
	  SPC->obj_spec->spec[index].flux   = SPC->obj_spec->spec[index].flux / cfactor;
	  SPC->obj_spec->spec[index].ferror = SPC->obj_spec->spec[index].ferror / cfactor;
	}
    }

}

/*
 * Function: get_nlin_corr
 *
 * Parameters:
 * @param nlincorr - the interpolator with the parameter
 * @param lamda    - the full spectrum to correct
 * @param cps      - the count rate
 *
 * Returns:
 * @return cfactor - the correction factor
 */
double
get_nlin_corr(interpolator *nlincorr, const double lambda, const double cps)
{
  double cfactor;
  double bbb;

  // evaluate the parameter
  bbb = eval_interp(nlincorr, lambda);

  // compute the correction factor
  cfactor = 1.0 - 2.0 * bbb + bbb * log10(cps);

  // return the correction factor
  return cfactor;
}


/* Modified from SUPOLAR, Copyright (c) Colorado School of Mines, 2011.*/
/* All rights reserved.                                                */

/* SUPOLAR_PS: $Revision: 1.9 $ ; $Date: Sat, 23 Jul 2016 00:10:26 -0700 */

#include <omp.h>
#include <stdio.h>
#include <gsl/gsl_statistics_float.h>
#include "su.h"
#include "segy.h"
#include "header.h"

/*********************** self documentation *****************************/
char *sdoc[] = {
"                                                                       ",
" SUPOLAR_PS - POLarization analysis of three-component data            ",
"            - to provide P and S arrivals                              ",
"            - Modification of SUPOLAR                                  ",
"                                                                       ",
" supolar_PS <stdin [optional parameters]                               ",
"                                                                       ",
" Required parameters:                                                  ",
"    -ntr (number of traces) needs to be specified in the SU file       ",
"                                                                       ",
" Optional parameters:                                                  ",
"    dt=(from header)  time sampling interval in seconds                ",
"    wl=0.1            correlation window length in seconds             ",
"    win=boxcar        correlation window shape, choose \"boxcar\",     ",
"                      \"hanning\", \"bartlett\", or \"welsh\"          ",
"    file=polar        base of output file name(s)                      ",
"    rl=1              1 = rectilinearity evaluating 2 eigenvalues,     ",
"                      2, 3 = rectilinearity evaluating 3 eigenvalues   ",
"    rlq=1.0           contrast parameter for rectilinearity            ",
"    dir=1             1 = 3 components of direction of polarization    ",
"                          (the only three-component output file)       ",
"    tau=0             1 = global polarization parameter                ",
"    ellip=0           1 = principal, subprincipal, and transverse      ",
"                          ellipticities e21, e31, and e32              ",
"    pln=0             1 = planarity measure                            ",
"    f1=0              1 = flatness or oblateness coefficient           ",
"    l1=0              1 = linearity coefficient                        ",
"    amp=0             1 = amplitude parameters: instantaneous,         ",
"                          quadratic, and eigenresultant ir, qr, and er ",
"    theta=0           1, 2, 3 = incidence angle of principal axis      ",
"    phi=0             1, 2, 3 = horizontal azimuth of principal axis   ",
"    angle=rad         unit of angles theta and phi, choose \"rad\",    ",
"                      \"deg\", or \"gon\"                              ",
"    all=0             1, 2, 3 = set all output flags to that value     ",
"    verbose=0         1 = echo additional information                  ",
"    kwl=5*(1/dt)      kurtosis window length, samples                  ",
"                                                                       ",
" Notes:                                                                ",
"    Three adjacent traces are considered as one three-component        ",
"    dataset.                                                           ",
"    Correct calculation of angles theta and phi requires the first of  ",
"    these traces to be the vertical component, followed by the two     ",
"    horizontal components (e.g. Z, N, E, or Z, inline, crossline).     ",
"    Significant signal energy on Z is necessary to resolve the 180 deg ",
"    ambiguity of phi (options phi=2,3 only).                           ",
"                                                                       ",
"    Each calculated polarization attribute is written into its own     ",
"    SU file. These files get the same base name (set with \"file=\")   ",
"    and the parameter flag as an extension (e.g. polar.rl).            ",
"                                                                       ",
"    In case of a tapered correlation window, the window length wl may  ",
"    have to be increased compared to the boxcar case, because of their ",
"    smaller effective widths (Bartlett, Hanning: 1/2, Welsh: 1/3).     ",
"                                                                       ",
" Range of values:                                                      ",
"    parameter     option  interval                                     ",
"    rl            1, 2    0.0 ... 1.0   (1.0: linear polarization)     ",
"    rl            3      -1.0 ... 1.0                                  ",
"    tau, l1       1       0.0 ... 1.0   (1.0: linear polarization)     ",
"    pln, f1       1       0.0 ... 1.0   (1.0: planar polarization)     ",
"    e21, e31, e32 1       0.0 ... 1.0   (0.0: linear polarization)     ",
"    theta         1      -pi/2... pi/2  rad                            ",
"    theta         2, 3    0.0 ... pi/2  rad                            ",
"    phi           1      -pi/2... pi/2  rad                            ",
"    phi           2      -pi  ... pi    rad   (see notes above)        ",
"    phi           3       0.0 ... 2 pi  rad   (see notes above)        ",
"                                                                       ",
"                                                                       ",
NULL};

/*
 * Author: Nils Maercklin,
 *         GeoForschungsZentrum (GFZ) Potsdam, Germany, 1998-2001.
 *         E-mail: nils@gfz-potsdam.de
 *
 *
 * References:
 *    Jurkevics, A., 1988: Polarization analysis of three-component
 *       array data. Bulletin of the Seismological Society of America,
 *       vol. 78, no. 5.
 *    Kanasewich, E. R., 1981: Time Sequence Analysis in Geophysics.
 *       The University of Alberta Press.
 *    Kanasewich, E. R., 1990: Seismic Noise Attenuation.
 *       Handbook of Geophysical Exploration, Pergamon Press, Oxford.
 *    Meyer, J. H. 1988: First Comparative Results of Integral and
 *       Instantaneous Polarization Attributes for Multicomponent Seismic
 *       Data. Institut Francais du Petrole.
 *    Press, W. H., Teukolsky, S. A., Vetterling, W. T., and Flannery, B. P.
 *       1996: Numerical Recipes in C - The Art of Scientific Computing.
 *       Cambridge University Press, Cambridge.
 *    Samson, J. C., 1973: Description of the Polarisation States of Vector
 *       Processes: Application to ULF Electromagnetic Fields.
 *       Geophysical Journal vol. 34, p. 403-419.
 *    Sheriff, R. E., 1991: Encyclopedic Dictionary of Exploration
 *       Geophysics. 3rd ed., Society of Exploration Geophysicists, Tulsa.
 *
 * Trace header fields accessed: ns, dt
 * Trace header fields modified: none
 */
/**************** end self doc *******************************************/

/* prototypes of functions used internally */
float calc_ellip(float *d, int d1, int d2);
float calc_er(float *d);
float calc_f1(float *d);
float calc_l1(float *d);
float calc_phi(float **v, int opt);
float calc_plan(float *d);
float calc_rl(float d[], float q, int opt);
float calc_tau(float *d);
float calc_theta(float **v, int opt);
float covar(float *data1, float *data2, int istart, int iwl, float *w);
float calc_pfilt(float rl, float theta);
float calc_sfilt(float rl, float theta);
float kurtosiswindow(float *data, float *data_kwl, int it, int kwl, int nt);
void ampparams(float **indata, float *data_ir, float *data_qr, int nt, int iwl);
void calc_dir(float **data3c_dir, float **v, int it);
void calc_window(float *w, int iwl, int iwin);
void fputdata(FILE *fileptr, FILE *headerptr, float outdata[], int nt);
void fputdata3c(FILE *fileptr, FILE *headerptr, float **outdata3c, int nt);

/* window shape identifiers */
#define WBOXCAR 0
#define WBARTLETT 1
#define WHANNING 2
#define WWELSH 3

segy tr;                /* SEG-Y record (trace and header) */

int main(int argc, char **argv)
{
    /* output file pointers */
    FILE *rlfp=NULL, *taufp=NULL, *e21fp=NULL;
    FILE *e31fp=NULL, *e32fp=NULL, *plnfp=NULL;
    FILE *f1fp=NULL, *l1fp=NULL;
    FILE *thetafp=NULL, *phifp=NULL, *dirfp=NULL;
    FILE *erfp=NULL, *irfp=NULL, *qrfp=NULL;
    FILE *headerfp=NULL, *pfiltfp=NULL, *sfiltfp=NULL;
    FILE *nfiltfp=NULL, *efiltfp=NULL;
 //   FILE *pkur=NULL, *skur=NULL;
        /* temporary file for trace headers */
			 /* (one 3C station only) */

    char *file=NULL;         /* base of output file name(s) */
    char *fname=NULL;        /* complete output file name */
    char *angle=NULL;        /* unit used for angles theta and phi */
    char *win=NULL;          /* shape of used time window */
    float fangle=0.0;   /* unit conversion factor applied to angles theta and phi */
    int iwin=0;           /* time window shape identifier */

    /* flags (see selfdoc) */
    int rl,theta,phi,tau,ellip,pln,f1,l1,dir,amp,verbose,all;

    int i,j,icomp;      /* indices for components (in loops) */
    int it;             /* index for time sample in main loop */
    int iwl;            /* correlation window length in samples */
    int nstat;          /* number of 3-component datasets */
    int nt;             /* number of time samples in one trace */
//    int kwl;            /* kurtosis window length in seconds */

    float **data3c;     /* three-component data ([1..3][0..nt-1]) */
    float **a;          /* covariance matrix (a[1..3][1..3]) */
    float **v;          /* eigenvectors of covariance matrix (v[1..3][1..3]) */
    float *d;           /* the corresponding eigenvalues (d[1..3]) */
    float *w;           /* time window weights for correlation window */
    float dt;           /* sampling interval in seconds */
    float rlq;          /* contrast factor of rectilinearity */
    float wl;           /* correlation window length in seconds */

    float *data_e21=NULL;    /* main ellipticity */
    float *data_e31=NULL;    /* second ellipticity */
    float *data_e32=NULL;    /* transverse ellipticity */
    float *data_er=NULL;     /* eigenresultant */
    float *data_f1=NULL;     /* flatness coefficient */
    float *data_ir=NULL;     /* instantaneous resultant */
    float *data_l1=NULL;     /* linearity coefficient */
    float *data_phi=NULL;    /* horizontal azimuth phi */
    float *data_pln=NULL;    /* planarity */
    float *data_qr=NULL;     /* quadratic resultant */
    float *data_rl=NULL;     /* rectilinearity factor */
    float *data_tau=NULL;    /* polarization parameter tau */
    float *data_theta=NULL;  /* inclination angle theta */
    float *data_pfilt=NULL;  /* P (vertical) polarization filter */
    float *data_sfilt=NULL;  /* S (horizontal) polarization filter */
    float *data_zfilt=NULL;  /* Z Filtered Trace */
    float *data_nfilt=NULL;  /* N Filtered Trace */
    float *data_efilt=NULL;  /* E Filtered Trace */
//    float *data_kwl=NULL;   /* Data for Kurtosis Window */
//    float *data_pkur=NULL;   /* Kurtosis detector for P */
//    float *data_skur=NULL;   /* Kurtosis detector for S */
    float **data3c_dir=NULL; /* 3 components of direction of polarization ([1..3][0..nt-1]) */

    /* initialize */
    initargs(argc, argv);
    requestdoc(1);

    /* get info from first trace */
    if(!gettr(&tr)) err("can't get first trace");
    nt = tr.ns;

    /* get parameters ... */
    if (!getparstring("file", &file)) file="polar";
    if (!getparstring("angle", &angle)) angle="rad";
    if (!getparstring("win", &win)) win="boxcar";
    if (!getparfloat("wl", &wl)) wl = 0.1;
    if (!getparfloat("dt", &dt)) dt = ((double) tr.dt)/1000000.0;
    if (!getparfloat("rlq", &rlq)) rlq = 1.0;
    if (!getparint("verbose", &verbose)) verbose = 0;
//    if (!getparint("kwl", &kwl)) kwl = 5 * ((int) 1/dt);

    /* ... and output flags */
    if (!getparint("all", &all))       all = 0;
    if (!getparint("rl", &rl))          rl = (all) ? all : 1;
    if (!getparint("dir", &dir))       dir = (all) ? 1 : 1;
    if (!getparint("theta", &theta)) theta = (all) ? all : 0;
    if (!getparint("phi", &phi))       phi = (all) ? all : 0;
    if (!getparint("tau", &tau))       tau = (all) ? 1 : 0;
    if (!getparint("ellip", &ellip)) ellip = (all) ? 1 : 0;
    if (!getparint("pln", &pln))       pln = (all) ? 1 : 0;
    if (!getparint("f1", &f1))          f1 = (all) ? 1 : 0;
    if (!getparint("l1", &l1))          l1 = (all) ? 1 : 0;
    if (!getparint("amp", &amp))       amp = (all) ? 1 : 0;

    checkpars();

    /* get time window shape */
    if      (STREQ(win, "boxcar"))   iwin=WBOXCAR;
    else if (STREQ(win, "bartlett")) iwin=WBARTLETT;
    else if (STREQ(win, "hanning"))  iwin=WHANNING;
    else if (STREQ(win, "welsh"))    iwin=WWELSH;
    else err("unknown win=%s", win);

    /* get unit conversion factor for angles */
    if      (STREQ(angle, "rad")) fangle=1.0;
    else if (STREQ(angle, "deg")) fangle=180.0/PI;
    else if (STREQ(angle, "gon")) fangle=200.0/PI;
    else err("unknown angle=%s", angle);

    /* convert seconds to samples */
    if (!dt) {
        dt = 0.004;
        warn("dt not set, assuming dt=0.004");
    }

    iwl = NINT(wl/dt);

    /* data validation */
    if (iwl<1) err("wl=%g must be positive", wl);
    if (iwl>nt) err("wl=%g too long for trace", wl);
    if (!strlen(file)) err("file= not set and default overridden");

    /* echo some information */
    if (verbose && (theta || phi)) warn("computing angles in %s", angle);
    if (verbose) warn("%s window length = %d samples\n", win, iwl);

    if (rl && theta) warn("computing filtered phase");

    /* open temporary file for trace headers */
    headerfp = etmpfile();

    /* set filenames and open files */
    fname = malloc( strlen(file)+7 );
    sprintf(fname, "%s.rl", file);    if (rl) rlfp = efopen(fname, "w");
    sprintf(fname, "%s.theta", file); if (theta) thetafp = efopen(fname, "w");
    sprintf(fname, "%s.phi", file);   if (phi) phifp = efopen(fname, "w");
    sprintf(fname, "%s.tau", file);   if (tau) taufp = efopen(fname, "w");
    sprintf(fname, "%s.e21", file);   if (ellip) e21fp = efopen(fname, "w");
    sprintf(fname, "%s.e31", file);   if (ellip) e31fp = efopen(fname, "w");
    sprintf(fname, "%s.e32", file);   if (ellip) e32fp = efopen(fname, "w");
    sprintf(fname, "%s.pln", file);   if (pln) plnfp = efopen(fname, "w");
    sprintf(fname, "%s.f1", file);    if (f1) f1fp = efopen(fname, "w");
    sprintf(fname, "%s.l1", file);    if (l1) l1fp = efopen(fname, "w");
    sprintf(fname, "%s.dir", file);   if (dir) dirfp = efopen(fname, "w");
    sprintf(fname, "%s.er", file);    if (amp) erfp = efopen(fname, "w");
    sprintf(fname, "%s.ir", file);    if (amp) irfp = efopen(fname, "w");
    sprintf(fname, "%s.qr", file);    if (amp) qrfp = efopen(fname, "w");
    sprintf(fname, "%s.pfilt", file); if (rl && theta) pfiltfp = efopen(fname, "w");
    sprintf(fname, "%s.sfilt", file); if (rl && theta) sfiltfp = efopen(fname, "w");
    sprintf(fname, "%s.nfilt", file); if (rl && theta) nfiltfp = efopen(fname, "w");
    sprintf(fname, "%s.efilt", file); if (rl && theta) efiltfp = efopen(fname, "w");
  //  sprintf(fname, "%s.pkur", file); if (rl && theta) pkur = efopen(fname, "w");
  //  sprintf(fname, "%s.skur", file); if (rl && theta) skur = efopen(fname, "w");
    free(fname);

    /* allocate space for input data and analysis matrices */
    /* index ranges used here: data3c[1..3][0..nt-1],      */
    /* a[1..3][1..3], v[1..3][1..3], d[1..3]               */
    data3c = ealloc2float(nt,3); data3c-=1;
    a = ealloc2float(3,3); a[0]-=1; a-=1;
    v = ealloc2float(3,3); v[0]-=1; v-=1;
    d = ealloc1float(3); d-=1;

    /* calculate time window weights */
    w = ealloc1float(iwl);
    memset((void *) w, 0, iwl*FSIZE);
    calc_window(w, iwl, iwin);

    /* allocate and zero out space for output data */
    if (rl) {
        data_rl = ealloc1float(nt);
        memset((void *) data_rl, 0, nt*FSIZE);
    }
    if (theta) {
        data_theta = ealloc1float(nt);
        memset((void *) data_theta, 0, nt*FSIZE);
    }
    if (phi) {
        data_phi = ealloc1float(nt);
        memset((void *) data_phi, 0, nt*FSIZE);
    }
    if (tau) {
        data_tau = ealloc1float(nt);
        memset((void *) data_tau, 0, nt*FSIZE);
    }
    if (ellip) {
        data_e21 = ealloc1float(nt);
        data_e31 = ealloc1float(nt);
        data_e32 = ealloc1float(nt);
        memset((void *) data_e21, 0, nt*FSIZE);
        memset((void *) data_e31, 0, nt*FSIZE);
        memset((void *) data_e32, 0, nt*FSIZE);
    }
    if (pln) {
        data_pln = ealloc1float(nt);
        memset((void *) data_pln, 0, nt*FSIZE);
    }
    if (f1) {
        data_f1 = ealloc1float(nt);
        memset((void *) data_f1, 0, nt*FSIZE);
    }
    if (l1) {
        data_l1 = ealloc1float(nt);
        memset((void *) data_l1, 0, nt*FSIZE);
    }
    if (amp) {
        data_er = ealloc1float(nt);
        data_ir = ealloc1float(nt);
        data_qr = ealloc1float(nt);
        memset((void *) data_er, 0, nt*FSIZE);
        memset((void *) data_ir, 0, nt*FSIZE);
        memset((void *) data_qr, 0, nt*FSIZE);
    }
    if (dir) {
        data3c_dir = ealloc2float(nt,3); data3c_dir-=1;
        for (i=1;i<=3;i++) memset((void *) data3c_dir[i], 0, nt*FSIZE);
    }
    if (rl && theta) {
        data_pfilt = ealloc1float(nt);
        memset((void *) data_pfilt, 0, nt*FSIZE);
        data_sfilt = ealloc1float(nt);
        memset((void *) data_pfilt, 0, nt*FSIZE);
/*        data_3Cfilt = ealloc2float(nt,3);
        for (i=1;i<=3;i++) memset((void *) data_3Cfilt[i], 0, nt*FSIZE); */
        data_zfilt = ealloc1float(nt);
        data_nfilt = ealloc1float(nt);
        data_efilt = ealloc1float(nt);
        memset((void *) data_zfilt, 0, nt*FSIZE);
        memset((void *) data_nfilt, 0, nt*FSIZE);
        memset((void *) data_efilt, 0, nt*FSIZE);

 //       data_pkur = ealloc1float(nt);
//        data_skur = ealloc1float(nt);
//        memset((void *) data_pkur, 0, nt*FSIZE);
//        memset((void *) data_skur, 0, nt*FSIZE);
        /* Allocate data for kurtosis window arrays */
//        data_kwl = ealloc1float(iwl);
//        memset((void *) data_kwl, 0, kwl*FSIZE);
    }

/* ************************ BEGIN CALCULATION ******************************* */

    /* loop over traces */
    icomp=0;
    nstat=0;
    // Need to convert this do while loop into a for loop so as to be easier
    // to parallelize
    warn("Trace Start Time: %d %d %d %d %d", tr.year, tr.day, tr.hour, tr.minute, tr.sec);
    do {
        /* store trace header in temporary file and read data */
        efwrite(&tr, HDRBYTES, 1, headerfp);
        icomp++;
        memcpy((void *)data3c[icomp], (const void *) tr.data, nt*FSIZE);

        /* process 3-component dataset */
        if (icomp==3) {
            erewind(headerfp);
            icomp = 0;
            nstat++;
            if (verbose)
                fprintf(stderr,"%s: analyzing station %d \r",argv[0], nstat);

            /* start loop over samples */

            for (it=iwl/2;it<nt-iwl/2;it++) {
                //warn("Sample %d", it);
                /* covariance matrix */

                for (i=1;i<=3;i++)
                {
                    for (j=i;j<=3;j++)
                    {
                        a[i][j]=a[j][i]=covar(data3c[i], data3c[j], it-iwl/2, iwl, w);
                    }
                }

                /* compute eigenvalues and vectors */
                eig_jacobi(a,d,v,3);
                sort_eigenvalues(d,v,3);

                /* polarization parameters */
                if (rl) data_rl[it]=calc_rl(d,rlq,rl);
                if (theta) data_theta[it]=calc_theta(v, theta) * fangle;
                if (phi) data_phi[it]=calc_phi(v, phi) * fangle;
                if (tau) data_tau[it]=calc_tau(d);
                if (ellip) {
                    data_e21[it]=calc_ellip(d,2,1);
                    data_e31[it]=calc_ellip(d,3,1);
                    data_e32[it]=calc_ellip(d,3,2);
                }
                if (pln) data_pln[it]=calc_plan(d);
                if (f1) data_f1[it]=calc_f1(d);
                if (l1) data_l1[it]=calc_l1(d);
                if (amp) data_er[it]=calc_er(d);
                if (dir) calc_dir(data3c_dir,v,it);
                if (rl && theta)
                {
                    data_zfilt[it] = data3c[1][it] * calc_pfilt(rl, theta);
                    data_nfilt[it] = data3c[2][it] * calc_sfilt(rl, theta);
                    data_efilt[it] = data3c[3][it] * calc_sfilt(rl, theta);
                    data_pfilt[it] = data_zfilt[it];
                    data_sfilt[it] = (data_nfilt[it] + data_efilt[it]) / 2;
//                    data_pkur[it] = kurtosiswindow(data_pfilt,data_kwl,it - kwl/2,kwl,nt);
//                    data_skur[it] = kurtosiswindow(data_sfilt,data_kwl,it - kwl/2,kwl,nt);
               }

            } /* end loop over samples */

            /* compute amplitude parameters */
            if (amp) ampparams(data3c, data_ir, data_qr, nt, iwl);

/* *************************** END CALCULATION ****************************** */
/* ***************************** BEGIN WRITE ******************************** */

            /* write polarization attributes to files */
            if (rl) fputdata(rlfp, headerfp, data_rl, nt);
            if (theta) fputdata(thetafp, headerfp, data_theta, nt);
            if (phi) fputdata(phifp, headerfp, data_phi, nt);
            if (tau) fputdata(taufp, headerfp, data_tau, nt);
            if (ellip) {
                fputdata(e21fp, headerfp, data_e21, nt);
                fputdata(e31fp, headerfp, data_e31, nt);
                fputdata(e32fp, headerfp, data_e32, nt);
            }
            if (pln) fputdata(plnfp, headerfp, data_pln, nt);
            if (f1) fputdata(f1fp, headerfp, data_f1, nt);
            if (l1) fputdata(l1fp, headerfp, data_l1, nt);
            if (amp) {
                fputdata(erfp, headerfp, data_er, nt);
                fputdata(irfp, headerfp, data_ir, nt);
                fputdata(qrfp, headerfp, data_qr, nt);
            }
            if (dir) fputdata3c(dirfp, headerfp, data3c_dir, nt);
            if (rl && theta)
            {
                fputdata(pfiltfp, headerfp, data_pfilt, nt);
                fputdata(sfiltfp, headerfp, data_sfilt, nt);
                fputdata(nfiltfp, headerfp, data_nfilt, nt);
                fputdata(efiltfp, headerfp, data_efilt, nt);
               // fputdata(pkur, headerfp, data_pkur, nt);
               // fputdata(skur, headerfp, data_skur, nt);
            }

/* ****************************** END WRITE ********************************* */

        } /* end of processing three-component dataset */

    } while (gettr(&tr)); /* end loop over traces */

    if (verbose) {
        fprintf(stderr,"\n");
        if (icomp) warn("last %d trace(s) skipped", icomp);
    }

    /* close files */
    efclose(headerfp);
    if (rl) efclose(rlfp);
    if (theta) efclose(thetafp);
    if (phi) efclose(phifp);
    if (tau) efclose(taufp);
    if (ellip) {
        efclose(e21fp);
        efclose(e31fp);
        efclose(e32fp);
    }
    if (pln) efclose(plnfp);
    if (f1) efclose(f1fp);
    if (l1) efclose(l1fp);
    if (amp) {
        efclose(erfp);
        efclose(irfp);
        efclose(qrfp);
    }
    if (dir) efclose(dirfp);
    if (rl && theta) {
        efclose(pfiltfp);
        efclose(sfiltfp);
//        efclose(pkur);
//        efclose(skur);
        }
    return(CWP_Exit());
}

/**********************************************************************/
/* Functions used internally                                          */
/**********************************************************************/

/* calculate time window weights for a smooth time window,        */
/* after Sheriff, 1991, p. 335; tapered windows are suggested by  */
/* Jurkevics, 1988                                                */

void calc_window(float *w, int iwl, int iwin)
{
    int i,j;      /* loop indices within window */
    float m;      /* half of time window (iwl/2) */

    m = (float) iwl/2;

    switch (iwin) {
        case WBOXCAR :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = 1.0;
            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
        case WBARTLETT :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = 1.0 - (fabs((float) i) / m);

            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
        case WHANNING :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = 0.5 + 0.5*cos(PI*fabs((float) i) / m);
            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
        case WWELSH :
            for (i=-iwl/2,j=0;i<iwl/2;i++,j++) {
                w[j] = fabs( fabs((float) i) / m - 1.0);
                w[j] *= w[j];
            }
            if (j<iwl) w[iwl-1] = w[0];
            break;
    }
}

/* covariance of two time sequences *data1 and *data2,   */
/* calculated in a window of length iwl,                 */
/* weighted with factors *w and starting at index istart */

float covar(float *data1, float *data2, int istart, int iwl, float *w)
{
    register int i, j;
    float cov=0.0;
    float mean1=0.0;
    float mean2=0.0;

    for (i=istart,j=0;i<(istart+iwl);i++,j++) {
        mean1 += data1[i] * w[j];
        mean2 += data2[i] * w[j];
    }
    mean1 = mean1 / (float) iwl;
    mean2 = mean2 / (float) iwl;

    for (i=istart,j=0;i<(istart+iwl);i++,j++) {
	/* modified by Mega G. Baitoff 6 May 2015 */
        /* cov += (data1[i]-mean1) * (data2[i]-mean2) * w[j]; */
        cov += (data1[i] * w[j] - mean1) * (data2[i] * w[j] - mean2);
    }
    cov = cov / (float) iwl;
    return cov;
}

float kurtosiswindow(float *data, float *data_kwl, int it, int kwl, int nt)
    {
        int i, j;
        for (i=it, j=0; i<(it+kwl); i++, j++)
        {
            data_kwl[j] = data[i];
        }
        return gsl_stats_float_kurtosis(data_kwl,1,kwl);
    }

/**********************************************************************/
/* Functions defining polarization attributes and angles              */
/**********************************************************************/

/* rectilinearity factor rl (3 different definitions) */

float calc_rl(float *d, float rlq, int opt)
{
    float rl;

    if (d[1]!=0.0) {
        switch (opt) {
            default:
                /* case 1: */
                /* rl definition after Kanasewich, 1981 */
                rl = 1 - pow(fabs(d[2]/d[1]), rlq);
                break;
            case 2:
                /* rl definition after Jurkevics, 1988 */
                /* case used by Z. Ross et. al., 2014 */
                rl = 1 - pow( 0.5*(fabs(d[2]/d[1]) + fabs(d[3]/d[1]) ), rlq);
                break;
            case 3:
                /* rl definition after Meyer, 1988 */
                rl = 1 - pow(fabs(d[2]/d[1]) + fabs(d[3]/d[1]), rlq);
                break;
        }
        return rl;
    }
    else {
        return 0.0;
    }
}

/* vertical (incidence) angle theta (different definitions) */
/* assumes a right handed coordinate system starting        */
/* with the vertical component Z (Z=1, N,R=2, E,T=3)        */

float calc_theta(float **v, int opt)
{
    float theta, horiz;

    switch (opt) {
        default:
            /* case 1: */
            /* definition after Kanasewich, 1981 */
            /* interval -pi/2 <= theta <= pi/2 */
            if (v[1][1]) {
                horiz = sqrt( v[2][1]*v[2][1] + v[3][1]*v[3][1] );
                theta = atan( horiz / v[1][1] );
            }
            else {
                theta = 0.0;
            }
            break;
        case 2:
        case 3:
            /* definition after Jurkevics, 1988 */
            /* interval 0.0 <= theta <= pi/2 */
            /* case used by Z. Ross et. al., 2014 */
            theta = acos( fabs(v[1][1]) );
            break;
    }
    return theta;
}

/* horizontal azimuth phi (different definitions)    */
/* assumes a right handed coordinate system starting */
/* with the vertical component Z (Z=1, N,R=2, E,T=3) */

/* Note: The SIGN function is introduced to resolve  */
/* the 180 deg ambiguity by taking the positive      */
/* vertical component of v[1][1] (Jurkevics, 1988).  */
#define VSIGN ( (v[1][1]<0.0) ? -1.0 : 1.0 )

float calc_phi(float **v, int opt)
{
    float phi;

    switch (opt) {
        default:
            /* case 1: */
            /* definition after Kanasewich, 1981 */
            /* interval -pi/2 <= phi <= pi/2 */
            if (v[1][2]) {
                phi = atan( v[3][1] / v[2][1] );
            }
            else {
                phi = (v[3][1]>0.0) ? 0.5*PI : -0.5*PI;
            }
            break;
        case 2:
        case 3:
            /* definitions after Jurkevics, 1988 */
            /* interval -pi <= phi <= pi */
            if (v[2][1]) {
                phi = atan2( v[3][1]*VSIGN, v[2][1]*VSIGN);
            }
            else {
                phi = (v[3][1]>0.0) ? 0.5*PI*VSIGN : -0.5*PI*VSIGN;
            }

            /* interval 0.0 <= phi <= 2*pi */
            if (phi<0.0 && opt==3) phi += 2.0*PI;
            break;
    }
    return phi;
}
#undef VSIGN

/* global polarization parameter tau (Samson, 1973) */

float calc_tau(float *d)
{
    float x1, x2, x3, x4, tau;

    if (d[1]) {
        x1 = pow(( 1 - fabs(d[2]/d[1])), 2.);
        x2 = pow(( 1 - fabs(d[3]/d[1])), 2.);
        x3 = pow(( fabs(d[2]/d[1]) - fabs(d[3]/d[1]) ), 2.);
        x4 = pow(( 1 + fabs(d[2]/d[1]) + fabs(d[3]/d[1])), 2.);
        tau = sqrt( (x1+x2+x3) / (2*x4) );
        return tau;
    }
    else
        return 0.0;
}

/* ellipticities e_ik */

float calc_ellip(float *d, int d1, int d2)
{
    float ellip;

    if (d[d2]) {
        ellip = sqrt( fabs(d[d1]/d[d2]) );
        return ellip;
    }
    else
        return 0.0;
}

/* planarity after Jurkevics, 1988 */

float calc_plan(float *d)
{
    float pln;

    if (d[1]+d[2]) {
        pln = 1.0 - 2.0*d[3] / (d[1] + d[2]);
        return pln;
    }
    else
        return 0.0;
}

/* flatness coefficient f1 after Benhama et. al, 1988 */

float calc_f1(float *d)
{
    float f1,x1,x2;

    x1 = 3.0 * calc_ellip(d,3,1);
    x2 = 1.0 + calc_ellip(d,2,1) + calc_ellip(d,3,1);
    f1 = 1.0 - x1 / x2;
    return f1;
}

/* linearity coefficient l1 */

float calc_l1(float *d)
{
    float l1,x1,x2;

    x1 = 3. * ( calc_ellip(d,2,1) + calc_ellip(d,3,1) );
    x2 = 2. * ( 1 + calc_ellip(d,2,1) + calc_ellip(d,3,1));
    l1 = 1. - x1 / x2;
    return l1;
}

/* direction cosines or directivity functions (3 components) */

void calc_dir(float **data3c_dir, float **v, int it)
{
    int i;
    for (i=1;i<=3;i++) {
        data3c_dir[i][it]=v[i][1];
    }
}

/* amplitude parameters */

/* eigenresultant */

float calc_er(float *d)
{
    float er;
    er = sqrt(fabs(d[1]));
    return er;
}

/* instantaneous and quadratic resultant (Meyer, 1988) */

void ampparams(float **indata, float *data_ir, float *data_qr, int nt, int iwl)
{
    int i,it;
    float sqrsum;

    for (it=0;it<nt;it++) {
        /* instantaneous resultant */
        data_ir[it] = sqrt(indata[1][it]*indata[1][it]+ \
            indata[2][it]*indata[2][it]+indata[3][it]*indata[3][it]);

        /* quadratic resultant */
        if ( (it >= iwl/2) && (it < nt-iwl/2) ) {
            sqrsum=0.0;
            for (i=it-iwl/2;i<(it+iwl);i++) {
                sqrsum += indata[1][i]*indata[1][i]+ \
                    indata[2][i]*indata[2][i]+indata[3][i]*indata[3][i];
            }
            data_qr[it] = sqrsum / (float) iwl;
        }
    }
}

/* P, S detection parameters */

/* P-Wave energy filter */
float calc_pfilt(float rl, float theta)
{
    float pfilt;
    pfilt = rl*cos(theta);
    return pfilt;
}

/* S-Wave energy filter */
float calc_sfilt(float rl, float theta)
{
    float sfilt;
    sfilt = rl * (1 - cos(theta));
    return sfilt;
}

/**********************************************************************/
/* Functions for data output                                          */
/**********************************************************************/

/* write one-component data into file */

void fputdata(FILE *fileptr, FILE *headerptr, float *outdata, int nt)
{
    efread(&tr, 1, HDRBYTES, headerptr);
    erewind(headerptr);

    memcpy((void *)tr.data, (const void *) outdata, nt*FSIZE);

    fputtr(fileptr, &tr);
}

/* write three-component data into file */

void fputdata3c(FILE *fileptr, FILE *headerptr, float **outdata3c, int nt)
{
    int i;

    for(i=1;i<=3;i++) {
        efread(&tr, 1, HDRBYTES, headerptr);

        memcpy((void *)tr.data, (const void *) outdata3c[i], nt*FSIZE);

        fputtr(fileptr, &tr);
    }
    erewind(headerptr);
}

/* END OF FILE */


#include <cblas.h>
#include <stdlib.h>
#include <string.h>
#include "loss.h"

void expit(const double *x, double *out, int x_len) {
    for (int i = 0; i < x_len; i++) {
        if (x[i] > 0) {
            out[i] = 1. / (1. + exp(-x[i]));
        } else {
            out[i] = 1. - 1. / (1. + exp(x[i]));
        }
    }
}

void log_logistic(const double *x, double *out, int x_len) {
    for (int i = 0; i < x_len; i++) {
        if (x[i] > 0.0) {
            out[i] = -log(1.0 + exp(-x[i]));
        } else {
            out[i] = x[i] - log(1.0 + exp(x[i]));
        }
    }
}

void logistic(const double *x, double *out, int x_len) {
    for (int i = 0; i < x_len; i++) {
        if (x[i] > 0) {
            out[i] = 1. / (1. + exp(-x[i]));
        } else {
            out[i] = 1. - 1. / (1. + exp(x[i]));
        }
    }
}

double log_sum_exp(const double *x, int x_len) {
    double max_x = x[0], out = 0.0;
    for (int i = 1; i < x_len; i++) {
        if (x[i] > max_x) {
            max_x = x[i];
        }
    }
    for (int i = 0; i < x_len; i++) {
        out += exp(x[i] - max_x);
    }
    return max_x + log(out);
}

void logistic_loss_grad(const double *w,
                        const double *x_tr,
                        const double *y_tr,
                        double *loss_grad,
                        double eta,
                        int n_samples,
                        int n_features) {
    for (int i = 0; i < n_features + 1; i++) {
        if (isnan(w[i])) {
            printf("%f\n", w[i]);
            printf("warning: loss grad error!\n");
            break;
        }
    }
    int i, n = n_samples, p = n_features;
    double intercept = w[p], sum_z0 = 0.0;
    loss_grad[0] = 0.0;
    double *yz = malloc(sizeof(double) * n);
    double *z0 = malloc(sizeof(double) * n);
    double *logistic = malloc(sizeof(double) * n);
    for (i = 0; i < n; i++) { /** calculate yz */
        yz[i] = intercept;
    }
    //x_tr^T*w+
    cblas_dgemv(CblasRowMajor, CblasNoTrans, n, p, 1., x_tr, p, w, 1, 1., yz, 1);
    for (i = 0; i < n; i++) {
        yz[i] *= y_tr[i];
    }
    expit(yz, z0, n); // calculate z0 and final intercept
    /** calculate logistic logistic[i] = 1/(1+exp(-y[i]*(xi^T*w+c)))*/
    log_logistic(yz, logistic, n);
    /** calculate loss of data fitting part*/
    for (i = 0; i < n; i++) {
        z0[i] = (z0[i] - 1.) * y_tr[i];
        sum_z0 += z0[i];
        loss_grad[0] -= logistic[i];
    }
    /**calculate loss of regularization part (it does not have intercept)*/
    loss_grad[0] += 0.5 * eta * cblas_ddot(p, w, 1, w, 1);
    /** calculate gradient of coefficients*/
    memcpy(loss_grad + 1, w, sizeof(double) * p);
    /** x^T*z0 + eta*w, where z0[i]=(logistic[i] - 1.)*yi*/
    cblas_dgemv(CblasRowMajor, CblasTrans,
                n, p, 1., x_tr, p, z0, 1, eta, loss_grad + 1, 1);
    /** calculate gradient of intercept part*/
    loss_grad[p + 1] = sum_z0; // intercept part
    free(logistic);
    free(z0);
    free(yz);
}

void logistic_loss_grad_sparse(const double *w,
                               const double *x_tr,
                               const double *y_tr,
                               double *loss_grad,
                               double eta,
                               int n_samples,
                               int n_features) {
    //sub_p is number of nonzeros in w.
    int i, n = n_samples, p = n_features, sub_p = 0;
    double intercept = w[p], sum_z0 = 0.0;
    loss_grad[0] = 0.0;
    double *yz = malloc(sizeof(double) * n);
    double *z0 = malloc(sizeof(double) * n);
    double *logistic = malloc(sizeof(double) * n);
    for (i = 0; i < n; i++) { /** calculate yz */
        yz[i] = intercept;
    }
    double nonzero_w[p];
    int nonzero_w_ind[p];
    for (i = 0; i < p; i++) {
        if (w[i] != 0.0) {
            nonzero_w_ind[sub_p] = i;
            nonzero_w[sub_p] = w[i];
            sub_p++;
        }
    }
    //x_tr^T*w+ to use the sparsity of w
    for (i = 0; i < n; i++) {
        double tmp_val = 0.0;
        for (int j = 0; j < sub_p; j++) {
            tmp_val += x_tr[i * p + nonzero_w_ind[j]] * nonzero_w[j];
        }
        yz[i] = yz[i] + tmp_val;
    }
    for (i = 0; i < n; i++) {
        yz[i] *= y_tr[i];
    }
    expit(yz, z0, n); // calculate z0 and final intercept
    /** calculate logistic logistic[i] = 1/(1+exp(-y[i]*(xi^T*w+c)))*/
    log_logistic(yz, logistic, n);
    /** calculate loss of data fitting part*/
    for (i = 0; i < n; i++) {
        z0[i] = (z0[i] - 1.) * y_tr[i];
        sum_z0 += z0[i];
        loss_grad[0] -= logistic[i];
    }
    /**calculate loss of regularization part (it does not have intercept)*/
    loss_grad[0] += 0.5 * eta * cblas_ddot(p, w, 1, w, 1);
    /** calculate gradient of coefficients*/
    memcpy(loss_grad + 1, w, sizeof(double) * p);
    /** x^T*z0 + eta*w, where z0[i]=(logistic[i] - 1.)*yi*/
    cblas_dgemv(CblasRowMajor, CblasTrans,
                n, p, 1., x_tr, p, z0, 1, eta, loss_grad + 1, 1);
    /** calculate gradient of intercept part*/
    loss_grad[p + 1] = sum_z0; // intercept part
    free(logistic);
    free(z0);
    free(yz);
}

void logistic_predict(const double *x_te,
                      const double *wt,
                      double *pred_prob,
                      double *pred_label,
                      double threshold,
                      int n,
                      int p) {
    openblas_set_num_threads(1);
    int i;
    cblas_dgemv(CblasRowMajor, CblasNoTrans,
                n, p, 1., x_te, p, wt, 1, 0., pred_prob, 1);
    for (i = 0; i < n; i++) {
        pred_prob[i] += wt[p]; // intercept
    }
    expit(pred_prob, pred_prob, n);
    for (i = 0; i < n; i++) {
        if (pred_prob[i] >= threshold) {
            pred_label[i] = 1.;
        } else {
            pred_label[i] = -1.;
        }
    }
}

void least_square_loss_grad(const double *w,
                            const double *x_tr,
                            const double *y_tr,
                            double *loss_grad,
                            double eta,
                            int n_samples,
                            int n_features) {
    for (int i = 0; i < n_features + 1; i++) {
        if (isnan(w[i])) {
            printf("%f\n", w[i]);
            printf("warning: loss grad error!\n");
            break;
        }
    }
    int i, n = n_samples, p = n_features;
    double *yz = malloc(sizeof(double) * n);
    double *w0 = malloc(sizeof(double) * n);
    cblas_dcopy(n, y_tr, 1, yz, 1); // y_tr --> yz
    cblas_dscal(p + 2, 0.0, loss_grad, 1);
    cblas_daxpy(p, eta, w, 1, loss_grad + 1, 1); // l2_lambda*w --> loss_grad
    //Order,TransA, M, N, alpha, A, lda, x, incX, beta, Y, incY
    // Y<-alpha*AX + beta*Y, where A=MxN, Y=N
    cblas_dgemv(CblasRowMajor, CblasNoTrans,
                n, p, 1., x_tr, p, w, 1, -1., yz, 1); //Xw - y
    for (i = 0; i < n; i++) { w0[i] = w[p]; }
    cblas_daxpy(n, 1., w0, 1, yz, 1);
    loss_grad[0] = cblas_ddot(n, yz, 1, yz, 1) / (2. * n);
    loss_grad[0] += (.5 * eta) * cblas_ddot(p, w, 1, w, 1); //loss
    cblas_dgemv(CblasRowMajor, CblasTrans,
                n, p, 1. / n, x_tr, p, yz, 1, 1., loss_grad + 1, 1);
    for (i = 0; i < n; i++) {
        w0[i] = w[p];
        loss_grad[p + 1] = yz[i];
    }
    loss_grad[p + 1] /= n;
    free(w0), free(yz);
}

void least_square_predict(const double *x_te,
                          const double *wt,
                          double *pred_prob,
                          double *pred_label,
                          double threshold,
                          int n_samples,
                          int p_features) {
    openblas_set_num_threads(1);
    int i, n = n_samples, p = p_features;
    cblas_dgemv(CblasRowMajor, CblasNoTrans,
                n, p, 1., x_te, p, wt, 1, 0., pred_prob, 1);
    for (i = 0; i < n; i++) {
        pred_prob[i] += wt[p]; // intercept
    }
    // this is not useful.
    for (i = 0; i < n; i++) {
        if (pred_prob[i] >= threshold) {
            pred_label[i] = 1.;
        } else {
            pred_label[i] = -1.;
        }
    }
}


#include <math.h>
#include <stdio.h>
#include <gsl/gsl_rng.h>
#include "matrix_operations.c"
#include "ellipse/parametric_function_roots.c"
#include "polynomial.c"

void characteristic_ellipse_matrix(double *X, double *R, double phi, double exponent)
{

    // rotation matrix
    double Q[2][2] = {{ 0 }};
    double Qt[2][2] = {{ 0 }};
    Q[0][0] = cos(phi);
    Q[0][1] = sin(phi);
    Q[1][0] = -sin(phi);
    Q[1][1] = cos(phi);
    matrix_transpose(&Qt[0][0], &Q[0][0], 2, 2);

    // radii matrix
    double O[2][2] = {{ 0 }};
    double diag_vals[2];
    diag_vals[0] = pow(R[0], exponent * -2.0);
    diag_vals[1] = pow(R[1], exponent * -2.0);
    set_diagonal(&O[0][0], diag_vals, 2, 2);

    // characteristic ellipse matrix
    double X_temp[2][2] = {{ 0 }};
    matrix_multiply(&X_temp[0][0], &O[0][0], &Q[0][0], 2, 2, 2);
    matrix_multiply(X, &Qt[0][0], &X_temp[0][0], 2, 2, 2);

}

double ellipse_overlap(double *rA, double *radiiA, double phiA, double *rB, double *radiiB, double phiB)
{

    // find XA^(-1) and XB^(1/2)
    double XA[2][2] = {{ 0 }};
    double XB[2][2] = {{ 0 }};
    characteristic_ellipse_matrix(&XA[0][0], &radiiA[0], phiA, -1.0);
    characteristic_ellipse_matrix(&XB[0][0], &radiiB[0], phiB, 0.5);

    // find A_AB
    double A_AB[2][2] = {{ 0 }};
    double A_temp[2][2] = {{ 0 }};
    matrix_multiply(&A_temp[0][0], &XA[0][0], &XB[0][0], 2, 2, 2);
    matrix_multiply(&A_AB[0][0], &XB[0][0], &A_temp[0][0], 2, 2, 2);

    // find r_AB
    double rAB[2];
    rAB[0] = rB[0] - rA[0];
    rAB[1] = rB[1] - rA[1];

    // find a_AB
    double a_AB[2];
    matrix_multiply(&a_AB[0], &XB[0][0], &rAB[0], 2, 2, 1);

    // extract elements of the matrix A_AB and a_AB
    double a11, a12, a21, a22;
    double b1, b2;
    a11 = A_AB[0][0];
    a12 = A_AB[0][1];
    a21 = A_AB[1][0];
    a22 = A_AB[1][1];
    b1 = a_AB[0];
    b2 = a_AB[1];

    // find coefficients for the parametric polynomial derivative used to find max
    double h[5];
    double z[8] = { 0, 0, 0, 0, 0, 0, 0, 0 };
    size_t n = 5;
    h[0] = h0(a11, a12, a21, a22, b1, b2);
    h[1] = h1(a11, a12, a21, a22, b1, b2);
    h[2] = h2(a11, a12, a21, a22, b1, b2);
    h[3] = h3(a11, a12, a21, a22, b1, b2);
    h[4] = h4(a11, a12, a21, a22, b1, b2);

    // find roots
    size_t m;
    m = find_roots(&z[0], &h[0], n);

    double F;
    int i;

    if (m > 1)
    {
        if (f(0, a11, a12, a21, a22, b1, b2) > f(1, a11, a12, a21, a22, b1, b2))
        {

           F = f(0, a11, a12, a21, a22, b1, b2);

        }
        else
        {

           F = f(1, a11, a12, a21, a22, b1, b2);

        }

        for(i=0; i<=m-2; i++)
        {

           if (( z[2 * i + 1] == 0 ) & (z[2 * i] > 0) & (z[2 * i] < 1))
           {

              F = f(z[2 * i], a11, a12, a21, a22, b1, b2);

           }

        }
    }
    else
    {
        F = 0.;
    }

    return F;

}

double container_square_overlap_potential(double *rA, double *radiiA, double phiA)
{

    double rB[2];
    double radiiB[2];
    double phiB;

    double top = 0;
    double bottom = 0;
    double left = 0;
    double right = 0;

    // top
    rB[0] = 0.5;
    rB[1] = 2.0;
    radiiB[0] = INFINITY;
    radiiB[1] = 1.0;
    phiB = 0.0;
    top = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    // bottom
    rB[0] = 0.5;
    rB[1] = -1.0;
    radiiB[0] = INFINITY;
    radiiB[1] = 1.0;
    phiB = 0.0;
    bottom = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    // left
    rB[0] = -1.0;
    rB[1] = 0.5;
    radiiB[0] = 1.0;
    radiiB[1] = INFINITY;
    phiB = 0.0;
    left = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    // right
    rB[0] = 2.0;
    rB[1] = 0.5;
    radiiB[0] = 1.0;
    radiiB[1] = INFINITY;
    phiB = 0.0;
    right = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

    return fminf(top, fminf(bottom, fminf(left, right)));

}

size_t rsa_align_square(double *x, double *y,
    size_t npoints, double *radius, double phi, int step_limit,
    unsigned long randSeed)
{

    // Setup GSL random number generator
    const gsl_rng_type * T;
    gsl_rng * r;
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    // Set the seed
    // srand ( time(NULL) );
    // unsigned long randSeed = rand();
    gsl_rng_set(r, randSeed);

    // arrays for overlap_potential functions
    double rA[2];
    double rB[2];
    double radiiA[2];
    double radiiB[2];
    double F;
    double xn = 0.;
    double yn = 0.;

    radiiA[0] = radius[0];
    radiiA[1] = radius[1];

    // Set the initial position
    double C;
    C = fmaxf(radius[0], radius[1]);

    F = 0;
    while (F < 1.)
    {

        xn = gsl_rng_uniform (r);
        yn = gsl_rng_uniform (r);

        rA[0] = xn;
        rA[1] = yn;

        F = container_square_overlap_potential(&rA[0], &radiiA[0], phi);

    }

    x[0] = xn;
    y[0] = yn;

    size_t valid_pts;
    int k, flag, step;

    step = 0;
    valid_pts = 1;
    while ((valid_pts < npoints) & (step < step_limit))
    {

        // Generate new ellipse inside the container
        F = 0;
        while (F < 1.)
        {

            xn = gsl_rng_uniform (r);
            yn = gsl_rng_uniform (r);

            rA[0] = xn;
            rA[1] = yn;

            F = container_square_overlap_potential(&rA[0], &radiiA[0], phi);

        }

        // Determine if new ellipse overlaps with existing ellipses
        flag = 1;
        for (k = 0; k < valid_pts; k++)
        {

            rA[0] = x[k];
            rA[1] = y[k];
            rB[0] = xn;
            rB[1] = yn;
            radiiA[0] = radius[0];
            radiiA[1] = radius[1];
            radiiB[0] = radius[0];
            radiiB[1] = radius[1];

            F = ellipse_overlap(&rA[0], &radiiA[0], phi, &rB[0], &radiiB[0], phi);

            if (F < 1.0)
            {

                flag = 0;
                break;

            }
        }
        if (flag == 1)
        {

           x[valid_pts] = xn;
           y[valid_pts] = yn;
           valid_pts += 1;

        }

        step += 1;

    }

    gsl_rng_free (r);

    return valid_pts;

}

size_t rsa_square(double *x, double *y,
    size_t npoints, double *radius, double *phi, int step_limit,
    unsigned long randSeed)
{

    // Setup GSL random number generator
    const gsl_rng_type * T;
    gsl_rng * r;
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    // Set the seed
    // srand ( time(NULL) );
    // unsigned long randSeed = rand();
    gsl_rng_set(r, randSeed);

    // arrays for overlap_potential functions
    double rA[2];
    double rB[2];
    double radiiA[2];
    double radiiB[2];
    double phiA;
    double phiB;
    double F;
    double xn = 0;
    double yn = 0;
    double phin = 0;

    radiiA[0] = radius[0];
    radiiA[1] = radius[1];

    // Set the initial position
    double C;
    C = fmaxf(radius[0], radius[1]);

    F = 0;
    while (F < 1.)
    {

        xn = gsl_rng_uniform (r);
        yn = gsl_rng_uniform (r);

        rA[0] = xn;
        rA[1] = yn;
        phiA = 2 * M_PI * gsl_rng_uniform (r);

        F = container_square_overlap_potential(&rA[0], &radiiA[0], phiA);

    }

    x[0] = xn;
    y[0] = yn;
    phi[0] = phiA;

    size_t valid_pts;
    int k, flag, step;

    step = 0;
    valid_pts = 1;
    while ((valid_pts < npoints) & (step < step_limit))
    {

        // Generate new ellipse inside the container
        F = 0;
        while (F < 1.)
        {

            xn = gsl_rng_uniform (r);
            yn = gsl_rng_uniform (r);
            phin = 2 * M_PI * gsl_rng_uniform (r);

            rA[0] = xn;
            rA[1] = yn;
            phiA = phin;

            F = container_square_overlap_potential(&rA[0], &radiiA[0], phiA);

        }

        // Determine if new ellipse overlaps with existing ellipses
        flag = 1;
        for (k = 0; k < valid_pts; k++)
        {

            rA[0] = x[k];
            rA[1] = y[k];
            phiA = phi[k];

            rB[0] = xn;
            rB[1] = yn;
            phiB = phin;

            radiiA[0] = radius[0];
            radiiA[1] = radius[1];
            radiiB[0] = radius[0];
            radiiB[1] = radius[1];

            F = ellipse_overlap(&rA[0], &radiiA[0], phiA, &rB[0], &radiiB[0], phiB);

            if (F < 1.0)
            {

                flag = 0;
                break;

            }
        }
        if (flag == 1)
        {

           x[valid_pts] = xn;
           y[valid_pts] = yn;
           phi[valid_pts] = phin;
           valid_pts += 1;

        }

        step += 1;

    }

    gsl_rng_free (r);

    return valid_pts;

}


/* interpolation/akima.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/* Author:  G. Jungman
 */
#include <config.h>
#include <stdlib.h>
#include <math.h>
#include <gsl/gsl_errno.h>
#include "integ_eval.h"
#include <gsl/gsl_interp.h>

typedef struct
{
  double * b;
  double * c;
  double * d;
  double * _m;
} akima_state_t;

/* common creation */
static void *
akima_alloc (size_t size)
{
  akima_state_t *state = (akima_state_t *) malloc (sizeof (akima_state_t));

  if (state == NULL)
    {
      GSL_ERROR_NULL("failed to allocate space for state", GSL_ENOMEM);
    }

  state->b = (double *) malloc (size * sizeof (double));

  if (state->b == NULL)
    {
      free (state);
      GSL_ERROR_NULL("failed to allocate space for b", GSL_ENOMEM);
    }

  state->c = (double *) malloc (size * sizeof (double));

  if (state->c == NULL)
    {
      free (state->b);
      free (state);
      GSL_ERROR_NULL("failed to allocate space for c", GSL_ENOMEM);
    }

  state->d = (double *) malloc (size * sizeof (double));

  if (state->d == NULL)
    {
      free (state->c);
      free (state->b);
      free (state);
      GSL_ERROR_NULL("failed to allocate space for d", GSL_ENOMEM);
    }

  state->_m = (double *) malloc ((size + 4) * sizeof (double));

  if (state->_m == NULL)
    {
      free (state->d);
      free (state->c);
      free (state->b);
      free (state);
      GSL_ERROR_NULL("failed to allocate space for _m", GSL_ENOMEM);
    }

  return state;
}

/* common calculation */
static void
akima_calc (const double x_array[], double b[],  double c[],  double d[], size_t size, double m[])
{
  size_t i;

  for (i = 0; i < (size - 1); i++)
    {
      const double NE = fabs (m[i + 1] - m[i]) + fabs (m[i - 1] - m[i - 2]);
      if (NE == 0.0)
        {
          b[i] = m[i];
          c[i] = 0.0;
          d[i] = 0.0;
        }
      else
        {
          const double h_i = x_array[i + 1] - x_array[i];
          const double NE_next = fabs (m[i + 2] - m[i + 1]) + fabs (m[i] - m[i - 1]);
          const double alpha_i = fabs (m[i - 1] - m[i - 2]) / NE;
          double alpha_ip1;
          double tL_ip1;
          if (NE_next == 0.0)
            {
              tL_ip1 = m[i];
            }
          else
            {
              alpha_ip1 = fabs (m[i] - m[i - 1]) / NE_next;
              tL_ip1 = (1.0 - alpha_ip1) * m[i] + alpha_ip1 * m[i + 1];
            }
          b[i] = (1.0 - alpha_i) * m[i - 1] + alpha_i * m[i];
          c[i] = (3.0 * m[i] - 2.0 * b[i] - tL_ip1) / h_i;
          d[i] = (b[i] + tL_ip1 - 2.0 * m[i]) / (h_i * h_i);
        }
    }
}

static int
akima_init (void * vstate, const double x_array[], const double y_array[],
            size_t size)
{
  akima_state_t *state = (akima_state_t *) vstate;

  double * m = state->_m + 2; /* offset so we can address the -1,-2
                                 components */

  size_t i;
  for (i = 0; i <= size - 2; i++)
    {
      m[i] = (y_array[i + 1] - y_array[i]) / (x_array[i + 1] - x_array[i]);
    }

  /* non-periodic boundary conditions */
  m[-2] = 3.0 * m[0] - 2.0 * m[1];
  m[-1] = 2.0 * m[0] - m[1];
  m[size - 1] = 2.0 * m[size - 2] - m[size - 3];
  m[size] = 3.0 * m[size - 2] - 2.0 * m[size - 3];

  akima_calc (x_array, state->b, state->c, state->d, size, m);

  return GSL_SUCCESS;
}

static int
akima_init_periodic (void * vstate,
                     const double x_array[],
                     const double y_array[],
                     size_t size)
{
  akima_state_t *state = (akima_state_t *) vstate;

  double * m = state->_m + 2; /* offset so we can address the -1,-2
                                 components */

  size_t i;
  for (i = 0; i <= size - 2; i++)
    {
      m[i] = (y_array[i + 1] - y_array[i]) / (x_array[i + 1] - x_array[i]);
    }

  /* periodic boundary conditions */
  m[-2] = m[size - 1 - 2];
  m[-1] = m[size - 1 - 1];
  m[size - 1] = m[0];
  m[size] = m[1];

  akima_calc (x_array, state->b, state->c, state->d, size, m);

  return GSL_SUCCESS;
}

static void
akima_free (void * vstate)
{
  akima_state_t *state = (akima_state_t *) vstate;

  free (state->b);
  free (state->c);
  free (state->d);
  free (state->_m);
  free (state);
}

static
int
akima_eval (const void * vstate,
            const double x_array[], const double y_array[], size_t size,
            double x,
            gsl_interp_accel * a,
            double *y)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t index;

  if (a != 0)
    {
      index = gsl_interp_accel_find (a, x_array, size, x);
    }
  else
    {
      index = gsl_interp_bsearch (x_array, x, 0, size - 1);
    }

  /* evaluate */
  {
    const double x_lo = x_array[index];
    const double delx = x - x_lo;
    const double b = state->b[index];
    const double c = state->c[index];
    const double d = state->d[index];
    *y = y_array[index] + delx * (b + delx * (c + d * delx));
    return GSL_SUCCESS;
  }
}

static int
akima_eval_deriv (const void * vstate,
                  const double x_array[], const double y_array[], size_t size,
                  double x,
                  gsl_interp_accel * a,
                  double *dydx)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t index;

  DISCARD_POINTER(y_array); /* prevent warning about unused parameter */

  if (a != 0)
    {
      index = gsl_interp_accel_find (a, x_array, size, x);
    }
  else
    {
      index = gsl_interp_bsearch (x_array, x, 0, size - 1);
    }

  /* evaluate */
  {
    double x_lo = x_array[index];
    double delx = x - x_lo;
    double b = state->b[index];
    double c = state->c[index];
    double d = state->d[index];
    *dydx = b + delx * (2.0 * c + 3.0 * d * delx);
    return GSL_SUCCESS;
  }
}

static
int
akima_eval_deriv2 (const void * vstate,
                   const double x_array[], const double y_array[], size_t size,
                   double x,
                   gsl_interp_accel * a,
                   double *y_pp)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t index;

  DISCARD_POINTER(y_array); /* prevent warning about unused parameter */

  if (a != 0)
    {
      index = gsl_interp_accel_find (a, x_array, size, x);
    }
  else
    {
      index = gsl_interp_bsearch (x_array, x, 0, size - 1);
    }

  /* evaluate */
  {
    const double x_lo = x_array[index];
    const double delx = x - x_lo;
    const double c = state->c[index];
    const double d = state->d[index];
    *y_pp = 2.0 * c + 6.0 * d * delx;
    return GSL_SUCCESS;
  }
}

static
int
akima_eval_integ (const void * vstate,
                  const double x_array[], const double y_array[], size_t size,
                  gsl_interp_accel * acc,
                  double a, double b,
                  double * result)
{
  const akima_state_t *state = (const akima_state_t *) vstate;

  size_t i, index_a, index_b;

  if (acc != 0)
    {
      index_a = gsl_interp_accel_find (acc, x_array, size, a);
      index_b = gsl_interp_accel_find (acc, x_array, size, b);
    }
  else
    {
      index_a = gsl_interp_bsearch (x_array, a, 0, size - 1);
      index_b = gsl_interp_bsearch (x_array, b, 0, size - 1);
    }

  *result = 0.0;

  /* interior intervals */

  for(i=index_a; i<=index_b; i++) {
    const double x_hi = x_array[i + 1];
    const double x_lo = x_array[i];
    const double y_lo = y_array[i];
    const double dx = x_hi - x_lo;
    if(dx != 0.0) {

      if (i == index_a || i == index_b)
        {
          double x1 = (i == index_a) ? a : x_lo;
          double x2 = (i == index_b) ? b : x_hi;
          *result += integ_eval (y_lo, state->b[i], state->c[i], state->d[i],
                                 x_lo, x1, x2);
        }
      else
        {
          *result += dx * (y_lo
                           + dx*(0.5*state->b[i]
                                 + dx*(state->c[i]/3.0
                                       + 0.25*state->d[i]*dx)));
        }
    }
    else {
      *result = 0.0;
      return GSL_EINVAL;
    }
  }

  return GSL_SUCCESS;
}

static const gsl_interp_type akima_type =
{
  "akima",
  5,
  &akima_alloc,
  &akima_init,
  &akima_eval,
  &akima_eval_deriv,
  &akima_eval_deriv2,
  &akima_eval_integ,
  &akima_free
};

const gsl_interp_type * gsl_interp_akima = &akima_type;

static const gsl_interp_type akima_periodic_type =
{
  "akima-periodic",
  5,
  &akima_alloc,
  &akima_init_periodic,
  &akima_eval,
  &akima_eval_deriv,
  &akima_eval_deriv2,
  &akima_eval_integ,
  &akima_free
};

const gsl_interp_type * gsl_interp_akima_periodic = &akima_periodic_type;