#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <math.h>
#include <unistd.h>
#include <limits.h>
#include <sys/time.h>
#include <zlib.h>

#ifdef INTEL_COMPILER
#include "mkl.h"
#else
#include <cblas.h>
#endif

//#define FIBS_UNIT 1000
#define DEFAULT_NDIGITS 10
#define SZBUF 1024
#define SZBYTE 256
#define N_GENOTYPES 4
#define NA_GENO_CHAR (N_GENOTYPES-1)
#define DEFAULT_ROW_SIZE 100000
#define DEFAULT_SIZE_MATRIX 1000000
#define DEFAULT_SIZE_HEADER 100000
#define DEFAULT_DELIMS " \t\r\n"
#define SZ_LONG_BUF 1000000
#define DEFAULT_TPED_NUM_HEADER_COLS 4
#define DEFAULT_TFAM_NUM_HEADER_COLS 6
#define DEFAULT_TPED_SNPID_INDEX 1
#define DEFAULT_PHENO_NUM_HEADER_COLS 2

struct HFILE {
  int gzflag;       // 1 if gz if used
  int wflag;        // r(0)/w(1) for plain, rb(0)/wb(1) for gz
  int nheadercols;  // # of header columns (0 if nrows=0)
  int nvaluecols;   // # of value cols (0 if nrows=0)
  int nrows;        // # of rows
  FILE* fp;         // plain file handle
  gzFile gzfp;      // gzip file handle
};

// Input routines
void close_file (struct HFILE* fhp);
struct HFILE open_file(char* filename, int gzflag, int wflag);
struct HFILE open_file_with_suffix(char* prefix, char* suffix, int gzflag, int wflag);
//void read_matrix_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, int symmetric, int* p_nmiss, unsigned char** matrix, char*** headers);
void read_matrix_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, int* p_nmiss, unsigned char** matrix, char*** headers);
unsigned char* tokenize_tped_line_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, char* lbuf, unsigned char* values, char** headers, int* p_nvalues, int* p_nmiss );

void emmax_error( const char* format, ... );
void print_help(void);
FILE* readfile(char* filename);

void print_help(void) {
  fprintf(stderr,"Usage: emmax_kin [tpedf]\n");
  fprintf(stderr,"Required parameters\n");
  fprintf(stderr,"\t[tpedf]     : tped file\n");
  fprintf(stderr,"Optional parameters\n");
  fprintf(stderr,"\t-d [# digits]  : precision of the kinship values (default : 10)\n");
  fprintf(stderr,"\t-M [float] : maximum memory in GB (default: 4.0)\n");
  fprintf(stderr,"\t-s : compute IBS kinship matrix (default is Balding-Nicholas)\n");
  fprintf(stderr,"\t-v : turn on verbose mode\n");
  fprintf(stderr,"\t-r : randomly fill missing genotypes (default is imputation by average)\n");
  fprintf(stderr,"\t-x : include non-autosomal chromosomes in computing kinship matrices\n");
  fprintf(stderr,"\t-S [int] : set random seed\n");
  fprintf(stderr,"\t-m [float] : MAF threshold (default is 0)\n");
  fprintf(stderr,"\t-c [float] : Call rate threshold (default is 0)\n");
}

//void read_matrix_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, int symmetric, int* p_nmiss, unsigned char** matrix, char*** headers) {
void read_matrix_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, int* p_nmiss, unsigned char** matrix, char*** headers) {
  char* lbuf = (char*) malloc(sizeof(char*) * SZ_LONG_BUF);
  int szmat = DEFAULT_SIZE_MATRIX;
  int szheader = DEFAULT_SIZE_HEADER;
  unsigned char* cmat = (unsigned char*) malloc(sizeof(unsigned char) * szmat );
  char** cheaders = (char**) malloc(sizeof(char*) * szheader );
  int nvalues, i, j, nmiss;

  fhp->nheadercols = nheadercols;
  nmiss = 0;

  while( tokenize_tped_line_with_col_headers(fhp, nheadercols, delims, lbuf, &cmat[fhp->nrows*fhp->nvaluecols], &cheaders[fhp->nrows*fhp->nheadercols], &nvalues, &nmiss) != NULL ) {
    if ( fhp->nrows == 1 ) {
      fhp->nvaluecols = nvalues;
    }
    else if ( fhp->nvaluecols != nvalues ) {
      emmax_error("The column size %d do not match to %d at line %d\n",nvalues,fhp->nvaluecols,fhp->nrows);
    }

    if ( (fhp->nrows+1)*(fhp->nvaluecols) > szheader ) {
      szheader *= 2;
      fprintf(stderr,"Header size is doubled to %d\n",szheader);
      cheaders = (char**) realloc( cheaders, sizeof(char*) * szheader );
    }

    if ( (fhp->nrows+1)*(fhp->nvaluecols) > szheader ) {
      szmat *= 2;
      fprintf(stderr,"Matrix size is doubled to %d\n",szmat);
      cmat = (unsigned char*) realloc( cmat, sizeof(unsigned char) * szmat );
    }
  }
  free(lbuf);

  *p_nmiss = nmiss;

  unsigned char* fmat = (unsigned char*) malloc(sizeof(unsigned char)*fhp->nrows*fhp->nvaluecols);
  char** fheaders = (char**) malloc(sizeof(char*)*fhp->nrows*fhp->nheadercols);
  for(i=0; i < fhp->nrows; ++i) {
    for(j=0; j < fhp->nvaluecols; ++j) {
      fmat[i+j*fhp->nrows] = cmat[i*fhp->nvaluecols+j];
    }
    for(j=0; j < fhp->nheadercols; ++j) {
      fheaders[i+j*fhp->nrows] = cheaders[i*fhp->nheadercols+j];
    }
  }
  free(cmat);
  free(cheaders);

  if ( matrix != NULL ) {
    if ( *matrix != NULL ) {
      free(*matrix);
    }
    *matrix = fmat;
  }

  if ( headers != NULL ) {
    if ( *headers != NULL ) {
      free(*headers);
    }
    *headers = fheaders;
  }
}

unsigned char* tokenize_tped_line_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, char* lbuf, unsigned char* values, char** headers, int* p_nvalues, int* p_nmiss ) {
  int j;
  char *token;
  unsigned char ctoken;

  char *ret = (fhp->gzflag == 1) ? gzgets(fhp->gzfp, lbuf, SZ_LONG_BUF) : fgets( lbuf, SZ_LONG_BUF, fhp->fp );
  int nmiss = 0;

  if ( ret == NULL ) {
    return NULL;
  }

  if ( fhp->nheadercols != nheadercols ) {
    emmax_error("# of header columns mismatch (%d vs %d) at line %d",fhp->nheadercols,nheadercols,fhp->nrows);
  }

  //fprintf(stderr,"tokenize-line called %s\n",lbuf);

  token = strtok(lbuf, delims);
  for( j=0; token != NULL; ++j ) {
    if ( j < nheadercols ) {
      headers[j] = strdup(token);
    }
    // if zero_miss_flag is set, assume the genotypes are encoded 0,1,2
    // Additively encodes the two genotypes in the following way
    // when (j-nheadercols) is even, 0->MISSING, add 1->0, 2->1
    // when (j-nheadercols) is odd, check 0-0 consistency, and add 1->0, 2->1
    else {
      ctoken = (unsigned char)(token[0]-'0');

      if ( ctoken > 2 ) {
	fprintf(stderr,"Unrecognized token %s\n",token);
	abort();
      }

      if ( (j-nheadercols) % 2 == 0 ) {
	values[(j-nheadercols)/2] = ctoken;
      }
      else {
	if ( ( ctoken > 0 ) && ( values[(j-nheadercols)/2] == 0 ) ) {
	  fprintf(stderr,"Unmatched token pair 0 %s\n",token);
	  abort();
	}
	else if ( ( ctoken == 0 ) && ( values[(j-nheadercols)/2] > 0 ) ) {
	  fprintf(stderr,"Unmatched token pair - %d 0\n",(int)values[(j-nheadercols)/2]);
	  abort();
	}
	values[(j-nheadercols)/2] = (unsigned char)(values[(j-nheadercols)/2]+ctoken);
      }
    }
    token = strtok(NULL, delims);
  }
  //fprintf(stderr,"tokenize-line ended %d %d\n",j,nheadercols);
  if ( (j-nheadercols) % 2 != 0 ) {
    fprintf(stderr,"Number of value tokens are not even %d\n",j-nheadercols);
    abort();
  }

  *p_nvalues = (j-nheadercols)/2;
  *p_nmiss += nmiss;
  ++(fhp->nrows);

  if ( j < nheadercols ) {
    fprintf(stderr,"Number of header columns are %d, but only %d columns were observed\n", nheadercols, j);
    abort();
  }

  return values;
}

// open_file_with_suffix()
// - [prefix].[suffix] : file name to open
// - gzflag : gzip flag (use gzfp if gzflag=1, otherwise use fp)
// - wflag : write flag (1 if write mode otherwise read mode
struct HFILE open_file_with_suffix(char* prefix, char* suffix, int gzflag, int wflag) {
  char filename[SZBUF];
  sprintf(filename,"%s.%s",prefix,suffix);
  return open_file(filename,gzflag,wflag);
}

// open_file()
// - filename : file name to open
// - gzflag : gzip flag (use gzfp if gzflag=1, otherwise use fp)
// - wflag : write flag (1 if write mode otherwise read mode)
struct HFILE open_file(char* filename, int gzflag, int wflag) {
  struct HFILE fh;
  fh.gzflag = gzflag;
  fh.wflag = wflag;
  fh.nheadercols = 0;
  fh.nvaluecols = 0;
  fh.nrows = 0;
  if ( gzflag == 1 ) {
    char* mode = (wflag == 1) ? "wb" : "rb";
    fh.gzfp = gzopen(filename,mode);
    fh.fp = NULL;

    if ( fh.gzfp == NULL ) {
      emmax_error("Cannot open file %s for reading",filename);
    }
  }
  else {
    char* mode = (wflag == 1) ? "w" : "r";
    fh.gzfp = (gzFile) NULL;
    fh.fp = fopen(filename,mode);

    if ( fh.fp == NULL ) {
      emmax_error("Cannot open file %s for writing",filename);
    }
  }
  return fh;
}

void emmax_error( const char* format, ... ) {
  va_list args;
  fprintf(stderr, "ERROR: ");
  va_start (args, format);
  vfprintf(stderr, format, args);
  va_end (args);
  fprintf(stderr,"\n");
  abort();
}

void close_file(struct HFILE* fhp) {
  if ( fhp->gzflag == 1 ) {
    gzclose(fhp->gzfp);
    fhp->gzfp = NULL;
  }
  else {
    fclose(fhp->fp);
    fhp->fp = NULL;
  }
}

