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

#ifndef __GSL_INTEGRATION_H__
#define __GSL_INTEGRATION_H__
#include <stdlib.h>
#include <gsl/gsl_math.h>

#undef __BEGIN_DECLS
#undef __END_DECLS
#ifdef __cplusplus
# define __BEGIN_DECLS extern "C" {
# define __END_DECLS }
#else
# define __BEGIN_DECLS /* empty */
# define __END_DECLS /* empty */
#endif

__BEGIN_DECLS

/* Workspace for adaptive integrators */

typedef struct
  {
    size_t limit;
    size_t size;
    size_t nrmax;
    size_t i;
    size_t maximum_level;
    double *alist;
    double *blist;
    double *rlist;
    double *elist;
    size_t *order;
    size_t *level;
  }
gsl_integration_workspace;

gsl_integration_workspace *
  gsl_integration_workspace_alloc (const size_t n);

void
  gsl_integration_workspace_free (gsl_integration_workspace * w);

/* Workspace for QAWS integrator */

typedef struct
{
  double alpha;
  double beta;
  int mu;
  int nu;
  double ri[25];
  double rj[25];
  double rg[25];
  double rh[25];
}
gsl_integration_qaws_table;

gsl_integration_qaws_table *
gsl_integration_qaws_table_alloc (double alpha, double beta, int mu, int nu);

int
gsl_integration_qaws_table_set (gsl_integration_qaws_table * t,
                                double alpha, double beta, int mu, int nu);

void
gsl_integration_qaws_table_free (gsl_integration_qaws_table * t);

/* Workspace for QAWO integrator */

enum gsl_integration_qawo_enum { GSL_INTEG_COSINE, GSL_INTEG_SINE };

typedef struct
{
  size_t n;
  double omega;
  double L;
  double par;
  enum gsl_integration_qawo_enum sine;
  double *chebmo;
}
gsl_integration_qawo_table;

gsl_integration_qawo_table *
gsl_integration_qawo_table_alloc (double omega, double L,
                                  enum gsl_integration_qawo_enum sine,
                                  size_t n);

int
gsl_integration_qawo_table_set (gsl_integration_qawo_table * t,
                                double omega, double L,
                                enum gsl_integration_qawo_enum sine);

int
gsl_integration_qawo_table_set_length (gsl_integration_qawo_table * t,
                                       double L);

void
gsl_integration_qawo_table_free (gsl_integration_qawo_table * t);

/* Definition of an integration rule */

typedef void gsl_integration_rule (const gsl_function * f,
                                   double a, double b,
                                   double *result, double *abserr,
                                   double *defabs, double *resabs);

void gsl_integration_qk15 (const gsl_function * f, double a, double b,
                           double *result, double *abserr,
                           double *resabs, double *resasc);

void gsl_integration_qk21 (const gsl_function * f, double a, double b,
                           double *result, double *abserr,
                           double *resabs, double *resasc);

void gsl_integration_qk31 (const gsl_function * f, double a, double b,
                           double *result, double *abserr,
                           double *resabs, double *resasc);

void gsl_integration_qk41 (const gsl_function * f, double a, double b,
                           double *result, double *abserr,
                           double *resabs, double *resasc);

void gsl_integration_qk51 (const gsl_function * f, double a, double b,
                           double *result, double *abserr,
                           double *resabs, double *resasc);

void gsl_integration_qk61 (const gsl_function * f, double a, double b,
                           double *result, double *abserr,
                           double *resabs, double *resasc);

void gsl_integration_qcheb (gsl_function * f, double a, double b,
                            double *cheb12, double *cheb24);

/* The low-level integration rules in QUADPACK are identified by small
   integers (1-6). We'll use symbolic constants to refer to them.  */

enum
  {
    GSL_INTEG_GAUSS15 = 1,      /* 15 point Gauss-Kronrod rule */
    GSL_INTEG_GAUSS21 = 2,      /* 21 point Gauss-Kronrod rule */
    GSL_INTEG_GAUSS31 = 3,      /* 31 point Gauss-Kronrod rule */
    GSL_INTEG_GAUSS41 = 4,      /* 41 point Gauss-Kronrod rule */
    GSL_INTEG_GAUSS51 = 5,      /* 51 point Gauss-Kronrod rule */
    GSL_INTEG_GAUSS61 = 6       /* 61 point Gauss-Kronrod rule */
  };

