/******************************************************************************
CosmoLike Configuration Space Covariances for Projected Galaxy 2-Point Statistics
https://github.com/CosmoLike/CosmoCov
by CosmoLike developers
******************************************************************************/

#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <assert.h>

#include <fftw3.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_sf_gamma.h>
#include <gsl/gsl_sf_legendre.h>
#include <gsl/gsl_sf_bessel.h>
#include <gsl/gsl_sf_trig.h>
#include <gsl/gsl_sf_erf.h>
#define NR_END 1
#define FREE_ARG char*

#define EXIT_MISSING_FILE(ein, purpose, filename) if (!ein) {fprintf(stderr, "Could not find %s file %s\n",purpose,filename);exit(1);}

static double darg __attribute__((unused)),maxarg1 __attribute__((unused)), maxarg2 __attribute__((unused));

#define FMAX(a,b) (maxarg1=(a), maxarg2=(b), (maxarg1) > (maxarg2) ? (maxarg1) : (maxarg2))
#define FMIN(a,b) (maxarg1=(a), maxarg2=(b), (maxarg1) < (maxarg2) ? (maxarg1) : (maxarg2))

//Note: SQR*SQR....gives undefined warning
static double sqrarg;
#define SQR(a) ((sqrarg=(a)) == 0.0 ? 0.0 : sqrarg*sqrarg)
static double cubearg;
#define CUBE(a) ((cubearg=(a)) == 0.0 ? 0.0 : cubearg*cubearg*cubearg)
static double pow4arg;
#define POW4(a) ((pow4arg=(a)) == 0.0 ? 0.0 : pow4arg*pow4arg*pow4arg*pow4arg)

void SVD_inversion(gsl_matrix *cov, gsl_matrix *inverseSVD,int Nmatrix);
double interpol2d(double **f, int nx, double ax, double bx, double dx, double x,   int ny, double ay, double by, double dy, double y, double lower, double upper);
double interpol2d_fitslope(double **f, int nx, double ax, double bx, double dx, double x, int ny, double ay, double by, double dy, double y, double lower);

double interpol(double *f, int n, double a, double b, double dx, double x, double lower, double upper);
double interpol_fitslope(double *f, int n, double a, double b, double dx, double x, double lower);
void free_double_vector(double *v, long nl, long nh);
long *long_vector(long nl, long nh);
int *int_vector(long nl, long nh);
double *create_double_vector(long nl, long nh);
void free_double_matrix(double **m, long nrl, long nrh, long ncl, long nch);
double **create_double_matrix(long nrl, long nrh, long ncl, long nch);
int line_count(char *filename);

void error(char *s);
void cdgamma(fftw_complex x, fftw_complex *res);

double int_gsl_integrate_insane_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter);
double int_gsl_integrate_high_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter);
double int_gsl_integrate_medium_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter);
double int_gsl_integrate_low_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter);
double int_gsl_integrate_cov_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter);

typedef struct {
     double pi;
     double pi_sqr;
     double twopi;
     double ln2;
     double arcmin;
     double lightspeed;
}con;

con constants = {
     3.14159265358979323846, //pi
     9.86960440108935861883, //pisqr
     6.28318530717958647693, //twopi
     0.69314718,
     2.90888208665721580e-4, //arcmin
     299792.458 //speed of light km/s
};

typedef struct {
     double low;
     double medium;
     double high;
     double insane;
}pre;

pre precision= {
  1e-2, //low
  1e-3, //medium
  1e-5, //high
  1e-7 //insane
};

typedef struct {
     double a_min;
     double k_min_mpc;
     double k_max_mpc;
     double k_max_mpc_class;
     double k_min_cH0;
     double k_max_cH0;
     double P_2_s_min;
     double P_2_s_max;
     double xi_via_hankel_theta_min;
     double xi_via_hankel_theta_max;
     double xi_3d_rmin;
     double xi_3d_rmax;
     double M_min;
     double M_max;
}lim;

lim limits = {
//	0.19, //a_min (in order to compute z=4 WFIRST)
    1./(1.+10.), //a_min (z=10, needed for CMB lensing)
	6.667e-6, //k_min_mpc
	1.e3, //k_max_mpc
	50., //k_max_mpc_class
	2.e-2, //k_min_cH0
	3.e+6,  //k_max_cH0
	0.1,//P_2_s_min
	1.0e5,//P_2_s_max
	3.0e-7,//xi_via_hankel_theta_min
	0.12, //xi_via_hankel_theta_max
    1.e-5,//xi_3d_rmin
    1.e+0,//xi_3d_rmax
	1.e+6, //M_min
	1.e+17,//M_max
};

typedef struct {
     int N_a ;
     int N_k_lin;
     int N_k_nlin;
     int N_ell;
     int N_theta;
     int N_thetaH;
     int N_S2;
     int N_DS;
     int N_norm;
     int N_r_3d;
     int N_k_3d;
     int N_a_halo;
}Ntab;

Ntab Ntable = {
100, //N_a
500, //N_k_lin
500, //N_k_nlin
200, //N_ell
200, //N_theta
2048, //N_theta for Hankel
1000, //N_S2
1000, //N_DS
50, //N_norm
50, //N_r_3d
25, //N_k_3d
20, //N_a_halo
};

struct cos{
     int ORDER ;
     double vt_max;
     double vt_min;
     double vt_bin_max;
     double vt_bin_min;
     int ni ;
     int nj ;
};

void SVD_inversion(gsl_matrix *cov, gsl_matrix *inverseSVD,int Nmatrix)
{
    int i,j;
    gsl_matrix *V=gsl_matrix_calloc(Nmatrix,Nmatrix);
    gsl_matrix *U=gsl_matrix_calloc(Nmatrix,Nmatrix);
    gsl_vector *S=gsl_vector_calloc(Nmatrix);
    gsl_vector *work=gsl_vector_calloc(Nmatrix);

    gsl_matrix_memcpy(U,cov);

    gsl_linalg_SV_decomp(U,V,S,work);

    for (i=0;i<Nmatrix;i++){
        gsl_vector *b=gsl_vector_calloc(Nmatrix);
        gsl_vector *x=gsl_vector_calloc(Nmatrix);
        gsl_vector_set(b,i,1.0);
        gsl_linalg_SV_solve(U,V,S,b,x);
        for (j=0;j<Nmatrix;j++){
            gsl_matrix_set(inverseSVD,i,j,gsl_vector_get(x,j));
        }
        gsl_vector_free(b);
        gsl_vector_free(x);
    }
    gsl_matrix_free(V);
    gsl_matrix_free(U);
    gsl_vector_free(S);
    gsl_vector_free(work);
}

void invert_matrix_colesky(gsl_matrix *A)
{
  gsl_linalg_cholesky_decomp (A);  // Adummy will be overwritten  */
  gsl_linalg_cholesky_invert (A);
}

double int_gsl_integrate_insane_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter)
{
  double res, err;
  gsl_integration_cquad_workspace *w = gsl_integration_cquad_workspace_alloc(niter);
  gsl_function F;
  F.function = func;
  F.params  = arg;
 gsl_integration_cquad(&F,a,b,0,precision.insane,w,&res,&err,0);
  if(NULL!=error)
    *error=err;
  gsl_integration_cquad_workspace_free(w);
  return res;
}

double int_gsl_integrate_high_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter)
{
  double res, err;
  gsl_integration_cquad_workspace *w = gsl_integration_cquad_workspace_alloc(niter);
  gsl_function F;
  F.function = func;
  F.params  = arg;
 gsl_integration_cquad(&F,a,b,0,precision.high,w,&res,&err,0);
  if(NULL!=error)
    *error=err;
  gsl_integration_cquad_workspace_free(w);
  return res;
}

double int_gsl_integrate_medium_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter)
{
  double res, err;
  gsl_integration_cquad_workspace *w = gsl_integration_cquad_workspace_alloc(niter);
  gsl_function F;
  F.function = func;
  F.params  = arg;
 gsl_integration_cquad(&F,a,b,0,precision.medium,w,&res,&err,0);
  if(NULL!=error)
    *error=err;
  gsl_integration_cquad_workspace_free(w);
  return res;
}

double int_gsl_integrate_low_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter)
{
  double res, err;
  gsl_integration_cquad_workspace *wcrude = gsl_integration_cquad_workspace_alloc(niter);
  gsl_function F;
  F.function = func;
  F.params  = arg;
 gsl_integration_cquad(&F,a,b,0,precision.low,wcrude,&res,&err,0);
  if(NULL!=error)
    *error=err;
  gsl_integration_cquad_workspace_free(wcrude);
  return res;
}

double int_gsl_integrate_cov_precision(double (*func)(double, void*),void *arg,double a, double b, double *error, int niter)
{
  double res, err;
  gsl_integration_cquad_workspace *wcrude = gsl_integration_cquad_workspace_alloc(niter);
  gsl_function F;
  F.function = func;
  F.params  = arg;
 gsl_integration_cquad(&F,a,b,1e-16,precision.high,wcrude,&res,&err,0);
  if(NULL!=error)
    *error=err;
  gsl_integration_cquad_workspace_free(wcrude);
  return res;
}

int line_count(char *filename)
{
  FILE *n;
  n = fopen(filename,"r");
  if (!n){printf("line_count: %s not found!\nEXIT!\n",filename);exit(1);}
  int ch =0, prev =0,number_of_lines = 0;

  do
    {
    prev = ch;
    ch = fgetc(n);
    if(ch == '\n')
    	number_of_lines++;
    } while (ch != EOF);
  fclose(n);
  // last line might not end with \n, but if previous character does, last line is empty
  if(ch != '\n' && prev !='\n' && number_of_lines != 0)
    number_of_lines++;
  return number_of_lines;
}

double **create_double_matrix(long nrl, long nrh, long ncl, long nch)
/* allocate a double matrix with subscript range m[nrl..nrh][ncl..nch] */
{
	long i, nrow=nrh-nrl+1,ncol=nch-ncl+1;
	double **m;

	/* allocate pointers to rows */
	m=(double **) calloc(nrow+NR_END,sizeof(double*));
	if (!m) error("allocation failure 1 in create_double_matrix()");
	m += NR_END;
	m -= nrl;

	/* allocate rows and set pointers to them */
	m[nrl]=(double *) calloc(nrow*ncol+NR_END,sizeof(double));
	if (!m[nrl]) error("allocation failure 2 in create_double_matrix()");
	m[nrl] += NR_END;
	m[nrl] -= ncl;

	for(i=nrl+1;i<=nrh;i++) m[i]=m[i-1]+ncol;

	/* return pointer to array of pointers to rows */
	return m;
}

void free_double_matrix(double **m, long nrl, long nrh, long ncl, long nch)
		/* free a double matrix allocated by create_double_matrix() */
{
	free((FREE_ARG) (m[nrl]+ncl-NR_END));
	free((FREE_ARG) (m+nrl-NR_END));
}

double *create_double_vector(long nl, long nh)
/* allocate a double vector with subscript range v[nl..nh] */
{
	double *v;

	v=(double *) calloc(nh-nl+1+NR_END,sizeof(double));
	if (!v) error("allocation failure in double vector()");
	return v-nl+NR_END;
}

int *int_vector(long nl, long nh)
			/* allocate a int vector with subscript range v[nl..nh] */
{
	int *v;

	v=(int *)calloc(nh-nl+1+NR_END, sizeof(int));
	if (!v) error("allocation failure in int vector()");
	return v-nl+NR_END;
}

long *long_vector(long nl, long nh)
			/* allocate a int vector with subscript range v[nl..nh] */
{
	long *v;

	v=(long *)calloc(nh-nl+1+NR_END,sizeof(long));
	if (!v) error("allocation failure in int vector()");
	return v-nl+NR_END;
}

void free_double_vector(double *v, long nl, long nh)
		/* free a double vector allocated with vector() */
{
	free((FREE_ARG) (v+nl-NR_END));
}

void error(char *s)
{
	printf("error:%s\n ",s);
	exit(1);
}

/* ============================================================ *
 * Interpolates f at the value x, where f is a double[n] array,	*
 * representing a function between a and b, stepwidth dx.	*
 * 'lower' and 'upper' are powers of a logarithmic power law	*
 * extrapolation. If no	extrapolation desired, set these to 0	*
 * ============================================================ */
double interpol(double *f, int n, double a, double b, double dx, double x, double lower, double upper)
{
	double r;
	int  i;
	if (x < a) {
		if (lower==0.) {
			//error("value too small in interpol");
			return 0.0;
		}
		return f[0] + lower*(x - a);
	}
	r = (x - a)/dx;
	i = (int)(floor(r));
	if (i+1 >= n) {
		if (upper==0.0) {
			if (i+1==n) {
				return f[i];  /* constant extrapolation */
			} else {
				//error("value too big in interpol");
				return 0.0;
			}
		} else {
			return f[n-1] + upper*(x-b); /* linear extrapolation */
		}
	} else {
		return (r - i)*(f[i+1] - f[i]) + f[i]; /* interpolation */
	}
}
double interpol_fitslope(double *f, int n, double a, double b, double dx, double x, double lower)
{
	double r;
	int  i,fitrange;
	if (x < a) {
		if (lower==0.) {
			//error("value too small in interpol");
			return 0.0;
		}
		return f[0] + lower*(x - a);
	}
	r = (x - a)/dx;
	i = (int)(floor(r));
	if (i+1 >= n) {
		if (n > 50){fitrange =5;}
		else{fitrange = (int)floor(n/10);}
		double upper = (f[n-1] - f[n-1-fitrange])/(dx*fitrange);
		return f[n-1] + upper*(x-b); /* linear extrapolation */

	} else {
		return (r - i)*(f[i+1] - f[i]) + f[i]; /* interpolation */
	}
}

/* ============================================================ *
 * like interpol, but f beeing a 2d-function			*
 * 'lower' and 'upper' are the powers of a power law extra-	*
 * polation in the second argument				*
 * ============================================================ */
