#include <stdio.h>
#include <stdlib.h>
#include <sys/time.h>
#include <assert.h>
#include <math.h>

#if  __INTEL_COMPILER
#include "mkl.h"
#else
#include <gsl/gsl_cblas.h>
#endif

#include "defs.h"
#include "utils.h"
#include "progressbar.h"

long naive(const double * restrict x1, const double * restrict y1, const int N, double * restrict d)
{
	long numcomputed=0;

	for(int i=0;i<N;i++) {
		for(int j=0;j<N;j++) {
			double *dist = d + i*N + j;
			const double *y0 = (const double *) &y1[j*NDIM];
			const double *x0 = (const double *) &x1[i*NDIM];
			const double dx = x0[0] - y0[0];
			const double dy = x0[1] - y0[1];
			const double dz = x0[2] - y0[2];
#ifndef SQRT_DIST
			*dist = dx*dx + dy*dy + dz*dz;
#else
			*dist = sqrt(dx*dx + dy*dy + dz*dz);
#endif
		}
 		numcomputed+=N;
	}

	return numcomputed;
}

long chunked(const double * restrict x1, const double * restrict y1, const int N, double * restrict d)
{
	long numcomputed = 0;
	const int block_size = 16;
	for(int i=0;i<N;i+=block_size) {
		const int block_size1 = (N-i) > block_size ? block_size:(N-i);
		for(int j=0;j<N;j+=block_size) {
			const int block_size2 = (N-j) > block_size ? block_size:(N-j);
			for(int ii=0;ii<block_size1;ii++) {
				const long index = (i+ii)*(long) N + j;
				double *dist = (double *) &d[index];
				for(int jj=0;jj<block_size2;jj++) {
					const double *y0 = (const double *) (&y1[(j+jj)*NDIM]);
					const double *x0 = (const double *) (&x1[(i+ii)*NDIM]);
					const double dx = x0[0] - y0[0];
					const double dy = x0[1] - y0[1];
					const double dz = x0[2] - y0[2];
					const double sqr_ds = dx*dx + dy*dy + dz*dz;
#ifndef SQRT_DIST
					dist[jj] = sqr_ds;
#else
					dist[jj] = sqrt(sqr_ds);
#endif
				}
				numcomputed += block_size2;
			}
		}
	}

	return numcomputed;
}

long compiler_vectorized_chunked(const double * restrict x1, const double * restrict y1, const int N, double * restrict d)
{
	long numcomputed = 0;
	const int block_size = 4;
	for(int i=0;i<N;i++) {
		int j;
		for(j=0;j<=(N-block_size);j+=block_size){
			double *dist = &d[i*N + j];
			for(int jj=0;jj<block_size;jj++) {
				const double *y0 = (const double *) (&y1[(j+jj)*NDIM]);
				const double *x0 = (const double *) (&x1[i*NDIM]);
				const double dx = x0[0] - y0[0];
				const double dy = x0[1] - y0[1];
				const double dz = x0[2] - y0[2];
				dist[jj] = dx*dx + dy*dy + dz*dz;
			}
			numcomputed+=block_size;
		}
		double *dist = (double *) &d[i*N + j];
		for(int jj=0;j<N;jj++,j++) {
			const double *y0 = (const double *) (&y1[j*NDIM]);
			const double *x0 = (const double *) (&x1[i*NDIM]);
			const double dx = x0[0] - y0[0];
			const double dy = x0[1] - y0[1];
			const double dz = x0[2] - y0[2];
			dist[jj] = dx*dx + dy*dy + dz*dz;
			numcomputed++;
		}
	}

#ifdef SQRT_DIST
	for(long index=0;index<numcomputed;index++){
		d[index] = sqrt(d[index]);
	}
#endif

	return numcomputed;
}

long blas_computed(const double * restrict x1, const double * restrict y1, const int N, double * restrict d)
{
	long numcomputed = (long) N * (long) N;
	const double alpha = -2.0;
	/* struct timeval t0,t1; */
	/* gettimeofday(&t0,NULL); */
	cblas_dgemm(CblasRowMajor,CblasNoTrans, CblasTrans,
							N, N, NDIM,
							alpha, x1, NDIM , y1, NDIM, 0.0, d, N);
	/* gettimeofday(&t1,NULL); */
	/* printf(" \t\t cblas_time = %6.2lf ms\n",1e3*(t1.tv_sec-t0.tv_sec) + 1e-3*(t1.tv_usec - t0.tv_usec)); */
	double *x0 = (double *) x1;
	for(int i=0;i<N;i++){
		const double sqr_x0_norm = x0[0]*x0[0] + x0[1]*x0[1] + x0[2]*x0[2];
		for(int j=0;j<N;j++) {
			const double *y0 = (const double *) &y1[j*NDIM];
			const double sqr_y0_norm = y0[0]*y0[0] + y0[1]*y0[1] + y0[2]*y0[2];
			double *dist = (double *) &d[i*N + j];
			*dist += (sqr_x0_norm + sqr_y0_norm);

#ifdef SQRT_DIST
			*dist = sqrt(*dist);
#endif
		}
		x0 += NDIM;
	}
	return numcomputed;
}

long check_result(double *dist, const double *x, const double *y, const int N)
{
	int bad=0;
	int numbadprinted=0;
	for(int i=0;i<N;i++) {
		for(int j=0;j<N;j++) {
			const double *tmp_x = (const double *) &(x[i*NDIM]);
			const double *tmp_y = (const double *) &(y[j*NDIM]);
			double sqr_ds = 0.0;
			for(int k=0;k<NDIM;k++) {
				const double dx = tmp_x[k] - tmp_y[k];
				sqr_ds += dx*dx;
			}
			const long index = i*(long) N + j;
			const double this_dist = dist[index];

#ifdef SQRT_DIST
			const double d = sqrt(sqr_ds);
#else
			const double d = sqr_ds;
#endif

			if(fabs(d - this_dist) > 1e-8) {
				if(numbadprinted < 20) {
					fprintf(stderr,"i = %d j = %d d = %0.8lf this_dist = %0.8lf \n",i,j,d,this_dist);
					numbadprinted++;
				}
				bad++;
			}
		}
	}

	return bad;
}

int main(int argc, char **argv)
{
	int numpart = 0;
	if(argc > 1 ) {
		numpart = atoi(argv[1]);
	} else {
		numpart = NELEMENTS;
	}

	const size_t numbytes = NDIM*numpart;
	double *x = calloc(sizeof(*x), numbytes);
	double *y = calloc(sizeof(*y), numbytes);
	const long totnpairs = (long) numpart * (long) numpart;
	double *dist = calloc(sizeof(*dist), totnpairs);
	assert(x != NULL && y != NULL && dist != NULL && "memory allocation failed");
  const char allfunction_names[][MAXLEN] = {"naive","chunked","compiler_vectorized_chunked","blas_computed"};
	const int ntests = sizeof(allfunction_names)/(sizeof(char)*MAXLEN);
	long (*allfunctions[]) (const double * restrict, const double * restrict, const int, double * restrict)      = {naive, chunked,compiler_vectorized_chunked,blas_computed};

	srand(seed);

	double function_best_mean_time[ntests],function_sigma_time[ntests],function_best_time_in_ms[ntests],function_best_mcycles[ntests];
	int function_niterations[ntests];
	for(int i=0;i<ntests;i++) {
		function_best_mean_time[i] = 1e16;
		function_sigma_time[i] = 0.0;
		function_best_time_in_ms[i] = 1e16;
		function_best_mcycles[i] = 1e16;
	}
#ifndef SQRT_DIST
	const long totflop = (long) numpart * (long) numpart * (8);
#else
	const long totflop = (long) numpart * (long) numpart * (8 + 10); //some hand-wavy thing saying that sqrt is 10 flop
#endif

	if(clustered_data == 0) {
		fill_array(x, numpart);
		fill_array(y, numpart);
	} else {
		assert(numpart <= max_galaxy_in_source_file && "Clustered data does not contain enough galaxies..please reduce numpart in defs.h");
		assert(NDIM == 3 && "Clustered galaxy data contains *exactly* 3 spatial dimensions");
		read_ascii(x, numpart, source_galaxy_file);
		read_ascii(y, numpart, source_galaxy_file);
	}

	const int64_t totniterations = repeat*ntests*(int64_t) max_niterations;
	int64_t numdone = 0;
	int interrupted=0;

	fprintf(stderr,"# Running benchmarks with N = %05d particles\n",numpart);
	init_my_progressbar(totniterations, &interrupted);
	for(int irep=0;irep<repeat;irep++) {
		for(int i=0;i<ntests;i++) {
			struct timeval t0,t1;
			double sum_x=0.0, sum_sqr_x=0.0;

			//warm-up
			gettimeofday(&t0,NULL);
			long ncomputed = (allfunctions[i]) (x, y, numpart, dist);
			gettimeofday(&t1,NULL);

			//check-result
			long numbad = check_result(dist, x, y, numpart);
			if(numbad != 0 || ncomputed != totnpairs) {
				fprintf(stderr,"ERROR: In function `%s' Number of incorrectly calculated distances = %ld out of a total of (%ld) possible pairs. Npairs computed = %ld\n", allfunction_names[i],numbad, totnpairs, ncomputed);
				goto cleanup;
			}

			double best_time_in_ms=1e16, best_time_in_megacycles=1e16;
			uint64_t start_cycles, end_cycles;
			const int64_t numdone_before_iter_loop = numdone;
			for(int iter=0;iter<max_niterations;iter++) {
				gettimeofday(&t0,NULL);
				start_cycles = rdtsc();
				(allfunctions[i]) (x, y, numpart, dist);
				end_cycles = rdtsc();
				gettimeofday(&t1,NULL);

				numdone++;
				my_progressbar(numdone,&interrupted);

				const double this_time_in_ms = 1e3*(t1.tv_sec-t0.tv_sec) +  1e-3*(t1.tv_usec - t0.tv_usec);
				if(this_time_in_ms < best_time_in_ms) best_time_in_ms = this_time_in_ms;
				const double this_time_in_megacycles = (end_cycles - start_cycles)/(1024.*1024.);
				if(this_time_in_megacycles < best_time_in_megacycles) best_time_in_megacycles = this_time_in_megacycles;

				sum_sqr_x += this_time_in_ms*this_time_in_ms;
				sum_x     += this_time_in_ms;

				if(max_niterations <= 10) {
					printf("     %-35s  %0.2lf \n",allfunction_names[i], this_time_in_ms);
				}

				if(best_time_in_ms < function_best_time_in_ms[i]) {
					function_best_time_in_ms[i] = best_time_in_ms;
				}

				if(best_time_in_megacycles < function_best_mcycles[i]) {
					function_best_mcycles[i] = best_time_in_megacycles;
				}

				const double mean_time  = sum_x/(iter+1.0);
				const double sigma_time = sqrt(sum_sqr_x/(iter+1.0) - mean_time*mean_time);
				if(mean_time < function_best_mean_time[i]) {
					function_best_mean_time[i] = mean_time;
				}
				function_niterations[i] = iter + 1;
				function_sigma_time[i] = sigma_time;

        //If the std.dev is small compared to typical runtime and
				//the code has run for more than XXX milli-seconds, then break
				if(sigma_time/mean_time < convergence_ratio && sum_x >= 300 && iter >= 5) {
					numdone = numdone_before_iter_loop + max_niterations;
					break;
				}
			}
		}//i loop over ntests
	}//irep loop over nrepeats
  finish_myprogressbar(&interrupted);

	printf("#######################################################\n");
	printf("##  Function                            Time (ms)      \n");
	printf("#######################################################\n");
	for(int i=0;i<ntests;i++) {
		const double mean_time = function_best_mean_time[i];
		const double sigma_time = function_sigma_time[i];
		const double gflops = (double) totflop/(1e-3*mean_time)/1e9;
		const double best_time_in_ms = function_best_time_in_ms[i];
		const double best_time_in_megacycles = function_best_mcycles[i];
		const int actual_niterations = function_niterations[i];
		printf(ANSI_COLOR_RED "# %-35s  %6.2lf +- %5.2lf " ANSI_COLOR_GREEN " (best -- %6.2lf ms, %6.2lf Mcycles) " ANSI_COLOR_RESET "," ANSI_COLOR_BLUE " %5.2lf GFlops [%04d iterations]" ANSI_COLOR_RESET "\n",
					 allfunction_names[i], mean_time, sigma_time, best_time_in_ms, best_time_in_megacycles, gflops, actual_niterations);

	}

cleanup:
	{
		free((void *) x);free((void *) y);free((void *) dist);
		exit(EXIT_FAILURE);
	}

	exit(EXIT_SUCCESS);
}


/* tsqr.c
 *
 * Copyright (C) 2015, 2019 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.
 */

/*
 * This module implements the sequential TSQR algorithm
 * described in
 *
 * [1] Demmel, J., Grigori, L., Hoemmen, M. F., and Langou, J.
 *     "Communication-optimal parallel and sequential QR and LU factorizations",
 *     UCB Technical Report No. UCB/EECS-2008-89, 2008.
 *
 * The algorithm operates on a tall least squares system:
 *
 * [ A_1 ] x = [ b_1 ]
 * [ A_2 ]     [ b_2 ]
 * [ ... ]     [ ... ]
 * [ A_k ]     [ b_k ]
 *
 * as follows:
 *
 * 1. Initialize
 *    a. [Q_1,R_1] = qr(A_1)
 *    b. z_1 = Q_1^T b_1
 * 2. Loop i = 2:k
 *    a. [Q_i,R_i] = qr( [ R_{i-1} ; A_i ] )
 *    b. z_i = Q_i^T [ z_{i-1} ; b_i ]
 * 3. Output:
 *    a. R = R_k
 *    b. Q^T b = z_k
 *
 * Step 2(a) is optimized to take advantage
 * of the sparse structure of the matrix
 */

#include <config.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_multilarge.h>

typedef struct
{
  size_t p;             /* number of columns of LS matrix */
  int nblocks;          /* number of blocks processed */
  double rnorm;         /* || b - A x || residual norm */
  int svd;              /* SVD of R factor has been computed */

  gsl_matrix *T;        /* block reflector matrix, p-by-p */
  gsl_matrix *R;        /* R factor, p-by-p */
  gsl_vector *QTb;      /* [ Q^T b ; b_i ], size p-by-1 */
  gsl_vector *work;     /* workspace, size p */
  gsl_vector *work3;    /* workspace, size 3*p */

  gsl_multifit_linear_workspace * multifit_workspace_p;
} tsqr_state_t;

