/***** hrmpg2.c     ************************************************
*      (modified from Hudson's ms_ss.c and rmhap_lik.c)

Generates samples with specified number of segregating sites and calculates
the distributions of the number of haplotypes (W) and the minimum number
of inferred recombination events (RM).  Does this for all specified
values of 4Nr.  Calculates MLE of 4Nr based on W alone, RM alone, and
both.  Note that likW[nsam] diverges, so we arbitrarily set
likW[nsam] = likW[nsam - 1].  Model of exp. growth considered.

Usage:

hrmpg2 npop samplesize[1] . . . samplesize[npop] nsites segsites Rmin Rmax pts alpha t0(if alpha!=0.0) (migrate if npop>1) howmany

where these arguments are:
  npop: Number of populations
  samplesize[i]: Number of individuals sampled from the i-th island
  nsites: Number of sites between which recombination can occur (# of bp's) 
  segsites: Number of polymorphic sites in the total sample
  Rmin: Minimum value of 4Nr considered.
  Rmax: Maximum value of 4Nr considered.
  pts: Number of different values of 4Nr considered.
  alpha: Growth rate
  t0: Time (in units of 4N generations, where N is the current pop. size)
      since the start of growth
  howmany: Howmany samples to generate.

Output consists of the command line arguments followed by the following 
variables and arrays:

  likW[a] = value of 4Nr that maximizes the likelihood of observing exactly 
            a different haplotypes.

  likRM[a] = value of 4Nr that maximizes the likelihood of observing RM = a.

  likRMW[a][b] = value of 4Nr that maximizes the likelihood of observing
                 RM = a and W = b.

  hap[a][b] = number of trials with 4Nr = Rmin + a * (Rmax - Rmin) / (pts - 1)
              with exactly b different haplotypes.

  RM[a][b] = number of trials with 4Nr = Rmin + a * (Rmax - Rmin) / (pts - 1)
             with RM = b.

  RMhap[a][b][c] = # of trials with 4Nr = Rmin + a * (Rmax - Rmin) / (pts - 1)
                   with RM = b and W = c.

*	
***************************************************************************/

#include "my.h"
#define SITESINC 10
#define SSMAX 1000
#include <time.h>

unsigned maxsites = SITESINC;

struct node{
  int abv;
  int ndes;
  float time;
};

struct segl {
  int beg;
  struct node *ptree;
  int next;
};

double *posit;
char **haplist; /* Dummy variable used in hapcount */

/* New function declarations */
int hapcount (char **p, int q, int r, int s); 
int min_rec (char **a, int b, int c, int d, int e);

