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

#include <config.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_errno.h>

/*
 * This module contains routines related to the Cholesky decomposition
 * of a complex Hermitian positive definite matrix.
 */

static void cholesky_complex_conj_vector(gsl_vector_complex *v);

/*
gsl_linalg_complex_cholesky_decomp()
  Perform the Cholesky decomposition on a Hermitian positive definite
matrix. See Golub & Van Loan, "Matrix Computations" (3rd ed),
algorithm 4.2.2.

Inputs: A - (input/output) complex postive definite matrix

Return: success or error

The lower triangle of A is overwritten with the Cholesky decomposition
*/

int
gsl_linalg_complex_cholesky_decomp(gsl_matrix_complex *A)
{
  const size_t N = A->size1;

  if (N != A->size2)
    {
      GSL_ERROR("cholesky decomposition requires square matrix", GSL_ENOTSQR);
    }
  else
    {
      size_t i, j;
      gsl_complex z;
      double ajj;

      for (j = 0; j < N; ++j)
        {
          z = gsl_matrix_complex_get(A, j, j);
          ajj = GSL_REAL(z);

          if (j > 0)
            {
              gsl_vector_complex_const_view aj =
                gsl_matrix_complex_const_subrow(A, j, 0, j);

              gsl_blas_zdotc(&aj.vector, &aj.vector, &z);
              ajj -= GSL_REAL(z);
            }

          if (ajj <= 0.0)
            {
              GSL_ERROR("matrix is not positive definite", GSL_EDOM);
            }

          ajj = sqrt(ajj);
          GSL_SET_COMPLEX(&z, ajj, 0.0);
          gsl_matrix_complex_set(A, j, j, z);

          if (j < N - 1)
            {
              gsl_vector_complex_view av =
                gsl_matrix_complex_subcolumn(A, j, j + 1, N - j - 1);

              if (j > 0)
                {
                  gsl_vector_complex_view aj =
                    gsl_matrix_complex_subrow(A, j, 0, j);
                  gsl_matrix_complex_view am =
                    gsl_matrix_complex_submatrix(A, j + 1, 0, N - j - 1, j);

                  cholesky_complex_conj_vector(&aj.vector);

                  gsl_blas_zgemv(CblasNoTrans,
                                 GSL_COMPLEX_NEGONE,
                                 &am.matrix,
                                 &aj.vector,
                                 GSL_COMPLEX_ONE,
                                 &av.vector);

                  cholesky_complex_conj_vector(&aj.vector);
                }

              gsl_blas_zdscal(1.0 / ajj, &av.vector);
            }
        }

      /* Now store L^H in upper triangle */
      for (i = 1; i < N; ++i)
        {
          for (j = 0; j < i; ++j)
            {
              z = gsl_matrix_complex_get(A, i, j);
              gsl_matrix_complex_set(A, j, i, gsl_complex_conjugate(z));
            }
        }

      return GSL_SUCCESS;
    }
} /* gsl_linalg_complex_cholesky_decomp() */

/*
gsl_linalg_complex_cholesky_solve()
  Solve A x = b where A is in cholesky form
*/

int
gsl_linalg_complex_cholesky_solve (const gsl_matrix_complex * cholesky,
                                   const gsl_vector_complex * b,
                                   gsl_vector_complex * x)
{
  if (cholesky->size1 != cholesky->size2)
    {
      GSL_ERROR ("cholesky matrix must be square", GSL_ENOTSQR);
    }
  else if (cholesky->size1 != b->size)
    {
      GSL_ERROR ("matrix size must match b size", GSL_EBADLEN);
    }
  else if (cholesky->size2 != x->size)
    {
      GSL_ERROR ("matrix size must match solution size", GSL_EBADLEN);
    }
  else
    {
      gsl_vector_complex_memcpy (x, b);

      /* solve for y using forward-substitution, L y = b */

      gsl_blas_ztrsv (CblasLower, CblasNoTrans, CblasNonUnit, cholesky, x);

      /* perform back-substitution, L^H x = y */

      gsl_blas_ztrsv (CblasLower, CblasConjTrans, CblasNonUnit, cholesky, x);

      return GSL_SUCCESS;
    }
} /* gsl_linalg_complex_cholesky_solve() */

/*
gsl_linalg_complex_cholesky_svx()
  Solve A x = b in place where A is in cholesky form
*/

int
gsl_linalg_complex_cholesky_svx (const gsl_matrix_complex * cholesky,
                                 gsl_vector_complex * x)
{
  if (cholesky->size1 != cholesky->size2)
    {
      GSL_ERROR ("cholesky matrix must be square", GSL_ENOTSQR);
    }
  else if (cholesky->size2 != x->size)
    {
      GSL_ERROR ("matrix size must match solution size", GSL_EBADLEN);
    }
  else
    {
      /* solve for y using forward-substitution, L y = b */

      gsl_blas_ztrsv (CblasLower, CblasNoTrans, CblasNonUnit, cholesky, x);

      /* perform back-substitution, L^H x = y */

      gsl_blas_ztrsv (CblasLower, CblasConjTrans, CblasNonUnit, cholesky, x);

      return GSL_SUCCESS;
    }
} /* gsl_linalg_complex_cholesky_svx() */

/********************************************
 *           INTERNAL ROUTINES              *
 ********************************************/

static void
cholesky_complex_conj_vector(gsl_vector_complex *v)
{
  size_t i;

  for (i = 0; i < v->size; ++i)
    {
      gsl_complex z = gsl_vector_complex_get(v, i);
      gsl_vector_complex_set(v, i, gsl_complex_conjugate(z));
    }
} /* cholesky_complex_conj_vector() */


// NIT_inspiral - code to rapidly compute extreme mass-ratio inspirals using self-force results
// Copyright (C) 2017  Niels Warburton
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

#ifndef __INTERPOLANT_H__
#define __INTERPOLANT_H__

#include <vector>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_interp2d.h>
#include <gsl/gsl_spline2d.h>

using namespace std;

typedef vector<double> Vector;

class Interpolant{
	public:
		// 1D interpolation
		Interpolant(Vector x, Vector f);
		double eval(double x);

		// 2D interpolation
		Interpolant(Vector x, Vector y, Vector f);
		double eval(double x, double y);

		// Destructor
		~Interpolant();

	private:
		int interp_type;	// Set to 1 for 1D interpolation and 2 for 2D interpolation

		gsl_spline *spline;
		gsl_spline2d *spline2d;
		gsl_interp_accel *xacc;
		gsl_interp_accel *yacc;
};
#endif // __INTERPOLANT_H__


/*
 * Simultaneous ML fitting of N 3D Gaussian kernels against a volumetric image
 * MATLAB interface in df_mlfitN.c
 *
 * TODO:
 * - how to handle background, still as a constant over the
 * patch?
 * - bg not handled yet.
 *
 */

#include <assert.h>
#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>

#include <gsl/gsl_multimin.h>
#include <gsl/gsl_statistics.h>

#include "blit3.h"
#include "mlfit.h"

#ifndef verbose
#define verbose 1
#endif

// When low on bugs, uncomment the following lines for speed
// Try with valgrind and cmocka first ...
//
// replace gsl_vector_get(v,i) by v->data[i*v->stride]
// #define GSL_RANGE_CHECK_OFF

// Globals
uint32_t maxIterations = 5000;
double convCriteria = 1e-6;

// Headers
int    localizeDot(double *, size_t, double *,  double *);
int    localize(double *, size_t, size_t, size_t, double *, size_t, size_t, double *);
double matrix_get(double * , size_t, size_t, size_t, double *);
double my_f (const gsl_vector *, void *);
int    unit_tests(void);

// Optimization constants
typedef struct {
  double * V; // Volumetric image
  double * W; // Temporary space for model during iterations
  double * D; // dots, x,y,z, nphot, sigmax, sigmay, sigmaz
  size_t Vm, Vn, Vp;
  size_t Dm, Dn;
  double * bg;
} optParams;

double my_f (const gsl_vector *v, void *params)
  // The function to optimize. mlsimplex controls/varies the variables
  // in v. params are to set up the context
{

  optParams *p = (optParams *) params;

  // Size of volumetric image
  const size_t Vm = p->Vm; // Size of image
  const size_t Vn = p->Vn;
  const size_t Vp = p->Vp;
  const size_t Dn = p->Dn; // number of dots
  const size_t Dm = p->Dm;  // and number of features
  double * D = p->D;
  const double * restrict V = p->V; // Volumetric data to fit against
  double * W = p->W; // Pre allocated memory for the model
  double * bg = p->bg;

#if verbose > 0
  printf("-> my_f, %d dots\n", (int) Dn);
#endif

  // Get the other parameters ...

  double photsum = 0;

  // Copy the cordinates and nphot at the current iteration to D
  for(size_t dd = 0; dd<Dn; dd++)
  {
#if verbose > 0
    printf("%lu: (%.1f %.1f %.1f) %.0f (%.1f %.1f %.1f)\n", dd, D[Dm*dd+0],D[Dm*dd+1], D[Dm*dd+2], D[Dm*dd+3], D[Dm*dd+4], D[Dm*dd+5], D[Dm*dd+6]);
#endif
    D[Dm*dd+0] = gsl_vector_get(v, 4*dd+0); // x
    D[Dm*dd+1] = gsl_vector_get(v, 4*dd+1); // y
    D[Dm*dd+2] = gsl_vector_get(v, 4*dd+2); // z
    D[Dm*dd+3] = gsl_vector_get(v, 4*dd+3); // nphot

#if verbose>0
    printf("%lu: xyz (%.1f %.1f %.1f) %.0f (%.1f %.1f %.1f)\n", dd, D[Dm*dd+0],D[Dm*dd+1], D[Dm*dd+2], D[Dm*dd+3], D[Dm*dd+4], D[Dm*dd+5], D[Dm*dd+6]);
#endif
    if(D[Dm*dd+3]<0)
      return INFINITY;

    photsum = photsum + D[7*dd+3];
  }

#if verbose > 0
  printf("phosum of dot parameters: %f\n", photsum);
#endif

  // Copy background map to model
  memcpy(W, bg, sizeof(double)*Vm*Vn*Vp);
  //  memset(W, 0, sizeof(double)*Vm*Vn*Vp); // W = 0;
  blit3g(W, Vm, Vn, Vp, D, Dn, mid_point,0);

  // set bg as well with blit3(g)
  for(size_t kk = 0; kk<Vn*Vm*Vp; kk++) {
    //W[kk]++;
  }

  double vmin = 10e99;
  double vmax = -10e99;
  double wmin = 10e99;
  double wmax = -10e99;
  double diffmax = 0;
  for(size_t kk = 0; kk<Vn*Vm*Vp; kk++)
  {
    vmin = GSL_MIN(vmin, V[kk]);
    vmax = GSL_MAX(vmax, V[kk]);
    wmin = GSL_MIN(wmin, W[kk]);
    wmax = GSL_MAX(wmax, W[kk]);
    diffmax = GSL_MAX(diffmax, fabs(V[kk]-W[kk]));
  }
#if verbose>0
  printf("V -- min: %f max %f\n", vmin, vmax);
  printf("W -- min: %f max %f\n", wmin, wmax);
  printf("diffmax: %f\n", diffmax);
#endif

  /*
   * from LL2PG.m
   *   model = x(3)+x(4)*gaussianInt2([x(1), x(2)], x(5), (size(patch, 1)-1)/2);
   *   mask = disk2d((size(patch,1)-1)/2);
   *   %L = -sum(sum(-(patch-model).^2./model - .5*log(model)));
   *   L = -sum(sum(mask.*(-(patch-model).^2./model - .5*log(model))));
   */

  double E = 0;
  for (size_t kk=0; kk<Vm*Vn*Vp; kk++) {
    //E+= (V[kk]-W[kk])*(V[kk]-W[kk])/V[kk] - .5*log(W[kk]);
    E+= (V[kk]-W[kk])*(V[kk]-W[kk]);
  }
  //  E = -E;

  /* Quadratic
   * for (size_t kk=0; kk<Rw*Rw; kk++)
   * E+= (GI[kk]-R[kk])*(GI[kk]-R[kk]);
   */
#if verbose >0
  printf("E: %f\n", E);
#endif
  return E;
}
double matrix_get(double * V, size_t Vm, size_t Vn, size_t Vp, double * D)
  // Safely get V(D), i.e. check bounds. Returns 0 if outside
  //
{

  double x = D[0];
  double y = D[1];
  double z = D[2];

  if(x<0)
    return 0;
  if(x>Vm-1)
    return 0;
  if(y<0)
    return 0;
  if(y>Vn-1)
    return 0;
  if(z<0)
    return 0;
  if(z>Vp-1)
    return 1;

  size_t pos = (size_t) nearbyint(x) + nearbyint(y)*Vm +  nearbyint(z)*Vm*Vn;
  assert(pos<Vm*Vn*Vp);
  return V[pos];
}

