/* spatial.c -- main driver program */
/* Kate Cowles */
/* generates fixed effects and spatially varying coefs as one block, 
   spatiotemp errors for each time as a separate block */
/* centering of all continuous covars */
/* spatiotemporal errors; geostat for 2-dim spatial, ar1 for temporal */
/* stores sum of fixed and random effects as mbetas */
/* 4/11/2001 */

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

#define NTV      0           /* number of time varying random coefs */
#define SEEDFILE "seedfile" 

void readrand();
void readdata();
void calcstuf();
void writchn(int iter, FILE *out, int thin);
void initchn();
void genz(int iter);
void genbeta(int iter);
void gentau(int iter);
void gentaub(int iter);
void gentaume(int iter);
void genphi(int iter);
void genphib(int iter);
void genrho(int iter);
void calcpred(int iter);
double calcdevmean();
void updtseed();
int cholanis(Matrixlt &l, double phi, Matrix &dsq);
double covfunc(double phi, double temp, int corrtype);
void cholar1(Matrixlt &l, double rho);
double compztz(Matrixlt &lar1inv, Matrixlt &lspatinv, Matrix &mz);

// global variables

   unsigned short seed16v[3] = {1, 2 ,3};
   char infile[128], inifile[128], outfile[128];

   int N, M, T, S, W,
       PM, PS, PT, NSV, NBETA, TOTBETA,
       CORRFUNC, MAXITERS, STARTCOMP, SKIP, REPS;
   double WIDTH, WIDTHDIV;

   Matrix adjmat,	// adjacency matrix of watershed
          mx,		// design mat for meas types
          sx,		// design mat for sites, 1st 2 are coords
          tx,		// design mat for times
          X,		// derived design mat
          Xs;		// design mat for spatially varying coefs

   int pppctr;
   Vector pppval;
  
   int **idx,		// index w/n watershed of which observation in mw vector
			// is first for each time
       *mm,		// which meas type is each obs
       *numtype,	// counters for total # of each type of meas
       ****numtypes,		// counters for # of each type of meas at each
       *starts,		// earliest and latest time for each watershed
       *startt,
       *stops,		// first and last site for each watershed
       *stopt,
       *ss,		// at which site was each obs meas'ed
       *tt,		// at which time was each obs meas'ed
       *vs,		// site index of mz
       *vt,		// time index of mz
       *wsi,		// which watershed each site is in
       *wss,		// tot num sites in watershed
       *wst;		// span of earliest to latest timest in watershd
			// watershed/site/time

       
   double logolddetbinv, logolddetbinv2,
          phi, rho, tau;
   Vector logolddet, mbeta, mu, mw, taume;
   Matrix *mz,		// spatiotemp errors
          *meanz,
          *dsq;		// distance matrix w/n ea watershed
   Matrixlt *lspat, *lar1, *lspatinv, *lar1inv;

   double phib, ztzb, taub, taub2, phib2, ztzb2;
   Vector mbetas,	// watershed specific slopes on time
          mbetas2;	// watershed specific intercepts
   Matrix spatbinv, spatbinv2;
   Matrixlt lspatbinv, lspatbinv2;
   
   double meantaub, meanphib, meantau, meanphi, meanrho,
          meantaub2, meanphib2, deviance, meandev;
   Vector ztz, xbar,
          meanmbeta, meanmbetas, meantaume, meanmbetas2;

   // in the above, mu will be used for the obs-specific means, computed
   // from the betas

// priors

   double aphi , bphi,			// uniform prior on phi
          aphib, bphib,			// uniform prior on phib and phib2
          a, b,				// gamma prior on tau
          ab, bb,			// gamma prior on taub
          ab2, bb2;			// gamma prior on taub2
   Vector c, d,				// gamma priors on taumes
          beta0,
          sigbinvb0;			// prior mean and prec mat of beta
   Matrix sigbinv;

   double rhomean, rhoprec,
          rholeft, rhoright;

   int movectr1, movectr2, movectr3, movectr4;


int main(int argc, char *argv[])
{
   int i, iter;
   double dic, devmean;
   FILE *out;
   
   strcpy(inifile, argv[0]);
   strcat(inifile, ".ini");
   strcpy(outfile, argv[0]);
   strcat(outfile, ".out");
   iter = 1;
   i = 0;
   while(iter) {
      i++;
      if(i >= argc) {
         error("main()", "must specify an input file", ERREXIT);
      } else if(argv[i][0] == '-') {
         switch(argv[i][1]) {
            case 'i':
               i++;
               if(i < argc) strcpy(inifile, argv[i]);
               else error("main()", "command line option requires a value", ERREXIT);
               break;
            case 'o':
               i++;
               if(i < argc) strcpy(outfile, argv[i]);
               else error("main()", "command line option requires a value", ERREXIT);
               break;
            default:
               error("main()", "unsupported command line option", ERREXIT);
         }
      } else {
          iter = 0;
      }
   }
   strcpy(infile, argv[i]);
   
   readrand();			// initialize random number seed
   printf("readdata \n");
   readdata();			// read data file
   printf("calcstuf \n");
   calcstuf();			// one-time calculations

   out = fopen(outfile, "w");
   if (out == NULL) {
      printf("Data file cannot be opened for writing.\n");
      exit(1);
   }

   initchn();			// initial values for parameters
   writchn(0, out, SKIP);
     
   for(iter = 1; iter < MAXITERS; iter++) {
      printf("iter %d \n", iter);
      genbeta(iter);
      genz(iter);
      gentau(iter); 
      gentaume(iter);
      genphi(iter);
      genrho(iter);
      gentaub(iter); 
      genphib(iter);
      calcpred(iter);
      writchn(iter, out, 1);
   }  

   fclose(out);

   printf(" movectr1 %d acceptance ratio %7.4f \n", movectr1,
          movectr1 / (double) MAXITERS);
   printf(" movectr2 %d acceptance ratio %7.4f \n", movectr2,
          movectr2 / (double) MAXITERS);
   printf(" movectr3 %d acceptance ratio %7.4f \n", movectr3,
          movectr3 / (double) MAXITERS);
   printf(" movectr4 %d acceptance ratio %7.4f \n", movectr4,
          movectr4 / (double) MAXITERS);

   printf(" posterior predictive p values \n");
   for (i = 0; i < N; i++)
      printf(" pppval[%d] %7.4f \n", i, pppval[i] / pppctr);

   devmean = calcdevmean();			// deviance from means of parms
   dic = 2.0 * meandev - devmean;
   printf("DIC %7.4f devmean %7.4f meandev %7.4f eff num parms %7.4f \n",
          dic, devmean, meandev, meandev - devmean);

   updtseed();
   
   return 0;

}