int main(int argc, char** argv) {
  int i, j, n, c, ac0, ac1, ac2, nmiss, nelems, nex, n_sum_nin, nin, nex_maf, nex_call_rate, nex_autosomal;
  int verbose, ndigits, tped_nheadercols, tfam_nheadercols, flag_autosomal, rand_fill_flag, ibs_flag, n_unit_lines;
  unsigned char *snprow;
  char *suffix, buf[SZBUF];
  double f, max_memory_GB, maf_thres, call_rate_thres, aaf, call_rate;
  //long *fibs_sums, *scores, mean_score;
  double *kin, *snpunit;
  char *tpedf, *delims, *lbuf;
  char **tfam_headers, **tped_headers;
  struct HFILE tpedh, tfamh, kinsh;
  struct timeval tv;

  // set default params
  gettimeofday(&tv, NULL);
  srand((unsigned int)tv.tv_usec);
  delims = DEFAULT_DELIMS;
  tped_nheadercols = DEFAULT_TPED_NUM_HEADER_COLS;
  tfam_nheadercols = DEFAULT_TFAM_NUM_HEADER_COLS;
  tped_headers = tfam_headers = NULL;
  tpedf = lbuf = 0;
  flag_autosomal = 1;
  rand_fill_flag = 0;
  ibs_flag = 0;
  verbose = 0;
  ndigits = DEFAULT_NDIGITS;
  max_memory_GB = 4.0; // 4.0GB
  maf_thres = 0.0;
  call_rate_thres = 0.0;

  // read arguments and update params
  while ((c = getopt(argc, argv, "d:rsc:vxS:M:m:c:")) != -1 ) {
    switch(c) {
    case 'd': // precision of digits
      ndigits = atoi(optarg);
      break;
    case 'r':
      rand_fill_flag = 1;
      break;
    case 's':
      ibs_flag = 1;
      break;
    case 'v':
      verbose = 1;
      break;
    case 'x':
      flag_autosomal = 0;
      break;
    case 'S':
      srand(atoi(optarg));
      break;
    case 'M':
      max_memory_GB = atof(optarg);
      break;
    case 'm':
      maf_thres = atof(optarg);
      break;
    case 'c':
      call_rate_thres = atof(optarg);
      break;
    default:
      fprintf(stderr,"Error : Unknown option unsigned character %c\n",c);
      abort();
    }
  }

  // Sanity check for the number of required parameters
  if ( argc != optind + 1 ) {
    print_help();
    abort();
  }

  // Read required parameters
  tpedf = argv[optind++];

  if ( verbose) fprintf(stderr,"\nReading TFAM file %s.tfam ....\n",tpedf);

  tfamh = open_file_with_suffix(tpedf, "tfam", 0, 0);
  //read_matrix_with_col_headers( &tfamh, tfam_nheadercols, delims, 0, &nmiss, NULL, &tfam_headers);
  read_matrix_with_col_headers( &tfamh, tfam_nheadercols, delims, &nmiss, NULL, &tfam_headers);
  n = tfamh.nrows; // n is # of individuals
  if ( verbose ) fprintf(stderr,"Identified %d individuals from TFAM file\n",n);

  // compute the # of lines to read together
  // we would need two n*n matrix, and n*m matrix
  // (n*n + n*m)*sizeof(double) < M*1e9
  // m < M*1e9/sizeof(double)/n - n
  n_unit_lines = (int)floor((max_memory_GB * 1.0e9 / sizeof(double) / n - n)/2)*2;
  n_sum_nin = 0;
#ifdef INTEL_COMPILER
  char cn = 'N', ct = 'T';
  double one = 1.;
#endif

  if ( verbose ) fprintf(stderr,"Setting # unit lines = %d to fit the memory requirement\n",n_unit_lines);

  snprow = (unsigned char*)malloc(sizeof(unsigned char)*n);
  tped_headers = (char**)malloc(sizeof(char*)*n);
  lbuf = (char*) malloc(sizeof(char*) * SZ_LONG_BUF);

  kin = (double*)calloc(n*n, sizeof(double));
  snpunit = (double*)malloc(n*n_unit_lines * sizeof(double));

  if ( verbose) fprintf(stderr,"Reading TPED file %s.tped ....\n",tpedf);

  tpedh = open_file_with_suffix(tpedf, "tped", 0, 0);
  tpedh.nheadercols = tped_nheadercols;

  nex_autosomal = nex_maf = nex_call_rate = 0;

  for ( i=0, nin=0, nex = 0; tokenize_tped_line_with_col_headers( &tpedh, tped_nheadercols, delims, lbuf, snprow, tped_headers, &nelems, &nmiss) != NULL; ++i) {
    if ( ( verbose ) && ( i % 10000 ) == 0 ) fprintf(stderr,"Reading %d SNPs\n",i);

    if ( ( flag_autosomal == 1 ) && ( ( atoi(tped_headers[0]) == 0 ) || ( atoi(tped_headers[0]) > 22 ) ) ) // if SNP is not in autosomal chromosomes
      //if ( ( flag_autosomal == 1 ) && ( atoi(tped_headers[0]) > 22 ) ) // if SNP is not in autosomal chromosomes
    {
      ++nex; // # excluded snps from the last unit
      ++nex_autosomal;
      continue;
    }

    if ( nelems != n ) {
      emmax_error("Number of values %d in line %d do not match to %d, the number of columns\n", nelems, tpedh.nvaluecols, n);
    }

  /*
    Perform rapid kinship generate IBS, BN, NCOR matrix
--------------------
IBS pairwise matrix
--------------------
i/j  0  1  2  3
0    NA NA NA NA
1    NA 2  1  0
2    NA 1  2  1
3    NA 0  1  2

In fact is what it does is
X = (m*n) genotype matrix (1,2,3 coded)
Xn = X-2
K = (t(Xn) %*% Xn)/(2*m)
------------------
* NA column may be just averaged or predicted based on r2 with previous SNP
  with a certain window size

---------------------
pairwise BN matrix
--------------------
X : (m*n) genotype matrix
Xn : (m*n) matrix each row standardized, missing assigned to 0
K = t(Xn) %*% Xn / L
--------------------
*/

    ac0 = ac1 = ac2 = 0;
    for(j=0; j < n; ++j) {
      if ( snprow[j] > 0 ) {
	if ( snprow[j] < 2 ) {
	  fprintf(stderr,"Invalid snprow[%d] value %d at line %d, individual %d\n",j,(int)snprow[j],i,j);
	}
	snprow[j] -= 2; // from 2,3,4 to 0,1,2 coding

	switch(snprow[j]) {
	case 0:
	  ++ac0;
	  break;
	case 1:
	  ++ac1;
	  break;
	case 2:
	  ++ac2;
	  break;
	default:
	  emmax_error("Unknown allele %s, converted to %d\n",buf,(int)snprow[j]);
	  break;
	}
      }
      else {
	snprow[j] = (unsigned char)NA_GENO_CHAR;
      }
    }

    call_rate = (double)(ac0+ac1+ac2)/(double)n;
    //fprintf(stderr,"CallRate = %lf\n", call_rate);
    if ( call_rate <= call_rate_thres ) {
      ++nex;
      ++nex_call_rate;
      continue;
    }
    aaf = (double)(ac1+2*ac2)/(double)(2*(ac0+ac1+ac2));
    if ( ( aaf <= maf_thres ) || ( 1.-aaf <= maf_thres ) ) {
      ++nex;
      ++nex_maf;
      continue;
    }

    if ( rand_fill_flag == 1 ) {
      for(j=0; j < n; ++j) {
	if ( snprow[j] == (unsigned char)NA_GENO_CHAR ) {
	  if ( (rand() / (double) RAND_MAX) > aaf ) {
	    if ( (rand() / (double) RAND_MAX) > aaf ) {
	      snprow[j] = (unsigned char)2;
	      //ac1 += 2;
	      ++ac2;
	    }
	    else {
	      snprow[j] = (unsigned char)1;
	      ++ac1;
	      //++ac0;
	      //++ac1;
	    }
	  }
	  else {
	    if ( (rand() / (double) RAND_MAX) > aaf ) {
	      snprow[j] = (unsigned char)1;
	      ++ac1;
	      //++ac0;
	      //++ac1;
	    }
	    else {
	      snprow[j] = (unsigned char)0;
	      ++ac0;
	      //ac0 += 2;
	    }
	  }
	}
      }
      aaf = (double)(ac1+2*ac2)/(double)(2*(ac0+ac1+ac2));
      //aaf = (double)ac1/(double)(ac0+ac1);
    }

    // copy current values to arrays
    // Xn = [-1,1] - two rows per SNP : IBS matrix
    // Xn ~ N(0,1) : BN matrix
    if ( ibs_flag == 1 ) {
      for(j=0; j < n; ++j) {
	if ( snprow[j] == (unsigned char)NA_GENO_CHAR ) {
	  snpunit[nin+  j*n_unit_lines] = 2.*aaf-1.;
	  snpunit[nin+1+j*n_unit_lines] = 2.*aaf-1.;
	}
	else {
	  if ( snprow[j] == 0 ) {
	    snpunit[nin+  j*n_unit_lines] = -1.;
	    snpunit[nin+1+j*n_unit_lines] = -1.;
	  }
	  else if ( snprow[j] == 1 ) {
	    snpunit[nin+  j*n_unit_lines] = 1.;
	    snpunit[nin+1+j*n_unit_lines] = -1.;
	  }
	  else if ( snprow[j] == 2 ) {
	    snpunit[nin+  j*n_unit_lines] = 1.;
	    snpunit[nin+1+j*n_unit_lines] = 1.;
	  }
	  else {
	    emmax_error("Invalid genotype %d\n",snprow[j]);
	  }
	}
      }
      nin += 2;
    }
    else {
      for(j=0; j < n; ++j) {
	if ( snprow[j] == (unsigned char)NA_GENO_CHAR ) {
	    snpunit[nin+j*n_unit_lines] = 0.;
	}
	else {
	  snpunit[nin+j*n_unit_lines] = ((double)snprow[j]-(aaf*2.))/sqrt(4*aaf*(1-aaf));
	}
      }
      ++nin;
    }

    // check if nin == n_unit_lines
    if ( nin >= n_unit_lines ) {
      if ( verbose ) {
	fprintf(stderr,"At SNP %d, intermediately computing kinship matrix with %d SNPs..\n",i,n_unit_lines);
      }
      // kin = kin + t(snpunit)%*%(snpunit)
#ifdef INTEL_COMPILER
      dgemm(&ct,&cn,&n,&n,&n_unit_lines,&one,snpunit,&n_unit_lines,snpunit,&n_unit_lines,&one,kin,&n);
#else
      cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans, n, n, n_unit_lines, 1.0, snpunit, n_unit_lines, snpunit, n_unit_lines, 1.0, kin, n);
#endif
      memset(snpunit, 0, sizeof(double)*n_unit_lines*n);

      n_sum_nin += nin;
      nin = 0;
    }
  }
  close_file(&tpedh);

  if ( verbose ) {
    fprintf(stderr,"Succesfully finished reading TPED file\n");
  }
  if ( nin > 0 ) {
    if ( verbose ) {
      fprintf(stderr,"Computing kinship matrix with the remaining %d SNPs..\n",nin);
    }
#ifdef INTEL_COMPILER
    dgemm(&ct,&cn,&n,&n,&n_unit_lines,&one,snpunit,&n_unit_lines,snpunit,&n_unit_lines,&one,kin,&n);
#else
    cblas_dgemm(CblasColMajor, CblasTrans, CblasNoTrans, n, n, n_unit_lines, 1.0, snpunit, n_unit_lines, snpunit, n_unit_lines, 1.0, kin, n);
#endif
    n_sum_nin += nin;
  }

  if ( ibs_flag == 1 ) {
    if ( rand_fill_flag == 1 ) {
      suffix = "rIBS.kinf";
    }
    else {
      suffix = "aIBS.kinf";
    }
  }
  else {
    if ( rand_fill_flag == 1 ) {
      suffix = "rBN.kinf";
    }
    else {
      suffix = "aBN.kinf";
    }
  }
  kinsh = open_file_with_suffix( tpedf, suffix, 0, 1 );

  if ( verbose ) fprintf(stderr,"Printing the kinship matrix to file %s.%s\n",tpedf,suffix);

  for(i=0; i < n; ++i) {
    for(j=0; j < n; ++j) {
      if ( j > 0 ) fprintf(kinsh.fp,"\t");
      f = (double)kin[i+j*n]/(double)(n_sum_nin);
      if ( ibs_flag == 1 ) {
	fprintf(kinsh.fp,"%-.*lf",ndigits,0.5*f+0.5);
      }
      else {
	fprintf(kinsh.fp,"%-.*lf",ndigits,f);
      }
    }
    fprintf(kinsh.fp,"\n");
  }
  close_file(&kinsh);
  free(snprow);
  free(kin);
  free(snpunit);
  free(lbuf);
  free(tped_headers);
  free(tfam_headers);
  return 0;
}


/*******************************************************************************  *
*
*       This file is part of the General Hidden Markov Model Library,
*       GHMM version __VERSION__, see http://ghmm.org
*
*       Filename: ghmm/ghmm/root_finder.c
*       Authors:  Achim Gaedke
*
*       Copyright (C) 1998-2004 Alexander Schliep
*       Copyright (C) 1998-2001 ZAIK/ZPR, Universitaet zu Koeln
*       Copyright (C) 2002-2004 Max-Planck-Institut fuer Molekulare Genetik,
*                               Berlin
*
*       Contact: schliep@ghmm.org
*
*       This library is free software; you can redistribute it and/or
*       modify it under the terms of the GNU Library General Public
*       License as published by the Free Software Foundation; either
*       version 2 of the License, or (at your option) any later version.
*
*       This library is distributed in the hope that it will be useful,
*       but WITHOUT ANY WARRANTY; without even the implied warranty of
*       MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
*       Library General Public License for more details.
*
*       You should have received a copy of the GNU Library General Public
*       License along with this library; if not, write to the Free
*       Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
*
*
*       This file is version $Revision: 2267 $
*                       from $Date: 2009-04-24 11:01:58 -0400 (Fri, 24 Apr 2009) $
*             last change by $Author: grunau $.
*
*******************************************************************************/

#ifdef HAVE_CONFIG_H
#  include "../config.h"
#endif  /*  */

#ifndef DO_WITH_GSL

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

#include "ghmm_internals.h"

double ghmm_zbrent_AB (double (*func) (double, double, double, double),
                   double x1, double x2, double tol, double A, double B,
                   double eps)
{
  fprintf (stderr, "Function ghmm_zbrent_AB() not implemented!\n");
  exit (1);

    /*
       double a, b, c;
       double fa, fb, fc;

       a = min(x1, x2);
       fa = (*func)(a);
       c = max(x1, x2);
       fc = (*func)(c);
       b = (c - a)/2.0;
       fb = (*func)(b);

       while (fabs(c - a) > tol + (tol * min(fabs(a), fabs(c))))
       {
       r = fb/fc;
       s = fb/fa;
       t = fa/fc;
       p = s * (t * (r - t) * (c - b) - (1.0 - r) * (b - a));
       q = (t - 1.0) * (r - 1.0) * (s - 1.0);

       x = b + p/q;

       if (x > a && x < c)
       {
       / * Accept interpolating point * /
       if (x < b)
       {
       c = b;
       fc = fb;
       b = x;
       fb = (*func)(b);
       }
       else if (x > b)
       {
       a = b;
       fa = fb;
       b = x;
       fb = (*func)(b);
       }
       }
       else
       {
       / * Use bisection * /
       }
       }
     */
}
#else   /*  */

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

/* struct for function pointer and parameters except 1st one */
  struct parameter_wrapper {
  double (*func) (double, double, double, double);
   double x2;
    double x3;
    double x4;
 } parameter;

/* calls the given function during gsl root solving iteration
   the first parameter variates, all other are kept constant
*/
double function_wrapper (double x, void *p)
{
  struct parameter_wrapper *param = (struct parameter_wrapper *) p;
  return param->func (x, param->x2, param->x3, param->x4);
}

/*
  this interface is used in sreestimate.c
 */