static void *tsqr_alloc(const size_t p);
static void tsqr_free(void *vstate);
static int tsqr_reset(void *vstate);
static int tsqr_accumulate(gsl_matrix * A, gsl_vector * b,
                           void * vstate);
static int tsqr_solve(const double lambda, gsl_vector * x,
                      double * rnorm, double * snorm,
                      void * vstate);
static int tsqr_rcond(double * rcond, void * vstate);
static int tsqr_lcurve(gsl_vector * reg_param, gsl_vector * rho,
                       gsl_vector * eta, void * vstate);
static const gsl_matrix * tsqr_R(const void * vstate);
static const gsl_vector * tsqr_QTb(const void * vstate);
static int tsqr_svd(tsqr_state_t * state);

/*
tsqr_alloc()
  Allocate workspace for solving large linear least squares
problems using the TSQR approach

Inputs: p - number of columns of LS matrix

Return: pointer to workspace
*/

static void *
tsqr_alloc(const size_t p)
{
  tsqr_state_t *state;

  if (p == 0)
    {
      GSL_ERROR_NULL("p must be a positive integer",
                     GSL_EINVAL);
    }

  state = calloc(1, sizeof(tsqr_state_t));
  if (!state)
    {
      GSL_ERROR_NULL("failed to allocate tsqr state", GSL_ENOMEM);
    }

  state->p = p;
  state->nblocks = 0;
  state->rnorm = 0.0;

  state->R = gsl_matrix_alloc(p, p);
  if (state->R == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate R matrix", GSL_ENOMEM);
    }

  state->QTb = gsl_vector_alloc(p);
  if (state->QTb == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate QTb vector", GSL_ENOMEM);
    }

  state->T = gsl_matrix_alloc(p, p);
  if (state->T == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate T matrix", GSL_ENOMEM);
    }

  state->work = gsl_vector_alloc(p);
  if (state->work == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate workspace vector", GSL_ENOMEM);
    }

  state->work3 = gsl_vector_alloc(3 * p);
  if (state->work3 == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate work3 vector", GSL_ENOMEM);
    }

  state->multifit_workspace_p = gsl_multifit_linear_alloc(p, p);
  if (state->multifit_workspace_p == NULL)
    {
      tsqr_free(state);
      GSL_ERROR_NULL("failed to allocate multifit workspace", GSL_ENOMEM);
    }

  return state;
}

static void
tsqr_free(void *vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;

  if (state->R)
    gsl_matrix_free(state->R);

  if (state->QTb)
    gsl_vector_free(state->QTb);

  if (state->T)
    gsl_matrix_free(state->T);

  if (state->work)
    gsl_vector_free(state->work);

  if (state->work3)
    gsl_vector_free(state->work3);

  if (state->multifit_workspace_p)
    gsl_multifit_linear_free(state->multifit_workspace_p);

  free(state);
}

static int
tsqr_reset(void *vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;

  gsl_matrix_set_zero(state->R);
  gsl_vector_set_zero(state->QTb);
  state->nblocks = 0;
  state->rnorm = 0.0;
  state->svd = 0;

  return GSL_SUCCESS;
}

/*
tsqr_accumulate()
  Add a new block of rows to the QR system

Inputs: A      - new block of rows, n-by-p
        b      - new rhs vector n-by-1
        vstate - workspace

Return: success/error

Notes:
1) On output, the upper triangular portion of state->R(1:p,1:p)
contains current R matrix

2) state->QTb(1:p) contains current Q^T b vector

3) A and b are destroyed
*/

static int
tsqr_accumulate(gsl_matrix * A, gsl_vector * b, void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;
  const size_t n = A->size1;
  const size_t p = A->size2;

  if (p != state->p)
    {
      GSL_ERROR("columns of A do not match workspace", GSL_EBADLEN);
    }
  else if (n != b->size)
    {
      GSL_ERROR("A and b have different numbers of rows", GSL_EBADLEN);
    }
  else if (state->nblocks == 0 && n < p)
    {
      GSL_ERROR ("n must be >= p", GSL_EBADLEN);
    }
  else if (state->nblocks == 0)
    {
      int status;
      gsl_matrix_view R = gsl_matrix_submatrix(A, 0, 0, p, p);
      gsl_vector_view QTb = gsl_vector_subvector(state->QTb, 0, p);
      gsl_vector_view b1 = gsl_vector_subvector(b, 0, p);

      /* this is the first matrix block A_1, compute its (dense) QR decomposition */

      /* compute QR decomposition of A */
      status = gsl_linalg_QR_decomp_r(A, state->T);
      if (status)
        return status;

      /* store upper triangular R factor in state->R */
      gsl_matrix_tricpy(CblasUpper, CblasNonUnit, state->R, &R.matrix);

      /* compute Q^T b and keep the first p elements */
      gsl_linalg_QR_QTvec_r(A, state->T, b, state->work);
      gsl_vector_memcpy(&QTb.vector, &b1.vector);

      if (n > p)
        {
          gsl_vector_view b2 = gsl_vector_subvector(b, p, n - p);
          state->rnorm = gsl_blas_dnrm2(&b2.vector);
        }
      else
        state->rnorm = 0.0;

      state->nblocks = 1;

      return GSL_SUCCESS;
    }
  else
    {
      int status;

      /* compute QR decomposition of [ R_{i-1} ; A_i ], accounting for
       * sparse structure */
      status = gsl_linalg_QR_UR_decomp(state->R, A, state->T);
      if (status)
        return status;

      /*
       * Compute:
       *
       * Q^T [ QTb_{i-1} ] = [ QTb_{i-1} - w ]
       *     [    b_i    ]   [  b_i - V~ w   ]
       *
       * where:
       *
       * w = T^T (QTb_{i-1} + V~^T b_i)
       *
       *       p
       * V = [ I  ] p
       *     [ V~ ] n
       */
      gsl_vector_memcpy(state->work, state->QTb);
      gsl_blas_dgemv(CblasTrans, 1.0, A, b, 1.0, state->work);                     /* w := w + V~^T b */
      gsl_blas_dtrmv(CblasUpper, CblasTrans, CblasNonUnit, state->T, state->work); /* w := T^T w */
      gsl_vector_sub(state->QTb, state->work);                                     /* QTb := QTb - w */

      /* update residual norm */
      gsl_blas_dgemv(CblasNoTrans, -1.0, A, state->work, 1.0, b);                  /* b := b - V~ w */
      state->rnorm = gsl_hypot(state->rnorm, gsl_blas_dnrm2(b));

      return GSL_SUCCESS;
    }
}

/*
tsqr_solve()
  Solve the least squares system:

chi^2 = || QTb - R x ||^2 + lambda^2 || x ||^2

using the SVD of R

Inputs: lambda - regularization parameter
        x      - (output) solution vector p-by-1
        rnorm  - (output) residual norm ||b - A x||
        snorm  - (output) solution norm ||x||
        vstate - workspace

Return: success/error
*/

static int
tsqr_solve(const double lambda, gsl_vector * x,
           double * rnorm, double * snorm,
           void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;

  if (x->size != state->p)
    {
      GSL_ERROR ("solution vector does not match workspace", GSL_EBADLEN);
    }
  else if (lambda < 0.0)
    {
      GSL_ERROR ("regularization parameter should be non-negative", GSL_EINVAL);
    }
  else
    {
      if (lambda == 0.0)
        {
          /* solve: R x = Q^T b */
          gsl_vector_memcpy(x, state->QTb);
          gsl_blas_dtrsv(CblasUpper, CblasNoTrans, CblasNonUnit, state->R, x);
          *rnorm = state->rnorm;
          *snorm = gsl_blas_dnrm2(x);
        }
      else
        {
          int status;

          /* compute SVD of R if not already computed */
          if (state->svd == 0)
            {
              status = tsqr_svd(state);
              if (status)
                return status;
            }

          status = gsl_multifit_linear_solve(lambda, state->R, state->QTb, x, rnorm, snorm,
                                             state->multifit_workspace_p);
          if (status)
            return status;

          *rnorm = gsl_hypot(*rnorm, state->rnorm);
        }

      return GSL_SUCCESS;
    }
}

/*
tsqr_lcurve()
  Compute L-curve of least squares system

Inputs: reg_param - (output) vector of regularization parameters
        rho       - (output) vector of residual norms
        eta       - (output) vector of solution norms
        vstate    - workspace

Return: success/error
*/

static int
tsqr_lcurve(gsl_vector * reg_param, gsl_vector * rho,
            gsl_vector * eta, void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;
  int status;
  size_t i;

  /* compute SVD of R if not already computed */
  if (state->svd == 0)
    {
      status = tsqr_svd(state);
      if (status)
        return status;
    }

  status = gsl_multifit_linear_lcurve(state->QTb, reg_param, rho, eta,
                                      state->multifit_workspace_p);

  /* now add contribution to rnorm from Q2 factor */
  for (i = 0; i < rho->size; ++i)
    {
      double *rhoi = gsl_vector_ptr(rho, i);
      *rhoi = gsl_hypot(*rhoi, state->rnorm);
    }

  return status;
}

static const gsl_matrix *
tsqr_R(const void * vstate)
{
  const tsqr_state_t *state = (const tsqr_state_t *) vstate;
  return state->R;
}

static const gsl_vector *
tsqr_QTb(const void * vstate)
{
  const tsqr_state_t *state = (const tsqr_state_t *) vstate;
  return state->QTb;
}

static int
tsqr_rcond(double * rcond, void * vstate)
{
  tsqr_state_t *state = (tsqr_state_t *) vstate;
  return gsl_linalg_tri_rcond(CblasUpper, state->R, rcond, state->work3);
}

/*
tsqr_svd()
  Compute the SVD of the upper triangular
R factor. This allows us to compute the upper/lower
bounds on the regularization parameter and compute
the matrix reciprocal condition number.

Inputs: state - workspace

Return: success/error
*/

static int
tsqr_svd(tsqr_state_t * state)
{
  int status;

  status = gsl_multifit_linear_svd(state->R, state->multifit_workspace_p);
  if (status)
    {
      GSL_ERROR("error computing SVD of R", status);
    }

  state->svd = 1;

  return GSL_SUCCESS;
}

static const gsl_multilarge_linear_type tsqr_type =
{
  "tsqr",
  tsqr_alloc,
  tsqr_reset,
  tsqr_accumulate,
  tsqr_solve,
  tsqr_rcond,
  tsqr_lcurve,
  tsqr_R,
  tsqr_QTb,
  tsqr_free
};

const gsl_multilarge_linear_type * gsl_multilarge_linear_tsqr = &tsqr_type;


// Copyright (c) 2017, Lawrence Livermore National Security, LLC. Produced at
// the Lawrence Livermore National Laboratory. LLNL-CODE-734707. All Rights
// reserved. See files LICENSE and NOTICE for details.
//
// This file is part of CEED, a collection of benchmarks, miniapps, software
// libraries and APIs for efficient high-order finite element and spectral
// element discretizations for exascale applications. For more information and
// source code availability see http://github.com/ceed.
//
// The CEED research is supported by the Exascale Computing Project 17-SC-20-SC,
// a collaborative effort of two U.S. Department of Energy organizations (Office
// of Science and the National Nuclear Security Administration) responsible for
// the planning and preparation of a capable exascale ecosystem, including
// software, applications, hardware, advanced system engineering and early
// testbed platforms, in support of the nation's exascale computing imperative.

//                        libCEED + PETSc Example: CEED BPs 3-6 with Multigrid
//
// This example demonstrates a simple usage of libCEED with PETSc to solve the
// CEED BP benchmark problems, see http://ceed.exascaleproject.org/bps.
//
// The code uses higher level communication protocols in DMPlex.
//
// Build with:
//
//     make multigrid [PETSC_DIR=</path/to/petsc>] [CEED_DIR=</path/to/libceed>]
//
// Sample runs:
//
//     multigrid -problem bp3
//     multigrid -problem bp4
//     multigrid -problem bp5 -ceed /cpu/self
//     multigrid -problem bp6 -ceed /gpu/cuda
//
//TESTARGS -ceed {ceed_resource} -test -problem bp3 -degree 3

/// @file
/// CEED BPs 1-6 multigrid example using PETSc
const char help[] = "Solve CEED BPs using p-multigrid with PETSc and DMPlex\n";

#include <stdbool.h>
#include <string.h>
#include <ceed.h>
#include <petsc.h>
#include <petscdmplex.h>
#include <petscksp.h>
#include <petscsys.h>

#include "bps.h"
#include "include/bpsproblemdata.h"
#include "include/petscmacros.h"
#include "include/petscutils.h"
#include "include/matops.h"
#include "include/structs.h"
#include "include/libceedsetup.h"

#if PETSC_VERSION_LT(3,12,0)
#ifdef PETSC_HAVE_CUDA
#include <petsccuda.h>
// Note: With PETSc prior to version 3.12.0, providing the source path to
//       include 'cublas_v2.h' will be needed to use 'petsccuda.h'.
#endif
#endif