/* **************** */
/*    readrand()    */
/* **************** */

void readrand()
{
   FILE *in;
   int i, tempseed[3];

   if ((in = fopen(SEEDFILE, "r")) != NULL) {
      for (i = 0; i < 3; i++) {
         fscanf(in, "%d", &tempseed[i]);
         seed16v[i] = tempseed[i];
      }
   } else {
      printf("Seed file cannot be opened for reading.\n");
      exit(1);
   }

   fclose(in);
   seed48(seed16v);
}


/* ************ */
/*   readdata   */
/* Read in data */
/* ************ */

void readdata()
{
   FILE *in;
   int i, j, k, l;
   double temp;

   finit(inifile, "CORRFUNC", &CORRFUNC, 1);
   finit(inifile, "MAXITERS", &MAXITERS, 1);
   finit(inifile, "SKIP", &SKIP, 1);
   finit(inifile, "REPS", &REPS, 1);
   STARTCOMP = (MAXITERS - 1) / 2;

   finit(inifile, "N", &N, 1);
   finit(inifile, "M", &M, 1);
   finit(inifile, "T", &T, 1);
   finit(inifile, "S", &S, 1);
   finit(inifile, "W", &W, 1);
   finit(inifile, "PM", &PM, 1);
   finit(inifile, "PS", &PS, 1);
   finit(inifile, "PT", &PT, 1);
   finit(inifile, "NSV", &NSV, 1);
   finit(inifile, "NBETA", &NBETA, 1);
   TOTBETA = NBETA + 2 * W * NSV;

   adjmat.dim(W, W);
   mx.dim(M, PM);
   sx.dim(S, PS);
   tx.dim(T, PT);
   X.dim(N, NBETA);
   Xs.dim(N, NSV);
   
   pppval.length(N);

   logolddet.length(W);
   mbeta.length(NBETA);
   mu.length(N);
   mw.length(N);
   taume.length(M);

   mbetas.length(W);
   mbetas2.length(W);
   spatbinv.dim(W, W);
   spatbinv2.dim(W, W);
   lspatbinv.dim(W);
   lspatbinv2.dim(W);
   ztz.length(W);
   xbar.length(NBETA);
   meanmbeta.length(NBETA);
   meanmbetas.length(W);
   meantaume.length(M);
   meanmbetas2.length(W);

   c.length(M);
   d.length(M);
   beta0.length(NBETA);
   sigbinvb0.length(NBETA);
   sigbinv.dim(NBETA, NBETA);

   numtypes = new int *** [W];
   starts = new int [W];
   startt = new int [W];
   stops = new int [W];
   stopt = new int [W];
   wss = new int [W];
   wst = new int [W];

   mm = new int [N];
   ss = new int [N];
   tt = new int [N];
   vs = new int [N];
   vt = new int [N];
   wsi = new int [S];
   numtype = new int [M];
   idx = new int * [W];
   
   mz = new Matrix [W];
   meanz = new Matrix [W];

   // initialize counters
   for(i = 0; i < M; i++) {
      numtype[i] = 0;
   }
   
   for(i = 0; i < W; i++) {
      startt[i] = T;
      stopt[i] = 0;
      starts[i] = S;
      stops[i] = 0;
   }
   
   if((in = fopen(infile, "r")) != NULL) {

      // obs data; must come in sorted by meas type
      printf("observation data \n");
      for(i = 0; i < N; i++) {
         fscanf(in, " %lf %d %d %d", &mw[i], &mm[i], &ss[i], &tt[i]);
         numtype[mm[i]]++;
      }

      // time data
      printf("time data \n");
      for(i = 0; i < T; i++) {
         for(j = 0; j < PT; j++) {
            fscanf(in, "%lf ", &tx[i][j]);
         }
      }
      
      cout << tx << endl << endl;

      // site data
      printf("site data \n");
      for(i = 0; i < S; i++) {
         for(j = 0; j < PS; j++)
            fscanf(in, "%lf ", &sx[i][j]);
         fscanf(in, "%d", &wsi[i]);
         if(i < starts[wsi[i]])
            starts[wsi[i]] = i;
         if(i > stops[wsi[i]])
            stops[wsi[i]] = i;
      }
      
      cout << sx << endl << endl;

      // pass through obs data again updating counters

      for(i = 0; i < N; i++) {
         if(tt[i] < startt[wsi[ss[i]]])
            startt[wsi[ss[i]]] = tt[i];
         if(tt[i] > stopt[wsi[ss[i]]])
            stopt[wsi[ss[i]]] = tt[i];
      }

      for(i = 0; i < W; i++) {
         wst[i] = stopt[i] - startt[i] + 1;
         idx[i] = new int [wst[i] + 1];
         for(j = 0; j < wst[i] + 1; j++)
            idx[i][j] = 0;
         wss[i] = stops[i] - starts[i] + 1;
         mz[i].dim(wst[i], wss[i]);
         meanz[i].dim(wst[i], wss[i]);
         numtypes[i] = new int ** [wst[i]];
         for(j = 0; j < wst[i]; j++) {
            numtypes[i][j] = new int * [wss[i]];
            for(k = 0; k < wss[i]; k++) {
               numtypes[i][j][k] = new int [M];
               for(l = 0; l < M; l++) {
                  numtypes[i][j][k][l] = 0;
               }
            }
         }
      }
       
      // one last pass through obs data

      for(i = 0; i < N; i++) {
         vs[i] = ss[i] - starts[wsi[ss[i]]];
         vt[i] = tt[i] - startt[wsi[ss[i]]];
         numtypes[wsi[ss[i]]][vt[i]][vs[i]][mm[i]]++;
         idx[wsi[ss[i]]][tt[i] - startt[wsi[ss[i]]] + 1]++;
      }

      for(i = 0; i < W; i++) {
         if(i > 0) {
            idx[i][0] = idx[i-1][wst[i-1]];
         }
         for(j = 1; j < wst[i] + 1; j++) {
            idx[i][j] += idx[i][j-1];
         }
      }

      // meas-type data
      printf("meas data \n");
      for(i = 0; i < M; i++) {
         for(j = 0; j < PM; j++)
            fscanf(in, "%lf ", &mx[i][j]);
      }
      cout << mx << endl << endl;

      // adjacency matrix
      for(i = 0; i < W; i++) {
         for(j = 0; j < W; j++) {
            fscanf(in, "%lf ", &adjmat[i][j]);
         }
      }
      cout << adjmat << endl << endl;

   } else {
      printf("File cannot be opened for reading.\n");
      exit(1);
   }
   fclose(in);

   // prior params
   finit(inifile, "aphi", &aphi, 1);
   finit(inifile, "bphi", &bphi, 1);
   finit(inifile, "aphib", &aphib, 1);
   finit(inifile, "bphib", &bphib, 1);
   finit(inifile, "WIDTH", &WIDTH, 1);
   finit(inifile, "WIDTHDIV", &WIDTHDIV, 1);
   finit(inifile, "rhomean", &rhomean, 1);
   finit(inifile, "rhoprec", &rhoprec, 1);
   finit(inifile, "rholeft", &rholeft, 1);
   finit(inifile, "rhoright", &rhoright, 1);
   finit(inifile, "a", &a, 1);
   finit(inifile, "b", &b, 1);
   finit(inifile, "ab", &ab, 1);
   finit(inifile, "bb", &bb, 1);
   finit(inifile, "ab2", &ab2, 1);
   finit(inifile, "bb2", &bb2, 1);
   finit(inifile, "c", c);
   finit(inifile, "d", d);
   finit(inifile, "beta0", beta0);
   finit(inifile, "sigbinv", sigbinv);
   cout << beta0 << endl << endl;
   cout << sigbinv << endl << endl;

   printf(" aphi %7.4f bphi %7.4f a %7.4f b %7.4f ab %7.4f bb %7.4f ab2 %7.4f bb2 %7.4f c[%d] %7.4f d[%d] %7.4f c[%d] %7.4f d[%d] %7.4f \n",
          aphi, bphi, a, b, ab, bb, ab2, bb2, 0, c[0], 0, d[0], 1, c[1], 1, d[1]);

}