double interpol2d(double **f,
		      int nx, double ax, double bx, double dx, double x,
		      int ny, double ay, double by, double dy, double y,
		      double lower, double upper)
{
	double t, dt, s, ds;
	int i, j;
	if (x < ax) {
		return 0.;
//		error("value too small in interpol2d");
	}

	if (x > bx) {
		return 0.;//
//		printf("%le %le\n",x,bx);
	//	error("value too big in interpol2d");
	}
	t = (x - ax)/dx;
	i = (int)(floor(t));
	dt = t - i;
	if (y < ay) {
		return ((1.-dt)*f[i][0] + dt*f[i+1][0]) + (y-ay)*lower;
	} else if (y > by) {
		return ((1.-dt)*f[i][ny-1] + dt*f[i+1][ny-1]) + (y-by)*upper;
	}
	s = (y - ay)/dy;
	j = (int)(floor(s));
	ds = s - j;
	if ((i+1==nx)&&(j+1==ny)) {
	  //printf("%d %d\n",i+1,j+1);
	  return (1.-dt)*(1.-ds)*f[i][j];

	}
	if (i+1==nx){
	  //printf("%d %d\n",i+1,j+1);
	  return (1.-dt)*(1.-ds)*f[i][j]+ (1.-dt)*ds*f[i][j+1];
	}
	if (j+1==ny){
	  //printf("%d %d\n",i+1,j+1);
	  return (1.-dt)*(1.-ds)*f[i][j]+ dt*(1.-ds)*f[i+1][j];
	}
	return (1.-dt)*(1.-ds)*f[i][j] +(1.-dt)*ds*f[i][j+1] + dt*(1.-ds)*f[i+1][j] + dt*ds*f[i+1][j+1];
}
double interpol2d_fitslope(double **f,
		      int nx, double ax, double bx, double dx, double x,
		      int ny, double ay, double by, double dy, double y,
		      double lower)
{
	double t, dt, s, ds, upper;
	int i, j, fitrange;
	if (x < ax) {
		return 0.;
//		error("value too small in interpol2d");
	}

	if (x > bx) {
		return 0.;
//	printf("%le %le\n",x,bx);
//	  error("value too big in interpol2d");
	}
	t = (x - ax)/dx;
	i = (int)(floor(t));
	dt = t - i;
	if (y < ay) {
		return ((1.-dt)*f[i][0] + dt*f[i+1][0]) + (y-ay)*lower;
	} else if (y > by) {
		if (ny > 25){fitrange =5;}
		else{fitrange = (int)floor(ny/5);}
		upper = ((1.-dt)*(f[i][ny-1] - f[i][ny-1-fitrange])+dt*(f[i+1][ny-1] - f[i+1][ny-1-fitrange]))/(dy*fitrange);
		return ((1.-dt)*f[i][ny-1] + dt*f[i+1][ny-1]) + (y-by)*upper;
	}
	s = (y - ay)/dy;
	j = (int)(floor(s));
	ds = s - j;
	if ((i+1==nx)&&(j+1==ny)) {
	  //printf("%d %d\n",i+1,j+1);
	  return (1.-dt)*(1.-ds)*f[i][j];

	}
	if (i+1==nx){
	  //printf("%d %d\n",i+1,j+1);
	  return (1.-dt)*(1.-ds)*f[i][j]+ (1.-dt)*ds*f[i][j+1];
	}
	if (j+1==ny){
	  //printf("%d %d\n",i+1,j+1);
	  return (1.-dt)*(1.-ds)*f[i][j]+ dt*(1.-ds)*f[i+1][j];
	}
	return (1.-dt)*(1.-ds)*f[i][j] +(1.-dt)*ds*f[i][j+1] + dt*(1.-ds)*f[i+1][j] + dt*ds*f[i+1][j+1];
}

void cdgamma(fftw_complex x, fftw_complex *res)
{
	double  xr, xi, wr, wi, ur, ui, vr, vi, yr, yi, t;

	xr = (double) x[0];
	xi = (double) x[1];

	if (xr<0) {
		wr = 1 - xr;
		wi = -xi;
	} else {
		wr = xr;
		wi = xi;
	}

	ur = wr + 6.00009857740312429;
	vr = ur * (wr + 4.99999857982434025) - wi * wi;
	vi = wi * (wr + 4.99999857982434025) + ur * wi;
	yr = ur * 13.2280130755055088 + vr * 66.2756400966213521 +
			0.293729529320536228;
	yi = wi * 13.2280130755055088 + vi * 66.2756400966213521;
	ur = vr * (wr + 4.00000003016801681) - vi * wi;
	ui = vi * (wr + 4.00000003016801681) + vr * wi;
	vr = ur * (wr + 2.99999999944915534) - ui * wi;
	vi = ui * (wr + 2.99999999944915534) + ur * wi;
	yr += ur * 91.1395751189899762 + vr * 47.3821439163096063;
	yi += ui * 91.1395751189899762 + vi * 47.3821439163096063;
	ur = vr * (wr + 2.00000000000603851) - vi * wi;
	ui = vi * (wr + 2.00000000000603851) + vr * wi;
	vr = ur * (wr + 0.999999999999975753) - ui * wi;
	vi = ui * (wr + 0.999999999999975753) + ur * wi;
	yr += ur * 10.5400280458730808 + vr;
	yi += ui * 10.5400280458730808 + vi;
	ur = vr * wr - vi * wi;
	ui = vi * wr + vr * wi;
	t = ur * ur + ui * ui;
	vr = yr * ur + yi * ui + t * 0.0327673720261526849;
	vi = yi * ur - yr * ui;
	yr = wr + 7.31790632447016203;
	ur = log(yr * yr + wi * wi) * 0.5 - 1;
	ui = atan2(wi, yr);
	yr = exp(ur * (wr - 0.5) - ui * wi - 3.48064577727581257) / t;
	yi = ui * (wr - 0.5) + ur * wi;
	ur = yr * cos(yi);
	ui = yr * sin(yi);
	yr = ur * vr - ui * vi;
	yi = ui * vr + ur * vi;
	if (xr<0) {
		wr = xr * 3.14159265358979324;
		wi = exp(xi * 3.14159265358979324);
		vi = 1 / wi;
		ur = (vi + wi) * sin(wr);
		ui = (vi - wi) * cos(wr);
		vr = ur * yr + ui * yi;
		vi = ui * yr - ur * yi;
		ur = 6.2831853071795862 / (vr * vr + vi * vi);
		yr = ur * vr;
		yi = ur * vi;
	}

	(*res)[0]=yr; (*res)[1]=yi;
}


#ifndef _PK_SLOPE_H_
#define _PK_SLOPE_H_

#include <fstream>
#include <string>
#include <vector>
#include <cmath>
#include <vector_types.h>
#include <gsl/gsl_deriv.h>
#include "cspline.h"
#include "file_check.h"

template <typename T> class pkSlope{
    int N;
    std::vector<T> k, Pk, n;
    cspline<T> slope;

    void get_pk_spline();

    void calculate_n();

    void get_n_spline();

    public:
        cspline<T> power;

        pkSlope();

        pkSlope(std::string in_pk_file);

        void initialize(std::string in_pk_file);

        void calculate();

        void get_device_spline(std::vector<float4> &d_spline);

        void get_device_spline(std::vector<double4> &d_spline);

};

template <typename T> double f(double x, void *params) {
    cspline<T> Pow = *(cspline<T> *)params;
    return Pow.evaluate(x);
}

template <typename T> void pkSlope<T>::get_pk_spline() {
    pkSlope<T>::power.initialize(pkSlope<T>::k, pkSlope<T>::Pk);
}

template <typename T> void pkSlope<T>::calculate_n() {
    gsl_function F;
    double result, err;

    F.function = &f<T>;
    F.params = &power;

    for (int i = 0; i < pkSlope<T>::N; ++i) {
        if (i == 0) {
            gsl_deriv_forward(&F, pkSlope<T>::k[i], 1e-8, &result, &err);
        } else if (i == pkSlope<T>::N - 1) {
            gsl_deriv_backward(&F, pkSlope<T>::k[i], 1e-8, &result, &err);
        } else {
            gsl_deriv_central(&F, pkSlope<T>::k[i], 1e-8, &result, &err);
        }
        pkSlope<T>::n.push_back(result);
    }
}

template <typename T> void pkSlope<T>::get_n_spline() {
    pkSlope<T>::slope.initialize(pkSlope<T>::k, pkSlope<T>::n);
}

template <typename T> pkSlope<T>::pkSlope() {
    pkSlope<T>::N = 0;
}

template <typename T> pkSlope<T>::pkSlope(std::string in_pk_file) {
    pkSlope<T>::initialize(in_pk_file);
}

template <typename T> void pkSlope<T>::initialize(std::string in_pk_file) {
    if (check_file_exists(in_pk_file)) {
        std::ifstream fin(in_pk_file);
        while(!fin.eof()) {
            T kt, pt;
            fin >> kt >> pt;
            if (!fin.eof()) {
                pkSlope<T>::k.push_back(log10(kt));
                pkSlope<T>::Pk.push_back(log10(pt));
            }
        }
        fin.close();

        pkSlope<T>::N = pkSlope<T>::k.size();
    }
}

template <typename T> void pkSlope<T>::calculate() {
    pkSlope<T>::get_pk_spline();
    pkSlope<T>::calculate_n();
    pkSlope<T>::get_n_spline();
}

template <typename T> void pkSlope<T>::get_device_spline(std::vector<float4> &d_spline) {
    pkSlope<T>::slope.set_pointer_for_device(d_spline);
}

template <typename T> void pkSlope<T>::get_device_spline(std::vector<double4> &d_spline) {
    pkSlope<T>::slope.set_pointer_for_device(d_spline);
}

#endif


#include <stdio.h>
#include <math.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_multiroots.h>
#include "ray.h"

#define pi M_PI

//#############################################################################
//#############################################################################
//Mode finder

//Data structure for julia cavity data and mode parameters
typedef struct {
    //Julia cavity info and functions -----------------------------------------
    //pointer to Boundary object
    void *bnd;
    //pointer to radius function
    double (*rfunc_p)(void *bnd,double theta);
    //pointer to (mutating) radius and normal vector angle function
    void (*rsys_p)(void *bnd, double theta, double results[]);
    //pointer to RefractiveIndex object
    void *idx;
    //pointer to refractive index value function
    double (*nfunc_p)(void *idx, double r, double theta);
    //pointer to (mutating) refractive index value and derivative function
    void (*nderiv_p)(void *idx, double r, double theta, double results[]);

    //mode-specific parameters ------------------------------------------------
    long order; //number of bounces in a period
    double rtol; //stopping criterion tolerance

    //preallocated arrays -----------------------------------------------------
    double *raypath_r;
    double *raypath_theta;
    long *bounceindices;
    double *bouncepts_chi;
    long *lengths;
    double *modebounces;
} modeinfo;

//#############################################################################
//Function for rootfinding
int rootfunc(const gsl_vector *x, void *params, gsl_vector *f){

    //params is pointer to modeinfo struct containing both cavity info and
    //simulation parameters
    modeinfo *mip = (modeinfo *)params;
    const long order = (*mip).order;
    double *raypath_r = (double *)(*mip).raypath_r;
    double *raypath_theta = (double *)(*mip).raypath_theta;
    long *bounceindices = (long *)(*mip).bounceindices;
    double *bouncepts_chi = (double *)(*mip).bouncepts_chi;
    long *lengths = (long *)(*mip).lengths;

    //Extract initial condition data (print for debugging)
    const double theta0 = gsl_vector_get(x,0), chi0 = gsl_vector_get(x,1);
    //printf("theta0 = %.8f, chi0 = %.8f\n",theta0,chi0);

    //Check for sensible values
    if(fabs(chi0) > pi/2){
        //Out of domain
        return GSL_EDOM;
    }

    //Setup initial ODE coordinate vector components
    double results[2];
    (*(*mip).rsys_p)((*mip).bnd,theta0,results);
    //result[0] = r0, results[1] = normang0
    const double n = (*(*mip).nfunc_p)((*mip).idx,results[0],theta0);
    const double phi0 = results[1] + chi0;
    const double pr0 = n*cos(phi0-theta0);
    const double ptheta0 = n*results[0]*sin(phi0-theta0);

    //Compute ray
    rayevolve(
        //Storage arrays
        raypath_r,raypath_theta,bounceindices,bouncepts_chi,lengths,
        //Initial conditions
        results[0],theta0,pr0,ptheta0,
        //Simulation parameters
        200.0,order+1,1e-12,1e-12,
        //Cavity parameters
        (*mip).bnd,(*mip).rfunc_p,(*mip).rsys_p,
        (*mip).idx,(*mip).nfunc_p,(*mip).nderiv_p);

    //Store difference in initial and final bounce locations (print for debugging)
    const double dtheta = fmod(raypath_theta[bounceindices[order]-1] - raypath_theta[bounceindices[0]-1] + pi, 2*pi) - pi;
    const double dchi = bouncepts_chi[order] - bouncepts_chi[0];
    gsl_vector_set(f,0,dtheta);
    gsl_vector_set(f,1,dchi);
    //printf("dtheta = %.8f, dchi = %.8f\n",dtheta,dchi);

    //Test stopping criterion
    if(gsl_multiroot_test_residual(f,(*mip).rtol) != GSL_CONTINUE){
        //Record mode bounce data in results array (to avoid running simulation again)
        double *modebounces = (double *)(*mip).modebounces;
        //theta values stored in first half of array (first column after reshaping)
        //chi values stored in second half of array (second column after reshaping)
        int i;
        for(i=0; i<order; i+=1){
            modebounces[i] = raypath_theta[bounceindices[i]-1];
            modebounces[order+i] = bouncepts_chi[i];
        }
    }

    return GSL_SUCCESS;
}

//#############################################################################
//Solver function
int findmode(
    //mode-specific variables
    long order, double modebounces[],
    //initial position on PSS
    double theta0, double chi0,
    //root-finding algorithm parameters
    double rtol, long maxiter,
    //Cavity parameters
    void *bnd, double (*rfunc_p)(void *bnd,double theta),
    void (*rsys_p)(void *bnd, double theta,double results[]),
    void *idx, double (*nfunc_p)(void *idx, double r, double theta),
    void (*nderiv_p)(void *idx, double r, double theta,double results[])){

        //Preallocate arrays for reuse
        long bounceindices[30], lengths[2];
        double raypath_r[50000],raypath_theta[50000],bouncepts_chi[30];

        //Condense input mode data into modeinfo structs
        modeinfo mi = {bnd,rfunc_p,rsys_p,idx,nfunc_p,nderiv_p,order,rtol,
            raypath_r,raypath_theta,bounceindices,bouncepts_chi,lengths,modebounces};

        //Store initial guess
        gsl_vector *x0 = gsl_vector_alloc(2);
        gsl_vector_set(x0,0,theta0);
        gsl_vector_set(x0,1,chi0);

        //Initiate solver
        const gsl_multiroot_fsolver_type *fs_t = gsl_multiroot_fsolver_dnewton;
        gsl_multiroot_fsolver *fs = gsl_multiroot_fsolver_alloc(fs_t,2);
        gsl_multiroot_function mrf = {rootfunc,2,&mi};
        gsl_multiroot_fsolver_set(fs,&mrf,x0);
        gsl_set_error_handler_off();

        int status, count = 0;

        //Iterate solver
        while(count != maxiter){
            status = gsl_multiroot_fsolver_iterate(fs);
            if(status != GSL_SUCCESS) break; //Some error has occurred

            gsl_vector *f = gsl_multiroot_fsolver_f(fs);
            if(gsl_multiroot_test_residual(f,rtol) != GSL_CONTINUE) break; //Done!

            count += 1;
        }

        //Free memory
        gsl_multiroot_fsolver_free(fs);
        gsl_vector_free(x0);

        //In case of error or aborted computation, indicate no result
        if((status != GSL_SUCCESS) || (count == maxiter)) return 1;

        return 0;
}