int localize(double * V, size_t Vm, size_t Vn, size_t Vp,
    double * D, size_t Dm, size_t Dn, double * F)
  // Localization for a dot roughly centered in V of size Vm x Vm
  // D[0], D[1], D[3] are the global coordinates of the dot
  // F are the fitted coordinates
{

  // avoid crashing matlab
  gsl_set_error_handler_off();

#if verbose > 0
  printf("-> Localize\n");
  printf("V: %lu %lu %lu\n", Vm, Vn, Vp);
  printf("D: %lu %lu\n", Dm, Dn);
#endif
  assert(Dm == 7); // x y z nphot sigmax sigmay sigmaz

  // Non-optimized parameters
  optParams par;
  par.V = V;
  par.Vm = Vm;
  par.Vn = Vn;
  par.Vp = Vp;
  par.Dn = Dn;
  par.Dm = Dm;
  par.W = malloc(Vm*Vn*Vp*sizeof(double));
  assert(par.W != NULL);
  par.D = D;
  par.bg = malloc(Vm*Vn*Vp*sizeof(double));

#if verbose > 0
  printf("Estimating local background\n");
#endif
  // Estimating a constant from the surroundings of the first dot
  // printf("Vm: %lu %Vn: %lu Vp: %lu, D: %f %f %f\n", Vm, Vn, Vp, D[0], D[1], D[2]);
  double bg = estimateBGV(V, Vm, Vn, Vp, D);
#if verbose>0
  printf("Background estimated to %f\n", bg);
#endif
  printf("Background estimated to %f\n", bg);

  for(size_t kk=0; kk<Vm*Vn*Vp; kk++)
    par.bg[kk] = bg;

  const gsl_multimin_fminimizer_type *T =
    gsl_multimin_fminimizer_nmsimplex2;
  gsl_multimin_fminimizer *s = NULL;
  gsl_vector *ss, *x;
  gsl_multimin_function minex_func;

  size_t iter = 0;
  int status;
  double size;

  /* Starting point */
#if verbose > 0
  printf("Setting up starting vector\n");
#endif
  const int nf = 4; // # of features to be optimized: x, y, z, nphot
  int nParameters = nf*Dn;
  x = gsl_vector_alloc (nParameters);
  for(size_t dd = 0; dd<Dn; dd++)
  {
    gsl_vector_set(x, dd*nf+0, D[dd*Dm+0]); // x position
    gsl_vector_set(x, dd*nf+1, D[dd*Dm+1]); // y position
    gsl_vector_set(x, dd*nf+2, D[dd*Dm+2]); // z position

    // number of photons
    // Since the blitting is scaled by the central pixel, the number
    // of photons can be estimated as the central value of each peak -
    // background from the surrounding
    double nphot = gsl_max(0, matrix_get(V, Vm, Vn, Vp, D+ dd*Dm)-par.bg[0]);
    printf("nphot: %f\n", nphot);
    gsl_vector_set(x, dd*nf+3, nphot);
  }
#if verbose > 0
  printf("Setting up initial step sizes\n");
#endif
  /* Set initial step sizes */
  ss = gsl_vector_alloc(nParameters);
  gsl_vector_set_all(ss, 0.01);
  for(size_t dd = 0; dd<Dn; dd++)
  {
    gsl_vector_set(ss, dd*nf+3, 5); // Number of photons
  }

#if verbose > 0
  printf("Testing the fitting with the initial data\n");
#endif
  my_f(x,(void*) &par);

  /* Initialize method and iterate */
  minex_func.n = nParameters;
  minex_func.f = my_f;
  minex_func.params = ∥
#if verbose > 0
  printf("Allocating for the solver\n");
#endif
  s = gsl_multimin_fminimizer_alloc (T, nParameters);
#if verbose > 0
  printf("Initializing solver\n");
#endif
  gsl_multimin_fminimizer_set (s, &minex_func, x, ss);

  do
  {
    iter++;
#if verbose>0
    printf("----------------------------- Iteration: %lu\n", iter);
#endif
    status = gsl_multimin_fminimizer_iterate(s);

    if (status)
      break;

    size = gsl_multimin_fminimizer_size(s);
    status = gsl_multimin_test_size(size, convCriteria);

    if (status == GSL_SUCCESS)
    {
#if verbose > 0
      printf ("converged to minimum at\n");
      printf ("%5lu x:%10.3e y:%10.3e z:%10.3e NP:%6.1f f() = %7.3f size = %10.3e\n",
          iter,
          gsl_vector_get (s->x, 0),
          gsl_vector_get (s->x, 1),
          gsl_vector_get (s->x, 2),
          gsl_vector_get (s->x, 3),
          s->fval, size);
#endif
    }
  }
  while (status == GSL_CONTINUE && iter < maxIterations);

  for(size_t dd = 0; dd<Dn; dd++)
  {
    F[5*dd+0] = gsl_vector_get (s->x, nf*dd+0); // x
    F[5*dd+1] = gsl_vector_get (s->x, nf*dd+1); // y
    F[5*dd+2] = gsl_vector_get (s->x, nf*dd+2); // z
    F[5*dd+3] = gsl_vector_get (s->x, nf*dd+3); // nphot
    F[5*dd+4] = status;                            // status
  }

  gsl_vector_free(x);
  gsl_vector_free(ss);
  gsl_multimin_fminimizer_free (s);

  free(par.W);
  free(par.bg);
  return status;
}

#ifndef _MAIN
#define _MAIN
int unit_tests(){
  double * V; // image
  int Vm = 102; int Vn = 107; int Vp = 60;
  double * D; // list of dots
  size_t Dm = 7; // number of features of the dots
  size_t Dn = 10; // number of dots
  double * F; // fitted dots

  printf("Image size: %dx%dx%d\n", Vm, Vn, Vp);
  printf("Localizing %lu dots\n", Dn);

  V = malloc(Vm*Vn*Vp*sizeof(double));
  D = malloc(Dm*Dn*sizeof(double));
  F = malloc(5*Dn*sizeof(double));

  // Initialize the data
  for(int kk=0; kk<Vm*Vn*Vp; kk++)
    V[kk] = rand_range(2,9);

  for(uint32_t kk=0; kk<Dn; kk++)
  {
    size_t pos = kk*Dm;
    D[pos] =   rand_range(6, Vm-7);
    D[pos+1] = rand_range(6, Vn-7);
    D[pos+2] = rand_range(6, Vp-7);
    D[pos+3] = 1000;
    D[pos+4] = 1.2;
    D[pos+5] = 1.2;
    D[pos+6] = 1.4;

    if(D[pos] < 6)
      D[pos] = 6;
    if(D[pos+1] < 6)
      D[pos+1] = 6;
    if(D[pos+2] < 6)
      D[pos+2] = 6;
#if verbose > 0
    printf("D %03d %f %f %f\n", kk, D[pos], D[pos+1], D[pos+2]);
#endif
  }

  for(uint32_t kk=0; kk<Dn; kk++)
  {
    size_t pos = nearbyint(D[kk*7]) + nearbyint(D[kk*7+1])*Vm + nearbyint(D[kk*7+2])*Vm*Vn;
    assert(pos<(size_t) Vm*Vn*Vp);
    assert(pos>0);
    V[pos] = 5;
  }

  // Run the optimization
  localize(V, Vm, Vn, Vp, D, Dm, Dn, F);

  // In next version, also supply clustering information

#if verbose >0
  for(size_t kk = 0; kk<Dm; kk++)
  {
    size_t Dpos = kk*7;
    size_t Fpos = kk*5;
    printf("%6lu (%f, %f, %f) -> (%f, %f, %f)\n", kk,
        D[Dpos], D[Dpos+1], D[Dpos+2],
        F[Fpos], F[Fpos+1], F[Fpos+2]);
  }
#endif
  free(F);
  free(D);
  free(V);
  return 0;
}

int main(int argc, char ** argv)
  // For testing, not used any more, see the MATLAB interface in
  // df_mlfit.c
{

  printf("%s\n", argv[0]);

  if(argc == 1)
    return unit_tests();

  return 0;
}
#endif


/**
  *
  * For reference: Pearson's rho...
  *
  *               sum( (x_i-E[x])(y_i-E[y]) )
  * rho = -------------------------------------------
  *       sqrt( sum((x_i-E[x])^2) sum((y_i-E[y])^2) )
  */

#include <stdio.h>
#include <stdlib.h>
#include <stdbool.h>
#include <limits.h>
#include <string.h>
#include <math.h>
#include <assert.h>
#include <alloca.h>

#include <gsl/gsl_cdf.h>

#include "stattest.h"
#include "mix.h"
#include "bvr.h"
#include "limits.h"

struct Pair {
	float cv; // Continuous Variable
	unsigned int dv; // Discrete Variable
};

static int _cmp_pair( const void *pvl, const void *pvr ) {
	const struct Pair *l = (const struct Pair *)pvl;
	const struct Pair *r = (const struct Pair *)pvr;
	const float delta = l->cv - r->cv;
	if( delta == 0.0 )
		return 0.0;
	else
		return  delta < 0.0 ? -1 : +1;
}

struct MixCovars {

	/**
	  * Maximum number of samples
	  */
	unsigned int SAMPLE_CAPACITY;

	/**
	  * All category labels pushed into this accumulator must
	  * be in [0,category_capacity).
	 */
	unsigned int CATEGORY_CAPACITY;

	/**
	  * Count of samples pushed since last cat_clear.
	  */
	unsigned int sample_count;

	/**
	  * This essentially bounds the allowed category LABELS.
	  * Any category label push'ed in must be in [0,expected_categories).
	  * Only this member, NOT category_capacity, should be used to size
	  * counts arrays!
	  */
	unsigned int expected_categories;

	/**
	  * This is an count of DISTINCT LABELS observed by push.
	  */
	unsigned int observed_categories;

	// Calculation state ///////////////////////////////////////////////////

	double mean_rank;  // necessarily of both covariates
	double sum_sq_dev; // used for Kruskal-Wallis and Spearman

	/**
	  * Following struct array is ONLY for Spearman rho calculation.
	  * Note that the number of non-empty categories after data entry
	  * says NOTHING about WHICH categories (indices) are non-zero.
	  * An earlier implementation was assuming category_count[0,1] > 0,
	  * BUT THIS NEED NOT BE TRUE! Thus, all this hack...
	  */
	struct {
		unsigned int index;
		unsigned int count;
		unsigned int meanRank;
	} edge[2];

	double sum_dev_prod; // used only for Spearman-rho

	// Buffer //////////////////////////////////////////////////////////////

	/**
	  * The TOTAL amount of malloc'ed space (for the purposes
	  * of fast clearing with memset).
	  */
	size_t SIZEOF_BUFFERS;

	/**
	  * Two buffers
	  */
	struct Pair *samples;

	/**
	  * A buffer allocated once in construction and reused for
	  * all subsequent analyses.
	  * Reallocation is not currently supported.
	  */
	unsigned int *category_count;
};

#ifdef _UNITTEST_MIX_
unsigned arg_min_mixb_count = 1;
#else
extern unsigned arg_min_mixb_count;
#endif

/**
  * Both minCat and maxCat assume that samples is non-empty,
  * but it's entirely possible that there were NO samples.
  */
static unsigned _minCat( struct MixCovars *co ) {
	unsigned int i = 0;
	while( i < co->expected_categories ) {
		if( co->category_count[i] > 0 )
			return i;
		i++;
	}
	return INT_MAX; // just to silence compiler.
}

static unsigned _maxCat( struct MixCovars *co ) {

	int i = co->expected_categories;
	while( i-- > 0 ) {
		if( co->category_count[i] > 0 )
			return (unsigned int)i;
	}
	return INT_MAX; // just to silence compiler
}

/**
  * Notice that all calcuations are performed without actually requiring
  * an array of ranks.
  */
static unsigned int _rank_sums( struct MixCovars *co,
	double *sums ) {

	const unsigned int N = co->sample_count;
	const bool PRECALC_SPEARMAN_PRODUCT
		= co->observed_categories == 2;
	const unsigned UPPER
		= co->edge[1].index;
	unsigned int ties = 0;
	unsigned int until = 0;
	float rank = 0;
	double diff;

	// The default pair::operator< automatically uses the first element.

	qsort( co->samples, co->sample_count, sizeof(struct Pair), _cmp_pair );

	for(unsigned int i = 0; i < N; i++ ) {

		const unsigned int cat
			= co->samples[i].dv;

		// The rank of the current sample is i+1 UNLESS we're in
		// the midst of a run of samples with the same value in
		// which case the last computed rank (a mean) is reused.

		if( ! ( i < until ) ) {

			until = i + 1;
			while( until < N
					&& ( co->samples[i].cv == co->samples[until].cv ) ) {
				ties++;
				until++;
			}

			if( i == 0 && until == N ) {
				// degenerate case: all ties.
			}

			// Mean of n integers starting on r:
			// [ (r+0) + (r+1) + (r+2) + ... + (r+(n-1)) ] / n
			// ... = \frac{nr + \sum_{i=0}^{n-1} i}{n}
			// ... = r + (1/n)[ (n-1)n / 2 ]
			// ... = r + (n-1)/2

			rank = (1.0+i) + ((until-i)-1)/2.0;

			// ...notice if the very next sample -is- different
			// then rank == i+1, the trivial case.
		}

		/**
		  * Warning: If the numeric feature is constant then rank==meanRank
		  * and diff == 0
		  */
		diff = rank - co->mean_rank;

		co->sum_sq_dev += ( diff * diff );

		if( PRECALC_SPEARMAN_PRODUCT ) {
			co->sum_dev_prod += ( diff * ( co->edge[UPPER==cat?1:0].meanRank - co->mean_rank ) );
		}

		sums[ cat ] += rank;
	}

	return ties;
}

#if defined(_UNITTEST_MIX_)
static void dbg_dump( struct MixCovars *co, FILE *fp ) {
	for(unsigned int i = 0; i < co->sample_count; i++ )
		fprintf( fp, "%.3f\t%d\n", co->samples[i].cv, co->samples[i].dv );
}
#endif

