/*
 * SPDX-License-Identifier:  BSD-3-Clause
 *
 * example_04-ArrayOfStructs-Naive-Omp-SIMD-Tiled.c :
 *      Example of SPH Density Calculation using a
 *      naive implementation of the main density loop,
 *      no neighbours earch, and Array of Structs (AoS)
 *      data layout, OpenMP parallelization and SIMD
 *      directives on the kernel and density calculation.
 *      This incorporates strip mining and exchange to
 *      implement cache blocking and support performance
 *      for large number of particles that would otherwise
 *      be lost.
 *
 * (C) Copyright 2021 José Hugo Elsas
 * Author: José Hugo Elsas <jhelsas@gmail.com>
 *
 * Command Line Options:
 *   -runs  <int>   : Set the number of repetitions (runs) for
 *                      calculating the density. The value of
 *                      the density is based on the last
 *                      iteration.
 *                    Default value: 1
 *   -run_seed <int>: Flag to set an alternative seed use for
 *                      for the PRNG. Instead of feeding seed
 *                      to the PRNG directly, it feeds
 *                      seed + iteration, as to generate different
 *                      configurations for each iteration.
 *                    Default value: 0 - (possible 0/1)
 *   -seed     <int>: Set the seed to use for the SPH particles
 *                      uniform position generation in the box
 *                    Default value: 123123123
 *
 *   -N        <int>: Set the number of SPH particles to be used
 *                    Default value: 1e5 = 100,000
 *   -h      <float>: Set the value of the smoothing kernel
 *                      parameter h, which corresponds to half
 *                      of the support of the kernel.
 *                    Default value: 0.05
 *
 *   -Nx       <int>: Set the number of Cells in the X direction
 *                    Default value: 10
 *   -Ny       <int>: Set the number of Cells in the Y direction
 *                    Default value: 10
 *   -Nz       <int>: Set the number of Cells in the Z direction
 *                    Default value: 10
 *
 *   -Xmin   <float>: Set the lower bound in the X direction for
 *                      the Cell Linked List box
 *                    Default value: 0.0
 *   -Ymin   <float>: Set the lower bound in the Y direction for
 *                    the Cell Linked List box
 *                      Default value: 0.0
 *   -Ymin   <float>: Set the lower bound in the Z direction for
 *                      the Cell Linked List box
 *                    Default value: 0.0
 *
 *   -Xmax   <float>: Set the lower bound in the X direction for
 *                      the Cell Linked List box
 *                    Default value: 1.0
 *   -Ymax   <float>: Set the lower bound in the Y direction for
 *                      the Cell Linked List box
 *                    Default value: 1.0
 *   -Zmax   <float>: Set the lower bound in the Z direction for
 *                      the Cell Linked List box
 *                    Default value: 1.0
 */

#include <math.h>
#include <ctype.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <limits.h>
#include <unistd.h>
#include <stdbool.h>
#include <sys/time.h>
#include <inttypes.h>

#include <omp.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_heapsort.h>

#include "sph_data_types.h"
#include "sph_linked_list.h"
#include "sph_utils.h"

#ifndef M_PI
#define M_PI (3.14159265358979323846)
#endif

#define COMPUTE_BLOCKS 1

int main_loop(int run, bool run_seed, int64_t N, double h, long int seed,
              linkedListBox *box, SPHparticle *lsph, double *times);

int compute_density_3d_naive_omp_simd_tiled(int N,double h,SPHparticle *lsph);

double w_bspline_3d_constant(double h);

#pragma omp declare simd
double w_bspline_3d_simd(double q);

int main(int argc, char **argv){
  bool run_seed = false;       // By default the behavior is is to use the same seed
  int err, runs = 1;           // By default the main loop only runs once
  long int seed = 123123123;   // The default seed is 123123123
  int64_t N = 100000;          // The default number of particles is N = 1e5 = 100,000
  double h=0.05;               // The default kernel smoothing length is h = 0.05
  linkedListBox *box;          // Uninitialized Box containing the cells for the cell linked list method
  SPHparticle *lsph;           // Uninitialized array of SPH particles

  box = (linkedListBox*)malloc(1*sizeof(linkedListBox)); // Create a box representing the entire 3d domain

  // allow for command line customization of the run
  arg_parse(argc,argv,&N,&h,&seed,&runs,&run_seed,box);  // Parse the command line options
                                                         // line arguments and override default values
  lsph = (SPHparticle*)malloc(N*sizeof(SPHparticle));    // Create an array of N particles

  double times[runs*COMPUTE_BLOCKS];

  for(int run=0;run<runs;run+=1)
    main_loop(run,run_seed,N,h,seed,box,lsph,times);

  bool is_cll = false;
  const char *prefix = "ex04,naive,AoS,omp,simd,tiled";
  print_time_stats(prefix,is_cll,N,h,seed,runs,lsph,box,times);
  print_sph_particles_density(prefix,is_cll,N,h,seed,runs,lsph,box);

  free(lsph);
  safe_free_box(box);

  return 0;
}

/*
 *  Function main_loop:
 *    Runs the main loop of the program, including the particle array generation,
 *    density calculation and the timings annotations.
 *
 *    Arguments:
 *       run <int>            : index (or value) or the present iteration
 *       run_seed <bool>      : boolean defining whether to use run index for seed or not
 *       N <int>              : Number of SPH particles to be used in the run
 *       h <double>           : Smoothing Length for the Smoothing Kernel w_bspline
 *       seed <long int>      : seed for GSL PRNG generator to generate particle positions
 *       box  <linkedListBox> : Box of linked list cells, encapsulating the 3d domain
 *       lsph <SPHparticle>   : Array (pointer) of SPH particles to be updated
 *       times <double>       : Array to store the computation timings to be updated
 *    Returns:
 *       0                    : error code returned
 *       lsph <SPHparticle>   : SPH particle array is updated in the rho field by reference
 *       times <double>       : Times is updated by reference
 */
int main_loop(int run, bool run_seed, int64_t N, double h, long int seed,
              linkedListBox *box, SPHparticle *lsph, double *times)
{
  int err;

  if(run_seed)
    err = gen_unif_rdn_pos_box(N,seed+run,box,lsph);
  else
    err = gen_unif_rdn_pos_box(N,seed,box,lsph);

  if(err)
    fprintf(stderr,"error in gen_unif_rdn_pos\n");

  // ------------------------------------------------------ //

  double t0,t1;

  t0 = omp_get_wtime();

  compute_density_3d_naive_omp_simd_tiled(N,h,lsph);        // Compute the density for all particles

  t1 = omp_get_wtime();

  // ------------------------------------------------------ //

  times[COMPUTE_BLOCKS*run+0] = t1-t0;                      // Only one component to measure time

  return 0;
}

/*
 *  Function compute_density_3d_naive_omp_simd_tiled:
 *    Computes the SPH density from the particles implementing a strip mine and exchange
 *    strategy to re-use data in cache over the direct loop. It executes calculations
 *    in parallel for the outer-most loop using openMP and SIMD in inner-most loop,
 *    though SIMD only for limited success.
 *
 *    Reference: https://en.wikipedia.org/wiki/Loop_nest_optimization
 *
 *    Arguments:
 *       N <int>              : Number of SPH particles to be used in the run
 *       h <double>           : Smoothing Length for the Smoothing Kernel w_bspline
 *       lsph <SPHparticle>   : Array (pointer) of SPH particles to be updated
 *    Returns:
 *       0                    : error code returned
 *       lsph <SPHparticle>   : SPH particle array is updated in the rho field by reference
 */
int compute_density_3d_naive_omp_simd_tiled(int N,double h,SPHparticle *lsph){
  const double inv_h = 1./h;                                       // Pre-invert the smoothing distance
  const double kernel_constant = w_bspline_3d_constant(h);         // Pre-compute the 3d normalization constant
  const int64_t STRIP = 500;                                       // Setting the size of the strip or block

  #pragma omp parallel for                                         // Run the iteration in Parallel
  for(int64_t ii=0;ii<N;ii+=1)                                     // Iterate
    lsph[ii].rho = 0.;                                             // Pre-initialize the density to zero

  #pragma omp parallel for                                         // Run the iteration in i in parallel
  for(int64_t i=0;i<N;i+=STRIP){                                   // Breaking up the i and j iterations in blocks
    for(int64_t j=0;j<N;j+=STRIP){                                 // of size STRIP to do data re-use and cache blocking
      for(int64_t ii=i;ii < ((i+STRIP<N)?(i+STRIP):N); ii+=1){     // Iterate a block over ii
        double xii = lsph[ii].r.x;                                 // Load the position in X for ii
        double yii = lsph[ii].r.y;                                 // Load the position in Y for ii
        double zii = lsph[ii].r.z;                                 // Load the position in Z for ii
        double rhoii = 0.0;                                        // Initialize partial density ii density to zero

        #pragma omp simd                                           // Hint at the compiler to vectorize this loop
        for(int64_t jj=j;jj < ((j+STRIP<N)?(j+STRIP):N); jj+=1 ){  // and iterate over the jj part of the block
          double q = 0.;                                           // initialize the distance variable

          double xij = xii-lsph[jj].r.x;                           // Load and subtract jj particle's X position component
          double yij = yii-lsph[jj].r.y;                           // Load and subtract jj particle's Y position component
          double zij = zii-lsph[jj].r.z;                           // Load and subtract jj particle's Z position component

          q += xij*xij;                                            // Add the jj contribution to the ii distance in X
          q += yij*yij;                                            // Add the jj contribution to the ii distance in Y
          q += zij*zij;                                            // Add the jj contribution to the ii distance in Z

          q = sqrt(q)*inv_h;                                       // Sqrt and normalizing the distance by the smoothing lengh

          rhoii += lsph[jj].nu*w_bspline_3d_simd(q);               // Add up the contribution from the jj particle
        }                                                          // to the intermediary density and then
        lsph[ii].rho += kernel_constant*rhoii;                     // add the intermediary density to the full density
      }
    }
  }

  return 0;
}

/*
 *  Function w_bspline_3d_constant:
 *    Returns the 3d normalization constant for the cubic b-spline SPH smoothing kernel
 *
 *    Arguments:
 *       h <double>           : Smoothing Length for the Smoothing Kernel w_bspline
 *    Returns:
 *       3d bspline normalization density <double>
 */
double w_bspline_3d_constant(double h){
  return 3./(2.*M_PI*h*h*h);  // 3d normalization value for the b-spline kernel
}

/*
 *  Function w_bspline_3d_simd:
 *    Returns the un-normalized value of the cubic b-spline SPH smoothing kernel
 *
 *    Arguments:
 *       q <double>           : Distance between particles normalized by the smoothing length h
 *    Returns:
 *       wq <double>          : Unnormalized value of the kernel
 *
 *    Observation:
 *       Why not else if(q<2.)?
 *       Because if you use "else if", the compiler refuses to vectorize,
 *       This results in a large slowdown, as of 2.5x slower for example_04
 */