/* Copyright 2017 Peter Williams and collaborators
 * Licensed under the MIT License.
 */

#include <Python.h>

#include <stdio.h>
#include <math.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_integration.h>
#include <gsl/gsl_poly.h>

#define COUNT(arr) (sizeof(arr) / sizeof(arr[0]))

typedef enum result_t {
    RESULT_OK = 0,
    RESULT_LAMBDA_FAILED = 1,
    RESULT_X_FAILED = 2,
    RESULT_GSL_ERROR = 3,
    RESULT_HIT_CRITICAL_POINT = 4,
} result_t;

typedef enum handedness_t {
    R_MODE = 1,
    L_MODE = -1,
} handedness_t;

typedef enum wave_filtering_t {
    WF_ALL_WAVES = 0,
    WF_NO_FORWARD_MODES = 1 << 0,
    WF_NO_BACKWARD_MODES = 1 << 1,
} wave_filtering_t;

typedef enum integration_type_t {
    ITYPE_BOUNCE_AVERAGED,
    ITYPE_LOCAL,
} integration_type_t;

typedef struct parameters_t {
    double E; /* dimensionless kinetic energy */
    double sin_alpha; /* sine of the pitch angle */
    double Omega_e; /* local electron gyrofrequency in rad/s */
    handedness_t handedness; /* handedness of wave; `s` in Summers/Shprits notation */
    double alpha_star; /* (cyclotron freq / plasma freq)**2 */
    double R; /* magnetic wave energy perturbation: (delta B / B)**2 */
    double x_m; /* center of wave frequency spectrum in units of cyclotron freq */
    double delta_x; /* width of wave frequency spectrum in units of cyclotron freq */
    double max_wave_latitude; /* maximum latitude at which waves are found, in radians */
    wave_filtering_t wave_filtering; /* which kinds of waves to filter out */
} parameters_t;

typedef struct state_t {
    parameters_t p;
    integration_type_t itype;
    result_t last_error;

    double gamma; /* Lorentz factor */
    double a; /* See Summers 2005, just after Eqn 24 */
    double b; /* See Summers 2005, just after Eqn 25 */
    double beta; /* velocity over speed of light */

    double latitude; /* current magnetic latitude, radians */
    double sin_alpha; /* current sin(alpha); varies since mu is conserved */
    double mu; /* current cos(alpha) */
    double Omega_e; /* current Omega_e; varies with latitude since it scales as B */

    int n_xy; /* Number of resonances encountered */
    double x[3]; /* x values of resonances */
    double y[3]; /* y values of resonances */
    double Q[3]; /* (1 - x cos alpha / y beta) */
    double R[3]; /* complex expression in summand */
} state_t;

typedef struct coefficients_t {
    double dimensionless_p; /* momentum over (m_0 * c) */
    double Daa;
    double err_Daa;
    double Dap_on_p;
    double err_Dpp_on_p2;
    double Dpp_on_p2;
    double err_Dap_on_p;
} coefficients_t;

static state_t *global_context = NULL;
static char global_err_msg[1024] = "";

static void
s05_error_handler(const char *reason, const char *file, int line, int gsl_errno)
{
    /* Go ahead and segfault if global_context is NULL. */

    int n;

    n = snprintf(global_err_msg, COUNT(global_err_msg),
                 "%s (%s:%d; %d = %s)", reason, file, line, gsl_errno, gsl_strerror(gsl_errno));
    if ((size_t) n > COUNT(global_err_msg))
        n = COUNT(global_err_msg) - 1;
    global_err_msg[n] = '\0';

    global_context->last_error = RESULT_GSL_ERROR;
}

void
summers2005_debug_hook(void)
{
    printf("Debug hook called.\n");
}

static gsl_integration_workspace *integ_workspace = NULL; /* We're lame and leak these */
static gsl_poly_complex_workspace *poly5_workspace = NULL;
static gsl_poly_complex_workspace *poly7_workspace = NULL;

const size_t INTEG_WS_SIZE = 512;
const double integ_abs_tol = 0.0;
const double integ_rel_tol = 1e-6;

const int lambda = -1; /* signifies that we're looking at electrons */
const double epsilon = 0.0005446; /* electron-to-proton mass ratio */
const double epsm1 =  -0.9994554; /* epsilon - 1; C is lame and can't do the math with consts! */
const double sigma = 2.; /* range of `x` that we consider */
const double pi_on_2nu = 0.89039; /* pi/2nu, nu = sqrt(pi)erf(sigma); see after Shprits 06 Eqn 5 */

static inline result_t
apply_latitude(double latitude, state_t *state)
{
    int i;

    if (state->last_error != RESULT_OK)
        return state->last_error;

    state->latitude = latitude;

    double scl = cos(latitude); /* sin of colatitude = cos of latitude */
    double r = sqrt(4 - 3 * scl * scl) / pow(scl, 6); /* B(lam) / B_eq */

    state->Omega_e = state->p.Omega_e * r;
    state->sin_alpha = state->p.sin_alpha * sqrt(r);
    state->mu = sqrt(1 - state->sin_alpha * state->sin_alpha); /* cos(alpha) */

    if (state->mu == 0.) /* happens at the very top of our bounce trajectory */
        state->mu = 1e-20;

    /* Find the critical `x` and `y` values. We do this following Appendix A
     * and Equation 24 of Summers 2005. */

    state->n_xy = 0;

    double c[5] = { 0 };
    double z[8] = { 0 };

    const int s = state->p.handedness;
    const double a = state->a;
    const double a2 = a * a;
    const double bm2 = pow(state->beta * state->mu, 2);

    c[0] = -a2 * epsilon;
    c[1] = a2 * s * epsm1 - 2 * a * epsilon;
    c[2] = a2 + 2 * a * s * epsm1 - epsilon + bm2 * (state->b + epsilon);
    c[3] = 2 * a + s * epsm1 - bm2 * s * epsm1;
    c[4] = 1 - bm2;

    gsl_poly_complex_solve(c, 5, poly5_workspace, z);

    if (state->last_error != RESULT_OK) {
        strncat(global_err_msg, " (while finding x/y roots)", COUNT(global_err_msg) - 1);
        return state->last_error;
    }

    for (i = 0; i < 4; i++) {
        if (z[2*i + 1] != 0.) /* Imaginary root? */
            continue;

        double x = z[2*i];

        if (x < 0) /* Non-physical root? */
            continue;

        if (x < state->p.x_m - sigma * state->p.delta_x || x > state->p.x_m + sigma * state->p.delta_x)
            continue; /* Out of waveband? */

        double y = (x + state->a) / (state->beta * state->mu);

        if ((state->p.wave_filtering & WF_NO_FORWARD_MODES) && y > 0)
            continue;

        if ((state->p.wave_filtering & WF_NO_BACKWARD_MODES) && y < 0)
            continue;

        state->x[state->n_xy] = x;
        state->y[state->n_xy] = y;
        state->n_xy++;
    }

    if (state->n_xy > 3) {
        snprintf(global_err_msg, COUNT(global_err_msg), "expect 0 to 3 X/Y solutions; got %d", state->n_xy);
        return RESULT_X_FAILED;
    }

    /* Now we calculate F(x,y) = dx/dy, following Appendix C of Summers 2005,
     * and related values that contribute to the sums in the integrands. `Q`
     * is `1 - x cosa / y beta` and `R` is the factor that is in common
     * between all three calculations. */

    for (i = 0; i < state->n_xy; i++) {
        double x = state->x[i];
        double y = state->y[i];

        c[0] = epsilon * (state->b + epsilon);
        c[1] = -0.5 * s * epsm1 * (state->b + 4 * epsilon);
        c[2] = 1 - 4 * epsilon + epsilon * epsilon;
        c[3] = 2 * s * epsm1;
        c[4] = 1.;

        double g = gsl_poly_eval(c, 5, x);
        double F = y * pow((x - s) * (x + s * epsilon), 2) / (x * g);

        state->Q[i] = 1 - x * state->mu / (y * state->beta);

        double k0 = pow((x - state->p.x_m) / state->p.delta_x, 2);
        double k1 = fabs(state->beta * state->mu - F);

        if (k1 < 1e-5) {
            snprintf(global_err_msg, COUNT(global_err_msg), "hit a critical point (%f, %f, %f)",
                     state->sin_alpha, x, y);
            return RESULT_HIT_CRITICAL_POINT;
        }

        state->R[i] = state->p.R * fabs(F) * exp(-k0) / (state->p.delta_x * k1);
    }

    return RESULT_OK;
}

typedef double (*gsl_signature)(double, void *);

static double
daa_integrand(double latitude, state_t *state)
{
    result_t r;

    if (state->last_error != RESULT_OK)
        return 0;

    if ((r = apply_latitude(latitude, state)) != RESULT_OK) {
        state->last_error = r;
        return 0;
    }

    int i;
    double v = 0;

    for (i = 0; i < state->n_xy; i++)
        v += state->R[i] * pow(state->Q[i], 2);

    if (state->itype == ITYPE_BOUNCE_AVERAGED)
        v *= state->mu * pow(cos(state->latitude), 7);

    return v * state->Omega_e;
}

static double
dap_on_p_integrand(double latitude, state_t *state)
{
    result_t r;

    if (state->last_error != RESULT_OK)
        return 0;

    if ((r = apply_latitude(latitude, state)) != RESULT_OK) {
        state->last_error = r;
        return 0;
    }

    int i;
    double v = 0;

    for (i = 0; i < state->n_xy; i++)
        v += state->R[i] * state->Q[i] * state->x[i] / state->y[i];

    /* Here we transform the sqrt(1 + 3 sin^2 lambda) term to use the cos
     * instead */

    double cos_lam = cos(state->latitude);
    double sa = state->sin_alpha;

    if (state->itype == ITYPE_BOUNCE_AVERAGED)
        v *= state->mu * cos_lam * sqrt(4 - 3 * cos_lam * cos_lam) / sa;

    return v * state->Omega_e * sa;
}

static double
dpp_on_p2_integrand(double latitude, state_t *state)
{
    result_t r;

    if (state->last_error != RESULT_OK)
        return 0;

    if ((r = apply_latitude(latitude, state)) != RESULT_OK) {
        state->last_error = r;
        return 0;
    }

    int i;
    double v = 0;

    for (i = 0; i < state->n_xy; i++)
        v += state->R[i] * pow(state->x[i] / state->y[i], 2);

    /* Same transform as in dap_on_p. */

    double cos_lam = cos(state->latitude);
    double sa = state->sin_alpha;

    if (state->itype == ITYPE_BOUNCE_AVERAGED)
        v *= cos_lam * sqrt(4 - 3 * cos_lam * cos_lam) / state->mu;

    return v * state->Omega_e * sa * sa;
}

static result_t
calc_coefficients(parameters_t *params, coefficients_t *coeffs)
{
    gsl_error_handler_t *prev_handler;
    gsl_function integrand;
    state_t state;
    double c[7] = { 0 };
    double z[12] = { 0 };
    double lambda_m = -1;
    double result, abserr, k;
    int i;

    if (params->sin_alpha > 0.99985) /* ~= 89 degrees */
        params->sin_alpha = 0.99985;

    global_context = &state;
    prev_handler = gsl_set_error_handler(s05_error_handler);

    if (integ_workspace == NULL) {
        integ_workspace = gsl_integration_workspace_alloc(INTEG_WS_SIZE);
        poly5_workspace = gsl_poly_complex_workspace_alloc(5);
        poly7_workspace = gsl_poly_complex_workspace_alloc(7);
    }

    state.p = *params;
    state.itype = ITYPE_BOUNCE_AVERAGED;
    state.last_error = RESULT_OK;

    /* Figure out the mirroring latitude; Shprits (2006) equation 10. */

    double f0 = pow(state.p.sin_alpha, 4);
    c[0] = -4 * f0;
    c[1] = 3 * f0;
    c[6] = 1.;

    gsl_poly_complex_solve(c, 7, poly7_workspace, z);

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while finding lambda_m)", COUNT(global_err_msg) - 1);
        return state.last_error;
    }

    for (i = 0; i < 6; i++) {
        if (z[2*i + 1] == 0. && z[2*i] > 0) {
            lambda_m = acos(sqrt(z[2*i]));
            break;
        }
    }

    if (!(lambda_m >= 0 && lambda_m <= M_PI_2)) {
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        snprintf(global_err_msg, COUNT(global_err_msg), "failed to compute lambda_m (%.16lf)", lambda_m);
        return RESULT_LAMBDA_FAILED;
    }

    /* It can happen that lambda_m is just a tiiiiiny bit too large such that
     * when we compute the sine of the pitch angle at lambda_m, we get
     * something bigger than one. Witness our beautiful workaround.
     */

    double cos_lat = cos(lambda_m);

    if (sqrt(4 - 3 * cos_lat * cos_lat) / (pow(cos_lat, 6) * state.p.sin_alpha * state.p.sin_alpha) > 1.)
        lambda_m -= 1e-7;

    /* If a latitude limit on waves is imposed, truncate lambda_m further. */

    if (lambda_m > state.p.max_wave_latitude)
        lambda_m = state.p.max_wave_latitude;

    /* Pre-compute quantities that do not depend on lambda. When lambda
     * changes, B and alpha change. Consequently Omega_e changes too. In our
     * model the other parameters stay fixed.
     */

    state.gamma = state.p.E + 1;
    state.a = ((int) state.p.handedness) * lambda / state.gamma;
    state.b = (1 + epsilon) / state.p.alpha_star;
    state.beta = sqrt(state.p.E * (state.p.E + 2)) / (state.p.E + 1);
    coeffs->dimensionless_p = state.gamma * state.beta;

    /* Shprits 2006 eqn 9, credited to Lencheck+ 1971 and Shultz & Lanzerotti 1974: */

    double s0 = 1.38 - 0.32 * (params->sin_alpha + sqrt(params->sin_alpha));
    double f1 = pi_on_2nu / (pow(state.p.E + 1, 2) * s0);

    /* Let's integrate! First, D_aa. */

    integrand.function = (gsl_signature) daa_integrand;
    integrand.params = &state;

    gsl_integration_qag(
        &integrand,
        0., /* lower limit */
        lambda_m, /* upper limit */
        integ_abs_tol,
        integ_rel_tol,
        INTEG_WS_SIZE, /* allocated workspace size */
        GSL_INTEG_GAUSS51, /* integration rule */
        integ_workspace, /* workspace */
        &result, /* output: value of the integral */
        &abserr /* estimated absolute error */
    );

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while computing Daa)", COUNT(global_err_msg) - 1);
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        return state.last_error;
    }

    k = f1 / (1 - state.p.sin_alpha * state.p.sin_alpha);
    coeffs->Daa = result * k;
    coeffs->err_Daa = abserr * k;

    /* D_ap/p. */

    integrand.function = (gsl_signature) dap_on_p_integrand;

    gsl_integration_qag(
        &integrand,
        0., /* lower limit */
        lambda_m, /* upper limit */
        integ_abs_tol,
        integ_rel_tol,
        INTEG_WS_SIZE, /* allocated workspace size */
        GSL_INTEG_GAUSS51, /* integration rule */
        integ_workspace, /* workspace */
        &result, /* output: value of the integral */
        &abserr /* estimated absolute error */
    );

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while computing Dap)", COUNT(global_err_msg) - 1);
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        return state.last_error;
    }

    k = -f1 * state.p.sin_alpha / (state.beta * sqrt(1 - state.p.sin_alpha * state.p.sin_alpha));
    coeffs->Dap_on_p = result * k;
    coeffs->err_Dap_on_p = abserr * fabs(k);

    /* D_pp/p^2. */

    integrand.function = (gsl_signature) dpp_on_p2_integrand;

    gsl_integration_qag(
        &integrand,
        0., /* lower limit */
        lambda_m, /* upper limit */
        integ_abs_tol,
        integ_rel_tol,
        INTEG_WS_SIZE, /* allocated workspace size */
        GSL_INTEG_GAUSS51, /* integration rule */
        integ_workspace, /* workspace */
        &result, /* output: value of the integral */
        &abserr /* estimated absolute error */
    );

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while computing Dpp)", COUNT(global_err_msg) - 1);
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        return state.last_error;
    }

    k = f1 * pow(state.beta, -2);
    coeffs->Dpp_on_p2 = result * k;
    coeffs->err_Dpp_on_p2 = abserr * k;

    /* Huzzah, all done. */

    gsl_set_error_handler(prev_handler);
    global_context = NULL;
    return RESULT_OK;
}