void
gsl_integration_qk (const int n, const double xgk[],
                    const double wg[], const double wgk[],
                    double fv1[], double fv2[],
                    const gsl_function *f, double a, double b,
                    double * result, double * abserr,
                    double * resabs, double * resasc);

int gsl_integration_qng (const gsl_function * f,
                         double a, double b,
                         double epsabs, double epsrel,
                         double *result, double *abserr,
                         size_t * neval);

int gsl_integration_qag (const gsl_function * f,
                         double a, double b,
                         double epsabs, double epsrel, size_t limit,
                         int key,
                         gsl_integration_workspace * workspace,
                         double *result, double *abserr);

int gsl_integration_qagi (gsl_function * f,
                          double epsabs, double epsrel, size_t limit,
                          gsl_integration_workspace * workspace,
                          double *result, double *abserr);

int gsl_integration_qagiu (gsl_function * f,
                           double a,
                           double epsabs, double epsrel, size_t limit,
                           gsl_integration_workspace * workspace,
                           double *result, double *abserr);

int gsl_integration_qagil (gsl_function * f,
                           double b,
                           double epsabs, double epsrel, size_t limit,
                           gsl_integration_workspace * workspace,
                           double *result, double *abserr);

int gsl_integration_qags (const gsl_function * f,
                          double a, double b,
                          double epsabs, double epsrel, size_t limit,
                          gsl_integration_workspace * workspace,
                          double *result, double *abserr);

int gsl_integration_qagp (const gsl_function * f,
                          double *pts, size_t npts,
                          double epsabs, double epsrel, size_t limit,
                          gsl_integration_workspace * workspace,
                          double *result, double *abserr);

int gsl_integration_qawc (gsl_function *f,
                          const double a, const double b, const double c,
                          const double epsabs, const double epsrel, const size_t limit,
                          gsl_integration_workspace * workspace,
                          double * result, double * abserr);

int gsl_integration_qaws (gsl_function * f,
                          const double a, const double b,
                          gsl_integration_qaws_table * t,
                          const double epsabs, const double epsrel,
                          const size_t limit,
                          gsl_integration_workspace * workspace,
                          double *result, double *abserr);

int gsl_integration_qawo (gsl_function * f,
                          const double a,
                          const double epsabs, const double epsrel,
                          const size_t limit,
                          gsl_integration_workspace * workspace,
                          gsl_integration_qawo_table * wf,
                          double *result, double *abserr);

int gsl_integration_qawf (gsl_function * f,
                          const double a,
                          const double epsabs,
                          const size_t limit,
                          gsl_integration_workspace * workspace,
                          gsl_integration_workspace * cycle_workspace,
                          gsl_integration_qawo_table * wf,
                          double *result, double *abserr);

/* Workspace for fixed-order Gauss-Legendre integration */

typedef struct
  {
    size_t n;         /* number of points */
    double *x;        /* Gauss abscissae/points */
    double *w;        /* Gauss weights for each abscissae */
    int precomputed;  /* high precision abscissae/weights precomputed? */
  }
gsl_integration_glfixed_table;

gsl_integration_glfixed_table * gsl_integration_glfixed_table_alloc (size_t n);

void gsl_integration_glfixed_table_free (gsl_integration_glfixed_table * t);

/* Routine for fixed-order Gauss-Legendre integration */

double gsl_integration_glfixed (const gsl_function *f,
                                double a,
                                double b,
                                const gsl_integration_glfixed_table * t);

/* Routine to retrieve the i-th Gauss-Legendre point and weight from t */

int gsl_integration_glfixed_point (double a,
                                   double b,
                                   size_t i,
                                   double *xi,
                                   double *wi,
                                   const gsl_integration_glfixed_table * t);

/* Cquad integration - Pedro Gonnet */

/* Data of a single interval */
typedef struct
{
  double a, b;
  double c[64];
  double fx[33];
  double igral, err;
  int depth, rdepth, ndiv;
} gsl_integration_cquad_ival;

/* The workspace is just a collection of intervals */
typedef struct
{
  size_t size;
  gsl_integration_cquad_ival *ivals;
  size_t *heap;
} gsl_integration_cquad_workspace;

gsl_integration_cquad_workspace *
gsl_integration_cquad_workspace_alloc (const size_t n);

void
gsl_integration_cquad_workspace_free (gsl_integration_cquad_workspace * w);

int
gsl_integration_cquad (const gsl_function * f, double a, double b,
		                   double epsabs, double epsrel,
		                   gsl_integration_cquad_workspace * ws,
		                   double *result, double *abserr, size_t * nevals);