/***************************************************************************
  * Publics
  */

void mix_destroy( void *pv ) {

	if( pv ) {
		struct MixCovars *co = (struct MixCovars *)pv;
		if( co->samples )
			free( co->samples );
		free( pv );
	}
}

/**
  * Pre-allocate a set of working buffers large enough for all anticipated
  * calculations (max feature length) and a struct to wrap them.
  */
void *mix_create( unsigned int sample_capacity, unsigned int category_capacity ) {

	struct MixCovars *co
		= calloc( 1, sizeof(struct MixCovars) );
	if( co ) {
		co->SAMPLE_CAPACITY   = sample_capacity;
		co->CATEGORY_CAPACITY = category_capacity;
		co->SIZEOF_BUFFERS
			= ( sample_capacity * sizeof(struct Pair) )
			+ ( category_capacity * sizeof(unsigned int) );

		// Allocate one large buffer and partition it up.
		co->samples
			= calloc( co->SIZEOF_BUFFERS, sizeof(char) );
		co->category_count
			= (unsigned int*)(co->samples
			+ sample_capacity);
		// If -anything- failed clean up any successes.
		if( NULL == co->samples ) {
			mix_destroy( co );
			return NULL;
		}
		return co;
	}
	return NULL;
}

void mix_clear( void *pv, unsigned expcat ) {

	struct MixCovars *co = (struct MixCovars *)pv;

	assert( expcat <= co->CATEGORY_CAPACITY );

	co->sample_count = 0;
	co->expected_categories = expcat;
	co->observed_categories = 0;
	co->mean_rank    = 0.0;
	co->sum_sq_dev   = 0.0;
	co->sum_dev_prod = 0.0;
	memset( co->edge, 0, sizeof(co->edge) );
	co->sum_dev_prod = 0.0;

	memset( co->samples, 0, co->SIZEOF_BUFFERS );
}

void mix_push( void *pv, float num, unsigned int cat ) {

	struct MixCovars *co = (struct MixCovars *)pv;

	assert( cat              < co->expected_categories );
	assert( co->sample_count < co->SAMPLE_CAPACITY     );

	if( 0 == co->category_count[ cat ] )
		co->observed_categories++;

	co->category_count[ cat ] += 1;

	// Not updating edges in here because the number of conditionals
	// executed for sample counts > 32 exceeds the work to find the
	// edges post-sample accumulation.
	co->samples[ co->sample_count ].cv = num;
	co->samples[ co->sample_count ].dv = cat;
	co->sample_count++;
}

size_t mix_size( void *pv ) {
	struct MixCovars *co = (struct MixCovars *)pv;
	return co->sample_count;
}

/**
  * - Establish whether or not Spearman is even sensible (according to
  *   whether or not the categorical variable is binary).
  * - Compute the mean rank used throughout subsequent calcs.
  * - Compute the mean ranks of the categorical values assuming
  *   they -are- binary.
  */
bool mix_complete( void *pv ) {

	struct MixCovars *co = (struct MixCovars *)pv;
	if( mix_degenerate( co ) ) return false;

	co->mean_rank = (1.0 + co->sample_count)/2.0;

	co->edge[1].index = _maxCat( co ); // ...this is used by kruskal_wallis, but

	assert( co->edge[1].index < co->expected_categories );

	// ...the remainder of the edge[] struct is only used if...

	if( co->observed_categories == 2 ) {

		// All the following is relevant only to Spearman rho calculation.

		co->edge[0].index = _minCat( co );
		assert( co->edge[0].index < co->expected_categories );

		co->edge[0].count = co->category_count[ co->edge[0].index ];
		co->edge[0].meanRank = MEAN_RANK_OF_TIES( 0, co->edge[0].count );

		co->edge[1].count = co->category_count[ co->edge[1].index ];

		if( co->edge[0].count < arg_min_mixb_count ||
			co->edge[1].count < arg_min_mixb_count )
			return false; // another degeneracy class

		co->edge[1].meanRank = MEAN_RANK_OF_TIES( co->edge[0].count, co->sample_count );

		assert( (co->edge[0].count + co->edge[1].count) == co->sample_count );
	}
	return true;
}

bool mix_degenerate( void *pv ) {
	struct MixCovars *co = (struct MixCovars *)pv;
	return co->observed_categories < 2 || co->sample_count < 2;
}

bool mix_categoricalIsBinary( void *pv ) {
	struct MixCovars *co = (struct MixCovars *)pv;
	return co->observed_categories == 2;
}

/**
 * Calculates the Kruskal-Wallis statistic and optionally a p-value.
 * Importantly, it does it in a one pass iteration over the data.
 */
int mix_kruskal_wallis( void *pv, struct Statistic *result ) {

	struct MixCovars *co = (struct MixCovars *)pv;
	const unsigned int N
		= co->sample_count ;
	const size_t SIZEOF_SUMS
		= co->expected_categories * sizeof(double);
	double *rank_sum
		= (double*)alloca( SIZEOF_SUMS );
	double numerator = 0.0;

	memset( rank_sum, 0, SIZEOF_SUMS   );

	result->extra_value[0] = _rank_sums( co, rank_sum );

	for(unsigned int i = 0; i <= co->edge[1].index; i++ ) {
		if( co->category_count[i] > 0 ) {
			rank_sum[i] /= co->category_count[i];
			double delta
				= rank_sum[i] - co->mean_rank;
			numerator
				+= ( co->category_count[i] * delta * delta );
		}
	}

	result->name
		= "Kruskal-Wallis_K";
	result->sample_count
		= N;
	result->value
		= (N-1) * ( numerator / co->sum_sq_dev );
	result->probability
		= gsl_cdf_chisq_Q( result->value, co->observed_categories-1 );

	return 0;
}

/**
 * P-value returned is ONE-SIDED. Double it for two-sided.
 */
#ifdef HAVE_MANN_WHITNEY
int mix_mann_whitney( void *pv, struct Statistic *result ) {

	struct MixCovars *co = (struct MixCovars *)pv;
	const size_t SIZEOF_SUMS
		= co->expected_categories * sizeof(double);
	double *rank_sum
		= (double*)alloca( SIZEOF_SUMS );
	memset( rank_sum, 0, SIZEOF_SUMS   );

	ms->ties
		= _rank_sums( co, rank_sum );

	// TODO: If too many ties, the following method of calculating U
	//       becomes invalid. What to do?

	const unsigned int SQ
		= category_count[0]*category_count[1];

	// If either group count is 0 we're done...

	if( SQ ) {

		const double U1
			= rank_sum[0] - (category_count[0]*(category_count[0]+1.0))/2.0;
		const double U2
			= rank_sum[1] - (category_count[1]*(category_count[1]+1.0))/2.0;

		assert( U1+U2 == SQ );

		ms->U = U1 < U2 ? U1 : U2;

		const double mean_U
			= SQ / 2.0;
		const double sigma_U
			= sqrt( SQ*( category_count[0] + category_count[1] + 1.0 ) / 12.0 );

		cs->P = gsl_cdf_ugaussian_P( ( ms->U - mean_U ) / sigma_U );
	}
	return 0;
}
#endif

/**
  * This calculation depends on a quantities calculated within
  * kruskal_wallis (or mann_whitney), so must be called after
  * either of them.
  * (This is also the only reason this method can be const.)
  */
double mix_spearman_rho( void *pv ) {

	struct MixCovars *co = (struct MixCovars *)pv;
	/**
	  * denominator requires sum of square deviation for
	  * the binary variables which collapses to...
	  */
	double SQD0 = (co->edge[0].meanRank - co->mean_rank ); SQD0 *= SQD0;
	double SQD1 = (co->edge[1].meanRank - co->mean_rank ); SQD1 *= SQD1;

	const double SUMSQD
		= co->edge[0].count * SQD0 + co->edge[1].count * SQD1;
	/**
	  * And the sum of squared deviation of the numeric covariate's
	  * ranks (sum_sq_dev) was (necessarily) calculated earlier!
	  */
	return co->sum_dev_prod / sqrt( co->sum_sq_dev * SUMSQD );
}

#ifdef _UNITTEST_MIX_

/**
 * This is intended to be exercised with the following R script
 * that generates a small table of grouped floating-point values
 * executes a Kruskal-Wallis test on it, and dumps the table
 * to a tab-delimited file with the test results as a comment on
 * the first line.
 *
 * x <- data.frame(
 * 		cat=as.integer(gl(3,4))-1,
 * 		num=c( runif(4)*10, runif(4)*20, 10+runif(4)*10 ) );
 * k <- with( x, kruskal.test( num, cat ) );
 * cat( sprintf( "# K=%f p-value=%f\n", k$statistic, k$p.value ), file="foo.tab" );
 * write.table( x, 'foo.tab', quote=F, sep='\t', row.names=F, col.names=F, append=TRUE );
 */
#include <err.h>

int main( int argc, char *argv[] ) {

	if( argc >= 3 ) {

		struct Statistic result;
		const int EXPCAT
			= atoi( argv[1] );
		struct MixCovars *accum
		   = mix_create( atoi( argv[2] ), MAX_CATEGORY_COUNT );
		FILE *fp
			= argc > 3
			? fopen( argv[3], "r" )
			: stdin;
		char *line = NULL;
		size_t n = 0;

		mix_clear( accum, EXPCAT );

		while( getline( &line, &n, fp ) > 0 ) {
			unsigned int cat;
			float        num;
			if( line[0] == '#' ) {
				fputs( line, stdout );
				continue;
			}
			if( 2 == sscanf( line, "%d\t%f\n", &cat, &num ) ) {
				mix_push( accum, num, cat );
			} else {
				fprintf( stderr, "Failure parsing line: %s", line );
			}
		}
		free( line );
		fclose( fp );

		// Always calculate Kruskal-Wallis since it's valid for any groups
		// count >= 2...

		if( mix_complete( accum ) ) {
			memset( &result, 0, sizeof(struct Statistic) );
			if( mix_kruskal_wallis( accum, &result ) == 0 ) {
				printf( "K=%f, p-value=%f", result.value, result.probability );
				if( mix_categoricalIsBinary( accum ) )
					printf( ", spearman=%f\n", mix_spearman_rho(accum) );
				else
					fputc( '\n', stdout );
			} else {
				printf( "error\n" );
			}
		}

#ifdef HAVE_MANN_WHITNEY
		// Do a Mann-Whitney, too, if there are only 2 groups.
		if( mix_categoricalIsBinary(accum) == 2 ) {
		}
#endif
		mix_destroy( accum );
	} else
		err( -1, "%s <categories> <sample count> [ <input file> ]", argv[0] );
	return 0;
}
#endif


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

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

#include "ccl.h"

ccl_cl_tracer_collection_t *ccl_cl_tracer_collection_t_new(int *status) {
  ccl_cl_tracer_collection_t *trc = NULL;
  trc = malloc(sizeof(ccl_cl_tracer_collection_t));
  if (trc == NULL)
    *status = CCL_ERROR_MEMORY;

  if (*status == 0) {
    trc->n_tracers = 0;
    // Currently CCL_MAX_TRACERS_PER_COLLECTION is hard-coded to 100.
    // It should be enough for any practical application with minimal memory overhead
    trc->ts = malloc(CCL_MAX_TRACERS_PER_COLLECTION*sizeof(ccl_cl_tracer_t *));
    if (trc->ts == NULL) {
      *status = CCL_ERROR_MEMORY;
      free(trc);
      trc = NULL;
    }
  }

  return trc;
}

void ccl_cl_tracer_collection_t_free(ccl_cl_tracer_collection_t *trc) {
  if (trc != NULL) {
    if (trc->ts != NULL)
      free(trc->ts);
    free(trc);
  }
}

void ccl_add_cl_tracer_to_collection(ccl_cl_tracer_collection_t *trc,
                                     ccl_cl_tracer_t *tr, int *status) {
  if (trc->n_tracers >= CCL_MAX_TRACERS_PER_COLLECTION) {
    *status = CCL_ERROR_MEMORY;
    return;
  }
  trc->ts[trc->n_tracers] = tr;
  trc->n_tracers++;
}

// Takes an array of z-dependent numbers and the corresponding z values
// and returns an array of a values and the corresponding a-dependent values.
// The order of the original arrays is assumed to be ascending in z, and
// the order of the returned arrays is swapped (so it has ascending a).
static void from_z_to_a(ccl_cosmology *cosmo,
                        int nz, double *z_arr, double *fz_arr,
                        double **a_arr, double **fa_arr, int *status) {
  *a_arr = NULL;
  *fa_arr = NULL;
  *a_arr = malloc(nz*sizeof(double));
  *fa_arr = malloc(nz*sizeof(double));

  if ((a_arr == NULL) || (fa_arr == NULL)) {
    *status = CCL_ERROR_MEMORY;
    free(*a_arr);
    free(*fa_arr);
    ccl_cosmology_set_status_message(
      cosmo, "ccl_tracers.c: from_z_to_a(): memory allocation error\n");
  }

  if (*status == 0) {
    int ia;
    // Populate array of scale factors in reverse order
    for (ia=0; ia < nz; ia++) {
      (*a_arr)[ia] = 1./(1+z_arr[nz-1-ia]);
      (*fa_arr)[ia] = fz_arr[nz-1-ia];
    }
  }
}

//Integrand for N(z) integrator
static double nz_integrand(double z, void *pars) {
  ccl_f1d_t *nz_f = (ccl_f1d_t *)pars;

  return ccl_f1d_t_eval(nz_f,z);
}