int main(int argc, char **argv) {
  PetscInt ierr;
  MPI_Comm comm;
  char filename[PETSC_MAX_PATH_LEN],
       ceed_resource[PETSC_MAX_PATH_LEN] = "/cpu/self";
  double my_rt_start, my_rt, rt_min, rt_max;
  PetscInt degree = 3, q_extra, *l_size, *xl_size, *g_size, dim = 3, fine_level,
           mesh_elem[3] = {3, 3, 3}, num_comp_u = 1, num_levels = degree, *level_degrees;
  PetscScalar *r;
  PetscScalar eps = 1.0;
  PetscBool test_mode, benchmark_mode, read_mesh, write_solution;
  PetscLogStage solve_stage;
  DM  *dm, dm_orig;
  SNES snes_dummy;
  KSP ksp;
  PC pc;
  Mat *mat_O, *mat_pr, mat_coarse;
  Vec *X, *X_loc, *mult, rhs, rhs_loc;
  PetscMemType mem_type;
  UserO *user_O;
  UserProlongRestr *user_pr;
  Ceed ceed;
  CeedData *ceed_data;
  CeedVector rhs_ceed, target;
  CeedQFunction qf_error, qf_restrict, qf_prolong;
  CeedOperator op_error;
  BPType bp_choice;
  CoarsenType coarsen;

  ierr = PetscInitialize(&argc, &argv, NULL, help);
  if (ierr) return ierr;
  comm = PETSC_COMM_WORLD;

  // Parse command line options
  ierr = PetscOptionsBegin(comm, NULL, "CEED BPs in PETSc", NULL); CHKERRQ(ierr);
  bp_choice = CEED_BP3;
  ierr = PetscOptionsEnum("-problem",
                          "CEED benchmark problem to solve", NULL,
                          bp_types, (PetscEnum)bp_choice, (PetscEnum *)&bp_choice,
                          NULL); CHKERRQ(ierr);
  num_comp_u = bp_options[bp_choice].num_comp_u;
  test_mode = PETSC_FALSE;
  ierr = PetscOptionsBool("-test",
                          "Testing mode (do not print unless error is large)",
                          NULL, test_mode, &test_mode, NULL); CHKERRQ(ierr);
  benchmark_mode = PETSC_FALSE;
  ierr = PetscOptionsBool("-benchmark",
                          "Benchmarking mode (prints benchmark statistics)",
                          NULL, benchmark_mode, &benchmark_mode, NULL);
  CHKERRQ(ierr);
  write_solution = PETSC_FALSE;
  ierr = PetscOptionsBool("-write_solution",
                          "Write solution for visualization",
                          NULL, write_solution, &write_solution, NULL);
  CHKERRQ(ierr);
  ierr = PetscOptionsScalar("-eps",
                            "Epsilon parameter for Kershaw mesh transformation",
                            NULL, eps, &eps, NULL);
  if (eps > 1 || eps <= 0) SETERRQ1(PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE,
                                      "-eps %D must be (0,1]", eps);
  degree = test_mode ? 3 : 2;
  ierr = PetscOptionsInt("-degree", "Polynomial degree of tensor product basis",
                         NULL, degree, &degree, NULL); CHKERRQ(ierr);
  if (degree < 1) SETERRQ1(PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE,
                             "-degree %D must be at least 1", degree);
  q_extra = bp_options[bp_choice].q_extra;
  ierr = PetscOptionsInt("-q_extra", "Number of extra quadrature points",
                         NULL, q_extra, &q_extra, NULL); CHKERRQ(ierr);
  ierr = PetscOptionsString("-ceed", "CEED resource specifier",
                            NULL, ceed_resource, ceed_resource,
                            sizeof(ceed_resource), NULL); CHKERRQ(ierr);
  coarsen = COARSEN_UNIFORM;
  ierr = PetscOptionsEnum("-coarsen",
                          "Coarsening strategy to use", NULL,
                          coarsen_types, (PetscEnum)coarsen,
                          (PetscEnum *)&coarsen, NULL); CHKERRQ(ierr);
  read_mesh = PETSC_FALSE;
  ierr = PetscOptionsString("-mesh", "Read mesh from file", NULL,
                            filename, filename, sizeof(filename), &read_mesh);
  CHKERRQ(ierr);
  if (!read_mesh) {
    PetscInt tmp = dim;
    ierr = PetscOptionsIntArray("-cells","Number of cells per dimension", NULL,
                                mesh_elem, &tmp, NULL); CHKERRQ(ierr);
  }
  ierr = PetscOptionsEnd(); CHKERRQ(ierr);

  // Set up libCEED
  CeedInit(ceed_resource, &ceed);
  CeedMemType mem_type_backend;
  CeedGetPreferredMemType(ceed, &mem_type_backend);

  // Setup DM
  if (read_mesh) {
    ierr = DMPlexCreateFromFile(PETSC_COMM_WORLD, filename, PETSC_TRUE, &dm_orig);
    CHKERRQ(ierr);
  } else {
    ierr = DMPlexCreateBoxMesh(PETSC_COMM_WORLD, dim, PETSC_FALSE, mesh_elem, NULL,
                               NULL, NULL, PETSC_TRUE, &dm_orig); CHKERRQ(ierr);
  }

  {
    DM dm_dist = NULL;
    PetscPartitioner part;

    ierr = DMPlexGetPartitioner(dm_orig, &part); CHKERRQ(ierr);
    ierr = PetscPartitionerSetFromOptions(part); CHKERRQ(ierr);
    ierr = DMPlexDistribute(dm_orig, 0, NULL, &dm_dist); CHKERRQ(ierr);
    if (dm_dist) {
      ierr = DMDestroy(&dm_orig); CHKERRQ(ierr);
      dm_orig = dm_dist;
    }
  }

  // Apply Kershaw mesh transformation
  ierr = Kershaw(dm_orig, eps); CHKERRQ(ierr);

  VecType vec_type;
  switch (mem_type_backend) {
  case CEED_MEM_HOST: vec_type = VECSTANDARD; break;
  case CEED_MEM_DEVICE: {
    const char *resolved;
    CeedGetResource(ceed, &resolved);
    if (strstr(resolved, "/gpu/cuda")) vec_type = VECCUDA;
    else if (strstr(resolved, "/gpu/hip/occa"))
      vec_type = VECSTANDARD; // https://github.com/CEED/libCEED/issues/678
    else if (strstr(resolved, "/gpu/hip")) vec_type = VECHIP;
    else vec_type = VECSTANDARD;
  }
  }
  ierr = DMSetVecType(dm_orig, vec_type); CHKERRQ(ierr);
  ierr = DMSetFromOptions(dm_orig); CHKERRQ(ierr);

  // Allocate arrays for PETSc objects for each level
  switch (coarsen) {
  case COARSEN_UNIFORM:
    num_levels = degree;
    break;
  case COARSEN_LOGARITHMIC:
    num_levels = ceil(log(degree)/log(2)) + 1;
    break;
  }
  ierr = PetscMalloc1(num_levels, &level_degrees); CHKERRQ(ierr);
  fine_level = num_levels - 1;

  switch (coarsen) {
  case COARSEN_UNIFORM:
    for (int i=0; i<num_levels; i++) level_degrees[i] = i + 1;
    break;
  case COARSEN_LOGARITHMIC:
    for (int i=0; i<num_levels - 1; i++) level_degrees[i] = pow(2,i);
    level_degrees[fine_level] = degree;
    break;
  }
  ierr = PetscMalloc1(num_levels, &dm); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &X); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &X_loc); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &mult); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &user_O); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &user_pr); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &mat_O); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &mat_pr); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &l_size); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &xl_size); CHKERRQ(ierr);
  ierr = PetscMalloc1(num_levels, &g_size); CHKERRQ(ierr);

  // Setup DM and Operator Mat Shells for each level
  for (CeedInt i=0; i<num_levels; i++) {
    // Create DM
    ierr = DMClone(dm_orig, &dm[i]); CHKERRQ(ierr);
    ierr = DMGetVecType(dm_orig, &vec_type); CHKERRQ(ierr);
    ierr = DMSetVecType(dm[i], vec_type); CHKERRQ(ierr);
    PetscInt dim;
    ierr = DMGetDimension(dm[i], &dim); CHKERRQ(ierr);
    ierr = SetupDMByDegree(dm[i], level_degrees[i], num_comp_u, dim,
                           bp_options[bp_choice].enforce_bc, bp_options[bp_choice].bc_func);
    CHKERRQ(ierr);

    // Create vectors
    ierr = DMCreateGlobalVector(dm[i], &X[i]); CHKERRQ(ierr);
    ierr = VecGetLocalSize(X[i], &l_size[i]); CHKERRQ(ierr);
    ierr = VecGetSize(X[i], &g_size[i]); CHKERRQ(ierr);
    ierr = DMCreateLocalVector(dm[i], &X_loc[i]); CHKERRQ(ierr);
    ierr = VecGetSize(X_loc[i], &xl_size[i]); CHKERRQ(ierr);

    // Operator
    ierr = PetscMalloc1(1, &user_O[i]); CHKERRQ(ierr);
    ierr = MatCreateShell(comm, l_size[i], l_size[i], g_size[i], g_size[i],
                          user_O[i], &mat_O[i]); CHKERRQ(ierr);
    ierr = MatShellSetOperation(mat_O[i], MATOP_MULT,
                                (void(*)(void))MatMult_Ceed); CHKERRQ(ierr);
    ierr = MatShellSetOperation(mat_O[i], MATOP_GET_DIAGONAL,
                                (void(*)(void))MatGetDiag); CHKERRQ(ierr);
    ierr = MatShellSetVecType(mat_O[i], vec_type); CHKERRQ(ierr);

    // Level transfers
    if (i > 0) {
      // Interp
      ierr = PetscMalloc1(1, &user_pr[i]); CHKERRQ(ierr);
      ierr = MatCreateShell(comm, l_size[i], l_size[i-1], g_size[i], g_size[i-1],
                            user_pr[i], &mat_pr[i]); CHKERRQ(ierr);
      ierr = MatShellSetOperation(mat_pr[i], MATOP_MULT,
                                  (void(*)(void))MatMult_Prolong);
      CHKERRQ(ierr);
      ierr = MatShellSetOperation(mat_pr[i], MATOP_MULT_TRANSPOSE,
                                  (void(*)(void))MatMult_Restrict);
      CHKERRQ(ierr);
      ierr = MatShellSetVecType(mat_pr[i], vec_type); CHKERRQ(ierr);
    }
  }
  ierr = VecDuplicate(X[fine_level], &rhs); CHKERRQ(ierr);

  // Print global grid information
  if (!test_mode) {
    PetscInt P = degree + 1, Q = P + q_extra;

    const char *used_resource;
    CeedGetResource(ceed, &used_resource);

    ierr = VecGetType(X[0], &vec_type); CHKERRQ(ierr);

    ierr = PetscPrintf(comm,
                       "\n-- CEED Benchmark Problem %d -- libCEED + PETSc + PCMG --\n"
                       "  PETSc:\n"
                       "    PETSc Vec Type                     : %s\n"
                       "  libCEED:\n"
                       "    libCEED Backend                    : %s\n"
                       "    libCEED Backend MemType            : %s\n"
                       "  Mesh:\n"
                       "    Number of 1D Basis Nodes (p)       : %d\n"
                       "    Number of 1D Quadrature Points (q) : %d\n"
                       "    Global Nodes                       : %D\n"
                       "    Owned Nodes                        : %D\n"
                       "    DoF per node                       : %D\n"
                       "  Multigrid:\n"
                       "    Number of Levels                   : %d\n",
                       bp_choice+1, vec_type, used_resource,
                       CeedMemTypes[mem_type_backend],
                       P, Q, g_size[fine_level]/num_comp_u, l_size[fine_level]/num_comp_u,
                       num_comp_u, num_levels); CHKERRQ(ierr);
  }

  // Create RHS vector
  ierr = VecDuplicate(X_loc[fine_level], &rhs_loc); CHKERRQ(ierr);
  ierr = VecZeroEntries(rhs_loc); CHKERRQ(ierr);
  ierr = VecGetArrayAndMemType(rhs_loc, &r, &mem_type); CHKERRQ(ierr);
  CeedVectorCreate(ceed, xl_size[fine_level], &rhs_ceed);
  CeedVectorSetArray(rhs_ceed, MemTypeP2C(mem_type), CEED_USE_POINTER, r);

  // Set up libCEED operators on each level
  ierr = PetscMalloc1(num_levels, &ceed_data); CHKERRQ(ierr);
  for (int i=0; i<num_levels; i++) {
    // Print level information
    if (!test_mode && (i == 0 || i == fine_level)) {
      ierr = PetscPrintf(comm,"    Level %D (%s):\n"
                         "      Number of 1D Basis Nodes (p)     : %d\n"
                         "      Global Nodes                     : %D\n"
                         "      Owned Nodes                      : %D\n",
                         i, (i? "fine" : "coarse"), level_degrees[i] + 1,
                         g_size[i]/num_comp_u, l_size[i]/num_comp_u); CHKERRQ(ierr);
    }
    ierr = PetscMalloc1(1, &ceed_data[i]); CHKERRQ(ierr);
    ierr = SetupLibceedByDegree(dm[i], ceed, level_degrees[i], dim, q_extra,
                                dim, num_comp_u, g_size[i], xl_size[i], bp_options[bp_choice],
                                ceed_data[i], i==(fine_level), rhs_ceed, &target);
    CHKERRQ(ierr);
  }

  // Gather RHS
  CeedVectorTakeArray(rhs_ceed, MemTypeP2C(mem_type), NULL);
  ierr = VecRestoreArrayAndMemType(rhs_loc, &r); CHKERRQ(ierr);
  ierr = VecZeroEntries(rhs); CHKERRQ(ierr);
  ierr = DMLocalToGlobal(dm[fine_level], rhs_loc, ADD_VALUES, rhs); CHKERRQ(ierr);
  CeedVectorDestroy(&rhs_ceed);

  // Create the restriction/interpolation QFunction
  CeedQFunctionCreateIdentity(ceed, num_comp_u, CEED_EVAL_NONE, CEED_EVAL_INTERP,
                              &qf_restrict);
  CeedQFunctionCreateIdentity(ceed, num_comp_u, CEED_EVAL_INTERP, CEED_EVAL_NONE,
                              &qf_prolong);

  // Set up libCEED level transfer operators
  ierr = CeedLevelTransferSetup(ceed, num_levels, num_comp_u, ceed_data,
                                level_degrees,
                                qf_restrict, qf_prolong); CHKERRQ(ierr);

  // Create the error QFunction
  CeedQFunctionCreateInterior(ceed, 1, bp_options[bp_choice].error,
                              bp_options[bp_choice].error_loc, &qf_error);
  CeedQFunctionAddInput(qf_error, "u", num_comp_u, CEED_EVAL_INTERP);
  CeedQFunctionAddInput(qf_error, "true_soln", num_comp_u, CEED_EVAL_NONE);
  CeedQFunctionAddOutput(qf_error, "error", num_comp_u, CEED_EVAL_NONE);

  // Create the error operator
  CeedOperatorCreate(ceed, qf_error, CEED_QFUNCTION_NONE, CEED_QFUNCTION_NONE,
                     &op_error);
  CeedOperatorSetField(op_error, "u", ceed_data[fine_level]->elem_restr_u,
                       ceed_data[fine_level]->basis_u, CEED_VECTOR_ACTIVE);
  CeedOperatorSetField(op_error, "true_soln",
                       ceed_data[fine_level]->elem_restr_u_i,
                       CEED_BASIS_COLLOCATED, target);
  CeedOperatorSetField(op_error, "error", ceed_data[fine_level]->elem_restr_u_i,
                       CEED_BASIS_COLLOCATED, CEED_VECTOR_ACTIVE);

  // Calculate multiplicity
  for (int i=0; i<num_levels; i++) {
    PetscScalar *x;

    // CEED vector
    ierr = VecZeroEntries(X_loc[i]); CHKERRQ(ierr);
    ierr = VecGetArray(X_loc[i], &x); CHKERRQ(ierr);
    CeedVectorSetArray(ceed_data[i]->x_ceed, CEED_MEM_HOST, CEED_USE_POINTER, x);

    // Multiplicity
    CeedElemRestrictionGetMultiplicity(ceed_data[i]->elem_restr_u,
                                       ceed_data[i]->x_ceed);
    CeedVectorSyncArray(ceed_data[i]->x_ceed, CEED_MEM_HOST);

    // Restore vector
    ierr = VecRestoreArray(X_loc[i], &x); CHKERRQ(ierr);

    // Creat mult vector
    ierr = VecDuplicate(X_loc[i], &mult[i]); CHKERRQ(ierr);

    // Local-to-global
    ierr = VecZeroEntries(X[i]); CHKERRQ(ierr);
    ierr = DMLocalToGlobal(dm[i], X_loc[i], ADD_VALUES, X[i]);
    CHKERRQ(ierr);
    ierr = VecZeroEntries(X_loc[i]); CHKERRQ(ierr);

    // Global-to-local
    ierr = DMGlobalToLocal(dm[i], X[i], INSERT_VALUES, mult[i]);
    CHKERRQ(ierr);
    ierr = VecZeroEntries(X[i]); CHKERRQ(ierr);

    // Multiplicity scaling
    ierr = VecReciprocal(mult[i]);
  }

  // Set up Mat
  for (int i=0; i<num_levels; i++) {
    // User Operator
    user_O[i]->comm = comm;
    user_O[i]->dm = dm[i];
    user_O[i]->X_loc = X_loc[i];
    ierr = VecDuplicate(X_loc[i], &user_O[i]->Y_loc); CHKERRQ(ierr);
    user_O[i]->x_ceed = ceed_data[i]->x_ceed;
    user_O[i]->y_ceed = ceed_data[i]->y_ceed;
    user_O[i]->op = ceed_data[i]->op_apply;
    user_O[i]->ceed = ceed;

    if (i > 0) {
      // Prolongation/Restriction Operator
      user_pr[i]->comm = comm;
      user_pr[i]->dmf = dm[i];
      user_pr[i]->dmc = dm[i-1];
      user_pr[i]->loc_vec_c = X_loc[i-1];
      user_pr[i]->loc_vec_f = user_O[i]->Y_loc;
      user_pr[i]->mult_vec = mult[i];
      user_pr[i]->ceed_vec_c = user_O[i-1]->x_ceed;
      user_pr[i]->ceed_vec_f = user_O[i]->y_ceed;
      user_pr[i]->op_prolong = ceed_data[i]->op_prolong;
      user_pr[i]->op_restrict = ceed_data[i]->op_restrict;
      user_pr[i]->ceed = ceed;
    }
  }

  // Setup dummy SNES for AMG coarse solve
  ierr = SNESCreate(comm, &snes_dummy); CHKERRQ(ierr);
  ierr = SNESSetDM(snes_dummy, dm[0]); CHKERRQ(ierr);
  ierr = SNESSetSolution(snes_dummy, X[0]); CHKERRQ(ierr);

  // -- Jacobian matrix
  ierr = DMSetMatType(dm[0], MATAIJ); CHKERRQ(ierr);
  ierr = DMCreateMatrix(dm[0], &mat_coarse); CHKERRQ(ierr);
  ierr = SNESSetJacobian(snes_dummy, mat_coarse, mat_coarse, NULL,
                         NULL); CHKERRQ(ierr);

  // -- Residual evaluation function
  ierr = SNESSetFunction(snes_dummy, X[0], FormResidual_Ceed,
                         user_O[0]); CHKERRQ(ierr);

  // -- Form Jacobian
  ierr = SNESComputeJacobianDefaultColor(snes_dummy, X[0], mat_O[0],
                                         mat_coarse, NULL); CHKERRQ(ierr);

  // Set up KSP
  ierr = KSPCreate(comm, &ksp); CHKERRQ(ierr);
  {
    ierr = KSPSetType(ksp, KSPCG); CHKERRQ(ierr);
    ierr = KSPSetNormType(ksp, KSP_NORM_NATURAL); CHKERRQ(ierr);
    ierr = KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT,
                            PETSC_DEFAULT); CHKERRQ(ierr);
  }
  ierr = KSPSetFromOptions(ksp); CHKERRQ(ierr);
  ierr = KSPSetOperators(ksp, mat_O[fine_level], mat_O[fine_level]);
  CHKERRQ(ierr);

  // Set up PCMG
  ierr = KSPGetPC(ksp, &pc); CHKERRQ(ierr);
  PCMGCycleType pcmg_cycle_type = PC_MG_CYCLE_V;
  {
    ierr = PCSetType(pc, PCMG); CHKERRQ(ierr);

    // PCMG levels
    ierr = PCMGSetLevels(pc, num_levels, NULL); CHKERRQ(ierr);
    for (int i=0; i<num_levels; i++) {
      // Smoother
      KSP smoother;
      PC smoother_pc;
      ierr = PCMGGetSmoother(pc, i, &smoother); CHKERRQ(ierr);
      ierr = KSPSetType(smoother, KSPCHEBYSHEV); CHKERRQ(ierr);
      ierr = KSPChebyshevEstEigSet(smoother, 0, 0.1, 0, 1.1); CHKERRQ(ierr);
      ierr = KSPChebyshevEstEigSetUseNoisy(smoother, PETSC_TRUE); CHKERRQ(ierr);
      ierr = KSPSetOperators(smoother, mat_O[i], mat_O[i]); CHKERRQ(ierr);
      ierr = KSPGetPC(smoother, &smoother_pc); CHKERRQ(ierr);
      ierr = PCSetType(smoother_pc, PCJACOBI); CHKERRQ(ierr);
      ierr = PCJacobiSetType(smoother_pc, PC_JACOBI_DIAGONAL); CHKERRQ(ierr);

      // Work vector
      if (i < num_levels - 1) {
        ierr = PCMGSetX(pc, i, X[i]); CHKERRQ(ierr);
      }

      // Level transfers
      if (i > 0) {
        // Interpolation
        ierr = PCMGSetInterpolation(pc, i, mat_pr[i]); CHKERRQ(ierr);
      }

      // Coarse solve
      KSP coarse;
      PC coarse_pc;
      ierr = PCMGGetCoarseSolve(pc, &coarse); CHKERRQ(ierr);
      ierr = KSPSetType(coarse, KSPPREONLY); CHKERRQ(ierr);
      ierr = KSPSetOperators(coarse, mat_coarse, mat_coarse); CHKERRQ(ierr);

      ierr = KSPGetPC(coarse, &coarse_pc); CHKERRQ(ierr);
      ierr = PCSetType(coarse_pc, PCGAMG); CHKERRQ(ierr);

      ierr = KSPSetOptionsPrefix(coarse, "coarse_"); CHKERRQ(ierr);
      ierr = PCSetOptionsPrefix(coarse_pc, "coarse_"); CHKERRQ(ierr);
      ierr = KSPSetFromOptions(coarse); CHKERRQ(ierr);
      ierr = PCSetFromOptions(coarse_pc); CHKERRQ(ierr);
    }

    // PCMG options
    ierr = PCMGSetType(pc, PC_MG_MULTIPLICATIVE); CHKERRQ(ierr);
    ierr = PCMGSetNumberSmooth(pc, 3); CHKERRQ(ierr);
    ierr = PCMGSetCycleType(pc, pcmg_cycle_type); CHKERRQ(ierr);
  }

  // First run, if benchmarking
  if (benchmark_mode) {
    ierr = KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, 1);
    CHKERRQ(ierr);
    ierr = VecZeroEntries(X[fine_level]); CHKERRQ(ierr);
    my_rt_start = MPI_Wtime();
    ierr = KSPSolve(ksp, rhs, X[fine_level]); CHKERRQ(ierr);
    my_rt = MPI_Wtime() - my_rt_start;
    ierr = MPI_Allreduce(MPI_IN_PLACE, &my_rt, 1, MPI_DOUBLE, MPI_MIN, comm);
    CHKERRQ(ierr);
    // Set maxits based on first iteration timing
    if (my_rt > 0.02) {
      ierr = KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, 5);
      CHKERRQ(ierr);
    } else {
      ierr = KSPSetTolerances(ksp, 1e-10, PETSC_DEFAULT, PETSC_DEFAULT, 20);
      CHKERRQ(ierr);
    }
  }

  // Timed solve
  ierr = VecZeroEntries(X[fine_level]); CHKERRQ(ierr);
  ierr = PetscBarrier((PetscObject)ksp); CHKERRQ(ierr);

  // -- Performance logging
  ierr = PetscLogStageRegister("Solve Stage", &solve_stage); CHKERRQ(ierr);
  ierr = PetscLogStagePush(solve_stage); CHKERRQ(ierr);

  // -- Solve
  my_rt_start = MPI_Wtime();
  ierr = KSPSolve(ksp, rhs, X[fine_level]); CHKERRQ(ierr);
  my_rt = MPI_Wtime() - my_rt_start;

  // -- Performance logging
  ierr = PetscLogStagePop();

  // Output results
  {
    KSPType ksp_type;
    PCMGType pcmg_type;
    KSPConvergedReason reason;
    PetscReal rnorm;
    PetscInt its;
    ierr = KSPGetType(ksp, &ksp_type); CHKERRQ(ierr);
    ierr = KSPGetConvergedReason(ksp, &reason); CHKERRQ(ierr);
    ierr = KSPGetIterationNumber(ksp, &its); CHKERRQ(ierr);
    ierr = KSPGetResidualNorm(ksp, &rnorm); CHKERRQ(ierr);
    ierr = PCMGGetType(pc, &pcmg_type); CHKERRQ(ierr);
    if (!test_mode || reason < 0 || rnorm > 1e-8) {
      ierr = PetscPrintf(comm,
                         "  KSP:\n"
                         "    KSP Type                           : %s\n"
                         "    KSP Convergence                    : %s\n"
                         "    Total KSP Iterations               : %D\n"
                         "    Final rnorm                        : %e\n",
                         ksp_type, KSPConvergedReasons[reason], its,
                         (double)rnorm); CHKERRQ(ierr);
      ierr = PetscPrintf(comm,
                         "  PCMG:\n"
                         "    PCMG Type                          : %s\n"
                         "    PCMG Cycle Type                    : %s\n",
                         PCMGTypes[pcmg_type],
                         PCMGCycleTypes[pcmg_cycle_type]); CHKERRQ(ierr);
    }
    if (!test_mode) {
      ierr = PetscPrintf(comm,"  Performance:\n"); CHKERRQ(ierr);
    }
    {
      PetscReal max_error;
      ierr = ComputeErrorMax(user_O[fine_level], op_error, X[fine_level], target,
                             &max_error); CHKERRQ(ierr);
      PetscReal tol = 5e-2;
      if (!test_mode || max_error > tol) {
        ierr = MPI_Allreduce(&my_rt, &rt_min, 1, MPI_DOUBLE, MPI_MIN, comm);
        CHKERRQ(ierr);
        ierr = MPI_Allreduce(&my_rt, &rt_max, 1, MPI_DOUBLE, MPI_MAX, comm);
        CHKERRQ(ierr);
        ierr = PetscPrintf(comm,
                           "    Pointwise Error (max)              : %e\n"
                           "    CG Solve Time                      : %g (%g) sec\n",
                           (double)max_error, rt_max, rt_min); CHKERRQ(ierr);
      }
    }
    if (benchmark_mode && (!test_mode)) {
      ierr = PetscPrintf(comm,
                         "    DoFs/Sec in CG                     : %g (%g) million\n",
                         1e-6*g_size[fine_level]*its/rt_max,
                         1e-6*g_size[fine_level]*its/rt_min);
      CHKERRQ(ierr);
    }
  }

  if (write_solution) {
    PetscViewer vtk_viewer_soln;

    ierr = PetscViewerCreate(comm, &vtk_viewer_soln); CHKERRQ(ierr);
    ierr = PetscViewerSetType(vtk_viewer_soln, PETSCVIEWERVTK); CHKERRQ(ierr);
    ierr = PetscViewerFileSetName(vtk_viewer_soln, "solution.vtu"); CHKERRQ(ierr);
    ierr = VecView(X[fine_level], vtk_viewer_soln); CHKERRQ(ierr);
    ierr = PetscViewerDestroy(&vtk_viewer_soln); CHKERRQ(ierr);
  }

  // Cleanup
  for (int i=0; i<num_levels; i++) {
    ierr = VecDestroy(&X[i]); CHKERRQ(ierr);
    ierr = VecDestroy(&X_loc[i]); CHKERRQ(ierr);
    ierr = VecDestroy(&mult[i]); CHKERRQ(ierr);
    ierr = VecDestroy(&user_O[i]->Y_loc); CHKERRQ(ierr);
    ierr = MatDestroy(&mat_O[i]); CHKERRQ(ierr);
    ierr = PetscFree(user_O[i]); CHKERRQ(ierr);
    if (i > 0) {
      ierr = MatDestroy(&mat_pr[i]); CHKERRQ(ierr);
      ierr = PetscFree(user_pr[i]); CHKERRQ(ierr);
    }
    ierr = CeedDataDestroy(i, ceed_data[i]); CHKERRQ(ierr);
    ierr = DMDestroy(&dm[i]); CHKERRQ(ierr);
  }
  ierr = PetscFree(level_degrees); CHKERRQ(ierr);
  ierr = PetscFree(dm); CHKERRQ(ierr);
  ierr = PetscFree(X); CHKERRQ(ierr);
  ierr = PetscFree(X_loc); CHKERRQ(ierr);
  ierr = PetscFree(mult); CHKERRQ(ierr);
  ierr = PetscFree(mat_O); CHKERRQ(ierr);
  ierr = PetscFree(mat_pr); CHKERRQ(ierr);
  ierr = PetscFree(ceed_data); CHKERRQ(ierr);
  ierr = PetscFree(user_O); CHKERRQ(ierr);
  ierr = PetscFree(user_pr); CHKERRQ(ierr);
  ierr = PetscFree(l_size); CHKERRQ(ierr);
  ierr = PetscFree(xl_size); CHKERRQ(ierr);
  ierr = PetscFree(g_size); CHKERRQ(ierr);
  ierr = VecDestroy(&rhs); CHKERRQ(ierr);
  ierr = VecDestroy(&rhs_loc); CHKERRQ(ierr);
  ierr = MatDestroy(&mat_coarse); CHKERRQ(ierr);
  ierr = KSPDestroy(&ksp); CHKERRQ(ierr);
  ierr = SNESDestroy(&snes_dummy); CHKERRQ(ierr);
  ierr = DMDestroy(&dm_orig); CHKERRQ(ierr);
  CeedVectorDestroy(&target);
  CeedQFunctionDestroy(&qf_error);
  CeedQFunctionDestroy(&qf_restrict);
  CeedQFunctionDestroy(&qf_prolong);
  CeedOperatorDestroy(&op_error);
  CeedDestroy(&ceed);
  return PetscFinalize();
}


