/* MCMCcoal.c
   Markov chain Monte Carlo coalescent program for population genetics data.
   It currently implements the JC69 model for analysing sequence data from multiple 
   loci.
   The input includes the number of species, the species tree, species 
   divergence times as well as \theta for current and ancestral populations.
   The program then simulates a genealogical tree and evolves sequences along 
   it under the JC69 model.  Different gene trees are thus used for different 
   replicates.

   Copyright by Ziheng Yang, since July 2002

   cc -o MCMCcoal -O3 MCMCcoal.c tools.c -lm
   cc -o MCMCcoal -march=athlon -mcpu=athlon -O4 -funroll-loops -fomit-frame-pointer -finline-functions MCMCcoal.c tools.c -lm

   cc -o MCMCcoal -mcpu=G5 -O4 -funroll-loops -fomit-frame-pointer -finline-functions MCMCcoal.c tools.c -lm

   cl -O2 MCMCcoal.c tools.c
*/

#ifdef MH
#define MetropolisHastings
#endif

#define NSPECIES      10       /* max # of species */
#define NS            500      /* max # of sequences */
#define NBRANCH       NS*2-2
#define MAXNSONS      10
#define NGENE         2000      /* max # of loci */
#define LSPNAME       30        /* # characters in sequence names */

#include "paml.h"
struct CommonInfo {
   char *z[2*NS-1], *spname[NS], seqf[256], cleandata, oldconP[NS*2-1];
   int ns, ls, lgene[1], model, clock, mcmc;
   int seqtype, ngene, *pose, npatt, np, readpattern;
   int ntime, ncatG, ncode, print, fix_rgene, posG[1];
   double alpha, pi[4], piG[1][4], *rates, rgene[1];
   double *conP;  /* not used */
   int    sconP, curconP;
   double *conPin[2], *conP0, *fpatt, space[10000];
}  com;
struct TREEB {
   int  nbranch, nnode, root, branches[NBRANCH][2];
   double lnL;
}  tree;
struct TREEN {
   int father, nson, sons[2], ibranch, ipop;  /* ibranch not used? */
   double branch, age, *conP, label;       /* age is uptodate but not branch */
   char fix_age;                       /* not used */
}  *nodes, *gnodes[NGENE], nodes_t[2*NS-1]; 
/*  nodes is a pointer, gnodes holds the gene trees, nodes_t is temp space
    gnodes is not allocated for the simulation program 
*/

/* sptree.nseqsp is working space, real thing in data.nseqsp[locus].
   npop is the number of theta's.  
   pops[] holds the node numbers in the species tree of the npop populations;
   the ancestral populations should be after the recent populations.  
   pop_pop_table[i][j] = 1 if pop i can coalescence in pop j and 0 if otherwise.
   son_age is needed when gaps instead of ages are used for speciation times
*/
struct SPECIESTREE {
   int nbranch, nnode, root, nspecies, nseqsp[NSPECIES]; 
   int npop, pops[NSPECIES*2-1], pop_pop_table[NSPECIES*2-1][NSPECIES*2-1];
   struct TREESPN {
      char name[LSPNAME*2];
      int father, nson, sons[2], son_age;
      double age, theta;
   } nodes[2*NSPECIES-1];
}  sptree;


#ifdef MetropolisHastings

struct DATA { /* locus-specific data and gene tree information, used in pGB_gammatau */
   int maxns, ngene, lgene[NGENE];
   int ns[NGENE], nseqsp[NGENE][NSPECIES], ls[NGENE], npatt[NGENE];
   int root[NGENE+1], conP_offset[NGENE];
   char *z[NGENE][NS], cleandata[NGENE];
   double *fpatt[NGENE], lnpGBi[NGENE], lnpDi[NGENE];
}  data;

/* speciation[0] = 0 for gaps and 1 for ages for tau in the prior
   speciation[1] = 0 for gaps and 1 for ages for outputting tau 
*/
struct MCMCPARAMETERS {
   int burnin, nsample, output, usedata, saveconP;
   int print, speciation[2], moveinnode, mixtheta_change;
   int mixtheta[NS*2-1];  /* 0: no update; 1: multiply -1: divide */
   double a_gamma[NS*3-2], b_gamma[NS*3-2], finetune[5];
}  mcmc; /* control parameters */

struct COALPOP {
   int nin[NSPECIES*2-1], ncoal[NSPECIES*2-1], nout[NSPECIES*2-1];
   double tj[NSPECIES*2-1][NS-1];
}  coal; /* data for coalescent in one population */

#endif

int ReadSpeciesTree(char *datafile);
int DownSptreeSetPops(int inode);
void GetRandomGtree(int locus);
void SimulateData(void);
int GetMem(void);
void FreeMem(void);
int ResetSpeciesGeneTree(int locus);
int Coalescence1Pop(int ispecies);
int ReadSeqData(char *seqfile, FILE*fout, int cleandata);
double lnpGB_thetatau(void);
double lnpData(double lnpDi[]);
double lnpD_locus(void);
int UseLocus(int locus, int copytreeconP, int useData, int setSeqName);
int AcceptLocus(int locus, int copyconP);
int SetParameters(double x[], int set_theta, int set_tau);
double lnpriorTheta (double theta[], double thetanew[]);
double lnpriorTau (double tau[], double taunew[]);
int MCMC(FILE* fout);
int UpdateGB_InternalNode(double* lnL, double finetune);
int UpdateGB_SPR(double* lnL, double finetune);
int UpdateTheta(double x[], double finetune, double space[]);
int UpdateTau(double x[], double *lnL, double finetune, double space[]);
int mixing(double x[], double* lnL, double finetune, double space[]);
int Reset_Thetamix(double var[]);
void GraftNode(int source, int target, double age, int ipop);
int MatchGTree(void);
void printGtree (int printBlength);
void checkGtree(void);
int getab_beta(void);

extern int noisy, IncludeNodeLabel;
char timestr[32];
double OLDAGE=99;
int LASTROUND=1, debug=0, testlnL=0, totalTimes, totalSeqs, NPMat=0;

#define REALSEQUENCE
#define NODESTRUCTURE
#include "treesub.c"

int main (int argc, char*argv[])
{
#ifdef MetropolisHastings
   char datafile[128]="MCMCcoal.dat";
#else
   char datafile[128]="MCcoal.dat";
#endif
   char outfile[128]="out", VStr[32]="Version 1.1";
   FILE *fout=(FILE*)gfopen(outfile,"w");
   int i;

   printf("MCMCcoal %s\n", VStr);
   starttime();
   if(argc>1) strcpy(datafile, argv[1]);
   com.cleandata=0;
   com.clock=1; com.ncode=4; com.model=0; noisy=0;
   FOR(i,4) com.pi[i]=1./4;
   ReadSpeciesTree(datafile);

#ifdef MetropolisHastings
   com.mcmc=1;
   if(fout==NULL) error2("outfile open error");
   ReadSeqData(com.seqf, fout, com.cleandata);
   GetMem();
   MCMC(fout);
   FreeMem();
#else
   com.mcmc=0;
   SimulateData(); 
#endif
   exit(0);
}


int ReadSpeciesTree (char* datafile)
{
/* this reads the control data file for simulation.
   It reads the species tree and initializes sptree.
   Note that (sptree.nseqsp[j] > 1) adds a population.
   This also reads the gamma prior parameters mcmc.a_gamma[] and mcmc.b_gamma[].
*/
   FILE *fdata=(FILE*)gfopen(datafile,"r");
   char line[2048]={'\0'};
   int i,j, lline=2048, son0;
   double a,b, percent[2];
   FILE*fseed=(FILE*)gfopen("SeedUsed", "w");

   fscanf(fdata, "%s%d", com.seqf, &i);
   if(i<0) i=abs(2*(int)time(NULL)+1);
   fprintf(fseed, "%d\n", i); fclose(fseed);

   SetSeed(i);
   fscanf(fdata, "%d", &com.ns);  
   if((sptree.nspecies=com.ns)>NSPECIES) error2("raise NSPECIES");

   nodes=nodes_t;  /* species tree is read into the temp space nodes_t */
   for(i=0; i<NS; i++) {  /* for both species & gene trees */
      if (com.spname[i]) free(com.spname[i]);
      com.spname[i]=(char*)malloc((LSPNAME+1)*sizeof(char));
      /* FOR(j,LSPNAME) com.spname[i][j]=0; */
   }
   for(i=0; i<sptree.nspecies; i++) {
      fscanf(fdata,"%s", com.spname[i]);
      strcpy(sptree.nodes[i].name, com.spname[i]);
   }
   
   for(i=0; i<sptree.nspecies; i++) {
      fscanf(fdata,"%d",&sptree.nseqsp[i]);
      if(sptree.nseqsp[i]<1) 
         puts("need at least 1 seq from each species?");
   }
   fgets(line, lline, fdata);
   if(sptree.nspecies>1) {
      ReadaTreeN (fdata, &i, &j, 0, 1);
      OutaTreeN(F0,1,1);  FPN(F0);
   }
   else {
      tree.root=nodes[0].nson=0;  tree.nnode=1; tree.nbranch=0; 
      nodes[0].age=0;
   }
   printf("%d species: ", sptree.nspecies);
   for(i=0; i<sptree.nspecies; i++) 
      printf(" %s (%d)",com.spname[i],sptree.nseqsp[i]);

   /* copy into sptree */
   sptree.nnode=tree.nnode;  sptree.nbranch=tree.nbranch; sptree.root=tree.root;
   for(i=0; i<sptree.nspecies*2-1; i++) {
      sptree.nodes[i].father=nodes[i].father;
      sptree.nodes[i].nson=nodes[i].nson;
      for(j=0;j<sptree.nodes[i].nson;j++) sptree.nodes[i].sons[j]=nodes[i].sons[j];
      sptree.nodes[i].age=nodes[i].age=nodes[i].branch;
      sptree.nodes[i].theta=nodes[i].label;
   }

   for(i=0,sptree.npop=0; i<sptree.nspecies; i++)
      if(sptree.nseqsp[i]>1)  sptree.pops[sptree.npop++]=i;
   for(; i<2*sptree.nspecies-1; i++) sptree.nodes[i].name[0]='\0';
   DownSptreeSetPops(sptree.root);

   printf("\npopulation by population table:\n");
   for(i=0; i<2*sptree.nspecies-1; i++) for(j=0; j<2*sptree.nspecies-1; j++) 
      sptree.pop_pop_table[i][j]=0;
   for(i=0; i<2*sptree.nspecies-1; i++) {
      for(j=i; ; j=sptree.nodes[j].father) {
         if(j>=sptree.nspecies || sptree.nseqsp[j]>1)
            sptree.pop_pop_table[i][j]=1;
         if(j==sptree.root) break;
      }
   }
   for(i=0; i<2*sptree.nspecies-1; i++,FPN(F0)) {
      printf("species %2d %-10s ", i+1,sptree.nodes[i].name);
      printf("age %.5f  theta %.5f  ", sptree.nodes[i].age, sptree.nodes[i].theta);
      for(j=0; j<2*sptree.nspecies-1; j++) 
         printf("%1d", sptree.pop_pop_table[i][j]);
      if(i<sptree.nspecies && sptree.nodes[i].age) 
         printf("\a  <-- age>0??");
   }
   com.np=sptree.npop+sptree.nspecies-1;
   if(!com.mcmc) {
      printf("\n%d theta parameters (populations) in the order:", sptree.npop);
      for(i=0; i<sptree.npop; i++)
         printf(" %2d (%s)", sptree.pops[i]+1, sptree.nodes[sptree.pops[i]].name);
      if(sptree.nspecies>1) {
         printf("\n%d species divergence times in the order:", sptree.nspecies-1);
         for(i=sptree.nspecies; i<sptree.nspecies*2-1; i++)
            printf(" %2d (%s)", i+1, sptree.nodes[i].name);
      }
   }

   /* count the number of sequences in gene tree com.ns or data.ns[] */
   for(i=0,com.ns=0; i<sptree.nspecies; i++) com.ns+=sptree.nseqsp[i];
   if(com.ns>NS) error2("raise NS in source");

#ifdef MetropolisHastings
   data.maxns=com.ns;
   fscanf(fdata,"%d", &mcmc.usedata);                                fgets(line,lline,fdata);
   fscanf(fdata,"%d%d%d", &mcmc.burnin,&mcmc.output,&mcmc.nsample);  fgets(line,lline,fdata);
   FOR(i,5) fscanf(fdata,"%lf", &mcmc.finetune[i]);                  fgets(line,lline,fdata);
   FOR(i,com.np) fscanf(fdata,"%lf", &mcmc.a_gamma[i]);              fgets(line,lline,fdata);
   FOR(i,com.np) fscanf(fdata,"%lf", &mcmc.b_gamma[i]);              fgets(line,lline,fdata);


   /* process gamma priors, note speciation times if mcmc.speciation[0]==0 */
   printf("\nGamma prior: mean +- SE (95%% CI) for theta's and tau's\n");
   for(j=0;j<sptree.nspecies*2-1;j++)  sptree.nodes[j].son_age=-1;
   for(j=0;j<sptree.nspecies;j++)  sptree.nodes[j].age=0;
   for(j=0; j<sptree.nspecies-1; j++) {
      i=j+sptree.npop; 
      sptree.nodes[sptree.nspecies+j].age=mcmc.a_gamma[i]/mcmc.b_gamma[i];
   }
   if(mcmc.speciation[0]==0) { /* use gaps for prior of speciation times */
      for(j=0;j<sptree.npop;j++)
         if((i=sptree.pops[j]) >= sptree.nspecies) {
            son0=sptree.nodes[i].sons[0];
            if(!sptree.nodes[son0].nson) son0=sptree.nodes[i].sons[1];
            else if (sptree.nodes[sptree.nodes[i].sons[1]].nson)
               puts("Be careful about priors for speciation times.");
            if(son0>=sptree.nspecies) {
               sptree.nodes[i].son_age=son0;
               sptree.nodes[i].age += sptree.nodes[son0].age;
            }
         }
   }
   for(j=0;j<com.np;j++) {
      a=mcmc.a_gamma[j]; b=mcmc.b_gamma[j];
      percent[0]=PointGamma(.025,a,b);  percent[1]=PointGamma(.975,a,b);
      if(j<sptree.npop)
         sprintf(line, "theta_%s ", sptree.nodes[sptree.pops[j]].name);
      else {
         i=sptree.nspecies+j-sptree.npop;
         sprintf(line, "tau_%s ", sptree.nodes[i].name);
         if(mcmc.speciation[0]==0) { /* gaps */
            son0=sptree.nodes[i].son_age;
            if(son0>=sptree.nspecies) 
               sprintf(line, "tau_%s - tau_%s ", sptree.nodes[i].name, sptree.nodes[son0].name);
         }
      }
      printf("%-26s%9.5f +-%8.5f (%.5f, %.5f) ",line,a/b,sqrt(a)/b, percent[0],percent[1]);
      if(mcmc.speciation[0]==0 && j>=sptree.npop) {
         i=sptree.nspecies+j-sptree.npop;
         printf(" tau_%-5s %.5f", sptree.nodes[i].name, sptree.nodes[i].age);
      }
      FPN(F0);
   }
   for(j=sptree.nspecies; j<sptree.nspecies*2-1; j++)
      if(j!=sptree.root && sptree.nodes[j].age>sptree.nodes[sptree.nodes[j].father].age)
         puts("\npriors for speciation dates are conflicting");
#endif
   fclose(fdata);
   return(0);
}