// Gets area of N(z) curve
static double get_nz_norm(ccl_cosmology *cosmo, ccl_f1d_t *nz_f,
                          double z0, double zf, int *status) {
  double nz_norm = -1, nz_enorm;

  // Get N(z) norm
  gsl_function F;
  gsl_integration_workspace *w = NULL;
  F.function = &nz_integrand;
  F.params = nz_f;

  w = gsl_integration_workspace_alloc(cosmo->gsl_params.N_ITERATION);

  if (w == NULL) {
    *status = CCL_ERROR_MEMORY;
    ccl_cosmology_set_status_message(
      cosmo, "ccl_tracers.c: get_nz_norm(): out of memory");
  }
  else {
    int gslstatus = gsl_integration_qag(
      &F, z0, zf, 0,
      cosmo->gsl_params.INTEGRATION_EPSREL,
      cosmo->gsl_params.N_ITERATION,
      cosmo->gsl_params.INTEGRATION_GAUSS_KRONROD_POINTS,
      w, &nz_norm, &nz_enorm);

    if (gslstatus != GSL_SUCCESS) {
      ccl_raise_gsl_warning(gslstatus, "ccl_tracers.c: get_nz_norm():");
      *status = CCL_ERROR_INTEG;
      ccl_cosmology_set_status_message(
        cosmo,
        "ccl_tracers.c: get_nz_norm(): "
        "integration error when normalizing N(z)\n");
    }
  }

  gsl_integration_workspace_free(w);

  return nz_norm;
}

static void from_z_to_chi(ccl_cosmology *cosmo, int nz, double *z_arr,
                          double *chi_arr, int *status) {
  for (int ichi=0; ichi < nz; ichi++)
    chi_arr[ichi] = ccl_comoving_radial_distance(cosmo, 1./(1+z_arr[ichi]), status);
}

void ccl_get_number_counts_kernel(ccl_cosmology *cosmo,
                                  int nz, double *z_arr, double *nz_arr,
                                  int normalize_nz,
                                  double *pchi_arr, int *status) {
  // Returns dn/dchi normalized to unit area from an unnormalized dn/dz.
  // Prepare N(z) spline
  ccl_f1d_t *nz_f = NULL;

  nz_f = ccl_f1d_t_new(nz, z_arr, nz_arr, 0, 0);
  if (nz_f == NULL) {
    *status = CCL_ERROR_SPLINE;
    ccl_cosmology_set_status_message(
      cosmo,
      "ccl_tracers.c: ccl_get_number_counts_kernel: "
      "error initializing spline\n");
  }

  // Get N(z) normalization
  double i_nz_norm = -1;
  if (*status == 0) {
    if (normalize_nz)
      i_nz_norm = 1./get_nz_norm(cosmo, nz_f, z_arr[0], z_arr[nz-1], status);
    else
      i_nz_norm = 1;
  }

  if (*status == 0) {
    // Populate arrays
    for(int ichi=0; ichi < nz; ichi++) {
      double a = 1./(1+z_arr[ichi]);
      double h = cosmo->params.h*ccl_h_over_h0(cosmo,a,status)/ccl_constants.CLIGHT_HMPC;
      // H(z) * dN/dz * 1/Ngal
      pchi_arr[ichi] = h*nz_arr[ichi]*i_nz_norm;
    }
  }

  ccl_f1d_t_free(nz_f);
}

//3 H0^2 Omega_M / 2
static double get_lensing_prefactor(ccl_cosmology *cosmo,int *status) {
  double hub = cosmo->params.h/ccl_constants.CLIGHT_HMPC;
  return 1.5*hub*hub*cosmo->params.Omega_m;
}

typedef struct {
  ccl_cosmology *cosmo;
  double z_max;
  double z_end;
  double chi_end;
  double i_nz_norm;
  ccl_f1d_t *nz_f;
  ccl_f1d_t *sz_f;
  int *status;
} integ_lensing_pars;

// Integrand for lensing kernel.
// Returns N(z) * (1 - 5*s(z)/2) * (chi(z)-chi) / chi(z)
static double lensing_kernel_integrand(double z, void *pars) {
  integ_lensing_pars *p = (integ_lensing_pars *)pars;
  double pz = ccl_f1d_t_eval(p->nz_f, z);
  double qz;
  if (p->sz_f == NULL) // No magnification factor
    qz = 1;
  else // With magnification factor
    qz = (1 - 2.5*ccl_f1d_t_eval(p->sz_f, z));

  if (z == 0)
    return pz * qz;
  else {
    double chi = ccl_comoving_radial_distance(p->cosmo, 1./(1+z), p->status);
    return (
      pz * qz *
      ccl_sinn(p->cosmo, chi-p->chi_end, p->status) /
      ccl_sinn(p->cosmo, chi, p->status));
  }
}

// Returns
// Integral[ p(z) * (1-5s(z)/2) * chi_end * (chi(z)-chi_end)/chi(z) , {z',z_end,z_max} ]
static double lensing_kernel_integrate(ccl_cosmology *cosmo,
                                       integ_lensing_pars *pars,
                                       gsl_integration_workspace *w) {
  int gslstatus = 0;
  double result, eresult;
  gsl_function F;
  F.function = &lensing_kernel_integrand;
  F.params = pars;
  gslstatus = gsl_integration_qag(
    &F, pars->z_end, pars->z_max, 0,
    cosmo->gsl_params.INTEGRATION_EPSREL,
    cosmo->gsl_params.N_ITERATION,
    cosmo->gsl_params.INTEGRATION_GAUSS_KRONROD_POINTS,
    w, &result, &eresult);

  if ((gslstatus != GSL_SUCCESS) || (*(pars->status))) {
    ccl_raise_gsl_warning(gslstatus, "ccl_tracers.c: lensing_kernel_integrate():");
    return -1;
  }

  return result * pars->i_nz_norm * pars->chi_end;
}

//Returns number of divisions on which
//the lensing kernel should be calculated
int ccl_get_nchi_lensing_kernel(int nz, double *z_arr, int *status) {
  int nchi;
  double dz = -1, z_max = -1;
  //Compute redshift step
  dz = (z_arr[nz-1]-z_arr[0])/(nz-1);

  //How many steps to z=0?
  return (int)(z_arr[nz-1]/dz+0.5);
}

//Return array with the values of chi at
//the which the lensing kernel will be
//calculated.
void ccl_get_chis_lensing_kernel(ccl_cosmology *cosmo,
                                 int nchi, double z_max,
                                 double *chis, int *status) {
  double dz = z_max/nchi;
  for(int ichi=0; ichi < nchi; ichi++) {
    double z = dz*ichi+1E-15;
    double a = 1./(1+z);
    chis[ichi] = ccl_comoving_radial_distance(cosmo, a, status);
  }
}

//Returns array with lensing kernel:
//3 * H0^2 * Omega_M / 2 / a *
// Integral[ p(z) * (1-5s(z)/2) * chi_end * (chi(z)-chi_end)/chi(z) ,
//          {z',z_end,z_max} ]
void ccl_get_lensing_mag_kernel(ccl_cosmology *cosmo,
                                int nz, double *z_arr, double *nz_arr,
                                int normalize_nz, double z_max,
                                int nz_s, double *zs_arr, double *sz_arr,
                                int nchi, double *chi_arr, double *wL_arr,
                                int *status) {
  ccl_f1d_t *nz_f = NULL;
  ccl_f1d_t *sz_f = NULL;
  gsl_integration_workspace *w;
  integ_lensing_pars *ipar;
  int local_status;

  // Prepare N(z) spline
  nz_f = ccl_f1d_t_new(nz, z_arr, nz_arr, 0, 0);
  if (nz_f == NULL) {
    *status = CCL_ERROR_SPLINE;
    ccl_cosmology_set_status_message(
      cosmo,
      "ccl_tracers.c: get_lensing_mag_kernel: error initializing spline\n");
  }

  // Get N(z) normalization
  double i_nz_norm = -1;
  if (*status == 0) {
    if (normalize_nz)
      i_nz_norm = 1./get_nz_norm(cosmo, nz_f, z_arr[0], z_arr[nz-1], status);
    else
      i_nz_norm = 1.;
  }

  // Prepare magnification bias spline if needed
  if (*status == 0) {
    if ((nz_s > 0) && (zs_arr != NULL) && (sz_arr != NULL)) {
      sz_f = ccl_f1d_t_new(nz_s, zs_arr, sz_arr, sz_arr[0], sz_arr[nz_s-1]);
      if (sz_f == NULL) {
        *status = CCL_ERROR_SPLINE;
        ccl_cosmology_set_status_message(
          cosmo,
          "ccl_tracers.c: get_lensing_mag_kernel: error initializing spline\n");
      }
    }
  }

  double lens_prefac = get_lensing_prefactor(cosmo, status);
  double chi, a, z, mgfac;
  int ichi;

  w = NULL;
  ipar = NULL;
  local_status = *status;

  if (local_status == 0) {
    ipar = malloc(sizeof(integ_lensing_pars));
    w = gsl_integration_workspace_alloc(cosmo->gsl_params.N_ITERATION);

    if ((ipar == NULL) || (w == NULL)) {
      local_status = CCL_ERROR_MEMORY;
    }
  }

  if (local_status == 0) {
    ipar->cosmo = cosmo;
    ipar->z_max = z_max;
    ipar->i_nz_norm = i_nz_norm;
    ipar->sz_f = sz_f;
    ipar->nz_f = nz_f;
    ipar->status = &local_status;
  }

  //Populate arrays
  for (ichi=0; ichi < nchi; ichi++) {
    if (local_status == 0) {
      chi = chi_arr[ichi];
      a = ccl_scale_factor_of_chi(cosmo, chi, &local_status);
      z = 1./a-1;
      // Add MG correction if needed
      mgfac = 1.0;
      if (fabs(cosmo->params.sigma_0))
        mgfac += ccl_Sig_MG(cosmo, a, &local_status);
      ipar->z_end = z;
      ipar->chi_end = chi;

      wL_arr[ichi] = lensing_kernel_integrate(cosmo, ipar, w)*(1+z)*lens_prefac*mgfac;
    } else {
      wL_arr[ichi] = NAN;
    }
  }

  gsl_integration_workspace_free(w);
  free(ipar);

  if (local_status != 0) {
    *status = CCL_ERROR_INTEG;
  }

  ccl_f1d_t_free(nz_f);
  ccl_f1d_t_free(sz_f);
}

// Returns kernel for CMB lensing
// 3H0^2Om/2 * chi * (chi_s - chi) / chi_s / a
void ccl_get_kappa_kernel(ccl_cosmology *cosmo, double chi_source,
                          int nchi, double *chi_arr,
                          double *wchi, int *status) {
  double lens_prefac = get_lensing_prefactor(cosmo, status) / ccl_sinn(cosmo, chi_source, status);

  for (int ichi=0; ichi < nchi; ichi++) {
    double chi = chi_arr[ichi];
    double a = ccl_scale_factor_of_chi(cosmo, chi, status);
    double mgfac = 1;
    // Add MG correction if needed
    if (fabs(cosmo->params.sigma_0))
      mgfac += ccl_Sig_MG(cosmo, a, status);
    wchi[ichi] = lens_prefac*(ccl_sinn(cosmo,chi_source-chi,status))*chi*mgfac/a;
  }
}

ccl_cl_tracer_t *ccl_cl_tracer_t_new(ccl_cosmology *cosmo,
                                     int der_bessel,
                                     int der_angles,
                                     int n_w, double *chi_w, double *w_w,
                                     int na_ka, double *a_ka,
                                     int nk_ka, double *lk_ka,
                                     double *fka_arr,
                                     double *fk_arr,
                                     double *fa_arr,
                                     int is_fka_log,
                                     int is_factorizable,
                                     int extrap_order_lok,
                                     int extrap_order_hik,
                                     int *status) {
  ccl_cl_tracer_t *tr = NULL;

  // Check der_bessel and der_angles are sensible
  if ((der_angles < 0) || (der_angles > 2)) {
    *status = CCL_ERROR_INCONSISTENT;
    ccl_cosmology_set_status_message(
      cosmo,
      "ccl_tracers.c: ccl_cl_tracer_new: der_angles must be between 0 and 2\n");
  }
  if ((der_bessel < -1) || (der_bessel > 2)) {
    *status = CCL_ERROR_INCONSISTENT;
    ccl_cosmology_set_status_message(
      cosmo,
      "ccl_tracers.c: ccl_cl_tracer_new: der_bessel must be between -1 and 2\n");
  }

  if (*status == 0) {
    tr = malloc(sizeof(ccl_cl_tracer_t));
    if (tr == NULL)
      *status = CCL_ERROR_MEMORY;
  }

  // Initialize everythin
  if (*status == 0) {
    tr->der_angles = der_angles;
    tr->der_bessel = der_bessel;
    tr->kernel = NULL; // Initialize these to NULL
    tr->transfer = NULL; // Initialize these to NULL
    tr->chi_min = 0;
    tr->chi_max = 1E15;
  }

  if (*status == 0) {
    // Initialize radial kernel
    if ((n_w > 0) && (chi_w != NULL) && (w_w != NULL)) {
      tr->kernel = ccl_f1d_t_new(n_w,chi_w,w_w,0,0);
      if (tr->kernel == NULL)
        *status=CCL_ERROR_MEMORY;
    }
  }

  // Find kernel edges
  if (*status == 0) {
    // If no radial kernel, set limits to zero and maximum distance
    if (tr->kernel == NULL) {
      tr->chi_min = 0;
      tr->chi_max = ccl_comoving_radial_distance(cosmo, cosmo->spline_params.A_SPLINE_MIN, status);
    }
    else {
      int ichi;
      double w_max = fabs(w_w[0]);

      // Find maximum of radial kernel
      for (ichi=0; ichi < n_w; ichi++) {
        if (fabs(w_w[ichi]) >= w_max)
          w_max = fabs(w_w[ichi]);
      }

      // Multiply by fraction
      w_max *= CCL_FRAC_RELEVANT;

      // Initialize as the original edges in case we don't find an interval
      tr->chi_min = chi_w[0];
      tr->chi_max = chi_w[n_w-1];

      // Find minimum
      for (ichi=0; ichi < n_w; ichi++) {
        if (fabs(w_w[ichi]) >= w_max) {
          tr->chi_min = chi_w[ichi];
          break;
        }
      }

      // Find maximum
      for (ichi=n_w-1; ichi >= 0; ichi--) {
        if (fabs(w_w[ichi]) >= w_max) {
          tr->chi_max = chi_w[ichi];
          break;
        }
      }
    }
  }

  if (*status == 0) {
    if ((fka_arr != NULL) || (fk_arr != NULL) || (fa_arr != NULL)) {
      tr->transfer = ccl_f2d_t_new(
        na_ka,a_ka, // na, a_arr
        nk_ka,lk_ka, // nk, lk_arr
        fka_arr, // fka_arr
        fk_arr, // fk_arr
        fa_arr, // fa_arr
        is_factorizable, // is factorizable
        extrap_order_lok, // extrap_order_lok
        extrap_order_hik, // extrap_order_hik
        ccl_f2d_constantgrowth, // extrap_linear_growth
        is_fka_log, // is_fka_log
        NULL, // growth (function)
        1, // growth_factor_0 -> will assume constant transfer function
        0, // growth_exponent
        ccl_f2d_3, // interp_type
        status);
      if (tr->transfer == NULL)
        *status=CCL_ERROR_MEMORY;
    }
  }

  return tr;
}