#ifndef __FFT_DECON_OPERATOR_H__
#define __FFT_DECON_OPERATOR_H__
#include <string>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_fft_complex.h>
#include "mspass/seismic/TimeWindow.h"
#include "mspass/utility/Metadata.h"
#include "mspass/seismic/CoreTimeSeries.h"
#include "mspass/algorithms/deconvolution/ComplexArray.h"
namespace mspass::algorithms::deconvolution{
/*! \brief Object to hold components needed in all fft based decon algorithms.

The fft based algorithms implemented here us the GNU Scientific Library
prime factorization fft algorithm.  Those methods require initialization
given length of the fft to load and store the factorization data.  This
object holds these for all such methods and recomputes them only when
needed for efficiency.  */
class FFTDeconOperator
{
public:
    FFTDeconOperator();
    FFTDeconOperator(const mspass::utility::Metadata& md);
    FFTDeconOperator(const FFTDeconOperator& parent);
    ~FFTDeconOperator();
    FFTDeconOperator& operator=(const FFTDeconOperator& parent);
    void changeparameter(const mspass::utility::Metadata& md);
    void change_size(const int nfft_new);
    void change_shift(const int shift) {
        sample_shift=shift;
    };
    int get_size(){return nfft;};
    int get_shift(){return sample_shift;};
    int operator_size() {
        return static_cast<int>(nfft);
    };
    int operator_shift() {
        return sample_shift;
    };
    double df(const double dt){
        double period;
        period=static_cast<double>(nfft)*dt;
        return 1.0/period;
    };
    /*! \brief Return inverse wavelet for Fourier methods.

    This is a helper to be used by the inverse_wavelet method for all
    Fourier based deconvolution methods.   It avoids repetitious code that
    would be required otherwise.  inverse_wavelet methods are only
    wrappers for this generic method.  See documentation for inverse_wavelet
    for description of tshift and t0parent.  */
    mspass::seismic::CoreTimeSeries FourierInverse(const ComplexArray& winv, const ComplexArray& sw,
   	const double dt, const double t0parent);

protected:
    int nfft;
    int sample_shift;
    gsl_fft_complex_wavetable *wavetable;
    gsl_fft_complex_workspace *workspace;
    ComplexArray winv;
};

/* This helper is best referenced here */

/*! \brief Circular buffer procedure.

 In the Fourier world circular vectors are an important thing
 to deal with because the fft is intimately connected with circlular
 things.   This routine can be used, for example, to time shift the
 time domain version of a signal after it was processed with an fft.

 \param d - is the input vector to be shifted
 \param i0 - is the wrap point.   On exit sample i0 of d will be sample 0.
   (Warning - this is C convention sample number)
*/

std::vector<double>  circular_shift(const std::vector<double>& d,const int i0);
/*! Derive fft length from a time window.

All deconvlution methods using an fft need to define nfft based on the
length of the working time series.   This procedure returns the size from
an input window and sample interval. */
int ComputeFFTLength(const mspass::seismic::TimeWindow w, const double dt);
/*! Derive fft length using parameters in a metadata object.

This procedure is basically a higher level version of the function of
the same name with a time window and sample interval argument.
This procedure extracts these using three parameter keys to extract
the real numbers form md:  deconvolution_data_window_start,
decon_window_end, and target dt. */
int ComputeFFTLength(const mspass::utility::Metadata& md);
/*! Returns next power of 2 larger than n.
 *
 * Some FFT implementations require the size of the input data vector be a power
 * of 2.  This routine can be used to define a buffer size satisfying that constraint.*/
extern "C" {
    unsigned int nextPowerOf2(unsigned int n);
}
}
#endif