int DownSptreeSetPops (int inode)
{
/* This traverse down the species tree to see which nodes are visited,
   and initializes npop, pops[], and sptree.nodes[].name.
   Populations where coalescent is possible (theta's) are collected into 
   sptree.pops[] in the post-order tree traversal.
*/
   int k,ison;

   for (k=0; k<sptree.nodes[inode].nson; k++) {
      ison=sptree.nodes[inode].sons[k];
      if(sptree.nodes[ison].nson)  
         DownSptreeSetPops(ison);
      strcat(sptree.nodes[inode].name, sptree.nodes[ison].name);
   }
   if(inode>=sptree.nspecies)
      sptree.pops[sptree.npop++]=inode;
   return(0);
}


/* used by Coalescence1Pop(). */
static int cNodeNumber, nin_G[NSPECIES*2-1], ins_G[NSPECIES*2-1][NS];


void GetRandomGtree(int locus)
{
/* generates a random gene tree in nodes[].  This does not initialize ipop.
*/
   int i;

   ResetSpeciesGeneTree(locus);
   Coalescence1Pop(sptree.root);
   /* NodeToBranch(); */
   tree.root=cNodeNumber-1;  nodes[tree.root].branch=0;  
   nodes[tree.root].father=-1;  tree.nnode=com.ns*2-1;
   FOR (i,tree.nnode) if(i!=tree.root) 
      nodes[i].branch=nodes[nodes[i].father].age-nodes[i].age;
}

int ResetSpeciesGeneTree (int locus)
{
/* Resets the species tree for Coalescence1Pop(), initializing nin_G, 
   ins_G[], ipop.
   Also resets cNodeNumber for constructing the gene tree.
   nodes[].ipop for tips in the gene tree are set already.
*/
   int is,j;

#ifdef MetropolisHastings
   com.ns=data.ns[locus];
   nodes=gnodes[locus]; 
   for(is=0; is<sptree.nspecies; is++) sptree.nseqsp[is]=data.nseqsp[locus][is];
#endif
   for(is=0; is<sptree.nspecies; is++) 
      nin_G[is] = sptree.nseqsp[is];
   for(is=0,cNodeNumber=0; is<sptree.nspecies; is++) {
      for(j=0; j<sptree.nseqsp[is]; j++,cNodeNumber++) {
         ins_G[is][j]=cNodeNumber;
         if(nin_G[is]>1) 
            sprintf(com.spname[cNodeNumber], "%s%d", sptree.nodes[is].name,j+1);
         else 
            sprintf(com.spname[cNodeNumber], "%s", sptree.nodes[is].name);
      }
   }
   for(; is<2*sptree.nspecies-1; is++) nin_G[is]=0;
   if(cNodeNumber!=com.ns)  error2("ns mismatch");
   for(j=0;j<com.ns;j++) ClearNode(j);
   return(0);
}

int Coalescence1Pop (int ispecies)
{
/* This simulates the coalescent process in population or species ispecies.
   It generates the random genealogy tree (possibly disconnected) with waiting 
   times.
   t: waiting time; T: node age
   This is used by GetRandomGtree().  
   nin_G[] and ins_G[] are set in ResetSpeciesGeneTree().
*/
   int j, k,k1,k2, father=sptree.nodes[ispecies].father;
   double t, T;

   for (k=0; k<sptree.nodes[ispecies].nson; k++)
      Coalescence1Pop(sptree.nodes[ispecies].sons[k]);
   T=sptree.nodes[ispecies].age;
   if(nin_G[ispecies]>1 && sptree.nodes[ispecies].theta==0)
      error2("theta=0 for a population");
   for (j=nin_G[ispecies]; j>1; j--,cNodeNumber++) {
      t=rndexp(sptree.nodes[ispecies].theta/(j*(j-1.)));  T+=t;
      if(ispecies!=sptree.root && T>sptree.nodes[father].age) break;

      k=(int)(j*rndu());  
      k1=ins_G[ispecies][k]; ins_G[ispecies][k]=ins_G[ispecies][j-1];
      k=(int)((j-1)*rndu());
      k2=ins_G[ispecies][k]; ins_G[ispecies][k]=cNodeNumber;

      nodes[cNodeNumber].nson=2;   nodes[cNodeNumber].ipop=ispecies;
      nodes[cNodeNumber].sons[0]=k1;
      nodes[cNodeNumber].sons[1]=k2;
      nodes[k1].father=nodes[k2].father=cNodeNumber;
      nodes[cNodeNumber].age=T;
   }
  /* outgoing lineages added to father's list */
   nin_G[ispecies]=j;
   if(ispecies==sptree.root && j>1) 
      error2("nin_Gtree > 1 at root");
   if(ispecies!=sptree.root) {
      for(k=0; k<j; k++) 
         ins_G[father][nin_G[father]++]=ins_G[ispecies][k];
   }
   return (0);
}


void printGtree (int printBlength)
{
   int i,j, ipop, ipopTrue;
   double t, tb[2];

   FOR (i,tree.nnode) if(i!=tree.root) {
      nodes[i].branch=nodes[nodes[i].father].age-nodes[i].age;
      if(nodes[i].branch<0) 
         printf("");
   }
   printf("\nns = %d  nnode = %d", com.ns, tree.nnode);
   printf("\n%7s%7s%11s (%s) %7s%7s","father","node","time","ipop","nson:","sons");
   FOR (i, tree.nnode) {
      t=nodes[i].age;
      ipop=nodes[i].ipop;
      tb[0]=sptree.nodes[ipop].age;  tb[1]=OLDAGE;
      if(ipop!=sptree.root) tb[1]=sptree.nodes[sptree.nodes[ipop].father].age;

      ipopTrue=(t>=tb[0] && t<tb[1]);
      printf ("\n%7d%7d %11.6f (%2d %s %11.6f):%7d  ",
         nodes[i].father, i, t, ipop, (ipopTrue?"true ":"false"), 
         sptree.nodes[ipop].age, nodes[i].nson);
      FOR(j, nodes[i].nson) printf (" %2d", nodes[i].sons[j]);
   }
   FPN(F0); OutaTreeN(F0,0,0); FPN(F0); OutaTreeN(F0,1,0); 
   if(printBlength) {
      FPN(F0); OutaTreeN(F0,1,1); FPN(F0); 
   }
}


int MatchGTree (void)
{
/* This tests the gene tree topology
*/
   return 
     (nodes[0].father==nodes[1].father
   && nodes[3].father==nodes[4].father 
   && nodes[nodes[3].father].father==tree.root
   && nodes[nodes[0].father].age<sptree.nodes[sptree.root].age
   && nodes[nodes[2].father].age<sptree.nodes[sptree.root].age
   && nodes[nodes[3].father].age<sptree.nodes[sptree.root].age
   );  /* P(gtree) = 0.237913 */
}



#ifndef MetropolisHastings

static int n123marks[][3]= {{1,2,3},{1,2,3},{2,3,1},{3,1,2}};

void p0124Fromb0b1 (int itree, double p[5], double b[2])
{
/* This calculates p0,p1,p2,p3,p4 for the 5 site patterns for 3 species, 
   given branch lengths b0 and b1.  b0 is the gap, and b1 is the distance 
   from the ancestor of 1 and 2 to species 1.
*/
   double e1, e2, e3;

   e1=exp(-4./3*b[1]);  e2=exp(-8./3*(b[0]+b[1]));  e3=e1*e2;  e1=e1*e1;
   p[0]      = (1. + 3*e1 +  6*e2 +  6*e3)/16;
   p[n123marks[itree][0]]  = (3. + 9*e1 -  6*e2 -  6*e3)/16;
   p[n123marks[itree][1]]=p[n123marks[itree][2]] 
      = (3. - 3*e1 +  6*e2 -  6*e3)/16;
   p[4]      = (6. - 6*e1 - 12*e2 + 12*e3)/16;
}