void ccl_cl_tracer_t_free(ccl_cl_tracer_t *tr) {
  if (tr != NULL) {
    if (tr->transfer != NULL)
      ccl_f2d_t_free(tr->transfer);
    if (tr->kernel != NULL)
      ccl_f1d_t_free(tr->kernel);
    free(tr);
  }
}

double ccl_cl_tracer_t_get_f_ell(ccl_cl_tracer_t *tr, double ell, int *status) {
  if (tr != NULL) {
    if (tr->der_angles == 1)
      return ell*(ell+1.);
    else if (tr->der_angles == 2) {
      if (ell <= 1) // This is identically 0
        return 0;
      else if (ell <= 10) // Use full expression in this case
        return sqrt((ell+2)*(ell+1)*ell*(ell-1));
      else {
        double lp1h = ell+0.5;
        double lp1h2 = lp1h*lp1h;
        if (ell <= 1000)  // This is accurate to 5E-5 for l>10
          return lp1h2*(1-1.25/lp1h2);
        else // This is accurate to 1E-6 for l>1000
          return lp1h2;
      }
    }
    else
      return 1;
  }
  else
    return 1;
}

double ccl_cl_tracer_t_get_kernel(ccl_cl_tracer_t *tr, double chi, int *status) {
  if (tr != NULL) {
    if (tr->kernel != NULL)
      return ccl_f1d_t_eval(tr->kernel, chi);
    else
      return 1;
  }
  else
    return 1;
}

double ccl_cl_tracer_t_get_transfer(ccl_cl_tracer_t *tr,
                                    double lk, double a, int *status) {
  if (tr != NULL) {
    if (tr->transfer != NULL)
      return ccl_f2d_t_eval(tr->transfer, lk, a, NULL, status);
    else
      return 1;
  }
  else
    return 1;
}


#ifndef Two_Group_Diffusion_h
#define Two_Group_Diffusion_h

#include <vector>

// Trilinos includes
#include "Kokkos_DefaultNode.hpp"
#include "Teuchos_ParameterList.hpp"
#include "Teuchos_RCP.hpp"
#include "Tpetra_CrsGraph.hpp"
#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_MultiVector.hpp"
#include "Tpetra_Operator.hpp"
#include "Tpetra_RowMatrix.hpp"
#include "Tpetra_Vector.hpp"

// SCALE includes
#include "Nemesis/comm/global.hh"

// vendored includes
#include <gsl/gsl>

// enrico includes
#include "Assembly_Model.h"
#include "Neutronics_Solver.h"
#include "Two_Group_Cross_Sections.h"

namespace enrico {

//===========================================================================//
/*!
 * \class Two_Group_Diffusion
 * \brief Two-group, 3D neutron diffusion solver.
 */
//===========================================================================//

class Two_Group_Diffusion : public Neutronics_Solver {
public:
  enum BC_TYPE { VACUUM, REFLECT };
  enum FACE { LO_X, HI_X, LO_Y, HI_Y, LO_Z, HI_Z };

  //@{
  //! Typedefs
  using XS = Two_Group_Cross_Sections;
  using XS_Data = XS::XS_Data;
  using Vec_Dbl = std::vector<double>;
  using Vec_Int = std::vector<int>;
  using Vec_BC = std::vector<BC_TYPE>;
  using SP_Assembly = std::shared_ptr<Assembly_Model>;
  using Pin_Map = std::vector<Assembly_Model::PIN_TYPE>;
  using ST = double;
  using LO = int;
  using GO = int;
  using NODE = KokkosClassic::DefaultNode::DefaultNodeType;
  using MV = Tpetra::MultiVector<ST, LO, GO, NODE>;
  using VECTOR = Tpetra::Vector<ST, LO, GO, NODE>;
  using OP = Tpetra::Operator<ST, LO, GO, NODE>;
  using MAP = Tpetra::Map<LO, GO, NODE>;
  using MATRIX = Tpetra::CrsMatrix<ST, LO, GO, NODE>;
  using GRAPH = Tpetra::CrsGraph<LO, GO, NODE>;
  using RCP_PL = Teuchos::RCP<Teuchos::ParameterList>;
  //@}

private:
  // >>> DATA
  SP_Assembly d_assembly;

  Vec_Dbl d_dx;
  Vec_Dbl d_dy;
  Vec_Dbl d_dz;
  Pin_Map d_pin_map;
  XS d_xs;
  size_t d_Nx, d_Ny, d_Nz, d_num_cells;

  // Solver parameters
  double d_tol;
  int d_max_iters;

  // Boundary conditions
  Vec_BC d_bcs;

  // Data stored over single group
  Teuchos::RCP<MAP> d_map;
  Teuchos::RCP<GRAPH> d_graph;
  Teuchos::RCP<MATRIX> d_A_therm, d_A_fast;
  Teuchos::RCP<VECTOR> d_x_therm, d_x_fast;
  Teuchos::RCP<VECTOR> d_b_therm, d_b_fast;

public:
  // Constructor
  Two_Group_Diffusion(SP_Assembly assembly, RCP_PL params, const Vec_Dbl& dz);

  // Solve
  virtual void solve(const Vec_Dbl& temperatures,
                     const Vec_Dbl& densities,
                     Vec_Dbl& power) override;

private:
  int cellid(int ix, int iy, int iz)
  {
    Expects(ix >= 0);
    Expects(ix < d_Nx);
    Expects(iy >= 0);
    Expects(iy < d_Ny);
    Expects(iz >= 0);
    Expects(iz < d_Nz);
    int cell = ix + d_Nx * (iy + d_Ny * iz);
    Ensures(cell >= 0);
    Ensures(cell < d_num_cells);
    return cell;
  }

  void build_matrices(const std::vector<XS_Data>& xs_data);

  // Compute 2-norm of vector
  double vec_norm(const Teuchos::ArrayRCP<double>& vec)
  {
    double nrm = 0.0;
    for (auto val : vec)
      nrm += val * val;
    return std::sqrt(nrm);
  }

  // Apply scale factor to vector data
  void scale_vec(Teuchos::ArrayRCP<double>& vec, double factor)
  {
    for (auto& val : vec)
      val *= factor;
  }
};

//---------------------------------------------------------------------------//
} // end namespace enrico

//---------------------------------------------------------------------------//
#endif // Two_Group_Diffusion_h

//---------------------------------------------------------------------------//
// end of Two_Group_Diffusion.h
//---------------------------------------------------------------------------//


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

#include <mpi.h>

/* Pull in the CBLAS header. */
#ifdef _MACOSX
#include <Accelerate/Accelerate.h>
#else
#ifdef _ATLAS
#include <cblas.h>
#else
#include <gsl_cblas.h>
#endif /* _ATLAS */
#endif /* _MACOSX */

/* These headers are provided by ScaleME. */
#include "ScaleME.h"

#include "precision.h"

#include "mlfma.h"
#include "direct.h"
#include "itsolver.h"
#include "util.h"

int matvec (cplx *out, cplx *in, cplx *cur, int id) {
	long i, nelt = (long)fmaconf.numbases * (long)fmaconf.bspboxvol;

	/* Compute the contrast pressure. */
#pragma omp parallel for default(shared) private(i)
	for (i = 0; i < nelt; ++i) cur[i] = in[i] * fmaconf.contrast[i];

	/* Reset the direct-interaction buffer and compute
	 * the matrix-vector product for the Green's matrix. */
	clrdircache();
	ScaleME_applyParFMA (cur, out);

	if (!id) return 0;

	/* Add in the identity portion. */
#pragma omp parallel for default(shared) private(i)
	for (i = 0; i < nelt; ++i) out[i] = fmaconf.cellvol * in[i] - out[i];

	return 0;
}

int gmres (cplx *rhs, cplx *sol, int guess,
		int mit, real tol, int quiet, augspace *aug) {
	long j, nelt = (long)fmaconf.numbases * (long)fmaconf.bspboxvol, lwork;
	int i, rank, one = 1, mred = mit;
	cplx *h, *v, *mvp, *beta, *vp, *hp, *s, cr,
		cone = 1., czero = 0., *azp, *zp;
	real rhn, err, *c;

	MPI_Comm_rank (MPI_COMM_WORLD, &rank);

	/* Allocate space for all required complex vectors. */
	lwork = (mit + 1) * (mit + nelt + 1) + nelt + mit;
	v = calloc (lwork, sizeof(cplx));	/* The Krylov subspace. */
	beta = v + nelt * (mit + 1);		/* The least-squares RHS. */
	mvp = beta + mit + 1;			/* Buffer for matrix-vector product. */
	h = mvp + nelt;				/* The upper Hessenberg matrix. */
	s = h + (mit + 1) * mit;		/* Givens rotation sines. */

	/* Allocate space for the Givens rotation cosines. */
	c = malloc (mit * sizeof(real));

	/* Compute the norm of the RHS for residual scaling. */
	rhn = parnorm(rhs, nelt);

	/* Compute the initial matrix-vector product for the input guess. */
	if (guess) matvec (v, sol, mvp, 1);

	/* Subtract from the RHS to form the residual. */
#pragma omp parallel for default(shared) private(j)
	for (j = 0; j < nelt; ++j) v[j] = rhs[j] - v[j];

	/* Zero the initial guess if one wasn't provided. */
	if (!guess) memset (sol, 0, nelt * sizeof(cplx));

	/* Find the norm of the initial residual. */
	err = parnorm(v, nelt);

	/* Construct the initial Arnoldi vector by normalizing the residual. */
#pragma omp parallel for default(shared) private(j)
	for (j = 0; j < nelt; ++j) v[j] /= err;

	/* Construct the vector beta for the minimization problem. */
	beta[0] = err;

	/* Report the RRE. */
	err /= rhn;
	if (!rank && !quiet) printf ("True residual: %g\n", err);

	/* The reduced number of iterations, if the space is augmented. */
	if (aug) mred = mit - aug->ntot;

	for (i = 0; i < mit && err > tol; ++i) {
		/* Point to the working space for this iteration. */
		vp = v + i * nelt;
		hp = h + i * (mit + 1);

		/* Compute the next expansion of the Krylov space. */
		if (!aug || i < mred) matvec (vp + nelt, vp, mvp, 1);
		else {
			/* Update with the next augmented vector. */
			azp = aug->az + nelt *
				((aug->nmax + aug->start + mred - i) % aug->nmax);
			/* Use the augmented space. */
			memcpy (vp + nelt, azp, nelt * sizeof(cplx));
		}

		/* Perform modified Gram-Schmidt to orthogonalize the basis. */
		/* This also builds the Hessenberg matrix column, including
		 * the 2-norm of the next basis vector. */
		cmgs (vp + nelt, hp, v, nelt, i + 1);

		/* Watch for breakdown. */
		if (cabs(hp[i + 1]) < REAL_EPSILON) {
			++i;
			break;
		}

		/* Apply previous Givens rotations to the Hessenberg column. */
		for (j = 0; j < i; ++j)
			ROT (&one, hp + j, &one, hp + j + 1, &one, c + j, s + j);

		/* Compute the Givens rotation for the current iteration. */
		LARTG (hp + i, hp + i + 1, c + i, s + i, &cr);
		/* Apply the current Givens rotation to the Hessenberg column. */
		hp[i] = cr;
		hp[i + 1] = 0;
		/* Perform the rotation on the vector beta. */
		ROT (&one, beta + i, &one, beta + i + 1, &one, c + i, s + i);

		/* Estimate the RRE for this iteration. */
		err = cabs(beta[i + 1]) / rhn;
		if (!rank && !quiet) printf ("GMRES(%d): %g\n", i, err);

		/* Flush the output buffers. */
		fflush (stdout);
		fflush (stderr);
	}

	/* If there were any GMRES iterations, update the solution. */
	if (i > 0 && aug) {
		/* Compute the optimum solution in the Krylov basis. */
		cplx *ys;
		ys = malloc(i * sizeof(cplx));
		memcpy (ys, beta, i * sizeof(cplx));
		TRSV (CblasColMajor, CblasUpper, CblasNoTrans,
				CblasNonUnit, i, h, mit + 1, ys, 1);

		/* Compute the next Krylov vector in the augmented space. */
		aug->start = (aug->start + 1) % aug->nmax;
		if (aug->ntot < aug->nmax) ++(aug->ntot);
		azp = aug->az + nelt * aug->start;

		beta[i] = 0;
		for (j = i - 1; j >= 0; --j) {
			/* Invert the Givens rotations. */
			s[j] = -s[j];
			ROT (&one, beta + j, &one, beta + j + 1, &one, c + j, s + j);
		}
		GEMV (CblasColMajor, CblasNoTrans, nelt, i + 1,
				&cone, v, nelt, beta, 1, &czero, azp, 1);

		/* Overwrite the Krylov subspace with the augmented subspace.
		 * The start pointer has already been incremented! */
		for (j = mred; j < i; ++j) {
			zp = aug->z + nelt *
				((aug->nmax + aug->start + mred - j - 1) % aug->nmax);
			memcpy(v + j * nelt, zp, nelt * sizeof(cplx));
		}

		/* Compute the solution update. */
		zp = aug->z + nelt * aug->start;
		GEMV (CblasColMajor, CblasNoTrans, nelt, i,
				&cone, v, nelt, ys, 1, &czero, zp, 1);
		for (j = 0; j < nelt; ++j) sol[j] += zp[j];

		free(ys);
	} else if (i > 0) {
		/* Compute the minimizer of the least-squares problem. */
		TRSV (CblasColMajor, CblasUpper, CblasNoTrans,
				CblasNonUnit, i, h, mit + 1, beta, 1);

		/* Compute the solution update in place. */
		GEMV (CblasColMajor, CblasNoTrans, nelt, i,
				&cone, v, nelt, beta, 1, &cone, sol, 1);
	}

	free (v);
	free (c);

	return i;
}