static result_t
calc_unaveraged_coefficients(parameters_t *params, coefficients_t *coeffs)
{
    gsl_error_handler_t *prev_handler;
    state_t state;

    global_context = &state;
    prev_handler = gsl_set_error_handler(s05_error_handler);

    if (integ_workspace == NULL) {
        integ_workspace = gsl_integration_workspace_alloc(INTEG_WS_SIZE);
        poly5_workspace = gsl_poly_complex_workspace_alloc(5);
        poly7_workspace = gsl_poly_complex_workspace_alloc(7);
    }

    state.p = *params;
    state.itype = ITYPE_LOCAL;
    state.last_error = RESULT_OK;

    /* The structure here is paralleling calc_coefficients for maintainability. */

    state.gamma = state.p.E + 1;
    state.a = ((int) state.p.handedness) * lambda / state.gamma;
    state.b = (1 + epsilon) / state.p.alpha_star;
    state.beta = sqrt(state.p.E * (state.p.E + 2)) / (state.p.E + 1);
    coeffs->dimensionless_p = state.gamma * state.beta;

    double f1 = pi_on_2nu / pow(state.p.E + 1, 2);

    /* Compute! It all works if we pretend latitude = 0 and set `itype` to LOCAL. */

    coeffs->Daa = daa_integrand(0., &state);

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while computing Daa)", COUNT(global_err_msg) - 1);
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        return state.last_error;
    }

    coeffs->Daa *= f1;
    coeffs->err_Daa = 0;

    /* D_ap/p. */

    coeffs->Dap_on_p = dap_on_p_integrand(0., &state);

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while computing Dap)", COUNT(global_err_msg) - 1);
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        return state.last_error;
    }

    coeffs->Dap_on_p *= f1 / state.beta;
    coeffs->err_Dap_on_p = 0;

    /* D_pp/p^2. */

    coeffs->Dpp_on_p2 = dpp_on_p2_integrand(0., &state);

    if (state.last_error != RESULT_OK) {
        strncat(global_err_msg, " (while computing Dpp)", COUNT(global_err_msg) - 1);
        gsl_set_error_handler(prev_handler);
        global_context = NULL;
        return state.last_error;
    }

    coeffs->Dpp_on_p2 *= f1 * pow(state.beta, -2);
    coeffs->err_Dpp_on_p2 = 0.;

    /* Huzzah, all done. */

    gsl_set_error_handler(prev_handler);
    global_context = NULL;
    return RESULT_OK;
}

static PyObject*
get_coeffs(PyObject *self, PyObject* args)
{
    int modespec, handspec, wave_filtering;
    parameters_t params;
    coefficients_t coeffs = { 0 };
    result_t r;

    if (!PyArg_ParseTuple(args, "iiidddddddd", &modespec, &handspec, &wave_filtering,
                          &params.E,
                          &params.sin_alpha,
                          &params.Omega_e,
                          &params.alpha_star,
                          &params.R,
                          &params.x_m,
                          &params.delta_x,
                          &params.max_wave_latitude))
        return NULL;

    if (handspec == 0)
        params.handedness = R_MODE;
    else if (handspec == 1)
        params.handedness = L_MODE;
    else {
        PyErr_SetString(PyExc_RuntimeError, "unexpected handedness magic constant");
        return NULL;
    }

    params.wave_filtering = (wave_filtering_t) wave_filtering;

    if (modespec == 0)
        r = calc_coefficients(&params, &coeffs);
    else if (modespec == 1)
        r = calc_unaveraged_coefficients(&params, &coeffs);
    else {
        PyErr_SetString(PyExc_RuntimeError, "unexpected mode magic constant");
        return NULL;
    }

    if (r != RESULT_OK) {
        PyErr_SetString(PyExc_RuntimeError, global_err_msg);
        return NULL;
    }

    return Py_BuildValue("ddddddd", coeffs.dimensionless_p,
                         coeffs.Daa, coeffs.err_Daa,
                         coeffs.Dap_on_p, coeffs.err_Dap_on_p,
                         coeffs.Dpp_on_p2, coeffs.err_Dpp_on_p2);
}

static PyMethodDef methods[] = {
    { "get_coeffs", get_coeffs, METH_VARARGS|METH_KEYWORDS,
      "(mode, handedness, wfilt, E, sa, Oe, a*, R, xm, dx, mwl) -> "
      "(dp, daa, udaa, dap/p, udap/p, dpp/p2, udpp/p2)" },
    { NULL, NULL, METH_NOARGS, NULL },
};

#if PY_MAJOR_VERSION >= 3
static struct PyModuleDef module_def = {
    PyModuleDef_HEAD_INIT,
    "_summers2005",
    NULL,
    0,
    methods,
    NULL,
    NULL,
    NULL,
    NULL
};

# define INIT_RET_TYPE PyObject *
#else
# define INIT_RET_TYPE void
#endif

INIT_RET_TYPE
PyInit__summers2005(void)
{
#if PY_MAJOR_VERSION >= 3
    return PyModule_Create(&module_def);
#else
    PyImport_AddModule("_summers2005");
    Py_InitModule("_summers2005", methods);
#endif
}


/* ============================================================ *
 * decomp_eb.c							*
 * Martin Kilbinger, Liping Fu 2008, 2009			*
 * ============================================================ */
#include "decomp_eb.h"
#include <gsl/gsl_sf_gamma.h>

/* === CFHTLS Wide 3rd data release, Fu&Kilbinger (2010) === */

/* S/N, Psi=19' eta=1/50 */
const double a_FK10_SN[N_FK10]        = {0.1197730890, -0.3881211865, 0.5212557875, -0.3440507036, 0.2761305382, -0.07286690971};

/* FoM, Psi=222', eta=1/10 */
const double a_FK10_FoM_eta10[N_FK10] = {0.009877788826, 0.1061397843, -0.4300211814, 0.5451016406, -0.3372272549, 0.1716983151};

/* FoM, Psi=222', eta=1/50 */
const double a_FK10_FoM_eta50[N_FK10] = {0.1239456383, -0.3881431858, 0.5579593467, -0.3679282338, 0.1540941993, 0.01293361618};

/* First-kind Chebyshev polynomial of order n */
#define EPS 1.0e-10
double Cheby(double x, int n, error **err)
{
   double Cn;

   testErrorRetVA(x<-1-EPS || x>1+EPS, mr_range, "x = %g out of range", *err, __LINE__, -1, x);

   if (x<-1) x = -1.0;
   if (x>1) x = 1.0;

   Cn = cos (n * acos (x));
   return Cn;
}
#undef EPS

/* Second-kind Chebyshev polynomial of order n */
#define EPS 1.0e-10
double Cheby2(double x, int n, error **err)
{
   double Un;

   testErrorRetVA(x<-1-EPS || x>1+EPS, mr_range, "x = %g out of range", *err, __LINE__, -1, x);

   if (x<-1) x = -1.0;
   if (x>1) x = 1.0;

   if (x == 1) {
      Un = n + 1;
   } else if (x == -1) {
      Un = pow(-1.0, n) * (n + 1.0);
   } else {
      Un = sin((n+1.0)*acos(x))/sin(acos (x));
   }

   return Un;
}
#undef EPS

/* Legendre polynomial of order n */
double Legen(double x, int n)
{

   switch (n) {
      case 0:
	 return 1.0;
      case 1:
	 return x;
      default:
	 return (2.0*n-1.0)/(double)n*x*Legen(x,n-1)-(n-1.0)/(double)n*Legen(x,n-2);
   }
}

/* General basis function of order n */
#define EPS 1.0e-10
double C(double x, int n, poly_t poly, error **err)
{
   double c;

   testErrorRetVA(x<-1-EPS || x>1+EPS, mr_range, "x = %g out of range", *err, __LINE__, -1, x);

   if (x<-1) x = -1.0;
   if (x>1) x = 1.0;

   switch (poly) {
      case cheby :
	 c = Cheby(x, n, err);
	 forwardError(*err, __LINE__, -1.0);
	 break;
      case cheby2 :
	 c = Cheby2(x, n, err);
	 forwardError(*err, __LINE__, -1.0);
	 break;
      case legen :
	 c = Legen(x, n);
	 break;
      default :
	 *err = addErrorVA(mr_poly, "Unknown polynomial type %d", *err, __LINE__, poly);
	 return -1.0;
   }

   return c;
}
#undef EPS

/* ============================================================ *
 * The filter function for the generalised ring statistics.	*
 * FK09 (11).							*
 * ============================================================ */
#define	 EPS 1.0e-6
double Tp(double x, const double *a, int N, poly_t poly, error **err)
{
   int n;
   double res, Cn, Cnm1=0.0, Cnm2;

   if (x<-1-EPS || x>+1+EPS) return 0;

   /* NEW! */
   //if (x<-1) x = -1;
   //if (x>+1) x = +1;

   testErrorRetVA(N<=0, mr_range, "N has to be larger than zero but is %d", *err, __LINE__, 0.0, N);

   testErrorRetVA(x>+1+EPS, mr_range, "x=%.10f out of range", *err, __LINE__, 0.0, x);
   testErrorRetVA(x<-1-EPS, mr_range, "x=%.10f out of range", *err, __LINE__, 0.0, x);

   if (poly==cheby || poly==cheby2) {

      Cnm2 = 1.0;
      if (poly==cheby)       Cnm1 = x;
      else if (poly==cheby2) Cnm1 = 2.0*x;

      res  = 0.0;
      res += a[0] * Cnm2;
      if (N==1) return res;
      res += a[1] * Cnm1;
      if (N==2) return res;

      for (n=2; n<N; n++) {
	 /* Chebyshev (both T,U) recurrence relation *
	  * C_n(x) = 2 x C_{n-1}(x) - C_{n-2}(x)     */
	 Cn   = 2.0*x*Cnm1 - Cnm2;
	 res += a[n] * Cn;
	 Cnm2 = Cnm1;
	 Cnm1 = Cn;
      }

   } else {

      for (n=0,res=0.0; n<N; n++) {
	 Cn   = C(x, n, poly, err);
	 forwardError(*err, __LINE__, -1);
	 res += a[n] * Cn;
      }

   }

   return res;
}
#undef EPS

/* ============================================================ *
 * FK09 (24).							*
 * ============================================================ */
double Fn0(double x, int n, poly_t poly, error **err)
{
   double fn0, denom=0.0;
   int nn;

   if (n<0) nn = -n;
   else nn = n;

   if (poly==cheby) {
      if (nn==0) return x + 1.0;
      if (nn==1) return 0.5*(x*x - 1.0);
   } else if (poly==cheby2) {
      if (n==-1) return 0.0;  /* n, not nn: no symmetry for cheby2! */
   } else {
      *err = addErrorVA(mr_poly, "Tm for polynomial type %d not defined", *err, __LINE__, poly);
      return 0.0;
   }

   fn0 = 0.0;

   if (nn%2==0) fn0 += 1.0;   /* n even */
   else fn0 -= 1.0;           /* n odd  */

   //printf("# ** 1 %g\n", fn0);

   fn0 += x*Cheby(x, nn, err);
   forwardError(*err, __LINE__, 0.0);

   //printf("# ** 2 %g\n", fn0);

   if (poly==cheby) fn0 += n*(1 - x*x)*Cheby2(x, n-1, err);
   else if (poly==cheby2) fn0 -= (1 - x*x)*Cheby2(x, n-1, err);
   forwardError(*err, __LINE__, 0.0);

   //printf("# ** 3 %g\n", fn0);

   if (poly==cheby) denom = 1.0 - (double)(n*n);
   else if (poly==cheby2) denom = 1.0+(double)n;         /* n, not nn! */

   fn0 /= denom;

   return fn0;
}

/* ============================================================ *
 * FK09 (26)							*
 * ============================================================ */
void Fnnu(double x, int n, poly_t poly, double Fn[], error **err)
{
   double f0, fp1, fp2, fp3, fm1, fm2, fm3;

   f0  = Fn0(x, n, poly, err);      forwardError(*err, __LINE__,);
   fp1 = Fn0(x, n+1, poly, err);    forwardError(*err, __LINE__,);
   fp2 = Fn0(x, n+2, poly, err);    forwardError(*err, __LINE__,);
   fp3 = Fn0(x, n+3, poly, err);    forwardError(*err, __LINE__,);
   fm1 = Fn0(x, n-1, poly, err);    forwardError(*err, __LINE__,);
   fm2 = Fn0(x, n-2, poly, err);    forwardError(*err, __LINE__,);
   fm3 = Fn0(x, n-3, poly, err);    forwardError(*err, __LINE__,);

   Fn[0] = f0;
   Fn[1] = 0.5*(fp1 + fm1);
   Fn[2] = 0.25*(fp2 + 2.0*f0 + fm2);
   Fn[3] = 0.125*(fp3 + 3.0*fp1 + 3.0*fm1 + fm3);
}