/* Romberg integration workspace and routines */

typedef struct
{
  size_t n;       /* maximum number of steps */
  double *work1;  /* workspace for a row of R matrix, size n */
  double *work2;  /* workspace for a row of R matrix, size n */
} gsl_integration_romberg_workspace;

gsl_integration_romberg_workspace *gsl_integration_romberg_alloc(const size_t n);
void gsl_integration_romberg_free(gsl_integration_romberg_workspace * w);
int gsl_integration_romberg(const gsl_function * f, const double a, const double b,
                            const double epsabs, const double epsrel, double * result,
                            size_t * neval, gsl_integration_romberg_workspace * w);

/* IQPACK related structures and routines */

typedef struct
{
  double alpha;
  double beta;
  double a;
  double b;
  double zemu;
  double shft;
  double slp;
  double al;
  double be;
} gsl_integration_fixed_params;

typedef struct
{
  int (*check)(const size_t n, const gsl_integration_fixed_params * params);
  int (*init)(const size_t n, double * diag, double * subdiag, gsl_integration_fixed_params * params);
} gsl_integration_fixed_type;

typedef struct
{
  size_t n;        /* number of nodes/weights */
  double *weights; /* quadrature weights */
  double *x;       /* quadrature nodes */
  double *diag;    /* diagonal of Jacobi matrix */
  double *subdiag; /* subdiagonal of Jacobi matrix */
  const gsl_integration_fixed_type * type;
} gsl_integration_fixed_workspace;

/* IQPACK integral types */
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_legendre;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_chebyshev;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_gegenbauer;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_jacobi;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_laguerre;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_hermite;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_exponential;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_rational;
GSL_VAR const gsl_integration_fixed_type * gsl_integration_fixed_chebyshev2;

gsl_integration_fixed_workspace *
gsl_integration_fixed_alloc(const gsl_integration_fixed_type * type, const size_t n,
                            const double a, const double b, const double alpha, const double beta);

void gsl_integration_fixed_free(gsl_integration_fixed_workspace * w);

size_t gsl_integration_fixed_n(const gsl_integration_fixed_workspace * w);

double *gsl_integration_fixed_nodes(const gsl_integration_fixed_workspace * w);

double *gsl_integration_fixed_weights(const gsl_integration_fixed_workspace * w);

int gsl_integration_fixed(const gsl_function * func, double * result,
                          const gsl_integration_fixed_workspace * w);

__END_DECLS

#endif /* __GSL_INTEGRATION_H__ */


/** @file */

#ifndef __CCL_F2D_H_INCLUDED__
#define __CCL_F2D_H_INCLUDED__

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

CCL_BEGIN_DECLS

//f2d extrapolation types for early times
typedef enum ccl_f2d_extrap_growth_t
{
  ccl_f2d_cclgrowth = 401, //Use CCL's linear growth
  ccl_f2d_customgrowth = 402, //Use a custom growth function
  ccl_f2d_constantgrowth = 403, //Use a constant growth factor
  ccl_f2d_no_extrapol = 404, //Do not extrapolate, just throw an exception
} ccl_f2d_extrap_growth_t;

//f2d interpolation types
typedef enum ccl_f2d_interp_t
{
  ccl_f2d_3 = 303, //Bicubic interpolation
} ccl_f2d_interp_t;

/**
 * Struct containing a 2D power spectrum
 */
typedef struct {
  double lkmin,lkmax; /**< Edges in log(k)*/
  double amin,amax; /**< Edges in a*/
  int is_factorizable; /**< Is this factorizable into k- and a-dependent functions? */
  int is_k_constant; /**< no k-dependence, just return 1*/
  int is_a_constant; /**< no a-dependence, just return 1*/
  int extrap_order_lok; /**< Order of extrapolating polynomial in log(k) for low k (0, 1 or 2)*/
  int extrap_order_hik; /**< Order of extrapolating polynomial in log(k) for high k (0, 1 or 2)*/
  ccl_f2d_extrap_growth_t extrap_linear_growth;  /**< Extrapolation type at high redshifts*/
  int is_log; /**< Do I hold the values of log(f(k,a))?*/
  double (*growth)(double); /**< Custom extrapolating growth function*/
  double growth_factor_0; /**< Constant extrapolating growth factor*/
  int growth_exponent; /**< Power to which growth should be exponentiated*/
  gsl_spline *fk; /**< Spline holding the values of the k-dependent factor*/
  gsl_spline *fa; /**< Spline holding the values of the a-dependent factor*/
  gsl_spline2d *fka; /**< Spline holding the values of f(k,a)*/
} ccl_f2d_t;