/* ********************* */
/*       calcstuf        */
/* One-time computations */
/* ********************* */

void calcstuf()
{
   int i, j, k;
   double maxdist, mindist, tol = 1e-4;

   pppval = 0.0;

   // Construct X, not centering continuous covariates
   xbar = 0.0;
   for (i = 0; i < N; i++) {
      for (j = 0; j < PM; j++) {
         // don't center intercepts
         X[i][j] = mx[mm[i]][j];
      }
      for (j = 0; j < PS; j++) {
         k = j + PM;
         X[i][k] = sx[ss[i]][j];
         xbar[k] += (X[i][k] - xbar[k]) / (i + 1.0); 
      }
      for (j = 0; j < PT; j++) {
         k = j + PM + PS;
         X[i][k] = tx[tt[i]][j];
         xbar[k] += (X[i][k] - xbar[k]) / (i + 1.0); 
      }

   }

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

   sigbinvb0 = sigbinv % beta0;

   // make separate matrix of squares and cross products for each watershed
   dsq = new Matrix [W];

   for (k = 0; k < W; k++) {
      dsq[k].dim(wss[k], wss[k]);
      mindist = 500.0;
      maxdist = 0.0;
      for (i = 0; i < wss[k]; i++)   
         for (j = 0; j < wss[k]; j++) {
            dsq[k][i][j]
               = 3956 * haversine(sx[i + starts[k]][0], sx[i + starts[k]][1],
                                  sx[j + starts[k]][0], sx[j + starts[k]][1])
                            / 10.0;
            dsq[k][i][j] *= dsq[k][i][j];

            if ((dsq[k][i][j] < tol) && (i != j))  dsq[k][i][j] = tol;

            if ((dsq[k][i][j] < mindist) && (i != j))  mindist = dsq[k][i][j];
            if (dsq[k][i][j] > maxdist)  maxdist = dsq[k][i][j];
         }

      printf("mindist %7.4f maxdist %7.4f \n", mindist, maxdist);
   }

   for (i = 0; i < S; i++) {
      for (j = 0 ; j < PS; j++) {
         sx[i][j] -= xbar[j + PM]; 
      }
   }

   for (i = 0; i < T; i++) {
      for (j = 0 ; j < PT; j++) {
         tx[i][j] -= xbar[j + PM + PS]; 
      }
   }

   for (i = 0; i < NBETA; i++)
      printf("xbar[%d] %7.4f \n", i, xbar[i]);

   // construct Xs
   for (i = 0; i < N; i++) {
      Xs[i][0] = tx[tt[i]][0];
   }

   // dimension the spatial covariance matrices
   lspat = new Matrixlt [W];
   lspatinv = new Matrixlt [W];
   lar1 = new Matrixlt [W];
   lar1inv = new Matrixlt [W];

   for (i = 0; i < W; i++) {
      lspat[i].dim(wss[i]);
      lspatinv[i].dim(wss[i]);
      lar1[i].dim(wst[i]);
      lar1inv[i].dim(wst[i]);
   }

}


/* ******* */
/* initchn */
/* ******* */