/* randist/mvgauss.c
 *
 * Copyright (C) 2016 Timothée Flutre, 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 <math.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_statistics.h>

static int multivar_vcov (const double data[], size_t d, size_t tda, size_t n,
                          double vcov[], size_t tda2);

/* Generate a random vector from a multivariate Gaussian distribution using
 * the Cholesky decomposition of the variance-covariance matrix, following
 * "Computational Statistics" from Gentle (2009), section 7.4.
 *
 * mu      mean vector (dimension d)
 * L       matrix resulting from the Cholesky decomposition of
 *         variance-covariance matrix Sigma = L L^T (dimension d x d)
 * result  output vector (dimension d)
 */
int
gsl_ran_multivariate_gaussian (const gsl_rng * r,
                               const gsl_vector * mu,
                               const gsl_matrix * L,
                               gsl_vector * result)
{
  const size_t M = L->size1;
  const size_t N = L->size2;

  if (M != N)
    {
      GSL_ERROR("requires square matrix", GSL_ENOTSQR);
    }
  else if (mu->size != M)
    {
      GSL_ERROR("incompatible dimension of mean vector with variance-covariance matrix", GSL_EBADLEN);
    }
  else if (result->size != M)
    {
      GSL_ERROR("incompatible dimension of result vector", GSL_EBADLEN);
    }
  else
    {
      size_t i;

      for (i = 0; i < M; ++i)
        gsl_vector_set(result, i, gsl_ran_ugaussian(r));

      gsl_blas_dtrmv(CblasLower, CblasNoTrans, CblasNonUnit, L, result);
      gsl_vector_add(result, mu);

      return GSL_SUCCESS;
    }
}

/* Compute the log of the probability density function at a given quantile
 * vector for a multivariate Gaussian distribution using the Cholesky
 * decomposition of the variance-covariance matrix.
 *
 * x       vector of quantiles (dimension d)
 * mu      mean vector (dimension d)
 * L       matrix resulting from the Cholesky decomposition of
 *         variance-covariance matrix Sigma = L L^T (dimension d x d)
 * result  output of the density (dimension 1)
 * work    vector used for intermediate computations (dimension d)
 */
int
gsl_ran_multivariate_gaussian_log_pdf (const gsl_vector * x,
                                       const gsl_vector * mu,
                                       const gsl_matrix * L,
                                       double * result,
                                       gsl_vector * work)
{
  const size_t M = L->size1;
  const size_t N = L->size2;

  if (M != N)
    {
      GSL_ERROR("requires square matrix", GSL_ENOTSQR);
    }
  else if (mu->size != M)
    {
      GSL_ERROR("incompatible dimension of mean vector with variance-covariance matrix", GSL_EBADLEN);
    }
  else if (x->size != M)
    {
      GSL_ERROR("incompatible dimension of quantile vector", GSL_EBADLEN);
    }
  else if (work->size != M)
    {
      GSL_ERROR("incompatible dimension of work vector", GSL_EBADLEN);
    }
  else
    {
      size_t i;
      double quadForm;        /* (x - mu)' Sigma^{-1} (x - mu) */
      double logSqrtDetSigma; /* log [ sqrt(|Sigma|) ] */

      /* compute: work = x - mu */
      for (i = 0; i < M; ++i)
        {
          double xi = gsl_vector_get(x, i);
          double mui = gsl_vector_get(mu, i);
          gsl_vector_set(work, i, xi - mui);
        }

      /* compute: work = L^{-1} * (x - mu) */
      gsl_blas_dtrsv(CblasLower, CblasNoTrans, CblasNonUnit, L, work);

      /* compute: quadForm = (x - mu)' Sigma^{-1} (x - mu) */
      gsl_blas_ddot(work, work, &quadForm);

      /* compute: log [ sqrt(|Sigma|) ] = sum_i log L_{ii} */
      logSqrtDetSigma = 0.0;
      for (i = 0; i < M; ++i)
        {
          double Lii = gsl_matrix_get(L, i, i);
          logSqrtDetSigma += log(Lii);
        }

      *result = -0.5*quadForm - logSqrtDetSigma - 0.5*M*log(2.0*M_PI);

      return GSL_SUCCESS;
    }
}

int
gsl_ran_multivariate_gaussian_pdf (const gsl_vector * x,
                                   const gsl_vector * mu,
                                   const gsl_matrix * L,
                                   double * result,
                                   gsl_vector * work)
{
  double logpdf;
  int status = gsl_ran_multivariate_gaussian_log_pdf(x, mu, L, &logpdf, work);

  if (status == GSL_SUCCESS)
    *result = exp(logpdf);

  return status;
}

/* Compute the maximum-likelihood estimate of the mean vector of samples
 * from a multivariate Gaussian distribution.
 *
 * Example from R (GPL): http://www.r-project.org/
 * (samples <- matrix(c(4.348817, 2.995049, -3.793431, 4.711934, 1.190864, -1.357363), nrow=3, ncol=2))
 * colMeans(samples) # 1.183478 1.515145
 */
int
gsl_ran_multivariate_gaussian_mean (const gsl_matrix * X, gsl_vector * mu_hat)
{
  const size_t M = X->size1;
  const size_t N = X->size2;

  if (N != mu_hat->size)
    {
      GSL_ERROR("mu_hat vector has wrong size", GSL_EBADLEN);
    }
  else
    {
      size_t j;

      for (j = 0; j < N; ++j)
        {
          gsl_vector_const_view c = gsl_matrix_const_column(X, j);
          double mean = gsl_stats_mean(c.vector.data, c.vector.stride, M);
          gsl_vector_set(mu_hat, j, mean);
        }

      return GSL_SUCCESS;
    }
}

/* Compute the maximum-likelihood estimate of the variance-covariance matrix
 * of samples from a multivariate Gaussian distribution.
 */
int
gsl_ran_multivariate_gaussian_vcov (const gsl_matrix * X, gsl_matrix * sigma_hat)
{
  const size_t M = X->size1;
  const size_t N = X->size2;

  if (sigma_hat->size1 != sigma_hat->size2)
    {
      GSL_ERROR("sigma_hat must be a square matrix", GSL_ENOTSQR);
    }
  else if (N != sigma_hat->size1)
    {
      GSL_ERROR("sigma_hat does not match X matrix dimensions", GSL_EBADLEN);
    }
  else
    {
      return multivar_vcov (X->data, N, X->tda, M, sigma_hat->data, sigma_hat->tda);
    }
}

/* Example from R (GPL): http://www.r-project.org/
 * (samples <- matrix(c(4.348817, 2.995049, -3.793431, 4.711934, 1.190864, -1.357363), nrow=3, ncol=2))
 * cov(samples) # 19.03539 11.91384 \n 11.91384  9.28796
 */
static int
multivar_vcov (const double data[], size_t d, size_t tda, size_t n,
               double vcov[], size_t tda2)
{
  size_t j1 = 0, j2 = 0;

  for (j1 = 0; j1 < d; ++j1)
    {
      vcov[j1 * tda2 + j1] = gsl_stats_variance(&(data[j1]), tda, n);
      for (j2 = j1 + 1; j2 < d; ++j2)
        {
          vcov[j1 * tda2 + j2] = gsl_stats_covariance(&(data[j1]), tda,
                                                      &(data[j2]), tda, n);
          vcov[j2 * tda2 + j1] = vcov[j1 * tda2 + j2];
        }
      }

  return GSL_SUCCESS;
}


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

void read_matrix(int** index, int** matrix, double scaling, int N_kw, char* input_fileName)
{
    FILE *fp = fopen(input_fileName, "r");
    fscanf(fp, "%*[^\n]\n");

    for (int ii = 0; ii < N_kw; ii++)
        fscanf(fp, "%d,", &((*index)[ii]));
    fscanf(fp, "%*[^\n]\n");

    int tmp;
    for (int ii = 0; ii < N_kw; ii++)
    {
        for (int jj = 0; jj < N_kw; jj++)
        {
            fscanf(fp, "%d,", &tmp);
            (*matrix)[ii*N_kw + jj] = (int) (scaling * tmp);
        }
        fscanf(fp, "%*[^\n]\n");
    }
    fclose(fp);
}

void pad_matrix(int** matrix_padded, int** matrix, int N_kw, int N_doc, int freq_max)
{
    // Initialising RNG
    const gsl_rng_type * T;
    gsl_rng * r;
    gsl_rng_env_setup();
    T = gsl_rng_default;
    r = gsl_rng_alloc(T);

    // perform padding on the keywords
    int ii, jj;
    #pragma omp parallel for private(ii)
    for (int ii = 0; ii < N_kw; ii++)
        (*matrix_padded)[ii*N_kw + ii] = 2*freq_max;

    // perform padding
    #pragma omp parallel for private(ii, jj)
    for (ii = 0; ii < N_kw; ii++)
    {
        for (jj = 0; jj < N_kw; jj++)
        {
            if (ii > jj)
            {
                int n1 = 2*freq_max - (*matrix)[ii*N_kw + ii];
                int n2 = 2*freq_max - (*matrix)[jj*N_kw + jj];
                int x1 = gsl_ran_hypergeometric(r, n1, 2*N_doc-n1, (*matrix_padded)[jj*N_kw + jj]);
                int x2 = gsl_ran_hypergeometric(r, n2, 2*N_doc-n2, (*matrix_padded)[ii*N_kw + ii]);

                (*matrix_padded)[ii*N_kw + jj] = (*matrix)[ii*N_kw + jj] + x1 + x2;
                (*matrix_padded)[jj*N_kw + ii] = (*matrix_padded)[ii*N_kw + jj];
            }
        }
    }
    gsl_rng_free(r);
}

void observe_matrix(gsl_matrix* matrix_obs, int** matrix_padded, int N_kw)
{
    // perform observed count generation
    for (int ii = 0; ii < N_kw; ii++)
        for (int jj = 0; jj < N_kw; jj++)
            gsl_matrix_set(matrix_obs, ii, jj, (double) ((*matrix_padded)[ii*N_kw + jj]));
}