/**
 * Create a ccl_f2d_t structure.
 * @param na number of elements in a_arr.
 * @param a_arr array of scale factor values at which the function is defined. The array should be ordered.
 * @param nk number of elements of lk_arr.
 * @param lk_arr array of logarithmic wavenumbers at which the function is defined (i.e. this array contains ln(k), NOT k). The array should be ordered.
 * @param fka_arr array of size na * nk containing the 2D function. The 2D ordering is such that fka_arr[ia*nk+ik] = f(k=exp(lk_arr[ik]),a=a_arr[ia]).
 * @param fk_arr array of size nk containing the k-dependent part of the function. Only relevant if is_factorizable is true.
 * @param fa_arr array of size na containing the a-dependent part of the function. Only relevant if is_factorizable is true.
 * @param is_factorizable if not 0, fk_arr and fa_arr will be used as 1-D arrays to construct a factorizable 2D function.
 * @param extrap_order_lok Order of the polynomial that extrapolates on wavenumbers smaller than the minimum of lk_arr. Allowed values: 0 (constant), 1 (linear extrapolation) and 2 (quadratic extrapolation). Extrapolation happens in ln(k).
 * @param extrap_order_hik Order of the polynomial that extrapolates on wavenumbers larger than the maximum of lk_arr. Allowed values: 0 (constant), 1 (linear extrapolation) and 2 (quadratic extrapolation). Extrapolation happens in ln(k).
 * @param extrap_linear_growth: ccl_f2d_extrap_growth_t value defining how the function with scale factors below the interpolation range. Allowed values: ccl_f2d_cclgrowth (scale with the CCL linear growth factor), ccl_f2d_customgrowth (scale with a custom function of redshift passed through `growth`), ccl_f2d_constantgrowth (scale by multiplying the function at the earliest available scale factor by a constant number, defined by `growth_factor_0`), ccl_f2d_no_extrapol (throw an error if the function is ever evaluated outside the interpolation range in a). Note that, above the interpolation range (i.e. for low redshifts), the function will be assumed constant.
 * @param is_fka_log: if not zero, `fka_arr` contains ln(f(k,a)) instead of f(k,a). If the function is factorizable, then `fk_arr` holds ln(K(k)) and `fa_arr` holds ln(A(a)), where f(k,a)=K(k)*A(a).
 * @param growth: custom growth function. Irrelevant if extrap_linear_growth!=ccl_f2d_customgrowth.
 * @param growth_factor_0: custom growth function. Irrelevant if extrap_linear_growth!=ccl_f2d_constantgrowth.
 * @param growth_exponent: power to which the extrapolating growth factor should be exponentiated when extrapolating (e.g. usually 2 for linear power spectra).
 * @param interp_type: 2D interpolation method. Currently only ccl_f2d_3 is implemented (bicubic interpolation).
 * @param status Status flag. 0 if there are no errors, nonzero otherwise.
 */
ccl_f2d_t *ccl_f2d_t_new(int na,double *a_arr,
			 int nk,double *lk_arr,
			 double *fka_arr,
			 double *fk_arr,
			 double *fa_arr,
			 int is_factorizable,
			 int extrap_order_lok,
			 int extrap_order_hik,
			 ccl_f2d_extrap_growth_t extrap_linear_growth,
			 int is_fka_log,
			 double (*growth)(double),
			 double growth_factor_0,
			 int growth_exponent,
			 ccl_f2d_interp_t interp_type,
			 int *status);

/**
 * Evaluate 2D function of k and a defined by ccl_f2d_t structure.
 * @param fka ccl_f2d_t structure defining f(k,a).
 * @param lk Natural logarithm of the wavenumber.
 * @param a Scale factor.
 * @param cosmo ccl_cosmology structure, only needed if evaluating f(k,a) at small scale factors outside the interpolation range, and if fka was initialized with extrap_linear_growth = ccl_f2d_cclgrowth.
 * @param status Status flag. 0 if there are no errors, nonzero otherwise.
 */
double ccl_f2d_t_eval(ccl_f2d_t *fka,double lk,double a,void *cosmo,
		      int *status);

/**
 * F2D structure destructor.
 * Frees up all memory associated with a f2d structure.
 * @param fka Structure to be freed.
 */
void ccl_f2d_t_free(ccl_f2d_t *fka);

CCL_END_DECLS

#endif