void initchn()
{
   double x;
   int j, l;
   
   finit(inifile, "mbeta", &x, 1);
   mbeta = x;
   finit(inifile, "mbetas", &x, 1);
   mbetas = x;
   finit(inifile, "mbetas2", &x, 1);
   mbetas2 = x;
   finit(inifile, "mz", &x, 1);
   for (j = 0; j < W; j++)
      mz[j] = x;
      
   finit(inifile, "taume", taume);
   finit(inifile, "tau", &tau, 1);
   finit(inifile, "taub", &taub, 1);
   finit(inifile, "taub2", &taub2, 1);
   finit(inifile, "phi", &phi, 1);
   finit(inifile, "phib", &phib, 1);
   finit(inifile, "phib2", &phib2, 1);
   finit(inifile, "rho", &rho, 1);

   // compute l

   for (l = 0; l < W; l++)
      for (j = 0; j < W; j++) {
         if (l == j) 
            spatbinv[l][j] = 1.0;
         else
            spatbinv[l][j] = -phib * adjmat[l][j];
      } 
     

   chol(spatbinv, lspatbinv);			// cholesky decomp of Sigma(b)
   logolddetbinv = 0.0;
   for (j = 0; j < W; j++) {
      logolddetbinv += log(lspatbinv[j][j]);	// actually square root of determinant
   }

   for (l = 0; l < W; l++)
      for (j = 0; j < W; j++) {
         if (l == j) 
            spatbinv2[l][j] = 1.0;
         else
            spatbinv2[l][j] = -phib2 * adjmat[l][j];
      } 
     

   chol(spatbinv2, lspatbinv2);			// cholesky decomp of Sigma(b)
   logolddetbinv2 = 0.0;
   for (j = 0; j < W; j++) {
      logolddetbinv2 += log(lspatbinv2[j][j]);	// actually square root of determinant
   }

   for (l = 0; l < W; l++) {
      // cholesky decomp of Sigma(B)
      cholanis(lspat[l], phi, dsq[l]);

      cholar1(lar1[l], rho);

      // compute sqrt of |Sigma(B)|
      logolddet[l] = 0.0;
      for (j = 0; j < wst[l]; j++)
         logolddet[l] += wss[l] * log(lar1[l][j][j]);
      for (j = 0; j < wss[l]; j++)
         logolddet[l] += wst[l] * log(lspat[l][j][j]);

   }

   // set mu

   for (j = 0; j < N; j++) {
       mu[j] = dot(X[j], mbeta)
               + tx[tt[j]][0] * mbetas[wsi[ss[j]]]
               + mbetas2[wsi[ss[j]]];
   }

   movectr1 = movectr2 = movectr3 = movectr4 = 0;

   pppctr = 0;
   pppval = 0.0;
   meanmbeta = 0.0;
   meanmbetas = 0.0;
   meanmbetas2 = 0.0;
   meantaume = 0.0;
   for (l = 0; l < W; l++)
      meanz[l] = 0.0;
   meantaub = meanphib = meantaub2 = meanphib2 = meantau = meanphi = meanrho
            = meandev = 0.0;

}


/* ******* */
/* writchn */
/* ******* */

void writchn(int iter, FILE *out, int thin)
{
   int i;

   if (iter == 0) {
      fprintf(out, "iter ");
      for (i = 0; i < NBETA; i++)
         fprintf(out, "beta%d ", i);
      for (i = 0; i <  W; i++)
         fprintf(out, "betas%d ", i);
      for (i = 0; i <  W; i++)
         fprintf(out, "betas02%d ", i);
      fprintf(out, "phi  rho tau phib taub phib2 taub2 ");
      for (i = 0; i < M; i++)
         fprintf(out, "taume%d ", i);
      fprintf(out, "\n");
   }
   
   if ((iter % thin) == 0) {
      fprintf(out, "%d ", iter);
      for (i = 0; i < NBETA; i++)
         fprintf(out, "%7.4f ", mbeta[i]);
      for (i = 0; i < W; i++)
         fprintf(out, "%7.4f ", mbetas[i] + mbeta[NBETA-1]);
      for (i = 0; i < W; i++)
         fprintf(out, "%7.4f ", mbetas2[i]);
      fprintf(out, " %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f ",
              phi, rho, tau, phib, taub, phib2, taub2);
      for (i = 0; i < M; i++)
         fprintf(out, "%7.4f ", taume[i]);
      fprintf(out, "\n");
   }

}


/* ******** */
/* genbeta  */
/* ******** */

void genbeta(int iter)
{
   int i, j, k, index1, index2;
   Vector betamean(TOTBETA), betamean2(TOTBETA), err(TOTBETA),
          wminz(N), wminz2(N);
   Matrix precmat(TOTBETA, TOTBETA);
   Matrixlt cholprec(TOTBETA);

   for (i = 0; i < N; i++) {
      // subtract off z
      wminz[i] = mw[i] - mz[wsi[ss[i]]][vt[i]][vs[i]];
      wminz2[i] = taume[mm[i]] * wminz[i];
   }

   // construct precision matrix coming from likelihood
   for (i = 0; i < N; i++) {
      for (j = 0; j < NBETA; j++) {
         for (k = 0; k < NBETA; k++)
            precmat[j][k] += X[i][j] * X[i][k] * taume[mm[i]];
      }
      index1 = NBETA + wsi[ss[i]] * NSV;
      for (j = 0; j < NSV; j++) {
         for (k = 0; k < NSV; k++) {
            precmat[index1 + j][index1 + k] += Xs[i][j] * Xs[i][k] * taume[mm[i]];
         }
         for (k = 0; k < NBETA; k++) {
            precmat[index1 + j][k] += Xs[i][j] * X[i][k] * taume[mm[i]];
         }
      }
      index2 = NBETA + W * NSV + wsi[ss[i]] * NSV;
      for (j = 0; j < NSV; j++) {
         for (k = 0; k < NSV; k++) {
            precmat[index2 + j][index2 + k] += 1.0 * 1.0 * taume[mm[i]];
         }
         for (k = 0; k < NBETA; k++) {
            precmat[index2 + j][k] += 1.0 * X[i][k] * taume[mm[i]];
         }
         for (k = 0; k < NSV; k++) {
            precmat[index2 + j][index1 + k] += 1.0 * Xs[i][k] * taume[mm[i]];
         }
      }
   }

   for (i = 0; i < TOTBETA; i++)
      for (j = i + 1; j < TOTBETA; j++)
         precmat[i][j] = precmat[j][i];

   // add priors
   for (i = 0; i < NBETA; i++)
      for (j = 0; j < NBETA; j++)
         precmat[i][j] += sigbinv[i][j];

   // this will have to be chged to accommodate NSV > 1
   for (i = 0; i < W * NSV; i++)
      for (j = 0; j < W; j++)
         precmat[NBETA + i][NBETA + j] += taub * spatbinv[i][j];

   for (i = 0; i < W * NSV; i++)
      for (j = 0; j <  W; j++)
         precmat[NBETA + W * NSV + i][NBETA + W * NSV + j] += taub2 * spatbinv2[i][j];

   for (i = 0; i < N; i++) {
      for (j = 0; j < NBETA; j++)
         betamean[j] += wminz2[i] * X[i][j];
      for (j = 0; j < NSV; j++)
         betamean[NBETA + wsi[ss[i]] * NSV + j] += wminz2[i] * Xs[i][j];
      for (j = 0; j < NSV; j++)
         betamean[NBETA + W * NSV + wsi[ss[i]] * NSV + j] += wminz2[i] * 1.0;
      for (j = 0; j < NTV; j++)
         betamean[NBETA + 2 * W * NSV + tt[i] * NTV + j] += wminz2[i] * Xs[i][j];
   }

   for (i = 0; i < NBETA; i++) {
      betamean[i] += sigbinvb0[i];
   }

   chol(precmat, cholprec);

   // puts mean of betas in betamean2 by using chol decomp to do least sqs
   cholls(cholprec, betamean, betamean2);

   // now add errors
   backsub(cholprec, rnorm(TOTBETA, 0.0, 1.0), err);

   betamean2 += err;
   for (i = 0; i < NBETA; i++)
      mbeta[i] = betamean2[i];

   for (i = 0; i < W * NSV; i++) {
      mbetas[i] = betamean2[i + NBETA];
      mbetas2[i] = betamean2[i + NBETA + W * NSV];
   }

   for (i = 0; i < N; i++) {
      mu[i] = dot(X[i], mbeta)
              + tx[tt[i]][0] * mbetas[wsi[ss[i]]]
              + mbetas2[wsi[ss[i]]];
   }

}


