#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_fit.h>

#include "cosmocalc.h"

double transfer_function(double k)
{
  static int initFlag = 1;
  static int currCosmoNum;
  static gsl_spline *cosmocalc_transfer_function_spline = NULL;
  static gsl_interp_accel *cosmocalc_transfer_function_acc = NULL;
  static double c0,c1;

  double transfer_function_table[COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH];
  double k_table[COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH];
  long i;
  double cov00,cov01,cov11,sumsq;

  if(initFlag == 1 || currCosmoNum != cosmoData.cosmoNum)
    {
      initFlag = 0;
      currCosmoNum = cosmoData.cosmoNum;

      for(i=0;i<COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH;++i)
        {
          k_table[i] = log(K_MAX/K_MIN)/(COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH-1.0)*((double) i) + log(K_MIN);
          if(cosmoData.useSmoothTransFunc)
	    transfer_function_table[i] = log(transfunct_eh98_smooth(exp(k_table[i])));
	  else if(cosmoData.useNoTransFunc)
	    transfer_function_table[i] = 0.0;
	  else
	    transfer_function_table[i] = log(transfunct_eh98(exp(k_table[i])));
        }

      //init the spline and accelerators
      if(cosmocalc_transfer_function_spline != NULL)
        gsl_spline_free(cosmocalc_transfer_function_spline);
      cosmocalc_transfer_function_spline = gsl_spline_alloc(GSL_SPLINE_TYPE,(size_t) (COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH));
      gsl_spline_init(cosmocalc_transfer_function_spline,k_table,transfer_function_table,(size_t) (COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH));
      if(cosmocalc_transfer_function_acc != NULL)
        gsl_interp_accel_reset(cosmocalc_transfer_function_acc);
      else
        cosmocalc_transfer_function_acc = gsl_interp_accel_alloc();

      gsl_fit_linear(k_table+COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH-COSMOCALC_TRANSFER_FUNCTION_FIT_LENGTH,(size_t) 1,
                     transfer_function_table+COSMOCALC_TRANSFER_FUNCTION_TABLE_LENGTH-COSMOCALC_TRANSFER_FUNCTION_FIT_LENGTH,(size_t) 1,
                     (size_t) COSMOCALC_TRANSFER_FUNCTION_FIT_LENGTH,&c0,&c1,&cov00,&cov01,&cov11,&sumsq);
    }

  if(k < K_MIN)
    return 1.0;
  else if(k < K_MAX)
    return exp(gsl_spline_eval(cosmocalc_transfer_function_spline,log(k),cosmocalc_transfer_function_acc));
  else
    return exp(c0+c1*log(k));
}

#define USEPAPER_EH98

double transfunct_eh98(double kin)
{
  //vars
  double k;
  double Tk,Tb,Tc,T0t,T0t11,T0t1bc;
  double omb,om0,omc,h;
  double theta2p7,f,ac,bc,s,q;
  double C,C11,C1bc,st,bb,ab,ksilk,bnode,Gy,y,zeq;
  double a1,a2,b1,b2,keq,Rd,zd,Req,b1d,b2d;

  //get cosmoparms
  omb = cosmoData.OmegaB;
  om0 = cosmoData.OmegaM;
  omc = cosmoData.OmegaM - cosmoData.OmegaB;
  h = cosmoData.h;

  //convert k from hMpc^-1 to Mpc^-1
  k = kin*h;

  //-----------
  //input parms
  //-----------
  theta2p7 = TCMB/2.7;

  //eqn 2
  zeq = 2.50e4*om0*h*h/(theta2p7*theta2p7*theta2p7*theta2p7);

  //eqn 3
  keq = 7.46e-2*om0*h*h/(theta2p7*theta2p7); //Mpc^-{1} (NOT h/Mpc)

  //eqn 4
  b1d = 0.313*pow(om0*h*h,-0.419)*(1.0 + 0.607*pow(om0*h*h,0.674));
  b2d = 0.238*pow(om0*h*h,0.223);
  zd = 1291.0*pow(om0*h*h,0.251)/(1.0 + 0.659*pow(om0*h*h,0.828))
    *(1.0 + b1d*pow(omb*h*h,b2d));

  //eqn 5
#ifdef USEPAPER_EH98
  Rd = 31.5*omb*h*h
    /(theta2p7*theta2p7*theta2p7*theta2p7)
    /((zd)/1e3); //matches wayne's paper, but not the code (zd -> 1+zd)
  //fprintf(stderr,"paper eh 98 1\n");
#else
  Rd = 31.5*omb*h*h
    /(theta2p7*theta2p7*theta2p7*theta2p7)
    /((1.0+zd)/1e3);
#endif
Req = 31.5*omb*h*h/(theta2p7*theta2p7*theta2p7*theta2p7)/(zeq/1e3);

  //eqn 6
  s = 2.0/3.0/keq*sqrt(6.0/Req)*
    log((sqrt(1.0 + Rd) + sqrt(Rd + Req))/(1.0 + sqrt(Req)));

  //eqn 7
  ksilk = 1.6*pow(omb*h*h,0.52)*pow(om0*h*h,0.73)*(1.0 + pow(10.4*om0*h*h,-0.95));

  //eqn 10
  q = k/13.41/keq;

  //eqn 11
  a1 = pow(46.9*om0*h*h,0.670)*(1.0 + pow(32.1*om0*h*h,-0.532));
  a2 = pow(12.0*om0*h*h,0.424)*(1.0 + pow(45.0*om0*h*h,-0.582));
  ac = pow(a1,-1.0*omb/om0)*pow(a2,-1.0*(omb/om0)*(omb/om0)*(omb/om0));

  //eqn 12
  b1 = 0.944/(1.0 + pow(458.0*om0*h*h,-0.708));
  b2 = pow(0.395*om0*h*h,-0.0266);
  bc = 1.0/(1.0 + b1*(pow(omc/om0,b2) - 1.0));

  //eqn 15
#ifdef USEPAPER_EH98
  y = (1.0 + zeq)/(1.0 + zd); //matches wayne's paper, but not the code (1+zeq -> zeq)
  //fprintf(stderr,"paper eh 98 2\n");
#else
  y = (zeq)/(1.0 + zd);
#endif
  Gy = y*(-6.0*sqrt(1.0 + y) + (2.0 + 3.0*y)*log((sqrt(1.0 + y) + 1.0)/(sqrt(1.0 + y) - 1.0)));

  //eqn 14
  ab = 2.07*keq*s*pow(1.0 + Rd,-3.0/4.0)*Gy;

  //----------------------------------
  // Get CDM part of transfer function
  //----------------------------------

  //eqn 18
  f = 1.0/(1.0 + (k*s/5.4)*(k*s/5.4)*(k*s/5.4)*(k*s/5.4));

  //eqn 20
  C = 14.2/ac + 386.0/(1.0 + 69.9*pow(q,1.08));

  //eqn 19
  T0t = log(M_E + 1.8*bc*q)/(log(M_E + 1.8*bc*q) + C*q*q);

  //eqn 17
  C1bc = 14.2 + 386.0/(1.0 + 69.9*pow(q,1.08));
  T0t1bc = log(M_E + 1.8*bc*q)/(log(M_E + 1.8*bc*q) + C1bc*q*q);
  Tc = f*T0t1bc + (1.0 - f)*T0t;

  //-------------------------------------
  // Get baryon part of transfer function
  //-------------------------------------

  //eqn 24
  bb = 0.5 + omb/om0 + (3.0 - 2.0*omb/om0)*sqrt((17.2*om0*h*h)*(17.2*om0*h*h) + 1.0);

  //eqn 23
  bnode = 8.41*pow(om0*h*h,0.435);

  //eqn 22
  st = s/pow(1.0 + (bnode/k/s)*(bnode/k/s)*(bnode/k/s),1.0/3.0);

  //eqn 21
  C11 = 14.2 + 386.0/(1.0 + 69.9*pow(q,1.08));
  T0t11 = log(M_E + 1.8*q)/(log(M_E + 1.8*q) + C11*q*q);
  Tb = (T0t11/(1.0 + (k*s/5.2)*(k*s/5.2))
	+ ab/(1.0 + (bb/k/s)*(bb/k/s)*(bb/k/s))/exp(pow(k/ksilk,1.4)))*sin(k*st)/(k*st);

  //------------------------
  // total transfer function
  //------------------------
  Tk = omb/om0*Tb + omc/om0*Tc;

  return Tk;
}

double transfunct_eh98_smooth(double kin)
{
  //vars
  double k,Tk;
  double omb,om0,omc,h;
  double theta2p7,s,q;
  double Gamma,alphaGamma,L0,C0;

  //get cosmoparms
  omb = cosmoData.OmegaB;
  om0 = cosmoData.OmegaM;
  omc = cosmoData.OmegaM - cosmoData.OmegaB;
  h = cosmoData.h;

  //convert k from hMpc^-1 to Mpc^-1
  k = kin*h;

  //-----------
  //input parms
  //-----------
  theta2p7 = TCMB/2.7;

  //eqn 26
  s = 44.5*log(9.83/om0/h/h)/sqrt(1.0 + 10.0*pow(omb*h*h,0.75));

  //eqn 31
  alphaGamma = 1.0 - 0.328*log(431.0*om0*h*h)*omb/om0 + 0.38*log(22.3*om0*h*h)*(omb/om0)*(omb/om0);

  //eqn 30
  Gamma = om0*h*(alphaGamma + (1.0 - alphaGamma)/(1.0 + pow(0.43*k*s,4.0)));

  //eqn 28
  q = kin*theta2p7*theta2p7/Gamma;

  //eqns 29
  C0 = 14.2 + 731.0/(1.0 + 62.5*q);
  L0 = log(2.0*exp(1.0) + 1.8*q);
  Tk = L0/(L0 + C0*q*q);

  return Tk;
}


/* linalg/balance.c
 *
 * Copyright (C) 2001 Brian Gough
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Balance a general matrix by scaling the columns
 *
 * B =  A D
 *
 * where D is a diagonal matrix
 */

#include <config.h>
#include <stdlib.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_blas.h>

#include "gsl_linalg.h"

int
gsl_linalg_balance_columns (gsl_matrix * A, gsl_vector * D)
{
  const size_t N = A->size2;
  size_t j;

  if (D->size != A->size2)
    {
      GSL_ERROR("length of D must match second dimension of A", GSL_EINVAL);
    }

  gsl_vector_set_all (D, 1.0);

  for (j = 0; j < N; j++)
    {
      gsl_vector_view A_j = gsl_matrix_column (A, j);

      double s = gsl_blas_dasum(&A_j.vector);

      double f = 1.0;

      if (s == 0.0)
        {
          gsl_vector_set (D, j, f);
          continue;
        }

      while (s > 1.0)
        {
          s /= 2.0;
          f *= 2.0;
        }

      while (s < 0.5)
        {
          s *= 2.0;
          f /= 2.0;
        }

      gsl_vector_set (D, j, f);

      if (f != 1.0)
        {
          gsl_blas_dscal(1.0/f, &A_j.vector);
        }
    }

  return GSL_SUCCESS;
}


/**
 *
 * @precisions normal z -> c d s
 *
 **/
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <cblas.h>
#include <lapacke.h>
#include <plasma.h>
#include <core_blas.h>
#include "auxiliary.h"

/*-------------------------------------------------------------------
 * Check the orthogonality of Q
 */