main(argc,argv)
     int argc;
     char *argv[];
{
  int nsam ,nsites, i, ns, howmany, npop, *config, pop, count, print_flag, narg ;
  double  theta, r, mig_rate, estimator, gst, alpha, within, between ;
  char **list, **cmatrix();
  FILE *pf, *fopen(), *pfseed;
  struct runsum  est_m, gst_m, within_m, between_m ;
  double ttot, alphag, t0;
  int segsites;
  unsigned short seedv[3], *pseed, *seed48();

  /* New variables, different from Hudson's program */
  
  int **hap, **RM, ***RMhap, a, b, c, d, x, pts, ctime, dummy, dummy2;
  int Hctr, Wctr, RMctr, RMWctr, Wmax, RMmax, RMWmax;
  double *likW, *likRM, **likRMW, R, Rmin, Rmax, C, chud();
  double Hsum, Hsqrsum, Hmsesum, Hmean, Hmeansqr, Hvar, Hmse;
  double Wsum, Wsqrsum, Wmsesum, Wmean, Wmeansqr, Wvar, Wmse;
  double RMsum, RMsqrsum, RMmsesum, RMmean, RMmeansqr, RMvar, RMmse;
  double RMWsum, RMWsqrsum, RMWmsesum, RMWmean, RMWmeansqr, RMWvar, RMWmse;
  time_t start, finish;
  clock_t cstart, cfinish;
  long seed1;
  int Min_rec();
  
  time (&start); cstart = clock();

  THINK_C_CL ;
  if (argc == 1) {
    printf("Usage: \nhrmpg2.1 npop samplesize[1] . . . samplesize[npop] nsites segsites Rmin Rmax pts alpha t0(if alpha!=0.0) (migrate if npop>1) howmany\n");
    exit(0);
  }

  npop = atoi( argv[1] );
  config = (int *)malloc((unsigned)((npop+1)*sizeof(int)));
  for( pop=0, narg=1 ; pop<npop; pop++) {
    config[pop] = atoi( argv[++narg] );
    if( config[pop] < 0 ) break;
  }
  nsites = atoi( argv[++narg]);
  segsites = atoi( argv[++narg] ) ;
  Rmin = atof (argv[++narg]);
  Rmax = atof (argv[++narg]);
  pts = atoi (argv[++narg]);
  alphag = atof( argv[++narg]);
  if( alphag != 0.0 ) t0 = atof( argv[++narg]);
  if( npop > 1 )  mig_rate = atof( argv[++narg]);
  else mig_rate = 0.0 ;
  howmany = atoi( argv[++narg]);

  pfseed = fopen("seedms","r");
  for(i=0;i<3;i++) fscanf(pfseed," %hd",seedv+i);
  fclose( pfseed);
  seed48( seedv );

  while( pop<npop ){ config[pop++]=0 ;  }
  pf = stdout ;
  alpha = npop/(npop-1.);
  alpha *= alpha ;
  nsam = 0 ;
  for(i=0;i<npop; i++) nsam += config[i] ;

  
  /* Initialize new variables */
  haplist = (char **) malloc ((unsigned) (nsam * sizeof (char *)));
  for (a=0; a<nsam; ++a)
    haplist[a]=(char *) malloc ((unsigned) ((segsites+1)*sizeof (char)));
  likW = (double *) malloc ((nsam+1) * sizeof (double));
  likRM = (double *) malloc (segsites * sizeof (double));
  likRMW = (double **) malloc (segsites * sizeof (double *));
  for (a=0; a<segsites; ++a)
    likRMW[a] = (double *) malloc ((nsam + 1) * sizeof (double));
  hap = (int **) malloc ((unsigned) (pts * sizeof (int *)));
  RM = (int **) malloc ((unsigned) pts * sizeof (int *));
  RMhap = (int ***) malloc ((unsigned) pts * sizeof (int **));
  for (a=0; a<pts; ++a) {
    hap[a] = (int *) malloc ((unsigned) ((nsam+1) * sizeof (int)));
    RM[a] = (int *) malloc ((unsigned) segsites * sizeof (int));
    RMhap[a] = (int **) malloc ((unsigned) segsites * sizeof (int *));
    for (b=0; b<segsites; ++b)
      RMhap[a][b] = (int *) malloc ((unsigned) (nsam+1) * sizeof (int));
  }

  
  for (a=0; a<pts; ++a) {
    for (b=0; b<=nsam; ++b) 
      hap[a][b] = 0;
    for (b=0; b<segsites; ++b) {
      RM[a][b] = 0;
      for (c=0; c<=nsam; ++c)
	RMhap[a][b][c] = 0;
    }
  }

  /* Print command line arguments in output */
  printf("\nhrmpg2 %d  ",npop);
  for(pop=0; pop<npop; pop++){   printf("%d ",config[pop]); }
  printf("%d ",nsites);
  printf("%d ", segsites);
  printf("%f %f %d %lf ", Rmin, Rmax, pts, alphag);
  if( alphag != 0.0 ) printf(" %lf  ",t0) ;
  if( npop > 1) { printf("%f  ",mig_rate);}
  printf("%d %d %d %d\n",howmany, seedv[0], seedv[1], seedv[2] );

  maxsites = segsites +2;
  list = cmatrix(nsam,maxsites+1);
  posit = (double *)malloc((unsigned)(maxsites*sizeof(double)));
  if(posit==NULL) perror("malloc error. main");
  
  for (x=0; x<pts; ++x) { /* Loop is for each value of 4Nr */
    r = Rmin + (x * (Rmax - Rmin))/ (float) (pts - 1);
    for (count=0; count<howmany; ++count)  {  /* Loop is for each trial */
      gensam(npop,nsam,config,nsites,segsites,r,mig_rate,list,alphag,t0);


	
/*  This prints out simulated sequences 
printf("\n");    
for (i=0; i<nsam; i++)
printf("%s\n", list[i]);
*/
      dummy = hapcount (list, nsam, 0, segsites - 1);
      ++hap[x][dummy];
      dummy2 = Min_rec (list, 0, segsites, nsam);
      ++RM[x][dummy2];
      ++RMhap[x][dummy2][dummy];
    }
  }

  pfseed = fopen("seedms","w");
  pseed = seed48(seedv);
  fprintf(pfseed,"%d %d %d\n",pseed[0], pseed[1],pseed[2]);


  for (a=2; a<nsam; ++a) {
    b=0;
    for (c=1; c<pts; ++c)
      if (hap[c][a]>hap[b][a])
	b = c;
    if (hap[b][a]>0)
      likW[a] = Rmin + (b * (Rmax - Rmin)) / (float) (pts - 1);
    else
      likW[a] = -1.0;
  }
  likW[nsam] = likW[nsam - 1];

  for (a=0; a<segsites; ++a) {
    b=0;
    for (c=1; c<pts; ++c)
      if (RM[c][a]>RM[b][a])
	b = c;
    if (RM[b][a]>0)
      likRM[a] = Rmin + (b * (Rmax - Rmin)) / (float) (pts - 1);
    else
      likRM[a] = -1.0;
  }



  for (d=0; d<segsites; ++d) {
    for (a=2; a<=nsam; ++a) {
      b=0;
      for (c=1; c<pts; ++c)
	if (RMhap[c][d][a]>RMhap[b][d][a])
	  b = c;
      if (RMhap[b][d][a]>0)
	likRMW[d][a] = Rmin + (b * (Rmax - Rmin)) / (float) (pts - 1);
      else
	likRMW[d][a] = -1.0;
    }
  }

  time (&finish); cfinish = clock();
  ctime = (cfinish - cstart) / CLOCKS_PER_SEC;
  printf("\n");
  printf("Runtime, ctime = %d %d\n\n", (finish - start), ctime);

  for (a=2; a<=nsam; ++a)
    printf("likW[%d] = %f\n", a, likW[a]);
  printf("\n");
  for (a=0; a<segsites; ++a)
    printf("likRM[%d] = %f\n", a, likRM[a]);
  printf("\n");
  for (a=0; a<segsites; ++a) {
    c = 0;
    for (b=2; b<=nsam; ++b) {
      printf("likRMW[%d][%d] = %f\t", a, b, likRMW[a][b]);
      ++c;
      if (c % 2 == 0)
	printf("\n");
    }
    printf("\n\n");
  }

  for (a=0; a<pts; ++a) {
    for (b=2; b<=nsam; ++b)
      printf("hap[%d][%d] = %d\t", a, b, hap[a][b]);
    printf("\n\n");
  }
  for (a=0; a<pts; ++a) {
    for (b=0; b<segsites; ++b)
      printf("RM[%d][%d] = %d\t", a, b, RM[a][b]);
    printf("\n\n");
  }
  for (a=0; a<pts; ++a) 
    for (b=0; b<segsites; ++b) {
      for (c=2; c<=nsam; ++c)
	printf("RMhap[%d][%d][%d] = %d\t", a, b, c, RMhap[a][b][c]);
      printf("\n\n");
    }
}