int bicgstab (cplx *rhs, cplx *sol,
		int guess, int mit, real tol, int quiet) {
	long j, nelt = (long)fmaconf.numbases * (long)fmaconf.bspboxvol;
	int i, rank;
	cplx *r, *rhat, *v, *p, *mvp, *t;
	cplx rho, alpha, omega, beta;
	real err, rhn;

	MPI_Comm_rank (MPI_COMM_WORLD, &rank);

	rho = alpha = omega = 1.;

	/* Allocate and zero the work arrays. */
	r = calloc (6L * nelt, sizeof(cplx));
	rhat = r + nelt;
	v = rhat + nelt;
	p = v + nelt;
	t = p + nelt;
	mvp = t + nelt;

	/* Compute the norm of the right-hand side for residual scaling. */
	rhn = parnorm(rhs, nelt);

	/* Compute the inital matrix-vector product for the input guess. */
	if (guess) matvec (r, sol, mvp, 1);

	/* Subtract from the RHS to form the residual. */
#pragma omp parallel for default(shared) private(j)
	for (j = 0; j < nelt; ++j) r[j] = rhs[j] - r[j];

	if (!guess) memset (sol, 0, nelt * sizeof(cplx));

	/* Copy the initial residual as the test vector. */
	memcpy (rhat, r, nelt * sizeof(cplx));

	/* Find the norm of the initial residual. */
	err = parnorm(r, nelt) / rhn;
	if (!rank && !quiet) printf ("True residual: %g\n", err);

	/* Run iterations until convergence or the maximum is reached. */
	for (i = 0; i < mit && err > tol; ++i) {
		/* Pre-compute portion of beta from previous iteration. */
		beta = alpha / (rho * omega);
		/* Compute rho for this iteration. */
		rho = pardot (rhat, r, nelt);
		/* Include the missing factor in beta. */
		beta *= rho;

		/* Update the search vector. */
#pragma omp parallel for default(shared) private(j)
		for (j = 0; j < nelt; ++j)
			p[j] = r[j] + beta * (p[j] - omega * v[j]);

		/* Compute the first search step, v = A * p. */
		matvec (v, p, mvp, 1);

		/* Compute the next alpha. */
		alpha = rho / pardot (rhat, v, nelt);

#pragma omp parallel for default(shared) private(j)
		for (j = 0; j < nelt; ++j) {
			/* Update the solution vector. */
			sol[j] += alpha * p[j];
			/* Update the residual vector. */
			r[j] -= alpha * v[j];
		}

		/* Compute the scaled residual norm and stop if convergence
		 * has been achieved. */
		err = parnorm(r, nelt) / rhn;
		if (!rank && !quiet) printf ("BiCG-STAB(%0.1f): %g\n", 0.5 + i, err);

		/* Flush the output buffers. */
		fflush (stdout);
		fflush (stderr);

		if (err < tol) break;

		/* Compute the next search step, t = A * r. */
		matvec (t, r, mvp, 1);

		/* Compute the update direction. */
		omega = pardot (t, r, nelt) / pardot (t, t, nelt);

		/* Update both the residual and the solution guess. */
#pragma omp parallel for default(shared) private(j)
		for (j = 0; j < nelt; ++j) {
			/* Update the solution vector. */
			sol[j] += omega * r[j];
			/* Update the residual vector. */
			r[j] -= omega * t[j];
		}

		/* Compute the scaled residual norm. */
		err = parnorm(r, nelt) / rhn;
		if (!rank && !quiet) printf ("BiCG-STAB(%d): %g\n", i + 1, err);

		fflush (stdout);
		fflush (stderr);
	}

	free (r);
	return i;
}


#ifndef TETRA_DOS_ELHAMILTONIAN_H
#define TETRA_DOS_ELHAMILTONIAN_H

#include <gsl/gsl_math.h>
#include <gsl/gsl_sf_trig.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include "environment.h"

void ElHamiltonian_Recip(Environment *env, double k[3], gsl_matrix_complex *H);

void ElHamiltonian(Environment *env, double k[3], gsl_matrix_complex *H);

gsl_complex EpsilonAE(Environment *env, double k[3]);

gsl_complex EpsilonBE(Environment *env, double k[3]);

gsl_complex EpsilonAO(Environment *env, double k[3]);

gsl_complex EpsilonBO(Environment *env, double k[3]);

#endif // TETRA_DOS_ELHAMILTONIAN_H


#include "cimple_c_from_py.h"
#include "setoper.h"
#include "cimple_safe_mode.h"
#include <cdd.h>
#include <gsl/gsl_matrix.h>

int main(){

    dd_set_global_constants();
    // Initialize state:
    system_dynamics *s_dyn;
    cost_function *f_cost;
    discrete_dynamics *d_dyn;
    current_state *now;

    system_alloc(&now, &s_dyn, &f_cost, &d_dyn);
    system_init(now, s_dyn, f_cost, d_dyn);

    double sec = 2;
    ACT(4, now, d_dyn, s_dyn, f_cost, sec);

    system_dynamics_free(s_dyn);
    discrete_dynamics_free(d_dyn);
    cost_function_free(f_cost);
    state_free(now);

    dd_free_global_constants();
    return 0;
}


/* TwoAdvSelfSims.c

Simulation calculating fixation probability of second beneficial allele, given existing
ben allele at initial frequency p. For use in the study "Linkage and the limits to natural
selection in partially selfing populations".

Simulation uses routines found with the GNU Scientific Library (GSL)
(http://www.gnu.org/software/gsl/)
Since GSL is distributed under the GNU General Public License
(http://www.gnu.org/copyleft/gpl.html), you must download it
separately from this file.

This program can be compiled with e.g. GCC using a command like:
gcc TwoAdvSelfSims -lm -lgsl -lgslcblas -I/usr/local/include -L/usr/local/lib TwoAdvSelfSims.c

Then run by executing:
./TwoAdvSelfSims N self rec ha sa hb sb p reps
Where:
- N is the population size
- self is the rate of self-fertilisation
- rec is recombination rate
- ha, hb is dominance at the original, introduced beneficial allele
- sa, sb is selection coefficient of the original, introduced beneficial allele
- p is the initial frequency of the first beneficial allele (when the second is introduced)
- reps is how many times the second allele should FIX before simulation stops
(the number of actual runs is greater due to stochastic loss of second allele)

Note that haplotypes are defined as:
x1 = ab
x2 = Ab
x3 = aB
x4 = AB

Genotypes defined as:
g11 = g1 = ab/ab
g12 = g2 = Ab/ab
g13 = g3 = aB/ab
g14 = g4 = AB/ab
g22 = g5 = Ab/Ab
g23 = g6 = Ab/aB
g24 = g7 = Ab/AB
g33 = g8 = aB/aB
g34 = g9 = aB/AB
g44 = g10 = AB/AB

Output files are the parameters;
followed by number of times each haplotype fixed;
followed by average total generations elapsed in each case;
Then total number of simulations ran;
Then fixation prob of allele, both unscaled and scaled to unlinked case,
along with 95% CI intervals for the latter case;
then number of allele fixations.

Note that 'fixation' DIFFERS depending on the inputs of sa, sb.
If sa >= sb (interference case) then 'fixation' counts as fixation of second allele on any genetic background.
If sa < sb (replacement case) then 'fixation' only considers fixation of second allele with neutral haplotype
(I.e. where the 'less fit' neutral allele at locus A fixes, instead of selected allele).

*/

/* Preprocessor statements */
#include <stdio.h>
#include <time.h>
#include <math.h>
#include <stddef.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_cdf.h>

/* Function prototypes */
void geninit(double *geninit, double FIS, const gsl_rng *r);
void selection(double *geninit);
void reproduction(double *geninit);
unsigned int hcheck(double *geninit, double *haps, unsigned int *hf, unsigned int stype);

/* Global variable declaration */
unsigned int N = 0;		/* Pop size */
double rec = 0;			/* Recombination rate */
double self = 0;		/* Rate of self-fertilisation */
double ha = 0;			/* Dominance of site A */
double sa = 0;			/* Fitness of site A */
double hb = 0;			/* Dominance of site B */
double sb = 0;			/* Fitness of site B */
double pee = 0;			/* Freq of initial sweep */

/* Main program */
int main(int argc, char *argv[]){
	unsigned int i; 				/* A counter */
	unsigned int reps;				/* Length of simulation (no. of introductions of neutral site) */
	unsigned int stype = 0;			/* What type of sim (replacement or hitch-hiking)? */
	unsigned int nsfix = 0;			/* sims where target type fixed */
	unsigned int nstot = 0;			/* total sims ran */
	unsigned int gens = 0;			/* Number gens elapsed */
	unsigned int isfin = 0;			/* Is sim finished? */
	unsigned int hf = 0;			/* The hap that fixed */
	double pf = 0;					/* Overall fix prob */
	double FIS = 0;					/* Wright's FIS */
	double StdFix = 0;				/* Standard Fixation prob if unlinked */
	double citop = 0;
	double cibot = 0;
	double nsfix2 = 0;
	double nstot2 = 0;
	char selfchar[10];
	char recchar[15];
	char hchar[10];
	char pchar[10];
	char fname[64];
	FILE *ofp_tr;

	/* GSL random number definitions */
	const gsl_rng_type * T;
	gsl_rng * r;

	/* This reads in data from command line. */
	if(argc != 10){
		fprintf(stderr,"Invalid number of input values.\n");
		exit(1);
	}
	N = strtod(argv[1],NULL);
	self = strtod(argv[2],NULL);
	rec = strtod(argv[3],NULL);
	ha = strtod(argv[4],NULL);
	sa = strtod(argv[5],NULL);
	hb = strtod(argv[6],NULL);
	sb = strtod(argv[7],NULL);
	pee = strtod(argv[8],NULL);
	reps = strtod(argv[9],NULL);

	if(sa >= sb){
		stype = 0;
	}else if(sa < sb){
		stype = 1;
	}
	FIS = self/(2.0-self);

	/* Arrays definition and memory assignment */
	double *genotype = calloc(10,sizeof(double));				/* Genotype frequencies */
	unsigned int *gensamp = calloc(10,sizeof(unsigned int));	/* New population samples */
	double *haps = calloc(4,sizeof(double));					/* Haplotypes */
	unsigned int *pfix = calloc(4,sizeof(unsigned int));		/* Haplotypes that fix */
	unsigned int *tfix = calloc(4,sizeof(unsigned int));		/* time that haps fix */

	/* create a generator chosen by the
    environment variable GSL_RNG_TYPE */

	gsl_rng_env_setup();
	if (!getenv("GSL_RNG_SEED")) gsl_rng_default_seed = time(0);
	T = gsl_rng_default;
	r = gsl_rng_alloc(T);

	nsfix = 0;
	nstot = 0;

    while(nsfix < reps){
    	nstot++;

		/* Initialising genotypes */
		geninit(genotype,FIS,r);
		gens = 0;
		isfin = 0;
		hf = 0;

		while(isfin == 0){
    		/* Selection routine */
    		selection(genotype);

	    	/* Reproduction routine */
    		reproduction(genotype);

    		/* Sampling based on new frequencies */
	       	gsl_ran_multinomial(r,10,N,genotype,gensamp);
       		for(i = 0; i < 10; i++){
    	   		*(genotype + i) = (*(gensamp + i))/(1.0*N);
	       	}

	       	gens++;
	       	isfin = hcheck(genotype, haps, &hf, stype);
       	}

       	if(isfin == 1){
       		(*(pfix + hf))++;
    	  	(*(tfix + hf)) += gens;
	       	if(stype == 0){
       			nsfix += (*(haps + 2) + *(haps + 3));
       		}else if(stype == 1){
    	   		nsfix += (*(haps + 2));
	       	}
       	}

	}	/* End of simulation */

	pf = nsfix/(1.0*nstot);
	StdFix = 2*sb*((hb+FIS-hb*FIS)/(1+FIS));		/* Fix prob of new allele if unlinked */

	nstot2 = nstot + 3.84;
	nsfix2 = (1.0/(nstot2))*(nsfix + (3.84/2.0));
	citop = nsfix2 + 1.96*sqrt((1.0/(nstot2))*nsfix2*(1.0-nsfix2));
	cibot = nsfix2 - 1.96*sqrt((1.0/(nstot2))*nsfix2*(1.0-nsfix2));
	citop = citop/(1.0*StdFix);
	cibot = cibot/(1.0*StdFix);

	/* Printing solutions to file */
	/* First, converting values to strings */
    sprintf(selfchar, "%0.2lf",self);
    sprintf(recchar, "%0.7lf",rec);
    sprintf(hchar, "%0.3lf",ha);
    sprintf(pchar, "%0.5lf",pee);

	strcpy(fname,"sim_self");
	strcat(fname,selfchar);
	strcat(fname,"rec");
	strcat(fname,recchar);
	strcat(fname,"h");
	strcat(fname,hchar);
	strcat(fname,"p");
	strcat(fname,pchar);
	strcat(fname,".sim");

	ofp_tr = fopen(fname,"a+");
	fprintf(ofp_tr,"%d %lf %lf %lf %lf %lf %lf %lf ",N,self,rec,ha,sa,hb,sb,pee);
	for(i = 0; i < 4; i++){
		fprintf(ofp_tr,"%d ",*(pfix + i));
	}
	for(i = 0; i < 4; i++){
		fprintf(ofp_tr,"%lf ",((*(tfix + i)))/(1.0*(*(pfix + i))));
	}
	fprintf(ofp_tr,"%d %lf %lf %lf %lf %d\n",nstot,pf,pf/(1.0*StdFix),citop,cibot,reps);
	fclose(ofp_tr);

	/* Freeing memory and wrapping up */
 	gsl_rng_free(r);
 	free(tfix);
 	free(pfix);
 	free(haps);
 	free(gensamp);
	free(genotype);
	return 0;
}