void PMismatch3s (void)
{
/* This calculate the tree mismatch probabilities for figure 3 of Yang
   (2002 Genetics).
*/
   int iHCG=sptree.root, iHC=3+4-iHCG; /* node numbers in sptree */
   double theta_HC=sptree.nodes[iHC].theta;
   double t_HC=sptree.nodes[iHC].age, t_HCG=sptree.nodes[iHCG].age;
   /* S: species tree; G: gene tree; E: estimated gene tree 
      G[gtree], E[gtree][etree] are counts of resolved (used) loci.
   */
   double SG, SE, GE, G[4], E[4][4], b[2], p0[5],p[5],nused, over,under; 
   int ii, nii=7, ls0[100]={200, 400, 500, 1000, 2000, 4000, 10000};
   int nr=10000000, ir, i,j, gtree=-1, n[5]={0}, etree,etree2,etree3,every=100;

   printf("\nPr{S-G mismatch} = %f from formula.\n",2./3*exp(-2*(t_HCG-t_HC)/theta_HC));
   puts("Ties in genetree are removed in the following calculation.");
   every=max2(every,nr/1000);
   for(ii=0; ii<nii; ii++) {
      com.ls=ls0[ii];
      printf("\n# sites? (Ctrl-C to break) "); fflush(stdout);
      scanf("%d", &com.ls);
      printf("%d sites, %dK replicates.\n", com.ls, nr/1000);
      FOR(i,4) G[i]=0; FOR(i,4) FOR(j,4) E[i][j]=0;
      for (ir=0,SG=SE=GE=0,nused=0; ir<nr; ir++) {
         GetRandomGtree(-1);
         if (nodes[2].father==tree.root)      gtree=(nodes[0].branch>t_HCG);/* (HC) */
         else if(nodes[0].father==tree.root)  gtree=2; /* (CG) */
         else if(nodes[1].father==tree.root)  gtree=3; /* (GH) */
         else    error2("binary tree?");
         b[1]=nodes[nodes[tree.root].sons[0]].age;
         b[0]=nodes[tree.root].age-b[1];
         p0124Fromb0b1 (gtree, p0, b);
         p[0]=p0[0]+p0[4]; p[1]=p[0]+p0[1]; p[2]=p[1]+p0[2]; p[3]=p[2]+p0[3]; p[4]=-1;
         /* matout(F0, p, 1, 4); */
//++pac         MultiNomial (com.ls, 4, p, n, NULL);
         MultiNomial2 (com.ls, 4, p, n, NULL);

         for(i=2,etree=1; i<=3; i++)   if(n[i]>n[etree]) etree=i;
         if(etree==1)      { etree2=2; etree3=3; }
         else if(etree==2) { etree2=3; etree3=1; }
         else              { etree2=1; etree3=2; }

         if(n[etree]!=n[etree2] && n[etree]!=n[etree3]) { /* exclude ties */
            nused++;
            G[gtree]++;
            E[gtree][etree]++;
            if(gtree>=2)  SG++;
            if(etree>=2)  SE++;
            if((gtree<=1 && etree!=1) || (gtree>=2 && etree!=gtree)) GE++;
         }
         if((ir+1)%every==0) {
            printf("%4.1f%% (%2d %2d %2d %2d) SG %.4f SE %.4f GE %.4f (+%.4f -%.4f) tie %.4f\n",
               (ir+1.)/nr*100, n[0],n[1],n[2],n[3], 
               SG/nused,SE/nused, GE/nused, 
               (E[0][2]+E[0][3] + E[1][2]+E[1][3])/nused,  /* T -> F */
               (E[2][1]+E[3][1])/nused,                    /* T -> F */
               (ir+1-nused)/(ir+1.));
         }
      }  /* for(ir) */
      if(G[0]+G[1]+G[2]+G[3]-nused!=0) error2("gtree counts incorrect.");
      SG/=nused; SE/=nused; GE/=nused; 
      FOR(i,4) FOR(j,4) E[i][j] /= G[i];
      FOR(i,4) G[i]/=nused;
      printf("\nfrequencies of gene trees 0123 (given ties are removed): ");
      matout2 (F0, &G[0], 1, 4, 9, 5);
      printf("transition probability matrix (gene tree 0123 -> MLtree (01)23):");
      matout2 (F0, &E[0][0], 4, 4, 9, 5);

      printf("\nThe following three should be equal:\n");
      printf("  (1) SE - SG = %.5f\n", SE-SG);
      printf("  (2) f(T0) * P0 = %.5f\n", G[0]*(E[0][2]+E[0][3]));
      over=G[0]*(E[0][2]+E[0][3])+G[1]*(E[1][2]+E[1][3]);
      under=G[2]*E[2][1]+G[3]*E[3][1];
      printf("  (3) over - under = %.5f - %.5f = %.5f\n", over,under,over-under);
      printf("\nf(2)*P2/2+f(3)*P3/2 = %.5f, which is GE - (over + under) = under.\n",
         G[2]*E[2][3]+G[3]*E[3][2]);
      printf("\nTime used: %s\n", printtime(timestr));
   }
   exit(0);
}

void SimulateData (void)
{
   char *treef="out.trees", savetree=0, makeseq=0;
   FILE *fseq=gfopen(com.seqf, "w"), *fout=gfopen(treef,"w");
   int nr, ir, i,j, variable_ns=0, nseqsp0[NSPECIES];
   double Tree1=0, mtH3=0, mtmrca=0;

   /* PMismatch3s(); */
   nr=20; com.ls=1000;  savetree=0, makeseq=1; com.cleandata=1;
   printf("\nnloci & ls? "); fflush(stdout);
   scanf("%d%d", &nr, &com.ls);

   com.alpha=0;  com.ncatG=8; for(i=0;i<4;i++) com.pi[i]=1./4;
   printf("\nsequence data go into %s, trees into %s.\n", com.seqf, treef);
   com.z[0]=(char*)malloc((2*com.ns-1)*com.ls*sizeof(char));
   for (i=1; i<2*com.ns-1; i++) com.z[i]=com.z[i-1]+com.ls;
   if(com.alpha) com.rates=(double*)malloc(com.ls*sizeof(double));
   if(com.z[0]==NULL || (com.alpha && com.rates==NULL)) error2("oom");

   if(variable_ns) FOR(j,sptree.nspecies) nseqsp0[j]=sptree.nseqsp[j];
   if(makeseq) fprintf(fseq, "%6d\n", nr);
   for (ir=0; ir<nr; ir++) {
      if(variable_ns)
         for( ; ; ) { /* make sure at least 2 sequences are in the data */
            for(j=0,com.ns=0;j<sptree.nspecies;j++) 
               com.ns+=sptree.nseqsp[j]=(int)(nseqsp0[j]*rndu()*0.95+0.5);
            if(com.ns>1) break;
         }
      GetRandomGtree(-1);
      if(savetree) { OutaTreeN(fout,1,1);  FPN(fout); }
/*
if(MatchGTree()) {
   Tree1++;
   mtH3=(mtH3*(Tree1-1)+nodes[nodes[2].father].age)/Tree1;
   mtmrca=(mtmrca*(Tree1-1)+nodes[tree.root].age)/Tree1;
}
if((ir+1)%10000==0) 
   printf("%d  %9.6f\tEh3%9.6f EtMRCA %9.6f\n", ir+1, Tree1/(ir+1),mtH3,mtmrca);
*/
      if(!makeseq) continue;
      if (com.alpha) {
         Rates4Sites (com.rates, com.alpha, com.ncatG, com.ls, 1, com.space);
         for (j=1; j<com.ls; j++) com.rates[j]+=com.rates[j-1];
         abyx (1/com.rates[com.ls-1], com.rates, com.ls);
      }
      FOR(i,com.ls) com.z[tree.root][i]=(char)(rndu()*4.);
      EvolveJC (tree.root);

      for(i=0;i<sptree.nspecies;i++) fprintf(fseq," %3d",sptree.nseqsp[i]); 
      FPN(fseq);
      printSeqs(fseq, NULL, NULL, 0);

   }  /* for(ir) */
   fclose(fseq); free(com.z[0]);
   if(com.alpha) free(com.rates);
   if(variable_ns) FOR(j,sptree.nspecies) sptree.nseqsp[j]=nseqsp0[j];
   exit(0);
}

void EvolveJC (int inode)
{
/* Special version of Evolve that works with JC69-like (poisson) model only.
   This can be used to generate amino acid sequences also.
   For each branch in the tree, this determines the number of mutations 
   and then assigns mutations to sites at random.
   When alpha>0, com.rates[] are the accumulative probabilities.
*/
   int is,j, h, nmut, imut, ison;
   double r;

   if (com.alpha && fabs(com.rates[com.ls-1]-1)>1e-4) 
      { printf ("rates c.d.f.: 1 = %.6f?\n", com.rates[com.ls-1]); exit(-1); }
   for (is=0; is<nodes[inode].nson; is++) {
      ison=nodes[inode].sons[is];
      FOR(h,com.ls) com.z[ison][h]=com.z[inode][h];
      nmut = rndpoisson (nodes[ison].branch*com.ls);
      for (imut=0; imut<nmut; imut++) {
         if (com.alpha==0) h=(int)(rndu()*com.ls);
         else 
            for (h=0,r=rndu(); h<com.ls; h++) if (r<com.rates[h]) break;
         j=(int)(rndu()*3);
         if (j>=com.z[ison][h]) j++;
         com.z[ison][h]=(char)j;
      }
      if (nodes[ison].nson) EvolveJC(ison);
   }
}


#else

int ReadSeqData(char *seqfile, FILE*fout, int cleandata)
{
/* Read sequences at each locus. This sets data.nseqsp[ig], but not gnodes[].ipop 
   which is set in GetMem(), as gnodes[] is allocated after this routine.

   use com.cleandata=1 to delete ambiguities at all loci.

*/
   FILE *fseq=gfopen(seqfile,"r");
   int locus, i, maxns[NSPECIES]={0};
   char clean0=cleandata;

   fscanf(fseq, "%d", &data.ngene);
   printf("\nReading sequence data..  %d loci\n", data.ngene);
   if(data.ngene>NGENE) error2("raise NGENE?");
   for(locus=0,data.maxns=0; locus<data.ngene; locus++) {
      fprintf(fout, "\n\n*** Locus %d ***\n", locus+1);
      for(i=0,data.ns[locus]=0; i<sptree.nspecies; i++) {
         fscanf(fseq, "%d", &data.nseqsp[locus][i]);
         data.ns[locus]+=data.nseqsp[locus][i];
         if(data.nseqsp[locus][i]>maxns[i]) maxns[i]=data.nseqsp[locus][i];
         if(data.nseqsp[locus][i]>1 && sptree.nseqsp[i]<=1) {
            printf("\nlocus %d species %d.", locus+1,i+1);
            error2("parameter file & sequence file do not match.");
         }
         fprintf(fout, "%s (%d) ", sptree.nodes[i].name, data.nseqsp[locus][i]);
      }
      FPN(fout);

      printf("locus %2d: %2d sequences (\n", locus+1, data.ns[locus]);
      FOR(i,sptree.nspecies) printf(" %2d", data.nseqsp[locus][i]);
      com.cleandata=clean0;
      ReadSeq (NULL, fseq, clean0);
      data.cleandata[locus]=com.cleandata;

      if(com.ns!=data.ns[locus]) error2("ns mismatch in seq data file");
      if(com.ns<=1) error2("1 seq not useful in sequence file");
      PatternJC69like (fout);
      data.maxns=max2(data.maxns,com.ns);
      data.ns[locus]=com.ns;  data.ls[locus]=com.ls;  data.npatt[locus]=com.npatt;
      data.fpatt[locus]=com.fpatt; com.fpatt=NULL;
      for(i=0; i<com.ns; i++) { data.z[locus][i]=com.z[i]; com.z[i]=NULL; }

      printf(" ), %3d patterns, %s", com.npatt,(com.cleandata?"clean":"messy"));
   }
   for(i=0,com.cleandata=1; i<data.ngene; i++) 
      if(!data.cleandata[i]) com.cleandata=0;

   fflush(fout);
   printf("\nmax number of sequences in species:");
   for(i=0; i<sptree.nspecies; i++) printf(" %2d", maxns[i]);  FPN(F0);
   /* Memory for com.z[] and pnodes is allocated later according to data.maxns */
   /*
   if(data.maxns>maxns0) puts("raise ns in parameter file?");  
   */
   if(data.maxns>NS) error2("raise NS in source file?");
   fclose(fseq);
   return(0);
}