/* ============================================================ *
 * FK09 (23)							*
 * ============================================================ */
#define EPS 1.0e-10
double alpha_explicit(double x, int n, double R, poly_t poly, error **err)
{
   double r, Fn[4], alpha;

   testErrorRet(poly!=cheby && poly!=cheby2, mr_poly,
		"Explicit expression only obtained so far for cheby1 and cheby2",
		*err, __LINE__, 0.0);

   if (x<-1.0+EPS) return 0.0;
   assert(x>-1.0);

   r = 3.0/dsqr(x + R);
   Fnnu(x, n, poly, Fn, err);
   forwardError(*err, __LINE__, 0.0);

   alpha  = R*(1.0-r*R*R)*Fn[0];
   alpha += (1.0 - 3.0*r*R*R)*Fn[1];
   alpha += -3.0*r*R*Fn[2];
   alpha += -r*Fn[3];
   alpha *= 4.0/3.0*r;

   return alpha;
}

/* ============================================================ *
 * The filter function T_- obtained from T_+, FK09 (20).	*
 * ============================================================ */
#define EPS 1.0e-10
double Tm(double x, const double *a, int N, poly_t poly, double R, error **err)
{
   double tm, c, alph;
   int n;

   testErrorRetVA(x<-1-EPS, mr_range, "x=%.20f smaller than -1", *err, __LINE__, 0.0, x);

   for (n=0,tm=0.0; n<N; n++) {
      c    = C(x, n, poly, err);                forwardError(*err, __LINE__, 0.0);
      alph = alpha_explicit(x, n, R, poly, err);
      forwardError(*err, __LINE__, 0.0);

      tm  += a[n]*(c + alph);
   }
   return tm;
}
#undef EPS

/* ================================================================== *
 * Returns the generalised ring statistic using a as the coefficients *
 * of the T_+ decomposition, see FK09 (4, 8, 11).		      *
 * The correlation functions xi+ and xi- are read from the            *
 * Nxi-dimensional arrays xip and xim, on angular scales th.          *
 * N is the size of the coefficient-array a.                          *
 * ================================================================== */
double RR_data(const double *xip, const double *xim, const double *th, const int Nxi,
	       double THETA_MIN, double THETA_MAX, const double *a, int N,
	       poly_t poly,  int pm, error **err)
{
   double theta, dtheta, res, A, B, x, summand;
   double dt1, dt2;
   int i;

   testErrorRetVA(abs(pm)!=1, mr_incompatible, "pm=%d not valid, has to be +1 or -1",
		  *err, __LINE__, 0.0, pm);

   testErrorRetVA(THETA_MIN<th[0], mr_range,
		  "THETA_MIN=%g' is smaller than minimum angular scale for xi+-, %g'",
		  *err, __LINE__, 0.0, THETA_MIN/arcmin, th[0]/arcmin);

   testErrorRetVA(THETA_MAX>th[Nxi-1], mr_range,
		  "THETA_MAX=%g' is larger than maximum angular scale for xi+-, %g'",
		  *err, __LINE__, 0.0, THETA_MAX/arcmin, th[Nxi-1]/arcmin);

   if (a==NULL) {
      *err = addError(mr_null, "Coefficients a=NULL", *err, __LINE__);
      return 0.0;
   }

   /* FK09 (8) */
   A = (THETA_MAX - THETA_MIN)/2.0;
   B = (THETA_MAX + THETA_MIN)/2.0;

   for (i=0,res=0.0; i<Nxi; i++) {
      theta  = th[i];
      if (theta < THETA_MIN) continue;
      else if (theta > THETA_MAX) break;

      /* theta[i] is the bin center */

      if (i==0) {
	 dt1 = (th[i+1] - th[i])/2.;
	 dtheta = 2*dt1;
      } else if (i==Nxi-1) {
	 dt2 = (th[i] - th[i-1])/2.;
	 dtheta = 2*dt2;
      } else {
	 dt1 = (th[i+1] - th[i])/2.;
	 dt2 = (th[i] - th[i-1])/2.;
	 dtheta = dt2 + dt1;
      }
      x = (theta-B)/A;

      if (pm==+1)  {
	 summand  = xip[i]*theta/dsqr(THETA_MAX);
	 summand *= Tp(x, a, N, poly, err);
	 forwardError(*err, __LINE__, -1.0);
      } else if (pm==-1) {
	 summand  = xim[i]*theta/dsqr(THETA_MAX);
	 summand *= Tm(x, a, N, poly, B/A, err);
	 forwardError(*err, __LINE__, -1.0);
      } else {
	 summand = 0.0;
      }

      res += summand*dtheta;
   }

   testErrorRet(!finite(res), math_infnan, "R is not finite", *err, __LINE__,
		  0.0);

   return res;
}

/* If cov_mode=outer, n, m are not used. For cov_mode=inner, a, N are not used.  *
 * If cov_mode=fromZ, the covariance using Z+ is returned and a, N are not used. */
#define EPS 1.0e-5
double cov_RR(const double *THETA_MIN, const double *THETA_MAX, const double *a, int N, poly_t poly,
	      const double *theta, const double *cov_xi, int Ntheta,
	      cov_mode_t cov_mode, int n, int m, double fac, error **err)
{
   double sum, dlogtheta, tmp, A[2], B[2], xi, xj, yi, yj, thetai, thetaj;
   int i, j;
   static double  **Cov_xi = NULL;

   /* Testing for a==NULL not possible because cov_RR is also called from get_cov_RR_inner_array */
   if (cov_mode==fromZ) {
#ifdef __MRING_H
      sum = cov_RR_Z(THETA_MIN, THETA_MAX, theta, cov_xi, Ntheta, err);
      forwardError(*err, __LINE__, 0.0);
      return sum;
#else
      *err = addError(mr_type, "cov_mode 'fromZ' not supported here", *err, __LINE__);
      return 0.0;
#endif
   }

   for (i=0; i<2; i++) {
      A[i] = (THETA_MAX[i] - THETA_MIN[i])/2.0;
      B[i] = (THETA_MAX[i] + THETA_MIN[i])/2.0;
   }

   dlogtheta = log(theta[1]) - log(theta[0]);

   for (i=0; i<2; i++) {
      testErrorRetVA(theta[0]>THETA_MIN[i], mr_range,
		     "Minumum of xi-covariance (%g=%g') larger than THETA_MIN (%g=%g')",
		     *err, __LINE__, -1, theta[0], theta[0]/arcmin, THETA_MIN[i], THETA_MIN[i]/arcmin);
      testErrorRetVA(theta[Ntheta-1]<THETA_MAX[i], mr_range,
		     "Maximum of xi-covariance scale (%g=%g') smaller than THETA_MAX (%g=%g')",
		     *err, __LINE__, -1, theta[Ntheta-1], theta[Ntheta-1]/arcmin, THETA_MAX[i],
		     THETA_MAX[i]/arcmin);
   }

   if (Cov_xi==NULL) {
      Cov_xi = sm2_matrix(0, Ntheta-1, 0, Ntheta-1, err);
      forwardError(*err, __LINE__, -1.0);
      for (i=0; i<Ntheta; i++) {
	 for (j=0; j<Ntheta; j++) {

	    Cov_xi[i][j] = cov_xi[i*Ntheta+j];

	 }

      }

      //if (fac<1.12)
	//fprintf(stderr, "WARNING: Cov-xi-factor = %g, too small for stable inversion of covariance!!!\n", fac);
   }

   sum = 0.0;
   //for (i=0; i<Ntheta; i++) {
   //thetai = theta[i];
   for (thetai=theta[0]; thetai<=theta[Ntheta-1]; thetai*=fac) {
      if (thetai<THETA_MIN[0] || thetai>THETA_MAX[0]) continue;
      xi = (thetai-B[0])/A[0];
      yi = thetai/THETA_MAX[0];

      //for (j=0; j<Ntheta; j++) {
      //thetaj = theta[j];
      for (thetaj=theta[0]; thetaj<=theta[Ntheta-1]; thetaj*=fac) {
	 if (thetaj<THETA_MIN[1] || thetaj>THETA_MAX[1]) continue;
	 xj = (thetaj-B[1])/A[1];
	 yj = thetaj/THETA_MAX[1];

	 tmp  = dsqr(dlogtheta);

	 if (cov_mode==outer) {
	    tmp *= Tp(xi, a, N, poly, err);
	    forwardError(*err, __LINE__, -1);
	    tmp *= Tp(xj, a, N, poly, err);
	    forwardError(*err, __LINE__, -1);
	 } else {
	    tmp *= C(xi, n, poly, err);  	    forwardError(*err, __LINE__, -1);
	    tmp *= C(xj, m, poly, err);  	    forwardError(*err, __LINE__, -1);
	 }

	 tmp *= dsqr(yi*yj);

	 //tmp *= cov_xi[i*Ntheta+j];
	 tmp *= sm2_interpol2d(Cov_xi, Ntheta, log(theta[0]), log(theta[Ntheta-1]), dlogtheta, log(thetai),
	 		       Ntheta, log(theta[0]), log(theta[Ntheta-1]), dlogtheta, log(thetaj), 0.0, 0.0, err);
	 forwardError(*err, __LINE__, -1.0);

	 sum += tmp;

      }

  }

   /* Comment the following two lines for faster code, but be aware that no check is done
    * whether the input covariance changes! */
   sm2_free_matrix(Cov_xi, 0, Ntheta-1, 0, Ntheta-1);
   Cov_xi = NULL;

   return sum;
}
#undef EPS

/* RR-Covariance using a diagonal xi-covariance.				 *
 * If cov_mode=outer, n, m are not used. For cov_mode=inner, a, N are not used.  *
 * If cov_mode=fromZ, the covariance using Z+ is returned and a, N are not used. */
#define EPS 1.0e-5
double cov_RR_diag_xi(const double *THETA_MIN, const double *THETA_MAX, const double *a, int N, poly_t poly,
		      const double *theta, const double *var_xi, int Ntheta,
		      int islog, error **err)
{
   double sum, dlogtheta, dtheta, tmp, A[2], B[2], xi, xj, yi, yj, thetai, thetaj;
   int i;

   for (i=0; i<2; i++) {
      A[i] = (THETA_MAX[i] - THETA_MIN[i])/2.0;
      B[i] = (THETA_MAX[i] + THETA_MIN[i])/2.0;
   }

   if (islog==0) {
      dtheta = theta[1] - theta[0];
      dlogtheta = 0.0;
   } else if (islog==1) {
      dlogtheta = log(theta[1]) - log(theta[0]);
      dtheta = 0.0;
   } else {
      *err = addErrorVA(mr_type, "Invalid flag islog = %d\n", *err, __LINE__, islog);
      return 0.0;
   }

   for (i=0; i<2; i++) {
      testErrorRetVA(theta[0]>THETA_MIN[i], mr_range,
		     "Minumum of xi-covariance (%g=%g') larger than THETA_MIN (%g=%g')",
		     *err, __LINE__, -1, theta[0], theta[0]/arcmin, THETA_MIN[i], THETA_MIN[i]/arcmin);
      testErrorRetVA(theta[Ntheta-1]<THETA_MAX[i], mr_range,
		     "Maximum of xi-covariance scale (%g=%g') smaller than THETA_MAX (%g=%g')",
		     *err, __LINE__, -1, theta[Ntheta-1], theta[Ntheta-1]/arcmin, THETA_MAX[i],
		     THETA_MAX[i]/arcmin);
   }

   sum = 0.0;
   for (i=0; i<Ntheta; i++) {
      thetai = theta[i];
      thetaj = thetai;

      if (thetai<THETA_MIN[0] || thetai>THETA_MAX[0]) continue;
      xi = (thetai-B[0])/A[0];
      yi = thetai/THETA_MAX[0];

      xj = (thetaj-B[1])/A[1];
      yj = thetaj/THETA_MAX[1];

      if (islog==0) {
	 //	 tmp = dsqr(dtheta) /thetai/thetai;
	 tmp = dsqr(dtheta)*yi*yj/THETA_MAX[0]/THETA_MAX[1];

      } else {

	 tmp  = dsqr(dlogtheta);
	 tmp *= dsqr(yi*yj);
      }

      tmp *= Tp(xi, a, N, poly, err);
      forwardError(*err, __LINE__, -1);
      tmp *= Tp(xj, a, N, poly, err);
      forwardError(*err, __LINE__, -1);

      tmp *= var_xi[i];
      //      tmp *= dsqr(yi*yj);

      sum += tmp;
  }

   return sum;
}
#undef EPS

/* ============================================================ *
 * Returns the chi^2 for a null test (RR=0).			*
 * ============================================================ */
double chi2_RB_null(const double *RB, const double *covRB, int NRB)
{
   int i, j;
   double c2;

   for (i=0,c2=0.0; i<NRB; i++) {
      for (j=0; j<NRB; j++) {
	 c2 += RB[i] * covRB[i*NRB+j] * RB[j];
      }
   }

   return c2;
}

/* ============================================================ *
 * Reads COSEBIs zeros and normalisation coefficients from file *
 * and returns polynomial coefficients.	File name is automatic, *
 * scale checks are done.					*
 * ============================================================ */
double *read_zeros_norm_cosebi_auto_check(double Psimin, double Psimax, const char *path, error **err)
{
   double *c_cosebi, psimin, psimax;
   char rname[1024];

   sprintf(rname, "%s/cosebi_tplog_rN_%d_%.3f_%.3f",
	   path == NULL ? "." : path, NMAX_COSEBI, Psimin/arcmin, Psimax/arcmin);
   c_cosebi = read_zeros_norm_cosebi(rname, &psimin, &psimax, err);
   forwardError(*err, __LINE__, NULL);

   testErrorRetVA(fabs(psimin-Psimin/arcmin)>EPSILON, mr_file,
		  "Inconsistent file %s, psimin=%g should be %g according to the file name",
		  *err, __LINE__, NULL, rname, psimin, Psimin/arcmin);
   testErrorRetVA(fabs(psimax-Psimax/arcmin)>EPSILON, mr_file,
		  "Inconsistent file %s, psimax=%g should be %g according to the file name",
		  *err, __LINE__, NULL, rname, psimax, Psimax/arcmin);

   return c_cosebi;
}

/* ============================================================ *
 * Reads COSEBIs zeros and normalisation coefficients from file *
 * and sets corresponding min and max scales.			*
 * Calculates polynomial coefficients c and returns them.       *
 * ============================================================ */