/* ******* */
/*  genz   */
/* ******* */

void genz(int iter)
{
   int i, j, k, l, n;
   double oneplusrhosq, sqrttau;
   Vector sime, err, hold, muvect;
   Matrix sigphinvm;
   Matrixlt linvm, cholprec;

   for (l = 0; l < W; l++) {

      // set-up
      n = wss[l];

      sime.length(n);
      err.length(n);
      hold.length(n);
      muvect.length(n);
      sigphinvm.dim(n, n);
      linvm.dim(n);
      cholprec.dim(n);

      lspatinv[l] = inv(lspat[l]);

      // compute tau^2 rho Sigma( \phi )^{-1}
      sqrttau = sqrt(tau * rho);
      for (i = 0; i < n; i++)
         for (j = 0; j <= i; j++)
            linvm[i][j] = sqrttau * lspatinv[l][i][j];

      prodt(linvm, sigphinvm);

      oneplusrhosq = 1.0 + rho * rho;

      // time 0
      sqrttau = sqrt(tau);
      for (i = 0; i < n; i++)
         for (j = 0; j <= i; j++)
            linvm[i][j] = sqrttau * lspatinv[l][i][j];
 
      // construct diagonal of Ks^t sigyinv Ks
      for (i = 0; i < n; i++) {
         sime[i] = 0.0;
         for (j = 0; j < M; j++) {
            sime[i] += taume[j] * numtypes[l][0][i][j];
         }
      }

      // sime is diagonal of inverse of measurement error covariance matrix
      // linvm is lower tri right sqrt of inverse of spatial covariance matrix
      // cholprec is right sq root of total precision matrix
      cholsum(sime, 1.0, linvm, cholprec); 

      // construct Ks^t sigyinv Y - X beta
      for (i = idx[l][0]; i < idx[l][1]; i++) {
         muvect[ss[i] - starts[l]] += taume[mm[i]] * (mw[i] - mu[i]);
      }

      // now add tau^2 Sigma(phi) ^{-1} rho z_1
      muvect += sigphinvm % mz[l][1];

      // solve normal eqns using chol decomp
      cholls(cholprec, muvect, mz[l][0]);
      if(isnan(mz[0][0][0])) {
         cout << "error 1";
         exit(0);
      }

      // This has put conditional mean of latent data in mz; add errors
      backsub(cholprec, rnorm(n, 0.0, 1.0), err);	// solve lower triangular system

      for (i = 0; i < n; i++) {
         mz[l][0][i] += err[i];
      }

      // times 1 to T-2
      sqrttau = sqrt(tau *  oneplusrhosq);
      for (i = 0; i < n; i++)
         for (j = 0; j <= i; j++)
            linvm[i][j] = sqrttau * lspatinv[l][i][j];

      for (k = 1; k < wst[l] - 1; k++) {

         // construct diagonal of Ks^t sigyinv Ks
         for (i = 0; i < n; i++) {
            sime[i] = 0 ;
            for (j = 0; j < M; j++)
               sime[i] += numtypes[l][k][i][j] * taume[j];
         }

         // sime is diagonal of inverse of measurement error covariance matrix
         // linvm is lower tri right sqrt of inverse of spatial covariance matrix
         // cholprec is right sq root of total precision matrix
         cholsum(sime, 1.0, linvm, cholprec);

         // construct Ks^t sigyinv Y - X beta
         muvect = 0.0;
         for (i = idx[l][k]; i < idx[l][k + 1]; i++) {
            muvect[ss[i] - starts[l]] += taume[mm[i]] * (mw[i] - mu[i]);
         }

         // now add tau^2 Sigma(phi) ^{-1} rho (z_(k - 1) + z_(k + 1))
         for (i = 0; i < n; i++)
            hold[i] = (mz[l][k - 1][i] + mz[l][k + 1][i]) / oneplusrhosq;

         muvect += sigphinvm % hold;

         // solve normal eqns using chol decomp
         cholls(cholprec, muvect, mz[l][k]);

         // This has put conditional mean of latent data in mz; add errors
         backsub(cholprec, rnorm(n, 0.0, 1.0), err);	// solve lower triangular system

         for (i = 0; i < n; i++) {
            mz[l][k][i] += err[i];
         }

      }

      // time T-1
      sqrttau = sqrt(tau);
      for (i = 0; i < n; i++)
         for (j = 0; j <= i; j++)
            linvm[i][j] = sqrttau * lspatinv[l][i][j];

      // construct diagonal of Ks^t sigyinv Ks */
      for (i = 0; i < n; i++) {
         sime[i] = 0 ;
         for (j = 0; j < M; j++)
            sime[i] += numtypes[l][wst[l] - 1][i][j] * taume[j];
      }

      // sime is diagonal of inverse of measurement error covariance matrix
      // linvm is lower tri right sqrt of inverse of spatial covariance matrix
      // cholprec is right sq root of total precision matrix
      cholsum(sime, 1.0, linvm, cholprec);

      // construct Ks^t sigyinv Y - X beta
      muvect = 0.0;
      for (i = idx[l][wst[l] - 1]; i < idx[l][wst[l]]; i++)
         muvect[ss[i] - starts[l]] += taume[mm[i]] *  (mw[i] - mu[i]) ; 

      // now add tau^2 Sigma(phi) ^{-1} rho z_1
      muvect += sigphinvm % mz[l][wst[l] - 2];

      // solve normal eqns using chol decomp
      cholls(cholprec, muvect, mz[l][wst[l] - 1]);

      // This has put conditional mean of latent data in mz; add errors
      backsub(cholprec, rnorm(n, 0.0, 1.0), err);	// solve lower triangular system

      for (i = 0; i < n ; i++) {
         mz[l][wst[l] - 1][i] += err[i];
      }

   }   // end of loop over all watershed

}