int GetMem (void)
{
/* This routine is called after the species tree and sequence data are read
   to allocate memory for gene trees (gnodes) and conditional probabilities 
   (conP).  The size of gnodes[locus] is determined using data.ns[locus], 
   no larger than nodes_t (static memory allocated using NS).  

   This routine also initializes ipop at tips in gene trees: gnodes[].ipop.
   It is not used by the simulation program.

   conP0 is the space for the tips.  Only one copy is needed as the conditional 
   probabilities for tips depend on the data and never change.  
   Space is allocated for a locus only if the locus has unclean (ambiguity) data, 
   that is, if(data.cleandata==0).
   The conditional probabilities for internal nodes are com.conPin[2],
   allocated according to data.ns[locus] and data.npatt[locus].
   Two copies of the space are allocated, hence the [2].  The copy used for the 
   current gene trees is conPin[com.curconP] while the copy for proposed gene
   trees is the alternative copy conPin[!com.curconP].  Even in the alternative 
   copy, the conP space for different loci do not overlap.  

   UpdateGB_InternalNode:
   UpdateGB_SPR:
     Those two proposal steps use UseLocus and AcceptLocus to copy between the 
     two copies of conPin.  They do not change com.curconP.

   UpdateTau:
   mixing:
     Those two steps use switchconPin to accept a locus, instead of memcpy to copy 
     the whole conPin space.

   UpdateTheta: This does not change the lnL.
*/
   /* sconP0: tips; sconP: internal nodes */
   int locus,i,j,k, sconP0, s1, stree;
   double *conP0, *conP;

   /* get mem for gnodes */
   for(locus=0,stree=0; locus<data.ngene; locus++)
      stree+=(data.ns[locus]*2-1)*sizeof(struct TREEN);
   if((gnodes[0]=(struct TREEN*)malloc(stree))==NULL) error2("oom gtree");
   for(locus=0,nodes=gnodes[0]; locus<data.ngene-1; locus++)
      gnodes[locus+1] = (nodes+=data.ns[locus]*2-1);
   for(locus=0; locus<data.ngene; locus++) {
      for(j=0; j<2*data.ns[locus]-1; j++) {
         gnodes[locus][j].nson=0; 
         gnodes[locus][j].age=0; 
         gnodes[locus][j].ibranch=-1;
      }
      for(i=0,k=0; i<sptree.nspecies; i++)  /* ipop for tips in gene tree */
         for(j=0; j<data.nseqsp[locus][i]; j++)
            gnodes[locus][k++].ipop=i;
   }

   /* get mem for conP (in nodes), and conP0 (tips) */
   if(mcmc.usedata) {
      data.conP_offset[0]=0;
      for(locus=0,com.sconP=sconP0=0; locus<data.ngene; locus++) {
         s1= com.ncode * data.npatt[locus];
         com.sconP += s1*(data.ns[locus]-1)*sizeof(double);
         if(!data.cleandata[locus])  sconP0 += s1*data.ns[locus]*sizeof(double);
         if(locus<data.ngene-1) 
            data.conP_offset[locus+1]=data.conP_offset[locus]+(data.ns[locus]-1)*s1;
      }
      if((com.conPin[0]=(double*)malloc(2*com.sconP+sconP0))==NULL) 
         error2("oom conP");
      com.conPin[1]      =com.conPin[0]+com.sconP/sizeof(double);
      if(sconP0) com.conP0=com.conPin[1]+com.sconP/sizeof(double);
      printf("%d bytes for conP, %d bytes for conP0, %d bytes for gene trees\n",
         2*com.sconP,sconP0,stree);

      /* set gnodes[locus][].conP for tips and internal nodes */
      com.curconP=0; conP=com.conPin[0]; conP0=com.conP0;
      for(locus=0; locus<data.ngene; locus++) {
         if(!data.cleandata[locus]) {
            for(j=0; j<data.ns[locus]; j++,conP0+=com.ncode*data.npatt[locus])
               gnodes[locus][j].conP = conP0;
            UseLocus(locus, 0, 1, 0);
            InitConditionalPNode ();
         }
         for(j=data.ns[locus]; j<data.ns[locus]*2-1; j++,conP+=com.ncode*data.npatt[locus])
            gnodes[locus][j].conP = conP;
      }
   }
   return(0);
}


void FreeMem(void)
{
   int locus, j;

   free(gnodes[0]);
   if(mcmc.usedata) free(com.conPin[0]);
   for(locus=0; locus<data.ngene; locus++) {
      free(data.fpatt[locus]);
      for(j=0;j<data.ns[locus]; j++)
         free(data.z[locus][j]);
   }
}

int UseLocus (int locus, int copytreeconP, int useData, int setSeqName)
{
/* This point *nodes to the gene tree at locus gnodes[locus] and set com.z[] 
   etc. for likelihood calculation for the locus.
   This also updates sptree.nseqsp[], which is used in UpdateGB_SPR() to 
   calculate ndesc[].
   if(copytreeconP), the gene tree gnodes[locus] is copied to nodes_t.  
   If (useData), com. and nodes[].conP are reset for data at that locus.
   Note that if (copytreeconP && useData), the conP for internal nodes point to 
   a fixed place (indicated by data.conP_offset[locus]) in the alternative space 
   com.conPin[!com.curconP].  Note that the conP for each locus uses different 
   space so that this routine can be used by all the proposal steps, some of 
   which operates on a per locus basis and some of which changes all loci.
   The conP for tips always point to com.conP0 (that is, when com.cleandata=0).
*/
   int i,j,k, s1=com.ncode*data.npatt[locus];
   double *conPt=com.conPin[!com.curconP]+data.conP_offset[locus];

   com.ns=data.ns[locus]; tree.root=data.root[locus]; tree.nnode=2*com.ns-1;
   FOR(i,sptree.nspecies) sptree.nseqsp[i] = data.nseqsp[locus][i];
   if(copytreeconP) { /* this preserves the old gene tree and conP. */
      nodes=nodes_t;
      memcpy(nodes_t, gnodes[locus], (com.ns*2-1)*sizeof(struct TREEN));
      if(useData) {
         for(i=com.ns; i<com.ns*2-1; i++)
            nodes_t[i].conP = conPt+(i-com.ns)*s1;
         memcpy(conPt, gnodes[locus][com.ns].conP, s1*(com.ns-1)*sizeof(double));
      }
   }
   else           /* this works on the old gene tree */
      nodes=gnodes[locus];
   if(useData) {
      com.cleandata=data.cleandata[locus];
      com.npatt=data.npatt[locus];  com.fpatt=data.fpatt[locus];
      for(i=0; i<com.ns; i++) com.z[i]=data.z[locus][i];
   }
   if(setSeqName) {
      FOR (i,tree.nnode) if(i!=tree.root) 
         nodes[i].branch=nodes[nodes[i].father].age-nodes[i].age;
      for(i=0,k=0;i<sptree.nspecies;i++) 
         FOR(j,sptree.nseqsp[i]) {
            if(sptree.nseqsp[i]>1) sprintf(com.spname[k++], "%s%d\0", sptree.nodes[i].name, j+1);
            else                   sprintf(com.spname[k++], "%s\0", sptree.nodes[i].name);
         }
   }
   return(0);
}

int AcceptLocus (int locus, int copyconP)
{
/* If (copytree && usedata), this copies nodes_t[].conP into gnodes[].conP for inner 
   nodes 
*/
   int i,ns=data.ns[locus], s1=com.ncode*data.npatt[locus];
   double *conPt;

   /* copies node_t into gnodes[locus] */
   data.root[locus]=tree.root;
   memcpy(gnodes[locus], nodes_t, (ns*2-1)*sizeof(struct TREEN));
   if(mcmc.usedata) {  /* reposition gnodes[][].conP */
      conPt=com.conPin[com.curconP]+data.conP_offset[locus];
      for(i=ns; i<ns*2-1; i++) gnodes[locus][i].conP = conPt+(i-ns)*s1;
      if(copyconP) {/* used in UpdateGB_InternalNode and UpdateGBTip only */
         conPt=com.conPin[!com.curconP]+data.conP_offset[locus];
         memcpy(gnodes[locus][ns].conP, conPt, s1*(ns-1)*sizeof(double));
      }
   }
   return(0);
}




void switchconPin(void)
{
/* This resets gnodes[locus].conP to the alternative com.conPin, to avoid 
   recalculation of conP, when a proposal is accepted in UpdateTau() or mixing().
*/
   int i,locus;
   double *conP;

   com.curconP=!com.curconP;
   conP=com.conPin[com.curconP];
   for(locus=0; locus<data.ngene; locus++)
      for(i=data.ns[locus]; i<data.ns[locus]*2-1; i++,conP+=com.ncode*data.npatt[locus])
         gnodes[locus][i].conP = conP;
}


int SetParameters(double x[], int set_theta, int set_tau)
{
/* x[] has node ages for species tree, and not gaps (inter-arrival times)
*/
   int j,dad;

   if(set_theta) 
      for(j=0; j<sptree.npop; j++)
         sptree.nodes[sptree.pops[j]].theta=x[j];
   if(set_tau) {
      for(j=sptree.nspecies; j<sptree.nnode; j++)
         sptree.nodes[j].age=x[sptree.npop+j-sptree.nspecies];

      for(j=sptree.nspecies; j<sptree.nnode; j++) {
         if(j==sptree.root) continue;
         dad=sptree.nodes[j].father;
         if(sptree.nodes[j].age>=sptree.nodes[dad].age)
            error2("species divergence times out of order.");
      }
   }
   return(0);
}

int InitConditionalPNode (void)
{
/* set partial likelihood at tips of the tree, when com.cleandata==0.
   Need testing if sequences in the data are ancestors.
*/
   static int times=0;
   int n=com.ncode, i,j,k,h;

   for (i=0;i<com.ns;i++) {
      FOR(h,com.npatt*n) nodes[i].conP[h]=0;
      for (h=0;h<com.npatt;h++) {
         k=strchr(BASEs,com.z[i][h])-BASEs; 
         if(k<0) 
            { printf("Character %d\n",com.z[i][h]); exit(-1); }
         else if(k<n)
            nodes[i].conP[h*n+k]=1;
         else 
            FOR(j,nBASEs[k]) 
               nodes[i].conP[h*n+ (strchr(BASEs,EquateNUC[k][j])-BASEs)]=1;
      }
   }
   return(0);
}