#include "mpi.h"
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <math.h>
#include <libiomp/omp.h>
#include <string.h>
#include <stdbool.h>

#ifdef __APPLE__

#include <cblas.h>

#define set_num_threads(x) openblas_set_num_threads(x)
#define get_num_threads() openblas_get_num_threads()
#else
#include <mkl.h>
#define set_num_threads(x) mkl_set_num_threads(x)
#define get_num_threads() mkl_get_num_threads()
#endif

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

    char *filename = NULL;
    int option = 0, omp_threads = 1;
    while ((option = getopt(argc, argv, "f:t:")) != -1) {
        switch (option) {
            case 'f':
                filename = optarg;
                break;
            case 't':
                omp_threads = atoi(optarg);
                break;
            default:
                printf("Usage: mpi_cg -f string -t num_threads \n");
                return 0;
        }
    }
    if (filename == NULL) {
        printf("Usage: mpi_cg -f string -t num_threads \n");
        return 0;
    }

    set_num_threads(omp_threads);
    omp_set_num_threads(omp_threads);

    int numTasks, rank;
    MPI_Init(&argc, &argv);
    MPI_Comm_size(MPI_COMM_WORLD, &numTasks);
    MPI_Comm_rank(MPI_COMM_WORLD, &rank);

    int qProcs = (int) sqrt(numTasks);
    if (qProcs - sqrt(numTasks) != 0) {
        if (rank == 0)
            printf("np must be a square number\n");
        MPI_Finalize();
        return 0;
    }

    int n, dims[2] = {qProcs, qProcs}, periods[2] = {0, 0}, reorder = 0, coords[2], rc;
    MPI_Comm cartComm;
    MPI_Cart_create(MPI_COMM_WORLD, 2, dims, periods, reorder, &cartComm);
    MPI_Comm_rank(cartComm, &rank);
    MPI_Cart_coords(cartComm, rank, 2, coords);

    MPI_File matFile;
    rc = MPI_File_open(cartComm, filename, MPI_MODE_RDONLY, MPI_INFO_NULL, &matFile);
    if (rc && (rank == 0)) {
        printf("Unable to open file %s\n", filename);
        fflush(stdout);
    }

    MPI_Status status;
    MPI_File_read(matFile, &n, 1, MPI_INT, &status);
    int blockSize = (n + qProcs - 1) / qProcs;  /* number of rows in _block_ */
    int lastBlock = n - (qProcs - 1) * blockSize;

    double *LocalA = (double *) malloc(blockSize * blockSize * sizeof(double));

    MPI_Datatype readFileType;
    MPI_Type_vector(blockSize, blockSize, n, MPI_DOUBLE, &readFileType);
    MPI_Type_commit(&readFileType);
    MPI_Offset offset = (MPI_Offset) (1 * sizeof(int) +
                                      sizeof(double) * (blockSize * coords[1] + blockSize * n * coords[0]));
    MPI_File_set_view(matFile, offset, MPI_DOUBLE, readFileType,
                      "native", MPI_INFO_NULL);

    MPI_File_read_at(matFile, 0, LocalA, blockSize * blockSize, MPI_DOUBLE, &status);
    MPI_Type_free(&readFileType);
    if (coords[0] == qProcs - 1) {
        memset(LocalA + lastBlock * blockSize, 0, (blockSize - lastBlock) * blockSize * sizeof(double));
    }
    if (coords[1] == qProcs - 1) {
        int kk = 0;
        for (kk = 0; kk < blockSize; ++kk) {
            memset(LocalA + kk * blockSize + lastBlock, 0, (blockSize - lastBlock) * sizeof(double));
        }
    }

    double *LocalB = (double *) malloc(blockSize * sizeof(double));
    memset(LocalB, 0, blockSize * sizeof(double));
    offset = (MPI_Offset) (1 * sizeof(int) + sizeof(double) * n * n);
    MPI_File_set_view(matFile, offset, MPI_DOUBLE, MPI_DOUBLE,
                      "native", MPI_INFO_NULL);
    offset = coords[1] * blockSize;
    MPI_File_read_at_all(matFile, offset, LocalB, coords[1] == qProcs - 1 ? lastBlock : blockSize, MPI_DOUBLE, &status);

    MPI_File_close(&matFile);
    double mpi_start = MPI_Wtime();
    MPI_Comm rowComm, colComm;
    MPI_Comm_split(MPI_COMM_WORLD, coords[0], rank, &rowComm);
    MPI_Comm_split(MPI_COMM_WORLD, coords[1], rank, &colComm);

    double tolerance = 1e-8;
    double alpha, beta, rho_new = 0.0, rho_old = 0.0;
    int iteration = 0;
    int world_root = numTasks - 1;
    double *loc_Vr = (double *) malloc(blockSize * sizeof(double));
    double *loc_Vp = (double *) malloc(blockSize * sizeof(double));
    double *loc_Vw = (double *) malloc(blockSize * sizeof(double));
    double *loc_Vx = (double *) malloc(blockSize * sizeof(double));

    int row_rank, col_rank;
    MPI_Comm_rank(colComm, &col_rank);
    MPI_Comm_rank(rowComm, &row_rank);

    //localr <- localB
    memset(loc_Vx, 0, blockSize * sizeof(double));
    cblas_dcopy(blockSize, LocalB, 1, loc_Vr, 1);
    double b_norm2 = 0;
    if (col_rank == qProcs - 1) {
        double loc_b_dot = cblas_ddot((row_rank == qProcs - 1 ? lastBlock : blockSize), loc_Vr, 1, loc_Vr, 1);
        MPI_Reduce(&loc_b_dot, &rho_new, 1, MPI_DOUBLE, MPI_SUM, qProcs - 1, rowComm);
        b_norm2 = sqrt(rho_new);
    }

    int loopBreak = 0;
    while (true) {
        if (rank == world_root) {
            loopBreak = (sqrt(rho_new) / b_norm2 < tolerance) ? 1 : 0;
        }
        MPI_Bcast(&loopBreak, 1, MPI_INT, world_root, MPI_COMM_WORLD);
        if (loopBreak > 0)
            break;

        iteration++;
        if (iteration == 1) {
            cblas_dcopy(blockSize, loc_Vr, 1, loc_Vp, 1);
        } else {
            if (rank == world_root) {
                beta = rho_new / rho_old;
            }
            MPI_Bcast(&beta, 1, MPI_DOUBLE, world_root, MPI_COMM_WORLD);
            cblas_dscal((row_rank == qProcs - 1 ? lastBlock : blockSize), beta, loc_Vp, 1);
            cblas_daxpy((row_rank == qProcs - 1 ? lastBlock : blockSize), 1.0, loc_Vr, 1, loc_Vp, 1);
        }

        cblas_dgemv(CblasRowMajor, CblasNoTrans, (col_rank == qProcs - 1 ? lastBlock : blockSize),
                    (row_rank == qProcs - 1 ? lastBlock : blockSize),
                    1.0, LocalA, blockSize, loc_Vp, 1, 0.0, loc_Vw, 1);

        if (col_rank == row_rank) {
            MPI_Reduce(MPI_IN_PLACE, loc_Vw, blockSize, MPI_DOUBLE, MPI_SUM, col_rank, rowComm);
        } else {
            MPI_Reduce(loc_Vw, loc_Vw, blockSize, MPI_DOUBLE, MPI_SUM, col_rank, rowComm);
        }

        MPI_Bcast(loc_Vw, (row_rank == qProcs - 1 ? lastBlock : blockSize), MPI_DOUBLE, row_rank, colComm);

        if (col_rank == qProcs - 1) {
            double loc_pw_dot = cblas_ddot((row_rank == qProcs - 1 ? lastBlock : blockSize), loc_Vp, 1, loc_Vw,
                                           1), glo_pw_dot;
            MPI_Reduce(&loc_pw_dot, &glo_pw_dot, 1, MPI_DOUBLE, MPI_SUM, qProcs - 1, rowComm);
            if (rank == world_root) {
                alpha = rho_new / glo_pw_dot;
            }
        }
        MPI_Bcast(&alpha, 1, MPI_DOUBLE, world_root, MPI_COMM_WORLD);

        cblas_daxpy((row_rank == qProcs - 1 ? lastBlock : blockSize), -alpha, loc_Vw, 1, loc_Vr, 1);

        if (coords[0] == qProcs - 1) {
            rho_old = rho_new;
            cblas_daxpy((row_rank == qProcs - 1 ? lastBlock : blockSize), alpha, loc_Vp, 1, loc_Vx, 1);
            double loc_r_dot = cblas_ddot(blockSize, loc_Vr, 1, loc_Vr, 1);

            MPI_Reduce(&loc_r_dot, &rho_new, 1, MPI_DOUBLE, MPI_SUM, qProcs - 1, rowComm);
        }
    }

    if (rank == world_root) {
        printf("mpi_cg_it: %d\t", iteration);
        printf("mpi_cg_time: %f\n", MPI_Wtime() - mpi_start);

    }
    MPI_Comm_free(&cartComm);
    MPI_Comm_free(&colComm);
    MPI_Comm_free(&rowComm);

    if (col_rank == qProcs - 1) {
        for (rc = 0; rc < (row_rank == qProcs - 1 ? lastBlock : blockSize); ++rc) {
            if (fabs(loc_Vx[rc] - 1) > 1e-3) {
                printf("MPI_CG_ERR %f \n", loc_Vx[rc]);
                break;
            }
        }
    }
    MPI_Finalize();
}


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