/* ********* */
/*   gentau  */
/* ********* */

void gentau(int iter)
{
   int i, n;
   double sumztz;

   sumztz = 0.0;
   n = 0;
   for (i = 0; i < W; i++) {
      lar1inv[i] = inv(lar1[i]);
      ztz[i] = compztz(lar1inv[i], lspatinv[i], mz[i]);
      sumztz += ztz[i];
      n += wss[i] * wst[i];
   }

   if (1) {
      tau = rgamma(a + n / 2.0, b + (1.0 - rho * rho) * sumztz / 2.0);
   }

   // note: we exit with ztz also computed */

}


/* ********** */
/*   gentaub  */
/* ********** */

/* note: we enter with spatbinv already computed */

void gentaub(int iter)
{
   Vector hold;

   if(1) {
      hold = spatbinv % mbetas;
      ztzb = dot(mbetas, hold);
      taub = rgamma(ab + W / 2.0, bb + ztzb / 2.0);
   }

   if(1) {
      hold = spatbinv2 % mbetas2;
      ztzb2 = dot(mbetas2, hold);
      taub2 = rgamma(ab2 + W / 2.0, bb2 + ztzb2 / 2.0 );
   }

   // note: we exit with ztzb also computed

}


/* ********* */
/*  gentaume */
/* ********* */

/* note: we enter with l, linv, and olddet already computed */

void gentaume(int iter)
{
   int i;
   double meresids;
   Vector wtw(M);

   for (i = 0; i < N; i++) {
      // subtract off spatially varying term involving time
      meresids = mw[i] - mu[i]
                 - mz[wsi[ss[i]]][vt[i]][vs[i]];
      wtw[mm[i]] += meresids * meresids;
   }

   for (i = 0; i < M; i++)			// for each meas type
      taume[i] = rgamma(c[i] + numtype[i] / 2.0, d[i] + wtw[i] / 2.0);

}


/* ******* */
/* genphi  */
/* ******* */

/* note: we enter with l, linv, olddet, zminmu, and ztz already computed */

void genphi(int iter)
{
   // Use Metrop-Hastings to generate new value of corr parm phi
   // no longer uses lgam func from CEPHES library, and requires linking in 
   // cprob.a from ~/libraries/cprob

   int i, info, j, k, l, moved;
   double alpha, left, lognewdet, newllik, newphi, newztz, oldllik, oldphi,
          oneminrhosq, ratio, right, width = WIDTH / WIDTHDIV, u;
   Vector hold;
   Matrixlt *newlspat, *newlspatinv;

   newlspat = new Matrixlt [W];
   newlspatinv = new Matrixlt [W];

   oneminrhosq = 1.0 - rho * rho;
   oldllik = 0.0;
   for (i = 0; i < W; i++) {
      oldllik +=  -logolddet[i] - oneminrhosq * tau * ztz[i] / 2.0;
      newlspat[i].dim(wss[i]);
      newlspatinv[i].dim(wss[i]);
   }
   oldphi = phi;
   moved = 0;

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

      // generate new candidate, uniform
      info = 1;
      while (info) {
         // keep generating new candidates until get one that
         // produces pos def Sigma(B)
         left = MAX(oldphi - width, aphi);
         right = MIN(oldphi + width, bphi);
         newphi = left + drand48() * (right - left);
         for (l = 0; l < W; l++) {
            info = cholanis(newlspat[l], newphi, dsq[l]);
            if (info)
               printf("info %d iter %d phi %7.4f oldphi %7.4f aphi %7.4f bphi %7.4f \n",
                      info, iter, newphi, oldphi, aphi, bphi );
         }
      }

      newllik = 0.0;
      for (l = 0; l < W; l++) {
         newlspatinv[l] = inv(newlspat[l]);
         newztz = compztz(lar1inv[l], newlspatinv[l], mz[l]);
         lognewdet = 0.0;
         for (j = 0; j < wst[l]; j++)
            lognewdet += wss[l] * log(lar1[l][j][j]);
         for (j = 0; j < wss[l]; j++)
            lognewdet += wst[l] * log(newlspat[l][j][j]);
         newllik += -lognewdet - oneminrhosq * tau * newztz / 2.0;
      }

      // ratio is the parts of the log of the MH acceptance ratio that come
      // from the prior and the cand gen density
      ratio = log((right - left) / 
                  (MIN(newphi + width, bphi) - MAX(newphi - width, aphi)));

      alpha = exp(ratio + newllik - oldllik);

      u = drand48();

      if (iter % 250 == 0) {
         printf("ratio %7.4f newllik %7.4f oldllik %7.4f \n", ratio, newllik, oldllik);
         printf("iter %d k %d alpha %7.4f u %7.4f newphi %7.4f oldphi %7.4f aphi %7.4f bphi %7.4f \n",
                iter, k, alpha, u, newphi, oldphi, aphi, bphi);
      }

      if (u < alpha) {
         // move accepted
         moved = 1;
         oldphi = newphi;
         for (l = 0; l < W; l++) {
            logolddet[l] = lognewdet; 
            ztz[l] = newztz;      
            for (i = 0; i < wss[l]; i++)
               for (j = 0; j <= i; j++) {
                  lspat[l][i][j] = newlspat[l][i][j];
                  lspatinv[l][i][j] = newlspatinv[l][i][j];
               }
         }
      }

   }

   phi = oldphi;
   if (moved)
      movectr1++;

   delete [] newlspat;
   delete [] newlspatinv;

}   


/* ******* */
/* genphib  */
/* ******* */