static double pij_JC69[2];

int ConditonalPNode (int inode)
{
   int n=com.ncode, i,j,k,h, ison;
   double t;

   FOR (i,nodes[inode].nson) {
      ison=nodes[inode].sons[i];
      if (nodes[ison].nson>0  && (!mcmc.saveconP || !com.oldconP[ison]))
         ConditonalPNode (nodes[inode].sons[i]);
   }
   FOR(i,com.npatt*n) nodes[inode].conP[i]=1;

   FOR (i, nodes[inode].nson) {
      ison=nodes[inode].sons[i];
      NPMat++;
      pijJC69 (pij_JC69, nodes[ison].branch);

      if (com.cleandata && nodes[ison].nson<1) /* tip */
         for (h=0; h<com.npatt; h++) 
            FOR(j,n) 
               nodes[inode].conP[h*n+j] *= pij_JC69[j!=com.z[ison][h]];
      else {
         for (h=0; h<com.npatt; h++) 
            FOR(j,n) {
               for(k=0,t=0; k<n; k++)    /* t is used as temp */
                  t += pij_JC69[j!=k] * nodes[ison].conP[h*n+k];
               nodes[inode].conP[h*n+j]*=t;
            }
      }
   }        /*  for (ison)  */
   return (0);
}

double lnpData (double lnpDi[])
{
/* This calculates the log likelihood, the log of the probability of the data 
   given gtree[nlocus] and blength[nlocus][2] for each locus.
   This updates gnodes[locus][].conP.
*/
   int j,locus;
   double lnL=0,y;

   if(mcmc.saveconP) FOR(j,data.maxns*2-1) com.oldconP[j]=0;
   for(locus=0; locus<data.ngene; locus++) {
      UseLocus(locus,0,1, 0);
      y=lnpD_locus();
      if(testlnL && fabs(lnpDi[locus]-y)>1e-10) 
         printf("\tlnLi %.6f != %.6f at locus %d\n", lnpDi[locus],y,locus+1);
      lnpDi[locus]=y;
      lnL+=y;
   }
   return(lnL);
}

double lnpD_locus (void)
{
/* This calculates ln{Di|Gi, Bi} using nodes[].age and tree.root.
   Note that nodes[].branch may not be kept up to date in the program.
*/
   int  h,i;
   double lnL=0, fh;

   if(!mcmc.usedata) return(0);
   FOR (i,tree.nnode) 
      if(i!=tree.root) 
         nodes[i].branch=nodes[nodes[i].father].age-nodes[i].age;

   ConditonalPNode (tree.root);

   for (h=0; h<com.npatt; h++) {
      for (i=0,fh=0; i<com.ncode; i++) 
         fh += com.pi[i]*nodes[tree.root].conP[h*com.ncode+i];
      lnL+=log(fh)*com.fpatt[h];
   }
   return (lnL);
}



double lnpGB_thetatau (void)
{
/* this calculates the prior of the gene tree and coalescent times (using tree 
   and nodes[]), when theta and tau are fixed from sptree.
   It initially collects information (coal) about the coalescent in each 
   population, such as nin, nout, coal, tj, by looping through all nodes[] 
   in the gene tree.  Then lnpGB is calculated by summing over the populations.
   Note that this does not use sptree.nseqsp[].
*/
   int i,j,k,inode, ipop, isonpop, nt;
   double lnp=0, starttime,endtime, *t,y;
   double small=0;
   /* double small=1e-12; */

   /* construct info (in, coal, out) for each population (the coal table) */
   for(i=0; i<sptree.nnode; i++) 
      { coal.nin[i]=coal.ncoal[i]=coal.nout[i]=0; }
   for(inode=0,coal.nout[sptree.root]=1; inode<tree.nnode; inode++) {
      ipop=nodes[inode].ipop;
      if(nodes[inode].nson==0) /* tip */
         coal.nin[ipop]++;
      else {                   /* internal node */
         coal.tj[ipop][coal.ncoal[ipop]++] = nodes[inode].age;
         for(i=0; i<nodes[inode].nson; i++) {
            isonpop=k=nodes[nodes[inode].sons[i]].ipop;
            if(isonpop==ipop) continue;
            for(; ; k=sptree.nodes[k].father) {
               if(k==isonpop) coal.nout[k]++;
               else if(k==ipop)  { coal.nin[k]++; break; }
               else { coal.nin[k]++; coal.nout[k]++; }
            }
         }
      }
      if(inode==tree.root && ipop!=sptree.root) {
         for(k=ipop; ; k=sptree.nodes[k].father) {
            if(k!=ipop)        coal.nin[k]++;
            if(k!=sptree.root) coal.nout[k]++;
            if(k==sptree.root) break;
         }
      }
   }
   for(k=0; k<sptree.npop; k++) { /* sort tj for each pop */
      ipop=sptree.pops[k]; 
      if(coal.nin[ipop]-coal.ncoal[ipop]!=coal.nout[ipop]) {
         printf("\nin out coal %d %d %d in ipop %d\n", 
            coal.nin[ipop], coal.nout[ipop], coal.ncoal[ipop], ipop);
         printGtree(1);
         error2("in out coal.");
      }
      nt=coal.ncoal[ipop]; t=coal.tj[ipop];
      if(nt>30)
         qsort(t, (size_t)nt, sizeof(double), comparedouble);
      else 
         for(i=0; i<nt-1; i++)  for (j=i+1; j<nt; j++)
            if (t[j]<t[i])  { y=t[i]; t[i]=t[j]; t[j]=y; } 
   }

   if(debug==1) {
      printf("\ncoalescence info in populations:\nspecies  in coal  out   tau     tj\n");
      for(i=0; i<sptree.nnode; i++,FPN(F0)) {
         printf("%4d: %5d%5d%5d %9.5f", i, coal.nin[i], coal.ncoal[i], coal.nout[i],sptree.nodes[i].age);
         for(j=0;j<coal.ncoal[i];j++) printf(" %9.5f", coal.tj[i][j]);
      }
      FPN(F0);
   }

   for(i=0; i<sptree.npop; i++) {
      ipop=sptree.pops[i]; t=coal.tj[ipop];
      starttime=sptree.nodes[ipop].age;
      if(ipop!=sptree.root) endtime=sptree.nodes[sptree.nodes[ipop].father].age;
      else                  endtime=-1;

      if(debug==1) printf("species %d: time: %8.4f > %8.4f ", ipop,starttime,endtime);

      if(coal.nin[ipop]>coal.nout[ipop]) {
         lnp+=(coal.nin[ipop]-coal.nout[ipop])*log(2/sptree.nodes[ipop].theta);
         for(j=coal.nin[ipop],k=0; j>coal.nout[ipop]; j--,k++) {
            y = t[k]- (k==0?starttime:t[k-1]);
            lnp -= j*(j-1)/sptree.nodes[ipop].theta*y;

            if(debug==1) printf(" j=%d: %9.5f", j,y);

            if(y<-small && noisy) {
               printGtree(1);
               printf("\n       tj = %.20f < 0 in ipop %d!\n", y, ipop);
               printf(" node age = %.20f\n", t[k]);
               printf("starttime = %.20f\n", starttime);
               getchar();
            }
         }
      }
      if(coal.nout[ipop]>1) {
         y = (coal.nout[ipop]==coal.nin[ipop] ? endtime-starttime : endtime-t[k-1]);
         lnp -= coal.nout[ipop]*(coal.nout[ipop]-1)/sptree.nodes[ipop].theta*y;

         if(debug==1) printf(" remained %.5f", y);
      }
      if(debug==1) FPN(F0);
   }
   if(debug==1) printf("\npGB=%.6g\n", exp(lnp));
   return(lnp);
}



int UpdateGB_InternalNode (double* lnL, double finetune)
{
/* This slides a node in the gene tree without changing the gene tree
   topology.  However, the new coalescent time tnew might be in different 
   population as the current population, so nodes[inode].ipop might change.
   The algorithm loops through all the internal nodes in the gene tree at each 
   locus.  The lower bound tb[0] is determined by age of the two sons 
   and also the age of the species that is ancestral to both sons.  The upper
   bound tb[1] is the father node's age.
*/
   int  accepted=0, locus, i,inode, sons[2],pops[2];
   int  ipopsource,ipoptarget, copytree;
   double lnacceptance, lnLd, lnpGBinew, lnpDinew, t,tnew;
   double tb[2];

   if(debug==2) puts("\nUpdateGB_InternalNode ");
   for(locus=0; locus<data.ngene; locus++) {

      if(debug==2) printf("\nlocus %d (ns = %d)\n", locus+1,data.ns[locus]);
      for(inode=data.ns[locus],copytree=1; inode<data.ns[locus]*2-1; inode++) {
         if(copytree) UseLocus(locus, 1, mcmc.usedata, 0); /* work on a copy */
         t=nodes[inode].age;
         ipopsource=nodes[inode].ipop;
         FOR(i,2) { sons[i]=nodes[inode].sons[i]; pops[i]=nodes[sons[i]].ipop; }
         tb[0]=max2(nodes[sons[0]].age, nodes[sons[1]].age);
         tb[1]=(inode==tree.root?OLDAGE:nodes[nodes[inode].father].age);

#if(1)   /* no change to ipop, remain in the same population */
         tb[0] = max2(tb[0],sptree.nodes[ipopsource].age);
         if(ipopsource!=sptree.root)
            tb[1] = min2(tb[1],sptree.nodes[sptree.nodes[ipopsource].father].age);
#else    /* possible change to ipop. Find the population ancestral to both sons */
         if(pops[0]!=pops[1]) { /* sptree.pops[] are in the right order */
            for(i=0; i<sptree.npop; i++) 
               if(sptree.pop_pop_table[pops[0]][sptree.pops[i]]
               && sptree.pop_pop_table[pops[1]][sptree.pops[i]])
                  break;
            tb[0] = max2(tb[0],sptree.nodes[sptree.pops[i]].age);
         }
#endif
         tnew=rnduab(t-finetune/2,t+finetune/2);
         tnew=reflect(tnew, tb[0], tb[1]);

         /* identify the target pop in which tnew is, by going up sptree */
         for(i=ipoptarget=pops[0]; ; i=sptree.nodes[i].father) {
            if(sptree.nodes[i].age>tnew) break;
            ipoptarget=i;
            if(i==sptree.root) break;
         }

         /* two things are updated in the loop */
         nodes[inode].age=tnew;
         nodes[inode].ipop=ipoptarget;

         if(mcmc.saveconP) {
            for(i=com.ns; i<com.ns*2-1; i++) com.oldconP[i]=1;
            for(i=inode; ; i=nodes[i].father)
               { com.oldconP[i]=0; if(i==tree.root) break; }
         }

         lnpGBinew = lnpGB_thetatau();
         lnpDinew = lnpD_locus();
         lnLd = lnpDinew-data.lnpDi[locus];
         lnacceptance = lnpGBinew - data.lnpGBi[locus] + lnLd;

         if(debug==2) {
            printf("node %2d tb (%8.5f %8.5f) t %7.5f > %7.5f lnLd %9.5f ",inode,tb[0],tb[1],t,tnew,lnLd);
            for(i=com.ns; i<com.ns*2-1; i++) printf("%d",com.oldconP[i]);
         }

         if(lnacceptance>=0 || rndu()<exp(lnacceptance)) {
            accepted++;
            *lnL+=lnLd;
            data.lnpDi[locus]=lnpDinew;  data.lnpGBi[locus]=lnpGBinew;
            AcceptLocus(locus, 1);
            copytree=0;

            if(debug==2) printf(" Y (%4d)\n", NPMat);

         }
         else {
            nodes[inode].age=t;
            nodes[inode].ipop=ipopsource;
            copytree=1;

            if(debug==2) printf(" N (%4d)\n", NPMat);
         }
      }
   }
   return(accepted);
}