int z_check_orthogonality(int M, int N, int LDQ, PLASMA_Complex64_t *Q)
{
    double alpha, beta;
    double normQ;
    int info_ortho;
    int i;
    int minMN = min(M, N);
    double eps;
    double *work = (double *)malloc(minMN*sizeof(double));

    eps = LAPACKE_dlamch_work('e');
    alpha = 1.0;
    beta  = -1.0;

    /* Build the idendity matrix USE DLASET?*/
    PLASMA_Complex64_t *Id = (PLASMA_Complex64_t *) malloc(minMN*minMN*sizeof(PLASMA_Complex64_t));
    memset((void*)Id, 0, minMN*minMN*sizeof(PLASMA_Complex64_t));
    for (i = 0; i < minMN; i++)
        Id[i*minMN+i] = (PLASMA_Complex64_t)1.0;

    /* Perform Id - Q'Q */
    if (M >= N)
        cblas_zherk(CblasColMajor, CblasUpper, CblasConjTrans, N, M, alpha, Q, LDQ, beta, Id, N);
    else
        cblas_zherk(CblasColMajor, CblasUpper, CblasNoTrans, M, N, alpha, Q, LDQ, beta, Id, M);

    normQ = LAPACKE_zlansy_work(LAPACK_COL_MAJOR, 'i', 'u', minMN, Id, minMN, work);

    printf("============\n");
    printf("Checking the orthogonality of Q \n");
    printf("||Id-Q'*Q||_oo / (N*eps) = %e \n",normQ/(minMN*eps));

    if ( isnan(normQ / (minMN * eps)) || (normQ / (minMN * eps) > 10.0) ) {
        printf("-- Orthogonality is suspicious ! \n");
        info_ortho=1;
    }
    else {
        printf("-- Orthogonality is CORRECT ! \n");
        info_ortho=0;
    }

    free(work); free(Id);

    return info_ortho;
}

/*------------------------------------------------------------
 *  Check the factorization QR
 */

int z_check_QRfactorization(int M, int N, PLASMA_Complex64_t *A1, PLASMA_Complex64_t *A2, int LDA, PLASMA_Complex64_t *Q)
{
    double Anorm, Rnorm;
    PLASMA_Complex64_t alpha, beta;
    int info_factorization;
    int i,j;
    double eps;

    eps = LAPACKE_dlamch_work('e');

    PLASMA_Complex64_t *Ql       = (PLASMA_Complex64_t *)malloc(M*N*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *Residual = (PLASMA_Complex64_t *)malloc(M*N*sizeof(PLASMA_Complex64_t));
    double *work              = (double *)malloc(max(M,N)*sizeof(double));

    alpha=1.0;
    beta=0.0;

    if (M >= N) {
        /* Extract the R */
        PLASMA_Complex64_t *R = (PLASMA_Complex64_t *)malloc(N*N*sizeof(PLASMA_Complex64_t));
        memset((void*)R, 0, N*N*sizeof(PLASMA_Complex64_t));
        LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,'u', M, N, A2, LDA, R, N);

        /* Perform Ql=Q*R */
        memset((void*)Ql, 0, M*N*sizeof(PLASMA_Complex64_t));
        cblas_zgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, M, N, N, CBLAS_SADDR(alpha), Q, LDA, R, N, CBLAS_SADDR(beta), Ql, M);
        free(R);
    }
    else {
        /* Extract the L */
        PLASMA_Complex64_t *L = (PLASMA_Complex64_t *)malloc(M*M*sizeof(PLASMA_Complex64_t));
        memset((void*)L, 0, M*M*sizeof(PLASMA_Complex64_t));
        LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,'l', M, N, A2, LDA, L, M);

    /* Perform Ql=LQ */
        memset((void*)Ql, 0, M*N*sizeof(PLASMA_Complex64_t));
        cblas_zgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, M, N, M, CBLAS_SADDR(alpha), L, M, Q, LDA, CBLAS_SADDR(beta), Ql, M);
        free(L);
    }

    /* Compute the Residual */
    for (i = 0; i < M; i++)
        for (j = 0 ; j < N; j++)
            Residual[j*M+i] = A1[j*LDA+i]-Ql[j*M+i];

    Rnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N, Residual, M, work);
    Anorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N, A2, LDA, work);

    if (M >= N) {
        printf("============\n");
        printf("Checking the QR Factorization \n");
        printf("-- ||A-QR||_oo/(||A||_oo.N.eps) = %e \n",Rnorm/(Anorm*N*eps));
    }
    else {
        printf("============\n");
        printf("Checking the LQ Factorization \n");
        printf("-- ||A-LQ||_oo/(||A||_oo.N.eps) = %e \n",Rnorm/(Anorm*N*eps));
    }

    if (isnan(Rnorm / (Anorm * N *eps)) || (Rnorm / (Anorm * N * eps) > 10.0) ) {
        printf("-- Factorization is suspicious ! \n");
        info_factorization = 1;
    }
    else {
        printf("-- Factorization is CORRECT ! \n");
        info_factorization = 0;
    }

    free(work); free(Ql); free(Residual);

    return info_factorization;
}

/*------------------------------------------------------------------------
 *  Check the factorization of the matrix A2
 */

int z_check_LLTfactorization(int N, PLASMA_Complex64_t *A1, PLASMA_Complex64_t *A2, int LDA, int uplo)
{
    double Anorm, Rnorm;
    PLASMA_Complex64_t alpha;
    int info_factorization;
    int i,j;
    double eps;

    eps = LAPACKE_dlamch_work('e');

    PLASMA_Complex64_t *Residual = (PLASMA_Complex64_t *)malloc(N*N*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *L1       = (PLASMA_Complex64_t *)malloc(N*N*sizeof(PLASMA_Complex64_t));
    PLASMA_Complex64_t *L2       = (PLASMA_Complex64_t *)malloc(N*N*sizeof(PLASMA_Complex64_t));
    double *work              = (double *)malloc(N*sizeof(double));

    memset((void*)L1, 0, N*N*sizeof(PLASMA_Complex64_t));
    memset((void*)L2, 0, N*N*sizeof(PLASMA_Complex64_t));

    alpha= 1.0;

    LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,' ', N, N, A1, LDA, Residual, N);

    /* Dealing with L'L or U'U  */
    if (uplo == PlasmaUpper){
        LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,'u', N, N, A2, LDA, L1, N);
        LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,'u', N, N, A2, LDA, L2, N);
        cblas_ztrmm(CblasColMajor, CblasLeft, CblasUpper, CblasConjTrans, CblasNonUnit, N, N, CBLAS_SADDR(alpha), L1, N, L2, N);
    }
    else{
        LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,'l', N, N, A2, LDA, L1, N);
        LAPACKE_zlacpy_work(LAPACK_COL_MAJOR,'l', N, N, A2, LDA, L2, N);
        cblas_ztrmm(CblasColMajor, CblasRight, CblasLower, CblasConjTrans, CblasNonUnit, N, N, CBLAS_SADDR(alpha), L1, N, L2, N);
    }

    /* Compute the Residual || A -L'L|| */
    for (i = 0; i < N; i++)
        for (j = 0; j < N; j++)
           Residual[j*N+i] = L2[j*N+i] - Residual[j*N+i];

    Rnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', N, N, Residual, N, work);
    Anorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', N, N, A1, LDA, work);

    printf("============\n");
    printf("Checking the Cholesky Factorization \n");
    printf("-- ||L'L-A||_oo/(||A||_oo.N.eps) = %e \n",Rnorm/(Anorm*N*eps));

    if ( isnan(Rnorm/(Anorm*N*eps)) || (Rnorm/(Anorm*N*eps) > 10.0) ){
        printf("-- Factorization is suspicious ! \n");
        info_factorization = 1;
    }
    else{
        printf("-- Factorization is CORRECT ! \n");
        info_factorization = 0;
    }

    free(Residual); free(L1); free(L2); free(work);

    return info_factorization;
}

/*--------------------------------------------------------------
 * Check the gemm
 */
double z_check_gemm(PLASMA_enum transA, PLASMA_enum transB, int M, int N, int K,
                   PLASMA_Complex64_t alpha, PLASMA_Complex64_t *A, int LDA,
                   PLASMA_Complex64_t *B, int LDB,
                   PLASMA_Complex64_t beta, PLASMA_Complex64_t *Cplasma,
                   PLASMA_Complex64_t *Cref, int LDC,
                   double *Cinitnorm, double *Cplasmanorm, double *Clapacknorm )
{
    PLASMA_Complex64_t beta_const = -1.0;
    double Rnorm;
    double *work = (double *)malloc(max(K,max(M, N))* sizeof(double));

    *Cinitnorm   = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N, Cref,    LDC, work);
    *Cplasmanorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N, Cplasma, LDC, work);

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

    *Clapacknorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N, Cref, LDC, work);

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

    Rnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N, Cref, LDC, work);

    free(work);

    return Rnorm;
}

/*--------------------------------------------------------------
 * Check the trsm
 */
double z_check_trsm(PLASMA_enum side, PLASMA_enum uplo, PLASMA_enum trans, PLASMA_enum diag,
                   int M, int NRHS, PLASMA_Complex64_t alpha,
                   PLASMA_Complex64_t *A, int LDA,
                   PLASMA_Complex64_t *Bplasma, PLASMA_Complex64_t *Bref, int LDB,
                   double *Binitnorm, double *Bplasmanorm, double *Blapacknorm )
{
    PLASMA_Complex64_t beta_const = -1.0;
    double Rnorm;
    double *work = (double *)malloc(max(M, NRHS)* sizeof(double));
    /*double eps = LAPACKE_dlamch_work('e');*/

    *Binitnorm   = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, NRHS, Bref,    LDB, work);
    *Bplasmanorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'm', M, NRHS, Bplasma, LDB, work);

    cblas_ztrsm(CblasColMajor, (CBLAS_SIDE)side, (CBLAS_UPLO)uplo,
                (CBLAS_TRANSPOSE)trans, (CBLAS_DIAG)diag, M, NRHS,
                CBLAS_SADDR(alpha), A, LDA, Bref, LDB);

    *Blapacknorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'm', M, NRHS, Bref, LDB, work);

    cblas_zaxpy(LDB * NRHS, CBLAS_SADDR(beta_const), Bplasma, 1, Bref, 1);

    Rnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'm', M, NRHS, Bref, LDB, work);
    Rnorm = Rnorm / *Blapacknorm;
    /* max(M,NRHS) * eps);*/

    free(work);

    return Rnorm;
}

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

double z_check_solution(int M, int N, int NRHS, PLASMA_Complex64_t *A, int LDA,
                      PLASMA_Complex64_t *B,  PLASMA_Complex64_t *X, int LDB,
                      double *anorm, double *bnorm, double *xnorm )
{
/*     int info_solution; */
    double Rnorm = -1.00;
    PLASMA_Complex64_t zone  =  1.0;
    PLASMA_Complex64_t mzone = -1.0;
    double *work = (double *)malloc(max(M, N)* sizeof(double));

    *anorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, N,    A, LDA, work);
    *xnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', M, NRHS, X, LDB, work);
    *bnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', N, NRHS, B, LDB, work);

    cblas_zgemm(CblasColMajor, CblasNoTrans, CblasNoTrans, M, NRHS, N, CBLAS_SADDR(zone), A, LDA, X, LDB, CBLAS_SADDR(mzone), B, LDB);

    Rnorm = LAPACKE_zlange_work(LAPACK_COL_MAJOR, 'i', N, NRHS, B, LDB, work);

    free(work);

    return Rnorm;
}


#include "parameters.h" // Includes all defined parameters
#include "axi.h" // Axisymmetric coordinates
#include "navier-stokes/centered.h" // To solve the Navier-Stokes
#include "two-phase.h" // Implements two-phase flow
// #include <omp.h> // For openMP parallel
#include <gsl/gsl_fit.h>

double * forces_array;
double * times_array;

double gradient = 0.5;
double intercept = 4.9;

int main () {

    forces_array = malloc(INTERP_NO * sizeof(double));
    times_array = malloc(INTERP_NO * sizeof(double));

    DT = 1e-2;

    run();

    free(forces_array);
    free(times_array);
}

event force(i++) {
    if (i < INTERP_NO) {
        forces_array[i] = gradient * t + intercept;
        times_array[i] = t;
    } else {
        double c0, c1, cov00, cov01, cov11, sumsq;
        gsl_fit_linear ( times_array, 1, forces_array, 1, INTERP_NO, &c0, &c1, \
            &cov00, &cov01, &cov11, &sumsq);

        fprintf(stderr, "t = %g, c0 = %g, c1 = %g\n", t, c0, c1);

        double current_force = gradient * t + intercept;
        double interp_force = c0 + c1 * t;
        fprintf(stderr, "f = %g, f_interp = %g", current_force, interp_force);
        fprintf(stderr, "\n");

        // #pragma omp critical
        for (int j = 0; j < INTERP_NO - 1; j++) {
            forces_array[j] = forces_array[j + 1];
            times_array[j] = times_array[j + 1];
        }
        forces_array[INTERP_NO - 1] = current_force;
        times_array[INTERP_NO - 1] = t;
    }
}