void genphib(int iter)
{
   int j, k, l, moved;
   double alpha, left, lognewdetbinv, newllik, newztzb, newphi,
          oldllik, oldphi, ratio, right, u;
   Vector hold(W);
   Matrix newspatbinv(W, W);
   Matrixlt newlinv(W);

   oldllik = logolddetbinv - taub * ztzb / 2.0;
   oldphi = phib;
   moved = 0;

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

      // generate new candidate, uniform
      left = MAX(oldphi - WIDTH, aphib);
      right = MIN(oldphi + WIDTH, bphib);
      newphi = left + drand48() * (right - left);

      // compute l
      for (l = 0; l < W; l++)
         for (j = 0; j < W; j++) {
            if (l == j)
               newspatbinv[l][j] = 1.0;
            else
               newspatbinv[l][j] = -newphi * adjmat[l][j];
         }
      chol(newspatbinv, newlinv);		// cholesky decomp of Sigma(b)inv
      lognewdetbinv = 0.0;
      for (j = 0; j < W; j++) {
         // actually square root of determinant
         lognewdetbinv += log(newlinv[j][j]);
      }

      hold = newspatbinv % mbetas;
      newztzb = dot(mbetas, hold);

      newllik =  lognewdetbinv - taub * newztzb / 2.0;

      // ratio is the parts of the log of the MH acceptance ratio that come
      // from the prior and the cand gen density
      ratio = log((right - left) /
                  (MIN(newphi + WIDTH, bphib) - MAX(newphi - WIDTH, aphib)));

      alpha = exp(ratio + newllik - oldllik);

      u = drand48();

      if (iter % 250 == 0)
         printf("iter %d k %d alpha %7.4f u %7.4f newphib %7.4f oldphib %7.4f aphi %7.4f bphi %7.4f \n",
                iter, k, alpha, u, newphi, oldphi, aphi, bphi);

      if (u < alpha) {
         // move accepted
         moved = 1;
         oldphi = newphi;
         logolddetbinv = lognewdetbinv;
         ztzb = newztzb;
         spatbinv = newspatbinv;
         lspatbinv = newlinv;
      }

   }

   phib = oldphi;
   if (moved)
      movectr2++;

   // now the same for phib2
   oldllik = logolddetbinv2 - taub2 * ztzb2 / 2.0;
   oldphi = phib2;

   moved = 0;

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

      // generate new candidate, uniform 
      left = MAX(oldphi - WIDTH, aphib);
      right = MIN(oldphi + WIDTH, bphib);
      newphi = left + drand48() * (right - left);

      // compute l
      for (l = 0; l < W; l++)
         for (j = 0; j < W; j++) {
            if (l ==j)
               newspatbinv[l][j] = 1.0;
            else
               newspatbinv[l][j] = -newphi * adjmat[l][j];
         }

      chol(newspatbinv, newlinv);		// cholesky decomp of Sigma(b)
      lognewdetbinv = 0.0;
      for (j = 0; j < W; j++) {
         // actually square root of determinant
         lognewdetbinv += log(newlinv[j][j]);
      }

      hold = newspatbinv % mbetas2;
      newztzb = dot(mbetas2, hold);

      newllik =  lognewdetbinv - taub2 * newztzb / 2.0;

      // ratio is the parts of the log of the MH acceptance ratio that come
      // from the prior and the cand gen density
      ratio = log((right - left) /
                  (MIN(newphi + WIDTH, bphib) - MAX(newphi - WIDTH, aphib)));

      alpha = exp(ratio + newllik - oldllik);

      u = drand48();

      if (iter % 250 == 0)
         printf("iter %d k %d alpha %7.4f u %7.4f newphib %7.4f oldphib %7.4f aphi %7.4f bphi %7.4f \n",
                iter, k, alpha, u, newphi, oldphi, aphi, bphi);

      if (u < alpha) {
         // move accepted
         moved = 1;
         oldphi = newphi;
         logolddetbinv2 = lognewdetbinv;
         ztzb2 = newztzb;
         spatbinv2 = newspatbinv;
         lspatbinv2 = newlinv;
      }

   }

   phib2 = oldphi;
   if (moved)
      movectr4++;

}


/* ******* */
/* genrho  */
/* ******* */

/* Note: we enter with lspat and lspatinv computed */

void genrho(int iter)
{
   double alpha2, denom, gewekemean, gewekestd, numer, rhohat, rhonew,
          rhonewsq, rhosq, u;
   Vector inprods(W);
   Matrix halves;

   int i, j, k, l, la, lb ;

   if (1) {
      numer = denom = 0.0;

      for (l = 0; l < W; l++) {
         halves.dim(wst[l], wss[l]);
         for (i = 0 ; i < wst[l]; i++) {
            halves[i] = lspatinv[l] % mz[l][i];
         }

         for (i = 1; i < wst[l] ; i++) {
            numer += dot(halves[i-1], halves[i]);
            denom += dot(halves[i-1], halves[i-1]);   
         }

         inprods[l] = dot(halves[0], halves[0]);
      }

      // generate candidate rhonew from truncated normal using geweke function
      gewekemean = (tau * numer + rhoprec * rhomean) /  
                   (tau * denom + rhoprec);
      gewekestd = 1.0 / sqrt(tau * denom + rhoprec);
      rhonew = rtnorm(gewekemean, gewekestd, rholeft, rhoright, false, false);

      // Metrop Hastings acceptance
      rhosq = rho * rho;
      rhonewsq = rhonew * rhonew;

      alpha2 = 0.0;

      for (l = 0; l < W; l++)
         alpha2 += (log(1.0 - rhonewsq) - log(1.0 - rhosq)) * wss[l] / 2.0
                   - tau * inprods[l] * (rhosq - rhonewsq) / 2.0;

      u = drand48();

      if (u < exp(alpha2)) {
         rho = rhonew;
         for (i = 0; i < W; i++) {
            cholar1(lar1[i], rho);
            logolddet[i] = 0.0;
            for (j = 0; j < wst[i]; j++)
               logolddet[i] += wss[i] * log(lar1[i][j][j]);
            for (j = 0; j < wss[i]; j++)
               logolddet[i] += wst[i] * log(lspat[i][j][j]);
         }
         movectr3++;
      }
     
   }

}


/* **************** */
/*    calcpred()    */
/* **************** */