double *read_zeros_norm_cosebi(const char *rname, double *psimin, double *psimax, error **err)
{
   FILE *F;
   int Nzeros, Ncoeff, k, n, off_c, off_R, nmax;
   ssize_t nread;
   double *Rn, *Norm, *c;
   char *str, *line=NULL;

   F = fopen_err(rname, "r", err);                  forwardError(*err, __LINE__, NULL);

   /* Read two header lines */
   line = malloc_err(1024*sizeof(char), err);       forwardError(*err, __LINE__, NULL);
   str = fgets(line, 1024, F);
   str = fgets(line, 1024, F);
   free(line);

   nread = fscanf(F, "%d  %lg %lg\n", &nmax, psimin, psimax);
   testErrorRet(nread != 3, mr_file, "File has wrong format.", *err, __LINE__, NULL);

   testErrorRetVA(nmax > NMAX_COSEBI, mr_range,
		  "COSEBI number of modes n=%d read from file %s cannot be larger than NMAX_COSEBI=%d",
		  *err, __LINE__, NULL, nmax, rname, NMAX_COSEBI);

   /* Number of zeros for nmax polynomials,               *
    * sum_{i=1}^{nmax+1} (i+1) = sum_{i=0}^{nmax+1} i - 1 */
   Nzeros = (nmax+1) * (nmax+2) / 2 - 1;

   //fprintf(stderr, "nmax = %d , Psi = [%g, %g]\n", nmax, *psimin, *psimax);
   //fprintf(stderr, "Reading %d Rn, %d Norm, from %s\n", Nzeros, nmax, rname);

   Rn   = malloc_err(sizeof(double)*Nzeros, err);   forwardError(*err, __LINE__, NULL);
   Norm = malloc_err(sizeof(double)*nmax, err);     forwardError(*err, __LINE__, NULL);

   /* Read zeros */
   for (k=0; k<Nzeros; k++) {
      fscanf(F, "%lg ", Rn+k);
   }
   /* Read normalisations */
   for (k=0; k<nmax; k++) {
      fscanf(F, "%lg ", Norm+k);
   }

   fclose(F);

   /* Calculate polynomial coefficients */
   Ncoeff = NMAX_COSEBI * (NMAX_COSEBI + 5) / 2;
   c =  malloc_err(sizeof(double)*Ncoeff, err);        forwardError(*err, __LINE__, NULL);

   /* Polynomial T_i of order i+1 */
   for (n=1,off_R=0; n<=nmax; n++) {

      off_c = n*(n + 3)/2 - 2;   // Offset for coefficients for polynomial #n

      for (k=0; k<n+2; k++) {
	 c[k + off_c] = sum_combinations(n+1-k, n+1, Rn+off_R, err);
	 forwardError(*err, __LINE__, NULL);
	 c[k + off_c] *= Norm[n-1];
	 if (k % 2 == 1) c[k + off_c] = -c[k + off_c];
	 if (n % 2 == 0) c[k + off_c] = -c[k + off_c];
      }

      off_R = (n+1)*(n+2)/2 - 1; // Ncoeff for index i, degree i+1
   }

   free(Rn);
   free(Norm);

   return c;
}

/* ============================================================ *
 * COSEBI logarithmis filter function, SEK10 (28). The          *
 * coefficients have to be precalculated from the zeros, output *
 * by Mathematica.						*
 * Not normalised. 						*
 * ============================================================ */
double Tplog_c(double z, const double *c, int n, error **err)
{
   int i, off_c;
   double res, zpowerm;
   const double *cn;

   off_c = n*(n + 3)/2 - 2; // Offset for coefficients for polynomial #n
   cn    = c + off_c;

   for (i=0,res=0.0,zpowerm=1.0; i<=n+1; i++) {
      res     += cn[i] * zpowerm;
      zpowerm *= z;
   }

   testErrorRet(!finite(res), math_infnan, "Tplog is not finite", *err, __LINE__, 0.0);

   return res;
}

/* ============================================================ *
 * Filter function, SEK10 (38, 39). Normalised.                 *
 * ============================================================ */
double Tmlog(double z, const double *c, int n, error **err)
{
   int m, off_c;
   double res, zpowerm;
   const double *cn;

   off_c = n*(n + 3)/2 - 2; // Offset for coefficients for polynomial #n
   cn    = c + off_c;

   for (m=0,res=0.0,zpowerm=1; m<=n; m++) { // n+1 MKDEBUG ??
      res += dnm(n, m ,cn) * zpowerm;
      zpowerm *= z;
   }

   res += an2(n, cn) * exp(-2.0*z) - an4(n, cn) * exp(-4.0*z);
   testErrorRet(!finite(res), math_infnan, "Tmlog is not finite", *err, __LINE__, 0.0);

   return res;
}

double an2(int n, const double *c)
{
   int j;
   double res, jfac, m2powerjp1;

   for (j=0,res=0.0,jfac=1.0,m2powerjp1=-2.0; j<=n+1; j++) {
      res        += c[j]*jfac / m2powerjp1;
      jfac       *= j + 1.0;      /* factorial(j) */
      m2powerjp1 *= -2.0;        /* (-2)^(j+1)   */
   }

   res *= 4.0;

   return res;
}

double an4(int n, const double *c)
{
   int j;
   double res, jfac, m4powerjp1;

   for (j=0,res=0.0,jfac=1.0,m4powerjp1=-4.0; j<=n+1; j++) {
      res        += c[j]*jfac / m4powerjp1;
      jfac       *= j + 1.0;
      m4powerjp1 *= -4.0;        /* (-4)^(j+1) */
   }

   res *= 12.0;

   return res;
}

/* dnm(n+1, n, c) = 0 */
double dnm(int n, const int m, const double *c)
{
   int j;
   double res, jfac, m2powermjm1, p2powermjm1, mfac;

   p2powermjm1 = 0.5;
   m2powermjm1 = -0.5;
   mfac = gsl_sf_fact(m);
   for (j=m,res=0.0,jfac=mfac; j<=n+1; j++) {
      res         += c[j] * jfac * m2powermjm1 * (3.0 * p2powermjm1 - 1.0);
      jfac        *= j + 1.0;        /* j! */
      m2powermjm1 /= -2.0;    /* (-2)^(m-j-1) */
      p2powermjm1 /= 2.0;     /* 2^(m-k-1)    */
   }

   res *= 4.0/mfac;
   res += c[m];

   return res;
}

double sum_combinations(int j, int n, const double *r, error **err)
{
   gsl_combination *c;
   size_t i, c_i;
   double sum, prod;

   if (n==0) return 1.0;

   c = gsl_combination_calloc(n, j);

   sum = 0.0;
   do {
      for (i=0,prod=1.0; i<j; i++) {
	 c_i =  gsl_combination_get(c, i);
	 prod *= r[c_i];
	 //printf(" %zu", c_i);
      }
      sum += prod;
      //printf("{"); gsl_combination_fprintf(stdout, c, " %u"); printf (" }\n");
   } while (gsl_combination_next(c) == GSL_SUCCESS);

   gsl_combination_free(c);

   return sum;
}

/* ============================================================ *
 * E-/B-mode correlation functions from COSEBIs, SEK10 (40).    *
 * The results are stored in xi_pm_EB = (E+, E-, B+, B-).       *
 * ============================================================ */
void xipmEB(double theta, double THETA_MIN, double THETA_MAX, const double *c,
	    const double *E, const double *B, int N, double xi_pm_EB[4], error **err)
{
   double Tp, Tm, z;
   int n;

   for (n=0; n<4; n++) {
      xi_pm_EB[n] = 0.0;
   }

   for (n=1; n<=N; n++) {
      z  = log(theta/THETA_MIN);
      Tp = Tplog_c(z, c, n, err);
      Tm = Tmlog(z, c, n, err);

      xi_pm_EB[0] += E[n] * Tp;   // E+
      xi_pm_EB[1] += E[n] * Tm;   // E-
      xi_pm_EB[2] += B[n] * Tp;   // B+
      xi_pm_EB[3] += B[n] * Tm;   // B-
   }

   for (n=0; n<4; n++) {
      xi_pm_EB[n] *= 2.0 / theta / (THETA_MAX - THETA_MIN);
   }
}


#include <math.h>
#include <float.h>
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <stdbool.h>

#include <cblas.h>

#include "funcsfa.h"

#define min(x, y) ((x)>(y)?(y):(x))
#define max(x, y) ((x)>(y)?(x):(y))

#define DEBUG 0
#define debug_print(fmt, ...) \
        do { if (DEBUG) fprintf(stderr, "%s:%4d:%25s(): " fmt, __FILE__, \
                                __LINE__, __func__, __VA_ARGS__); } while (0)

double eps_ev = 1e-20;

// **************************************
// Lapack
// **************************************

// Lapack Fortran DGESV interface
void
dgesv_(int *n, int *nrhs, double *a, int *lda, int *ipiv, double *b, int *ldb,
       int *info);
// DGESV wrapper, for solving linear equations.
static int
dgesv(int n, int nrhs, double *a, int lda, int *ipiv, double *b, int ldb)
{
  int info = -1;
  dgesv_(&n, &nrhs, a, &lda, ipiv, b, &ldb, &info);
  return info;
}

// Lapack Fortran DGESVD interface
void
dgesvd_(char *jobu, char *jobvt, int *m, int *n, double *a, int *lda,
        double *s, double *u, int *ldu, double *vt, int *ldvt, double *work,
        int *lwork, int *info);
// DGESVD wrapper, for Singular Value Decomposition.
static int
dgesvd(char jobu, char jobvt, int m, int n, double *a, int lda, double *s,
       double *u, int ldu, double *vt, int ldvt)
{
  int info = -1;
  int lwork = -1;
  double optimal_lwork;
  double *work;
  // Query optimal work size
  dgesvd_(&jobu, &jobvt, &m, &n, a, &lda,
      s, u, &ldu, vt, &ldvt, &optimal_lwork,
      &lwork, &info);
  if (info != 0) {
    return info;
  }
  lwork = (int) optimal_lwork;

  work = malloc(sizeof(double)*lwork);

  dgesvd_(&jobu, &jobvt, &m, &n, a, &lda,
      s, u, &ldu, vt, &ldvt, work,
      &lwork, &info);

  free(work);
  return info;
}

// **************************************
// Objects
// **************************************

funcsfa_Factorization*
funcsfa_Factorization_alloc(size_t n_features, size_t n_factors,
                            size_t n_samples, size_t n_datatypes,
                            size_t *n_features_split, double *data)
{
  size_t i;
  funcsfa_Factorization *m;
  size_t sum_n_features_split = 0;

  debug_print("%s\n", "Allocating SFA object");

  m = malloc(sizeof(funcsfa_Factorization));
  m->n_features = n_features;
  m->n_factors = n_factors;
  m->n_samples = n_samples;
  m->n_datatypes = n_datatypes;
  m->n_features_split = malloc(sizeof(size_t)*n_datatypes);
  for(i=0; i<n_datatypes; i++)
  {
    m->n_features_split[i] = n_features_split[i];
    sum_n_features_split += n_features_split[i];
  }
  if (sum_n_features_split != n_features)
  {
    debug_print("Total number of features (%zd) does not match sum (%zd)\n",
        n_features, sum_n_features_split);
    free(m->n_features_split);
    free(m);
    return NULL;
  }
  m->data = data;
  m->coefficients = malloc(sizeof(double)*n_features*n_factors);
  m->factors = malloc(sizeof(double)*n_factors*n_samples);
  m->factor_cov = malloc(sizeof(double)*n_factors*n_factors);
  m->residual_var = malloc(sizeof(double)*n_features);

  return(m);
}

void
funcsfa_Factorization_free(funcsfa_Factorization *m)
{
  debug_print("%s\n", "Freeing SFA object");

  if (m != NULL)
  {
    free(m->n_features_split);
    free(m->coefficients);
    free(m->factors);
    free(m->factor_cov);
    free(m->residual_var);
  }
  free(m);
}

funcsfa_Monitor*
funcsfa_Monitor_alloc(int max_iter)
{
  funcsfa_Monitor *mon;
  int i;

  debug_print("%s\n", "Allocating Monitor object");

  mon = malloc(sizeof(funcsfa_Monitor));
  if (mon == NULL) {
    return NULL;
  }
  mon->n_iter = 0;
  mon->reconstruction_error = malloc(sizeof(double)*(max_iter+1));
  mon->max_diff_factors = malloc(sizeof(double)*(max_iter+1));
  mon->max_diff_coefficients = malloc(sizeof(double)*(max_iter+1));
  if ((mon->reconstruction_error == NULL) ||
      (mon->max_diff_factors == NULL) ||
      (mon->max_diff_coefficients == NULL))
  {
    funcsfa_Monitor_free(mon);
    return NULL;
  }
  for (i = 0; i < max_iter+1; i++)
  {
    mon->reconstruction_error[i] = NAN;
    mon->max_diff_factors[i] = NAN;
    mon->max_diff_coefficients[i] = NAN;
  }

  return mon;
}

void
funcsfa_Monitor_free(funcsfa_Monitor *mon)
{
  debug_print("%s\n", "Freeing Monitor object");

  if (mon != NULL)
  {
    free(mon->reconstruction_error);
    free(mon->max_diff_factors);
    free(mon->max_diff_coefficients);
  }
  free(mon);
}

/** Bunch of memory for temporary matrices.
 * @private
 */
typedef struct
{
  /** Integer vector n_factors long */
  int *factors_int;
  /** Double vector n_features long */
  double *features_double;
  /** Double matrix n_features x n_factors 1 */
  double *features_factors_double_a;
  /** Double matrix n_features x n_factors 2 */
  double *features_factors_double_b;
  /** Double matrix n_factors x n_factors 1 */
  double *factors_factors_double_a;
  /** Double matrix n_factors x n_factors 2 */
  double *factors_factors_double_b;

  /** Precomputed value of diag(data * t(data)) */
  double *diag_data_square;
  /** Precomputed value of sum(diag(data * t(data))) per data type */
  double *dt_var;
  /** Factors of previous iteration */
  double *prev_factors;
  /** Coefficients of previous iteration */
  double *prev_coefficients;
} Workspace;

static Workspace*
ws_alloc(funcsfa_Factorization *m)
{
  Workspace *ws;

  debug_print("%s\n", "Allocating Workspace");

  ws = malloc(sizeof(Workspace));
  ws->factors_int = malloc(sizeof(int)*m->n_factors);
  ws->features_double = malloc(sizeof(double)*m->n_features);

  ws->features_factors_double_a = malloc(sizeof(double)*m->n_features*m->n_factors);
  ws->features_factors_double_b = malloc(sizeof(double)*m->n_features*m->n_factors);
  ws->factors_factors_double_a = malloc(sizeof(double)*m->n_factors*m->n_factors);
  ws->factors_factors_double_b = malloc(sizeof(double)*m->n_factors*m->n_factors);

  ws->diag_data_square = malloc(sizeof(double)*m->n_features);
  ws->dt_var = malloc(sizeof(double)*m->n_factors);
  ws->prev_factors = malloc(sizeof(double)*m->n_factors*m->n_samples);
  ws->prev_coefficients = malloc(sizeof(double)*m->n_features*m->n_factors);

  return ws;
}