/* Initialising genotypes */
void geninit(double *geninit, double FIS, const gsl_rng *r){

	unsigned int htype = 0;		/* Type of het assignment */

	double *ptype = calloc(3,sizeof(double));
	unsigned int *ctype = calloc(3,sizeof(double));

	/* Routine to determine initial background of second mutant, given first is at frequency p. */
	*(ptype + 0) = (1-pee)*(1-pee) + FIS*pee*(1-pee);
	*(ptype + 1) = 2*pee*(1-pee)*(1-FIS);
	*(ptype + 2) = pee*pee + FIS*pee*(1-pee);

	/* First initialise baseline freqs */
	*(geninit + 0) = *(ptype + 0);
	*(geninit + 1) = *(ptype + 1);
	*(geninit + 2) = 0;
	*(geninit + 3) = 0;
	*(geninit + 4) = *(ptype + 2);
	*(geninit + 5) = 0;
	*(geninit + 6) = 0;
	*(geninit + 7) = 0;
	*(geninit + 8) = 0;
	*(geninit + 9) = 0;

	/* Then decide where to add new mutant */
	gsl_ran_multinomial(r,3,1,ptype,ctype);
	if(*(ctype + 0) == 1){
		*(geninit + 0) -= 1/(1.0*N);
		*(geninit + 2) += 1/(1.0*N);
	}else if(*(ctype + 1) == 1){
		*(geninit + 1) -= 1/(1.0*N);
	  	htype = gsl_ran_bernoulli(r,0.5);
	  	if(htype == 0){
	  		*(geninit + 5) += 1/(1.0*N);
	  	}else if(htype == 1){
	  		*(geninit + 3) += 1/(1.0*N);
	  	}
	}else if(*(ctype + 2) == 1){
		*(geninit + 4) -= 1/(1.0*N);
		*(geninit + 6) += 1/(1.0*N);
	}

	free(ctype);
	free(ptype);

}	/* End of gen initiation routine */

/* Selection routine */
void selection(double *geninit){
	/* Fitness of each genotype */
	double W11, W12, W13, W14, W22, W23, W24, W33, W34, W44;
	double Wmean;				/* Mean fitness */

	W11 = 1;
	W12 = 1 + ha*sa;
	W13 = 1 + hb*sb;
	W14 = 1 + ha*sa + hb*sb;
	W22 = 1 + sa;
	W23 = 1 + ha*sa + hb*sb;
	W24 = 1 + sa + hb*sb;
	W33 = 1 + sb;
	W34 = 1 + ha*sa + sb;
	W44 = 1 + sa + sb;

	/* Mean fitness calculation */
	Wmean = ((*(geninit + 0))*W11) + ((*(geninit + 1))*W12) + ((*(geninit + 2))*W13) + ((*(geninit + 3))*W14) + ((*(geninit + 4))*W22) + ((*(geninit + 5))*W23) + ((*(geninit + 6))*W24) + ((*(geninit + 7))*W33) + ((*(geninit + 8))*W34) + ((*(geninit + 9))*W44);

	/* Changing frequencies by selection */
	*(geninit + 0) = ((*(geninit + 0))*W11)/Wmean;
	*(geninit + 1) = ((*(geninit + 1))*W12)/Wmean;
	*(geninit + 2) = ((*(geninit + 2))*W13)/Wmean;
	*(geninit + 3) = ((*(geninit + 3))*W14)/Wmean;
	*(geninit + 4) = ((*(geninit + 4))*W22)/Wmean;
	*(geninit + 5) = ((*(geninit + 5))*W23)/Wmean;
	*(geninit + 6) = ((*(geninit + 6))*W24)/Wmean;
	*(geninit + 7) = ((*(geninit + 7))*W33)/Wmean;
	*(geninit + 8) = ((*(geninit + 8))*W34)/Wmean;
	*(geninit + 9) = ((*(geninit + 9))*W44)/Wmean;

}	/* End of selection routine */

/* Reproduction routine */
void reproduction(double *geninit){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	/* Haplotypes */
	double x1, x2, x3, x4;

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Baseline change in haplotype frequencies */
	x1 = g11s + (g12s + g13s + g14s)/2.0 - ((g14s - g23s)*rec)/2.0;
	x2 = g22s + (g12s + g23s + g24s)/2.0 + ((g14s - g23s)*rec)/2.0;
	x3 = g33s + (g13s + g23s + g34s)/2.0 + ((g14s - g23s)*rec)/2.0;
	x4 = g44s + (g14s + g24s + g34s)/2.0 - ((g14s - g23s)*rec)/2.0;

	/* Change in SEXUAL frequencies (both outcrossing and selfing) */
	*(geninit + 0) = (g11s + (g12s + g13s + g14s*pow((1 - rec),2) + g23s*pow(rec,2))/4.0)*self + (1 - self)*pow(x1,2);
	*(geninit + 4) = (g22s + (g12s + g24s + g23s*pow((1 - rec),2) + g14s*pow(rec,2))/4.0)*self + (1 - self)*pow(x2,2);
	*(geninit + 7) = (g33s + (g13s + g34s + g23s*pow((1 - rec),2) + g14s*pow(rec,2))/4.0)*self + (1 - self)*pow(x3,2);
	*(geninit + 9) = (g44s + (g24s + g34s + g14s*pow((1 - rec),2) + g23s*pow(rec,2))/4.0)*self + (1 - self)*pow(x4,2);
	*(geninit + 1) = ((g12s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x1*x2;
	*(geninit + 2) = ((g13s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x1*x3;
	*(geninit + 3) = ((g14s*pow((1 - rec),2) + g23s*pow(rec,2))*self)/2.0 + 2.0*(1 - self)*x1*x4;
	*(geninit + 5) = ((g23s*pow((1 - rec),2) + g14s*pow(rec,2))*self)/2.0 + 2.0*(1 - self)*x2*x3;
	*(geninit + 6) = ((g24s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x2*x4;
	*(geninit + 8) = ((g34s + (g14s + g23s)*(1 - rec)*rec)*self)/2.0 + 2.0*(1 - self)*x3*x4;

}	/* End of reproduction routine */

/* Has any allele fixed or not? */
unsigned int hcheck(double *geninit, double *haps, unsigned int *hf, unsigned int stype){
	/* Fed-in genotype frequencies (for ease of programming) */
	double g11s, g12s, g13s, g14s, g22s, g23s, g24s, g33s, g34s, g44s;
	unsigned int retval = 0;

	/* Initial definition of genotypes */
	g11s = *(geninit + 0);
	g12s = *(geninit + 1);
	g13s = *(geninit + 2);
	g14s = *(geninit + 3);
	g22s = *(geninit + 4);
	g23s = *(geninit + 5);
	g24s = *(geninit + 6);
	g33s = *(geninit + 7);
	g34s = *(geninit + 8);
	g44s = *(geninit + 9);

	/* Calculation of haplotypes */
	*(haps + 0) = g11s + (g12s + g13s + g14s)/2.0;
	*(haps + 1) = g22s + (g12s + g23s + g24s)/2.0;
	*(haps + 2) = g33s + (g13s + g23s + g34s)/2.0;
	*(haps + 3) = g44s + (g14s + g24s + g34s)/2.0;

/* 	printf("Haps are %lf %lf %lf %lf\n",*(haps + 0),*(haps + 1),*(haps + 2),*(haps + 3));*/

	if(*(haps + 0) == 1){
		retval = 1;
		*hf = 0;
	}
	else if(*(haps + 1) == 1){
		retval = 1;
		*hf = 1;
	}
	else if(*(haps + 2) == 1){
		retval = 1;
		*hf = 2;
	}
	else if(*(haps + 3) == 1){
		retval = 1;
		*hf = 3;
	}else if(stype == 1){
		if( (*(haps + 2) + *(haps + 3) ) == 0 ){
			retval = 2;
			*hf = 0;
		}
	}

	return retval;

}	/* End of hap check routine */

/* End of program */


/**
 *
 * @file testing_zgemm.c
 *
 *  PLASMA testing routines
 *  PLASMA is a software package provided by Univ. of Tennessee,
 *  Univ. of California Berkeley and Univ. of Colorado Denver
 *
 * @version 2.6.0
 * @author Emmanuel Agullo
 * @author Mathieu Faverge
 * @date 2010-11-15
 * @precisions normal z -> c d s
 *
 **/
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>

#include <plasma.h>
#include <cblas.h>
#include <lapacke.h>
#include <core_blas.h>
#include "testing_zmain.h"

#undef REAL
#define COMPLEX

static int check_solution(PLASMA_enum transA, PLASMA_enum transB, int M, int N, int K,
                          PLASMA_Complex64_t alpha, PLASMA_Complex64_t *A, int LDA,
                          PLASMA_Complex64_t *B, int LDB,
                          PLASMA_Complex64_t beta, PLASMA_Complex64_t *Cref, PLASMA_Complex64_t *Cplasma, int LDC);

int testing_zgemm(int argc, char **argv)
{
    /* Check for number of arguments*/
    if ( argc != 8) {
        USAGE("GEMM", "alpha beta M N K LDA LDB LDC",
              "   - alpha  : alpha coefficient\n"
              "   - beta   : beta coefficient\n"
              "   - M      : number of rows of matrices A and C\n"
              "   - N      : number of columns of matrices B and C\n"
              "   - K      : number of columns of matrix A / number of rows of matrix B\n"
              "   - LDA    : leading dimension of matrix A\n"
              "   - LDB    : leading dimension of matrix B\n"
              "   - LDC    : leading dimension of matrix C\n");
        return -1;
    }

    PLASMA_Complex64_t alpha = (PLASMA_Complex64_t) atol(argv[0]);
    PLASMA_Complex64_t beta = (PLASMA_Complex64_t) atol(argv[1]);
    int M     = atoi(argv[2]);
    int N     = atoi(argv[3]);
    int K     = atoi(argv[4]);
    int LDA   = atoi(argv[5]);
    int LDB   = atoi(argv[6]);
    int LDC   = atoi(argv[7]);

    double eps;
    int info_solution;
    int i, j, ta, tb;
    int LDAxK = LDA*max(M,K);
    int LDBxN = LDB*max(K,N);
    int LDCxN = LDC*N;

    PLASMA_Complex64_t *A      = (PLASMA_Complex64_t *)malloc(LDAxK*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *B      = (PLASMA_Complex64_t *)malloc(LDBxN*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *C      = (PLASMA_Complex64_t *)malloc(LDCxN*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *Cinit  = (PLASMA_Complex64_t *)malloc(LDCxN*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *Cfinal = (PLASMA_Complex64_t *)malloc(LDCxN*sizeof(PLASMA_Complex64_t));

    /* Check if unable to allocate memory */
    if ((!A)||(!B)||(!Cinit)||(!Cfinal)){
        printf("Out of Memory \n ");
        return -2;
    }

    eps = LAPACKE_dlamch_work('e');

    printf("\n");
    printf("------ TESTS FOR PLASMA ZGEMM ROUTINE -------  \n");
    printf("            Size of the Matrix %d by %d\n", M, N);
    printf("\n");
    printf(" The matrix A is randomly generated for each test.\n");
    printf("============\n");
    printf(" The relative machine precision (eps) is to be %e \n",eps);
    printf(" Computational tests pass if scaled residuals are less than 10.\n");

    /*----------------------------------------------------------
     *  TESTING ZGEMM
     */

    /* Initialize A, B, C */
    LAPACKE_zlarnv_work(IONE, ISEED, LDAxK, A);
    LAPACKE_zlarnv_work(IONE, ISEED, LDBxN, B);
    LAPACKE_zlarnv_work(IONE, ISEED, LDCxN, C);

#ifdef COMPLEX
    for (ta=0; ta<3; ta++) {
        for (tb=0; tb<3; tb++) {
#else
    for (ta=0; ta<2; ta++) {
        for (tb=0; tb<2; tb++) {
#endif
            for ( i = 0; i < M; i++)
                for (  j = 0; j < N; j++)
                    Cinit[LDC*j+i] = C[LDC*j+i];
            for ( i = 0; i < M; i++)
                for (  j = 0; j < N; j++)
                    Cfinal[LDC*j+i] = C[LDC*j+i];

            /* PLASMA ZGEMM */
            PLASMA_zgemm(trans[ta], trans[tb], M, N, K, alpha, A, LDA, B, LDB, beta, Cfinal, LDC);

            /* Check the solution */
            info_solution = check_solution(trans[ta], trans[tb], M, N, K,
                                           alpha, A, LDA, B, LDB, beta, Cinit, Cfinal, LDC);

            if (info_solution == 0) {
                printf("***************************************************\n");
                printf(" ---- TESTING ZGEMM (%s, %s) ............... PASSED !\n", transstr[ta], transstr[tb]);
                printf("***************************************************\n");
            }
            else {
                printf("************************************************\n");
                printf(" - TESTING ZGEMM (%s, %s) ... FAILED !\n", transstr[ta], transstr[tb]);
                printf("************************************************\n");
            }
        }
    }
#ifdef _UNUSED_
    }}
#endif
    free(A); free(B); free(C);
    free(Cinit); free(Cfinal);

    return 0;
}

/*--------------------------------------------------------------
 * Check the solution
 */

static int check_solution(PLASMA_enum transA, PLASMA_enum transB, int M, int N, int K,
                          PLASMA_Complex64_t alpha, PLASMA_Complex64_t *A, int LDA,
                          PLASMA_Complex64_t *B, int LDB,
                          PLASMA_Complex64_t beta, PLASMA_Complex64_t *Cref, PLASMA_Complex64_t *Cplasma, int LDC)
{
    int info_solution;
    double Anorm, Bnorm, Cinitnorm, Cplasmanorm, Clapacknorm, Rnorm, result;
    double eps;
    PLASMA_Complex64_t beta_const;

    double *work = (double *)malloc(max(K,max(M, N))* sizeof(double));
    int Am, An, Bm, Bn;

    beta_const  = -1.0;

    if (transA == PlasmaNoTrans) {
        Am = M; An = K;
    } else {
        Am = K; An = M;
    }
    if (transB == PlasmaNoTrans) {
        Bm = K; Bn = N;
    } else {
        Bm = N; Bn = K;
    }

    Anorm       = LAPACKE_zlange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), Am, An, A,       LDA, work);
    Bnorm       = LAPACKE_zlange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), Bm, Bn, B,       LDB, work);
    Cinitnorm   = LAPACKE_zlange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), M,  N,  Cref,    LDC, work);
    Cplasmanorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), M,  N,  Cplasma, LDC, work);

    cblas_zgemm(CblasColMajor, (CBLAS_TRANSPOSE)transA, (CBLAS_TRANSPOSE)transB, M, N, K,
                CBLAS_SADDR(alpha), A, LDA, B, LDB, CBLAS_SADDR(beta), Cref, LDC);

    Clapacknorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), M, N, Cref, LDC, work);

    cblas_zaxpy(LDC * N, CBLAS_SADDR(beta_const), Cplasma, 1, Cref, 1);

    Rnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, lapack_const(PlasmaInfNorm), M, N, Cref, LDC, work);

    eps = LAPACKE_dlamch_work('e');

    printf("Rnorm %e, Anorm %e, Bnorm %e, Cinitnorm %e, Cplasmanorm %e, Clapacknorm %e\n",
           Rnorm, Anorm, Bnorm, Cinitnorm, Cplasmanorm, Clapacknorm);

    result = Rnorm / ((Anorm + Bnorm + Cinitnorm) * N * eps);
    printf("============\n");
    printf("Checking the norm of the difference against reference ZGEMM \n");
    printf("-- ||Cplasma - Clapack||_oo/((||A||_oo+||B||_oo+||C||_oo).N.eps) = %e \n",
           result);

    if (  isnan(Rnorm) || isinf(Rnorm) || isnan(result) || isinf(result) || (result > 10.0) ) {
         printf("-- The solution is suspicious ! \n");
         info_solution = 1;
    }
    else {
         printf("-- The solution is CORRECT ! \n");
         info_solution= 0 ;
    }

    free(work);

    return info_solution;
}