event end(t = 1) {
    fprintf(stderr, "Ended\n");
}


/*
 * testgen.c
 * Patrick Alken
 *
 * Compile: gcc -g -O2 -Wall -o testgen testgen.c -lm -lgsl -llapack -lf77blas -lcblas -latlas -lg2c
 *
 * Usage: testgen [options]
 *
 * -i             : incremental matrices
 * -z             : compute Schur vectors and test them
 * -n size        : size of matrices
 * -l lower-bound : lower bound for matrix elements
 * -u upper-bound : upper bound for matrix elements
 * -c num         : number of matrices to solve
 */

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

#include <gsl/gsl_math.h>
#include <gsl/gsl_eigen.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_ieee_utils.h>
#include <gsl/gsl_complex.h>
#include <gsl/gsl_complex_math.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_sort.h>
#include <gsl/gsl_sort_vector.h>

typedef struct
{
  gsl_eigen_gen_workspace *gen_p;
  gsl_matrix *A;
  gsl_matrix *B;
  gsl_vector_complex *alpha;
  gsl_vector *beta;
  gsl_vector_complex *evals;

  gsl_matrix *Q;
  gsl_matrix *Z;
  int compute_schur;

  size_t n_evals;
} gen_workspace;

gen_workspace *gen_alloc(size_t n, int compute_schur);
void gen_free(gen_workspace *w);
int gen_proc(gen_workspace *w);

typedef struct
{
  gsl_matrix *A;
  gsl_matrix *B;
  gsl_matrix *Q;
  gsl_matrix *Z;
  int N;

  char jobvsl;
  char jobvsr;
  char sort;
  int selctg;
  int lda;
  int ldb;
  int sdim;
  double *alphar;
  double *alphai;
  gsl_vector *beta;
  int ldvsr;
  int lwork;
  int info;
  double *work;

  gsl_vector_complex *evals;
  gsl_vector_complex *alpha;
  size_t n_evals;
} lapack_workspace;

lapack_workspace *lapack_alloc(const size_t n);
void lapack_free(lapack_workspace *w);
int lapack_proc(lapack_workspace *w);

void dgges_(char *jobvsl, char *jobvsr, char *sort, int *selctg, int *n,
            double *a, int *lda, double *b, int *ldb, int *sdim,
            double *alphar, double *alphai, double *beta, double *vsl,
            int *ldvsl, double *vsr, int *ldvsr, double *work, int *lwork,
            int *bwork, int *info);

/*
 * Global variables
 */
unsigned long count = 0;

/*
 * Prototypes
 */

void make_random_matrix(gsl_matrix *m, gsl_rng *r, int lower, int upper);
void make_random_integer_matrix(gsl_matrix *m, gsl_rng *r, int lower,
                                int upper);
void make_start_matrix(gsl_matrix *m, int lower);
int inc_matrix (gsl_matrix *m, int lower, int upper);
void output_matrix(gsl_matrix *m);
void print_matrix(gsl_matrix *m, const char *str);
int test_evals(gsl_vector_complex *obs, gsl_vector_complex *expected,
               gsl_matrix *A, gsl_matrix *B,
               const char *obsname, const char *expname);
int test_alpha(gsl_vector_complex *obs, gsl_vector_complex *expected,
               gsl_matrix *A, gsl_matrix *B, const char *obsname,
               const char *expname);
int test_beta(gsl_vector *obs, gsl_vector *expected,
              gsl_matrix *A, gsl_matrix *B, const char *obsname,
              const char *expname);
void test_schur(gsl_matrix *A, gsl_matrix *S, gsl_matrix *Q, gsl_matrix *Z);
void print_vector(gsl_vector_complex *eval, const char *str);
int cmp(double a, double b);
int compare(const void *a, const void *b);
void sort_complex_vector(gsl_vector_complex *v);

gen_workspace *
gen_alloc(size_t n, int compute_schur)
{
  gen_workspace *w;

  w = (gen_workspace *) calloc(1, sizeof(gen_workspace));

  w->gen_p = gsl_eigen_gen_alloc(n);

  w->A = gsl_matrix_alloc(n, n);
  w->B = gsl_matrix_alloc(n, n);
  w->alpha = gsl_vector_complex_alloc(n);
  w->beta = gsl_vector_alloc(n);
  w->evals = gsl_vector_complex_alloc(n);
  w->compute_schur = compute_schur;

  if (compute_schur)
    {
      w->Q = gsl_matrix_alloc(n, n);
      w->Z = gsl_matrix_alloc(n, n);
      gsl_eigen_gen_params(1, 1, 0, w->gen_p);
    }

  return (w);
} /* gen_alloc() */

void
gen_free(gen_workspace *w)
{
  if (w->gen_p)
    gsl_eigen_gen_free(w->gen_p);

  if (w->A)
    gsl_matrix_free(w->A);

  if (w->B)
    gsl_matrix_free(w->B);

  if (w->alpha)
    gsl_vector_complex_free(w->alpha);

  if (w->beta)
    gsl_vector_free(w->beta);

  if (w->evals)
    gsl_vector_complex_free(w->evals);

  if (w->Q)
    gsl_matrix_free(w->Q);

  if (w->Z)
    gsl_matrix_free(w->Z);

  free(w);
}

int
gen_proc(gen_workspace *w)
{
  int s;

  s = gsl_eigen_gen_QZ(w->A, w->B, w->alpha, w->beta, w->Q, w->Z, w->gen_p);

  w->n_evals = w->gen_p->n_evals;

  return s;
} /* gen_proc() */

lapack_workspace *
lapack_alloc(const size_t n)
{
  lapack_workspace *w;
  double work[1];

  w = (lapack_workspace *) calloc(1, sizeof(lapack_workspace));

  w->A = gsl_matrix_alloc(n, n);
  w->B = gsl_matrix_alloc(n, n);
  w->Q = gsl_matrix_alloc(n, n);
  w->Z = gsl_matrix_alloc(n, n);
  w->alphar = malloc(n * sizeof(double));
  w->alphai = malloc(n * sizeof(double));
  w->beta = gsl_vector_alloc(n);
  w->alpha = gsl_vector_complex_alloc(n);
  w->evals = gsl_vector_complex_alloc(n);

  w->N = (int) n;
  w->n_evals = 0;

  w->jobvsl = 'N';
  w->jobvsr = 'N';
  w->sort = 'N';
  w->info = 0;

  w->lwork = -1;
  dgges_(&w->jobvsl,
         &w->jobvsr,
         &w->sort,
         (int *) 0,
         &w->N,
         w->A->data,
         (int *) &w->A->tda,
         w->B->data,
         (int *) &w->B->tda,
         &w->sdim,
         w->alphar,
         w->alphai,
         w->beta->data,
         w->Q->data,
         (int *) &w->Q->tda,
         w->Z->data,
         (int *) &w->Z->tda,
         work,
         &w->lwork,
         (int *) 0,
         &w->info);

  w->lwork = (int) work[0];
  w->work = malloc(w->lwork * sizeof(double));

  return (w);
} /* lapack_alloc() */

void
lapack_free(lapack_workspace *w)
{
  if (w->A)
    gsl_matrix_free(w->A);

  if (w->B)
    gsl_matrix_free(w->B);

  if (w->Q)
    gsl_matrix_free(w->Q);

  if (w->Z)
    gsl_matrix_free(w->Z);

  if (w->work)
    free(w->work);

  if (w->alphar)
    free(w->alphar);

  if (w->alphai)
    free(w->alphai);

  if (w->beta)
    gsl_vector_free(w->beta);

  if (w->alpha)
    gsl_vector_complex_free(w->alpha);

  if (w->evals)
    gsl_vector_complex_free(w->evals);

  free(w);
} /* lapack_free() */

int
lapack_proc(lapack_workspace *w)
{
  dgges_(&w->jobvsl,
         &w->jobvsr,
         &w->sort,
         (int *) 0,
         &w->N,
         w->A->data,
         (int *) &w->A->tda,
         w->B->data,
         (int *) &w->B->tda,
         &w->sdim,
         w->alphar,
         w->alphai,
         w->beta->data,
         w->Q->data,
         (int *) &w->Q->tda,
         w->Z->data,
         (int *) &w->Z->tda,
         w->work,
         &w->lwork,
         (int *) 0,
         &w->info);

  return (w->info);
} /* lapack_proc() */

/**********************************************
 * General routines
 **********************************************/

void
make_random_matrix(gsl_matrix *m, gsl_rng *r, int lower, int upper)
{
  size_t i, j;
  size_t N = m->size1;

  for (i = 0; i < N; ++i)
  {
    for (j = 0; j < N; ++j)
    {
      gsl_matrix_set(m,
                     i,
                     j,
                     gsl_rng_uniform(r) * (upper - lower) + lower);
    }
  }
} /* make_random_matrix() */

void
make_random_integer_matrix(gsl_matrix *m, gsl_rng *r, int lower, int upper)
{
  size_t i, j;
  size_t N = m->size1;

  for (i = 0; i < N; ++i)
  {
    for (j = 0; j < N; ++j)
    {
      double a = gsl_rng_uniform(r) * (upper - lower) + lower;
      gsl_matrix_set(m, i, j, floor(a));
    }
  }
} /* make_random_integer_matrix() */

void
make_start_matrix(gsl_matrix *m, int lower)

{
  size_t i, j;
  size_t N = m->size1;

  for (i = 0; i < N; ++i)
    for (j = 0; j < N; ++j)
      gsl_matrix_set(m, i, j, (double)lower);
} /* make_start_matrix() */

int
inc_matrix (gsl_matrix *m, int lower, int upper)
{
  size_t i = 0;
  size_t N = m->size1 * m->size2;
  int carry = 1;

  for (i = 0; carry > 0 && i < N; i++)
    {
      double v = m->data[i] + carry;
      carry = (v > upper) ? 1 : 0;
      m->data[i] = (v > upper) ? lower : v;
    }

  return carry;
} /* inc_matrix() */

void
output_matrix(gsl_matrix *m)
{
  size_t i, j;
  size_t N = m->size1;
  size_t M = m->size2;

  for (i = 0; i < N; ++i)
  {
    for (j = 0; j < M; ++j)
      {
        printf("%10.18e%s",
        /*printf("%10.18e%s",*/
               gsl_matrix_get(m, i, j),
               (j < M - 1) ? "," : ";\n");
      }
  }
}

void
print_matrix(gsl_matrix *m, const char *str)

{
  size_t i, j;
  size_t N = m->size1;
  size_t M = m->size2;
  gsl_matrix_view v;
  size_t rows, cols;
  size_t r, c;
  char buf[100];

  /*print_octave(m, str);
  return;*/

  /*rows = GSL_MIN(15, N);*/
  rows = N;
  cols = GSL_MIN(15, N);
  /*cols = N;*/

  for (i = 0; i < N; i += rows)
  {
    for (j = 0; j < M; j += cols)
    {
      r = GSL_MIN(rows, N - i);
      c = GSL_MIN(cols, N - j);

      v = gsl_matrix_submatrix(m, i, j, r, c);

      sprintf(buf, "%s(%u:%u,%u:%u)",
              str,
              i + 1,
              i + r,
              j + 1,
              j + c);

      printf("%s = [\n", buf);

      output_matrix(&v.matrix);

      printf("]\n");
    }
  }
} /* print_matrix() */