/* These tests are based on the NIST Statistical Reference Datasets
   See http://www.nist.gov/itl/div898/strd/index.html for more
   information. */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_test.h>
#include <gsl/gsl_multifit.h>
#include <gsl/gsl_multifit_nlin.h>
#include <gsl/gsl_blas.h>

#include <gsl/gsl_ieee_utils.h>

#include "test_longley.c"
#include "test_filip.c"
#include "test_pontius.c"
#include "test_brown.c"
#include "test_enso.c"
#include "test_kirby2.c"
#include "test_hahn1.c"
#include "test_nelson.c"
#include "test_fn.c"
#include "test_estimator.c"

void
test_lmder (gsl_multifit_function_fdf * f, double x0[],
            double * X, double F[], double * cov);

void
test_fdf (const char * name, gsl_multifit_function_fdf * f,
          double x0[], double x[], double sumsq,
          double sigma[]);

int
main (void)
{
  gsl_ieee_env_setup();

  test_longley();
  test_filip();
  test_pontius();
  test_estimator();

  {
    gsl_multifit_function_fdf f = make_fdf (&brown_f, &brown_df, &brown_fdf,
                                            brown_N, brown_P, 0);

    test_lmder(&f, brown_x0, &brown_X[0][0], brown_F, &brown_cov[0][0]);
  }

  {
    gsl_multifit_function_fdf f = make_fdf (&enso_f, &enso_df, &enso_fdf,
                                            enso_N, enso_P, 0);

    test_fdf("nist-ENSO", &f, enso_x0, enso_x, enso_sumsq, enso_sigma);
  }

  {
    gsl_multifit_function_fdf f = make_fdf (&kirby2_f, &kirby2_df, &kirby2_fdf,
                                            kirby2_N, kirby2_P, 0);

    test_fdf("nist-kirby2", &f, kirby2_x0, kirby2_x, kirby2_sumsq, kirby2_sigma);
  }

  {
    gsl_multifit_function_fdf f = make_fdf (&hahn1_f, &hahn1_df, &hahn1_fdf,
                                            hahn1_N, hahn1_P, 0);

    test_fdf("nist-hahn1", &f, hahn1_x0, hahn1_x, hahn1_sumsq, hahn1_sigma);
  }

#ifdef JUNK
  {
    gsl_multifit_function_fdf f = make_fdf (&nelson_f, &nelson_df, &nelson_fdf,
                                            nelson_N, nelson_P, 0);

    test_fdf("nist-nelson", &f, nelson_x0, nelson_x, nelson_sumsq, nelson_sigma);
  }
#endif

  /* now summarize the results */

  exit (gsl_test_summary ());
}