#ifndef _PARAMETERS_H_
#define _PARAMETERS_H_
#include <string>
#include <fstream>
#include <iostream>
#include <cstdlib>
#include "../Headers/voxel.h"
#include <math.h>
#include <gsl/gsl_rng.h>

using namespace std;

class Parameters{
//Atributs
	public:
	int it; //Contador
	int nt;	// Nombre de passos total
   	int ntout; // Nombre de pasos per l'output
	int neq; //Number of equations
	float dt; // Pas de temps
	float t;

	string Dir_Output;
    string Dir_Input;
    string NumPar_Dat;

	string Data_Dat;
	string Init_Cond_Dat;
   	int nNodes;
   	int nNodes_inh;
    string Kex_Dat;
    string Kin_Dat;

	float ** Kex; //Coupling matrix for excitatory population
	float ** Kin; //Coupling matrix for inhibitory population
	voxel * V; //Voxel

	ofstream fout_data;
	ofstream fout_init_cond;

        public:
        Parameters();
        ~Parameters();

	//Mètodes

	void Read_Kex_Dat(void);
	void Read_Kin_Dat(void);
	void RefreshDelay(void);
	void Read_NumPar_Dat(void);

	void WriteData();
	void WriteInitialConditions(int);
	void openfiles();
	void closefiles();

        //private:

};

#endif


#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <omp.h>
#include <timerlib.h>
#ifndef NOBLAS
#ifdef MKL
#include <mkl_cblas.h>
#else
#include <cblas.h>
#endif
#endif
#include "gdrdgemm.h"

#define NMAT 2048
void gdr_check_and_restart(double a[][NMAT],
			   double b[][NMAT],
			   double c[][NMAT])
{
    int try =0;
    static int initialized = 0;
    if (initialized) return;
    while(1){
	int i,j;
	for(i=0;i<NMAT;i++){
	    for(j=0;j<NMAT;j++){
		a[i][j]=0;
		b[i][j]=i*NMAT+j;
		c[i][j]=0;
	    }
	}
	for(i=0;i<NMAT;i++)a[i][i]=1;
	cblas_dgemm( CblasRowMajor, CblasNoTrans, CblasNoTrans,
		     NMAT,NMAT, NMAT, 1.0, a, NMAT, b, NMAT, 0.0, c, NMAT);
	cblas_dgemm( CblasRowMajor, CblasNoTrans, CblasNoTrans,
		     NMAT,NMAT, NMAT, 1.0, a, NMAT, b, NMAT, 0.0, c, NMAT);
	mygdrdgemm(NMAT, NMAT, NMAT, 1.0, (double*)a, NMAT,
		   (double*)b, NMAT, 0.0, (double*) c, NMAT);
	int err = 0;
	for(i=0;i<NMAT;i++){
	    for(j=0;j<NMAT;j++){
		if (b[i][j] != c[i][j]){
		    err ++;
		}
	    }
	}
	if (err == 0){
	    fprintf(stderr,"gdr_check_and_restart passed %d\n", try);
	    initialized=1;
	    return;
	}
	try++;
	fprintf(stderr, "gdr_check_and_restart, err=%d try=%d\n", err, try);
	gdr_free();
	gdr_init();
    }
}

void gdrblas_dgemm
(
#ifndef MKL
   const enum CBLAS_ORDER             ORDER,
   const enum CBLAS_TRANSPOSE             TRANSA,
   const enum CBLAS_TRANSPOSE             TRANSB,
#else
   const CBLAS_ORDER             ORDER,
   const CBLAS_TRANSPOSE             TRANSA,
   const CBLAS_TRANSPOSE             TRANSB,
#endif
   const int                        M,
   const int                        N,
   const int                        K,
   const double                     ALPHA,
   const double *                   A,
   const int                        LDA,
   const double *                   B,
   const int                        LDB,
   const double                     BETA,
   double *                         C,
   const int                        LDC
)
{
    int NOTA, NOTB,gdrdoneflag;
    double alpha = ALPHA, beta = BETA;
    int F77M=M, F77N=N, F77K=K, F77lda=LDA, F77ldb=LDB, F77ldc=LDC;

    if(      TRANSA == CblasNoTrans ){
	NOTA = 1;
    }else{
	NOTA = 0;
    }
    if(      TRANSB == CblasNoTrans ){
	NOTB = 1;
    }else{
	NOTB = 0;
    }

    if( ORDER == CblasColMajor ){
	gdrdoneflag=0;
	gdr_dgemm_(&NOTA,&NOTB,&F77M,&F77N,&F77K, &alpha, &beta,
		   &F77lda, &F77ldb, &F77ldc,A,B,C, &gdrdoneflag);
    }  else  {
	gdrdoneflag=0;
	gdr_dgemm_(&NOTA,&NOTB,&F77N,&F77M,&F77K, &alpha, &beta,
		   &F77ldb, &F77lda, &F77ldc,B,A,C, &gdrdoneflag);
    }

    if(gdrdoneflag!=1){
	cblas_dgemm(ORDER, TRANSA, TRANSB,
		    M, N, K, ALPHA, A, LDA, B, LDB, BETA, C, LDC);
    }
    /*
     * End of HPL_dgemm
     */
}
void dumpcmat(int m, int n, int nc, double c[][nc])
{
    static int callcount = 0;
    static FILE* fid;
    if (callcount == 0){
	fid = fopen("/tmp/matdata", "w");
    }
    callcount ++;
    if(callcount < 8){
	fprintf(fid,"\nPrint CMAT callcount=%d\n",callcount);
	int i, j;
	for(i=0;i<m;i++){
	    fprintf(fid,"\ni=%d\n", i);
	    for(j=0;j<n;j++){
		if ((j%8)==0 )fprintf(fid,"\n%5d:", j);
		fprintf(fid," %20.12e",c[i][j]);
	    }
	}
    }
}
double touchcmat(int m, int n, int nc, double c[][nc])
{
    double sum=0;
    int i, j;
    for(i=0;i<m;i++){
	for(j=0;j<n;j++){
	    sum += c[i][j]*c[i][j];
	}
    }
    return sum;
}

double ssum = 0.0;

void mygdrdgemm(int m,
		int n,
		int k,
		double alpha,
		double * a,
		int na,
		double * b,
		int nb,
		double beta,
		double * c,
		int nc)
{
    int nota=1, notb=1;
    int gdrdoneflag = 0;
    static int first_call = 1;
    if (first_call){
	gdrdgemm_set_procname(MP_myprocid());
	first_call=0;
	init_current_time();
    }
    //    char str[128];
    //    sprintf(str,"before sums = %25.20e %25.20e %25.20e",
    //	    touchcmat(m,k,na,a),  touchcmat(k,n,nb,b), touchcmat(m,n,nc,c));
    //    MP_message(str);
    //    dprintf(9,"mygdrdgemm omp_max_threads=%d procs=%d\n",
    //	    omp_get_max_threads(),omp_get_num_procs());
    double zero=0.0;
    //    int tmp=0;
    //    gdr_dgemm_(&nota, &notb, &n, &m, &k, &zero, &beta, &nb, &na, &nc,
    //	       b, a, c, &tmp);
    gdr_dgemm_(&nota, &notb, &n, &m, &k, &alpha, &beta, &nb, &na, &nc,
	       b, a, c, &gdrdoneflag);
    //	cblas_dgemm( CblasRowMajor, CblasNoTrans, CblasNoTrans,
    //		     m,n, k, alpha, a, na, b, nb, beta, c, nc);
    //	gdrdoneflag=1;
    //    fprintf(stderr,"gdrflag=%d\n", gdrdoneflag);
    if(gdrdoneflag!=1){
	cblas_dgemm( CblasRowMajor, CblasNoTrans, CblasNoTrans,
		     m,n, k, alpha, a, na, b, nb, beta, c, nc);
    }
    //    dumpcmat(m,n,nc,c);
    //    sprintf(str,"after sums = %25.20e %25.20e %25.20e",
    //	    touchcmat(m,k,na,a),  touchcmat(k,n,nb,b), touchcmat(m,n,nc,c));
    //    MP_message(str);
}