void permutation_generation(int* idx1, int* idx2, int** permutation_tmp, int** permutation, int** permutation_inv, int N_kw, int N_obs)
{
    *idx1 = rand() % N_obs;
    *idx2 = -1;
    int idx_old = (*permutation)[*idx1];
    int idx_new = rand() % N_kw;

    (*permutation_tmp)[*idx1] = idx_new;

    if ((*permutation_inv)[idx_new] >= 0)
    {
        *idx2 = (*permutation_inv)[idx_new];
        (*permutation_tmp)[*idx2] = idx_old;
    }
}


#pragma once
#include <thrust/host_vector.h>
#include <sstream>
#include "IsingLattice2D.h"
#include "class_mc_io.h"
#include "Matrix.h"
#include "MemTimeTester.h"
extern "C" {
#include "random.h"
}
#include <gsl/gsl_sf_zeta.h>
#include <cmath>
#include "obs_calc_fast.cuh"

double real_trigamma_cmplxarg(double x, double y, int m, int l);

class GeneralLRW{
	//This is a more lightweight version of the long range wolff class, designed to be
	//more general - hopefully to admit the use of user-defined "test_spins" functions
	//so that it can be used for many different MC simulations.  It does not hold the
	//lattice in order to keep the lattice easily available for GPU computations
	std::vector<std::vector<double>> interactions;//interaction matrix
	std::vector<std::vector<double>> interaction_sum;//values from interaction sum integral
	std::vector<double> site_sums;//absolute sum of interactions at each site
	double mag;//current magnetization
	double h; //applied field
	std::vector<spin> buffer;//holds the indices of spins to check
	std::vector<std::vector<int>> cluster;//0 at indices that are not in cluster, 1 at indices that are
	int cluster_size;//size of the cluster during a given MC step
	int cluster_mag;//total magnetization of the cluster
	void set_spin_boson_model_wc_exp(class_mc_params);//version for finite cutoff frequency with exponential tail
	void set_spin_boson_model_wc_hard(class_mc_params);//version for finite cutoff frequency with sharp cutoff
	void set_spin_boson_model(class_mc_params);//create interaction matrix for the spin boson model with infinite cutoff frequency
	void set_interaction_sum();
	void set_site_sums();
	MemTimeTester timer;

	inline double kernel_integrand(double alpha, int Nt, double rbv, double taub, double u){
		//evaluate the integrand of the kernel of the spin-boson effective interactions
		return 2.0 * alpha / Nt / Nt * u * cos(rbv*u)*(exp(u*(-taub)) + exp(u*(taub - 1.0)))/(1 - exp(-u));
	}

	double simpson_3_8(double alpha, int Nt, double bomc, double rbv, double taub, int N_points){
		//perform an integration of the spin-boson kernel with a hard, finite cutoff bomc.  Evaluated at
		//site distance rbv = r/(beta*v), taub = tau/beta, with number of time slices Nt and number of
		//integration slices N_points (this method will multiply N_points by 3 and use simpson's 3/8 rule)
		double du = (bomc)/(3.0*N_points);
		double result = 3.0*du*(4.0*alpha / Nt / Nt + kernel_integrand(alpha, Nt, rbv, taub, bomc))/8.0;
		for (int i = 1; i < 3*N_points; ++i){
			if(3*(i/3) == i){
				result += 0.75*du*kernel_integrand(alpha, Nt, rbv, taub, i*du);
			}
			else{
				result += 9.0*du*kernel_integrand(alpha, Nt, rbv, taub, i*du)/8.0;
			}
		}
		return result;
	}

public:
	GeneralLRW(class_mc_params);

	void print_timers(){ timer.print_timers(); }

	void step(IsingLattice2D& lat);

    bool step_one_site(IsingLattice2D& lat, double *prev_action, thrust::host_vector<double>& corr_ref);

	bool step_one_site_prealloc(IsingLattice2D& lat, double *prev_action, thrust::host_vector<double>& corr_ref,
			cufftHandle *full_forward_plan, cufftHandle *full_backward_plan, cufftHandle *onesite_forward_plan, cufftHandle *onesite_backward_plan,
			cudaError cuda_status, cufftDoubleReal *full_state_rs, cufftDoubleComplex *full_state_ft, cufftDoubleReal *onesite_state_rs, cufftDoubleComplex *onesite_state_ft);

	bool step_one_site_prealloc_lowmem(IsingLattice2D& lat, double *prev_action, thrust::host_vector<double>& corr_ref,
			thrust::host_vector<double>& ,thrust::host_vector<double>& ,thrust::host_vector<double>& ,thrust::host_vector<double>& ,
			thrust::host_vector<double>& ,thrust::host_vector<double>& ,
			cufftHandle *full_forward_plan, cufftHandle *full_backward_plan, cufftHandle *onesite_forward_plan, cufftHandle *onesite_backward_plan,
			cudaError cuda_status, cufftDoubleReal *full_state_rs, cufftDoubleComplex *full_state_ft, cufftDoubleReal *onesite_state_rs, cufftDoubleComplex *onesite_state_ft);

	bool metropolis_step(IsingLattice2D& lat, double *action);

	bool fast_metropolis_step(IsingLattice2D& lat, double *action);

	void test_spins(spin, IsingLattice2D& lat);

    void test_spins_one_site(spin, IsingLattice2D& lat);

	void set_mag(IsingLattice2D& lat){
		mag = 0;
		for (int i = 0; i < lat.get_Lx(); ++i){
			for (int j = 0; j < lat.get_Ly(); ++j){
				mag += lat.get_spin(i, j);
			}
		}
		mag = ((double) mag) / lat.get_N();
	}

	double get_mag(){
		return mag;
	}

	double get_cluster_size() { return (double)cluster_size; }

    void get_thrust_interactions(thrust::host_vector<double>& result){
		if(result.size() != interactions.size()*interactions[0].size()){
			result.resize(interactions.size()*interactions[0].size());
		}
		for (int i = 0; i < interactions.size(); ++i){
			for (int j = 0; j < interactions[i].size(); ++j){
				result[i*interactions[i].size() + j] = interactions[i][j];
			}
		}
    }

	std::vector<double> get_interaction_vector(){
		std::vector<double> result(interactions.size()*interactions[0].size());
		int Ly = interactions[0].size();
		for (int x = 0; x < interactions.size(); ++x){
			for (int y = 0; y < Ly; ++y){
				result[x*Ly + y] = interactions[x][y];
			}
		}
		return result;
	}

	std::vector<std::vector<double>> get_interactions(){
		return interactions;
	}

    double calc_mqt0_abs(IsingLattice2D&, double);

    double calc_mqt0_2_abs(IsingLattice2D&, double);

    double calc_mqt0_4_abs(IsingLattice2D&, double);

    double calc_mqt0_gpu(thrust::host_vector<double>&, double, int);

    double calc_locw1(IsingLattice2D&, double);

	double calc_sx(IsingLattice2D&);

	double calc_space_kinks(IsingLattice2D& lat);

	double calc_sz_stagger(IsingLattice2D&);

	double calc_xmag(IsingLattice2D&);

	double calc_xmag2(IsingLattice2D&);

	double calc_xmag4(IsingLattice2D&);

    double calc_mag(IsingLattice2D&);

	double calc_loc(IsingLattice2D&);

	double calc_loc2(IsingLattice2D&);

	double calc_loc4(IsingLattice2D&);

	double calc_s1s2(IsingLattice2D&);

	double calc_action_slow(IsingLattice2D&);

	double calc_point_action_slow(IsingLattice2D&,int, int);

	double calc_point_action_fast(IsingLattice2D&, int, int);

	void print_site_sums(){
		std::cout << "Site sums:\n" << vec2str(site_sums) << "\n";
	}

	void print_interactions() {
		std::stringstream outstring;
		outstring << "Interactions:\n";
		for (int i = 0; i < interactions.size(); ++i){
			outstring << vec2str(interactions[i]) << "\n";
		}
		std::cout << outstring.str() << "\n";
	}

	void print_interaction_sum() {
		std::stringstream outstring;
		outstring << "Interaction sum:\n";
		for (int i = 0; i < interaction_sum.size(); ++i){
			outstring << vec2str(interaction_sum[i]) << "\n";
		}
		std::cout << outstring.str() << "\n";
	}

	std::string get_int_string() {
		std::stringstream ss;
		ss << "Interactions:\n" ;
		for (int i = 0; i < interactions.size(); ++i){
			ss << vec2str(interactions[i]) << "\n";
		}
		return ss.str();
	}

};


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

#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_linalg.h>

#define L 1 //1 2 3 4
#define LL 2 //2^(1 2 3 4)
#define LLL 4 //2^(2 3 4 5)

#define N 8

#define M 1 //# of trials
#define R 1 //the number of different initial conditions
#define S 100 //the number of different rules
#define T1 100 //transient
#define T2 1000 //time window used
int **mat;
int component[N],cnum;

void findcomponent(int node)
{
	int i;

	for (i=0;i<N;i++){
		if (mat[node][i]==1 && component[i]==0){
			component[i]=cnum;
			findcomponent(i);
		}
	}
}

double dabs(double x)
{
	if (x<0.0)return -x;
	else return x;
}