int
test_evals(gsl_vector_complex *obs, gsl_vector_complex *expected,
           gsl_matrix *A, gsl_matrix *B, const char *obsname,
           const char *expname)
{
  size_t N = expected->size;
  size_t i, k;
  double max, max_abserr, max_relerr;

  max = 0.0;
  max_abserr = 0.0;
  max_relerr = 0.0;
  k = 0;

  for (i = 0; i < N; ++i)
    {
      gsl_complex z = gsl_vector_complex_get(expected, i);
      max = GSL_MAX_DBL(max, gsl_complex_abs(z));
    }

  for (i = 0; i < N; ++i)
    {
      gsl_complex z_obs = gsl_vector_complex_get(obs, i);
      gsl_complex z_exp = gsl_vector_complex_get(expected, i);

      double x_obs = GSL_REAL(z_obs);
      double y_obs = GSL_IMAG(z_obs);
      double x_exp = GSL_REAL(z_exp);
      double y_exp = GSL_IMAG(z_exp);

      double abserr_x = fabs(x_obs - x_exp);
      double abserr_y = fabs(y_obs - y_exp);
      double noise = max * GSL_DBL_EPSILON * N * N;

      max_abserr = GSL_MAX_DBL(max_abserr, abserr_x + abserr_y);

      if (abserr_x < noise && abserr_y < noise)
        continue;

      if (abserr_x > 1.0e-6 || abserr_y > 1.0e-6)
        ++k;
    }

    if (k)
      {
        printf("==== CASE %lu ===========================\n\n", count);

        print_matrix(A, "A");
        print_matrix(B, "B");

        printf("=== eval - %s ===\n", expname);
        print_vector(expected, expname);

        printf("=== eval - %s ===\n", obsname);
        print_vector(obs, obsname);

        printf("max abserr = %g  max relerr = %g\n", max_abserr, max_relerr);

        printf("=========================================\n\n");
      }

    return k;
} /* test_evals() */

int
test_alpha(gsl_vector_complex *obs, gsl_vector_complex *expected,
           gsl_matrix *A, gsl_matrix *B, const char *obsname,
           const char *expname)
{
  size_t N = expected->size;
  size_t i, k;
  double max, max_abserr, max_relerr;

  max = 0.0;
  max_abserr = 0.0;
  max_relerr = 0.0;
  k = 0;

  for (i = 0; i < N; ++i)
    {
      gsl_complex z = gsl_vector_complex_get(expected, i);
      max = GSL_MAX_DBL(max, gsl_complex_abs(z));
    }

  for (i = 0; i < N; ++i)
    {
      gsl_complex z_obs = gsl_vector_complex_get(obs, i);
      gsl_complex z_exp = gsl_vector_complex_get(expected, i);

      double x_obs = GSL_REAL(z_obs);
      double y_obs = fabs(GSL_IMAG(z_obs));
      double x_exp = GSL_REAL(z_exp);
      double y_exp = fabs(GSL_IMAG(z_exp));

      double abserr_x = fabs(x_obs - x_exp);
      double abserr_y = fabs(y_obs - y_exp);
      double noise = max * GSL_DBL_EPSILON * N * N;

      max_abserr = GSL_MAX_DBL(max_abserr, abserr_x + abserr_y);

      if (abserr_x < noise && abserr_y < noise)
        continue;

      if (abserr_x > 1.0e-6 || abserr_y > 1.0e-6)
        ++k;
    }

    if (k)
      {
        printf("==== CASE %lu ===========================\n\n", count);

        print_matrix(A, "A");
        print_matrix(B, "B");

        printf("=== alpha - %s ===\n", expname);
        print_vector(expected, expname);

        printf("=== alpha - %s ===\n", obsname);
        print_vector(obs, obsname);

        printf("max abserr = %g  max relerr = %g\n", max_abserr, max_relerr);

        printf("=========================================\n\n");
      }

    return k;
} /* test_alpha() */

int
test_beta(gsl_vector *obs, gsl_vector *expected,
          gsl_matrix *A, gsl_matrix *B, const char *obsname,
          const char *expname)
{
  size_t N = expected->size;
  size_t i, k;
  double max, max_abserr, max_relerr;

  max = 0.0;
  max_abserr = 0.0;
  max_relerr = 0.0;
  k = 0;

  for (i = 0; i < N; ++i)
    {
      double z = gsl_vector_get(expected, i);
      max = GSL_MAX_DBL(max, fabs(z));
    }

  for (i = 0; i < N; ++i)
    {
      double v_obs = gsl_vector_get(obs, i);
      double v_exp = gsl_vector_get(expected, i);
      double abserr = fabs(v_obs - v_exp);
      double noise = max * GSL_DBL_EPSILON * N * N;

      max_abserr = GSL_MAX_DBL(max_abserr, abserr);

      if (abserr < noise)
        continue;

      if (abserr > 1.0e-6)
        ++k;
    }

    if (k)
      {
        printf("==== CASE %lu ===========================\n\n", count);

        print_matrix(A, "A");
        print_matrix(B, "B");

        printf("=== beta - %s ===\n", expname);
        printf("%s = [\n", expname);
        gsl_vector_fprintf(stdout, expected, "%.12e");
        printf("]\n");

        printf("=== beta - %s ===\n", obsname);
        printf("%s = [\n", obsname);
        gsl_vector_fprintf(stdout, obs, "%.12e");
        printf("]\n");

        printf("max abserr = %g  max relerr = %g\n", max_abserr, max_relerr);

        printf("=========================================\n\n");
      }

    return k;
} /* test_beta() */

/* test if A = Q S Z^t */
void
test_schur(gsl_matrix *A, gsl_matrix *S, gsl_matrix *Q, gsl_matrix *Z)
{
  const size_t N = A->size1;
  gsl_matrix *T1, *T2;
  size_t i, j, k;
  double lhs, rhs;
  double abserr;

  T1 = gsl_matrix_alloc(N, N);
  T2 = gsl_matrix_alloc(N, N);

  /* compute T1 = S Z^t */
  gsl_blas_dgemm(CblasNoTrans,
                 CblasTrans,
                 1.0,
                 S,
                 Z,
                 0.0,
                 T1);

  /* compute T2 = Q T1 = Q S Z^t */
  gsl_blas_dgemm(CblasNoTrans,
                 CblasNoTrans,
                 1.0,
                 Q,
                 T1,
                 0.0,
                 T2);

  k = 0;
  for (i = 0; i < N; ++i)
    {
      for (j = 0; j < N; ++j)
        {
          lhs = gsl_matrix_get(A, i, j);
          rhs = gsl_matrix_get(T2, i, j);

          abserr = fabs(lhs - rhs);

          if (abserr > 1.0e-6)
            ++k;
        }
    }

  if (k)
    {
      printf("==== CASE %lu ===========================\n\n", count);

      print_matrix(A, "A");

      printf("=== Schur Form matrix ===\n");
      print_matrix(S, "S");

      printf("=== Left Schur matrix ===\n");
      print_matrix(Q, "Q");

      printf("=== Right Schur matrix ===\n");
      print_matrix(Z, "Z");

      printf("=== Q S Z^t ===\n");
      print_matrix(T1, "Q S Z^t");

      printf("=== A - Q S Z^t ===\n");
      gsl_matrix_sub(T2, A);
      print_matrix(T1, "A - Q S Z^t");

      printf("=========================================\n\n");
    }

  gsl_matrix_free(T1);
  gsl_matrix_free(T2);
} /* test_schur() */

void
print_vector(gsl_vector_complex *eval, const char *str)
{
  size_t N = eval->size;
  size_t i;
  gsl_complex z;

  printf("%s = [\n", str);

  for (i = 0; i < N; ++i)
    {
      z = gsl_vector_complex_get(eval, i);
      printf("%.18e %.18e;\n", GSL_REAL(z), GSL_IMAG(z));
    }

  printf("]\n");
} /* print_vector() */

int
cmp(double a, double b)
{
  return ((a > b) ? 1 : ((a < b) ? -1 : 0));
} /* cmp() */

int
compare(const void *a, const void *b)
{
  const double *x = a;
  const double *y = b;
  int r1 = cmp(y[0], x[0]);
  int r2 = cmp(y[1], x[1]);

  if (!gsl_finite(x[0]))
    return 1;
  if (!gsl_finite(y[0]))
    return -1;

  if (fabs(x[0] - y[0]) < 1.0e-8)
    {
      /* real parts are very close to each other */
      return r2;
    }
  else
    {
      return r1 ? r1 : r2;
    }
} /* compare() */

void
sort_complex_vector(gsl_vector_complex *v)
{
  qsort(v->data, v->size, 2 * sizeof(double), &compare);
} /* sort_complex_vector() */

int
main(int argc, char *argv[])
{
  gen_workspace *gen_workspace_p;
  lapack_workspace *lapack_workspace_p;
  size_t N;
  int c;
  int lower;
  int upper;
  int incremental;
  size_t nmat;
  gsl_matrix *A, *B;
  gsl_rng *r;
  int s;
  int compute_schur;
  size_t i;

  gsl_ieee_env_setup();
  gsl_rng_env_setup();

  N = 30;
  lower = -10;
  upper = 10;
  incremental = 0;
  nmat = 0;
  compute_schur = 0;

  while ((c = getopt(argc, argv, "ic:n:l:u:z")) != (-1))
    {
      switch (c)
        {
          case 'i':
            incremental = 1;
            break;

          case 'n':
            N = strtol(optarg, NULL, 0);
            break;

          case 'l':
            lower = strtol(optarg, NULL, 0);
            break;

          case 'u':
            upper = strtol(optarg, NULL, 0);
            break;

          case 'c':
            nmat = strtoul(optarg, NULL, 0);
            break;

          case 'z':
            compute_schur = 1;
            break;

          case '?':
          default:
            printf("usage: %s [-i] [-z] [-n size] [-l lower-bound] [-u upper-bound] [-c num]\n", argv[0]);
            exit(1);
            break;
        } /* switch (c) */
    }

  A = gsl_matrix_alloc(N, N);
  B = gsl_matrix_alloc(N, N);
  gen_workspace_p = gen_alloc(N, compute_schur);
  lapack_workspace_p = lapack_alloc(N);

  r = gsl_rng_alloc(gsl_rng_default);

  if (incremental)
    {
      make_start_matrix(A, lower);

      /* we need B to be non-singular */
      make_random_integer_matrix(B, r, lower, upper);
    }

  fprintf(stderr, "testing N = %d", N);
  if (incremental)
    fprintf(stderr, " incrementally");
  else
    fprintf(stderr, " randomly");

  fprintf(stderr, " on element range [%d, %d]", lower, upper);

  if (compute_schur)
    fprintf(stderr, ", with Schur vectors");

  fprintf(stderr, "\n");

  while (1)
    {
      if (nmat && (count >= nmat))
        break;

      ++count;

      if (!incremental)
        {
          make_random_matrix(A, r, lower, upper);
          make_random_matrix(B, r, lower, upper);
        }
      else
        {
          s = inc_matrix(A, lower, upper);
          if (s)
            break; /* all done */

          make_random_integer_matrix(B, r, lower, upper);
        }

      /*if (count != 89120)
        continue;*/

      /* make copies of matrices */
      gsl_matrix_memcpy(gen_workspace_p->A, A);
      gsl_matrix_memcpy(gen_workspace_p->B, B);
      gsl_matrix_transpose_memcpy(lapack_workspace_p->A, A);
      gsl_matrix_transpose_memcpy(lapack_workspace_p->B, B);

      /* compute eigenvalues with LAPACK */
      s = lapack_proc(lapack_workspace_p);

      if (s != GSL_SUCCESS)
        {
          printf("LAPACK failed, case %lu\n", count);
          exit(1);
        }

#if 0
      print_matrix(A, "A");
      print_matrix(B, "B");
      gsl_matrix_transpose(lapack_workspace_p->A);
      gsl_matrix_transpose(lapack_workspace_p->B);
      print_matrix(lapack_workspace_p->A, "S_lapack");
      print_matrix(lapack_workspace_p->B, "T_lapack");
#endif

      /* compute eigenvalues with GSL */
      s = gen_proc(gen_workspace_p);

      if (s != GSL_SUCCESS)
        {
          printf("=========== CASE %lu ============\n", count);
          printf("Failed to converge: found %u eigenvalues\n",
                 gen_workspace_p->n_evals);
          print_matrix(A, "A");
          print_matrix(B, "B");
          print_matrix(gen_workspace_p->A, "Af");
          print_matrix(gen_workspace_p->B, "Bf");
          print_matrix(lapack_workspace_p->A, "Ae");
          print_matrix(lapack_workspace_p->B, "Be");
          exit(1);
        }

#if 0
      print_matrix(gen_workspace_p->A, "S_gsl");
      print_matrix(gen_workspace_p->B, "T_gsl");
#endif

      /* compute alpha / beta vectors */
      for (i = 0; i < N; ++i)
        {
          double beta;
          gsl_complex alpha, z;

          beta = gsl_vector_get(gen_workspace_p->beta, i);
          if (beta == 0.0)
            GSL_SET_COMPLEX(&z, GSL_POSINF, GSL_POSINF);
          else
            {
              alpha = gsl_vector_complex_get(gen_workspace_p->alpha, i);
              z = gsl_complex_div_real(alpha, beta);
            }

          gsl_vector_complex_set(gen_workspace_p->evals, i, z);

          beta = gsl_vector_get(lapack_workspace_p->beta, i);
          GSL_SET_COMPLEX(&alpha,
                          lapack_workspace_p->alphar[i],
                          lapack_workspace_p->alphai[i]);

          if (beta == 0.0)
            GSL_SET_COMPLEX(&z, GSL_POSINF, GSL_POSINF);
          else
            z = gsl_complex_div_real(alpha, beta);

          gsl_vector_complex_set(lapack_workspace_p->evals, i, z);
          gsl_vector_complex_set(lapack_workspace_p->alpha, i, alpha);
        }

#if 0
      gsl_sort_vector(gen_workspace_p->beta);
      gsl_sort_vector(lapack_workspace_p->beta);
      sort_complex_vector(gen_workspace_p->alpha);
      sort_complex_vector(lapack_workspace_p->alpha);

      s = test_alpha(gen_workspace_p->alpha,
                     lapack_workspace_p->alpha,
                     A,
                     B,
                     "gen",
                     "lapack");
      s = test_beta(gen_workspace_p->beta,
                    lapack_workspace_p->beta,
                    A,
                    B,
                    "gen",
                    "lapack");
#endif
#if 1
      sort_complex_vector(gen_workspace_p->evals);
      sort_complex_vector(lapack_workspace_p->evals);

      s = test_evals(gen_workspace_p->evals,
                     lapack_workspace_p->evals,
                     A,
                     B,
                     "gen",
                     "lapack");
#endif

      if (compute_schur)
        {
          test_schur(A,
                     gen_workspace_p->A,
                     gen_workspace_p->Q,
                     gen_workspace_p->Z);
          test_schur(B,
                     gen_workspace_p->B,
                     gen_workspace_p->Q,
                     gen_workspace_p->Z);
        }
    }

  gsl_matrix_free(A);
  gsl_matrix_free(B);
  gen_free(gen_workspace_p);
  lapack_free(lapack_workspace_p);

  if (r)
    gsl_rng_free(r);

  return 0;
} /* main() */