void
test_lmder (gsl_multifit_function_fdf * f, double x0[],
            double * X, double F[], double * cov)
{
  const gsl_multifit_fdfsolver_type *T;
  gsl_multifit_fdfsolver *s;

  const size_t n = f->n;
  const size_t p = f->p;

  int status;
  size_t iter = 0, i;

  gsl_vector_view x = gsl_vector_view_array (x0, p);

  T = gsl_multifit_fdfsolver_lmsder;
  s = gsl_multifit_fdfsolver_alloc (T, n, p);
  gsl_multifit_fdfsolver_set (s, f, &x.vector);

  do
    {
      status = gsl_multifit_fdfsolver_iterate (s);

      for (i = 0 ; i < p; i++)
        {
          gsl_test_rel (gsl_vector_get (s->x, i), X[p*iter+i], 1e-5,
                        "lmsder, iter=%u, x%u", iter, i);
        }

      gsl_test_rel (gsl_blas_dnrm2 (s->f), F[iter], 1e-5,
                    "lmsder, iter=%u, f", iter);

      iter++;
    }
  while (iter < 20);

  {
    size_t i, j;
    gsl_matrix * covar = gsl_matrix_alloc (4, 4);
    gsl_multifit_covar (s->J, 0.0, covar);

    for (i = 0; i < 4; i++)
      {
        for (j = 0; j < 4; j++)
          {
            gsl_test_rel (gsl_matrix_get(covar,i,j), cov[i*p + j], 1e-7,
                          "gsl_multifit_covar cov(%d,%d)", i, j) ;
          }
      }

    gsl_matrix_free (covar);
  }

  gsl_multifit_fdfsolver_free (s);

}