#define swap2(a,b,y) { y=a; a=b; b=y; }

void GraftNode(int source, int target, double age, int ipop)
{
/* This performs surgery on the gene tree (nodes) to cut branch source 
   (branch ancestral to node source) and regrafts it to branch target.
   It assumes binary rooted tree.
   age & ipop are for the target node for attaching (father).
   After the operation, the following branches are generated:
     grandpa->sib (at the source)
     father->source, father->target, targetfather->father (at the target)

   IMPORTANT: This keeps the node number for tree.root unchanged, as required 
   by UpdateGBTip.
*/
   int i,k, father, grandpa=-1, sib, targetfather, oldroot=tree.root;
   double y;

   father=nodes[source].father;
   nodes[father].age=age;  nodes[father].ipop=ipop;
   if(father!=tree.root) grandpa=nodes[father].father;
   sib=nodes[father].sons[0]+nodes[father].sons[1]-source;
   targetfather=nodes[target].father;

   if(target!=sib && target!=father) { /* change to tree topology */
      /* cut at source */
      if (father==tree.root) {         /* old root is lost */
         tree.root=sib;  nodes[sib].father=-1;
      }
      else {
         if(nodes[grandpa].sons[i=0]!=father) i=1;
         nodes[grandpa].sons[i]=sib; nodes[sib].father=grandpa;
      }

      /* regraft */
      nodes[father].father=nodes[target].father;
      nodes[target].father=father;
      nodes[father].sons[0]=target;
      nodes[father].sons[1]=source;
      if (target==tree.root) {  /* new root is born */
         tree.root=father;
      }
      else
         for(i=0; i<2; i++) {
            if(nodes[targetfather].sons[i]==target) 
               { nodes[targetfather].sons[i]=father; break; }
         }
   }
   if(tree.root!=oldroot) { /* swap everything except .conP */
      /* if(debug==3)  printGtree(1); */
      swap2(nodes[tree.root].ipop, nodes[oldroot].ipop, k);
      swap2(nodes[tree.root].age, nodes[oldroot].age, y);
      FOR(i,2) {
         nodes[nodes[tree.root].sons[i]].father=oldroot;
         nodes[nodes[oldroot].sons[i]].father=tree.root;
      }
      FOR(i,2) swap2(nodes[tree.root].sons[i], nodes[oldroot].sons[i], k);
      FOR(i,2) if(nodes[oldroot].sons[i]==oldroot) nodes[oldroot].sons[i]=tree.root;
      k=nodes[oldroot].father;
      FOR(i,2) if(nodes[k].sons[i]==oldroot) nodes[k].sons[i]=tree.root;
      if((nodes[tree.root].father=k)==tree.root)
         nodes[tree.root].father=oldroot;
      nodes[oldroot].father=-1;

      father=nodes[source].father;
      if(grandpa==tree.root)    grandpa=oldroot;
      else if(grandpa==oldroot) grandpa=tree.root;
      tree.root=oldroot;
   }
   if(mcmc.usedata && mcmc.saveconP) {
      for(i=com.ns; i<com.ns*2-1; i++) com.oldconP[i]=1;
      for(i=father; ; i=nodes[i].father)
         { com.oldconP[i]=0; if(i==tree.root) break; }
      for(i=grandpa; i!=-1; i=nodes[i].father)
         { com.oldconP[i]=0; if(i==tree.root) break; }
   }
}

int UpdateGB_SPR (double* lnL, double finetune)
{
/* This removes a node (source) in the gene tree and regrafts it to a feasible 
   branch.  There is no change to the node times inside the subtree.  
   The step cycles through all loci and all nodes except the root.
   ndesc[] holds the number of descendent tips at the locus of each population.
   For example, ndesc[sptree.root]==com.ns.  This is used to identify feasible 
   populations (species) to determine range[0] for tnew.  It is used to deal 
   with loci at which some species are missing.  ndesc[] is calculated from 
   sptree.nseqsp[species], which is updated for the locus in UseLocus().  

   Note that inode is the root node for the subtree being moved, species is 
   the population inode is in, and t and tnew are the time at which the subtree
   coalesce with the rest of the gtree, that is, the age of inode's father.  
   To determine the lower bound range[0] for tnew, the following requirements 
   are considered:
   (a) tnew must be older than the age of the subtree: nodes[inode].age.
   (b) The population of the target node must be ancestral to or the same as 
       the population of the subtree: (sptree.pop_pop_table[species][ipop]==1).
   (c) There must be other lineages outside the subtree passing that population
       to which the subtree can join.  For tips (inode<com.ns), this requirement 
       is satisfied if (ndesc[ipop]>1).  For internal nodes, (ndesc[ipop]>1) is 
       necessary but not sufficient, and the algorithm loops through all branches 
       (nodes) in the gene tree, and identifies those outside subtree whose 
       father node is older than range[0].  For each of such branches, we 
       record the youngest age at which the subtree can join.
*/
   int accepted=0,nsource, ntarget,targets[NS*2-1],target, copytree;
   int locus, i,j, inode, species, ipop, ipopsource, ipoptarget, father;
   int ndesc[2*NSPECIES-1]; /* #descendents for ancestral species */
   double lnacceptance=0, lnLd, lnpGBinew,lnpDinew;
   double y,range[2], t,tnew;

   if(debug==3) puts("\nUpdateGB_SPR ");
   for(locus=0; locus<data.ngene; locus++) {
      UseLocus(locus, 0, 0, 0); /* this resets sptree.nseqsp[] */

      if(debug==3) {
         printf("\nlocus %d (ns = %d)\n", locus+1,data.ns[locus]);
         for(i=sptree.nspecies;i<2*sptree.nspecies-1;i++) 
            printf("sp.time[%d] = %.6f\n", i,sptree.nodes[i].age);
         printGtree(1);
      }
      FOR(i,2*sptree.nspecies-1) ndesc[i]=0;
      FOR(species,sptree.nspecies) { /* update ndesc[] for each locus */
         ndesc[species]=sptree.nseqsp[species];
         for(i=sptree.nspecies;i<2*sptree.nspecies-1;i++)
            if(sptree.pop_pop_table[species][i]) ndesc[i]+=sptree.nseqsp[species];
      }
      if(ndesc[sptree.root]!=com.ns) error2("ndesc");

      for(inode=0,copytree=1; inode<(mcmc.moveinnode?com.ns*2-1:com.ns); inode++) {
         /* NOTE: tree.root should not be changed in GraftNode() */
         if(inode==tree.root) continue; 
         if(copytree) UseLocus(locus, 1, mcmc.usedata, 0);
         species=nodes[inode].ipop;
         father=nodes[inode].father;  ipopsource=nodes[father].ipop;
         t=nodes[father].age;

         /* locate lower bound range[0] for tnew.  See notes above */
         range[0]=nodes[inode].age;  range[1]=OLDAGE;
         for(i=0; i<sptree.npop; i++) {
            ipop=sptree.pops[i];
            if(sptree.pop_pop_table[species][ipop] && ndesc[ipop]>1)
               { range[0]=max2(range[0],sptree.nodes[ipop].age); break; }
         }
         /* For moving internal nodes, find min feasible speciation time (y) 
            in remaining subtree.  See note (c) above. */
         if(inode>=com.ns && sptree.nspecies>1) {
            for(i=0,y=sptree.nodes[sptree.root].age; i<tree.nnode; i++) {
               if(i==tree.root || nodes[i].father==inode || nodes[nodes[i].father].age<=range[0])
                  continue;
               /* Is node i a descendent of inode? */
               for(j=i; j!=tree.root&&j!=inode; j=nodes[j].father) ;
               if(j==inode) continue;  /* skip node i if it is a descendent */
               for(ipop=nodes[i].ipop; ipop!=sptree.root; ipop=sptree.nodes[ipop].father)
                  if(sptree.pop_pop_table[species][ipop])
                     { y=min2(y,sptree.nodes[ipop].age); break; }
               if(y<range[0]) break;  /* no point for checking further */
            }
            range[0] = max2(range[0],y);
         }

         tnew=rnduab(t-finetune/2,t+finetune/2);
         tnew=reflect(tnew, range[0], range[1]);

         /* identify the target pop in which tnew is, by going up sptree */
         for(ipop=ipoptarget=species; ; ipop=sptree.nodes[ipop].father) {
            if(sptree.nodes[ipop].age>tnew) break;
            ipoptarget=ipop;
            if(ipop==sptree.root) break;
         }

         if(debug==3) printf("inode %2d father %2d range (%6.4f %7.4f) time %7.5f (%2d) >%7.5f (%2d) ",
                              inode,father,range[0],range[1],t,ipopsource,tnew,ipoptarget);

         /* count and identify feasible target branches. Watch out for root */
         ntarget=0;  nsource=1;
         if(tnew>=nodes[tree.root].age)
            targets[ntarget++]=tree.root;
         else
            for(i=0; i<tree.nnode; i++) { /* Is node i a possible target? */
               if(nodes[i].age<=tnew && nodes[nodes[i].father].age>tnew
                  && sptree.pop_pop_table[nodes[i].ipop][ipoptarget]
                  && i!=inode && i!=tree.root) 
                  targets[ntarget++]=i;
            }
         if(father!=tree.root) 
            for(i=0; i<tree.nnode; i++) { /* Is node i a possible source? */
               if(nodes[i].age<=t && nodes[nodes[i].father].age>t
                  && sptree.pop_pop_table[nodes[i].ipop][ipopsource]
                  && i!=inode && i!=tree.root && i!=father 
                  && nodes[i].father!=father)
                  nsource++;
            }

         if(nsource<1 || ntarget<1) {
            printf("\nsource target %d %d", nsource,ntarget); 
            printf("\nlocus %d node %d\n", locus, inode); 
            printGtree(1);
            exit(-1); 
         }
         target=targets[(int)(ntarget*rndu())];
         GraftNode(inode,target,tnew, ipoptarget);

         if(debug==3) {
            printf("line %d > %2d (", nsource,ntarget);
            for(i=0; i<ntarget; i++) printf(" %2d", targets[i]);
            printf(") > %2d ", target);
            printGtree(1);
         }
         lnacceptance = log((double)ntarget/nsource);
         lnpGBinew = lnpGB_thetatau();
         lnacceptance += lnpGBinew - data.lnpGBi[locus];
         lnpDinew = lnpD_locus();
         lnLd = lnpDinew-data.lnpDi[locus];
         lnacceptance += lnLd;

         if(debug==3) { 
            printf(" lnLd %9.5f ", lnLd);
            for(i=com.ns; i<com.ns*2-1; i++) printf("%d",com.oldconP[i]); 
         }

         if(lnacceptance>=0 || rndu()<exp(lnacceptance)) {
            accepted++;
            *lnL+=lnLd;
            data.lnpDi[locus]=lnpDinew;
            data.lnpGBi[locus]=lnpGBinew;
            AcceptLocus(locus, 1);
            copytree=0;
            if(debug==3) printf(" Y (%4d)\n", NPMat);
         }
         else {
            copytree=1;
            if(debug==3) printf(" N (%4d)\n", NPMat);
         }
      }
   }
   return(accepted);
}