// Copyright (c) 2021 Stig Rune Sellevag
//
// This file is distributed under the MIT License. See the accompanying file
// LICENSE.txt or http://www.opensource.org/licenses/mit-license.php for terms
// and conditions.

#ifndef SCILIB_LINALG_MATRIX_DECOMPOSITION_H
#define SCILIB_LINALG_MATRIX_DECOMPOSITION_H

#ifdef USE_MKL
#include <mkl.h>
#else
#ifdef __clang__
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wreturn-type-c-linkage"
#endif
#include <lapacke.h>
#ifdef __clang__
#pragma clang diagnostic pop
#endif
#endif

#include <scilib/mdarray.h>
#include <scilib/linalg_impl/lapack_types.h>
#include <experimental/mdspan>
#include <exception>
#include <cassert>
#include <type_traits>

namespace Sci {
namespace Linalg {

namespace stdex = std::experimental;

// LU factorization.
template <class Layout>
inline void lu(Sci::Matrix_view<double, Layout> a,
               Sci::Vector_view<BLAS_INT, Layout> ipiv)
{
    static_assert(a.is_contiguous());
    static_assert(ipiv.is_contiguous());

    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));

    assert(static_cast<BLAS_INT>(ipiv.size()) >= std::min(m, n));

    auto matrix_layout = LAPACK_ROW_MAJOR;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = LAPACK_COL_MAJOR;
        lda = m;
    }

    BLAS_INT info =
        LAPACKE_dgetrf(matrix_layout, m, n, a.data(), lda, ipiv.data());
    if (info < 0) {
        throw std::runtime_error("dgetrf: illegal input parameter");
    }
    if (info > 0) {
        throw std::runtime_error("dgetrf: U matrix is singular");
    }
}

template <class Layout, class Allocator_a, class Allocator_ipiv>
inline void lu(Sci::Matrix<double, Layout, Allocator_a>& a,
               Sci::Vector<BLAS_INT, Layout, Allocator_ipiv>& ipiv)
{
    lu(a.view(), ipiv.view());
}

// QR factorization.
template <class Layout>
inline void qr(Sci::Matrix_view<double, Layout> a,
               Sci::Matrix_view<double, Layout> q,
               Sci::Matrix_view<double, Layout> r)
{
    assert(q.extent(0) == a.extent(0) && q.extent(1) == a.extent(1));
    assert(r.extent(0) == a.extent(0) && r.extent(1) == a.extent(1));

    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));

    auto matrix_layout = LAPACK_ROW_MAJOR;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = LAPACK_COL_MAJOR;
        lda = m;
    }
    Sci::copy(a, q);
    Sci::Vector<double, Layout> tau(std::min(m, n));

    // Compute QR factorization:

    BLAS_INT info =
        LAPACKE_dgeqrf(matrix_layout, m, n, q.data(), lda, tau.data());
    if (info != 0) {
        throw std::runtime_error("dgeqrf failed");
    }

    // Compute Q:

    info = LAPACKE_dorgqr(matrix_layout, m, n, n, q.data(), lda, tau.data());
    if (info != 0) {
        throw std::runtime_error("dorgqr failed");
    }

    // Compute R:

    matrix_product(transposed(q), a, r);
    transposed(q);
}

template <class Layout, class Allocator>
inline void qr(Sci::Matrix<double, Layout, Allocator>& a,
               Sci::Matrix<double, Layout, Allocator>& q,
               Sci::Matrix<double, Layout, Allocator>& r)
{
    qr(a.view(), q.view(), r.view());
}

// Singular value decomposition.
template <class Layout>
inline void svd(Sci::Matrix_view<double, Layout> a,
                Sci::Vector_view<double, Layout> s,
                Sci::Matrix_view<double, Layout> u,
                Sci::Matrix_view<double, Layout> vt)
{
    const BLAS_INT m = static_cast<BLAS_INT>(a.extent(0));
    const BLAS_INT n = static_cast<BLAS_INT>(a.extent(1));
    const BLAS_INT ldu = m;
    const BLAS_INT ldvt = n;

    assert(static_cast<BLAS_INT>(s.extent(0)) == std::min(m, n));
    assert(static_cast<BLAS_INT>(u.extent(0)) == m);
    assert(static_cast<BLAS_INT>(u.extent(1)) == ldu);
    assert(static_cast<BLAS_INT>(vt.extent(0)) == n);
    assert(static_cast<BLAS_INT>(vt.extent(1)) == ldvt);

    auto matrix_layout = LAPACK_ROW_MAJOR;
    BLAS_INT lda = n;

    if constexpr (std::is_same_v<Layout, stdex::layout_left>) {
        matrix_layout = LAPACK_COL_MAJOR;
        lda = m;
    }

    Sci::Vector<double, Layout> superb(std::min(m, n) - 1);

    BLAS_INT info =
        LAPACKE_dgesvd(matrix_layout, 'A', 'A', m, n, a.data(), lda, s.data(),
                       u.data(), ldu, vt.data(), ldvt, superb.data());
    if (info != 0) {
        throw std::runtime_error("dgesvd failed");
    }
}

template <class Layout, class Allocator>
inline void svd(Sci::Matrix<double, Layout, Allocator>& a,
                Sci::Vector<double, Layout, Allocator>& s,
                Sci::Matrix<double, Layout, Allocator>& u,
                Sci::Matrix<double, Layout, Allocator>& vt)
{
    svd(a.view(), s.view(), u.view(), vt.view());
}

} // namespace Linalg
} // namespace Sci

#endif // SCILIB_LINALG_MATRIX_DECOMPOSITION_H


#pragma once

#include <string>
#include <vector>

#include <boost/random.hpp>
#include <cblas.h>

namespace con {

  using std::vector;
  using std::cout;
  using std::cerr;
  using std::endl;
  using std::string;

  static boost::mt19937 rng(0);

  #define BUG(x) std::cout<<#x<<" = "<<(x)<<std::endl;

  typedef double Real;

  typedef vector<Real> Vec;

  Real sigmoid(const Real &z) {
    return 1.0 / (1.0 + std::exp(-z));
  }

  Real derivativeSigmoid(const Real &v) {
    return v * (1.0 - v);
  }

  Real sqr(const Real &x) {
    return x * x;
  }

  Real randomize(const Real &min, const Real &max) {
    boost::uniform_real<Real> dst(min, max);
    return dst(rng);
  }

  void randomizeVec(const Real &min, const Real &max, Vec *a) {
    for (auto it = a->begin(); it != a->end(); it++) {
      (*it) = randomize(min, max);
    }
  }

  void gaussianRng(const Real &mean, const Real &std, Vec *a) {
    boost::normal_distribution<Real> random_distribution(mean, std);
    boost::variate_generator<boost::mt19937, boost::normal_distribution<Real>>
      variate_generator(rng, random_distribution);

    for (auto it = a->begin(); it != a->end(); it++) {
      *it = variate_generator();
    }
  }

  void clear(Vec *a) {
    for (auto it = a->begin(); it != a->end(); it++) {
      (*it) = 0;
    }
  }

  void clear(vector<Vec> *a) {
    for (int i = 0; i < a->size(); i++) {
      clear(&a->at(i));
    }
  }

  void print(const Vec &a) {
    for (auto x : a) {
      cout << x << " ";
    }
    cout << endl;
  }

  void reshape(const int num, const int width, const int height, const int depth, vector<Vec> *a) {
    a->resize(num);
    for (int i = 0; i < num; i++) {
      a->at(i).resize(width * height * depth);
    }
  }

  int ceilDiv(const int &x, const int &y) {
    return (x + y - 1) / y;
  }

  void gemm(
      const CBLAS_TRANSPOSE &TransA, const CBLAS_TRANSPOSE &TransB,
      const int &M, const int &N, const int &K,
      const Real &alpha, const Vec &vecA, const Vec &vecB,
      const Real &beta, Vec *vecC) {

    const Real *A = &vecA[0];
    const Real *B = &vecB[0];
    Real *C = &vecC->at(0);

    int lda = (TransA == CblasNoTrans) ? K : M;
    int ldb = (TransB == CblasNoTrans) ? N : K;

    cblas_dgemm(CblasRowMajor, TransA, TransB, M, N, K, alpha, A, lda, B,
        ldb, beta, C, N);
  }

  void gemv(
      const CBLAS_TRANSPOSE &TransA,
      const int &M, const int &N,
      const Real &alpha, const Vec &matA, const Vec &vecX,
      const double beta, Vec *vecY) {

    const Real *A = &matA[0];
    const Real *x = &vecX[0];
    Real *y = &vecY->at(0);

    cblas_dgemv(CblasRowMajor, TransA, M, N, alpha, A, N, x, 1, beta, y, 1);
  }

  void vexp(const int &n, const Vec &input, Vec *output) {
    // vsExp(n, &input[0], &output->at(0));
    for (int i = 0; i < n; i++) {
      output->at(i) = exp(input[i]);
    }
  }

  void vdiv(const int &n, const Vec &a, const Vec &b, Vec *c) {
    // vsDiv(n, &a[0], &b[0], &c->at(0));
    for (int i = 0; i < n; i++) {
      c->at(i) = a[i] / b[i];
    }
  }

  void copy(const int &n, const Vec &input, Vec *output) {
    for (int i = 0; i < n; i++) {
      output->at(i) = input[i];
    }
  }

  void ones(const int &n, Vec *v) {
    v->resize(n);
    std::fill(v->begin(), v->end(), 1.0);
  }
}


#ifndef IFT_COMMON_H_
#define IFT_COMMON_H_

#include <omp.h>
#include <stdlib.h>
#include <stdio.h>
#if !defined(__APPLE__)
	#include <malloc.h>
#endif
#include <math.h>
#include <string.h>
#include <limits.h>
#include <float.h>
#include <sys/time.h>
#include <time.h>
#include <unistd.h>
#include <assert.h>
#include <cblas.h>
#include <dirent.h>
#include <regex.h>