#pragma omp declare simd
double w_bspline_3d_simd(double q){                                // Use as input the normalized distance
  double wq = 0.0;
  double wq1 = (0.6666666666666666 - q*q + 0.5*q*q*q);             // The first polynomial of the spline
  double wq2 = 0.16666666666666666*(2.-q)*(2.-q)*(2.-q);           // The second polynomial of the spline

  if(q<2.)                                                         // If the distance is below 2
    wq = wq2;                                                      // Use the 2nd polynomial for the spline

  if(q<1.)                                                         // If the distance is below 1
    wq = wq1;                                                      // Use the 1nd polynomial for the spline

  return wq;                                                       // return which ever value corresponds to the distance
}


/*
** LISA, 2nd level using design matrix
**
** G.Lohmann, July 2018
*/
#include <viaio/Vlib.h>
#include <viaio/VImage.h>
#include <viaio/mu.h>
#include <viaio/option.h>

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

#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_cdf.h>
#include <gsl/gsl_math.h>

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

extern gsl_matrix *PseudoInv(gsl_matrix *A,gsl_matrix *B);
extern gsl_matrix *XRead2ndLevel(VString);
extern double t2z(double t,double df);

double EstimateVariance(gsl_matrix *XInv,gsl_vector *con)
{
  int nbeta = XInv->size1;
  double var=0;

  gsl_matrix *bcov = gsl_matrix_calloc(nbeta,nbeta);
  gsl_blas_dgemm(CblasNoTrans,CblasTrans,1.0,XInv,XInv,0.0,bcov);

  gsl_vector *tmp  = gsl_vector_alloc(nbeta);
  gsl_blas_dgemv(CblasTrans,1.0,bcov,con,0.0,tmp);
  gsl_blas_ddot (tmp,con,&var);
  double sigma = sqrt(var);
  gsl_matrix_free(bcov);
  gsl_vector_free(tmp);

  return sigma;
}

/*
** get effective degrees of freedom
*/
double DoF(gsl_matrix *X,gsl_matrix *XInv,double *xtrace)
{
  int i,nimages = X->size1;

  gsl_matrix *R = gsl_matrix_calloc (nimages,nimages);
  gsl_matrix *P = gsl_matrix_calloc (nimages,nimages);
  gsl_blas_dgemm(CblasNoTrans,CblasNoTrans,1.0,X,XInv,0.0,P);
  gsl_matrix_set_identity(R);
  gsl_matrix_sub(R,P);

  double trace = 0;
  for (i=0; i<nimages; i++)
    trace += gsl_matrix_get(R,i,i);
  gsl_blas_dgemm(CblasNoTrans,CblasNoTrans,1.0,R,R,0.0,P);
  gsl_matrix_free(R);

  double trace2 = 0;
  for (i=0; i<nimages; i++) {
    trace2 += gsl_matrix_get(P,i,i);
  }
  gsl_matrix_free(P);

  (*xtrace) = trace;
  double df = (trace*trace) / trace2;
  return df;
}

/*
** general linear regression, 2nd level
*/
void GLM2(VImage *src,gsl_matrix *X,gsl_vector *contrast,int *permflag,int *permtable,int *signtable,int signswitch,VImage dest)
{
  int i,j,ip,b,r,c;
  double t=0,z=0,sum=0,tsigma=0,d=0,err=0,var=0,u=0,tiny=TINY;

  int ncols   = VImageNColumns(src[0]);
  int nrows   = VImageNRows(src[0]);
  int nslices = VImageNBands(src[0]);
  int nimages = X->size1;
  VFillImage(dest,VAllBands,0);

  /* permutation or sign switching applied to design matrix */
  gsl_matrix *XP = gsl_matrix_calloc(X->size1,X->size2);
  gsl_matrix_memcpy(XP,X);

  if (signswitch >= 0) {   /* sign switching  (one-sample test) */
    for (i=0; i<X->size1; i++) {
      if (signtable[i] < 0) {
	u = gsl_matrix_get(X,i,signswitch);
	gsl_matrix_set(XP,i,signswitch,-u);
      }
    }
  }

  else {       /*  permutations  */
    for (i=0; i<X->size1; i++) {
      ip = permtable[i];
      for (j=0; j<X->size2; j++) {
	if (permflag[j] > 0) {  /* do not permute columns containing nuisance covariates */
	  u = gsl_matrix_get(X,i,j);
	  gsl_matrix_set(XP,ip,j,u);
	}
      }
    }
  }

  /* pseudo inverse */
  gsl_matrix *XInv = PseudoInv(XP,NULL);

  /* get DoF and variance */
  double trace = 0;
  double df = DoF(XP,XInv,&trace);
  double sigma = EstimateVariance(XInv,contrast);
  if (df < tiny) VError(" Zero degrees of freedom");
  if (sigma < tiny) VError(" No variance in design/contrast");

  /* main loop */
  gsl_set_error_handler_off();
  gsl_vector *y    = gsl_vector_alloc (nimages);
  gsl_vector *yz   = gsl_vector_alloc (nimages);
  gsl_vector *beta = gsl_vector_alloc (X->size2);

  for (b=0; b<nslices; b++) {
    for (r=0; r<nrows; r++) {
      for (c=0; c<ncols; c++) {

	int nonzero=0;
	for (j=0; j<nimages; j++) {
	  y->data[j] = VGetPixel(src[j],b,r,c);
	  if (fabs(y->data[j]) > 0) nonzero++;
	}
	if (nonzero < nimages-2) continue;

	/* compute beta's */
	gsl_blas_dgemv(CblasNoTrans,1.0,XInv,y,0.0,beta);

	/* residuals */
	gsl_blas_dgemv(CblasNoTrans,1.0,XP,beta,0.0,yz);
	err = 0;
	for (j=0; j<nimages; j++) {
	  d = y->data[j] - yz->data[j];
	  err += d*d;
	}

	/* get z-value */
	t = z = sum = 0;
	gsl_blas_ddot (beta,contrast,&sum);
	if (fabs(sum) < tiny) continue;
	var = err / trace;
	tsigma = sqrt(var) * sigma;
	if (tsigma > tiny) {
	  t = sum / tsigma;
	  z = t2z(t,df);
	  if (sum < 0) z = -z;
	}
	VPixel(dest,b,r,c,VFloat) = z;
      }
    }
  }

  gsl_vector_free(y);
  gsl_vector_free(yz);
  gsl_vector_free(beta);
  gsl_matrix_free(XP);
  gsl_matrix_free(XInv);
}


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_cdf.h>

#include <omp.h>

#include "util.h"