static void
ws_init(Workspace *ws, funcsfa_Factorization *m)
{
  double *d_d;
  size_t dt_i, feature_start, n_features, i;

  debug_print("%s\n", "Initializing Workspace");

  d_d = malloc(sizeof(double)*m->n_features*m->n_features);

  cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans,
      m->n_features, m->n_features, m->n_samples, 1,
      m->data, m->n_samples,
      m->data, m->n_samples,
      0, d_d, m->n_features);

  for (i = 0; i < m->n_features; i++)
  {
    ws->diag_data_square[i] = d_d[i*m->n_features + i];
  }
  feature_start = 0;
  for(dt_i = 0; dt_i < m->n_datatypes; dt_i++)
  {
    ws->dt_var[dt_i] = 0;
    n_features = m->n_features_split[dt_i];
    for(i = feature_start; i < (feature_start + n_features); i++)
    {
      ws->dt_var[dt_i] += ws->diag_data_square[i];
    }
    ws->dt_var[dt_i] /= n_features;
    feature_start = feature_start + n_features;
  }

  free(d_d);
}

static void
ws_free(Workspace *ws)
{
  debug_print("%s\n", "Freeing Workspace");

  free(ws->factors_int);
  free(ws->features_double);
  free(ws->features_factors_double_a);
  free(ws->features_factors_double_b);
  free(ws->factors_factors_double_a);
  free(ws->factors_factors_double_b);

  free(ws->diag_data_square);
  free(ws->dt_var);
  free(ws->prev_factors);
  free(ws->prev_coefficients);

  free(ws);
}

// **************************************
// Utility Functions
// **************************************

static double
max_diff(double *a, double *b, size_t length)
{
  size_t i;
  double md, d;

  md = 0;
  for (i = 0; i < length; i++)
  {
    d = fabs(a[i] - b[i]);
    if (d > md)
    {
      md = d;
    }
  }
  return md;
}

static bool
finite_array(double *a, size_t length)
{
  for (size_t i=0; i<length; i++)
  {
    if (!isfinite(a[i]))
    {
      return false;
    }
  }
  return true;
}

static bool
positive_array(double *a, size_t length)
{
  for (size_t i=0; i<length; i++)
  {
    if (!(a[i] > 0.0))
    {
      return false;
    }
  }
  return true;
}

#define check_finite(a, length) \
    if (DEBUG && !finite_array(a, length)) {\
      debug_print("%s\n", "Error, non-finite values in matrix");}

#define check_positive(a, length) \
    if (DEBUG && !positive_array(a, length)) {\
      debug_print("%s\n", "Error, non-positive values in matrix");}

// **************************************
// Algorithm
// **************************************

static int
expectation_factors(funcsfa_Factorization *m, Workspace *ws)
{
  size_t i, j;
  double *r_c, *c_r_c, *c_r_c_i, *o, *c_o;
  double factor_mean, factor_var, factor_sd, deviance;
  int *ipiv;
  int info;

  debug_print("%s\n", "Computing Expectation and Covariance of Factors");
  check_finite(m->coefficients, m->n_features*m->n_factors);
  check_finite(m->residual_var, m->n_features);
  check_positive(m->residual_var, m->n_features);

  /* Compute $\Psi^{-1} B$
   * r_c = diag(1/residual_var) * coefficients */
  r_c = ws->features_factors_double_a;
  for (i=0; i < m->n_features; i++)
  {
    for (j=0; j < m->n_factors; j++)
    {
      r_c[j * m->n_features + i] =
          m->coefficients[j * m->n_features + i] / m->residual_var[i];
    }
  }

  /* Compute $B^T \Psi^{-1} B + I$
   * c_r_c = t(coefficients) * r_c + I */
  c_r_c = ws->factors_factors_double_a;
  cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans,
      m->n_factors, m->n_factors, m->n_features, 1,
      m->coefficients, m->n_features,
      r_c, m->n_features,
      0, c_r_c, m->n_factors);
  for (i = 0; i < m->n_factors; i++)
  {
    c_r_c[i*m->n_factors + i] += 1;
  }

  /* Compute $O = \Psi^{-1}B(B^T \Psi^{-1} B + I)^{-1}$
   * o = r_c*solve(c_r_c) */
  c_r_c_i = ws->factors_factors_double_b;
  ipiv = ws->factors_int;
  for (i = 0; i < m->n_factors; i++)
  {
    for (j = 0; j < m->n_factors; j++)
    {
      c_r_c_i[j * m->n_factors + i] = (i == j);
    }
  }

  info = dgesv(m->n_factors, m->n_factors,
      c_r_c, m->n_factors, ipiv,
      c_r_c_i, m->n_factors);
  if (info != 0) {
    debug_print("Lapack error: %d\n", info);
    return -1;
  }
  c_r_c = NULL;
  o = ws->features_factors_double_b;
  cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans,
     m->n_features, m->n_factors, m->n_factors, 1,
     r_c, m->n_features,
     c_r_c_i, m->n_factors,
     0, o, m->n_features);

  /* Compute $E[Z|X] = (\Psi^{-1}B(B^T\Psi^{-1}B+I)^{-1})^TX$
   * factors = t(o) * t(X)
   */
  cblas_dgemm(CblasColMajor, CblasTrans, CblasTrans,
      m->n_factors, m->n_samples, m->n_features, 1,
      o, m->n_features,
      m->data, m->n_samples,
      0, m->factors, m->n_factors);

  /* Compute $B^T\Psi^{-1}B(B^T\Psi^{-1}B+I)^{-1}$
   * c_o = t(coefficients) * o
   */
  c_o = ws->factors_factors_double_a;
  cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans,
      m->n_factors, m->n_factors, m->n_features, 1,
      m->coefficients, m->n_features,
      o, m->n_features,
      0, c_o, m->n_factors);

  /*
   * Compute $E[ZZ^T|X] = I - (B^T\Psi^{-1}B(B^T\Psi^{-1}B+I)^{-1})^T + E[Z|X]E[Z|X]^T$
   * factor_cov = factors * t(factors) + I - c_o
   */
  cblas_dgemm(CblasColMajor, CblasNoTrans, CblasTrans,
      m->n_factors, m->n_factors, m->n_samples, 1,
      m->factors, m->n_factors,
      m->factors, m->n_factors,
      0, m->factor_cov, m->n_factors);

  for (i = 0; i < m->n_factors; i++)
  {
    for (j = 0; j < m->n_factors; j++)
    {
      m->factor_cov[j * m->n_factors + i] += (i==j) - c_o[j * m->n_factors + i];
    }
  }

  debug_print("Factor value: %g\n", m->factors[0]);

  // Rescale factors
  for (i = 0; i < m->n_factors; i++)
  {
    factor_mean = 0;
    for (j = 0; j < m->n_samples; j++)
    {
      factor_mean += m->factors[i + j*m->n_factors];
    }
    factor_mean /= m->n_samples;
    debug_print("  Factor mean %zd: %g\n", i, factor_mean);
    factor_var = 0;
    for (j = 0; j < m->n_samples; j++)
    {
      deviance = m->factors[i + j*m->n_factors] - factor_mean;
      factor_var += deviance*deviance;
    }
    factor_var /= m->n_samples;
    debug_print("  Factor variation %zd: %g\n", i, factor_var);
    for (j = 0; j < m->n_factors; j++)
    {
      m->factor_cov[i + j*m->n_factors] /= factor_var;
    }
    factor_sd = sqrt(factor_var);
    for (j = 0; j < m->n_samples; j++)
    {
      m->factors[i + j * m->n_factors] /= factor_sd;
    }
  }

  return 0;
}

static int
maximize(funcsfa_Factorization *m, double *lambdas, Workspace *ws)
{
  const size_t n_lambda = 2;
  size_t dt_i, f_i, lambda_i, k;
  double *ZXt, *BZZt;
  double unpenalized_coeff, l1, l2, dt_var;
  size_t n_features, feature_start;

  debug_print("%s\n", "Estimating Coefficients");

  check_finite(m->factor_cov, m->n_factors*m->n_factors);
  check_finite(m->factors, m->n_factors*m->n_samples);

  ZXt = ws->features_factors_double_a;
  BZZt = ws->features_double;
  // Compute E[Z|X]X^T
  cblas_dgemm(CblasColMajor, CblasNoTrans, CblasNoTrans,
      m->n_factors, m->n_features, m->n_samples, 1,
      m->factors, m->n_factors,
      m->data, m->n_samples,
      0, ZXt, m->n_factors);

  // Coordinate Descent over factors
  for(k = 0; k < m->n_factors; k++)
  {
    debug_print("\tCoordinate descent for factor: %zd\n", k);
    // Compute $BE[ZZ^T|X]_(\cdot k)$
    cblas_dgemv(CblasColMajor, CblasNoTrans,
      m->n_features, m->n_factors, 1.0,
      m->coefficients, m->n_features,
      m->factor_cov+(k*m->n_factors), 1,
      0.0, BZZt, 1);

    lambda_i = 0;
    feature_start = 0;
    for(dt_i = 0; dt_i < m->n_datatypes; dt_i++)
    {
      n_features = m->n_features_split[dt_i];
      l1 = lambdas[lambda_i];
      l2 = lambdas[lambda_i+1];
      for(f_i = feature_start; f_i < (feature_start + n_features); f_i++)
      {
        // Compute: $$
        //  \hat B_{(ij)} \leftarrow \frac
        //    {\operatorname{S}(
        //      \hat B_{(ij)} + E[Z|X]_{(j \cdot)} X_{(i \cdot)}^T -
        //      \hat B_{(i\cdot)} \operatorname{E}[ZZ^T|X]_{(\cdot j)},
        //      l_1)}
        //    {1+l_2}
        // $$
        unpenalized_coeff = m->coefficients[f_i + (k * m->n_features)] +
                            ((ZXt[k + (f_i * m->n_factors)] - BZZt[f_i])
                             / m->n_samples);
        m->coefficients[f_i + (k * m->n_features)] =
          copysign(fmax(0.0, fabs(unpenalized_coeff) - l1),
                   unpenalized_coeff) /
          (1 + l2);
      }
      feature_start = feature_start + n_features;
      lambda_i = lambda_i + n_lambda;
    }
  }

  debug_print("%s\n", "Estimating Residual Variance");

  // Residual variance per data type
  feature_start = 0;
  for(dt_i = 0; dt_i < m->n_datatypes; dt_i++)
  {
    dt_var = 0;
    n_features = m->n_features_split[dt_i];
    for(f_i = feature_start; f_i < (feature_start + n_features); f_i++)
    {
      for(k = 0; k < m->n_factors; k++)
      {
        dt_var += (ZXt[k + (f_i * m->n_factors)] *
                   m->coefficients[f_i + (k * m->n_features)]);
      }
    }
    dt_var = (ws->dt_var[dt_i] - (dt_var / n_features)) / m->n_samples;
    for(f_i = feature_start; f_i < (feature_start + n_features); f_i++)
    {
      if (dt_var > eps_ev) {
        m->residual_var[f_i] = dt_var;
      } else {
        m->residual_var[f_i] = eps_ev;
      }
    }
    debug_print("\tResidual Variance %zd: %g\n", dt_i, dt_var);
    feature_start = feature_start + n_features;
  }
  return 0;
}

static int
init_svd(funcsfa_Factorization *m, Workspace *ws)
{
  int info;
  size_t i, j;
  double *x, *s, *vt, *u;
  double scale, f, explained_var, data_var;
  size_t total_n_factors;

  debug_print("%s\n", "Initializing with SVD");

  total_n_factors = min(m->n_features, m->n_samples);

  x = malloc(sizeof(double) * m->n_features * m->n_samples);
  s = malloc(sizeof(double) * total_n_factors);
  vt = malloc(sizeof(double) * total_n_factors * m->n_features);
  u = malloc(sizeof(double) * total_n_factors * m->n_samples);
  if (x == NULL || s == NULL || vt == NULL)
  {
    free(x);
    free(vt);
    free(s);
    free(u);
    return -1;
  }

  for (i=0; i < m->n_features * m->n_samples; i++)
  {
    x[i] = m->data[i];
  }

  /* vt = svd(X).V.t */
  debug_print("%s\n", "\tStarting Lapack SVD");
  info = dgesvd('S', 'S',
      m->n_samples, m->n_features,
      x, m->n_samples,
      s, u, m->n_samples,
      vt, total_n_factors);
  if (info != 0) {
    debug_print("Lapack error: %d\n", info);
    free(x);
    free(vt);
    free(u);
    free(s);
    return -1;
  }
  debug_print("%s\n", "\tDone Lapack SVD");

  scale = 0;
  for (i=0; i < m->n_factors; i++)
  {
    for (j=0; j < m->n_samples; j++)
    {
      f = u[j + i*m->n_samples] * s[i];
      m->factors[i + j*m->n_factors] = f;
      scale += f*f;
    }
  }
  scale /= m->n_factors * m->n_samples;
  scale = sqrt(scale);
  for (i=0; i < m->n_factors; i++)
  {
    for (j=0; j < m->n_features; j++)
    {
      m->coefficients[j + i*m->n_features] = vt[i + j*total_n_factors] * scale;
    }
  }
  for (i=0; i < m->n_factors; i++)
  {
    for (j=0; j < m->n_samples; j++)
    {
      m->factors[i + j*m->n_factors] /= scale;
    }
  }

  explained_var = 0;
  for (i=0; i < m->n_factors; i++)
  {
    explained_var += s[i]*s[i];
  }
  explained_var /= m->n_samples * m->n_features;
  data_var = 0;
  for (i=0; i < m->n_features; i++)
  {
    data_var += ws->diag_data_square[i];
  }
  data_var /= m->n_samples * m->n_features;
  debug_print("\tData Var: %g\n", data_var);
  explained_var = data_var - explained_var;
  explained_var = fmax(0, explained_var) + eps_ev;
  for (i = 0; i < m->n_features; i++) {
    m->residual_var[i] = explained_var;
  }
  debug_print("\tResidual Variance: %g\n", explained_var);

  free(x);
  free(vt);
  free(u);
  free(s);
  return 0;
}

double
calc_reconstruction_error(funcsfa_Factorization *m, Workspace *ws)
{
  int i, j;
  double error, e;
  double *prediction;

  error = 0.0;
  for (i = 0; i < m->n_samples; i++)
  {
    prediction = ws->features_double;
    cblas_dgemv(CblasColMajor, CblasNoTrans,
        m->n_features, m->n_factors,
        1.0, m->coefficients, m->n_features,
        m->factors + (i*m->n_factors), 1,
        0.0, prediction, 1);
    for (j = 0; j < m->n_features; j++)
    {
      e = m->data[i + (j * m->n_samples)] - prediction[j];
      error += e*e;
    }
  }
  return error;
}