/*
 * Common data types
 */

#define INFINITY_INT  INT_MAX
#define INFINITY_FLT  FLT_MAX
#define INFINITY_DBL  DBL_MAX
#define INFINITY_LDBL LDBL_MAX

typedef struct timeval timer;

typedef unsigned char  uchar;
typedef unsigned short ushort;
typedef unsigned int   uint;
typedef unsigned long long ullong;

typedef struct ift_band {
  float *val;
} iftBand;

typedef struct ift_vector {
  float x,y,z;
} iftVector, iftPoint;

typedef struct ift_voxel {
  int x,y,z;
} iftVoxel;

typedef struct ift_dcomplex
{
  double r;
  double i;
} iftComplex;

typedef struct file_list {
  char **filesRoutes;
  char **filesNames;
  int    n;
} fileList;

typedef struct ift_name_metric_pair {
  char *name;
  float metric;
} iftNameMetricPair;

/**
 * Common definitions
 */

#define IFT_RANDOM_SEED (unsigned int) 213344
#define MAXWEIGHT     4095.0
#define AXIS_X 0
#define AXIS_Y 1
#define AXIS_Z 2
#define PI          3.1415926536
#define INTERIOR    0
#define EXTERIOR    1
#define BOTH        2
#define WHITE       0
#define GRAY        1
#define BLACK       2
#define NIL        -1
#define INCREASING  1
#define DECREASING  0
#define Epsilon     1E-05

/**
 * Common operations
 */

#ifndef MAX
#define MAX(x,y) (((x) > (y))?(x):(y))
#endif

#ifndef MIN
#define MIN(x,y) (((x) < (y))?(x):(y))
#endif

#define ROUND(x) ((x < 0)?(int)(x-0.5):(int)(x+0.5))
#define SIGN(x) ((x >= 0)?1:-1)

/**
 * Common functions to allocate memory
 */

char     *iftAllocCharArray(int n);
uchar    *iftAllocUCharArray(int n);
short    *iftAllocShortArray(int n);
ushort   *iftAllocUShortArray(int n);
uint     *iftAllocUIntArray(int n);
ullong   *iftAllocULLongArray(int n);
int      *iftAllocIntArray(int n);
float    *iftAllocFloatArray(int n);
double   *iftAllocDoubleArray(int n);
iftComplex  *iftAllocComplexArray(int n);
long double *iftAllocLongDoubleArray(int n);

void      iftPrintFloatArray(float* v, int n);
/**
 * Error messages
 */

#define MSG1  "Cannot allocate memory space"
#define MSG2  "Cannot open file"

/**
 *  Error message msg is printed in function func and the program
 *  exits abnormally.
 */

void iftError(char *msg,char *func);

/**
 *  Warning message msg is printed in function func and the program
 *  continues.
 */

void iftWarning(char *msg,char *func);

/**
 *  The contents of a and b are interchanged.
 */

void iftSwitchValues(int *a, int *b);
void iftSSwitchValues(char *a, char *b, int size);
void iftFSwitchValues(float *a, float *b);
void iftDSwitchValues(double *a, double *b);
void iftSwitchVoxels(iftVoxel *u, iftVoxel *v);

/**
 * Returns a random integer number between low and high.
 */

int iftRandomInteger (int low, int high);

/*
 * Randomly selects nelems of the set [low, high]
 */
int *iftRandomIntegers (int low, int high, int nelems);

/**
 * Randomly selects a normal distributed (N(0,1)) float number
 */

float iftRandomNormalFloat(void);

/**
 * Returns the distance between P0 and the line from P1 to P2, whose
 * size is P1P2
 */

float iftVoxelLineDist2D(iftVoxel P0, iftVoxel P1, iftVoxel P2, float P1P2);

/**
 * Returns the position of P0 with respect to the line from P1 to
 * P2. Negative values indicate left side, 0 indicates on the line,
 * and positive values indicate right side.
 */

int iftVoxelLinePosition2D (iftVoxel P0, iftVoxel P1, iftVoxel P2);

/**
 * Returns initial time
 */

timer *iftTic(void);

/**
 * Returns final time
 */

timer *iftToc(void);

/**
 * Computes the difference in ms from the initial time to the final time
 */

float iftCompTime(timer *tic, timer *toc);

/**
 * Generates seed for rand(), used in iftRandomInteger.
 */

void iftRandomSeed(unsigned int);

/**
 * Returns the factorial of a number or NIL in case of overflow
 */

long double iftFactorial(int n);

/**
 * Returns the limit to avoid overflow in factorial computation
 */

int iftFactorialLimit(void);

float iftVoxelDistance(iftVoxel u, iftVoxel v);
int iftSquaredVoxelDistance(iftVoxel u, iftVoxel v);

float iftPointDistance(iftPoint u, iftPoint v);

int   iftVoxelSquareDistance(iftVoxel u, iftVoxel v);

float iftInnerProduct(iftVector a, iftVector b);

iftVector iftCrossProduct(iftVector a, iftVector b);

char iftCollinearPoints(iftPoint P1, iftPoint P2, iftPoint P3);

char iftCollinearVoxels(iftVoxel P1, iftVoxel P2, iftVoxel P3);

iftVector iftNormalizeVector(iftVector v);

void iftNormalizeFloatArray(float *array, int nelems);

float iftVectorMagnitude(iftVector v);

void iftRemoveCarriageReturn(char *token); /* useful to get rid of the
					      carriage return and the
					      line feed characteres
					      introduced by DOS
					      systems when reading
					      strings from ASCII
					      files */

void   iftWriteFloatArray(float *v, int size, char *filename);
float *iftReadFloatArray(char *filename, int *size);

/**
 * Evaluates the sigmoid function, with x = value.
 * Alfa controls the decay of the function.
 */

float iftSigmoidalValue(float value, float alfa);

void iftVerifyToken(FILE *fp, char *token, char *function);
void iftReadIntValue(FILE *fp, int *value, char *token, char *function);
void iftReadIntValues(FILE *fp, int **value, int nvalues, char *token, char *function);
void iftWriteIntValue(FILE *fp, int value, char *token);
void iftWriteIntValues(FILE *fp, int *value, int nvalues, char *token);
void iftReadFloatValue(FILE *fp, float *value, char *token, char *function);
void iftReadFloatValues(FILE *fp, float **value, int nvalues, char *token, char *function);
void iftWriteFloatValue(FILE *fp, float value, char *token);
void iftWriteFloatValues(FILE *fp, float *value, int nvalues, char *token);
void iftReadDoubleValue(FILE *fp, double *value, char *token, char *function);
void iftReadDoubleValues(FILE *fp, double **value, int nvalues, char *token, char *function);
void iftWriteDoubleValue(FILE *fp, double value, char *token);
void iftWriteDoubleValues(FILE *fp, double *value, int nvalues, char *token);
void iftSkipComments(FILE *fp);
char iftVoxelsAreEqual(iftVoxel u1, iftVoxel u2);
char iftPointsAreEqual(iftPoint u1, iftPoint u2);

/**
 * Common function to handle arrays
 */

void iftCopyFloatArray(float *array1, float *array2, int nelems);
void iftCopyIntArray(int *array1, int *array2, int nelems);
int *iftMergeIntArray(int *array1, int n1, int *array2, int n2, int *nelems);
int *iftIntArrayOfUniqueElemsTransform(int *array, int *n);
fileList             *iftCreateFileList(void);
void                 iftDestroyFileList(fileList **list);
fileList             *iftGetFiles(char *dirname, char *type);

/* These functions are currently used to communicate with numpy */
void iftWriteRawIntArray(char *filename, int *array, int n);
int* iftReadRawIntArray(char *filename, int n);

float iftMean(float *x, int n);
float iftVar(float *x, int n);
float iftCov(float *x, float *y, int n);

int iftAlmostZero(float x);

#endif


#include <jni.h>
#include <assert.h>
#include <lapacke.h>

/* Cholesky */

JNIEXPORT jint Java_JAMAJni_CholeskyDecomposition_dpotrf (JNIEnv *env, jclass klass, jint matrix_layout, jchar uplo, jint n, jdoubleArray a, jint lda){

    double *aElems;
    int info;

    aElems = (*env)-> GetDoubleArrayElements (env, a, NULL);

    assert(aElems);

    info = LAPACKE_dpotrf((int) matrix_layout, (char) uplo, (lapack_int) n, aElems, (lapack_int) lda);

    (*env)-> ReleaseDoubleArrayElements (env, a, aElems, 0);

    return info;

}

JNIEXPORT jint Java_JAMAJni_CholeskyDecomposition_dpotri(JNIEnv *env, jclass klass, jint matrix_layout, jchar uplo, jint n, jdoubleArray a, jint lda){

    double *aElems;
    int info;

    aElems = (*env)-> GetDoubleArrayElements (env, a, NULL);

    assert(aElems);

    info = LAPACKE_dpotri((int) matrix_layout, (char) uplo, (lapack_int) n, aElems, (lapack_int) lda);

    (*env)-> ReleaseDoubleArrayElements (env, a, aElems, 0);

    return info;

}

JNIEXPORT jint Java_JAMAJni_CholeskyDecomposition_dpotrs (JNIEnv *env, jclass klass, jint matrix_layout, jchar uplo, jint n, jint nrhs, jdoubleArray a, jint lda, jdoubleArray b, jint ldb){

    double *aElems, *bElems;
    int info;

    aElems = (*env)-> GetDoubleArrayElements (env, a, NULL);
    bElems = (*env)-> GetDoubleArrayElements (env, b, NULL);

    assert(aElems && bElems);

    info = LAPACKE_dpotrs ((int) matrix_layout, (char) uplo, (lapack_int) n, (lapack_int) nrhs, aElems, (lapack_int) lda, bElems, (lapack_int) ldb);

    (*env)-> ReleaseDoubleArrayElements (env, a, aElems, JNI_ABORT);
    (*env)-> ReleaseDoubleArrayElements (env, b, bElems, 0);

    return info;

}


#include "../include/matrix.h"

#ifndef MATREX_NO_BLAS

#include <cblas.h>

void
matrix_dot(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasNoTrans,
    MX_ROWS(first),
    MX_COLS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );
}

void
matrix_dot_and_add(
  const float alpha, const Matrix first, const Matrix second, const Matrix third, Matrix result
) {
  const uint64_t data_size = MX_ROWS(first) * MX_COLS(second) + 2;

  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasNoTrans,
    MX_ROWS(first),
    MX_COLS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );

  for(uint64_t index = 2; index < data_size; index += 1) {
    result[index] += third[index];
  }
}

void
matrix_dot_and_apply(
  const float alpha, const Matrix first, const Matrix second, const char *function_name, Matrix result
) {
  const math_func_ptr_t func = math_func_from_name(function_name);

  const uint64_t data_size = MX_ROWS(first) * MX_COLS(second) + 2;

  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasNoTrans,
    MX_ROWS(first),
    MX_COLS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );

  for(uint64_t index = 2; index < data_size; index += 1) {
    result[index] = func(result[index]);
  }
}

void
matrix_dot_nt(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  MX_SET_ROWS(result, MX_ROWS(first));
  MX_SET_COLS(result, MX_ROWS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasNoTrans,
    CblasTrans,
    MX_ROWS(first),
    MX_ROWS(second),
    MX_COLS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );
}

void
matrix_dot_tn(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  MX_SET_ROWS(result, MX_COLS(first));
  MX_SET_COLS(result, MX_COLS(second));

  cblas_sgemm(
    CblasRowMajor,
    CblasTrans,
    CblasNoTrans,
    MX_COLS(first),
    MX_COLS(second),
    MX_ROWS(first),
    alpha,
    first + 2,
    MX_COLS(first),
    second + 2,
    MX_COLS(second),
    0.0,
    result + 2,
    MX_COLS(result)
  );
}

#else

void
matrix_dot(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + k*MX_COLS(second) + c];
      result[elem_offset] *= alpha;
    }
}

void
matrix_dot_and_add(
  const float alpha, const Matrix first, const Matrix second, const Matrix third, Matrix result
) {
  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = third[elem_offset];
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + k*MX_COLS(second) + c];
      result[elem_offset] *= alpha;
    }
}