int main(int argc, char *argv[]) {
    char* input_fileName1 = argv[1];
    char* input_fileName2 = argv[2];
    int N_doc_bg = atoi(argv[3]);
    int N_kw = atoi(argv[4]);
    int N_obs = atoi(argv[5]);
    int N_iter = atoi(argv[6]);
    int m = atoi(argv[7]);
    double p = 0.89;
    double q = atof(argv[8]);
    char* output_fileName = argv[9];
    int N_doc = 480000/2;

    int* matrix = (int*) malloc(sizeof(int) * N_obs * N_obs);
    int* matrix_bg = (int*) malloc(sizeof(int) * N_kw * N_kw);
    int* matrix_padded = (int*) malloc(sizeof(int) * N_obs * N_obs);
    int* true_index = (int*) malloc(sizeof(int) * N_kw);
    int* permutation = (int*) malloc(sizeof(int) * N_obs);
    gsl_matrix* matrix_obs;

    // Setup
    for (int round = 0; round < 10; round++)
    {
        char input_fileName1_extend[40];
        char input_fileName2_extend[40];
        sprintf(input_fileName1_extend, "%s%d", input_fileName1, round);
        sprintf(input_fileName2_extend, "%s%d", input_fileName2, round);

        struct timeval tv1,tv2;
        gettimeofday(&tv1, NULL);
        read_matrix(&true_index, &matrix_bg, 1.0*N_doc/N_doc_bg, N_kw, input_fileName2_extend);
        read_matrix(&true_index, &matrix, 1.0, N_obs, input_fileName1_extend);
        gettimeofday(&tv2, NULL);
        printf("Reading done: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));
        fflush(stdout);

        for (int iter = 0; iter < 10; iter++)
        {
            printf("Run %d\n", iter);
            matrix_obs = gsl_matrix_alloc(N_obs, N_obs);

            gettimeofday(&tv1, NULL);
            pad_matrix(&matrix_padded, &matrix, m, p, q, N_obs, N_doc);
            gettimeofday(&tv2, NULL);
            printf("Padding done: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));
            fflush(stdout);

            gettimeofday(&tv1, NULL);
            observe_matrix(matrix_obs, &matrix_padded, N_obs);
            gettimeofday(&tv2, NULL);
            printf("Observed matrix generated: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));
            fflush(stdout);

            // Permute observed matrix randomly and attack

            gettimeofday(&tv1, NULL);
            attack(matrix_obs, &matrix_bg, &permutation, m, p, q, N_kw, N_obs, N_doc, N_iter);
            gettimeofday(&tv2, NULL);
            printf("Main attack done: %f.\n", (double) ((tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec)) / N_iter);
            fflush(stdout);

            //for (int ii = 0; ii < 20; ii++)
            //   printf("%d, %d, %d, %d, %d\n", permutation[ii], true_index[ii], matrix[ii*N_obs+ii], matrix_bg[true_index[ii]*N_kw+true_index[ii]], matrix_bg[permutation[ii]*N_kw+permutation[ii]]);

            char output_fileName_full[40];
            sprintf(output_fileName_full, "%s%d-%d", output_fileName, round, iter);
            print_result(output_fileName_full, &permutation, &true_index, N_obs);
            //sprintf(output_fileName_full, "%s-%d-full", output_fileName, iter);
            //print_full_result(output_fileName_full, permutation_matrix, N_kw);
        }
    }

    free(matrix);
    free(matrix_bg);
    free(matrix_padded);
    gsl_matrix_free(matrix_obs);
    return(0);
}

double log_score(int idx1, int idx2, gsl_matrix* matrix_obs, int** matrix, int** permutation, int m, double p, double q, int N_kw, int N_doc)
{
    int idx1_m = (*permutation)[idx1];
    int idx2_m = (*permutation)[idx2];

    double score = 0.0;
    if (idx1 == idx2)
    {
        int N1 = m * (*matrix)[idx1_m*N_kw + idx2_m];
        double N1_mean = p * N1;
        double N1_var = p * (1-p) * N1 + (m * N_doc - N1) * N1 * 1.0 / N_doc / m;

        int N2 = m * (N_doc - (*matrix)[idx1_m*N_kw + idx2_m]);
        double N2_mean = q * N2;
        double N2_var = q * (1-q) * N2 + (m * N_doc - N2) * N2 * 1.0 / N_doc / m;

        double N3_var = 1.0 * m / N_doc * (*matrix)[idx1_m*N_kw + idx2_m] * (N_doc - (*matrix)[idx1_m*N_kw + idx2_m]);

        int k = (int) gsl_matrix_get(matrix_obs, idx1, idx2);

        score = gsl_cdf_gaussian_P(floor(k - N1_mean - N2_mean) + 0.5, sqrt(N1_var + N2_var + N3_var));
        score -= gsl_cdf_gaussian_P(floor(k - N1_mean - N2_mean) - 0.5, sqrt(N1_var + N2_var + N3_var));
    }

    else
    {
        int N1 = m * (*matrix)[idx1_m*N_kw + idx2_m];
        double N1_mean = p * p * N1;
        double N1_var  = p * p * (1-p) * (1-p)* N1;

        int N2 = m * ((*matrix)[idx1_m*N_kw + idx1_m] + (*matrix)[idx2_m*N_kw + idx2_m] - 2*(*matrix)[idx1_m*N_kw + idx2_m]);
        double N2_mean = p * q * N2;
        double N2_var  = p * q * (1-p) * (1-q) * N2;

        int N3 = m * N_doc - N1 - N2;
        double N3_mean = q * q * N3;
        double N3_var  = q * q * (1-q) * (1-q) * N3;

        double N4_var = 1.0 * m / N_doc * (*matrix)[idx1_m*N_kw + idx2_m] * (N_doc - (*matrix)[idx1_m*N_kw + idx2_m]);

        int M = (int) gsl_matrix_get(matrix_obs, idx1, idx2);

        score = gsl_cdf_gaussian_P(floor(M - N1_mean - N2_mean - N3_mean) + 0.5, sqrt(N1_var + N2_var + N3_var + N4_var));
        score -= gsl_cdf_gaussian_P(floor(M - N1_mean - N2_mean - N3_mean) - 0.5, sqrt(N1_var + N2_var + N3_var + N4_var));
    }

    if (score == 0)
        return(-500.0);
    return(log(score));
}

void attack(gsl_matrix* matrix_obs, int** matrix, int** permutation, int m, double p, double q, int N_kw, int N_obs, int N_doc, int N_iter)
{
    // Initialise data structures
    double* score_matrix = (double*) malloc(sizeof(double) * N_kw * N_kw);
    double* score_row1 = (double*) malloc(sizeof(double) * N_kw);
    double* score_row2 = (double*) malloc(sizeof(double) * N_kw);

    int* permutation_tmp = (int*) malloc(sizeof(int) * N_obs);
    int* permutation_inv = (int*) malloc(sizeof(int) * N_kw);

    // Initialising RNG
    const gsl_rng_type * T;
    gsl_rng * r;
    gsl_rng_env_setup();
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    struct timeval tv1,tv2;
    gettimeofday(&tv1, NULL);

    solution_initial(permutation, matrix_obs, matrix, m, p, q, N_kw, N_obs, N_doc);

    for (int ii = 0; ii < N_obs; ii++)
        permutation_tmp[ii] = (*permutation)[ii];
    for (int ii = 0; ii < N_kw; ii++)
        permutation_inv[ii] = -1;
    for (int ii = 0; ii < N_obs; ii++)
        permutation_inv[permutation_tmp[ii]] = ii;

    // Compute initial score
    #pragma omp parallel for shared(score_matrix, matrix_obs, matrix)
    for (int ii = 0; ii < N_obs * N_obs; ii++)
        score_matrix[ii] = log_score((int) (ii / N_obs), ii % N_obs, matrix_obs, matrix, permutation, m, p, q, N_kw, N_doc);
    gettimeofday(&tv2, NULL);
    printf("Initial score computed: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));

    // Iterations of simulated annealing
    double temp = (double) N_kw;
    int N_stuck = 0;
    for (int iter = 0; iter < N_iter; iter++)
    {
        /* Status code */
	    if (iter % (N_iter / 10) == 0)
        {
            gettimeofday(&tv1, NULL);
            printf("Iteration: %d, %d, %d.\n", iter, N_stuck, (int) (tv1.tv_sec - tv2.tv_sec));
            fflush(stdout);
            gettimeofday(&tv2, NULL);
        }

	    if (N_stuck >= (N_iter / 20))
            iter = N_iter;

        int idx1, idx2;
        permutation_generation(&idx1, &idx2, &permutation_tmp, permutation, &permutation_inv, matrix_obs, matrix, m, p, q, N_kw, N_obs, N_doc);
        if (idx1 == idx2)
        {
            N_stuck++;
            continue;
        }

        int ii = 0;
        #pragma omp parallel for shared(score_row1)
        for (ii = 0; ii < N_obs; ii++)
            score_row1[ii] = log_score(idx1, ii, matrix_obs, matrix, &permutation_tmp, m, p, q, N_kw, N_doc);

        if (idx2 >= 0)
            #pragma omp parallel for shared(score_row2)
            for (ii = 0; ii < N_obs; ii++)
                score_row2[ii] = log_score(idx2, ii, matrix_obs, matrix, &permutation_tmp, m, p, q, N_kw, N_doc);

        double score_diff = 0;
        for (int ii = 0; ii < N_obs; ii++)
            score_diff += score_row1[ii];
        for (int ii = 0; ii < N_obs; ii++)
            score_diff -= score_matrix[idx1*N_obs + ii];
        if (idx2 >= 0)
        {
            for (int ii = 0; ii < N_obs; ii++)
                score_diff += score_row2[ii];
            for (int ii = 0; ii < N_obs; ii++)
                score_diff -= score_matrix[idx2*N_obs + ii];
        }

        // compute difference in score, with exponentiation
        score_diff = score_diff / temp;

        if (score_diff < -40)
            score_diff = 0;
        else if (score_diff > 0)
            score_diff = 1.01;
        else
            score_diff = exp(score_diff);

        if (gsl_ran_flat(r, 0, 1) < score_diff)
        {
            // Update the scores
            for (int ii = 0; ii < N_obs; ii++)
                score_matrix[idx1*N_obs + ii] = score_row1[ii];
            for (int ii = 0; ii < N_obs; ii++)
                score_matrix[ii*N_obs + idx1] = score_row1[ii];
            if (idx2 >= 0)
            {
                for (int ii = 0; ii < N_obs; ii++)
                    score_matrix[idx2*N_obs + ii] = score_row2[ii];
                for (int ii = 0; ii < N_obs; ii++)
                    score_matrix[ii*N_obs + idx2] = score_row2[ii];
            }

            // Update the permutation
            permutation_inv[(*permutation)[idx1]] = -1;
            (*permutation)[idx1] = permutation_tmp[idx1];
            permutation_inv[permutation_tmp[idx1]] = idx1;
            if (idx2 >= 0)
            {
                (*permutation)[idx2] = permutation_tmp[idx2];
                permutation_inv[permutation_tmp[idx2]] = idx2;
            }
            N_stuck = 0;
        }
        else
        {
            // Update the permutation
            permutation_tmp[idx1] = (*permutation)[idx1];
            if (idx2 >= 0)
                permutation_tmp[idx2] = (*permutation)[idx2];
            N_stuck += 1;
        }

        temp *= 0.995;
    }

    free(score_matrix);
    free(score_row1);
    free(score_row2);
    gsl_rng_free(r);
}

void print_result(char* output_fileName, int** permutation, int** true_index, int N_obs)
{
    FILE* fp = fopen(output_fileName, "w");
    int count = 0;
    int count2 = 0;
    for (int ii = 0; ii < N_obs; ii++)
        if ((*permutation)[ii] == (*true_index)[ii])
            count++;

    fprintf(fp, "%d\n", count);
    fclose(fp);
    printf("Success: %d/%d.\n", count, N_obs);
}

void print_full_result(char* output_fileName, int** permutation, int** true_index, int N_obs)
{
    FILE* fp = fopen(output_fileName, "w");
    int count = 0;
    for (int ii = 0; ii < N_obs; ii++)
        if ((*permutation)[ii] == (*true_index)[ii])
            count++;

    fprintf(fp, "%d\n", count);
    fclose(fp);
}


#include <viaio/Vlib.h>
#include <viaio/VImage.h>
#include <viaio/mu.h>
#include <viaio/option.h>

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

#include <gsl/gsl_cblas.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_statistics.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_wavelet.h>
#include <gsl/gsl_wavelet2d.h>

extern void gsl_sort_vector(gsl_vector *);
#define ABS(x) ((x) > 0 ? (x) : -(x))

void WriteTest(gsl_matrix *data)
{
  FILE *fp=0;
  VAttrList alist;
  VImage tmp;
  double u;
  int i,j;

  tmp = VCreateImage(1,data->size1,data->size2,VFloatRepn);
  for (i = 0; i < data->size1; i++) {
    for (j = 0; j < data->size2; j++) {
      u = gsl_matrix_get(data,i,j);
      VPixel(tmp,0,i,j,VFloat) = u;
    }
  }
  alist = VCreateAttrList();
  VAppendAttr(alist,"image",NULL,VImageRepn,tmp);
  fp = fopen("qtest.v","w");
  VWriteFile (fp,alist);
  fclose(fp);
  exit(0);
}

double
estimateSigma(gsl_matrix *data)
{
  int i,j,i2,j2,n,nn;
  gsl_vector *vec=NULL;
  double median=0,med0,u,xmax,tiny=1.0e-6;

  i2 = data->size1/2;
  j2 = data->size2/2;

  nn = 3 * (i2+1) * (j2+1);
  vec = gsl_vector_calloc(nn);

  /*
  ** first pass, median, no absolute values
  */
  xmax = VRepnMaxValue(VDoubleRepn);
  gsl_vector_set_all(vec,xmax);

  n = 0;
  for (i=0; i<i2; i++) {
    for (j=j2; j<data->size2; j++) {
      u = gsl_matrix_get(data,i,j);
      if (ABS(u) > tiny) gsl_vector_set(vec,n++,u);
    }
  }

  for (i=i2; i<data->size1; i++) {
    for (j=0; j<j2; j++) {
      u = gsl_matrix_get(data,i,j);
      if (ABS(u) > tiny) gsl_vector_set(vec,n++,u);
    }
  }

  for (i=i2; i<data->size1; i++) {
    for (j=j2; j<data->size2; j++) {
      u = gsl_matrix_get(data,i,j);
      if (ABS(u) > tiny) gsl_vector_set(vec,n++,u);
    }
  }

  gsl_sort_vector(vec);
  u  = (double)n/2.0;
  nn = (int)u;
  med0 = 0.5*(gsl_vector_get(vec,nn) + gsl_vector_get(vec,nn+1));

  /*
  ** median, 2nd pass
  */
  xmax = VRepnMaxValue(VDoubleRepn);
  gsl_vector_set_all(vec,xmax);

  n = 0;
  for (i=0; i<i2; i++) {
    for (j=j2; j<data->size2; j++) {
      u = gsl_matrix_get(data,i,j);
      if (ABS(u) <= tiny) continue;
      u = ABS(u - med0);
      gsl_vector_set(vec,n++,u);
    }
  }

  for (i=i2; i<data->size1; i++) {
    for (j=0; j<j2; j++) {
      u = gsl_matrix_get(data,i,j);
      if (ABS(u) <= tiny) continue;
      u = ABS(u - med0);
      gsl_vector_set(vec,n++,u);
    }
  }

  for (i=i2; i<data->size1; i++) {
    for (j=j2; j<data->size2; j++) {
      u = gsl_matrix_get(data,i,j);
      if (ABS(u) <= tiny) continue;
      u = ABS(u - med0);
      gsl_vector_set(vec,n++,u);
    }
  }

  gsl_sort_vector(vec);
  u  = (double)n/2.0;
  nn = (int)u;
  median = 0.5*(gsl_vector_get(vec,nn) + gsl_vector_get(vec,nn+1));

  return (median/0.455);
}

double sgn(double x)
{
  if (x < 0) return -1;
  else if (x > 0) return 1;
  else return 0;
}

void
Filtering(gsl_matrix *data, int type, double intensity, double sigma)
{
  int i,j,n;
  double d,dd,alpha;

  n = data->size1;

  double tau = sigma * intensity * sqrt(2 * log10((double) (n*n)));

  for (i=0; i < data->size1; i++) {
    for (j=0; j < data->size2; j++) {
      d = gsl_matrix_get(data,i,j);

      /* Wiener-Filter */
      if (type == 0) {
	alpha = ((d * d) - intensity * (sigma * sigma)) / (d * d);
	if (alpha < 0) alpha = 0;
	d *= alpha;
      }

      /* Soft-Threshold */
      else if (type == 1) {
	if (ABS(d) < tau)
	  dd = 0;
	else
	  dd = sgn(d) * (ABS(d) - tau);
	d = dd;
      }

      /* Hard-Threshold */
      else if (type == 2) {
	if (ABS(d) < tau) d = 0;
      }

      gsl_matrix_set(data,i,j,d);
    }
  }
}

void
Shift(gsl_matrix *data,gsl_matrix *shiftdata,int ishift,int jshift)
{
  int i,j,ii,jj,n;
  double u;

  gsl_matrix_set_zero(shiftdata);

  n = data->size1;
  for (i=0; i<n; i++) {
    ii = i+ishift;
    if (ii < 0 || ii >= n) continue;
    for (j=0; j<n; j++) {
      jj = j+jshift;
      if (jj < 0 || jj >= n) continue;
      u = gsl_matrix_get(data,i,j);
      gsl_matrix_set(shiftdata,ii,jj,u);
    }
  }
}

void
AddShift(gsl_matrix *data,gsl_matrix *shiftdata,int ishift,int jshift)
{
  int i,j,ii,jj,n;
  double u,v;
  n = data->size1;

  for (i=0; i<n; i++) {
    ii = i+ishift;
    if (ii < 0 || ii >= n) continue;
    for (j=0; j<n; j++) {
      jj = j+jshift;
      if (jj < 0 || jj >= n) continue;
      u = gsl_matrix_get(data,i,j);
      v = gsl_matrix_get(shiftdata,ii,jj);
      u = u+v;
      gsl_matrix_set(data,i,j,u);
    }
  }
}

VImage
VWavelets(VImage src,VImage dest,int filtertype,int wtype,int nshift,double level)
{
  static gsl_wavelet_workspace *work=NULL;
  static gsl_wavelet *w=NULL;
  static gsl_matrix *data=NULL,*tmp=NULL,*shiftdata=NULL;
  int nslices,nrows,ncols,i,j,slice,n,size,dim;
  double u,sum1,sum2,nx,mean,sig,sigma=0,avesig;
  extern VFloat VGetPixelValue (VImage,int);

  /*
  ** read image info
  */
  nslices = VImageNBands(src);
  nrows = VImageNRows(src);
  ncols = VImageNColumns(src);

  if (dest == NULL)
    dest = VCreateImage(nslices,nrows,ncols,VPixelRepn(src));
  VCopyImageAttrs (src, dest);
  VFillImage(dest,VAllBands,0);

  size = nrows;
  if (nrows < ncols) size = ncols;
  n = (int) ceil(log10((double) size) / log10(2.0));
  dim = (int) pow(2.0, (double) n);

  /*
  ** ini wavelets transform
  */
  if (work == NULL) {

    switch (wtype) {
    case 0:
      w = gsl_wavelet_alloc (gsl_wavelet_haar_centered,2);
      break;
    case 1:
      w = gsl_wavelet_alloc (gsl_wavelet_daubechies_centered,4);
      break;
    case 2:
      w = gsl_wavelet_alloc (gsl_wavelet_daubechies_centered,6);
      break;
    case 3:
      w = gsl_wavelet_alloc (gsl_wavelet_daubechies_centered,8);
      break;
    case 4:
      w = gsl_wavelet_alloc (gsl_wavelet_bspline_centered,103);
      break;
    case 5:
      w = gsl_wavelet_alloc (gsl_wavelet_bspline_centered,202);
      break;
    default:
      VError("illegal wavelet type");
    }

    work = gsl_wavelet_workspace_alloc (dim);
    data = gsl_matrix_calloc(dim,dim);
    shiftdata = gsl_matrix_calloc(dim,dim);
    tmp = gsl_matrix_calloc(dim,dim);
  }

  avesig = 0;
  for (slice=0; slice<nslices; slice++) {

    /*
    ** read data
    */
    sum1 = sum2 = nx = 0;
    for (i=0; i<nrows; i++) {
      for (j=0; j<ncols; j++) {
	u = VGetPixel(src,slice,i,j);
	sum1 += u;
	sum2 += u*u;
	nx++;
      }
    }
    mean = sum1/nx;
    sig = sqrt((sum2 - nx * mean * mean) / (nx - 1.0));

    if (sig < 1.0e-3) {
      VFillImage(dest,slice,0);
      continue;
    }

    gsl_matrix_set_zero(data);
    for (i=0; i<nrows; i++) {
      for (j=0; j<ncols; j++) {
	u = VGetPixel(src,slice,i,j);
	u = (u-mean)/sig;
	gsl_matrix_set(data,i,j,u);
      }
    }
    gsl_matrix_memcpy(tmp,data);

    /*
    ** get sigma estimate (Donoho)
    */
    if (gsl_wavelet2d_nstransform_matrix_forward(w,data,work) != GSL_SUCCESS)
      VError(" gsl_wavelet2d_nstransform_matrix_forward failed");

    sigma = estimateSigma(data);
    /* fprintf(stderr," slice: %3d  %f\n",slice,sigma); */
    avesig += sigma;

    Filtering(data,filtertype,level,sigma);
    gsl_wavelet2d_nstransform_matrix_inverse(w,data,work);

    /*
    ** translation invariance, shift 4/8 algorithm
    */
    nx = 1;
    for (i=-nshift; i<=nshift; i++) {
      for (j=-nshift; j<=nshift; j++) {
	if (i == 0 && j == 0) continue;
	Shift(tmp,shiftdata,i,j);
	gsl_wavelet2d_nstransform_matrix_forward(w,shiftdata,work);
	Filtering(shiftdata,filtertype,level,sigma);
	gsl_wavelet2d_nstransform_matrix_inverse(w,shiftdata,work);
	AddShift(data,shiftdata,i,j);
	nx++;
      }
    }
    gsl_matrix_scale(data,1.0/nx);

    /*
    ** output
    */
    for (i=0; i<nrows; i++) {
      for (j=0; j<ncols; j++) {
	u = gsl_matrix_get(data,i,j);
	u = sig*u + mean;
	if (u < VPixelMinValue(dest)) u = VPixelMinValue(dest);
	if (u > VPixelMaxValue(dest)) u = VPixelMaxValue(dest);
	VSetPixel(dest,slice,i,j,u);
      }
    }
  }
  avesig /= (double)nslices;
  /* fprintf(stderr," kappa: %f\n",sqrt(avesig) * 1.75); */

  /*
  gsl_matrix_free(shiftdata);
  gsl_matrix_free(tmp);
  gsl_wavelet_free (w);
  gsl_wavelet_workspace_free (work);
  free (data);
  */
  return dest;
}

VDictEntry WDict[] = {
  { "haar", 0 },
  { "daub4", 1 },
  { "daub6", 2 },
  { "daub8", 3 },
  { "bspline103", 4 },
  { "bspline202", 5 },
  { NULL }
};

int main (int argc,char *argv[])
{
  static VShort ftype  = 0;
  static VShort wtype  = 0;
  static VShort nshift = 1;
  static VDouble level = 0.5;
  static VOptionDescRec  options[] = {
    {"wavelet",VShortRepn,1,(VPointer) &wtype,VOptionalOpt,WDict,"wavelet type"},
    {"filter",VShortRepn,1,(VPointer) &ftype,VOptionalOpt,NULL,"filter type"},
    {"shift",VShortRepn,1,(VPointer) &nshift,VOptionalOpt,NULL,"shift window size"},
    {"level",VDoubleRepn,1,(VPointer) &level,VOptionalOpt,NULL,"level"}
  };
  FILE *in_file,*out_file;
  VAttrList list=NULL;
  VAttrListPosn posn;
  VImage src=NULL,dest=NULL;
  int n=0,nimages=0;
  char *prg=GetLipsiaName("vdenoise");
  fprintf (stderr, "%s\n", prg);

  VParseFilterCmd (VNumber (options),options,argc,argv,&in_file,&out_file);

  if (! (list = VReadFile (in_file, NULL))) exit (1);
  fclose(in_file);

  nimages = 0;
  for (VFirstAttr (list, & posn); VAttrExists (& posn); VNextAttr (& posn)) {
    if (VGetAttrRepn (& posn) == VImageRepn) nimages++;
  }

  n = 0;
  for (VFirstAttr (list, & posn); VAttrExists (& posn); VNextAttr (& posn)) {
    if (VGetAttrRepn (& posn) != VImageRepn) continue;
    VGetAttrValue (& posn, NULL,VImageRepn, & src);
    fprintf(stderr," image %4d of %d\r",n,nimages);
    dest = VWavelets(src,dest,(int)ftype,(int)wtype,(int)nshift,level);
    src = VCopyImage(dest,src,VAllBands);
    VSetAttrValue (& posn, NULL,VImageRepn,src);
    n++;
  }
  if (src == NULL) VError(" no input image found");

  VHistory(VNumber(options),options,prg,&list,&list);
  if (! VWriteFile (out_file, list)) exit (1);
  fprintf (stderr, "%s: done.\n", argv[0]);
  return 0;
}


#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>
#include <ccl.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>

//Structure containing info about a model parameter
typedef struct {
  int index;
  double value;
} LSS_param;

//Structure defining one of the tracers being correlated
//Note that each LSS tracer contains many redshift bins
//(whereas a CCL_ClTracer corresponds to a single bin)
typedef struct {
  int n_bins; //Number of redshift bins

  int *num_z; //Number of z-values for the n(z) of each bin (dims -> [nbins])
  double **z_nz_arr; //z-values for the n(z) of each bin    (dims -> [nbins,num_z])
  double **n_nz_arr; //n(z) of each bin                     (dims -> [nbins,num_z])

  int n_nodes_b; //Number of nodes defining b(z)
  double *z_bz_arr; //Redshift values for each node (dims -> [n_nodes_b])
  LSS_param *b_bz_arr; //Bias values for each node  (dims -> [n_nodes_b])
  int interp_scheme_b; //Interpolation scheme for b(z)

  //Add similar stuff for magnification bias, 2nd-order bias etc.,

  //Add similar stuff for photo-z uncertainties?

  //Flags describing theoretical calculation
  int flag_w_rsd; //Include RSDs?
  int flag_w_mag; //Include magnification?
  int non_linear_type; //Scheme for non-linearities?

  //Array of CCL_ClTracers for each redshift bin (stored here to avoid recomputation?)
  CCL_ClTracer *t; // (dims -> [n_bins])
} LSS_tracer_info;

//Structure containing info needed to compute the LSS likelihood
typedef struct {
  //Cosmological parameters
  LSS_param par_Oc;
  LSS_param par_Ob;
  LSS_param par_Ok;
  LSS_param par_h;
  LSS_param par_s8;
  LSS_param par_ns;
  LSS_param par_w0;
  LSS_param par_wa;

  //Tracers to correlate
  LSS_tracer_info *tr1;
  LSS_tracer_info *tr2;

  //Probably a lot of stuff missing here
  //...
} LSS_likelihood_workspace;

//Structure defining a set of power spectra
typedef struct {
  //Number of bins for each tracer
  int nbin_1;
  int nbin_2;

  //Number of bandpowers for each cross-power spectrum
  int **n_bpw; // (dim -> [nbin_1,nbin_2])
  //Effective ell for each bandpower
  double ***l_bpw; // (dim -> [nbin_1,nbin_2,n_bpw[nbin_1,nbin_2]])
  //Effective C_ell for each bandpower
  gsl_vector *cl_bpw; // (flattened [nbin_1,nbin_2,n_bpw[nbin_1,nbin_2]])
  //This will probably have to be more general, with a weight-per-ell
  //array for each bandpower (e.g. to account for pseudo-Cl binning).
} LSS_2pt_Cell;

//Something like the above should be defined for the precision matrix,
//or it could be read directly into a 2D matrix

//////////
// Functions defined in lss_cell_io.c
//
//Read power spectrum
LSS_2pt_Cell *lss_read_cell(char *fname);
//Write power spectrum
void lss_write_cell(char *fname,LSS_2pt_Cell *cl);
//Read precision matrix
gsl_matrix *lss_read_precision(char *fname);
//Write precision matrix
void lss_write_precision(char *fname,gsl_matrix *prec,LSS_2pt_Cell *cl);
//Destructor
void lss_free_cell(LSS_2pt_Cell *cl);

//////////
// Functions defined in lss_tracers.c
//
//Creates new tracer
LSS_tracer_info *lss_tracer_new(int n_bins,
				int *num_z,double **z_nz_arr,double **n_nz_arr,
				int n_nodes_b,double *z_bz_arr,double *b_bz_arr,
				int flag_w_rsd,int flag_w_mag,int non_linear_type);
//Destructor
void lss_tracer_free(LSS_tracer_info *tr);

//////////
// Functions defined in lss_cell_th.c
//
//Compute theory power spectrum (and return as gsl_vector)
gsl_vector *lss_cell_theory(ccl_cosmology *cosmo,LSS_likelihood_workspace *w,LSS_2pt_Cell *cl);

//////////
// Functions defined in lss_cell_like.c
//
//Computes likelihood for a set of cosmological and nuisance parameters (encoded in params)
double lss_likelihood(int n_par,double *params,LSS_2pt_Cell *data,gsl_matrix *prec,LSS_likelihood_workspace *w);
//Creates a new LSS_likelihood_workspace
LSS_likelihood_workspace *lss_like_workspace_new(ccl_cosmology *cosmo,LSS_tracer_info *tr1,LSS_tracer_info *tr2);
//Destructor
void lss_like_workspace_free(LSS_likelihood_workspace *w);


/*
 * Command-line numerical polynomial equation solver using the GNU Scientific Library (GSL).
 * GSL is available from: "http://www.gnu.org/software/gsl/".
 *
 * Copyright (C) 2008-2011 George Gesslein II

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.

This library 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
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

The chief copyright holder can be contacted at gesslein@mathomatic.org, or
George Gesslein II, P.O. Box 224, Lansing, NY  14882-0224  USA.

 */

/*
 * To compile:

./compile.roots

 * or

gcc -O3 -Wall -o roots roots.c -lgsl -lgslcblas

 * This program nicely solves any degree polynomial
 * with real coefficients by calling the
 * GSL function gsl_poly_complex_solve().
 * This file is also useful for testing
 * and as an example of using the GSL from C.
 * The results are double precision floating point values
 * that are sometimes accurate to 14 digits.
 * Complex numbers are output if successful.
 * Here is an example of it solving a cubic polynomial:

$ roots 1 1 1 1
The 3 approximate floating point solutions of:
+x^3 +x^2 +x +1 = 0
are:

x = +0 +1*i
x = +0 -1*i
x = -1
$

Try this for a large amount of error:
roots -1 -1 1 1

Proof the GSL isn't perfect.

 */

#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <float.h>
#include <math.h>
#include <errno.h>
#include <gsl/gsl_poly.h>
#include <gsl/gsl_errno.h>

#define EPSILON	0.00000000000005	/* a good value for doubles */
#define HELP	1			/* display helpful text */

void	display_root(int i);

char	*prog_name = "roots";		/* name of this program */
double	*a, *z;				/* input and output arrays, respectively */
int	precision = DBL_DIG - 1;	/* display precision, it is not useful to set this higher than 15 */

void
usage(void)
{
  printf("\n%s version 1.0 - numerical polynomial equation solver\n", prog_name);
  printf("Uses the GNU Scientific Library (GSL).\n");
  printf("\nSolves polynomial = 0 when given all real coefficients of the polynomial.\n");
  printf("Double precision floating point math is used, accurate to about 14 digits.\n");
  printf("\nUsage: %s highest-power-coefficient ... constant-term\n", prog_name);
  printf("\nThe coefficients must be decimal, floating point, real numbers.\n");
  printf("For example, if 4 real numbers are given, there will be 3 complex number\n");
  printf("results or \"roots\" that are all valid solutions to polynomial = 0.\n");
}

int
main(int argc, char **argv)
{
  int i, highest_power;
  char *cp, *arg;
  gsl_poly_complex_workspace *w;

  if (argc <= 1) {
    fprintf(stderr, "%s: The polynomial coefficients must be specified on the command line.\n", prog_name);
    usage();
    exit(2);
  }

  highest_power = argc - 2;
  a = calloc(highest_power + 1, sizeof(double)); /* allocate real double input array */
  z = calloc(2 * highest_power, sizeof(double)); /* allocate complex double output array */

/* parse the command line into the coefficient array a[] */
  for (i = 0; i < argc - 1; i++) {
    arg = argv[argc-i-1];
    errno = 0;
    a[i] = strtod(arg, &cp);
    if (arg == cp || *cp) {
      fprintf(stderr, "%s: Argument \"%s\" is not a floating point number.\n", prog_name, arg);
      usage();
      exit(2);
    }
    if (errno) {
      fprintf(stderr, "%s: Argument \"%s\" is out of range.\n", prog_name, arg);
      exit(2);
    }
  }

#if	HELP
/* nicely display the actual polynomial equation we are solving */
  printf("The %d approximate floating point solutions of:\n", highest_power);
  for (i = highest_power; i >= 0; i--) {
    if (a[i]) {
      if (i && a[i] == 1.0) {
	printf("+x");
      } else {
        printf("%+.*g", precision, a[i]);
        if (i) {
	  printf("*x");
        }
      }
      if (i > 1) {
        printf("^%d", i);
      }
      printf(" ");
    }
  }
  printf("= 0\nare:\n\n");
#endif

/* solve the polynomial equation */
  w = gsl_poly_complex_workspace_alloc(highest_power + 1);
  if (gsl_poly_complex_solve(a, highest_power + 1, w, z) != GSL_SUCCESS) {
    fprintf(stderr, "%s: GSL approximation failed.\n", prog_name);
    exit(1);
  }
  gsl_poly_complex_workspace_free(w);

/* display all solutions */
  for (i = 0; i < highest_power; i++) {
#ifdef	EPSILON /* zero out relatively very small values (which are floating point error) */
    if (fabs(z[2*i] * EPSILON) > fabs(z[2*i+1]))
      z[2*i+1] = 0.0;
    else if (fabs(z[2*i+1] * EPSILON) > fabs(z[2*i]))
      z[2*i] = 0.0;
#endif
#if	HELP
    printf("x = ");
#endif
    display_root(i);
    printf("\n");
  }

  exit(0);
}

void
display_root(int i)
{
  printf("%+.*g", precision, z[2*i]);		/* output real part */
  if (z[2*i+1])
    printf(" %+.*g*i", precision, z[2*i+1]);	/* output imaginary part */
}


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/time.h>
#include <math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_blas.h>

#include <omp.h>

#include "util.h"

int main(int argc, char *argv[]) {
    char* input_fileName1 = argv[1];
    char* input_fileName2 = argv[2];
    int N_doc_bg = atoi(argv[3]);
    int N_kw = atoi(argv[4]);
    int N_obs = atoi(argv[5]);
    int N_iter = atoi(argv[6]);
    double lambda = atof(argv[7]);
    char* output_fileName = argv[8];
    int N_doc = 480000/2;

    int* matrix = (int*) malloc(sizeof(int) * N_obs * N_obs);
    int* matrix_bgi = (int*) malloc(sizeof(int) * N_kw * N_kw);
    int* matrix_padded = (int*) malloc(sizeof(int) * N_obs * N_obs);
    int* true_index = (int*) malloc(sizeof(int) * N_kw);
    int* permutation = (int*) malloc(sizeof(int) * N_obs);
    gsl_matrix* matrix_obs;

    for (int round = 0; round < 10; round++)
    {
        char input_fileName1_extend[40];
        char input_fileName2_extend[40];
        sprintf(input_fileName1_extend, "%s%d", input_fileName1, round);
        sprintf(input_fileName2_extend, "%s%d", input_fileName2, round);

        // Setup
        struct timeval tv1,tv2;
        gettimeofday(&tv1, NULL);
        read_matrix(&true_index, &matrix_bgi, 1.0*N_doc/N_doc_bg, N_kw, input_fileName2_extend);
        read_matrix(&true_index, &matrix, 1.0, N_obs, input_fileName1_extend);
        gettimeofday(&tv2, NULL);
        printf("Reading done: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));
        fflush(stdout);

        for (int iter = 0; iter < 10; iter++)
        {
            printf("Run %d\n", iter);
            matrix_obs = gsl_matrix_alloc(N_obs, N_obs);

            gettimeofday(&tv1, NULL);
            pad_matrix(&matrix_padded, &matrix, lambda, N_obs, N_doc);
            gettimeofday(&tv2, NULL);
            printf("Padding done: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));
            fflush(stdout);

            gettimeofday(&tv1, NULL);
            observe_matrix(matrix_obs, &matrix_padded, N_obs);
            gettimeofday(&tv2, NULL);
            printf("Observed matrix generated: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));
            fflush(stdout);

            gettimeofday(&tv1, NULL);
            attack(matrix_obs, &matrix_bgi, &permutation, N_kw, N_obs, N_doc, N_iter);
            gettimeofday(&tv2, NULL);
            printf("Main attack done: %d.\n",  (tv2.tv_sec - tv1.tv_sec));
            fflush(stdout);

            char output_fileName_full[40];
            sprintf(output_fileName_full, "%s%d-%d", output_fileName, round, iter);
            print_result(output_fileName_full, &permutation, &true_index, N_obs);
            //sprintf(output_fileName_full, "%s%d-%d-full", output_fileName, round, iter);
            //print_full_result(output_fileName_full, &permutation, &true_index, N_obs);
        }
    }

    free(matrix);
    free(matrix_padded);
    gsl_matrix_free(matrix_obs);
    return(0);
}

double log_score(int idx1, int idx2, gsl_matrix* matrix_obs, int** matrix, int** permutation, int N_kw, int N_doc)
{
    if (idx1 == idx2)
        return(0.0);

    int idx1_m = (*permutation)[idx1];
    int idx2_m = (*permutation)[idx2];

    double mean1, mean2, mean3;
    mean3 = 0;
    double var = 0;

    if (gsl_matrix_get(matrix_obs, idx1, idx1) > (*matrix)[idx1_m*N_kw + idx1_m])
    {
        mean1 = (*matrix)[idx1_m*N_kw + idx2_m];
        double x1, x2;
        x1 = gsl_matrix_get(matrix_obs, idx1, idx1) - (*matrix)[idx1_m*N_kw + idx1_m];
        x2 = gsl_matrix_get(matrix_obs, idx2, idx2) - (*matrix)[idx1_m*N_kw + idx2_m];
        mean3 = x1 / N_doc * x2;
        var += x1 / N_doc * x2 * (N_doc - x2) / N_doc;
    }
    else
    {
        mean1 = (*matrix)[idx1_m*N_kw + idx2_m] * gsl_matrix_get(matrix_obs, idx1, idx1) / (*matrix)[idx1_m*N_kw + idx1_m];
        var   += mean1 / (*matrix)[idx1_m*N_kw + idx1_m] * ((*matrix)[idx1_m*N_kw + idx1_m] - gsl_matrix_get(matrix_obs, idx1, idx1));
    }

    if (gsl_matrix_get(matrix_obs, idx2, idx2) > (*matrix)[idx2_m*N_kw + idx2_m])
    {
        mean2 = mean1;
        double x1, x2;
        x1 = gsl_matrix_get(matrix_obs, idx2, idx2) - (*matrix)[idx2_m*N_kw + idx2_m];
        x2 = gsl_matrix_get(matrix_obs, idx1, idx1) - (*matrix)[idx1_m*N_kw + idx2_m];
        mean3 += x1 / N_doc * x2;
        var += x1 / N_doc * x2 * (N_doc - x2) / N_doc;
    }
    else
    {
        mean2 = mean1 * gsl_matrix_get(matrix_obs, idx2, idx2) / (*matrix)[idx2_m*N_kw + idx2_m];
        var  += mean2 / (*matrix)[idx2_m*N_kw + idx2_m] * ((*matrix)[idx2_m*N_kw + idx2_m] - gsl_matrix_get(matrix_obs, idx2, idx2));
    }

    var += 1.0 * (*matrix)[idx1_m*N_kw + idx2_m] / N_doc * (N_doc - (*matrix)[idx2_m*N_kw + idx2_m]);

    double score = gsl_ran_gaussian_pdf(mean2 + mean3 - gsl_matrix_get(matrix_obs, idx1, idx2), sqrt(var));

    if (score == 0)
        return(-500.0);
    return(log(score));
}

void attack(gsl_matrix* matrix_obs, int** matrix, int** permutation, int N_kw, int N_obs, int N_doc, int N_iter)
{
    // Initialise data structures
    double* score_matrix = (double*) malloc(sizeof(double) * N_obs * N_obs);
    double* score_row1 = (double*) malloc(sizeof(double) * N_obs);
    double* score_row2 = (double*) malloc(sizeof(double) * N_obs);

    int* permutation_tmp = (int*) malloc(sizeof(int) * N_obs);
    int* permutation_inv = (int*) malloc(sizeof(int) * N_kw);

    // Initialise permutations
    for (int ii = 0; ii < N_obs; ii++)
        (*permutation)[ii] = ii;
    for (int ii = 0; ii < N_obs; ii++)
        permutation_tmp[ii] = ii;
    for (int ii = 0; ii < N_obs; ii++)
        permutation_inv[ii] = ii;
    for (int ii = N_obs; ii < N_kw; ii++)
        permutation_inv[ii] = -1;

    // Initialising RNG
    const gsl_rng_type * T;
    gsl_rng * r;
    gsl_rng_env_setup();
    T = gsl_rng_default;
    r = gsl_rng_alloc (T);

    struct timeval tv1,tv2;
    gettimeofday(&tv1, NULL);

    // Compute initial score
    #pragma omp parallel for shared(score_matrix, matrix_obs, matrix)
    for (int ii = 0; ii < N_obs * N_obs; ii++)
        score_matrix[ii] = log_score((int) (ii / N_obs), ii % N_obs, matrix_obs, matrix, permutation, N_kw, N_doc);
    gettimeofday(&tv2, NULL);
    printf("Initial score computed: %f.\n", (double) (tv2.tv_usec - tv1.tv_usec) / 1000000 + (double) (tv2.tv_sec - tv1.tv_sec));

    // Iterations of simulated annealing
    double temp = (double) N_kw;
    int N_stuck = 0;
    for (int iter = 0; iter < N_iter; iter++)
    {
	    /* Status code */
	    if (iter % (N_iter / 10) == 0)
        {
            gettimeofday(&tv1, NULL);
            printf("Iteration: %d, %d, %d.\n", iter, N_stuck, (int) (tv1.tv_sec - tv2.tv_sec));
            fflush(stdout);
            gettimeofday(&tv2, NULL);
        }

        if (N_stuck >= 20000)
            iter = N_iter;

	    /* Main code */
        int idx1, idx2;
        permutation_generation(&idx1, &idx2, &permutation_tmp, permutation, &permutation_inv, N_kw, N_obs);

        int ii = 0;
        #pragma omp parallel for shared(score_row1)
        for (ii = 0; ii < N_obs; ii++)
            score_row1[ii] = log_score(idx1, ii, matrix_obs, matrix, &permutation_tmp, N_kw, N_doc);

        if (idx2 >= 0)
            #pragma omp parallel for shared(score_row2)
            for (ii = 0; ii < N_obs; ii++)
                score_row2[ii] = log_score(idx2, ii, matrix_obs, matrix, &permutation_tmp, N_kw, N_doc);

        double score_diff = 0;
        for (int ii = 0; ii < N_obs; ii++)
            score_diff += score_row1[ii];
        for (int ii = 0; ii < N_obs; ii++)
            score_diff -= score_matrix[idx1*N_obs + ii];
        if (idx2 >= 0)
        {
            for (int ii = 0; ii < N_obs; ii++)
                score_diff += score_row2[ii];
            for (int ii = 0; ii < N_obs; ii++)
                score_diff -= score_matrix[idx2*N_obs + ii];
        }

        // compute difference in score, with exponentiation
        score_diff = score_diff / temp;

        if (score_diff < -40)
            score_diff = 0;
        else if (score_diff > 0)
            score_diff = 1.01;
        else
            score_diff = exp(score_diff);

        if (gsl_ran_flat(r, 0, 1) < score_diff)
        {
            // Update the scores
            for (int ii = 0; ii < N_obs; ii++)
                score_matrix[idx1*N_obs + ii] = score_row1[ii];
            for (int ii = 0; ii < N_obs; ii++)
                score_matrix[ii*N_obs + idx1] = score_row1[ii];
            if (idx2 >= 0)
            {
                for (int ii = 0; ii < N_obs; ii++)
                    score_matrix[idx2*N_obs + ii] = score_row2[ii];
                for (int ii = 0; ii < N_obs; ii++)
                    score_matrix[ii*N_obs + idx2] = score_row2[ii];
            }

            // Update the permutation
            permutation_inv[(*permutation)[idx1]] = -1;
            (*permutation)[idx1] = permutation_tmp[idx1];
            permutation_inv[permutation_tmp[idx1]] = idx1;
            if (idx2 >= 0)
            {
                (*permutation)[idx2] = permutation_tmp[idx2];
                permutation_inv[permutation_tmp[idx2]] = idx2;
            }
            N_stuck = 0;
        }
        else
        {
            // Update the permutation
            permutation_tmp[idx1] = (*permutation)[idx1];
            if (idx2 >= 0)
                permutation_tmp[idx2] = (*permutation)[idx2];
            N_stuck += 1;
        }

        temp *= 0.995;
    }

    free(score_matrix);
    free(score_row1);
    free(score_row2);
    gsl_rng_free(r);
}

void print_result(char* output_fileName, int** permutation, int** true_index, int N_obs)
{
    FILE* fp = fopen(output_fileName, "w");
    int count = 0;
    for (int ii = 0; ii < N_obs; ii++)
        if ((*permutation)[ii] == (*true_index)[ii])
            count++;

    fprintf(fp, "%d\n", count);
    fclose(fp);
    printf("Success: %d/%d.\n", count, N_obs);
}

void print_full_result(char* output_fileName, int** permutation, int** true_index, int N_obs)
{
    FILE* fp = fopen(output_fileName, "w");
    int count = 0;
    for (int ii = 0; ii < N_obs; ii++)
        fprintf(fp, "%d, %d\n", (*permutation)[ii], (*true_index)[ii]);
    fclose(fp);
}


/* randist/beta.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000, 2007 James Theiler, Brian Gough
 *
 * 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 <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_sf_gamma.h>

/* The beta distribution has the form

   p(x) dx = (Gamma(a + b)/(Gamma(a) Gamma(b))) x^(a-1) (1-x)^(b-1) dx

   The method used here is the one described in Knuth */

double
gsl_ran_beta (const gsl_rng * r, const double a, const double b)
{
  if ( (a <= 1.0) && (b <= 1.0) )
    {
      double U, V, X, Y;
      while (1)
        {
          U = gsl_rng_uniform_pos(r);
          V = gsl_rng_uniform_pos(r);
          X = pow(U, 1.0/a);
          Y = pow(V, 1.0/b);
          if ((X + Y ) <= 1.0)
            {
              if (X + Y > 0)
                {
                  return X/ (X + Y);
                }
              else
                {
                  double logX = log(U)/a;
                  double logY = log(V)/b;
                  double logM = logX > logY ? logX: logY;
                  logX -= logM;
                  logY -= logM;
                  return exp(logX - log(exp(logX) + exp(logY)));
                }
            }
        }
    }
  else
    {
      double x1 = gsl_ran_gamma (r, a, 1.0);
      double x2 = gsl_ran_gamma (r, b, 1.0);
      return x1 / (x1 + x2);
    }
}

double
gsl_ran_beta_pdf (const double x, const double a, const double b)
{
  if (x < 0 || x > 1)
    {
      return 0 ;
    }
  else
    {
      double p;

      double gab = gsl_sf_lngamma (a + b);
      double ga = gsl_sf_lngamma (a);
      double gb = gsl_sf_lngamma (b);

      if (x == 0.0 || x == 1.0)
        {
	  if (a > 1.0 && b > 1.0)
	    {
	      p = 0.0;
	    }
	  else
	    {
	      p = exp (gab - ga - gb) * pow (x, a - 1) * pow (1 - x, b - 1);
	    }
        }
      else
        {
          p = exp (gab - ga - gb + log(x) * (a - 1)  + log1p(-x) * (b - 1));
        }

      return p;
    }
}


/* rng/ran2.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 James Theiler, Brian Gough
 *
 * 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 2 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_rng.h>

/* This is an implementation of the algorithm used in Numerical
   Recipe's ran2 generator.  It is a L'Ecuyer combined recursive
   generator with a 32-element shuffle-box.

   As far as I can tell, in general the effects of adding a shuffle
   box cannot be proven theoretically, so the period of this generator
   is unknown.

   The period of the underlying combined generator is O(2^60). */

static inline unsigned long int ran2_get (void *vstate);
static double ran2_get_double (void *vstate);
static void ran2_set (void *state, unsigned long int s);

static const long int m1 = 2147483563, a1 = 40014, q1 = 53668, r1 = 12211;
static const long int m2 = 2147483399, a2 = 40692, q2 = 52774, r2 = 3791;

#define N_SHUFFLE 32
#define N_DIV (1 + 2147483562/N_SHUFFLE)

typedef struct
  {
    unsigned long int x;
    unsigned long int y;
    unsigned long int n;
    unsigned long int shuffle[N_SHUFFLE];
  }
ran2_state_t;

static inline unsigned long int
ran2_get (void *vstate)
{
  ran2_state_t *state = (ran2_state_t *) vstate;

  const unsigned long int x = state->x;
  const unsigned long int y = state->y;

  long int h1 = x / q1;
  long int t1 = a1 * (x - h1 * q1) - h1 * r1;

  long int h2 = y / q2;
  long int t2 = a2 * (y - h2 * q2) - h2 * r2;

  if (t1 < 0)
    t1 += m1;

  if (t2 < 0)
    t2 += m2;

  state->x = t1;
  state->y = t2;

  {
    unsigned long int j = state->n / N_DIV;
    long int delta = state->shuffle[j] - t2;
    if (delta < 1)
      delta += m1 - 1;
    state->n = delta;
    state->shuffle[j] = t1;
  }

  return state->n;
}

static double
ran2_get_double (void *vstate)
{
  float x_max = 1 - 1.2e-7f ; /* Numerical Recipes version of 1-FLT_EPS */

  float x = ran2_get (vstate) / 2147483563.0f ;

  if (x > x_max)
    return x_max ;

  return x ;
}

static void
ran2_set (void *vstate, unsigned long int s)
{
  ran2_state_t *state = (ran2_state_t *) vstate;
  int i;

  if (s == 0)
    s = 1;	/* default seed is 1 */

  state->y = s;

  for (i = 0; i < 8; i++)
    {
      long int h = s / q1;
      long int t = a1 * (s - h * q1) - h * r1;
      if (t < 0)
	t += m1;
      s = t;
    }

  for (i = N_SHUFFLE - 1; i >= 0; i--)
    {
      long int h = s / q1;
      long int t = a1 * (s - h * q1) - h * r1;
      if (t < 0)
	t += m1;
      s = t;
      state->shuffle[i] = s;
    }

  state->x = s;
  state->n = s;

  return;
}

static const gsl_rng_type ran2_type =
{"ran2",			/* name */
 2147483562,			/* RAND_MAX */
 1,				/* RAND_MIN */
 sizeof (ran2_state_t),
 &ran2_set,
 &ran2_get,
 &ran2_get_double};

const gsl_rng_type *gsl_rng_ran2 = &ran2_type;


/* ode-initval/test_odeiv.c
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 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., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Author:  G. Jungman
 */
#include <config.h>
#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include <gsl/gsl_test.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_ieee_utils.h>
#include "gsl_odeiv.h"

int rhs_linear (double t, const double y[], double f[], void *params);
int jac_linear (double t, const double y[], double *dfdy, double dfdt[],
                void *params);
int rhs_sin (double t, const double y[], double f[], void *params);
int jac_sin (double t, const double y[], double *dfdy, double dfdt[],
             void *params);
int rhs_exp (double t, const double y[], double f[], void *params);
int jac_exp (double t, const double y[], double *dfdy, double dfdt[],
             void *params);
int rhs_stiff (double t, const double y[], double f[], void *params);
int jac_stiff (double t, const double y[], double *dfdy, double dfdt[],
               void *params);
void test_stepper_linear (const gsl_odeiv_step_type * T, double h,
                         double base_prec);
void test_stepper_sin (const gsl_odeiv_step_type * T, double h, double base_prec);
void test_stepper_exp (const gsl_odeiv_step_type * T, double h, double base_prec);
void test_stepper_stiff (const gsl_odeiv_step_type * T, double h, double base_prec);
void test_evolve_system_flat (gsl_odeiv_step * step,
                              const gsl_odeiv_system * sys,
                              double t0, double t1, double hstart,
                              double y[], double yfin[],
                              double err_target, const char *desc);

void test_evolve_system (const gsl_odeiv_step_type * T,
                         const gsl_odeiv_system * sys,
                         double t0, double t1, double hstart,
                         double y[], double yfin[],
                         double err_target, const char *desc);

void test_evolve_exp (const gsl_odeiv_step_type * T, double h, double err);
void test_evolve_sin (const gsl_odeiv_step_type * T, double h, double err);
void test_evolve_stiff1 (const gsl_odeiv_step_type * T, double h, double err);
void test_evolve_stiff5 (const gsl_odeiv_step_type * T, double h, double err);

/* RHS for a + b t */

int
rhs_linear (double t, const double y[], double f[], void *params)
{
  f[0] = 0.0;
  f[1] = y[0];
  return GSL_SUCCESS;
}

int
jac_linear (double t, const double y[], double *dfdy, double dfdt[],
	    void *params)
{
  gsl_matrix dfdy_mat;
  dfdy_mat.size1 = 2;
  dfdy_mat.size2 = 2;
  dfdy_mat.tda = 2;
  dfdy_mat.data = dfdy;
  dfdy_mat.block = 0;
  gsl_matrix_set (&dfdy_mat, 0, 0, 0.0);
  gsl_matrix_set (&dfdy_mat, 0, 1, 0.0);
  gsl_matrix_set (&dfdy_mat, 1, 0, 1.0);
  gsl_matrix_set (&dfdy_mat, 1, 1, 0.0);
  dfdt[0] = 0.0;
  dfdt[1] = 0.0;
  return GSL_SUCCESS;
}

gsl_odeiv_system rhs_func_lin = {
  rhs_linear,
  jac_linear,
  2,
  0
};

/* RHS for sin(t),cos(t) */

int
rhs_sin (double t, const double y[], double f[], void *params)
{
  f[0] = -y[1];
  f[1] = y[0];
  return GSL_SUCCESS;
}

int
jac_sin (double t, const double y[], double *dfdy, double dfdt[],
	 void *params)
{
  gsl_matrix dfdy_mat;
  dfdy_mat.data = dfdy;
  dfdy_mat.size1 = 2;
  dfdy_mat.size2 = 2;
  dfdy_mat.tda = 2;
  dfdy_mat.block = 0;
  gsl_matrix_set (&dfdy_mat, 0, 0, 0.0);
  gsl_matrix_set (&dfdy_mat, 0, 1, -1.0);
  gsl_matrix_set (&dfdy_mat, 1, 0, 1.0);
  gsl_matrix_set (&dfdy_mat, 1, 1, 0.0);
  dfdt[0] = 0.0;
  dfdt[1] = 0.0;
  return GSL_SUCCESS;
}

gsl_odeiv_system rhs_func_sin = {
  rhs_sin,
  jac_sin,
  2,
  0
};

/* RHS for a exp(t)+ b exp(-t) */

int
rhs_exp (double t, const double y[], double f[], void *params)
{
  f[0] = y[1];
  f[1] = y[0];
  return GSL_SUCCESS;
}

int
jac_exp (double t, const double y[], double *dfdy, double dfdt[],
	 void *params)
{
  gsl_matrix dfdy_mat;
  dfdy_mat.data = dfdy;
  dfdy_mat.size1 = 2;
  dfdy_mat.size2 = 2;
  dfdy_mat.tda = 2;
  dfdy_mat.block = 0;
  gsl_matrix_set (&dfdy_mat, 0, 0, 0.0);
  gsl_matrix_set (&dfdy_mat, 0, 1, 1.0);
  gsl_matrix_set (&dfdy_mat, 1, 0, 1.0);
  gsl_matrix_set (&dfdy_mat, 1, 1, 0.0);
  dfdt[0] = 0.0;
  dfdt[1] = 0.0;
  return GSL_SUCCESS;
}

gsl_odeiv_system rhs_func_exp = {
  rhs_exp,
  jac_exp,
  2,
  0
};

/* RHS for stiff example */

int
rhs_stiff (double t, const double y[], double f[], void *params)
{
  f[0] = 998.0 * y[0] + 1998.0 * y[1];
  f[1] = -999.0 * y[0] - 1999.0 * y[1];
  return GSL_SUCCESS;
}

int
jac_stiff (double t, const double y[], double *dfdy, double dfdt[],
	   void *params)
{
  gsl_matrix dfdy_mat;
  dfdy_mat.data = dfdy;
  dfdy_mat.size1 = 2;
  dfdy_mat.size2 = 2;
  dfdy_mat.tda = 2;
  dfdy_mat.block = 0;
  gsl_matrix_set (&dfdy_mat, 0, 0, 998.0);
  gsl_matrix_set (&dfdy_mat, 0, 1, 1998.0);
  gsl_matrix_set (&dfdy_mat, 1, 0, -999.0);
  gsl_matrix_set (&dfdy_mat, 1, 1, -1999.0);
  dfdt[0] = 0.0;
  dfdt[1] = 0.0;
  return GSL_SUCCESS;
}

gsl_odeiv_system rhs_func_stiff = {
  rhs_stiff,
  jac_stiff,
  2,
  0
};

void
test_stepper_linear (const gsl_odeiv_step_type * T, double h,
		     double base_prec)
{
  int s = 0;
  double y[2];
  double yerr[2];
  double t;
  double del;
  double delmax = 0.0;
  int count = 0;

  gsl_odeiv_step *stepper = gsl_odeiv_step_alloc (T, 2);

  y[0] = 1.0;
  y[1] = 0.0;

  for (t = 0.0; t < 4.0; t += h)
    {
      gsl_odeiv_step_apply (stepper, t, h, y, yerr, 0, 0, &rhs_func_lin);
      del = fabs ((y[1] - (t + h)) / y[1]);
      delmax = GSL_MAX_DBL (del, delmax);
      if (del > (count + 1.0) * base_prec)
	{
	  printf ("  LINEAR(%20.17g)  %20.17g  %20.17g  %8.4g\n", t + h, y[1],
		  t + h, del);
	  s++;
	}
      count++;
    }

  gsl_test (s, "%s, linear [0,4], max relative error = %g",
	    gsl_odeiv_step_name (stepper), delmax);

  gsl_odeiv_step_free (stepper);
}

void
test_stepper_sin (const gsl_odeiv_step_type * T, double h, double base_prec)
{
  int s = 0;
  double y[2];
  double yerr[2];
  double t;
  double del;
  double delmax = 0.0;
  int count = 0;

  gsl_odeiv_step *stepper = gsl_odeiv_step_alloc (T, 2);

  y[0] = 1.0;
  y[1] = 0.0;

  for (t = 0.0; t < M_PI; t += h)
    {
      int stat;
      double sin_th = sin (t + h);
      gsl_odeiv_step_apply (stepper, t, h, y, yerr, 0, 0, &rhs_func_sin);
      del = fabs ((y[1] - sin_th) / sin_th);
      delmax = GSL_MAX_DBL (del, delmax);
      {
	if (t < 0.5 * M_PI)
	  {
	    stat = (del > (count + 1.0) * base_prec);
	  }
	else if (t < 0.7 * M_PI)
	  {
	    stat = (del > 1.0e+04 * base_prec);
	  }
	else if (t < 0.9 * M_PI)
	  {
	    stat = (del > 1.0e+06 * base_prec);
	  }
	else
	  {
	    stat = (del > 1.0e+09 * base_prec);
	  }
	if (stat != 0)
	  {
	    printf ("  SIN(%22.18g)  %22.18g  %22.18g  %10.6g\n", t + h, y[1],
		    sin_th, del);
	  }
	s += stat;
      }
      count++;
    }
  if (delmax > 1.0e+09 * base_prec)
    {
      s++;
      printf ("  SIN(0 .. M_PI)  delmax = %g\n", delmax);
    }

  gsl_test (s, "%s, sine [0,pi], max relative error = %g",
	    gsl_odeiv_step_name (stepper), delmax);

  delmax = 0.0;
  for (; t < 100.5 * M_PI; t += h)
    {
      gsl_odeiv_step_apply (stepper, t, h, y, yerr, 0, 0, &rhs_func_sin);
      del = fabs (y[1] - sin (t));
      delmax = GSL_MAX_DBL (del, delmax);
      count++;
    }

  if (del > count * 2.0 * base_prec)
    {
      s++;
      printf ("  SIN(%22.18g)  %22.18g  %22.18g  %10.6g\n", t + h, y[1],
	      sin (t), del);
    }

  gsl_test (s, "%s, sine [pi,100.5*pi], max absolute error = %g",
	    gsl_odeiv_step_name (stepper), delmax);

  gsl_odeiv_step_free (stepper);
}

void
test_stepper_exp (const gsl_odeiv_step_type * T, double h, double base_prec)
{
  int s = 0;
  double y[2];
  double yerr[2];
  double t;
  double del, delmax = 0.0;
  int count = 0;

  gsl_odeiv_step *stepper = gsl_odeiv_step_alloc (T, 2);

  y[0] = 1.0;
  y[1] = 1.0;

  for (t = 0.0; t < 20.0; t += h)
    {
      double ex = exp (t + h);
      gsl_odeiv_step_apply (stepper, t, h, y, yerr, 0, 0, &rhs_func_exp);
      del = fabs ((y[1] - ex) / y[1]);
      delmax = GSL_MAX_DBL (del, delmax);
      if (del > (count + 1.0) * 2.0 * base_prec)
	{
	  printf ("  EXP(%20.17g)  %20.17g  %20.17g  %8.4g\n", t + h, y[1],
		  ex, del);
	  s++;
	}
      count++;
    }

  gsl_test (s, "%s, exponential [0,20], max relative error = %g",
	    gsl_odeiv_step_name (stepper), delmax);

  gsl_odeiv_step_free (stepper);
}

void
test_stepper_stiff (const gsl_odeiv_step_type * T, double h, double base_prec)
{
  int s = 0;
  double y[2];
  double yerr[2];
  double t;
  double del, delmax = 0.0;
  int count = 0;

  gsl_odeiv_step *stepper = gsl_odeiv_step_alloc (T, 2);

  y[0] = 1.0;
  y[1] = 0.0;

  for (t = 0.0; t < 20.0; t += h)
    {
      gsl_odeiv_step_apply (stepper, t, h, y, yerr, NULL, NULL,
			    &rhs_func_stiff);

      if (t > 0.04)
	{
	  double arg = t + h;
	  double e1 = exp (-arg);
	  double e2 = exp (-1000.0 * arg);
	  double u = 2.0 * e1 - e2;
	  /* double v = -e1 + e2; */
	  del = fabs ((y[0] - u) / y[0]);
	  delmax = GSL_MAX_DBL (del, delmax);

	  if (del > (count + 1.0) * 100.0 * base_prec)
	    {
	      printf ("  STIFF(%20.17g)  %20.17g  %20.17g  %8.4g\n", arg,
		      y[0], u, del);
	      s++;
	    }
	}
      count++;
    }

  gsl_test (s, "%s, stiff [0,20], max relative error = %g",
	    gsl_odeiv_step_name (stepper), delmax);

  gsl_odeiv_step_free (stepper);
}

void
test_evolve_system_flat (gsl_odeiv_step * step,
			 const gsl_odeiv_system * sys,
			 double t0, double t1, double hstart,
			 double y[], double yfin[],
			 double err_target, const char *desc)
{
  int s = 0;
  double frac;

  double t = t0;
  double h = hstart;

  gsl_odeiv_evolve *e = gsl_odeiv_evolve_alloc (sys->dimension);

  while (t < t1)
    {
      gsl_odeiv_evolve_apply (e, NULL, step, sys, &t, t1, &h, y);
    }

  frac = fabs ((y[1] - yfin[1]) / yfin[1]) + fabs ((y[0] - yfin[0]) / yfin[0]);

  if (frac > 2.0 * e->count * err_target)
    {
      printf ("FLAT t = %.5e  y0 = %g y1= %g\n", t, y[0], y[1]);
      s++;
    }

  gsl_test (s, "%s, %s, evolve, no control, max relative error = %g",
	    gsl_odeiv_step_name (step), desc, frac);

  gsl_odeiv_evolve_free (e);
}

void
test_evolve_system (const gsl_odeiv_step_type * T,
		    const gsl_odeiv_system * sys,
		    double t0, double t1, double hstart,
		    double y[], double yfin[],
		    double err_target, const char *desc)
{
  int s = 0;
  double frac;

  double t = t0;
  double h = hstart;

  gsl_odeiv_step * step = gsl_odeiv_step_alloc (T, sys->dimension);

  gsl_odeiv_control *c =
    gsl_odeiv_control_standard_new (0.0, err_target, 1.0, 1.0);
  gsl_odeiv_evolve *e = gsl_odeiv_evolve_alloc (sys->dimension);

  while (t < t1)
    {
      gsl_odeiv_evolve_apply (e, c, step, sys, &t, t1, &h, y);
      /* printf ("SYS  t = %.18e h = %g y0 = %g y1= %g\n", t, h, y[0], y[1]); */
    }

  frac = fabs ((y[1] - yfin[1]) / yfin[1]) + fabs ((y[0] - yfin[0]) / yfin[0]);

  if (frac > 2.0 * e->count * err_target)
    {
      printf ("SYS  t = %.5e h = %g y0 = %g y1= %g\n", t, h, y[0], y[1]);
      s++;
    }

  gsl_test (s, "%s, %s, evolve, standard control, relative error = %g",
	    gsl_odeiv_step_name (step), desc, frac);

  gsl_odeiv_evolve_free (e);
  gsl_odeiv_control_free (c);
  gsl_odeiv_step_free (step);
}

void
test_evolve_exp (const gsl_odeiv_step_type * T, double h, double err)
{
  double y[2];
  double yfin[2];

  y[0] = 1.0;
  y[1] = 1.0;
  yfin[0] = exp (10.0);
  yfin[1] = yfin[0];
  test_evolve_system (T, &rhs_func_exp, 0.0, 10.0, h, y, yfin, err, "exp [0,10]");
}

void
test_evolve_sin (const gsl_odeiv_step_type * T, double h, double err)
{
  double y[2];
  double yfin[2];

  y[0] = 1.0;
  y[1] = 0.0;
  yfin[0] = cos (2.0);
  yfin[1] = sin (2.0);
  test_evolve_system (T, &rhs_func_sin, 0.0, 2.0, h, y, yfin, err, "sine [0,2]");
}

void
test_evolve_stiff1 (const gsl_odeiv_step_type * T, double h, double err)
{
  double y[2];
  double yfin[2];

  y[0] = 1.0;
  y[1] = 0.0;
  {
    double arg = 1.0;
    double e1 = exp (-arg);
    double e2 = exp (-1000.0 * arg);
    yfin[0] = 2.0 * e1 - e2;
    yfin[1] = -e1 + e2;
  }
  test_evolve_system (T, &rhs_func_stiff, 0.0, 1.0, h, y, yfin, err,
                      "stiff [0,1]");
}

void
test_evolve_stiff5 (const gsl_odeiv_step_type * T, double h, double err)
{
  double y[2];
  double yfin[2];

  y[0] = 1.0;
  y[1] = 0.0;
  {
    double arg = 5.0;
    double e1 = exp (-arg);
    double e2 = exp (-1000.0 * arg);
    yfin[0] = 2.0 * e1 - e2;
    yfin[1] = -e1 + e2;
  }
  test_evolve_system (T, &rhs_func_stiff, 0.0, 5.0, h, y, yfin, err,
                      "stiff [0,5]");
}

int
main (void)
{
  int i;

  struct ptype
  {
    const gsl_odeiv_step_type *type;
    double h;
  }
  p[20];

  p[0].type = gsl_odeiv_step_rk2;
  p[0].h = 1.0e-03;
  p[1].type = gsl_odeiv_step_rk2imp;
  p[1].h = 1.0e-03;
  p[2].type = gsl_odeiv_step_rk4;
  p[2].h = 1.0e-03;
  p[3].type = gsl_odeiv_step_rk4imp;
  p[3].h = 1.0e-03;
  p[4].type = gsl_odeiv_step_rkf45;
  p[4].h = 1.0e-03;
  p[5].type = gsl_odeiv_step_rk8pd;
  p[5].h = 1.0e-03;
  p[6].type = gsl_odeiv_step_rkck;
  p[6].h = 1.0e-03;
  p[7].type = gsl_odeiv_step_bsimp;
  p[7].h = 0.1;
  p[8].type = gsl_odeiv_step_gear1;
  p[8].h = 1.0e-03;
  p[9].type = gsl_odeiv_step_gear2;
  p[9].h = 1.0e-03;
  p[10].type = 0;

  gsl_ieee_env_setup ();

  for (i = 0; p[i].type != 0; i++)
    {
      test_stepper_linear (p[i].type, p[i].h, GSL_SQRT_DBL_EPSILON);
      test_stepper_sin (p[i].type, p[i].h / 10.0, GSL_SQRT_DBL_EPSILON);
      test_stepper_exp (p[i].type, p[i].h / 10.0, GSL_SQRT_DBL_EPSILON);
      test_stepper_stiff (p[i].type, p[i].h / 10.0, GSL_SQRT_DBL_EPSILON);
    }

  for (i = 0; p[i].type != 0; i++)
    {
      test_evolve_exp (p[i].type, p[i].h, GSL_SQRT_DBL_EPSILON);
      test_evolve_sin (p[i].type, p[i].h, GSL_SQRT_DBL_EPSILON);
      test_evolve_stiff1 (p[i].type, p[i].h, 1e-5);
      test_evolve_stiff5 (p[i].type, p[i].h, 1e-5);
    }

  exit (gsl_test_summary ());
}