double ghmm_zbrent_AB (double (*func) (double, double, double, double), double x1,
                  double x2, double tol, double A, double B, double eps)
{

    /* initialisation of wrapper structure */
  struct parameter_wrapper param;
  gsl_function f;

#ifdef HAVE_GSL_INTERVAL /* gsl_interval vanished with version 0.9 */
    gsl_interval x;

#endif  /*  */
    gsl_root_fsolver * s;
  double tolerance;
  int success = 0;
  double result = 0;
  param.func = func;
  param.x2 = A;
  param.x3 = B;
  param.x4 = eps;
  f.function = &function_wrapper;
  f.params = (void *) &param;
  tolerance = tol;

    /* initialisation */
#ifdef HAVE_GSL_INTERVAL
# ifdef GSL_ROOT_FSLOVER_ALLOC_WITH_ONE_ARG
    x.lower = x1;
  x.upper = x2;
  s = gsl_root_fsolver_alloc (gsl_root_fsolver_brent);
  gsl_root_fsolver_set (s, &f, x);

# else
    s = gsl_root_fsolver_alloc (gsl_root_fsolver_brent, &f, x);

# endif
#else   /* gsl_interval vanished with version 0.9 */
    s = gsl_root_fsolver_alloc (gsl_root_fsolver_brent);
  gsl_root_fsolver_set (s, &f, x1, x2);

#endif  /*  */

    /* iteration */
    do {
    success = gsl_root_fsolver_iterate (s);
    if (success == GSL_SUCCESS)
       {

#ifdef HAVE_GSL_INTERVAL
        gsl_interval new_x;
      new_x = gsl_root_fsolver_interval (s);
      success = gsl_root_test_interval (new_x, tolerance, tolerance);

#else   /* gsl_interval vanished with version 0.9 */
      double x_up;
      double x_low;
      (void) gsl_root_fsolver_iterate (s);
      x_up = gsl_root_fsolver_x_upper (s);
      x_low = gsl_root_fsolver_x_lower (s);
      success = gsl_root_test_interval (x_low, x_up, tolerance, tolerance);

#endif  /*  */
      }
  } while (success == GSL_CONTINUE);

    /* result */
    if (success != GSL_SUCCESS)
     {
    gsl_error ("solver failed", __FILE__, __LINE__, success);
    }

  else
     {
    result = gsl_root_fsolver_root (s);
    }

    /* destruction */
    gsl_root_fsolver_free (s);
  return result;
}

#endif  /*  */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <iniparser.h>
#include "prepmt/prepmt_hudson96.h"
#ifdef PARMT_USE_INTEL
#include <mkl_cblas.h>
#else
#include <cblas.h>
#endif
#include "cps.h"
#include "iscl/array/array.h"
#include "iscl/fft/fft.h"
#include "iscl/memory/memory.h"
#include "iscl/os/os.h"

/*!
 * @brief Reads the ini file for Computer Programs in Seismology hudson96
 *        forward modeling variables.
 *
 * @param[in] iniFile    Name of ini file.
 * @param[in] section    Section of ini file to read.
 *
 * @param[out] parms     The hudson96 parameters.
 *
 * @result 0 indicates success.
 *
 * @author Ben Baker, ISTI
 *
 */