void
matrix_dot_and_apply(
  const float alpha, const Matrix first, const Matrix second, const char *function_name, Matrix result
) {
  const math_func_ptr_t func = math_func_from_name(function_name);

  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + k*MX_COLS(second) + c];
      result[elem_offset] = func(alpha * result[elem_offset]);
    }
}

void
matrix_dot_nt(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  const int64_t rows = MX_ROWS(first);
  const int64_t cols = MX_ROWS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_COLS(first); k++)
        result[elem_offset] += first[2 + r*MX_COLS(first) + k] * second[2 + c*MX_COLS(second) + k];
      result[elem_offset] *= alpha;
    }
}

void
matrix_dot_tn(const float alpha, const Matrix first, const Matrix second, Matrix result) {
  const int64_t rows = MX_COLS(first);
  const int64_t cols = MX_COLS(second);

  MX_SET_ROWS(result, rows);
  MX_SET_COLS(result, cols);

  for (int64_t r = 0; r < rows; r++)
    for (int64_t c = 0; c < cols; c++) {
      const int64_t elem_offset = 2 + r*cols + c;
      result[elem_offset] = 0.0;
      for (int64_t k = 0; k < MX_ROWS(first); k++)
        result[elem_offset] += first[2 + r + k*rows]*second[2 + c + k*cols];
      result[elem_offset] *= alpha;
    }
}

#endif


/*
 * This file is part of MXE. See LICENSE.md for licensing information.
 */

/*  taken from http://www.netlib.org/lapack/lapacke.html */

/*  Calling CGEQRF and CUNGQR to compute Q with workspace querying */

#include <stdio.h>
#include <stdlib.h>
#include <lapacke_utils.h>
#include <cblas.h>

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

   lapack_complex_float *a,*tau,*r,*work,one,zero,query;
   lapack_int info,m,n,lda,lwork;
   int i,j;
   float err;
   m = 10;   n = 5;   lda = m;
   one = lapack_make_complex_float(1.0,0.0);
   zero= lapack_make_complex_float(0.0,0.0);
   a = calloc(m*n,sizeof(lapack_complex_float));
   r = calloc(n*n,sizeof(lapack_complex_float));
   tau = calloc(m,sizeof(lapack_complex_float));
   for(j=0;j<n;j++)
      for(i=0;i<m;i++)
         a[i+j*m] = lapack_make_complex_float(i+1,j+1);
   info = LAPACKE_cgeqrf_work(LAPACK_COL_MAJOR,m,n,a,lda,tau,&query,-1);
   lwork = (lapack_int)query;
   info = LAPACKE_cungqr_work(LAPACK_COL_MAJOR,m,n,n,a,lda,tau,&query,-1);
   lwork = MAX(lwork,(lapack_int)query);
   work = calloc(lwork,sizeof(lapack_complex_float));
   info = LAPACKE_cgeqrf_work(LAPACK_COL_MAJOR,m,n,a,lda,tau,work,lwork);
   info = LAPACKE_cungqr_work(LAPACK_COL_MAJOR,m,n,n,a,lda,tau,work,lwork);
   for(j=0;j<n;j++)
      for(i=0;i<n;i++)
         r[i+j*n]=(i==j)?-one:zero;
   cblas_cgemm(CblasColMajor,CblasConjTrans,CblasNoTrans,
               n,n,m,&one,a,lda,a,lda,&one,r,n);
   err=0.0;
   for(i=0;i<n;i++)
      for(j=0;j<n;j++)
         err=MAX(err,cabs(r[i+j*n]));
   printf("error=%e\n",err);
   free(work);
   free(tau);
   free(r);
   free(a);
   return(info);
}


#include "allelefreq.h"
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <gsl/gsl_sf_psi.h>
#include <gsl/gsl_errno.h>

void P_update_simple(const uint8_t* G, const double* zetabeta, const double* zetagamma, const double* xi, const double* beta, const double* gamma, double* var_beta, double* var_gamma, long N, long L, long K)
{
    uint8_t genotype;
    long idx, n, l, k;
    double theta_beta_sum, theta_gamma_sum;
    double *var_beta_tmp, *var_gamma_tmp;

    var_beta_tmp = (double*) malloc(K * sizeof(double));
    var_gamma_tmp = (double*) malloc(K * sizeof(double));

    // loop over loci
    for (l=0; l<L; l++) {

        for (k=0; k<K; k++) {

            var_beta_tmp[k] = 0.0;
            var_gamma_tmp[k] = 0.0;
        }

        // loop over samples
        for (n=0; n<N; n++) {

            genotype = G[n*L+l];

            // missing data do not contribute
            if (genotype!=3) {

                // compute xi*zeta_{beta,gamma}
                theta_beta_sum = 0.0;
                theta_gamma_sum = 0.0;
                for (k=0; k<K; k++) {
                    theta_beta_sum += xi[n*K+k] * zetabeta[l*K+k];
                    theta_gamma_sum += xi[n*K+k] * zetagamma[l*K+k];
                }

                // increment var_{beta,gamma}_tmp
                for (k=0; k<K; k++) {
                    var_beta_tmp[k] += (double) genotype * xi[n*K+k] / theta_beta_sum;
                    var_gamma_tmp[k] += (double) (2-genotype) * xi[n*K+k] / theta_gamma_sum;
                }
            }
        }

        // compute var_{beta,gamma}
        for (k=0; k<K; k++) {
            idx = l*K+k;
            var_beta[idx] = beta[idx] + zetabeta[idx] * var_beta_tmp[k];
            var_gamma[idx] = gamma[idx] + zetagamma[idx] * var_gamma_tmp[k];
        }
    }

    free( var_beta_tmp );
    free( var_gamma_tmp );
}

void P_update_logistic(const double* Dvarbeta, const double* Dvargamma, const double* mu, const double* Lambda, double* var_beta, double* var_gamma, double mintol, long L, long K)
{
    /*
    `Dvarbeta` and `Dvargamma` are a function of the values of `var_beta` and `var_gamma`. This
    dependence, however, is explicit on a set of latent populations assignments which in-turn
    depend on the variational parameters estimated at the previous step. So, the variables
    `Dvarbeta` and `Dvargamma` do not have to be updated.
    */

    long l, k, idx, numvar, update;
    long iter = 0;
    double tol = 10.0;
    double tmptol;
    double beta, gamma;
    double A, pbetagamma, pbeta, pgamma, ppbeta, ppgamma;
    double A_1, B_1, C_1, A_2, B_2, C_2;

    /*
    Iterate until succesive estimates are sufficiently similar
    or the number of iterations exceeds 1000.
    */
    while (tol>mintol && iter<1000) {

        numvar = 0;
        tol = 0.0;

        // loop over loci
        for (l=0; l<L; l++) {

            // only update a locus, if all its variational parameters
            // satisfy positivity constraints.
            update = 1;
            for (k=0; k<K; k++) {
                if (var_beta[l*K+k]<=0 || var_gamma[l*K+k]<=0) {
                    update = 0;
                }
            }

            if (update==1) {

                // loop over populations
                for (k=0; k<K; k++) {

                    idx = l*K+k;
                    // compute pseudo-hyperparameters
                    pbetagamma = gsl_sf_psi_1(var_beta[idx]+var_gamma[idx]);
                    pbeta = gsl_sf_psi_1(var_beta[idx]);
                    pgamma = gsl_sf_psi_1(var_gamma[idx]);
                    ppbeta = gsl_sf_psi_n(2, var_beta[idx]);
                    ppgamma = gsl_sf_psi_n(2, var_gamma[idx]);

                    A_1 = pbeta-pbetagamma;
                    B_1 = -1.*pbetagamma;
                    A_2 = -1.*pbetagamma;
                    B_2 = pgamma-pbetagamma;
                    A = (gsl_sf_psi(var_beta[idx]) - gsl_sf_psi(var_gamma[idx]) - mu[l])*Lambda[k];
                    C_1 = -1.*A*pbeta - 0.5*Lambda[k]*ppbeta;
                    C_2 = A*pgamma - 0.5*Lambda[k]*ppgamma;

                    beta = (C_1*B_2-C_2*B_1)/(A_1*B_2-A_2*B_1);
                    gamma = (C_1*A_2-C_2*A_1)/(B_1*A_2-B_2*A_1);

                    // compute var_{beta,gamma}
                    tmptol = var_beta[idx];
                    var_beta[idx] = beta + Dvarbeta[idx];
                    tol += fabs(var_beta[idx]-tmptol);

                    tmptol = var_gamma[idx];
                    var_gamma[idx] = gamma + Dvargamma[idx];
                    tol += fabs(var_gamma[idx]-tmptol);

                    numvar += 1;

                }
            }
        }

        // compute convergence tolerance
        tol = 0.5*tol/numvar;
        iter += 1;
    }

    // printf("tol = %.8f, iter = %lu\n",tol,iter);
}


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdbool.h>
#include <math.h>
#include <time.h>
#include <fftw3.h>
#include <assert.h>
#include <complex.h>
#include <gsl/gsl_math.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_matrix.h>
#include <gsl/gsl_vector.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_linalg.h>
#include <gsl/gsl_permutation.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_eigen.h>

#include "gslCompute.h"
#include "MVNsampling.h"
#include "Input.h"
#include "Basis.h"
#include "SampSetUp.h"
#include "Output_t0.h"
#include "SampInit.h"
#include "SampZinter.h"
#include "SampTheta.h"
#include "SampPhi2Inv.h"
#include "SampLoading.h"
#include "SampLatent.h"
#include "SampMu.h"
#include "SampAlpha.h"
#include "SampBeta.h"
#include "SampGamma.h"
#include "Sampweights.h"
#include "SampErr_e.h"
#include "SampErr_zeta.h"
#include "SampErr_eps.h"
#include "testing_training.h"
#include "SampMean.h"
#include "SampOutput.h"

#include <R.h>
#include <Rinternals.h>
//#include <Rmath.h>