int hapcount (char **list, int nsam, int first, int last)
{
  int a, b, c, hlist = 0;

  for (b=0; b<nsam; ++b) {
    for (c=0; c<hlist; ++c) {
      for (a=first; a<=last && list[b][a] == haplist[c][a]; ++a)
	;
      if (a == (last + 1))
	break;
    }
    if (c==hlist) {
      for (a=first; a<=last; ++a)
	haplist[c][a] = list[b][a];
      hlist++;
    }
  }
  return (hlist);
}

        int 
gensam( npop,nsam,inconfig, nsites, segsites, r,mig_rate,list, alphag,t0)
        int npop,nsam,inconfig[], nsites , segsites ;
        char **list;
        double  r, mig_rate,  alphag, t0;
{
        int nsegs, h, i, k, j, seg, ns, start, end, len, segsit ;
        struct segl *seglst, *segtre_mig() ; /* used to be: [MAXSEG];  */
        double nsinv,  tseg, tt, ttime() ;
        double *pk ;
        int *ss ;


        nsinv = 1./nsites;
        seglst = segtre_mig(npop,nsam,inconfig,nsites,r,mig_rate, &nsegs, alphag,t0) ;

        pk = (double *)malloc((unsigned)(nsegs*sizeof(double)));
        ss = (int *)malloc((unsigned)(nsegs*sizeof(int)));
        if( (pk==NULL) || (ss==NULL) ) perror("malloc error. gensam.2");


        tt = 0.0 ;
        for( seg=0, k=0; k<nsegs; seg=seglst[seg].next, k++) { 
                end = ( k<nsegs-1 ? seglst[seglst[seg].next].beg -1 : nsites-1 )
;
                start = seglst[seg].beg ;
                len = end - start + 1 ;
                tseg = len/(double)nsites ;
               pk[k] = ttime(seglst[seg].ptree,nsam)*tseg ;
                tt += pk[k] ;
                }

	for (k=0;k<nsegs;k++) pk[k] /= tt ;
        mnmial(segsites,nsegs,pk,ss);
        ns = 0 ;
        for( seg=0, k=0; k<nsegs; seg=seglst[seg].next, k++) { 
                end = ( k<nsegs-1 ? seglst[seglst[seg].next].beg -1 : nsites-1 )
;
                start = seglst[seg].beg ;
                len = end - start + 1 ;
                tseg = len/(double)nsites;
                make_gametes(nsam,seglst[seg].ptree,tt*pk[k]/tseg, ss[k], ns, list);

                free(seglst[seg].ptree) ;
                locate(ss[k],start*nsinv, len*nsinv,posit+ns);   
                ns += ss[k] ;
                }
        for(i=0;i<nsam;i++) list[i][ns] = '\0' ;
        free(pk);
        free(ss);
        return( ns ) ;
}