int prepmt_hudson96_readHudson96Parameters(const char *iniFile,
                                           const char *section,
                                           struct hudson96_parms_struct *parms)
{
    const char *s;
    char vname[256];
    dictionary *ini;
    //------------------------------------------------------------------------//
    cps_setHudson96Defaults(parms);
    if (!os_path_isfile(iniFile))
    {
        fprintf(stderr, "%s: ini file: %s does not exist\n", __func__, iniFile);
        return -1;
    }
    ini = iniparser_load(iniFile);
    if (ini == NULL)
    {
        fprintf(stderr, "%s: Cannot parse ini file\n", __func__);
        return -1;
    }
    // Teleseismic model
    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:modeltel", section);
    s = iniparser_getstring(ini, vname, NULL);
    if (s != NULL)
    {
        strcpy(parms->modeltel, s);
    }
    else
    {
        strcpy(parms->modeltel, "tak135sph.mod\0");
    }
    // Receiver model
    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:modelrec", section);
    s = iniparser_getstring(ini, vname, NULL);
    if (s != NULL)
    {
        strcpy(parms->modelrec, s);
    }
    else
    {
        strcpy(parms->modelrec, parms->modelrec);
    }
    // Source model
    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:modelsrc", section);
    s = iniparser_getstring(ini, vname, NULL);
    if (s != NULL)
    {
        strcpy(parms->modelsrc, s);
    }
    else
    {
        strcpy(parms->modelsrc, parms->modelsrc);
    }
    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:hs", section);
    parms->hs = iniparser_getdouble(ini, vname, 0.0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dt", section);
    parms->dt = iniparser_getdouble(ini, vname, 1.0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:npts", section);
    parms->npts = iniparser_getint(ini, vname, 1024);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:gcarc", section);
    parms->gcarc = iniparser_getdouble(ini, vname, 50.0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:offset", section);
    parms->offset = iniparser_getdouble(ini, vname, 10.0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dosrc", section);
    parms->dosrc = iniparser_getboolean(ini, vname, 1);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dorec", section);
    parms->dorec = iniparser_getboolean(ini, vname, 1);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dotel", section);
    parms->dotel = iniparser_getboolean(ini, vname, 1);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dop", section);
    parms->dop = iniparser_getboolean(ini, vname, 1);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dokjar", section);
    parms->dokjar = iniparser_getboolean(ini, vname, 1);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:loffsetdefault", section);
    parms->loffsetdefault = iniparser_getboolean(ini, vname, 1);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:verbose", section);
    parms->verbose = iniparser_getboolean(ini, vname, 0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:utstar", section);
    parms->utstar = iniparser_getdouble(ini, vname, -12345.0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:dottonly", section);
    parms->dottonly = iniparser_getboolean(ini, vname, 0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:zsrc", section);
    parms->zsrc = iniparser_getdouble(ini, vname, 100.0);

    memset(vname, 0, 256*sizeof(char));
    sprintf(vname, "%s:zrec", section);
    parms->zrec = iniparser_getdouble(ini, vname, 60.0);
    // Free ini dictionary
    iniparser_freedict(ini);
    return 0;
}
//============================================================================//
/*!
 * @brief Computes the fundamental fault Green's functions for the teleseismic
 *        body waves with hudson96.
 *
 * @param[in] nobs            Number of observations.
 * @param[in] obs             Observations with source location and
 *                            receiver location.  This is an array of length
 *                            nobs.
 * @param[in] luseCrust1      If true then use crust1.0 models at source
 *                            and receiver.
 * @param[in] crust1Dir       If luseCrust1 is true then this is the name
 *                            of the crust1.0 directory.  If NULL and
 *                            luseCrust1 is used then it will default to
 *                            libcps configuration.
 * @param[in] luseSrcModel    If true then use the local source velocity
 *                            at the source.  This supersedes the crust1.0
 *                            source velocity if it is defined.
 * @param[in] sourceModel     If luseSrcModel is true then this is the name
 *                            of the CPS style source model to read.
 *
 * @param[out] telmod         Holds the teleseismic model (ak135).
 * @param[out] srcmod         Holds the local source model.
 * @param[in,out] recmod      On input contains sufficient space and should
 *                            be an array of length nobs.
 *                            On output holds the receiver models.
 *
 * @param[out] ierr           0 indicates success.
 *
 * @author Ben Baker, ISTI
 *
 * @bug I need better handling of the offset.
 *
 */
int hudson96_getModels(const int nobs, const struct sacData_struct *obs,
                       const bool luseCrust1, const char *crust1Dir,
                       const bool luseSrcModel, const char *sourceModel,
                       struct vmodel_struct *telmod,
                       struct vmodel_struct *srcmod,
                       struct vmodel_struct *recmod)
{
    double *lats, *lons, evla, evlo;
    int ierr, iobs;
    ierr = 0;
    // Set the teleseismic model
    memset(srcmod, 0, sizeof(struct vmodel_struct));
    memset(telmod, 0, sizeof(struct vmodel_struct));
    cps_globalModel_ak135f(telmod);
    if (luseCrust1)
    {
        fprintf(stdout, "%s: Reading crust1.0...\n", __func__);
        lats = memory_calloc64f(nobs);
        lons = memory_calloc64f(nobs);
        ierr = sacio_getFloatHeader(SAC_FLOAT_EVLA, obs[0].header, &evla);
        if (ierr != 0)
        {
            fprintf(stderr, "%s: Error - evla not set\n", __func__);
            return ierr;
        }
        ierr = sacio_getFloatHeader(SAC_FLOAT_EVLO, obs[0].header, &evlo);
        if (ierr != 0)
        {
            fprintf(stderr, "%s: Error - evlo not set\n", __func__);
            return ierr;
        }
        for (iobs=0; iobs<nobs; iobs++)
        {
            sacio_getFloatHeader(SAC_FLOAT_STLA, obs[iobs].header,
                                 &lats[iobs]);
            sacio_getFloatHeader(SAC_FLOAT_STLO, obs[iobs].header,
                                 &lons[iobs]);
        }
        ierr = cps_crust1_getCrust1ForHerrmann(crust1Dir, false,
                                               evla, evlo,
                                               lats, lons,
                                               nobs,
                                               srcmod, recmod);
        if (ierr != 0)
        {
            fprintf(stderr, "%s: Failed to load crust1.0 model\n", __func__);
            return ierr;
        }
        memory_free64f(&lats);
        memory_free64f(&lons);
    }
    // Use teleseismic model for receiver model
    else
    {
        fprintf(stdout, "%s: Setting receiver models to teleseismic model\n",
                __func__);
        for (iobs=0; iobs<nobs; iobs++)
        {
            cps_utils_copyVmodelStruct(*telmod, &recmod[iobs]);
        }
    }
    // Use source model
    if (luseSrcModel)
    {
        fprintf(stdout, "%s: Using source model: %s\n", __func__, sourceModel);
        cps_utils_freeVmodelStruct(srcmod);
        ierr = cps_getmod(sourceModel, srcmod);
        if (ierr != 0)
        {
            fprintf(stderr, "%s: Error loading srcmod\n", __func__);
            return ierr;
        }
    }
    fprintf(stdout, "%s: Computing Green's functions...\n", __func__);
    // Otherwise source to teleseismic model
    if (!luseSrcModel && !luseCrust1)
    {
        fprintf(stdout, "%s: Setting source model to teleseismic model\n",
                __func__);
        cps_utils_copyVmodelStruct(*telmod, srcmod);
    }
    return 0;
}
//============================================================================//
/*!
 * @brief Computes the fundamental fault Green's functions for the teleseismic
 *        body waves with hudson96.
 *
 * @param[in] hudson96Parms   hudson96 forward modeling parameters.
 * @param[in] hpulse96Parms   hpulse96 forward modeling parameters.
 * @param[in] telmod          Holds the teleseismic model (ak135).
 * @param[in] srcmod          Holds the local source model.
 * @param[in] recmod          Holds each receiver model.
 *                            This is an array of dimension [nobs].
 * @param[in] ntstar          Number of t*'s in grid.
 * @param[in] tstars          Array of dimension [ntstar] with the attenuation
 *                            factors.
 * @param[in] ndepth          Number of source depths in grid.
 * @param[in] depths          Array of dimension [ndepth] with the source
 *                            depths specified in (km).
 * @param[in] nobs            Number of observations.
 * @param[in] obs             Observed waveforms which contain the receiver
 *                            locations, components of motion, first
 *                            arrival time, sampling period, and number of
 *                            points in waveform.
 *
 * @param[out] ierr           0 indicates success.
 *
 * @author Ben Baker, ISTI
 *
 * @bug I need better handling of the offset.
 *
 */
struct sacData_struct *prepmt_hudson96_computeGreensFF(
    const struct hudson96_parms_struct hudson96Parms,
    const struct hpulse96_parms_struct hpulse96Parms,
    const struct vmodel_struct telmod,
    const struct vmodel_struct srcmod,
    const struct vmodel_struct *recmod,
    const int ntstar, const double *__restrict__ tstars,
    const int ndepth, const double *__restrict__ depths,
    const int nobs, const struct sacData_struct *obs, int *ierr)
{
    struct hudson96_parms_struct hudson96ParmsWork;
    struct hpulse96_parms_struct hpulse96ParmsWork;
    struct hpulse96_data_struct zresp;
    struct sacData_struct *sacFFGrns;
    struct hwave_greens_struct *ffGrns;
    char sncl[64], phaseName[8];
    double *dptr, cmpaz, cmpinc, dt, gcarc, offset0,
           pickTime, xnorm;
    int i, idep, ierrAll, ierr1, ierr2, indx, iobs, iobs0, ip, it,
        factor, kndx, nloop, npts;
    bool lzero[10], lfound, lsh;
    enum isclError_enum isclError;
    const enum sacHeader_enum pickVars[11]
       = {SAC_FLOAT_A,
          SAC_FLOAT_T0, SAC_FLOAT_T1, SAC_FLOAT_T2, SAC_FLOAT_T3,
          SAC_FLOAT_T4, SAC_FLOAT_T5, SAC_FLOAT_T6, SAC_FLOAT_T7,
          SAC_FLOAT_T8, SAC_FLOAT_T9};
    const enum sacHeader_enum pickTypes[11]
       = {SAC_CHAR_KA,
          SAC_CHAR_KT0, SAC_CHAR_KT1, SAC_CHAR_KT2, SAC_CHAR_KT3,
          SAC_CHAR_KT4, SAC_CHAR_KT5, SAC_CHAR_KT6, SAC_CHAR_KT7,
          SAC_CHAR_KT8, SAC_CHAR_KT9};
    // Here's the idea: I want to input magnitudes with units of N-m so:
    // N-m -> Dyne-dm (1.e+7)
    // CPS internally scales from Dyne-cm to cm (1.e+20)
    // Finally, I want outputs proprotional to m (1.e-2)
    const double xmom = 1.0;     // no confusing `relative' magnitudes
    const double xcps = 1.e-20;  // convert dyne-cm mt to output cm
    const double cm2m = 1.e-2;   // cm to meters
    const double dcm2nm = 1.e+7; // magnitudes intended to be specified in
                                 // Dyne-cm but I work in N-m
    // Given a M0 in Newton-meters get a seismogram in meters
    const double xscal = xmom*xcps*cm2m*dcm2nm;
    const char *cfaults[10] = {"ZDS\0", "ZSS\0", "ZDD", "ZEX\0",
                               "RDS\0", "RSS\0", "RDD", "REX\0",
                               "TDS\0", "TSS\0"};
    const int npTypes = 11;
    const int nrDist = 1;
    const int nrDepths = 1;
    const int nsDepths = 1;
    //------------------------------------------------------------------------//
    //
    // Quick error checks
    *ierr = 0;
    sacFFGrns = NULL;
    if (nobs < 1 || obs == NULL ||
        ntstar < 1 || tstars == NULL ||
        ndepth < 1 || depths == NULL)
    {
        *ierr = 1;
        if (nobs < 1){fprintf(stderr, "%s: Error no observations\n", __func__);}
        if (obs == NULL){fprintf(stderr, "%s: Error obs is NULL\n", __func__);}
        if (ntstar < 1){fprintf(stderr, "%s: Error no t*'s\n", __func__);}
        if (tstars == NULL)
        {
            fprintf(stderr, "%s: Error tstars is NULL\n", __func__);
        }
        if (ndepth < 1){fprintf(stderr, "%s: Error no depths\n", __func__);}
        if (depths == NULL)
        {
            fprintf(stderr, "%s: Error depths is NULL\n", __func__);
        }
        return sacFFGrns;
    }
    // Set the modeling structures
    memset(&hudson96ParmsWork, 0, sizeof(struct hudson96_parms_struct));
    memset(&hpulse96ParmsWork, 0, sizeof(struct hpulse96_parms_struct));
    cps_setHudson96Defaults(&hudson96ParmsWork);
    cps_setHpulse96Defaults(&hpulse96ParmsWork);
    cps_utils_copyHudson96ParmsStruct(hudson96Parms, &hudson96ParmsWork);
    cps_utils_copyHpulse96ParmsStruct(hpulse96Parms, &hpulse96ParmsWork);
    offset0 = hudson96Parms.offset;
    // Set space for output
    nloop = nobs*ndepth*ntstar;
    sacFFGrns = (struct sacData_struct *)
                 calloc((size_t) (10*nloop), sizeof(struct sacData_struct));
    // Loop on the distances, depths, t*'s and compute greens functions
    iobs0 =-1;
    ierrAll = 0;
    for (indx=0; indx<nloop; indx++)
    {
        ffGrns = NULL;
        // Convert 3D grid into loop indices (i->n1 is fast; k->n3 is slow)
        *ierr = prepmt_hudson96_grd2ijk(indx,
                                        ntstar, ndepth, nobs,
                                        &it, &idep, &iobs);
        if (*ierr != 0)
        {
            fprintf(stderr, "%s: Failed to convert to grid\n", __func__);
            break;
        }
        // New observation (station) -> may need to update velocity model
        if (iobs != iobs0)
        {
            iobs0 = iobs;
        }
        *ierr = sacio_getFloatHeader(SAC_FLOAT_DELTA,
                                     obs[iobs].header, &dt);
        if (*ierr != 0)
        {
            fprintf(stderr, "%s: Could not get sampling period\n", __func__);
            break;
        }
        *ierr = sacio_getFloatHeader(SAC_FLOAT_GCARC,
                                     obs[iobs].header, &gcarc);
        if (*ierr != 0)
        {
            fprintf(stderr, "%s: Could not get gcarc\n", __func__);
            break;
        }
        *ierr = sacio_getIntegerHeader(SAC_INT_NPTS,
                                       obs[iobs].header, &npts);
        if (*ierr != 0)
        {
            fprintf(stderr, "%s: Could not get npts\n", __func__);
            break;
        }
        // Tie waveform modeling to first pick type
        lfound = false;
        hudson96ParmsWork.dop = true;
        for (ip=0; ip<npTypes; ip++)
        {
            ierr1 = sacio_getCharacterHeader(pickTypes[ip],
                                             obs[iobs].header, phaseName);
            ierr2 = sacio_getFloatHeader(pickVars[ip],
                                         obs[iobs].header, &pickTime);
            if (ierr1 == 0 && ierr2 == 0)
            {
                lfound = true;
                if (strncasecmp(phaseName, "P", 1) == 0)
                {
                    hudson96ParmsWork.dop = true;
                }
                else if (strncasecmp(phaseName, "S", 1) == 0)
                {
                    hudson96ParmsWork.dop = false;
                }
                else
                {
                    fprintf(stderr, "%s: Can't classify phase: %s\n",
                            __func__, phaseName);
                    lfound = false;
                }
                break;
            }
        }
        if (!lfound)
        {
            fprintf(stdout, "%s: No pick available - won't be able to align\n",
                    __func__);
            continue;
        }
        memset(&zresp, 0, sizeof(struct hpulse96_data_struct));
        hudson96ParmsWork.hs = depths[idep];
        hudson96ParmsWork.dt = dt;
        hudson96ParmsWork.gcarc = gcarc;
        hudson96ParmsWork.utstar = tstars[it];
        factor = 2;
        if (hudson96ParmsWork.dt < 0.1){factor = 4;}
        hudson96ParmsWork.npts = MAX(256,
                                     MIN(2048,
                                         factor*fft_nextpow2(MAX(1, npts), ierr)));
        hudson96ParmsWork.offset = fmin(pickTime, offset0); //TODO: fmin or fmax?
//printf("%s %f\n", phaseName, pickTime);
        // Try to pad this out a little
        hudson96ParmsWork.offset = fmax(hudson96ParmsWork.offset,
                                  (double) (hudson96ParmsWork.npts - 1)*dt*0.2);
//        printf("%e %d\n", hudson96ParmsWork.offset/dt, npts);
/*
        if (pickTime < hudson96ParmsWork.offset)
        {
            hudson96ParmsWork.offset = (double) (int) (pickTime/dt + 0.5)*dt;
        }
*/
        zresp = hudson96_interface(hudson96ParmsWork,
                                   telmod, recmod[iobs], srcmod, ierr);
        if (*ierr != 0)
        {
            fprintf(stderr, "%s: Error calling hudson96\n", __func__);
            ierrAll = ierrAll + 1;
            goto NEXT_OBS;
        }
        ffGrns = hpulse96_interface(nrDist, nrDepths, nsDepths,
                                    hpulse96ParmsWork,
                                    &zresp, ierr);
        if (*ierr != 0)
        {
            fprintf(stderr, "%s: Error calling hpulse96 interface\n", __func__);
            ierrAll = ierrAll + 1;
            goto NEXT_OBS;
        }
        // Set the fundamental fault Green's functions
        for (i=0; i<10; i++)
        {
            dptr = NULL;
            lsh = false;
            lzero[i] = false;
            if (i == 0)
            {
                dptr = ffGrns->zds;
            }
            else if (i == 1)
            {
                dptr = ffGrns->zss;
            }
            else if (i == 2)
            {
                dptr = ffGrns->zdd;
            }
            else if (i == 3)
            {
                dptr = ffGrns->zex;
            }
            else if (i == 4)
            {
                dptr = ffGrns->rds;
                cmpinc = 90.0;
            }
            else if (i == 5)
            {
                dptr = ffGrns->rss;
                cmpinc = 90.0;
            }
            else if (i == 6)
            {
                dptr = ffGrns->rdd;
                cmpinc = 90.0;
            }
            else if (i == 7)
            {
                dptr = ffGrns->rex;
                cmpinc = 90.0;
            }
            else if (i == 8)
            {
                dptr = ffGrns->tds;
                lsh = true;
                cmpaz = 90.0;
                cmpinc = 90.0;
            }
            else if (i == 9)
            {
                dptr = ffGrns->tss;
                lsh = true;
                cmpaz = 90.0;
                cmpinc = 90.0;
            }
            kndx = indx*10 + i;
            sacio_setDefaultHeader(&sacFFGrns[kndx].header);
            sacFFGrns[kndx].header.lhaveHeader = true;
            sacio_setIntegerHeader(SAC_INT_NPTS, ffGrns->npts,
                                   &sacFFGrns[kndx].header);
            sacio_setIntegerHeader(SAC_INT_NZYEAR, 1970,
                                   &sacFFGrns[kndx].header);
            sacio_setIntegerHeader(SAC_INT_NZJDAY, 1,
                                   &sacFFGrns[kndx].header);
            sacio_setIntegerHeader(SAC_INT_NZHOUR, 0,
                                   &sacFFGrns[kndx].header);
            sacio_setIntegerHeader(SAC_INT_NZMIN, 0,
                                   &sacFFGrns[kndx].header);
            sacio_setIntegerHeader(SAC_INT_NZSEC, 0,
                                   &sacFFGrns[kndx].header);
            sacio_setIntegerHeader(SAC_INT_NZMSEC, 0,
                                   &sacFFGrns[kndx].header);
            sacio_setBooleanHeader(SAC_BOOL_LCALDA, false,
                                   &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_DELTA, ffGrns->dt,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_GCARC, ffGrns->dist/111.195,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_DIST, ffGrns->dist,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_EVDP, depths[idep],
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_STEL, ffGrns->stelel,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_AZ,  0.0,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_BAZ, 180.0,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_O, -ffGrns->t0,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_B, ffGrns->t0,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_CMPAZ, cmpaz,
                                 &sacFFGrns[kndx].header);
            sacio_setFloatHeader(SAC_FLOAT_CMPINC, cmpinc,
                                 &sacFFGrns[kndx].header);
            if (hudson96ParmsWork.dop)
            {
                sacio_setFloatHeader(SAC_FLOAT_A, ffGrns->timep,
                                     &sacFFGrns[kndx].header);
                sacio_setCharacterHeader(SAC_CHAR_KA, "P",
                                         &sacFFGrns[kndx].header);
            }
            else
            {
                if (!lsh)
                {
                    sacio_setFloatHeader(SAC_FLOAT_A, ffGrns->timesv,
                                         &sacFFGrns[kndx].header);
                }
                else
                {
                    sacio_setFloatHeader(SAC_FLOAT_A, ffGrns->timesh,
                                         &sacFFGrns[kndx].header);
                }
                sacio_setCharacterHeader(SAC_CHAR_KA, "S",
                                         &sacFFGrns[kndx].header);
            }
            sacio_setCharacterHeader(SAC_CHAR_KO, "O\0",
                                     &sacFFGrns[kndx].header);
            sacio_setCharacterHeader(SAC_CHAR_KEVNM, "SYNTHETIC\0",
                                     &sacFFGrns[kndx].header);
            sacio_setCharacterHeader(SAC_CHAR_KCMPNM, cfaults[i],
                                     &sacFFGrns[kndx].header);
            sacFFGrns[kndx].npts = ffGrns->npts;
            sacFFGrns[kndx].data = sacio_malloc64f(ffGrns->npts);
            if (dptr != NULL)
            {
                array_copy64f_work(ffGrns->npts, dptr, sacFFGrns[kndx].data);
            }
            else
            {
                array_set64f_work(ffGrns->npts, 0.0, sacFFGrns[kndx].data);
            }
            xnorm = cblas_dnrm2(sacFFGrns[kndx].npts, sacFFGrns[kndx].data, 1);
            if (xnorm == 0.0){lzero[i] = true;}
            // Handle scaling
            cblas_dscal(sacFFGrns[kndx].npts, xscal, sacFFGrns[kndx].data, 1);
            dptr = NULL;
        } // Loop on fundamnetal faults
        if (array_sum8l(10, lzero, &isclError) == 10)
        {
            memset(sncl, 0, 64*sizeof(char));
            sprintf(sncl, "%s.%s.%s.%s",
                    obs[iobs].header.knetwk, obs[iobs].header.kstnm,
                    obs[iobs].header.kcmpnm, obs[iobs].header.khole);
            fprintf(stdout,
                    "%s: Warning Greens fn %s w/ npts %d at gcarc=%e is zero\n",
                    __func__, sncl, hudson96ParmsWork.npts, gcarc);
        }
NEXT_OBS:;
        if (ffGrns != NULL)
        {
            cps_utils_freeHwaveGreensStruct(ffGrns);
            free(ffGrns);
        }
        cps_utils_freeHpulse96DataStruct(&zresp);
    }
    return sacFFGrns;
}
//============================================================================//
/*!
 * @brief Converts a 1D index to 3D grid indices:
 *
 *        Changes:
 *        for (igrd=0; igrd<n1*n2*n3; igrd++)
 *        {
 *
 *        }
 *        To:
 *        for (k=0; k<n3; k++)
 *        {
 *            for (j=0; j<n2; j++)
 *            {
 *                for (i=0; i<n1; i++)
 *                {
 *                    igrd = k*n2*n1 + j*n1 + i;
 *                }
 *            }
 *        }
 *
 * @param[in] igrd  Flattened C numbered grid index: igrd = k*n2*n1 + j*n1 + i.
 * @param[in] n1    This is the fastest loop iterator (innermost loop).
 * @param[in] n2    This is the middle loop iterator (intermediate loop).
 * @param[in] n3    This is the slowest loop iterator (outermost loop).
 * @param[out] i    C index in inner loop [0,n1).
 * @param[out] j    C index in intermediate loop [0,n2).
 * @param[out] k    C index in outermost loop [0,n3).
 *
 * @result 0 indicates success.
 *
 */
int prepmt_hudson96_grd2ijk(const int igrd,
                            const int n1, const int n2, const int n3,
                            int *i, int *j, int *k)
{
    int ierr, n12;
    ierr = 0;
    n12 = n1*n2;
    *k = (igrd)/n12;
    *j = (igrd - *k*n12)/n1;
    *i =  igrd - *k*n12 - *j*n1;
    if (*i < 0 || *i > n1 - 1){ierr = ierr + 1;}
    if (*j < 0 || *j > n2 - 1){ierr = ierr + 1;}
    if (*k < 0 || *k > n3 - 1){ierr = ierr + 1;}
    return ierr;
}
//============================================================================//
/*!
 * @brief Maps the observation, depth, and t* to the global index.
 *
 * @param[in] ndepth    Number of depths.
 * @param[in] ntstar    Number of t*'s.
 * @param[in] iobs      C indexed observation number [0,nobs-1]
 * @param[in] idep      C indexed depth number [0,ndepth-1]
 * @param[in] it        C indexed t* number [0,ntstar-1]
 *
 * @result If >= 0 then this (iobs, idep, it) index in the Green's functions
 *         table.
 *
 * @author Ben Baker, ISTI
 *
 */
int prepmt_hudson96_observationDepthTstarToIndex(
    const int ndepth, const int ntstar,
    const int iobs, const int idep, const int it)
{
    int indx;
    indx = iobs*ntstar*ndepth + idep*ntstar + it;
    return indx;
}


/* monte/gsl_monte_vegas.h
 *
 * Copyright (C) 1996, 1997, 1998, 1999, 2000 Michael Booth
 *
 * 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

/* header for the gsl "vegas" routines.  Mike Booth, May 1998 */

#ifndef __GSL_MONTE_VEGAS_H__
#define __GSL_MONTE_VEGAS_H__

#include <gsl/gsl_rng.h>
#include <gsl/gsl_monte.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

enum {GSL_VEGAS_MODE_IMPORTANCE = 1,
      GSL_VEGAS_MODE_IMPORTANCE_ONLY = 0,
      GSL_VEGAS_MODE_STRATIFIED = -1};

typedef struct {
  /* grid */
  size_t dim;
  size_t bins_max;
  unsigned int bins;
  unsigned int boxes; /* these are both counted along the axes */
  double * xi;
  double * xin;
  double * delx;
  double * weight;
  double vol;

  double * x;
  int * bin;
  int * box;

  /* distribution */
  double * d;

  /* control variables */
  double alpha;
  int mode;
  int verbose;
  unsigned int iterations;
  int stage;

  /* scratch variables preserved between calls to vegas1/2/3  */
  double jac;
  double wtd_int_sum;
  double sum_wgts;
  double chi_sum;
  double chisq;

  double result;
  double sigma;

  unsigned int it_start;
  unsigned int it_num;
  unsigned int samples;
  unsigned int calls_per_box;

  FILE * ostream;

} gsl_monte_vegas_state;

int gsl_monte_vegas_integrate(gsl_monte_function * f,
                              double xl[], double xu[],
                              size_t dim, size_t calls,
                              gsl_rng * r,
                              gsl_monte_vegas_state *state,
                              double* result, double* abserr);

gsl_monte_vegas_state* gsl_monte_vegas_alloc(size_t dim);

int gsl_monte_vegas_init(gsl_monte_vegas_state* state);

void gsl_monte_vegas_free (gsl_monte_vegas_state* state);

__END_DECLS

#endif /* __GSL_MONTE_VEGAS_H__ */


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

#include "infmcmc.h"

#define nj1 32
#define nk1 32

void infmcmc_initChain(INFCHAIN *C, const int nj, const int nk) {
  const int maxk      = (nk >> 1) + 1;
  const int sspectral = sizeof(fftw_complex) * nj * maxk;
  const int sphysical = sizeof(double      ) * nj * nk;
  //const int obsVecMem = sizeof(double      ) * sizeObsVector;
  FILE *fp;
  unsigned long int seed;

  // Set up variables
  C->nj = nj;
  C->nk = nk;
  //C->sizeObsVector = sizeObsVector;
  C->currentIter = 0;
  C->accepted = 0;
  C->_shortTimeAccProbAvg = 0.0;
  C->_bLow = 0.0;
  C->_bHigh = 1.0;

  // Allocate a ton of memory
  C->currentPhysicalState      = (double *)malloc(sphysical);
  C->avgPhysicalState          = (double *)malloc(sphysical);
  C->varPhysicalState          = (double *)malloc(sphysical);
  C->proposedPhysicalState     = (double *)malloc(sphysical);
  C->_M2                       = (double *)malloc(sphysical);
  //C->currentStateObservations  = (double *)malloc(obsVecMem);
  //C->proposedStateObservations = (double *)malloc(obsVecMem);
  //C->data                      = (double *)malloc(obsVecMem);

  C->currentSpectralState  = (fftw_complex *)fftw_malloc(sspectral);
  C->avgSpectralState      = (fftw_complex *)fftw_malloc(sspectral);
  C->priorDraw             = (fftw_complex *)fftw_malloc(sspectral);
  C->proposedSpectralState = (fftw_complex *)fftw_malloc(sspectral);

  memset(C->currentPhysicalState,  0, sphysical);
  memset(C->avgPhysicalState,      0, sphysical);
  memset(C->varPhysicalState,      0, sphysical);
  memset(C->proposedPhysicalState, 0, sphysical);
  memset(C->_M2,                   0, sphysical);
  memset(C->currentSpectralState,  0, sspectral);
  memset(C->avgSpectralState,      0, sspectral);
  memset(C->proposedSpectralState, 0, sspectral);

  C->accProb = 0.0;
  C->avgAccProb = 0.0;
  C->logLHDCurrentState = 0.0;

  /*
   * Set some default values
   */
  C->alphaPrior = 3.0;
  C->rwmhStepSize = 1e-4;
  C->priorVar = 1.0;
  C->priorStd = 1.0;

  C->r = gsl_rng_alloc(gsl_rng_taus2);

  fp = fopen("/dev/urandom", "rb");

  if (fp != NULL) {
    fread(&seed, sizeof(unsigned long int), 1, fp);
    gsl_rng_set(C->r, seed);
    fclose(fp);
    printf("Using random seed\n");
  }
  else {
    gsl_rng_set(C->r, 0);
    printf("Using zero seed\n");
  }

  C->_c2r = fftw_plan_dft_c2r_2d(nj, nk, C->proposedSpectralState, C->proposedPhysicalState, FFTW_MEASURE);
  C->_r2c = fftw_plan_dft_r2c_2d(nj, nk, C->currentPhysicalState, C->currentSpectralState, FFTW_MEASURE);
}

void infmcmc_freeChain(INFCHAIN *C) {
  // Free all allocated memory used by the chain
  free(C->currentPhysicalState);
  free(C->avgPhysicalState);
  free(C->varPhysicalState);
  free(C->proposedPhysicalState);
  free(C->_M2);
  //free(C->currentStateObservations);
  //free(C->proposedStateObservations);

  fftw_free(C->currentSpectralState);
  fftw_free(C->avgSpectralState);
  fftw_free(C->priorDraw);
  fftw_free(C->proposedSpectralState);

  gsl_rng_free(C->r);
}

void infmcmc_resetChain(INFCHAIN *C) {
  infmcmc_freeChain(C);
  infmcmc_initChain(C, C->nj, C->nk);
}

void infmcmc_writeChain(const INFCHAIN *C, FILE *fp) {
  const int s = C->nj * C->nk;

  fwrite(&(C->nj),                 sizeof(int),    1,         fp);
  fwrite(&(C->nk),                 sizeof(int),    1,         fp);
  fwrite(&(C->currentIter),        sizeof(int),    1,         fp);
  fwrite(C->currentPhysicalState,  sizeof(double), s,         fp);
  fwrite(C->avgPhysicalState,      sizeof(double), s,         fp);
  fwrite(C->varPhysicalState,      sizeof(double), s,         fp);
  fwrite(&(C->logLHDCurrentState), sizeof(double), 1,         fp);
  fwrite(&(C->accProb),            sizeof(double), 1,         fp);
  fwrite(&(C->avgAccProb),         sizeof(double), 1,         fp);
}

void infmcmc_writeChainInfo(const INFCHAIN *C, FILE *fp) {
  fwrite(&(C->nj),             sizeof(int), 1, fp);
  fwrite(&(C->nk),             sizeof(int), 1, fp);
}

void infmcmc_writeVFChain(const INFCHAIN *U, const INFCHAIN *V, FILE *fp) {
  const int s = U->nj * U->nk;

  fwrite(U->currentPhysicalState,  sizeof(double), s,         fp);
  fwrite(U->avgPhysicalState,      sizeof(double), s,         fp);
  fwrite(U->varPhysicalState,      sizeof(double), s,         fp);
  fwrite(V->currentPhysicalState,  sizeof(double), s,         fp);
  fwrite(V->avgPhysicalState,      sizeof(double), s,         fp);
  fwrite(V->varPhysicalState,      sizeof(double), s,         fp);
  fwrite(&(U->logLHDCurrentState), sizeof(double), 1,         fp);
  fwrite(&(U->accProb),            sizeof(double), 1,         fp);
  fwrite(&(U->avgAccProb),         sizeof(double), 1,         fp);
}

void infmcmc_printChain(INFCHAIN *C) {
  printf("Iteration %d\n", C->currentIter);
  printf("-- Length is         %d x %d\n", C->nj, C->nk);
  printf("-- llhd val is       %lf\n", C->logLHDCurrentState);
  printf("-- Acc. prob is      %.10lf\n", C->accProb);
  printf("-- Avg. acc. prob is %.10lf\n", C->avgAccProb);
  printf("-- Beta is           %.10lf\n\n", C->rwmhStepSize);
  //finmcmc_printCurrentState(C);
  //finmcmc_printAvgState(C);
  //finmcmc_printVarState(C);
}

void randomPriorDraw(INFCHAIN *C) {
  int j, k;
  const int maxk = (C->nk >> 1) + 1;
  const int nko2 = C->nk >> 1;
  const int njo2 = C->nj >> 1;
  double xrand, yrand, c;
  //const double one = 1.0;

  c = 4.0 * M_PI * M_PI;

  for(j = 0; j < C->nj; j++) {
    for(k = 0; k < maxk; k++) {
      xrand = gsl_ran_gaussian_ziggurat(C->r, C->priorStd);
      if((j == 0) && (k == 0)) {
        C->priorDraw[0] = 0.0;
      }
      else if((j == njo2) && (k == nko2)) {
        C->priorDraw[maxk*njo2+nko2] = /*C->nj */ xrand / pow(c * ((j * j) + (k * k)), (double)C->alphaPrior/2.0);
      }
      else if((j == 0) && (k == nko2)) {
        C->priorDraw[nko2] = /*C->nj */ xrand / pow(c * ((j * j) + (k * k)), (double)C->alphaPrior/2.0);
      }
      else if((j == njo2) && (k == 0)) {
        C->priorDraw[maxk*njo2] = /*C->nj */ xrand / pow(c * ((j * j) + (k * k)), (double)C->alphaPrior/2.0);
      }
      else {
        xrand /= sqrt(2.0);
        yrand = gsl_ran_gaussian_ziggurat(C->r, C->priorStd) / sqrt(2.0);
        C->priorDraw[maxk*j+k] = /*C->nj */ (xrand + I * yrand) / pow(c * ((j * j) + (k * k)), (double)C->alphaPrior/2.0);
        if(j > njo2) {
          C->priorDraw[maxk*j+k] = conj(C->priorDraw[maxk*(C->nj-j)+k]);
        }
      }
    }
  }
}

void randomDivFreePriorDraw(INFCHAIN *C1, INFCHAIN *C2) {
  int j, k;
  const int maxk = (C1->nk >> 1) + 1;
  const int njo2 = C2->nj >> 1;
  const int nko2 = C1->nk >> 1;
  double modk;

  randomPriorDraw(C1);

  for (j = 0; j < C1->nj; j++) {
    for (k = 0; k < maxk; k++) {
      if (j == 0 && k == 0) {
        C1->priorDraw[0] = 0.0;
        C2->priorDraw[0] = 0.0;
        continue;
      }

      modk = sqrt(j * j + k * k);
      if (j < njo2) {
        C2->priorDraw[maxk*j+k] = -j * C1->priorDraw[maxk*j+k] / modk;
      }
      else if (j > njo2) {
        C2->priorDraw[maxk*j+k] = -(j - C1->nj) * C1->priorDraw[maxk*j+k] / modk;
      }
      else {
        C2->priorDraw[maxk*j+k] = 0.0;
      }

      if (k < nko2) {
        C1->priorDraw[maxk*j+k] *= k / modk;
      }
      else {
        C1->priorDraw[maxk*j+k] = 0.0;
      }
    }
  }
}

void infmcmc_seedWithPriorDraw(INFCHAIN *C) {
  const int size = sizeof(fftw_complex) * C->nj * ((C->nk >> 1) + 1);
  fftw_complex *uk = (fftw_complex *)fftw_malloc(size);
  const fftw_plan p = fftw_plan_dft_c2r_2d(C->nj, C->nk, uk, C->currentPhysicalState, FFTW_ESTIMATE);

  randomPriorDraw(C);
  memcpy(uk, C->priorDraw, size);

  //fixme: put into current spectral state

  fftw_execute(p);
  fftw_destroy_plan(p);
  fftw_free(uk);
}

void infmcmc_seedWithDivFreePriorDraw(INFCHAIN *C1, INFCHAIN *C2) {
  const int size = sizeof(fftw_complex) * C1->nj * ((C1->nk >> 1) + 1);
  fftw_complex *uk = (fftw_complex *)fftw_malloc(size);
  const fftw_plan p = fftw_plan_dft_c2r_2d(C1->nj, C1->nk, uk, C1->currentPhysicalState, FFTW_ESTIMATE);

  randomDivFreePriorDraw(C1, C2);

  memcpy(uk, C1->priorDraw, size);
  fftw_execute_dft_c2r(p, uk, C1->currentPhysicalState);

  memcpy(uk, C2->priorDraw, size);
  fftw_execute_dft_c2r(p, uk, C2->currentPhysicalState);

  memcpy(C1->currentSpectralState, C1->priorDraw, size);
  memcpy(C2->currentSpectralState, C2->priorDraw, size);

  fftw_destroy_plan(p);
  fftw_free(uk);
}

void infmcmc_proposeRWMH(INFCHAIN *C) {
  int j, k;
  const int maxk = (C->nk >> 1) + 1;
  const int N = C->nj * C->nk;
  const double sqrtOneMinusBeta2 = sqrt(1.0 - C->rwmhStepSize * C->rwmhStepSize);
  double *u = (double *)malloc(sizeof(double) * C->nj * C->nk);
  fftw_complex *uk = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * C->nj * maxk);

  memcpy(u, C->currentPhysicalState, sizeof(double) * C->nj * C->nk);
  fftw_execute_dft_r2c(C->_r2c, u, C->currentSpectralState);

  // Draw from prior distribution
  randomPriorDraw(C);

  for(j = 0; j < C->nj; j++) {
    for(k = 0; k < maxk; k++) {
      C->proposedSpectralState[maxk*j+k] = sqrtOneMinusBeta2 * C->currentSpectralState[maxk*j+k] + C->rwmhStepSize * C->priorDraw[maxk*j+k];
      C->proposedSpectralState[maxk*j+k] /= N;
    }
  }

  memcpy(uk, C->proposedSpectralState, sizeof(fftw_complex) * C->nj * maxk);
  fftw_execute_dft_c2r(C->_c2r, uk, C->proposedPhysicalState);
  fftw_free(uk);
  free(u);
}

void infmcmc_adaptRWMHStepSize(INFCHAIN *C, double inc) {
  // Adapt to stay in 20-30% range.
  int adaptFreq = 100;
  double rate;

  if (C->currentIter > 0 && C->currentIter % adaptFreq == 0) {
    rate = (double) C->_shortTimeAccProbAvg / adaptFreq;

    if (rate < 0.2) {
      //C->_bHigh = C->rwmhStepSize;
      //C->rwmhStepSize = (C->_bLow + C->_bHigh) / 2.0;
      C->rwmhStepSize -= inc;
    }
    else if (rate > 0.3) {
      //C->_bLow = C->rwmhStepSize;
      //C->rwmhStepSize = (C->_bLow + C->_bHigh) / 2.0;
      C->rwmhStepSize += inc;
    }

    C->_shortTimeAccProbAvg = 0.0;
  }
  else {
    C->_shortTimeAccProbAvg += C->accProb;
  }
}

void infmcmc_proposeDivFreeRWMH(INFCHAIN *C1, INFCHAIN *C2) {
  int j, k;
  const int maxk = (C1->nk >> 1) + 1;
  const int N = C1->nj * C1->nk;
  const double sqrtOneMinusBeta2 = sqrt(1.0 - C1->rwmhStepSize * C1->rwmhStepSize);
  double *u = (double *)malloc(sizeof(double) * C1->nj * C1->nk);
  fftw_complex *uk = (fftw_complex *)fftw_malloc(sizeof(fftw_complex) * C1->nj * maxk);

  memcpy(u, C1->currentPhysicalState, sizeof(double) * C1->nj * C1->nk);
  fftw_execute_dft_r2c(C1->_r2c, u, C1->currentSpectralState);

  memcpy(u, C2->currentPhysicalState, sizeof(double) * C2->nj * C2->nk);
  fftw_execute_dft_r2c(C2->_r2c, u, C2->currentSpectralState);

  // Draw from prior distribution
  randomDivFreePriorDraw(C1, C2);

  for(j = 0; j < C1->nj; j++) {
    for(k = 0; k < maxk; k++) {
      C1->currentSpectralState[maxk*j+k] /= N;
      C2->currentSpectralState[maxk*j+k] /= N;
      C1->proposedSpectralState[maxk*j+k] = sqrtOneMinusBeta2 * C1->currentSpectralState[maxk*j+k] + C1->rwmhStepSize * C1->priorDraw[maxk*j+k];
      //C1->proposedSpectralState[maxk*j+k] /= N;
      C2->proposedSpectralState[maxk*j+k] = sqrtOneMinusBeta2 * C2->currentSpectralState[maxk*j+k] + C2->rwmhStepSize * C2->priorDraw[maxk*j+k];
      //C2->proposedSpectralState[maxk*j+k] /= N;
    }
  }

  memcpy(uk, C1->proposedSpectralState, sizeof(fftw_complex) * C1->nj * maxk);
  fftw_execute_dft_c2r(C1->_c2r, uk, C1->proposedPhysicalState);
  memcpy(uk, C2->proposedSpectralState, sizeof(fftw_complex) * C1->nj * maxk);
  fftw_execute_dft_c2r(C2->_c2r, uk, C2->proposedPhysicalState);

  fftw_free(uk);
  free(u);
}

void infmcmc_updateAvgs(INFCHAIN *C) {
  int j, k;
  double deltar;
  fftw_complex deltac;
  C->currentIter++;

  /*
   Update physical vectors
   */
  if (C->currentIter == 1) {
    for (j = 0; j < C->nj; j++) {
      for (k = 0; k < C->nk; k++) {
        deltar = C->currentPhysicalState[C->nk * j + k] - C->avgPhysicalState[C->nk * j + k];
        C->avgPhysicalState[C->nk * j + k] += (deltar / C->currentIter);
        C->_M2[C->nk * j + k] += (deltar * (C->currentPhysicalState[C->nk * j + k] - C->avgPhysicalState[C->nk * j + k]));
        C->varPhysicalState[C->nk * j + k] = -1.0;
      }
    }
  }
  else {
    for (j = 0; j < C->nj; j++) {
      for (k = 0; k < C->nk; k++) {
        deltar = C->currentPhysicalState[C->nk * j + k] - C->avgPhysicalState[C->nk * j + k];
        C->avgPhysicalState[C->nk * j + k] += (deltar / C->currentIter);
        C->_M2[C->nk * j + k] += (deltar * (C->currentPhysicalState[C->nk * j + k] - C->avgPhysicalState[C->nk * j + k]));
        C->varPhysicalState[C->nk * j + k] = C->_M2[C->nk * j + k] / (C->currentIter - 1);
      }
    }
  }

  /*
   Update spectral vectors
  */
  if (C->currentIter == 1) {
    for (j = 0; j < C->nj; j++) {
      for (k = 0; k < C->nk/2 + 1; k++) {
        deltac = C->currentSpectralState[(C->nk/2 + 1) * j + k] - C->avgSpectralState[(C->nk/2 + 1) * j + k];
        C->avgSpectralState[(C->nk/2 + 1) * j + k] += (deltac / C->currentIter);
      }
    }
  }
  else {
    for (j = 0; j < C->nj; j++) {
      for (k = 0; k < C->nk/2 + 1; k++) {
        deltac = C->currentSpectralState[(C->nk/2 + 1) * j + k] - C->avgSpectralState[(C->nk/2 + 1) * j + k];
        C->avgSpectralState[(C->nk/2 + 1) * j + k] += (deltac / C->currentIter);
      }
    }
  }

  /*
   Update scalars
   */
  C->avgAccProb += ((C->accProb - C->avgAccProb) / C->currentIter);
}

void infmcmc_updateRWMH(INFCHAIN *C, double logLHDOfProposal) {
  double alpha;

  alpha = exp(C->logLHDCurrentState - logLHDOfProposal);

  if (alpha > 1.0) {
    alpha = 1.0;
  }
  // fixme: set acc prob here instead
  if (gsl_rng_uniform(C->r) < alpha) {
    memcpy(C->currentSpectralState, C->proposedSpectralState, sizeof(fftw_complex) * C->nj * ((C->nk >> 1) + 1));
    memcpy(C->currentPhysicalState, C->proposedPhysicalState, sizeof(double) * C->nj * C->nk);
    C->accProb = alpha;
    C->logLHDCurrentState = logLHDOfProposal;
  }

  infmcmc_updateAvgs(C);
}

void infmcmc_updateVectorFieldRWMH(INFCHAIN *C1, INFCHAIN *C2, double logLHDOfProposal) {
  double alpha;

  // log likelihoods will be the same for both chains
	alpha = exp(C1->logLHDCurrentState - logLHDOfProposal);

  if (alpha > 1.0) {
    alpha = 1.0;
  }

  C1->accepted = 0;
  C2->accepted = 0;
  C1->accProb = alpha;
  C2->accProb = alpha;

  if (gsl_rng_uniform(C1->r) < alpha) {
    C1->accepted = 1;
    C2->accepted = 1;
    memcpy(C1->currentSpectralState, C1->proposedSpectralState, sizeof(fftw_complex) * C1->nj * ((C1->nk >> 1) + 1));
    memcpy(C1->currentPhysicalState, C1->proposedPhysicalState, sizeof(double) * C1->nj * C1->nk);
    C1->logLHDCurrentState = logLHDOfProposal;
    memcpy(C2->currentSpectralState, C2->proposedSpectralState, sizeof(fftw_complex) * C2->nj * ((C2->nk >> 1) + 1));
    memcpy(C2->currentPhysicalState, C2->proposedPhysicalState, sizeof(double) * C2->nj * C2->nk);
    C2->logLHDCurrentState = logLHDOfProposal;
  }

  infmcmc_updateAvgs(C1);
  infmcmc_updateAvgs(C2);
}

void infmcmc_setRWMHStepSize(INFCHAIN *C, double beta) {
  C->rwmhStepSize = beta;
}

void infmcmc_setPriorAlpha(INFCHAIN *C, double alpha) {
  C->alphaPrior = alpha;
}

void infmcmc_setPriorVar(INFCHAIN *C, double var) {
  C->priorVar = var;
  C->priorStd = sqrt(var);
}

double L2Field(fftw_complex *uk, int nj, int nk) {
  int j, k;
  const int maxk = (nk >> 1) + 1;
  double sum = 0.0;

  for (j = 0; j < nj; j++) {
    for (k = 0; k < maxk; k++) {
      sum += cabs(uk[maxk*j+k]) * cabs(uk[maxk*j+k]);
    }
  }

  return 2.0 * sum;
}

double infmcmc_L2Current(INFCHAIN *C) {
  return L2Field(C->currentSpectralState, C->nj, C->nk);
}

double infmcmc_L2Proposed(INFCHAIN *C) {
  return L2Field(C->proposedSpectralState, C->nj, C->nk);
}

double infmcmc_L2Prior(INFCHAIN *C) {
  return L2Field(C->priorDraw, C->nj, C->nk);
}

void randomPriorDrawOLD(gsl_rng *r, double PRIOR_ALPHA, fftw_complex *randDrawCoeffs) {
  int j, k;
  double xrand, yrand, c;//, scale;
  //const double one = 1.0;

  c = 4.0 * M_PI * M_PI;

  for(j = 0; j < nj1; j++) {
    for(k = 0; k < nk1/2 + 1; k++) {
      if((j == 0) && (k == 0)) {
        randDrawCoeffs[0] = 0.0;
      }
      else if((j == nj1/2) && (k == nk1/2)) {
        xrand = gsl_ran_gaussian_ziggurat(r, 1.0);
        randDrawCoeffs[(nk1/2 + 1) * nj1/2 + nk1/2] = xrand / pow((c * ((j * j) + (k * k))), (double)PRIOR_ALPHA/2.0);
      }
      else if((j == 0) && (k == nk1/2)) {
        xrand = gsl_ran_gaussian_ziggurat(r, 1.0);
        randDrawCoeffs[nk1/2] = xrand / pow((c * ((j * j) + (k * k))), (double)PRIOR_ALPHA/2.0);
      }
      else if((j == nj1/2) && (k == 0)) {
        xrand = gsl_ran_gaussian_ziggurat(r, 1.0);
        randDrawCoeffs[(nk1/2 + 1) * nj1/2] = xrand / pow((c * ((j * j) + (k * k))), (double)PRIOR_ALPHA/2.0);
      }
      else {
        xrand = gsl_ran_gaussian_ziggurat(r, 1.0) / sqrt(2.0);
        yrand = gsl_ran_gaussian_ziggurat(r, 1.0) / sqrt(2.0);
        randDrawCoeffs[(nk1/2 + 1) * j + k] = (xrand + I * yrand) / pow((c * ((j * j) + (k * k))), (double)PRIOR_ALPHA/2.0);
        if(j > nj1/2) {
          randDrawCoeffs[(nk1/2 + 1) * j + k] = conj(randDrawCoeffs[(nk1/2+1)*(nj1-j)+k]);
        }
      }
    }
  }

  /*
  for(j = 1; j < nj1 / 2; j++) {
    for(k = 0; k < nk1 / 2 + 1; k++) {
      xrand = gsl_ran_gaussian_ziggurat(r, 1.0) / M_SQRT2;
      yrand = gsl_ran_gaussian_ziggurat(r, 1.0) / M_SQRT2;
      scale = pow(c * ((j * j) + (k * k)), (double)PRIOR_ALPHA/2.0);
      randDrawCoeffs[(nk1/2 + 1) * j + k] = (xrand + I * yrand) / scale;
      randDrawCoeffs[(nk1/2 + 1) * (nj1-j) + k] = (xrand - I * yrand) / scale;
    }
  }

  for(k = 1; k < nk1 / 2 + 1; k++) {
    xrand = gsl_ran_gaussian_ziggurat(r, 1.0) / M_SQRT2;
    randDrawCoeffs[k] = (xrand + I * yrand) / pow(c * (k * k), (double)PRIOR_ALPHA/2.0);
  }

  for(k = 0; k < nk1/2; k++) {
    xrand = gsl_ran_gaussian_ziggurat(r, 1.0) / M_SQRT2;
    randDrawCoeffs[(nk1/2+1)*(nj1/2)+k] = (xrand + I * yrand) / pow(c * ((nj1 * nj1 / 4) + (k * k)), (double)PRIOR_ALPHA/2.0);
  }

  randDrawCoeffs[0] = 0.0;

  xrand = gsl_ran_gaussian_ziggurat(r, 1.0);
  randDrawCoeffs[(nk1/2+1)*(nj1/2)+(nk1/2)] = xrand / pow(c * ((j * j) + (k * k)), (double)PRIOR_ALPHA/2.0);
  */
}

void setRWMHStepSize(CHAIN *C, double stepSize) {
  C->rwmhStepSize = stepSize;
}

void resetAvgs(CHAIN *C) {
  /*
    Sets avgPhysicalState to currentPhysicalState and
         avgSpectralState to currentSpectralState
  */
  const size_t size_doublenj = sizeof(double) * C->nj;
  memcpy(C->avgPhysicalState, C->currentPhysicalState, size_doublenj * C->nk);
  memcpy(C->avgSpectralState, C->currentSpectralState, size_doublenj * ((C->nk >> 1) + 1));
}

void resetVar(CHAIN *C) {
  // Sets varPhysicalState to 0
  memset(C->varPhysicalState, 0, sizeof(double) * C->nj * C->nk);
}

void proposeIndependence(CHAIN *C) {
  const int maxk = (C->nk >> 1) + 1;

  memcpy(C->proposedSpectralState, C->priorDraw, sizeof(fftw_complex) * C->nj * maxk);
}

double lsqFunctional(const double * const data, const double * const obsVec, const int obsVecSize, const double obsStdDev) {
  int i;
  double temp1, sum1 = 0.0;

  for(i = 0; i < obsVecSize; i++) {
    temp1 = (data[i] - obsVec[i]);
    sum1 += temp1 * temp1;
  }

  return sum1 / (2.0 * obsStdDev * obsStdDev);
}

void acceptReject(CHAIN *C) {
  const double phi1 = lsqFunctional(C->data, C->currentStateObservations, C->sizeObsVector, C->obsStdDev);
  const double phi2 = lsqFunctional(C->data, C->proposedStateObservations, C->sizeObsVector, C->obsStdDev);
  double tempAccProb = exp(phi1 - phi2);

  if(tempAccProb > 1.0) {
    tempAccProb = 1.0;
  }

  if(gsl_rng_uniform(C->r) < tempAccProb) {
    memcpy(C->currentSpectralState, C->proposedSpectralState, sizeof(fftw_complex) * C->nj * ((C->nk >> 1) + 1));
    memcpy(C->currentPhysicalState, C->proposedPhysicalState, sizeof(double) * C->nj * C->nk);
    memcpy(C->currentStateObservations, C->proposedStateObservations, sizeof(double) * C->sizeObsVector);

    C->accProb = tempAccProb;
    C->currentLSQFunctional = phi2;
  }
  else {
    C->currentLSQFunctional = phi1;
  }
}

//
//void updateChain(CHAIN *C) {
//C->currentIter++;
//void updateAvgs(CHAIN *C);
//void updateVar(CHAIN *C);
/*
int main(void) {
  int nj = 32, nk = 32, sizeObsVector = 0;
  unsigned long int randseed = 0;

  FILE *fp;
  CHAIN *C;

  C = (CHAIN *)malloc(sizeof(CHAIN));

  initChain(C, nj, nk, sizeObsVector, randseed);

  randomPriorDraw2(C);

  fp = fopen("2.dat", "w");
  fwrite(C->priorDraw, sizeof(fftw_complex), nj * (nk / 2 + 1), fp);
  fclose(fp);

  return 0;
}
*/


/**
 * \author Sylvain Marsat, University of Maryland - NASA GSFC
 *
 * \brief C header for functions computing coeffcients of Not-A-Knot and Quadratic splines in matrix form.
 *
 */

#ifndef _FRESNEL_H
#define _FRESNEL_H

#define _XOPEN_SOURCE 500

#ifdef __GNUC__
#define UNUSED __attribute__ ((unused))
#else
#define UNUSED
#endif

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <complex.h>
#include <time.h>
#include <unistd.h>
#include <getopt.h>
#include <stdbool.h>
#include <string.h>

#include <gsl/gsl_errno.h>
#include <gsl/gsl_bspline.h>
#include <gsl/gsl_blas.h>
#include <gsl/gsl_min.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_complex.h>

#include "constants.h"
#include "struct.h"

#if defined(__cplusplus)
#define complex _Complex
extern "C" {
#elif 0
} /* so that editors will match preceding brace */
#endif

double complex ComputeInt(
  gsl_matrix* splinecoeffsAreal,         /*  */
  gsl_matrix* splinecoeffsAimag,         /*  */
  gsl_matrix* splinecoeffsphase);        /*  */

double complex ComputeIntCase1a(
  const double complex* coeffsA,         /* */
  const double p1,                       /* */
  const double p2,                       /* */
  const double scale);                   /* */
double complex ComputeIntCase1b(
  const double complex* coeffsA,         /* */
  const double p1,                       /* */
  const double p2,                       /* */
  const double scale);                   /* */
double complex ComputeIntCase2(
  const double complex* coeffsA,         /* */
  const double p1,                       /* */
  const double p2);                      /* */
double complex ComputeIntCase3(
  const double complex* coeffsA,         /* */
  const double p1,                       /* */
  const double p2);                      /* */
double complex ComputeIntCase4(
  const double complex* coeffsA,         /* */
  const double p1,                       /* */
  const double p2);                      /* */

#if 0
{ /* so that editors will match succeeding brace */
#elif defined(__cplusplus)
}
#endif

#endif /* _FRESNEL_H */


/* multifit/work.c
 *
 * Copyright (C) 2000 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., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 */

#include <config.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_multifit.h>

gsl_multifit_linear_workspace *
gsl_multifit_linear_alloc (size_t n, size_t p)
{
  gsl_multifit_linear_workspace *w;

  w = (gsl_multifit_linear_workspace *)
    malloc (sizeof (gsl_multifit_linear_workspace));

  if (w == 0)
    {
      GSL_ERROR_VAL ("failed to allocate space for multifit_linear struct",
                     GSL_ENOMEM, 0);
    }

  w->n = n;                     /* number of observations */
  w->p = p;                     /* number of parameters */

  w->A = gsl_matrix_alloc (n, p);

  if (w->A == 0)
    {
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for A", GSL_ENOMEM, 0);
    }

  w->Q = gsl_matrix_alloc (p, p);

  if (w->Q == 0)
    {
      gsl_matrix_free (w->A);
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for Q", GSL_ENOMEM, 0);
    }

  w->QSI = gsl_matrix_alloc (p, p);

  if (w->QSI == 0)
    {
      gsl_matrix_free (w->Q);
      gsl_matrix_free (w->A);
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for QSI", GSL_ENOMEM, 0);
    }

  w->S = gsl_vector_alloc (p);

  if (w->S == 0)
    {
      gsl_matrix_free (w->QSI);
      gsl_matrix_free (w->Q);
      gsl_matrix_free (w->A);
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for S", GSL_ENOMEM, 0);
    }

  w->t = gsl_vector_alloc (n);

  if (w->t == 0)
    {
      gsl_vector_free (w->S);
      gsl_matrix_free (w->QSI);
      gsl_matrix_free (w->Q);
      gsl_matrix_free (w->A);
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for t", GSL_ENOMEM, 0);
    }

  w->xt = gsl_vector_calloc (p);

  if (w->xt == 0)
    {
      gsl_vector_free (w->t);
      gsl_vector_free (w->S);
      gsl_matrix_free (w->QSI);
      gsl_matrix_free (w->Q);
      gsl_matrix_free (w->A);
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for xt", GSL_ENOMEM, 0);
    }

  w->D = gsl_vector_calloc (p);

  if (w->D == 0)
    {
      gsl_vector_free (w->D);
      gsl_vector_free (w->t);
      gsl_vector_free (w->S);
      gsl_matrix_free (w->QSI);
      gsl_matrix_free (w->Q);
      gsl_matrix_free (w->A);
      free (w);
      GSL_ERROR_VAL ("failed to allocate space for xt", GSL_ENOMEM, 0);
    }

  return w;
}

void
gsl_multifit_linear_free (gsl_multifit_linear_workspace * work)
{
  gsl_matrix_free (work->A);
  gsl_matrix_free (work->Q);
  gsl_matrix_free (work->QSI);
  gsl_vector_free (work->S);
  gsl_vector_free (work->t);
  gsl_vector_free (work->xt);
  gsl_vector_free (work->D);
  free (work);
}


//! \file heat_driver.h
//! Driver for Magnolia's heat transfer solver
#ifndef ENRICO_SURROGATE_HEAT_DRIVER_H
#define ENRICO_SURROGATE_HEAT_DRIVER_H

#include "enrico/geom.h"
#include "enrico/heat_fluids_driver.h"

#include <gsl/gsl>
#include <mpi.h>
#include <pugixml.hpp>
#include <xtensor/xtensor.hpp>

#include <cstddef>

namespace enrico {

//! Struct containing geometric information for a flow channel
struct Channel {
  //! Channel index
  int index_;

  //! Channel flow area
  double area_;

  //! Vector of rod IDs connected to this channel, all with a fractional perimeter
  //! in contact with the channel equal to 0.25
  std::vector<std::size_t> rod_ids_;
};

//! Struct containing geometry information for a cylindrical solid rod
struct Rod {
  //! Rod index
  int index_;

  //! Rod cladding outer radius
  double clad_outer_radius_;

  //! Rod cladding inner radius
  double clad_inner_radius_;

  //! Rod pellet radius
  double pellet_radius_;

  //! Vector of channel IDs connected to this rod, all with a fractional perimeter
  //! in contact with the rod equal to 0.25
  std::vector<std::size_t> channel_ids_;
};

//! Class to construct flow channels for a Cartesian lattice of pins
class ChannelFactory {
public:
  ChannelFactory(double pitch, double rod_radius)
    : pitch_(pitch)
    , radius_(rod_radius)
    , interior_area_(pitch_ * pitch_ - M_PI * radius_ * radius_)
  {}

  //! Make a corner subchannel connected to given rods
  Channel make_corner(const std::vector<std::size_t>& rods) const
  {
    Channel c;
    c.index_ = index_++;
    c.area_ = 0.25 * interior_area_;
    c.rod_ids_ = rods;
    return c;
  }

  //! Make an edge subchannel connected to given rods
  Channel make_edge(const std::vector<std::size_t>& rods) const
  {
    Channel c;
    c.index_ = index_++;
    c.area_ = 0.5 * interior_area_;
    c.rod_ids_ = rods;
    return c;
  }

  //! Make an interior subchannel connected to given rods
  Channel make_interior(const std::vector<std::size_t>& rods) const
  {
    Channel c;
    c.index_ = index_++;
    c.area_ = interior_area_;
    c.rod_ids_ = rods;
    return c;
  }

private:
  //! rod pitch
  double pitch_;

  //! rod outer radius
  double radius_;

  //! interior channel flow area, which is proportional to flow areas for all
  //! other channel types
  double interior_area_;

  //! index of constructed channel
  static int index_;
};

class RodFactory {
public:
  RodFactory(double clad_OR, double clad_IR, double pellet_OR)
    : clad_outer_r_(clad_OR)
    , clad_inner_r_(clad_IR)
    , pellet_outer_r_(pellet_OR)
  {}

  //! Make a rod connected to given channels
  Rod make_rod(const std::vector<std::size_t>& channels) const
  {
    Rod r;
    r.index_ = index_++;
    r.clad_inner_radius_ = clad_inner_r_;
    r.clad_outer_radius_ = clad_outer_r_;
    r.pellet_radius_ = pellet_outer_r_;
    r.channel_ids_ = channels;
    return r;
  }

private:
  //! Cladding outer radius
  double clad_outer_r_;

  //! Cladding inner radius
  double clad_inner_r_;

  //! Pellet outer radius
  double pellet_outer_r_;

  //! Index of constructed rod
  static int index_;
};

/**
 * Class providing surrogate thermal-hydraulic solution for a Cartesian
 * bundle of rods with upwards-flowing coolant. A conduction model is used
 * for the solid phase, with axial conduction neglected. The solid phase is
 * linked to the fluid phase by conjugate heat transfer, which is treated
 * here with a pseudo-steady-state approach where the power entering the fluid
 * matches the power in the rod at that axial elevation. It is assumed that
 * there is zero thermal resistance between the rod and the fluid.
 *
 * The fluid solution is obtained with a very simplified "subchannel" method
 * that neglects all crossflow terms between channels such that the method
 * is more akin to a single-pin analysis in a coolant-centered basis. The
 * enthalpy is solved by simply axial energy balance, while the axial momentum
 * equation is solved for pressure (the mass flow rate in each channel being
 * fixed) while neglecting friction effects.
 */
class SurrogateHeatDriver : public HeatFluidsDriver {
public:
  //! Initializes heat-fluids surrogate with the given MPI communicator.
  //!
  //! \param comm  The MPI communicator used to initialze the surrogate
  //! \param node  XML node containing settings for surrogate
  explicit SurrogateHeatDriver(MPI_Comm comm, pugi::xml_node node);

  //! Verbosity options for printing simulation results
  enum class verbose { NONE, LOW, HIGH };

  bool has_coupling_data() const final { return comm_.rank == 0; }

  //! Get the number of local mesh elements
  //! \return Number of local mesh elements
  int n_local_elem() const override;

  //! Get the number of global mesh elements
  //! \return Number of global mesh elements
  std::size_t n_global_elem() const override;

  //! Set the heat source for a given local element
  //!
  //! \param local_elem A local element ID
  //! \param heat A heat source term
  //! \return Error code
  int set_heat_source_at(int32_t local_elem, double heat) override;

  //! Solves the heat-fluids surrogate solver
  void solve_step() final;

  void solve_heat();

  void solve_fluid();

  //! Returns Number of rings in fuel and clad
  std::size_t n_rings() const { return n_fuel_rings_ + n_clad_rings_; }

  //! Returns cladding inner radius
  double clad_inner_radius() const { return clad_inner_radius_; }

  //! Returns cladding outer radius
  double clad_outer_radius() const { return clad_outer_radius_; }

  //! Returns pellet outer radius
  double pellet_radius() const { return pellet_radius_; }

  //! Returns number of fuel rings
  std::size_t n_fuel_rings() const { return n_fuel_rings_; }

  //! Returns number of clad rings
  std::size_t n_clad_rings() const { return n_clad_rings_; }

  //! Returns number of pins in x-direction
  std::size_t n_pins_x() const { return n_pins_x_; }

  //! Returns number of pins in y-direction
  std::size_t n_pins_y() const { return n_pins_y_; }

  //! Returns pin pitch
  double pin_pitch() const { return pin_pitch_; }

  //! Returns inlet temperature boundary condition in [K]
  double inlet_temperature() const { return inlet_temperature_; }

  //! Returns inlet mass flowrate boundary condition in [kg/s]
  double mass_flowrate() const { return mass_flowrate_; }

  //! Returns maximum number of subchannel iterations
  std::size_t max_subchannel_its() const { return max_subchannel_its_; }

  //! Returns subchannel convergence tolerance for enthalpy
  double subchannel_tol_h() const { return subchannel_tol_h_; }

  //! Returns subchannel convergence tolerance for pressure
  double subchannel_tol_p() const { return subchannel_tol_p_; }

  //! Returns convergence tolerance for solid energy equation
  double heat_tol() const { return heat_tol_; }

  //! Write data to VTK
  void write_step(int timestep, int iteration) final;

  //! Returns solid temperature in [K] for given region
  double solid_temperature(std::size_t pin, std::size_t axial, std::size_t ring) const;

  //! Returns fluid density in [g/cm^3] for given region
  double fluid_density(std::size_t pin, std::size_t axial) const;

  //! Returns fluid temperature in [K] for given region
  double fluid_temperature(std::size_t pin, std::size_t axial) const;

  // Data on fuel pins
  xt::xtensor<double, 2> pin_centers_; //!< (x,y) values for center of fuel pins
  xt::xtensor<double, 1> z_;           //!< Bounding z-values for axial segments
  std::size_t n_axial_;                //!< number of axial segments
  std::size_t n_azimuthal_{4};         //!< number of azimuthal segments

  //! Total number of pins
  std::size_t n_pins_;

  // Dimensions for a single fuel pin axial segment
  double clad_outer_radius_;     //!< clad outer radius in [cm]
  double clad_inner_radius_;     //!< clad inner radius in [cm]
  double pellet_radius_;         //!< fuel pellet radius in [cm]
  std::size_t n_fuel_rings_{20}; //!< number of fuel rings
  std::size_t n_clad_rings_{2};  //!< number of clad rings

  //!< Channels in the domain
  std::vector<Channel> channels_;

  //!< Rods in the domain
  std::vector<Rod> rods_;

  //! Mass flowrate for coolant-centered channels; this is determine by distributing
  //! a total inlet mass flowrate among the channels based on the fractional flow area.
  xt::xtensor<double, 1> channel_flowrates_;

  // solver variables and settings
  xt::xtensor<double, 4>
    source_; //!< heat source for each (pin, axial segment, ring, azimuthal segment)
  xt::xtensor<double, 1> r_grid_clad_; //!< radii of each clad ring in [cm]
  xt::xtensor<double, 1> r_grid_fuel_; //!< radii of each fuel ring in [cm]

  //! Cross-sectional areas of rings in fuel and cladding
  xt::xtensor<double, 1> solid_areas_;

  // visualization
  std::string viz_basename_{
    "heat_surrogate"}; //!< base filename for visualization files (default: magnolia)
  std::string viz_iterations_{
    "none"};                    //!< visualization iterations to write (none, all, final)
  std::string viz_data_{"all"}; //!< visualization data to write
  std::string viz_regions_{"all"}; //!< visualization regions to write
  size_t vtk_radial_res_{20};      //!< radial resolution of resulting vtk files

private:
  //! Get temperature of local mesh elements
  //! \return Temperature of local mesh elements in [K]
  std::vector<double> temperature_local() const override;

  //! Get density of local mesh elements
  //! \return Density of local mesh elements in [g/cm^3]
  std::vector<double> density_local() const override;

  //! States whether each local region is in fluid
  //! \return For each local region, 1 if region is in fluid and 0 otherwise
  std::vector<int> fluid_mask_local() const override;

  //! Get centroids of local mesh elements
  //! \return Centroids of local mesh elements
  std::vector<Position> centroid_local() const override;

  //! Get volumes of local mesh elements
  //! \return Volumes of local mesh elements
  std::vector<double> volume_local() const override;

  //! Create internal arrays used for heat equation solver
  void generate_arrays();

  //! Channel index in terms of row, column index
  int channel_index(int row, int col) const { return row * (n_pins_x_ + 1) + col; }

  //! Rod power at a given node in a given pin, computed by integrating the heat source
  //! (assumed constant in each ring) over the pin.
  //! \param pin   pin index
  //! \param axial axial index
  double rod_axial_node_power(const int pin, const int axial) const;

  //! Diagnostic function to assess whether the mass is conserved by the subchannel
  //! solver by comparing the mass flowrate in each axial plane (at cell-centered
  //! positions) to the specified inlet mass flowrate.
  //! \param rho density in a cell-centered basis
  //! \param u   axial velocity in a face-centered basis
  bool is_mass_conserved(const xt::xtensor<double, 2>& rho,
                         const xt::xtensor<double, 2>& u) const;

  //! Diagnostic function to assess whether the energy is conserved by the subchannel
  //! solver by comparing the energy deposition in each channel in each axial plane
  //! (at cell-centered positions) to the powers of the rods connected to that channel.
  //! \param rho density in a cell-centered basis
  //! \param u   axial velocity in a face-centered basis
  //! \param h   enthalpy in a face-centered basis
  //! \param q   powers in each channel in a cell-centered basis
  bool is_energy_conserved(const xt::xtensor<double, 2>& rho,
                           const xt::xtensor<double, 2>& u,
                           const xt::xtensor<double, 2>& h,
                           const xt::xtensor<double, 2>& q) const;

  //!< solid temperature in [K] for each (pin, axial segment, ring)
  xt::xtensor<double, 3> solid_temperature_;

  //! Flow areas for coolant-centered channels
  xt::xtensor<double, 1> channel_areas_;

  //! Fluid temperature in a rod-centered basis indexed by rod ID and axial ID
  xt::xtensor<double, 2> fluid_temperature_;

  //! Fluid density in [g/cm^3] in a rod-centered basis indexed by rod ID and axial ID
  xt::xtensor<double, 2> fluid_density_;

  //! Number of pins in the x-direction in a Cartesian grid
  std::size_t n_pins_x_;

  //! Number of pins in the y-direction in a Cartesian grid
  std::size_t n_pins_y_;

  //! Pin pitch, assumed the same for the x and y directions
  double pin_pitch_;

  //! Inlet fluid temperature [K]
  double inlet_temperature_;

  //! Mass flowrate of fluid into the domain [kg/s]
  double mass_flowrate_;

  //! Number of channels
  std::size_t n_channels_;

  //! Maximum number of iterations for subchannel solution, set to a default value
  //! of 100 if not set by the user
  int max_subchannel_its_ = 100;

  //! Convergence tolerance on enthalpy for the subchannel solution for use in
  //! convergence based on the L-1 norm, set to a default value of 1e-2
  double subchannel_tol_h_ = 1e-2;

  //! Convergence tolerance on pressure for the subchannel solution for use in
  //! convergence based on the L-1 norm, set to a default value of 1e-2
  double subchannel_tol_p_ = 1e-2;

  //! Convergence tolerance for solid temperature solution, set to a default value
  //! of 1e-4
  double heat_tol_ = 1e-4;

  //! Gravitational acceleration
  const double g_ = 9.81;

  //! Verbosity setting for printing simulation results; defaults to NONE
  verbose verbosity_ = verbose::NONE;

}; // end SurrogateHeatDriver

} // namespace enrico

#endif // ENRICO_SURROGATE_HEAT_DRIVER_H


/* 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 Foundation, Inc., 59 Temple Place, Suite
   330, Boston, MA 02111-1307 USA

   Original implementation was copyright (C) 1997 Makoto Matsumoto and
   Takuji Nishimura. Coded by Takuji Nishimura, considering the
   suggestions by Topher Cooper and Marc Rieffel in July-Aug. 1997, "A
   C-program for MT19937: Integer version (1998/4/6)"

   This implementation copyright (C) 1998 Brian Gough. I reorganized
   the code to use the module framework of GSL.  The license on this
   implementation was changed from LGPL to GPL, following paragraph 3
   of the LGPL, version 2.

   The seeding procedure has been updated to match the 10/99 release
   of MT19937.

   The original code included the comment: "When you use this, send an
   email to: matumoto@math.keio.ac.jp with an appropriate reference to
   your work".

   Makoto Matsumoto has a web page with more information about the
   generator, http://www.math.keio.ac.jp/~matumoto/emt.html.

   The paper below has details of the algorithm.

   From: Makoto Matsumoto and Takuji Nishimura, "Mersenne Twister: A
   623-dimensionally equidistributerd uniform pseudorandom number
   generator". ACM Transactions on Modeling and Computer Simulation,
   Vol. 8, No. 1 (Jan. 1998), Pages 3-30

   You can obtain the paper directly from Makoto Matsumoto's web page.

   The period of this generator is 2^{19937} - 1.

*/

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

static inline unsigned long int mt_get (void *vstate);
static double mt_get_double (void *vstate);
static void mt_set (void *state, unsigned long int s);

#define N 624	/* Period parameters */
#define M 397

/* most significant w-r bits */
static const unsigned long UPPER_MASK = 0x80000000UL;

/* least significant r bits */
static const unsigned long LOWER_MASK = 0x7fffffffUL;

typedef struct
  {
    unsigned long mt[N];
    int mti;
  }
mt_state_t;

static inline unsigned long
mt_get (void *vstate)
{
  mt_state_t *state = (mt_state_t *) vstate;

  unsigned long k ;
  unsigned long int *const mt = state->mt;

#define MAGIC(y) (((y)&0x1) ? 0x9908b0dfUL : 0)

  if (state->mti >= N)
    {	/* generate N words at one time */
      int kk;

      for (kk = 0; kk < N - M; kk++)
	{
	  unsigned long y = (mt[kk] & UPPER_MASK) | (mt[kk + 1] & LOWER_MASK);
	  mt[kk] = mt[kk + M] ^ (y >> 1) ^ MAGIC(y);
	}
      for (; kk < N - 1; kk++)
	{
	  unsigned long y = (mt[kk] & UPPER_MASK) | (mt[kk + 1] & LOWER_MASK);
	  mt[kk] = mt[kk + (M - N)] ^ (y >> 1) ^ MAGIC(y);
	}

      {
	unsigned long y = (mt[N - 1] & UPPER_MASK) | (mt[0] & LOWER_MASK);
	mt[N - 1] = mt[M - 1] ^ (y >> 1) ^ MAGIC(y);
      }

      state->mti = 0;
    }

  /* Tempering */

  k = mt[state->mti];
  k ^= (k >> 11);
  k ^= (k << 7) & 0x9d2c5680UL;
  k ^= (k << 15) & 0xefc60000UL;
  k ^= (k >> 18);

  state->mti++;

  return k;
}

static double
mt_get_double (void * vstate)
{
  return mt_get (vstate) / 4294967296.0 ;
}

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

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

  /* This is the October 1999 version of the seeding procedure. It
     was updated by the original developers to avoid the periodicity
     in the simple congruence originally used.

     Note that an ANSI-C unsigned long integer arithmetic is
     automatically modulo 2^32 (or a higher power of two), so we can
     safely ignore overflow. */

#define LCG(x) ((69069 * x) + 1) &0xffffffffUL

  for (i = 0; i < N; i++)
    {
      state->mt[i] = s & 0xffff0000UL;
      s = LCG(s);
      state->mt[i] |= (s &0xffff0000UL) >> 16;
      s = LCG(s);
    }

  state->mti = i;
}

/* This is the original version of the seeding procedure, no longer
   used but available for compatibility with the original MT19937. */

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

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

  state->mt[0] = s & 0xffffffffUL;

#define LCG1998(n) ((69069 * n) & 0xffffffffUL)

  for (i = 1; i < N; i++)
    state->mt[i] = LCG1998 (state->mt[i - 1]);

  state->mti = i;
}

static const gsl_rng_type mt_type =
{"mt19937",			/* name */
 0xffffffffUL,			/* RAND_MAX  */
 0,			        /* RAND_MIN  */
 sizeof (mt_state_t),
 &mt_set,
 &mt_get,
 &mt_get_double};

static const gsl_rng_type mt_1998_type =
{"mt19937_1998",		/* name */
 0xffffffffUL,			/* RAND_MAX  */
 0,			        /* RAND_MIN  */
 sizeof (mt_state_t),
 &mt_1998_set,
 &mt_get,
 &mt_get_double};

const gsl_rng_type *gsl_rng_mt19937 = &mt_type;
const gsl_rng_type *gsl_rng_mt19937_1998 = &mt_1998_type;

/* MT19937 is the default generator, so define that here too */

const gsl_rng_type *gsl_rng_default = &mt_type;
unsigned long int gsl_rng_default_seed = 0;