void calcpred(int iter)
{
   int i;
   double normval, resid, temp;
   Vector logsqrttau(M), stddev(M);

   for (i = 0; i < M; i++) {
      stddev[i] = 1.0 / sqrt(taume[i]);
      logsqrttau[i] = 0.5 * (log(taume[i]) - log(2.0 * PI));
   }
   if (iter > STARTCOMP) {
      pppctr++;
      deviance = 0.0;
      for (i = 0; i < N; i++) {
         normval = rnorm(0.0, 1.0);
         temp = mu[i] + mz[wsi[ss[i]]][vt[i]][vs[i]];
         if ((temp + stddev[mm[i]] * normval) > mw[i])
            pppval[i]++;
         resid = mw[i] - temp;
         deviance += logsqrttau[mm[i]]
                     - 0.5 * taume[mm[i]] * resid * resid;
      }
      deviance *= -2.0;
      for (i = 0; i < NBETA; i++)
         meanmbeta[i] += (mbeta[i] - meanmbeta[i]) / pppctr;
      for (i = 0; i < W; i++) {
         meanmbetas[i] += (mbetas[i] - meanmbetas[i]) / pppctr;
         meanmbetas2[i] += (mbetas2[i] - meanmbetas2[i]) / pppctr;
      }
      for (i = 0; i < M; i++)
         meantaume[i] += (taume[i] - meantaume[i]) / pppctr;
      for (i = 0; i < W; i++)
         meanz[i] += (mz[i] - meanz[i]) / pppctr;
      meantaub += (taub - meantaub) / pppctr;
      meanphib += (phib - meanphib) / pppctr;
      meantaub2 += (taub2 - meantaub2) / pppctr;
      meanphib2 += (phib2 - meanphib2) / pppctr;
      meantau += (tau - meantau) / pppctr;
      meanphi += (phi - meanphi) / pppctr;
      meanrho += (rho - meanrho) / pppctr;
      meandev += (deviance - meandev) / pppctr;
   }

}


/* **************** */
/*   calcdevmean    */
/* **************** */

double calcdevmean()
{
   int i;
   double resid, retval, temp;
   Vector logsqrttau(M);

   retval = 0.0;
   for (i = 0; i < M; i++) {
      logsqrttau[i] = 0.5 * (log(meantaume[i]) - log(2.0 * PI));
   }
     
   for (i = 0; i < N; i++) {
      mu[i] = dot(meanmbeta, X[i]);
      temp = mu[i] + tx[tt[i]][0] * meanmbetas[wsi[ss[i]]] +
             + meanmbetas2[wsi[ss[i]]]
             + meanz[wsi[ss[i]]][vt[i]][vs[i]];
      resid = mw[i] - temp;
      retval += logsqrttau[mm[i]] - 0.5 * meantaume[mm[i]] * resid * resid;
   }
   retval *= -2.0;

   return(retval);
}


/* **************** */
/*     updtseed     */
/* Update seed file */
/* **************** */

void updtseed()
{
   FILE *out;
   int i;
   unsigned short *myptr;

   myptr = seed48(seed16v);
   out = fopen(SEEDFILE, "w");
   for (i = 0; i < 3; i++)
      fprintf(out, "%10u", *(myptr + i));
   fprintf(out, "\n");
   fclose(out);
}


/* *********** */
/*   cholanis  */
/* *********** */

/* Choleski decomposition of Sigma(B), as defined by covariance func */
/* lower triangular factorization will be returned in l.*/
/* n is dimension of Sigma(B) */
/* Uses sums of squares and cross-products matrix computed in calcstuf and
   defined as a global variable. */

int cholanis(Matrixlt &l, double phi, Matrix &dsq)
{
   int i, j, k, n = l.nrow();
   double sum;

   for(i = 0; i < n; i++) {
      for(j = 0; j <= i; j++) {
         sum = covfunc(phi, dsq[j][i], CORRFUNC);
         for(k = 0; k < j; k++) {
            sum -= l[i][k] * l[j][k];
         }
         if(i != j) {
            l[i][j] = sum / l[j][j];
         } else {
            if(sum <= 0.0)
               return 1;
            l[i][i] = sqrt(sum);
         }
      }
   }
   
   return 0;
   
}


/* *********** */
/*   covfunc   */
/* *********** */

/* temp is squared distance, phi is corr param */

double covfunc(double phi, double temp, int corrtype)
{
   double parm, retval, tol = 1.0e-12;

   switch (corrtype) {
      case 1:				// exponential
      {
         parm = phi * sqrt(temp);
         if (parm > tol)
            retval = exp(-parm);
         else
            retval = 1.0;
         break;
      }
      case 2:				// Gaussian
      {
         parm = phi * temp;
         if (parm > tol)
            retval = exp(-parm);
         else
            retval = 1.0;
         break;
      }
      case 3:				// Cauchy
      {
         parm = phi * temp;
         retval =  1.0 / (1.0 + parm);
         break;
      }
      case 4:				// Spherical
      {
         if (sqrt(temp) > 1.0 / phi)
            retval = 0.0;
         else {
            parm = phi * sqrt(temp);
            retval = 0.5 * (parm * parm * parm - 3 * parm + 2.0);
         }
         break;
      }
   }

   return retval;

}


/* ********** */
/*   cholar1  */
/* ********** */

/* Choleski decomposition of AR1 corr mat */
/* rho is corr coef */

void cholar1(Matrixlt &l, double rho)
{

   int i, j, n = l.nrow();
   double s1minrhosq;
   Vector powrho(n);

   s1minrhosq = sqrt(1.0 - rho * rho);

   powrho[0] = 1.0;
   for (i = 1; i < n; i++)
      powrho[i] = rho * powrho[i-1] ;

   for (i = 0; i < n; i++) {
      l[i][0] = powrho[i];
      for (j = 1; j <= i; j++)
         l[i][j] = s1minrhosq * powrho[i - j];
   }

}


/* *********** */
/*   compztz   */
/* compute ztz */
/* *********** */

// computes quadratic form [ mz^T times its prec mat times mz] exploiting
// identities from Harville

double compztz(Matrixlt &lar1inv, Matrixlt &lspatinv, Matrix &mz)
{
   Matrix hold;

   // multiply T x S matrix times lspatinv^T
   hold = lar1inv % mz % t(lspatinv);

   // return tr(hold^T hold)
   return sum(hold * hold);

}