/* allocates space for gametes (character strings) */
	char **
cmatrix(nsam,len)
	int nsam, len;
{
	int i;
	char **m;

	if( ! ( m = (char **) malloc( (unsigned)( nsam*sizeof( char* )) ) ) )
	   perror("alloc error in cmatrix") ;
	for( i=0; i<nsam; i++) {
		if( ! ( m[i] = (char *) malloc( (unsigned) (len*sizeof( char )) )))
			perror("alloc error in cmatric. 2");
		}
	return( m );
}



	int
locate(n,beg,len,ptr)
	int n;
	double beg, len, *ptr;
{
	int i;

	ordran(n,ptr);
	for(i=0; i<n; i++)
		ptr[i] = beg + ptr[i]*len ;

}


/**** ordran.c  ***/
/*
	int
ordran(n,pbuf)
	int n;
	double pbuf[];
{
	ranvec(n,pbuf);
	order(n,pbuf);
	return;
}


	int
ranvec(n,pbuf)
	int n;
	double pbuf[];
{
	int i;
	double ran1();

	for(i=0; i<n; i++)
		pbuf[i] = ran1();

	return;
}

	int
order(n,pbuf)
	int n;
	double pbuf[];
{
	int gap, i, j;
	double temp;

	for( gap= n/2; gap>0; gap /= 2)
	   for( i=gap; i<n; i++)
		for( j=i-gap; j>=0 && pbuf[j]>pbuf[j+gap]; j -=gap) {
		   temp = pbuf[j];
		   pbuf[j] = pbuf[j+gap];
		   pbuf[j+gap] = temp;
		   }
}
*/