int main()
{
	FILE *fp,*fq,*fe,*fg,*fh,*fc,*fa;

	int i,j,k,l,t,m,r,s;
	int flag,flag1,flag2,tmp,tmp1,tmp2,tmp3,tmp4,ttmp;
	double dtmp;

	int x[N],xx[N][L+1],y[N];
	int in[N];
	int **nb,**rule;

	int check[N],sum,check_sum;

	double p1[N][LL],p2[N][LLL];
	double ****p3,****p4;
	int num1,num2,num3;

	double indeg[N],outdeg[N];

	double p=0.5;

	double q;
	int kav=2;
	int **adj,**adjflag;

	double *llap;
	double **lap,**edge,**grad,**circ;
	double div[N],sol[N];
	double diag[N];
	double **uu,**vv,**lap_inv;

	double el2_av,gl2_av,cl2_av,cratio_av;
	double el2,gl2,cl2,cratio;

	double *curl,*vp;
	int trinum;
	int **tri;

	double *ccurcur;
	double **curcur;
	double *diag2;
	double **uu2,**vv2,**curcur_inv;

	double **lcirc,**gcirc;

	double hl2_av,lcl2_av,gcratio_av,lcratio_av;
	double hl2,lcl2,gcratio,lcratio;

	double et,gt,ht,ct;
	double et_av,gt_av,ht_av,ct_av;

	char filename[100];

	gsl_matrix *v_gsl=gsl_matrix_alloc(N,N);
	gsl_vector *s_gsl=gsl_vector_alloc(N);
	gsl_vector *work=gsl_vector_alloc(N);
	gsl_matrix_view lap_gsl;

	gsl_matrix_view cur_gsl;

	const gsl_rng_type * TYPE;
	gsl_rng * ran;

	gsl_rng_env_setup();
	TYPE=gsl_rng_default;
	ran=gsl_rng_alloc(TYPE);
	gsl_rng_set(ran,2);

	adj=malloc(sizeof(int*)*N);
	adjflag=malloc(sizeof(int*)*N);

	mat=malloc(sizeof(int*)*N);
	lap=malloc(sizeof(double*)*N);
	edge=malloc(sizeof(double*)*N);
	grad=malloc(sizeof(double*)*N);
	circ=malloc(sizeof(double*)*N);
	lcirc=malloc(sizeof(double*)*N);
	gcirc=malloc(sizeof(double*)*N);

	uu=malloc(sizeof(double*)*N);
	vv=malloc(sizeof(double*)*N);
	lap_inv=malloc(sizeof(double*)*N);

	llap=malloc(sizeof(double)*(N*N));

	p3=malloc(sizeof(double***)*N);
	p4=malloc(sizeof(double***)*N);

	for (i=0;i<N;i++){
		adj[i]=malloc(sizeof(int)*N);
		adjflag[i]=malloc(sizeof(int)*N);

		mat[i]=malloc(sizeof(int)*N);

		lap[i]=malloc(sizeof(double)*N);
		edge[i]=malloc(sizeof(double)*N);
		grad[i]=malloc(sizeof(double)*N);
		circ[i]=malloc(sizeof(double)*N);
		lcirc[i]=malloc(sizeof(double)*N);
		gcirc[i]=malloc(sizeof(double)*N);

		uu[i]=malloc(sizeof(double)*N);
		vv[i]=malloc(sizeof(double)*N);
		lap_inv[i]=malloc(sizeof(double)*N);

		p3[i]=malloc(sizeof(double**)*N);
		p4[i]=malloc(sizeof(double**)*N);
		for (j=0;j<N;j++){
			p3[i][j]=malloc(sizeof(double*)*LL);
			p4[i][j]=malloc(sizeof(double*)*LLL);
			for (k=0;k<LL;k++){
				p3[i][j][k]=malloc(sizeof(double)*LL);
			}
			for (k=0;k<LLL;k++){
				p4[i][j][k]=malloc(sizeof(double)*LL);
			}
		}
	}

	sprintf(filename,"rtn_hodge_sw_av_n%d_k%d.txt",N,kav);
	fp=fopen(filename,"w");

	sprintf(filename,"rtn_hodge_sw_trial_n%d_k%d.txt",N,kav);
	fq=fopen(filename,"w");

	for (q=0.1;q<0.11;q=q+0.1){

		el2_av=0.0; gl2_av=0.0; cl2_av=0.0;
		cratio_av=0.0;

		hl2_av=0.0; lcl2_av=0.0;
		gcratio_av=0.0; lcratio_av=0.0;

		et_av=0.0; gt_av=0.0; ht_av=0.0; ct_av=0.0;

		for (m=0;m<M;m++){

			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					adj[i][j]=0;
					adjflag[i][j]=0;
				}
			}
			for (i=0;i<N;i++){
				for (k=-kav;k<kav+1;k++){
					if (k!=0){
						adj[(i+k+N)%N][i]=1;
						adj[i][(i+k+N)%N]=1;
					}
				}
			}

			for (i=0;i<N;i++){
				for (k=-kav;k<kav+1;k++){
					if (k!=0 && gsl_rng_uniform(ran)<q){
						adjflag[(i+k+N)%N][i]=1;
						adjflag[i][(i+k+N)%N]=1;
					}
				}
			}

			for (i=0;i<N;i++){
				for (j=i+1;j<N;j++){
					if (adjflag[i][j]==1){
						if (gsl_rng_uniform(ran)<0.5){
							while (1){
								tmp=(int)((double)N*gsl_rng_uniform(ran));

								if (j!=tmp && adj[tmp][j]==0){
									adj[i][j]=0;
									adj[j][i]=0;
									adj[tmp][j]=1;
									adj[j][tmp]=1;
									break;
								}
							}
						} else{
							while (1){
								tmp=(int)((double)N*gsl_rng_uniform(ran));

								if (i!=tmp && adj[i][tmp]==0){
									adj[i][j]=0;
									adj[j][i]=0;
									adj[tmp][i]=1;
									adj[i][tmp]=1;
									break;
								}
							}
						}
					}
				}
			}

			for (i=0;i<N;i++){
				for (j=i+1;j<N;j++){
					if (adj[i][j]==1){
						if (gsl_rng_uniform(ran)<0.5)adj[i][j]=0;
						else adj[j][i]=0;
					}
				}
			}

			for (i=0;i<N;i++){
				in[i]=0;
				for (j=0;j<N;j++){
					in[i]+=adj[j][i];
				}
			}

			nb=malloc(sizeof(int*)*N);
			rule=malloc(sizeof(int*)*N);
			for (i=0;i<N;i++){
				nb[i]=malloc(sizeof(int)*in[i]);
				rule[i]=malloc(sizeof(int)*in[i]);
			}

			//in-neighbors
			for (i=0;i<N;i++){
				k=0;
				for (j=0;j<N;j++){
					if (adj[j][i]==1){
						nb[i][k]=j;
						k++;
					}
				}
			}
			//

			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					edge[i][j]=0.0;
				}
			}

			for (s=0;s<S;s++){
				//random generation of rules
				for (i=0;i<N;i++){
					for (j=0;j<in[i];j++){
					//	rule[i][j]=gsl_ran_gaussian(ran,sigma);
						if (gsl_rng_uniform(ran)<0.5)rule[i][j]=1;
						else rule[i][j]=-1;
					}
				}

				for (k=0;k<N;k++){
					for (i=0;i<LL;i++){
						p1[k][i]=0.0;
					}

					for (i=0;i<LLL;i++){
						p2[k][i]=0.0;
					}

					for (l=0;l<in[k];l++){
						for (i=0;i<LL;i++){
							for (j=0;j<LL;j++){
								p3[k][nb[k][l]][i][j]=0.0;
							}
						}
						for (i=0;i<LLL;i++){
							for (j=0;j<LL;j++){
								p4[k][nb[k][l]][i][j]=0.0;
							}
						}
					}
				}

				for (r=0;r<R;r++){

					//initial state
					for (i=0;i<N;i++){
						if (gsl_rng_uniform(ran)<p)x[i]=1;
						else x[i]=-1;
					}

					//transient
					for (t=0;t<T1;t++){

						for (i=0;i<N;i++){
							tmp=0;
							for (j=0;j<in[i];j++){
								tmp+=rule[i][j]*x[nb[i][j]];
							}
							if (tmp<0)y[i]=-1;
							else y[i]=1;
						}

						for (i=0;i<N;i++){
							x[i]=y[i];
						}
					}

					for (i=0;i<N;i++){
						for (j=0;j<L;j++){
							xx[i][j]=0;
						}
						xx[i][L]=x[i];
					}

					//time window used to calculation
					for (t=0;t<T2+L-1;t++){
						for (i=0;i<N;i++){
							tmp=0;
							for (j=0;j<in[i];j++){
								tmp+=rule[i][j]*x[nb[i][j]];
							}
							if (tmp<0)y[i]=-1;
							else y[i]=1;

							for (j=0;j<L;j++){
								xx[i][j]=xx[i][j+1];
							}
							xx[i][L]=y[i];
						}

						if (t>=L-1){
							for (i=0;i<N;i++){
								num1=0; num2=0;
								for (k=0;k<L;k++){
									num1+=((xx[i][k]+1)/2)*(int)pow(2.0,k);
									num2+=((xx[i][k]+1)/2)*(int)pow(2.0,k);
								}
								num2+=((xx[i][L]+1)/2)*(int)pow(2.0,L);

								p1[i][num1]+=1.0;
								p2[i][num2]+=1.0;

								for (j=0;j<in[i];j++){
									num3=0;
									for (k=0;k<L;k++){
										num3+=((xx[nb[i][j]][k]+1)/2)*(int)pow(2.0,k);
									}

									p3[i][nb[i][j]][num1][num3]+=1.0;
									p4[i][nb[i][j]][num2][num3]+=1.0;
								}
							}
						}

						for (i=0;i<N;i++){
							x[i]=y[i];
						}
					}
				}//end of r-loop

				for (k=0;k<N;k++){
					for (i=0;i<LL;i++){
						p1[k][i]=p1[k][i]/(double)(T2*R);
					}

					for (i=0;i<LLL;i++){
						p2[k][i]=p2[k][i]/(double)(T2*R);
					}

					for (l=0;l<in[k];l++){
						for (i=0;i<LL;i++){
							for (j=0;j<LL;j++){
								p3[k][nb[k][l]][i][j]=p3[k][nb[k][l]][i][j]/(double)(T2*R);
							}
						}
						for (i=0;i<LLL;i++){
							for (j=0;j<LL;j++){
								p4[k][nb[k][l]][i][j]=p4[k][nb[k][l]][i][j]/(double)(T2*R);
							}
						}
					}
				}

				//construction of edge flow
				for (i=0;i<N;i++){
					for (j=0;j<in[i];j++){
						for (num2=0;num2<LLL;num2++){
							for (num3=0;num3<LL;num3++){
								if (p4[i][nb[i][j]][num2][num3]>0.0)edge[nb[i][j]][i]+=p4[i][nb[i][j]][num2][num3]*(log(p4[i][nb[i][j]][num2][num3])/log(2.0) + log(p1[i][num2%LL])/log(2.0) - log(p2[i][num2])/log(2.0) - log(p3[i][nb[i][j]][num2%LL][num3])/log(2.0));
							}
						}
					}
				}
			}//end of s-loop

			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					edge[i][j]=edge[i][j]/(double)S;
				}
			}

			for (i=0;i<N;i++){
				for (j=i+1;j<N;j++){
					dtmp=edge[i][j]-edge[j][i];
					edge[i][j]=dtmp;
					edge[j][i]=-dtmp;
				}
			}

			//finding the number of components
			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					mat[i][j]=0;
				}
			}

			for (i=0;i<N;i++){
				for (j=0;j<in[i];j++){
					if (i!=nb[i][j]){
						mat[nb[i][j]][i]=1;
						mat[i][nb[i][j]]=1;
					}
				}
			}

			tmp=0;
			for (i=0;i<N;i++){
				for (j=i+1;j<N;j++){
					if (mat[i][j]==1)tmp++;
				}
			}
			et=(double)tmp;

			for (i=0;i<N;i++){
				component[i]=0;
			}

			cnum=1; component[0]=cnum;
			findcomponent(0);
			for (i=1;i<N;i++){
				if (component[i]==0){
					cnum++;
					component[i]=cnum;
					findcomponent(i);
				}
			}

			gt=(double)(N-cnum);

			//graph Laplacian
			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					lap[i][j]=0.0;
				}
			}

			for (i=0;i<N;i++){
				lap[i][i]=0.0;
				for (j=0;j<N;j++){
					lap[i][i]+=(double)mat[i][j];
				}
			}

			for (i=0;i<N;i++){
				for (j=0;j<in[i];j++){
					if (i!=nb[i][j]){
						lap[nb[i][j]][i]=-1.0;
						lap[i][nb[i][j]]=-1.0;
					}
				}
			}

			//generalized inverse of lap
			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					llap[i*N+j]=lap[i][j];
				}
			}

			lap_gsl=gsl_matrix_view_array(llap,N,N);
			gsl_linalg_SV_decomp(&lap_gsl.matrix,v_gsl,s_gsl,work);

			for (i=0;i<N-cnum;i++){
				diag[i]=1.0/gsl_vector_get(s_gsl,i);
			}
			for (i=N-cnum;i<N;i++){
				diag[i]=0.0;
			}

			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					uu[i][j]=gsl_matrix_get(&lap_gsl.matrix,i,j);
					vv[i][j]=gsl_matrix_get(v_gsl,i,j);
				}
			}

			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					lap_inv[i][j]=0.0;
					for (k=0;k<N;k++){
						lap_inv[i][j]+=vv[i][k]*diag[k]*uu[j][k];
					}
				}
			}
			//

			//divergence of edge flow
			for (i=0;i<N;i++){
				div[i]=0.0;
				for (j=0;j<N;j++){
					div[i]+=edge[i][j];
				}
			}

			//solution (negative potential)
			for (i=0;i<N;i++){
				sol[i]=0.0;
				for (j=0;j<N;j++){
					sol[i]-=lap_inv[i][j]*div[j];
				}
			}

			//gradient flow
			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					grad[i][j]=sol[j]-sol[i];
				}
			}

			//circular flow
			for (i=0;i<N;i++){
				for (j=0;j<N;j++){
					circ[i][j]=edge[i][j]-grad[i][j];
				}
			}

			el2=0.0; gl2=0.0; cl2=0.0;
			for (i=0;i<N;i++){
				for (j=i+1;j<N;j++){
					if (lap[i][j]<0.0){
						el2+=edge[i][j]*edge[i][j];
						gl2+=grad[i][j]*grad[i][j];
						cl2+=circ[i][j]*circ[i][j];
					}
				}
			}

			if (el2>0.0)cratio=cl2/el2;
			else cratio=0.0;

			el2_av+=el2/(double)M;
			gl2_av+=gl2/(double)M;
			cl2_av+=cl2/(double)M;
			cratio_av+=cratio/(double)M;

			//decomposing circular flow into harmonic and curl flows
			trinum=0;
			for (i=0;i<N;i++){
				for (j=i+1;j<N;j++){
					for (k=j+1;k<N;k++){
						if (mat[i][j]*mat[j][k]*mat[k][i]!=0)trinum++;
					}
				}
			}

			if (trinum>0){
				curl=malloc(sizeof(double)*trinum);
				tri=malloc(sizeof(int*)*trinum);
				for (i=0;i<trinum;i++){
					tri[i]=malloc(sizeof(int)*3);
				}

				l=0;
				for (i=0;i<N;i++){
					for (j=i+1;j<N;j++){
						for (k=j+1;k<N;k++){
							if (mat[i][j]*mat[j][k]*mat[k][i]!=0){
								tri[l][0]=i; tri[l][1]=j; tri[l][2]=k;
								curl[l]=circ[i][j]+circ[j][k]-circ[i][k];
								l++;
							}
						}
					}
				}

				ccurcur=malloc(sizeof(double)*trinum*trinum);
				diag2=malloc(sizeof(double)*trinum);
				vp=malloc(sizeof(double)*trinum);

				curcur=malloc(sizeof(double*)*trinum);
				uu2=malloc(sizeof(double*)*trinum);
				vv2=malloc(sizeof(double*)*trinum);
				curcur_inv=malloc(sizeof(double*)*trinum);
				for (i=0;i<trinum;i++){
					curcur[i]=malloc(sizeof(double)*trinum);
					uu2[i]=malloc(sizeof(double)*trinum);
					vv2[i]=malloc(sizeof(double)*trinum);
					curcur_inv[i]=malloc(sizeof(double)*trinum);
				}

				for (i=0;i<trinum;i++){
					for (j=0;j<trinum;j++){
						curcur[i][j]=0.0;
						if (tri[i][0]==tri[j][0] && tri[i][1]==tri[j][1])curcur[i][j]+=1.0;
						if (tri[i][0]==tri[j][0] && tri[i][1]==tri[j][2])curcur[i][j]-=1.0;
						if (tri[i][0]==tri[j][1] && tri[i][1]==tri[j][2])curcur[i][j]+=1.0;
						if (tri[i][1]==tri[j][0] && tri[i][2]==tri[j][1])curcur[i][j]+=1.0;
						if (tri[i][1]==tri[j][0] && tri[i][2]==tri[j][2])curcur[i][j]-=1.0;
						if (tri[i][1]==tri[j][1] && tri[i][2]==tri[j][2])curcur[i][j]+=1.0;
						if (tri[i][0]==tri[j][0] && tri[i][2]==tri[j][1])curcur[i][j]-=1.0;
						if (tri[i][0]==tri[j][0] && tri[i][2]==tri[j][2])curcur[i][j]+=1.0;
						if (tri[i][0]==tri[j][1] && tri[i][2]==tri[j][2])curcur[i][j]-=1.0;
					}
				}

				gsl_matrix *v_gsl2=gsl_matrix_alloc(trinum,trinum);
				gsl_vector *s_gsl2=gsl_vector_alloc(trinum);
				gsl_vector *work2=gsl_vector_alloc(trinum);

				//generalized inverse of curcur
				for (i=0;i<trinum;i++){
					for (j=0;j<trinum;j++){
						ccurcur[i*trinum+j]=curcur[i][j];
					}
				}

				cur_gsl=gsl_matrix_view_array(ccurcur,trinum,trinum);
				gsl_linalg_SV_decomp(&cur_gsl.matrix,v_gsl2,s_gsl2,work2);

				tmp=0;
				for (i=0;i<trinum;i++){
					dtmp=gsl_vector_get(s_gsl2,i);
					if (dtmp>1e-8)diag2[i]=1.0/dtmp;
					else {
						diag2[i]=0.0;
						tmp++;
					}
				}

				ct=(double)(trinum-tmp);
				ht=et-gt-ct;

				for (i=0;i<trinum;i++){
					for (j=0;j<trinum;j++){
						uu2[i][j]=gsl_matrix_get(&cur_gsl.matrix,i,j);
						vv2[i][j]=gsl_matrix_get(v_gsl2,i,j);
					}
				}

				for (i=0;i<trinum;i++){
					for (j=0;j<trinum;j++){
						curcur_inv[i][j]=0.0;
						for (k=0;k<trinum;k++){
							curcur_inv[i][j]+=vv2[i][k]*diag2[k]*uu2[j][k];
						}
					}
				}

				for (i=0;i<trinum;i++){
					vp[i]=0.0;
					for (j=0;j<trinum;j++){
						vp[i]+=curcur_inv[i][j]*curl[j];
					}
				}

				for (i=0;i<N;i++){
					lcirc[i][i]=0.0;
					for (j=i+1;j<N;j++){
						lcirc[i][j]=0.0;
						if (mat[i][j]!=0){
							for (k=0;k<trinum;k++){
								if (tri[k][0]==i && tri[k][1]==j)lcirc[i][j]+=vp[k];
								if (tri[k][0]==i && tri[k][2]==j)lcirc[i][j]-=vp[k];
								if (tri[k][1]==i && tri[k][2]==j)lcirc[i][j]+=vp[k];
							}
						}
					}
				}

				for (i=0;i<N;i++){
					for (j=0;j<i;j++){
						lcirc[i][j]=-lcirc[j][i];
					}
				}

				for (i=0;i<N;i++){
					for (j=0;j<N;j++){
						gcirc[i][j]=circ[i][j]-lcirc[i][j];
					}
				}

				hl2=0.0; lcl2=0.0;
				for (i=0;i<N;i++){
					for (j=i+1;j<N;j++){
						if (lap[i][j]<0.0){
							hl2+=gcirc[i][j]*gcirc[i][j];
							lcl2+=lcirc[i][j]*lcirc[i][j];
						}
					}
				}

				fe=fopen("edgeflow.txt","w");
				fg=fopen("gradientflow.txt","w");
				fh=fopen("harmonicflow.txt","w");
				fc=fopen("curlflow.txt","w");
				fa=fopen("adjmatrix.txt","w");
				for (i=0;i<N;i++){
					for (j=0;j<N;j++){
						if (adj[i][j]==1)fprintf(fa,"1 ");
						else fprintf(fa,"0 ");

						if (lap[i][j]<0.0){
							fprintf(fe,"%lf ",edge[i][j]);
							fprintf(fg,"%lf ",grad[i][j]);
							fprintf(fh,"%lf ",gcirc[i][j]);
							fprintf(fc,"%lf ",lcirc[i][j]);
						} else{
							fprintf(fe,"%lf ",0.0);
							fprintf(fg,"%lf ",0.0);
							fprintf(fh,"%lf ",0.0);
							fprintf(fc,"%lf ",0.0);
						}
					}
					fprintf(fe,"\n");
					fprintf(fg,"\n");
					fprintf(fh,"\n");
					fprintf(fc,"\n");
					fprintf(fa,"\n");
				}
				fclose(fe);
				fclose(fg);
				fclose(fh);
				fclose(fc);
				fclose(fa);

				gsl_matrix_free(v_gsl2);
				gsl_vector_free(s_gsl2);
				gsl_vector_free(work2);

				free(ccurcur);
				free(diag2);
				free(vp);

				if (curcur){
					for (i=0;i<trinum;i++){
						free(curcur[i]);
					}
				}
				if (uu2){
					for (i=0;i<trinum;i++){
						free(uu2[i]);
					}
				}
				if (vv2){
					for (i=0;i<trinum;i++){
						free(vv2[i]);
					}
				}
				if (curcur_inv){
					for (i=0;i<trinum;i++){
						free(curcur_inv[i]);
					}
				}

			} else{
				ct=0.0;
				ht=et-gt-ct;

				for (i=0;i<N;i++){
					for (j=0;j<N;j++){
						gcirc[i][j]=circ[i][j];
					}
				}

				hl2=0.0; lcl2=0.0;
				for (i=0;i<N;i++){
					for (j=i+1;j<N;j++){
						if (lap[i][j]<0.0){
							hl2+=gcirc[i][j]*gcirc[i][j];
						}
					}
				}
			}

			if (cl2>0.0)gcratio=hl2/el2;
			else gcratio=0.0;

			hl2_av+=hl2/(double)M;
			lcl2_av+=lcl2/(double)M;
			gcratio_av+=gcratio/(double)M;

			lcratio=cratio-gcratio;
			lcratio_av+=lcratio/(double)M;

			et_av+=et/(double)M;
			gt_av+=gt/(double)M;
			ht_av+=ht/(double)M;
			ct_av+=ct/(double)M;

			if (m%1==0)printf("%lf %d %d %lf %lf %lf %lf %lf %lf\n",q,m,cnum,el2,gl2,hl2,lcl2,gcratio,lcratio);
			fprintf(fq,"%lf %d %d %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf\n",q,m,cnum,el2,gl2,hl2,lcl2,gcratio,lcratio,et,gt,ht,ct);

			if (nb){
				for (i=0;i<N;i++){
					if (nb[i])free(nb[i]);
				}
				free(nb);
			}

			if (rule){
				for (i=0;i<N;i++){
					if (rule[i])free(rule[i]);
				}
				free(rule);
			}
		}//end of m-loop

		fprintf(fp,"%lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf %.15lf\n",q,el2_av,gl2_av,hl2_av,lcl2_av,gcratio_av,lcratio_av,et_av,gt_av,ht_av,ct_av);
	}//end of q-loop

	fclose(fp);
	fclose(fq);

	return 0;
}


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <gbpLib.h>
#include <gbpMath.h>
#include <gbpCosmo_linear_theory.h>
#include <gsl/gsl_multimin.h>