int UpdateTheta (double x[], double finetune, double space[])
{
/* This updates theta's one by one.   Note theta=x.
   This step does not require calculation of the likelihood.
*/
   int ipop, j, locus, accepted=0;
   double xnew[2*NSPECIES-1], lnacceptance;
   double *lnpGBinew=space, *a=mcmc.a_gamma, *b=mcmc.b_gamma;

   for(ipop=0; ipop<sptree.npop; ipop++) {
      /* prior and proposal ratios */
      for(j=0; j<sptree.npop; j++) xnew[j]=x[j];

      /* for(j=0; j<sptree.npop; j++) { */ 
      for(j=ipop; j==ipop; j++) {  /* changes one theta */
         xnew[j]=rnduab(x[j]-finetune/2, x[j]+finetune/2);
         if(xnew[j]<0) xnew[j]*=-1;
      }
      lnacceptance = lnpriorTheta(x, xnew);

      SetParameters (xnew, 1, 0);
      for(locus=0; locus<data.ngene; locus++) {
         UseLocus(locus, 1, 0, 0);
         lnpGBinew[locus] = lnpGB_thetatau();
         lnacceptance += lnpGBinew[locus] - data.lnpGBi[locus];
      }
      if(lnacceptance>=0 || rndu()<exp(lnacceptance)) {
         for(locus=0; locus<data.ngene; locus++)
            data.lnpGBi[locus]=lnpGBinew[locus];
         for(j=0; j<sptree.npop; j++) x[j]=xnew[j];
         accepted++;
      }
      else
         SetParameters (x, 1, 0);
   }
   return(accepted);
}



int UpdateTau (double x[], double *lnL, double finetune, double space[])
{
/* This updates speciation dates tau using the rubber-band algorithm.
   ntj[] are counts of nodes below and above tau (m and n in the paper).
*/
   int k, itau, is, inode, locus, accepted=0;
   int ntj[2], ntj_locus[2];
   double *tau=x+sptree.npop, taunew[NSPECIES-1], taub[2], t, taufactor[2];
   double lnacceptance=0, lnLd;
   double *lnpGBinew=space, *lnpDinew=space+data.ngene;

   if(debug==5) puts("\nUpdateTau ");
   FOR(is,sptree.nspecies-1) taunew[is]=tau[is];
   for(is=sptree.nspecies; is<sptree.nnode; is++) {
      itau=is-sptree.nspecies;
      lnLd=0;   taub[0]=0;  taub[1]=OLDAGE;
      if(sptree.nodes[is].nson)
         taub[0]=max2(sptree.nodes[sptree.nodes[is].sons[0]].age,
                      sptree.nodes[sptree.nodes[is].sons[1]].age);
      if(is!=sptree.root) 
         taub[1]=sptree.nodes[sptree.nodes[is].father].age;
      taunew[itau]=rnduab(tau[itau]-finetune/2,tau[itau]+finetune/2);
      taunew[itau]=reflect(taunew[itau],taub[0],taub[1]);
      FOR(k,2) 
         taufactor[k]=(taunew[itau]-taub[k])/(tau[itau]-taub[k]);
      sptree.nodes[is].age=taunew[itau];

      lnacceptance = lnpriorTau(tau, taunew);
      
      if(debug==5) printf("species %d taub: %8.5f %8.5f tau: %8.5f%8.5f", is,taub[0],taub[1],tau[itau],taunew[itau]); 
	  
      ntj[0]=ntj[1]=0;
      for(locus=0; locus<data.ngene; locus++) {
         UseLocus(locus, 1, mcmc.usedata, 0);  /* copy gtree & conP */
         ntj_locus[0]=ntj_locus[1]=0;
         if(mcmc.saveconP) for(k=com.ns; k<2*com.ns-1; k++) com.oldconP[k]=1;

         for(inode=com.ns; inode<tree.nnode; inode++) {
            t=nodes[inode].age;
            if (t>=taub[0] && t<taub[1]
            && (nodes[inode].ipop==is || sptree.nodes[nodes[inode].ipop].father==is)) {
               k = (t>=tau[itau] && is!=sptree.root); /* k=0: below; 1: above */
               nodes[inode].age = t = taub[k]+taufactor[k]*(t-taub[k]);
               ntj_locus[k]++;
               if (t<taub[0] || (t<sptree.nodes[sptree.root].age && t>taub[1]))
                  error2("tj out of range");

               if(mcmc.saveconP)
                  for(k=inode; ; k=nodes[k].father)
                     { com.oldconP[k]=0; if(k==tree.root) break; }
            }
         }

         ntj[0]+=ntj_locus[0]; ntj[1]+=ntj_locus[1];
         lnpGBinew[locus] = lnpGB_thetatau();
         if(ntj_locus[0]+ntj_locus[1]) {
            lnpDinew[locus] = lnpD_locus();
            lnLd += lnpDinew[locus]-data.lnpDi[locus];
         }
         else
            lnpDinew[locus] = data.lnpDi[locus];
         lnacceptance += lnpGBinew[locus] - data.lnpGBi[locus];

         if(debug==5) printf(" (%d %d) ",ntj_locus[0],ntj_locus[1]);

      }
      lnacceptance += lnLd + ntj[0]*log(taufactor[0]) + ntj[1]*log(taufactor[1]);

      if(debug==5) {
         printf(" ntj: %d %d", ntj[0],ntj[1]);
         printf(" lnLd %9.5f ", lnLd);
         for(k=com.ns; k<com.ns*2-1; k++) printf("%d",com.oldconP[k]);
      }

      if(lnacceptance>=0 || rndu()<exp(lnacceptance)) {
         accepted++;
         for(locus=0; locus<data.ngene; locus++) {
            data.lnpDi[locus]=lnpDinew[locus];  
            data.lnpGBi[locus]=lnpGBinew[locus];
            UseLocus(locus, 0, 0, 0);

            for(inode=com.ns; inode<tree.nnode; inode++) { /* redo changes */
               t=nodes[inode].age;
               if (t>=taub[0] && t<taub[1]
               && (nodes[inode].ipop==is || sptree.nodes[nodes[inode].ipop].father==is)) {
                  k = (t>=tau[itau] && is!=sptree.root);
                  nodes[inode].age = taub[k]+taufactor[k]*(t-taub[k]);
               }
            }
         }
         tau[itau] = taunew[itau];
         *lnL+=lnLd;
         if(mcmc.usedata) switchconPin();
         if(debug==5) printf(" Y (%4d)\n", NPMat);
      }
      else {
         sptree.nodes[is].age = taunew[itau] = tau[itau];
         if(debug==5) printf(" N (%4d)\n", NPMat);
      }
   }
   return(accepted);
}


int Reset_Thetamix(double var[])
{
/* This reset the mixing directions for the theta parameters, according to 
   their correlations with the tau parameters (ages).
   mcmc.mixtheta[]: 0: no change; 1: multiply -1: divide
*/
   int i,j, ntheta=sptree.npop, changed=0, jrmax;
   double r[(NSPECIES*3-1)*(NSPECIES*3-1)], rmax;
   double smallr=0.2;

   printf("\n\nCorrelation matrix (age, not gap, is always used.)");
   FOR(i,com.np) FOR(j,com.np) 
      r[i*com.np+j]=var[i*com.np+j]/sqrt(var[i*com.np+i]*var[j*com.np+j]);
   for(i=0;i<com.np;i++,FPN(F0))  FOR(j,i)
      printf(" %8.5f", r[i*com.np+j]);
   FOR(i,ntheta) {  /* going through theta's */
      for(j=ntheta,rmax=0,jrmax=0; j<com.np; j++) /* going through tau's */
         if(fabs(r[i*com.np+j])>rmax) { rmax=fabs(r[i*com.np+j]); jrmax=j; }
      mcmc.mixtheta[i] = (rmax<smallr ? 0 : r[i*com.np+jrmax] > 0 ? 1 : -1);
      if(mcmc.mixtheta[i]!=1) changed=1;
   }
   if(changed) {
      printf("\nUpdating theta's in the mixing step: ");
      FOR(i,ntheta) {
         printf("%3d (%s)",mcmc.mixtheta[i], 
            (mcmc.mixtheta[i]==0?"no change":mcmc.mixtheta[i]==1?"multiply":"divide"));
      }
      FPN(F0); FPN(F0);
   }
   return(0);
}

int mixing (double x[], double* lnL, double finetune, double space[])
{
/* This multiplies all tau, and all branch lengths by a constant c.  For the 
   theta, the behavior depends on mcmc.mixtheta[].  If this is 0, theta is 
   not changed; if it is 1, theta is multiplied by c, and if it is -1, theta 
   is divided by c.
*/
   int locus, j, k, ntheta=sptree.npop, ntau=sptree.nspecies-1;
   int nmultiply=0, ndivide=0;
   double *xnew=space, *lnpDinew=space+com.np, *lnpGBinew=lnpDinew+data.ngene;
   double c, logc, lnacceptance=0, lnLd, *a=mcmc.a_gamma, *b=mcmc.b_gamma;

   for(j=0; j<ntheta; j++) 
      if(mcmc.mixtheta[j]==1)       nmultiply++;
      else if(mcmc.mixtheta[j]==-1) ndivide++;

   c=exp(finetune*(rndu()-.5));  logc=log(c);
   for(locus=0,k=0; locus<data.ngene; locus++)  k+=data.ns[locus]-1;
   lnacceptance=(ntau+k+nmultiply-ndivide)*logc;
   for(j=0;j<com.np;j++) xnew[j]=x[j]*c;
   for(j=0;j<ntheta;j++) 
      if(mcmc.mixtheta[j]==0)       xnew[j]=x[j];
      else if(mcmc.mixtheta[j]==-1) xnew[j]=x[j]/c;

   lnacceptance+=lnpriorTheta(x, xnew)+lnpriorTau(x+ntheta, xnew+ntheta);
   SetParameters(xnew,1,1);

   if(mcmc.saveconP) FOR(j,data.maxns*2-1) com.oldconP[j]=0;
   for(locus=0,lnLd=0; locus<data.ngene; locus++) {
      UseLocus(locus, 1, mcmc.usedata, 0);  /* copy gtree & conP */
      for(j=com.ns; j<tree.nnode; j++)  nodes[j].age*=c;
      lnpGBinew[locus]=lnpGB_thetatau();
      lnacceptance += lnpGBinew[locus]-data.lnpGBi[locus];

      lnpDinew[locus] = lnpD_locus();
      lnLd += lnpDinew[locus]-data.lnpDi[locus];
      lnacceptance += lnpDinew[locus]-data.lnpDi[locus];
   }

   if(lnacceptance>0 || rndu()<exp(lnacceptance)) { /* accept */
      FOR(j,com.np) x[j]=xnew[j];
      FOR(locus,data.ngene) {
         UseLocus(locus, 0, 1, 0);
         FOR(j,tree.nnode)  { nodes[j].branch*=c;  nodes[j].age*=c; }
         data.lnpDi[locus]=lnpDinew[locus];
         data.lnpGBi[locus]=lnpGBinew[locus];
      }
      if(mcmc.usedata) switchconPin();
      *lnL += lnLd;
      return(1);
   }
   else /* restore sptree */
      SetParameters(x,1,1);

   return(0);
}