SEXP csblf(SEXP outputpath, SEXP seed, SEXP nburnin, SEXP niter){
    gsl_set_error_handler_off();
    //////////////////////////////////////////////////////////
    ///// GSL Random Number Generator Initialization /////
    //////////////////////////////////////////////////////////
    const gsl_rng_type * T;
    gsl_rng * r;
    gsl_rng_env_setup();
    T = gsl_rng_default;
    r = gsl_rng_alloc(T);
    gsl_rng_set(r, asInteger(seed));
    //////////////////////////////////////////////////////////
    ///// Input Data /////
    //////////////////////////////////////////////////////////
    // Input/Output data directory
    char *inpathx, *inpath, *outpath; //, *outputpath;
    inpathx = (char *)calloc(500,sizeof(char));
    inpath = (char *)calloc(500,sizeof(char));
    outpath = (char *)calloc(500,sizeof(char));
    char *outppath = CHAR(asChar(outputpath));

    char dataSource[20] = "RealData"; // data from simulation or read data source
    int sim = strcmp(dataSource, "Simulation") == 0 ? 1 : 0;
    strcat(inpathx, outppath);
    strcat(inpathx, "Data/");
    strcat(inpath, inpathx);
    strcat(outpath, outppath);
    strcat(outpath, "Result/");

    // Input data
    struct Inputdata data;
    data = input(inpathx, inpath, sim);
    int L = data.sizes[0]; // number of parcels; restricted to be 1
    int nobs = data.sizes[1]; // number of training observations
    int P = data.sizes[3]; // number of imaging predictors
    int nts = data.sizes[4]; // number of observations for test

    ///// Basis Function /////
    // Define basis functions
    struct BasisFunc BF;
    double bandwidth = 1.0/10.0;
    int dd = 6.0;
    BF = genBasis(L, outpath, data, bandwidth, dd);
    int M = BF.M;
    //////////////////////////////////////////////////////////
    ///// MCMC /////
    //////////////////////////////////////////////////////////
    // Number of latent factors
    int K ;
    K = strcmp(dataSource, "Simulation") == 0 ? 20 : 9;
    // Length of chain
    int iter = asInteger(niter); // total iterations
    int burnin = asInteger(nburnin); // iterations after burnin
    // Parameters
    struct Sampling PostSamp;
    PostSamp = setupSamp(M, nobs, nts, L, P, K, BF, data);
    // Initialization
    bool printInit = false;
    //set_initial2(L, nobs, nts, K, P, PostSamp, data, BF, r, outpath, printInit, inpath);
    set_initial(L, nobs, nts, K, P, PostSamp, data, BF, r, outpath, printInit);
    clock_t start, end;
    start = clock();
    int *singular = (int *)calloc(2, sizeof(int));
    int t;
    // Start MCMC
    Rprintf("************ Start MCMC ************\n");
    for(t=1; t<=iter; t++){
        if (t % 50 == 0) {
            Rprintf("**** t=%d *****\n", t);
        }
        // Post sampling of Zinter
        Zinter_samp(nobs, L, K, PostSamp, data, BF, r, singular);
        err2inv_u_samp(L, PostSamp, data, r);
        err2inv_e_samp(nobs, L, PostSamp, data, BF, r);
        // Post sampling of theta
        theta_samp(nobs, L, K, PostSamp, data, BF, r, singular);
        // Post sampling of Phi2Inv
        phi2inv_samp(K, nobs, L, PostSamp, data, BF, r);
        // Post sampling of Loadings and its hyperparameters
        load_samp(K, L, nobs, PostSamp, BF, r);
        // Update latent variables
        latent_samp(L, K, nobs, PostSamp, data, BF, r);
        // Update mu
        mu_samp(L, K, nobs, PostSamp, data, BF, r);
        // Update alpha
        alpha_samp(nobs, K, P, L, PostSamp, data, BF, r);
        // Update beta
        beta_samp_approx(nobs, L, K, PostSamp, data, BF, r);
        // Update Gamma
        gamma_samp(P, nobs, K, L, PostSamp, data, BF, r);
        // Update Error
        err2inv_zeta_samp(nobs, L, K, PostSamp, data, BF, r);
        err2inv_eps_samp(nobs, L, K, PostSamp, data, BF, r);
        ///////// Posterior estimations and predictions /////////
        if(t>(iter-burnin) && t%1==0){
            // Training set
            est_training(nobs, P, L, K, PostSamp, data, BF);
            // test
            if(nts>0){
                est_testing(nts, nobs, P, L, K, PostSamp, data, BF);
            }
            // For posterior mean
            samp_mean(L, K, nobs, nts, P, PostSamp, data, BF);
        }
        // Write posterior samplings
        if(t%50==0){
            output_samp(outpath, t, burnin, nobs, nts, M, K, L, P, singular, PostSamp, data, BF);
        }
        // Write posterior mean
        if(t==iter){
            output_mean(L, K, nobs, nts, P, outpath, PostSamp, data, BF, burnin);
        }
    }
    end = clock();
    double mcmc_time_used = ((double) (end - start)) / CLOCKS_PER_SEC;
    Rprintf("MCMC: %.2f secs\n", mcmc_time_used);
    ////////////////////////////////////////////////////////////
    ///// Release Memory /////
    ////////////////////////////////////////////////////////////
    Rprintf("Release Memory.\n");
    gsl_rng_free(r);
    free(outpath);
    free(inpath);
    free(inpathx);
    free(singular);
    // Input Data
    free(data.sizes);
    free(data.parcel_len);
    free(data.parcel_len_sum);
    free(data.axes);
    free(data.Z);
    free(data.X);
    int l;
    for(l=0; l<L; l++){
        free(data.Zl[l]);
        free(data.Xl[l]);
    }
    if(nts > 0){
        free(data.Z_test);
        free(data.X_test);
        for(l=0; l<L; l++){
            free(data.Zl_test[l]);
            free(data.Xl_test[l]);
            free(PostSamp.xb_test[l]);
        }
        free(data.Zl_test);
        free(data.Xl_test);
        free(PostSamp.xb_test);
    }
    // Basis Functions
    for(l=0; l<L; l++){
        free(BF.basis[l]);
        free(BF.kernel_loc[l]);
    }
    free(BF.basis);
    free(BF.kernel_loc);
    free(BF.Ml);
    // Post Sampling
    for(l=0; l<L; l++){
        free(PostSamp.basis2[l]);
        free(PostSamp.zb[l]);
        free(PostSamp.zb2[l]);
        free(PostSamp.xb[l]);
        free(PostSamp.rxb[l]);

        free(PostSamp.Zinter[l]);
        free(PostSamp.theta[l]);
        free(PostSamp.load[l]);
        free(PostSamp.loadstar[l]);
        free(PostSamp.latent[l]);
        free(PostSamp.latentstar[l]);
        free(PostSamp.latent2[l]);

        free(PostSamp.beta[l]);
        free(PostSamp.betastar[l]);
        free(PostSamp.alpha[l]);
        free(PostSamp.alphastar[l]);
        free(PostSamp.xba[l]);
        free(PostSamp.xbar[l]);
        free(PostSamp.mustar[l]);
        free(PostSamp.sumX[l]);
    }

    free(PostSamp.basis2);
    free(PostSamp.zb);
    free(PostSamp.zb2);
    free(PostSamp.xb);
    free(PostSamp.rxb);
    free(PostSamp.Zinter);
    free(PostSamp.theta);
    free(PostSamp.load);
    free(PostSamp.loadstar);
    free(PostSamp.latent);
    free(PostSamp.latentstar);
    free(PostSamp.latent2);
    free(PostSamp.Phi2Inv);
    free(PostSamp.beta);
    free(PostSamp.betastar);
    free(PostSamp.alpha);
    free(PostSamp.alphastar);
    free(PostSamp.xba);
    free(PostSamp.xbar);
    free(PostSamp.mustar);
    free(PostSamp.sumX);
    free(PostSamp.err2inv_e);
    free(PostSamp.err2inv_zeta);
    free(PostSamp.err2inv_eps);
    free(PostSamp.err2inv_u);

    for(l=0; l<L; l++){
        free(PostSamp.mean_Zinter[l]);
        free(PostSamp.mean_theta[l]);
        free(PostSamp.mean_load[l]);
        free(PostSamp.mean_loadstar[l]);
        free(PostSamp.mean_latent[l]);
        free(PostSamp.mean_latentstar[l]);
        free(PostSamp.mean_beta[l]);
        free(PostSamp.mean_betastar[l]);
        free(PostSamp.mean_alpha[l]);
        free(PostSamp.mean_alphastar[l]);
        free(PostSamp.mean_mu[l]);
    }
    free(PostSamp.mean_Zinter);
    free(PostSamp.mean_theta);
    free(PostSamp.mean_load);
    free(PostSamp.mean_loadstar);
    free(PostSamp.mean_latent);
    free(PostSamp.mean_latentstar);
    free(PostSamp.mean_Phi2Inv);
    free(PostSamp.mean_beta);
    free(PostSamp.mean_betastar);
    free(PostSamp.mean_alpha);
    free(PostSamp.mean_alphastar);
    free(PostSamp.mean_mu);
    free(PostSamp.mean_gamma);
    free(PostSamp.mean_gamma_update);
    free(PostSamp.mean_w);
    free(PostSamp.err2inv_e_mean);
    free(PostSamp.err2inv_zeta_mean);
    free(PostSamp.err2inv_eps_mean);
    free(PostSamp.err2inv_u_mean);

    if(nts > 0){
        for(l=0; l<L; l++){
            free(PostSamp.Zinter_test[l]);
            free(PostSamp.latent_test[l]);
            free(PostSamp.theta_test[l]);
            free(PostSamp.outcome_test[l]);
            free(PostSamp.Zinter_test_mean[l]);
            free(PostSamp.latent_test_mean[l]);
            free(PostSamp.theta_test_mean[l]);
            free(PostSamp.outcome_test_mean[l]);
            free(PostSamp.outcome_test_mean2[l]);
        }
        free(PostSamp.Zinter_test);
        free(PostSamp.latent_test);
        free(PostSamp.theta_test);
        free(PostSamp.outcome_test);
        free(PostSamp.Zinter_test_mean);
        free(PostSamp.latent_test_mean);
        free(PostSamp.theta_test_mean);
        free(PostSamp.outcome_test_mean);
        free(PostSamp.outcome_test_mean2);
    }

    for(l=0; l<L; l++){
        free(PostSamp.latent_train[l]);
        free(PostSamp.latent_train_err[l]);
        free(PostSamp.theta_train[l]);
        free(PostSamp.theta_train_err[l]);
        free(PostSamp.outcome_train[l]);
        free(PostSamp.latent_train_mean[l]);
        free(PostSamp.theta_train_mean[l]);
        free(PostSamp.outcome_train_mean[l]);
        free(PostSamp.outcome_train_mean2[l]);
    }
    free(PostSamp.latent_train);
    free(PostSamp.latent_train_err);
    free(PostSamp.theta_train);
    free(PostSamp.theta_train_err);
    free(PostSamp.outcome_train);
    free(PostSamp.latent_train_mean);
    free(PostSamp.theta_train_mean);
    free(PostSamp.outcome_train_mean);
    free(PostSamp.outcome_train_mean2);

    Rprintf("MCMC complete.\n");
    return R_NilValue;
}


/* specfunc/gsl_sf_dilog.h
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Gerard Jungman
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or (at
 * your option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 */

/* Author:  G. Jungman */

#ifndef __GSL_SF_DILOG_H__
#define __GSL_SF_DILOG_H__

#include <gsl/gsl_sf_result.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

/* Real part of DiLogarithm(x), for real argument.
 * In Lewin's notation, this is Li_2(x).
 *
 *   Li_2(x) = - Re[ Integrate[ Log[1-s] / s, {s, 0, x}] ]
 *
 * Note that Im[Li_2(x)] = { 0 for x <= 1, -Pi*log(x) for x > 1 }
 */
int     gsl_sf_dilog_e(const double x, gsl_sf_result * result);
double     gsl_sf_dilog(const double x);

/* DiLogarithm(z), for complex argument z = r Exp[i theta].
 */
int gsl_sf_complex_dilog_e(const double r, double theta, gsl_sf_result * result_re, gsl_sf_result * result_im);

__END_DECLS

#endif /* __GSL_SF_DILOG_H__ */


#include <stdio.h>
     #include <math.h>
     #include <gsl/gsl_sort.h>
     #include <gsl/gsl_wavelet.h>
     #include <gsl/gsl_math.h>

     int
     main (int argc, char **argv)
     {
       int i, j, k, n = 256, nc = 20;
       double *data = malloc (n * sizeof (double));
       double *abscoeff = malloc (n * sizeof (double));
       size_t *p = malloc (n * sizeof (size_t));

       FILE * f;
       gsl_wavelet *w;
       gsl_wavelet_workspace *work;

       w = gsl_wavelet_alloc (gsl_wavelet_daubechies, 4);
       work = gsl_wavelet_workspace_alloc (n);

       f = fopen (argv[1], "r");
       for (i = 0; i < n; i++)
         {
           fscanf (f, "%lg", &data[i]);
         }
       fclose (f);

       gsl_wavelet_transform_forward (w, data, 1, n, work);

       for (i = 0; i < n; i++)
         {
           abscoeff[i] = fabs (data[i]);
	   printf ("abscoeff[%d] = %g\n", i, abscoeff[i]);
         }

       printf ("(-1,0) = %g\n\n", abscoeff[0]);
       i = 1;
       for (j = 0; j < 8; j++)
	 {
	   for (k = 0; k < gsl_pow_int(2, j); k++)
	     {
	       printf ("(%d,%d) = %g\n", j, k, abscoeff[i]);
	       i++;
	     }
	   printf("\n");
	 }

       /* ************** */
       j = 0;
       for (i = 0; i < n; ++i)
	 j += abscoeff[i];
       for (i = 0; i < n; ++i)
	 abscoeff[i] /= j;
       /* ************** */

       gsl_sort_index (p, abscoeff, 1, n);

       for (i = 0; (i + nc) < n; i++)
	 {
	   printf ("p[%d] = %ld\n", i, p[i]);
	   data[p[i]] = 0;
	 }

       gsl_wavelet_transform_inverse (w, data, 1, n, work);

       for (i = 0; i < n; i++)
         {
           printf ("%g\n", data[i]);
         }

       gsl_wavelet_free (w);
       gsl_wavelet_workspace_free (work);

       free (data);
       free (abscoeff);
       free (p);
       return 0;
     }