int
funcsfa(funcsfa_Factorization *m, double eps, int max_iter,
        char regularizations[], double *lambdas, funcsfa_Monitor *mon)
{
  int iter = 0;
  int err = 0;
  size_t i;
  double diff = DBL_MAX;
  double diff_f, diff_c;
  double rec_error;
  int const LEN_PREV_REC_ERROR = 10;
  double prev_rec_error[LEN_PREV_REC_ERROR];
  int prev_rec_error_i=0;

  Workspace *ws;

  for (i=0; i<m->n_datatypes; i++)
  {
    if (regularizations[i] != 'e')
    {
      debug_print("%s\n", "Only elastic net regularization is supported");
      return -4;
    }
  }

  ws = ws_alloc(m);
  ws_init(ws, m);

  debug_print("%s\n", "Starting Factorization");

  err = init_svd(m, ws);
  if (err != 0)
  {
    return(1);
  }

  rec_error = calc_reconstruction_error(m, ws);
  if (mon != NULL)
  {
    mon->reconstruction_error[0] = rec_error;
    mon->max_diff_factors[0] = 0;
    mon->max_diff_coefficients[0] = 0;
  }
  for (i=0; i<LEN_PREV_REC_ERROR; i++) {
    prev_rec_error[i] = 0.0;
  }

  /* EM Loop */
  while (iter < max_iter && diff > eps) {
    debug_print("Iteration %d\n", iter);
    memcpy(ws->prev_factors, m->factors,
        sizeof(double)*m->n_factors*m->n_samples);
    memcpy(ws->prev_coefficients, m->coefficients,
        sizeof(double)*m->n_features*m->n_factors);
    err = expectation_factors(m, ws);
    if (err != 0)
    {
      debug_print("Error in computing expectation of factors: %d\n", err);
      ws_free(ws);
      return(2);
    }
    err = maximize(m, lambdas, ws);
    if (err != 0)
    {
      debug_print("Error in maximization: %d\n", err);
      ws_free(ws);
      return(3);
    }

    diff_f = max_diff(m->factors,
        ws->prev_factors, m->n_factors*m->n_samples);
    diff_c = max_diff(m->coefficients, ws->prev_coefficients,
        m->n_features*m->n_factors);
    debug_print("Diff (factor/coefficients): %g; %g\n", diff_f, diff_c);
    prev_rec_error[prev_rec_error_i] = rec_error;
    prev_rec_error_i += 1;
    if (prev_rec_error_i > LEN_PREV_REC_ERROR) {
      prev_rec_error_i = 0;
    }
    rec_error = calc_reconstruction_error(m, ws);
    debug_print("Reconstruction Error: %g\n", rec_error);
    diff = 0.0;
    for (i=0; i<LEN_PREV_REC_ERROR; i++) {
      diff += fabs(prev_rec_error[i] - rec_error);
    }
    diff /= LEN_PREV_REC_ERROR;
    debug_print("Diff (Reconstruction Error): %g\n", diff);

    if (mon != NULL)
    {
      mon->reconstruction_error[iter+1] = rec_error;
      mon->max_diff_factors[iter+1] = diff_f;
      mon->max_diff_coefficients[iter+1] = diff_c;
      mon->n_iter = iter;
    }
    iter = iter + 1;
  }
  ws_free(ws);

  return 0;
}


/*
** Implementation of LISA algorithm
** for statistical inference of fMRI images
**
** 2nd level inference (one-sample test)
**
** G.Lohmann, 2017
*/
#include <viaio/Vlib.h>
#include <viaio/file.h>
#include <viaio/mu.h>
#include <viaio/option.h>
#include <viaio/os.h>
#include <viaio/VImage.h>

#include <gsl/gsl_cdf.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_histogram.h>

#include <math.h>
#include <stdio.h>
#include <string.h>

#ifdef _OPENMP
#include <omp.h>
#endif /*_OPENMP*/

#define ABS(x) ((x) > 0 ? (x) : -(x))

extern void   VIsolatedVoxels(VImage src,float threshold);
extern void   VHistogram(gsl_histogram *histogram,VString filename);
extern void   VCheckImage(VImage src);
extern void   FDR(VImage src,VImage dest,gsl_histogram *nullhist,gsl_histogram *realhist,double);
extern double ttest1(double *data1,int n);
extern void   ImageStats(VImage src,double *,double *,double *hmin,double *hmax);
extern void   VBilateralFilter(VImage src,VImage dest,int radius,double var1,double var2,int);
extern double VImageVar(VImage src);
extern void   VImageCount(VImage src);
extern void   VGetHistRange(VImage src,double *hmin,double *hmax);
extern void   VZScale(VImage src,float mode,float stddev);
extern float  VGetMode(VImage src);
extern double t2z(double,double);
extern void   HistoUpdate(VImage,gsl_histogram *);

/* generate permutation table */
int **genperm(gsl_rng *rx,int n,int numperm)
{
  int i,j;
  int **table = (int **) VCalloc(numperm,sizeof(int *));
  for (i = 0; i < numperm; i++) {
    table[i] = (int *) VCalloc(n,sizeof(int));
    for (j=0; j<n; j++) {
      table[i][j] = 0;
      if (gsl_ran_bernoulli (rx,(double)0.5) == 1) table[i][j] = 1;
    }
  }
  return table;
}

/* onesample t-test */
void OnesampleTest(VImage *src1,VImage dest,int *permtable,int n)
{
  int i,k,b,r,c,nslices,nrows,ncols;
  double nx,ave,var,u,t,z,df;
  double tiny=1.0e-8;
  extern void avevar(double *data,int n,double *a,double *v);

  nslices = VImageNBands(src1[0]);
  nrows   = VImageNRows(src1[0]);
  ncols   = VImageNColumns(src1[0]);
  VFillImage(dest,VAllBands,0);
  double *data = (double *) VCalloc(n,sizeof(double));

  for (b=0; b<nslices; b++) {
    for (r=0; r<nrows; r++) {
      for (c=0; c<ncols; c++) {
	k = 0;
	for (i=0; i<n; i++) {
	  u = (double)VPixel(src1[i],b,r,c,VFloat);
	  if (ABS(u) > tiny) {
	    if (permtable[i] > 0) u = -u;
	    data[k] = u;
	    k++;
	  }
	}
	if (k < n-2) continue;
	avevar(data,k,&ave,&var);
	if (var < tiny) continue;
	nx   = (double)k;
	t    = sqrt(nx) * ave/sqrt(var);
	df   = nx - 1.0;
	z    = t2z(t,df);
	if (t < 0) z = -z;
	VPixel(dest,b,r,c,VFloat) = z;
      }
    }
  }
  VFree(data);
}

int main (int argc, char *argv[])
{
  static VArgVector in_files1;
  static VString  out_filename="";
  static VString  mask_filename="";
  static VFloat   alpha = 0.05;
  static VShort   radius = 2;
  static VFloat   rvar = 2.0;
  static VFloat   svar = 2.0;
  static VShort   numiter = 2;
  static VShort   numperm = 5000;
  static VLong    seed = 99402622;
  static VBoolean centering = FALSE;
  static VBoolean cleanup = TRUE;
  static VShort   nproc = 0;
  static VOptionDescRec options[] = {
    {"in", VStringRepn, 0, & in_files1, VRequiredOpt, NULL,"Input files" },
    {"out", VStringRepn, 1, & out_filename, VRequiredOpt, NULL,"Output file" },
    {"alpha",VFloatRepn,1,(VPointer) &alpha,VOptionalOpt,NULL,"FDR significance level"},
    {"perm",VShortRepn,1,(VPointer) &numperm,VOptionalOpt,NULL,"Number of permutations"},
    {"mask", VStringRepn, 1, (VPointer) &mask_filename, VRequiredOpt, NULL, "Mask"},
    {"seed",VLongRepn,1,(VPointer) &seed,VOptionalOpt,NULL,"Seed for random number generation"},
    {"radius",VShortRepn,1,(VPointer) &radius,VOptionalOpt,NULL,"Bilateral parameter (radius in voxels)"},
    {"rvar",VFloatRepn,1,(VPointer) &rvar,VOptionalOpt,NULL,"Bilateral parameter (radiometric)"},
    {"svar",VFloatRepn,1,(VPointer) &svar,VOptionalOpt,NULL,"Bilateral parameter (spatial)"},
    {"filteriterations",VShortRepn,1,(VPointer) &numiter,VOptionalOpt,NULL,"Bilateral parameter (number of iterations)"},
    {"cleanup",VBooleanRepn,1,(VPointer) &cleanup,VOptionalOpt,NULL,"Whether to remove isloated voxels"},
    {"j",VShortRepn,1,(VPointer) &nproc,VOptionalOpt,NULL,"Number of processors to use, '0' to use all"},
  };

  FILE *fp=NULL;
  VString in_filename;
  VAttrList list1=NULL,out_list=NULL,geolist=NULL;
  VImage *src1;
  int i,nimages,npix=0;
  char *prg_name=GetLipsiaName("vlisa_onesample");
  fprintf (stderr, "%s\n", prg_name);

  /* parse command line */
  if (! VParseCommand (VNumber (options), options, & argc, argv)) {
    VReportUsage (argv[0], VNumber (options), options, NULL);
    exit (EXIT_FAILURE);
  }
  if (argc > 1) {
    VReportBadArgs (argc, argv);
    exit (EXIT_FAILURE);
  }

  /* omp-stuff */
#ifdef _OPENMP
  int num_procs=omp_get_num_procs();
  if (nproc > 0 && nproc < num_procs) num_procs = nproc;
  fprintf(stderr," using %d cores\n",(int)num_procs);
  omp_set_num_threads(num_procs);
#endif /* _OPENMP */

  /* images  */
  VImage mask = VReadImageFile(mask_filename);
  if (mask==NULL) VError("Error reading mask file %s",mask_filename);

  nimages = in_files1.number;
  src1 = (VImage *) VCalloc(nimages,sizeof(VImage));
  for (i = 0; i < nimages; i++) {
    in_filename = ((VString *) in_files1.vector)[i];
    list1   = VReadAttrList(in_filename,0L,TRUE,FALSE);
    VMaskMinval(list1,mask,0.0);

    src1[i] = VReadImage(list1);
    if (src1[i] == NULL) VError(" no input image found");
    if (VPixelRepn(src1[i]) != VFloatRepn) VError(" input pixel repn must be float");
    if (i == 0) npix = VImageNPixels(src1[i]);
    else if (npix != VImageNPixels(src1[i])) VError(" inconsistent image dimensions");

    /* use geometry info from 1st file */
    if (geolist == NULL) {
      geolist = VGetGeoInfo(list1);
      double *D = VGetGeoPixdim(geolist,NULL);
      if (fabs(D[0]-4.0) < TINY) VError(" The input must be a list of separate 3D images, not one 4D file");
    }
  }

  /* ini random permutations */
  size_t n = nimages;
  gsl_rng_env_setup();
  const gsl_rng_type *T = gsl_rng_default;
  gsl_rng *rx = gsl_rng_alloc(T);
  gsl_rng_set(rx,(unsigned long int)seed);
  int nperm=0;
  int **permtable = genperm(rx,(int)n,(int)numperm);

  /* estimate null variance based on first 30 permutations */
  double hmin=0,hmax=0;
  float stddev=1.0;
  if (numperm > 0) {
    int tstperm = 30;
    if (tstperm > numperm) tstperm = numperm;
    double varsum=0,nx=0;

#pragma omp parallel for shared(src1,permtable) schedule(dynamic)
    for (nperm = 0; nperm < tstperm; nperm++) {
      VImage zmap = VCreateImageLike(src1[0]);
      OnesampleTest(src1,zmap,permtable[nperm],nimages);
#pragma omp critical
      {
	varsum += VImageVar(zmap);
	nx++;
      }
      VDestroyImage(zmap);
    }
    double meanvar = varsum/nx;
    stddev = sqrt(meanvar);
  }

  /* no permutation */
  int *nopermtable = (int *) VCalloc(n,sizeof(int));
  VImage dst1  = VCreateImageLike (src1[0]);
  VImage zmap1 = VCreateImageLike(src1[0]);
  OnesampleTest(src1,zmap1,nopermtable,nimages);

  if (numperm == 0) {
    double z = VImageVar(zmap1);
    stddev = (float)(sqrt(z)); /* update stddev */
  }
  float mode=0;
  if (centering) mode = VGetMode(zmap1);
  if (numperm > 0) VZScale(zmap1,mode,stddev);
  VBilateralFilter(zmap1,dst1,(int)radius,(double)(rvar),(double)svar,(int)numiter);

  /* ini histograms */
  VGetHistRange(dst1,&hmin,&hmax);
  size_t nbins = 20000;
  gsl_histogram *hist0 = gsl_histogram_alloc (nbins);
  gsl_histogram_set_ranges_uniform (hist0,hmin,hmax);
  gsl_histogram *histz = gsl_histogram_alloc (nbins);
  gsl_histogram_set_ranges_uniform (histz,hmin,hmax);
  HistoUpdate(dst1,histz);

  /* random permutations */
#pragma omp parallel for shared(src1,permtable) schedule(dynamic)
  for (nperm = 0; nperm < numperm; nperm++) {
    if (nperm%20 == 0) fprintf(stderr," perm  %4d  of  %d\r",nperm,(int)numperm);

    VImage zmap = VCreateImageLike(src1[0]);
    VImage dst  = VCreateImageLike (zmap);
    OnesampleTest(src1,zmap,permtable[nperm],nimages);
    float mode=0;
    if (centering) mode = VGetMode(zmap);
    VZScale(zmap,mode,stddev);
    VBilateralFilter(zmap,dst,(int)radius,(double)(rvar),(double)svar,(int)numiter);

#pragma omp critical
    {
      HistoUpdate(dst,hist0);
    }
    VDestroyImage(dst);
    VDestroyImage(zmap);
  }

  /* apply fdr */
  VImage fdrimage = VCopyImage (dst1,NULL,VAllBands);
  if (numperm > 0) {
    FDR(dst1,fdrimage,hist0,histz,(double)alpha);

    if (cleanup && alpha < 1.0) {
      VIsolatedVoxels(fdrimage,(float)(1.0-alpha));
    }
    if (alpha < 1.0) VImageCount(fdrimage);
  }

  /* write output to disk */
  out_list = VCreateAttrList ();
  VHistory(VNumber(options),options,prg_name,&list1,&out_list);
  VSetGeoInfo(geolist,out_list);
  VAppendAttr (out_list,"image",NULL,VImageRepn,fdrimage);
  fp = VOpenOutputFile (out_filename, TRUE);
  if (! VWriteFile (fp, out_list)) exit (1);
  fclose(fp);
  fprintf (stderr, "\n%s: done.\n", argv[0]);
  exit(0);
}