void checkGtree (void)
{
   int locus, inode,j, ipop, father;
   double t, tb[2];
   /* double small=1e-12, verysmall=1e-17; */
   double small=0, verysmall=0;

   FOR(locus,data.ngene) {
      UseLocus(locus, 0, 0, 1);
      for(inode=com.ns; inode<tree.nnode; inode++) {
         t=nodes[inode].age;   father=nodes[inode].father;
         ipop=nodes[inode].ipop;
         tb[0]=sptree.nodes[ipop].age;  tb[1]=OLDAGE;
         for(j=0; j<nodes[inode].nson; j++) 
            tb[0]=max2(tb[0], nodes[nodes[inode].sons[j]].age);
         if(ipop!=sptree.root) 
            tb[1]=sptree.nodes[sptree.nodes[ipop].father].age;
         if(inode!=tree.root) 
            tb[1]=min2(tb[1],nodes[nodes[inode].father].age);

         if(tb[1]<tb[0] || t<tb[0]-small || t>tb[1]+small) {
            printf("\nlocus %d node %d t=%9.6f tb: (%9.6f %9.6f)", locus,inode,t,tb[0],tb[1]);
            printf("\nspecies tree and gene tree:\n");
            for(j=0; j<2*sptree.nspecies-1; j++,FPN(F0)) {
               printf("species %d %-12s ", j+1,sptree.nodes[j].name);
               printf("age %.4f  theta %.4f  ", sptree.nodes[j].age, sptree.nodes[j].theta);
            }
            printGtree(1);
            printf("\ninconsistent gene tree at locus %d node %d\ttime %.5f (%.5f %.5f)\n",
               locus,inode,t,tb[0],tb[1]);
            puts("Enter to continue. ");
            getchar();
         }
         else if(t<tb[0]+verysmall)
            nodes[inode].age=tb[0]+verysmall;
         else if(t>tb[1]-verysmall)
            nodes[inode].age=tb[1]-verysmall;
      }
   }
}

double lnpriorTheta (double theta[], double thetanew[])
{
/* this calculates the prior ratio for theta's.
*/
   int j;
   double lnp=0, *a=mcmc.a_gamma, *b=mcmc.b_gamma;

   for(j=0; j<sptree.npop; j++) {
      if(fabs(thetanew[j]-theta[j])>1e-99) 
         lnp += (a[j]-1)*log(thetanew[j]/theta[j]) - b[j]*(thetanew[j]-theta[j]);
   }
   return(lnp);
}

double lnpriorTau (double tau[], double taunew[])
{
/* This calculates the log prior for speciation times using mcmc.a_gamma 
   and mcmc.b_gamma.  Note that tau[], used in the MCMC algorithm, are 
   always ages, while the gamma prior might be specified for the gaps.  
   The variable transfomation has Jacobi = 1.
*/
   int i,j;
   double lnp=0, gap, gapnew;
   double *a=mcmc.a_gamma+sptree.npop, *b=mcmc.b_gamma+sptree.npop;

   if(mcmc.speciation[0]==1) { /* ages */
      for(j=0; j<sptree.nspecies-1; j++)
         if(fabs(taunew[j]-tau[j])>1e-20) 
            lnp += (a[j]-1)*log(taunew[j]/tau[j]) - b[j]*(taunew[j]-tau[j]);
   }
   else  {                     /* gaps */
      for(j=0; j<sptree.nspecies-1; j++) {
         gap=tau[j]; gapnew=taunew[j];
         i=sptree.nodes[sptree.nspecies+j].son_age - sptree.nspecies;
         if(i>=0) { gap-=tau[i]; gapnew-=taunew[i]; }
         if(fabs(gapnew-gap)>1e-20) 
            lnp += (a[j]-1)*log(gapnew/gap) - b[j]*(gapnew-gap);
      }
   }
   return (lnp);
}


int DownSptreeSetTime (int inode)
{
/* This traverse down the species tree to specify the node ages, starting from 
   the root age.
*/
   int j,ison, dad;

   for (j=0; j<sptree.nodes[inode].nson; j++) {
      ison=sptree.nodes[inode].sons[j];
      if(sptree.nodes[ison].nson) {
         dad=sptree.nodes[ison].father;
         sptree.nodes[ison].age = sptree.nodes[dad].age*(.5+.5*rndu());
         DownSptreeSetTime(ison);
      }
   }
   return(0);
}

int MCMC (FILE* fout)
{
   char mcmcout[32]="mcmc.out";
   FILE *fmcmcout=gfopen(mcmcout,"w");
   int locus, i,j,k, ir, nused=0;
   /* x[] is current state of chain, and y[] and my[] are for output.   */
   double x[NSPECIES*3-2],  mx[NSPECIES*3-2], varx[(NSPECIES*3-1)*(NSPECIES*3-1)];
   double xo[NSPECIES*3-1], mxo[NSPECIES*3-1]; /* xo for output */
   double mtmrca[NGENE]={0};
   double *tau=x+sptree.npop, lnL, Ptj=0,Ptip=0,Ptheta=0,Ptau=0,Pmix=0,total=0;

   noisy=3;
   mcmc.print=1;
   mcmc.moveinnode=1;     /* moves internal nodes in the gene tree as well */
   mcmc.speciation[0]=0;  /* =0 for gaps; 1 for ages for the prior */
   mcmc.speciation[1]=0;  /* =0 for gaps; 1 for ages for the output */
   if(mcmc.speciation[0]==1) mcmc.speciation[1]=1;
   mcmc.mixtheta_change=1;
   mcmc.saveconP=1;
   if(!mcmc.usedata) mcmc.saveconP=0;
   FOR(j,sptree.npop) mcmc.mixtheta[j]=1; /* 0: no update, 1: multiply 2: divide */

   FOR(j,data.maxns*2-1) com.oldconP[j]=0;
   printf("\n%d burnin, sampled every %d, %d samples\n", 
           mcmc.burnin, mcmc.output, mcmc.nsample);
   if(mcmc.usedata) puts("Approximating posterior, using sequence data");
   else             puts("Approximating prior, not using sequence data");

   printf("(Settings: cleandata=%d print=%d saveconP=%d moveinnode=%d mixtheta_change=%d)\n",
          com.cleandata, mcmc.print, mcmc.saveconP, mcmc.moveinnode, mcmc.mixtheta_change);

   /* get initial values.  
   Initials values for speciation ages are not done properly */

   for(j=0; j<com.np; j++) x[j]=0.1;
   for(j=0; j<sptree.npop; j++) x[j]=rndgamma(mcmc.a_gamma[j])/mcmc.b_gamma[j];
   for( ; j<com.np; j++) x[j]=rndu()*rndu()*0.1; 

   
   if(sptree.nspecies>1) 
      sptree.nodes[sptree.root].age = rndu();
   DownSptreeSetTime(sptree.root);
   for(j=sptree.nspecies; j<sptree.nspecies*2-1; j++)
      x[sptree.npop+j-sptree.nspecies]=sptree.nodes[j].age;

   /*
   printf("initials (np = %d)? (note times are ages) ", com.np);
   for(j=0;j<com.np;j++) scanf("%lf", &x[j]);
   */

   printf("\nStarting MCMC... ");
   printf("Initial parameters (gene trees generated from the prior):\n");
   for(j=0;j<com.np;j++) printf("%10.5f",x[j]);
   for(locus=0,totalTimes=totalSeqs=0; locus<data.ngene; locus++) 
      { totalTimes+=data.ns[locus]-1; totalSeqs+=data.ns[locus]; }
   /* initialize likelihood and prior, for each locus */
   SetParameters(x,1,1);

   for(locus=0,lnL=0; locus<data.ngene; locus++) {
      GetRandomGtree(locus);  data.root[locus]=tree.root;
      data.lnpGBi[locus]=lnpGB_thetatau();
   }
   lnL=lnpData(data.lnpDi);
   printf("\nlnL0 =%12.3f\n", lnL);
   zero(mxo,com.np);  zero(mx,com.np);  zero(varx,com.np*com.np);

   for(ir=-mcmc.burnin,nused=0; ir<mcmc.output*mcmc.nsample; ir++) {
      if(mcmc.usedata && mcmc.mixtheta_change && sptree.nspecies>1 
         && (ir+1)==(mcmc.output*mcmc.nsample/10))
         Reset_Thetamix(varx);  /* this resets the mixing step */
      if(ir==0) {  /* reset after burnin */
         total=0; Ptj=Ptip=Ptheta=Ptau=Pmix=0;
         nused=0; zero(mxo,com.np); zero(mx,com.np); zero(varx,com.np*com.np); 
      }
      total++;

      Ptj+=UpdateGB_InternalNode(&lnL, mcmc.finetune[0]);
      Ptip+=UpdateGB_SPR(&lnL, mcmc.finetune[1]);
      Ptheta+=UpdateTheta(x, mcmc.finetune[2], com.space);
      if(sptree.nspecies>1)
         Ptau+=UpdateTau(x, &lnL, mcmc.finetune[3], com.space);
      Pmix+=mixing(x, &lnL, mcmc.finetune[4], com.space);
      /* checkGtree(); */

      nused++;
      FOR(j,com.np) xo[j]=x[j];
      if(mcmc.speciation[1]==0)  /* print gaps */
         for(j=0;j<sptree.nspecies-1; j++) {
            k=sptree.nodes[sptree.nspecies+j].son_age-sptree.nspecies;
            if(k>0) xo[sptree.npop+j]=x[sptree.npop+j]-x[sptree.npop+k];
         }
      FOR(j,com.np) mxo[j]=(mxo[j]*(nused-1)+xo[j])/nused;
      FOR(j,com.np) mx[j]=(mx[j]*(nused-1)+x[j])/nused;
      FOR(i,com.np) FOR(j,com.np) 
         varx[i*com.np+j] += (x[i]-mx[i])*(x[j]-mx[j]) * (nused-1.)/nused;
      if(sptree.nspecies==1)
         FOR(j,data.ngene)
            mtmrca[j]=(mtmrca[j]*(nused-1)+gnodes[j][data.root[j]].age)/nused;
      if(mcmc.print && ir>=0 && (ir+1)%mcmc.output==0) {
         for(j=0;j<com.np; j++) fprintf(fmcmcout," %7.6f",xo[j]);
         if(sptree.nspecies==1)
            for(j=0;j<data.ngene; j++)
               fprintf(fmcmcout," %7.6f", gnodes[j][data.root[j]].age);
         fprintf(fmcmcout," %10.3f\n",lnL);
      }
      if((ir+1)%max2(mcmc.output, mcmc.output*mcmc.nsample/10000)==0) {

         printf("%3.0f%% %5.2f%5.2f%5.2f%5.2f%5.2f \n",
                 (ir+1.)/(mcmc.nsample*mcmc.output)*100.,
                 Ptj/(total*totalTimes), 
                 Ptip/(total*(mcmc.moveinnode?2*totalTimes:totalSeqs)), 
                 Ptheta/(total*sptree.npop),
                 (sptree.nspecies>1?Ptau/(total*(sptree.nspecies-1)):0), Pmix/total);
         for(j=0; j<com.np; j++)
            { printf(" %7.5f", mxo[j]);  if(j==sptree.npop-1) printf(" "); }

         if(sptree.nspecies==1)
            for(j=0,printf(" |");j<min2(4,data.ngene); j++)
               printf(" %7.5f", mtmrca[j]);
         if(mcmc.usedata) printf(" %4.2f", lnL);


         if(mcmc.output*mcmc.nsample>50 && (ir+1)%(mcmc.output*mcmc.nsample/25)==0) {
            printf("  %s\n", printtime(timestr));
            testlnL=1;
            if(fabs(lnL-lnpData(data.lnpDi))>1e-6) error2("lnL not right"); 
            testlnL=0;
         }
      }
   }
   fclose(fmcmcout);
   printf("\nTime used: %s", printtime(timestr));

   if(mcmc.print) {
      printf("\n\nSummarizing, time reset.");
      fprintf(fout,"\n\nSummary of MCMC results\n");
      DescriptiveStatistics(fout, mcmcout, 50, 20);
   }
   return(0);
}

#endif