typedef struct init_gbpCosmo2gbpCosmo_integrand_params_struct init_gbpCosmo2gbpCosmo_integrand_params;
struct init_gbpCosmo2gbpCosmo_integrand_params_struct {
    double                     inv_s;
    double                     z_source;
    double                     z_target;
    double                     R_1;
    double                     R_2;
    cosmo_info **              cosmo_source;
    cosmo_info **              cosmo_target;
    int                        n_int;
    gsl_integration_workspace *wspace;
};

double init_gbpCosmo2gbpCosmo_integrand(double R, void *params_in) {
    init_gbpCosmo2gbpCosmo_integrand_params *params       = (init_gbpCosmo2gbpCosmo_integrand_params *)params_in;
    double                                   inv_s        = params->inv_s;
    double                                   z_source     = params->z_source;
    double                                   z_target     = params->z_target;
    cosmo_info **                            cosmo_source = params->cosmo_source;
    cosmo_info **                            cosmo_target = params->cosmo_target;
    return (pow(1. - (sigma_R(cosmo_source, inv_s * R, z_source, PSPEC_LINEAR_TF, PSPEC_ALL_MATTER)) /
                         (sigma_R(cosmo_target, R, z_target, PSPEC_LINEAR_TF, PSPEC_ALL_MATTER)),
                2.) /
            R);
}

double init_gbpCosmo2gbpCosmo_minimize_function(const gsl_vector *v_i, void *params_in) {
    // Set variable integrand parameters
    init_gbpCosmo2gbpCosmo_integrand_params *params = (init_gbpCosmo2gbpCosmo_integrand_params *)params_in;
    params->inv_s                                   = gsl_vector_get(v_i, 0);
    params->z_target                                = gsl_vector_get(v_i, 1);

    // Perform integral to minimize
    gsl_function integrand;
    double       delta_i;
    double       abs_error;
    integrand.function = init_gbpCosmo2gbpCosmo_integrand;
    integrand.params   = params_in;
    gsl_integration_qag(&integrand, params->R_1, params->R_2, 0, 1e-3, params->n_int, GSL_INTEG_GAUSS61, params->wspace, &delta_i, &abs_error);
    return (delta_i / take_ln(params->R_2 / params->R_1));
}

void init_gbpCosmo2gbpCosmo(cosmo_info **           cosmo_source,
                            cosmo_info **           cosmo_target,
                            double                  z_min,
                            double                  M_min,
                            double                  M_max,
                            gbpCosmo2gbpCosmo_info *gbpCosmo2gbpCosmo) {
    SID_log("Initializing cosmology scaling...", SID_LOG_OPEN | SID_LOG_TIMER);
    SID_set_verbosity(SID_SET_VERBOSITY_RELATIVE, -1);

    // Store some infor in the gbpCosmo2gbpCosmo_info structure
    gbpCosmo2gbpCosmo->M_min        = M_min;
    gbpCosmo2gbpCosmo->M_max        = M_max;
    gbpCosmo2gbpCosmo->z_min        = z_min;
    gbpCosmo2gbpCosmo->cosmo_source = (*cosmo_source);
    gbpCosmo2gbpCosmo->cosmo_target = (*cosmo_target);

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

    // Starting point
    x = gsl_vector_alloc(2);
    gsl_vector_set(x, 0, 1.);    // inv_s
    gsl_vector_set(x, 1, z_min); // z_scaled

    // Set initial step sizes to 1
    ss = gsl_vector_alloc(2);
    gsl_vector_set_all(ss, 1.0);

    // Set parameters
    init_gbpCosmo2gbpCosmo_integrand_params params;
    params.cosmo_source = cosmo_source;
    params.cosmo_target = cosmo_target;
    params.z_source     = z_min;
    params.R_1          = R_of_M(M_min, *cosmo_source);
    params.R_2          = R_of_M(M_max, *cosmo_source);
    params.inv_s        = gsl_vector_get(x, 0);
    params.z_target     = gsl_vector_get(x, 1);
    params.n_int        = 100;
    params.wspace       = gsl_integration_workspace_alloc(params.n_int);

    // Initialize method
    minex_func.n      = 2;
    minex_func.f      = init_gbpCosmo2gbpCosmo_minimize_function;
    minex_func.params = (void *)(&params);
    s                 = gsl_multimin_fminimizer_alloc(T, 2);
    gsl_multimin_fminimizer_set(s, &minex_func, x, ss);

    // Perform minimization
    double size;
    int    status;
    size_t iter     = 0;
    size_t iter_max = 200;
    do {
        iter++;
        status = gsl_multimin_fminimizer_iterate(s);
        if(status)
            SID_exit_error("Error encountered during minimisation in init_gbpCosmo2gbpCosmo() (status=%d).",
                           SID_ERROR_LOGIC, status);
        size   = gsl_multimin_fminimizer_size(s);
        status = gsl_multimin_test_size(size, 1e-2);
    } while(status == GSL_CONTINUE && iter <= iter_max);
    if(status != GSL_SUCCESS)
        SID_exit_error("Failed to converge during minimisation in init_gbpCosmo2gbpCosmo() (status=%d,iter=%d).",
                       SID_ERROR_LOGIC, status, iter);

    // Finalize results
    double Omega_M_source           = ((double *)ADaPS_fetch(*cosmo_source, "Omega_M"))[0];
    double H_Hubble_source          = 1e2 * ((double *)ADaPS_fetch(*cosmo_source, "h_Hubble"))[0];
    double Omega_M_target           = ((double *)ADaPS_fetch(*cosmo_target, "Omega_M"))[0];
    double H_Hubble_target          = 1e2 * ((double *)ADaPS_fetch(*cosmo_target, "h_Hubble"))[0];
    gbpCosmo2gbpCosmo->s_L          = 1. / gsl_vector_get(s->x, 0);
    gbpCosmo2gbpCosmo->s_M          = (Omega_M_target * H_Hubble_target) / (Omega_M_source * H_Hubble_source) * pow((gbpCosmo2gbpCosmo->s_L), 3.);
    gbpCosmo2gbpCosmo->z_min_scaled = gsl_vector_get(s->x, 1);
    ;

    // Calculate growth factors needed for
    //    determining redshift mappings
    gbpCosmo2gbpCosmo->D_prime_z_min = linear_growth_factor(z_min, cosmo_target);
    gbpCosmo2gbpCosmo->D_z_scaled    = linear_growth_factor(gbpCosmo2gbpCosmo->z_min_scaled, cosmo_source);
    gbpCosmo2gbpCosmo->D_ratio       = gbpCosmo2gbpCosmo->D_prime_z_min / gbpCosmo2gbpCosmo->D_z_scaled;

    // Clean-up
    gsl_vector_free(x);
    gsl_vector_free(ss);
    gsl_multimin_fminimizer_free(s);
    gsl_integration_workspace_free(params.wspace);
    SID_set_verbosity(SID_SET_VERBOSITY_DEFAULT);
    SID_log("Done.", SID_LOG_CLOSE);
}


#ifndef DIFFEQ_SIM_H
#define DIFFEQ_SIM_H

#include <gsl/gsl_odeiv.h>
#include <gsl/gsl_errno.h>
#include <iostream>
#include <vector>

using namespace std;

class DiffEq_Sim {
    private:
        double t;      //initial time
        double h;      //time step
        double tmax;   //max time
        double hmin;

    public:
        DiffEq_Sim() {
            t    = 0.0;      //initial time
            h    = 0.1;     //time step
            tmax = 2000;
            hmin = 0.2;
        };

        ~DiffEq_Sim() {};

        int nbins;
        double* y;

        void printY() { for(int i=0; i < nbins; i++) { cout << y[i] << " ";} cout << endl; }

        vector<double> get_state() {
            vector<double> C;
            C.assign(y, y + nbins);
            return C;
        }

        double get_time() { return t; }

        virtual void initialize() {}
        virtual void derivative(const double y[], double dydt[]){}

        static int function(double t, double const y[], double dydt[], void *params) {
            DiffEq_Sim* model = static_cast <DiffEq_Sim*> (params);
            model->derivative(y, dydt);
            return GSL_SUCCESS;
        }

        int run_simulation() {
            gsl_odeiv_evolve*  e = gsl_odeiv_evolve_alloc(nbins);
            gsl_odeiv_control* c = gsl_odeiv_control_y_new(1e-20, 0);
            gsl_odeiv_step*    s = gsl_odeiv_step_alloc(gsl_odeiv_step_rkf45, nbins);
            gsl_odeiv_system sys = {function, NULL, nbins, this };
            while (t < tmax) {  //convergence check here
                int status = gsl_odeiv_evolve_apply(e, c, s, &sys, &t, tmax, &h, y);
                if (status != GSL_SUCCESS) { return status; }
            }
            return 0;
        }

       int step_simulation( double stepsize ) {
            gsl_odeiv_evolve*  e = gsl_odeiv_evolve_alloc(nbins);
            gsl_odeiv_control* c = gsl_odeiv_control_y_new(1e-5, 0);
            gsl_odeiv_step*    s = gsl_odeiv_step_alloc(gsl_odeiv_step_rkf45, nbins);
            gsl_odeiv_system sys = {function, NULL, nbins, this };

            double tstop = t+stepsize;
            while (t < tstop) {
                int status = gsl_odeiv_evolve_apply(e, c, s, &sys, &t, tstop, &h, y);
                if (status != GSL_SUCCESS) { return status; }
            }
            return 0;
        }

        /*
        double* advance_simulation(double I_lim) {
            while (t < 2000) {  //convergence check here
                int status = gsl_odeiv_evolve_apply(e, c, s, &sys, &t, t1, &h, y);
                if (y[1] < I_lim) { return y; }
            }
            printY();
            return y;
        }
*/
};

#endif