void
test_fdf (const char * name, gsl_multifit_function_fdf * f,
          double x0[], double x_final[],
          double f_sumsq, double sigma[])
{
  const gsl_multifit_fdfsolver_type *T;
  gsl_multifit_fdfsolver *s;

  const size_t n = f->n;
  const size_t p = f->p;

  int status;
  size_t iter = 0;

  gsl_vector_view x = gsl_vector_view_array (x0, p);

  T = gsl_multifit_fdfsolver_lmsder;
  s = gsl_multifit_fdfsolver_alloc (T, n, p);
  gsl_multifit_fdfsolver_set (s, f, &x.vector);

  do
    {
      status = gsl_multifit_fdfsolver_iterate (s);

#ifdef DEBUG
       printf("iter = %d  status = %d  |f| = %.18e x = \n",
         iter, status, gsl_blas_dnrm2 (s->f));

         gsl_vector_fprintf(stdout, s->x, "%.8e");
#endif
      status = gsl_multifit_test_delta (s->dx, s->x, 0.0, 1e-7);

      iter++;
    }
  while (status == GSL_CONTINUE && iter < 1000);

  {
    size_t i;
    gsl_matrix * covar = gsl_matrix_alloc (p, p);
    gsl_multifit_covar (s->J, 0.0, covar);

    for (i = 0 ; i < p; i++)
      {
        gsl_test_rel (gsl_vector_get (s->x, i), x_final[i], 1e-5,
                      "%s, lmsder, x%u", name, i);
      }

    {
      double s2 = pow(gsl_blas_dnrm2 (s->f), 2.0);

      gsl_test_rel (s2, f_sumsq, 1e-5, "%s, lmsder, |f|^2", name);

      for (i = 0; i < p; i++)
        {
          double ei = sqrt(s2/(n-p))*sqrt(gsl_matrix_get(covar,i,i));
          gsl_test_rel (ei, sigma[i], 1e-4,
                        "%s, sigma(%d)", name, i) ;
        }
    }

    gsl_matrix_free (covar);
  }

  /* Check that there is no hidden state, restarting should
     produce identical results. */

  {
    int status0, status1;
    size_t i;
    gsl_multifit_fdfsolver *t = gsl_multifit_fdfsolver_alloc (T, n, p);
    gsl_multifit_fdfsolver_set (t, f, &x.vector);

    /* do a few extra iterations to stir things up */

    gsl_multifit_fdfsolver_set (s, f, &x.vector);

    for (i = 0; i < 3; i++)
      {
        gsl_multifit_fdfsolver_iterate (s);
      }

    gsl_multifit_fdfsolver_set (s, f, &x.vector);

    do
      {
        status0 = gsl_multifit_fdfsolver_iterate (s);
        status1 = gsl_multifit_fdfsolver_iterate (t);

        gsl_test_int(status0, status1, "%s, lmsder status after set iter=%u", name, iter);

        for (i = 0; i < p; i++) {
          double sxi = gsl_vector_get(s->x,i);
          double txi = gsl_vector_get(t->x,i);
#ifdef DEBUG
          printf("%d %g %g\n", i, sxi, txi);
#endif
          gsl_test_rel(sxi, txi, 1e-15, "%s, lmsder after set, %u/%u", name, iter, i);
        }

#ifdef DEBUG
        printf("iter = %d  status = %d  |f| = %.18e x = \n",
               iter, status, gsl_blas_dnrm2 (s->f));

        gsl_vector_fprintf(stdout, s->x, "%.8e");
#endif
        status0 = gsl_multifit_test_delta (s->dx, s->x, 0.0, 1e-7);
        status1 = gsl_multifit_test_delta (t->dx, s->x, 0.0, 1e-7);

        gsl_test_int(status0, status1, "%s, lmsder test delta status after set iter=%u", name, iter);

        iter++;
      }
    while (status1 == GSL_CONTINUE && iter < 1000);

    gsl_multifit_fdfsolver_free (t);
  }

  gsl_multifit_fdfsolver_free (s);
}