#include <sstream>
#include "EM.h++"
#include "EM.c++"

/* g++ emTrain.c++ -o emTrain */

#define MAXLINE 94096



int   writeoften = 0;
int   iters = 100000;
int   M     = 2;
char  infile[256];
char  outfile[256];
double TOL=1e-8;

void
parse_args (int argc, char **argv)
{
  int i;

  strcpy(infile, "");
  strcpy(outfile, "em.model");

  // Set the command line arguments
  for(i=1; i<argc; i++)
    {
      if (!strcmp(argv[i], "-complexity"))
	M         = atoi(argv[++i]);
      else if (!strcmp(argv[i], "-tolerance"))
	TOL       = atof(argv[++i]);
      else if (!strcmp(argv[i], "-iters"))
	iters     = atoi(argv[++i]);
      else if (!strcmp(argv[i], "-writeoften"))
	writeoften = 1;
      else if  (!strcmp(argv[i], "-in"))
	strcpy(infile, argv[++i]);
      else if (!strcmp(argv[i], "-out"))
	strcpy(outfile, argv[++i]);
      else 
	{
	  fprintf(stderr,"Usage: %s -in <fname> -complexity <int> -out <fname> -iters <int> -tolerance <float> -writeoften \n",argv[0]);
	  fprintf(stderr,"\n");
	  fprintf(stderr," Gaussian Mixture Modeling with EM.\n Example for running with 3 Gaussians \n./emTrain -in data.txt -complexity 3 -out model.txt\n The output format is a text file with M Gaussians stored as: \n mixing proportion of Gaussian 1 D times on a line\n mean of Gaussian 1 as D scalars on a line \n D by D covariance matrix of Gaussian 1 as D lines of D scalars \n Gaussian 2...\n");
	  fprintf(stderr,"\n Tony Jebara 1997\n");
	  exit (2);
	}
    }
}

int
main(int argc, char *argv[]) 
{
  int     N, D, i;
  int     n, d, d1, m;
  FILE    *fp;
  char    line[MAXLINE+1];
  char    b2[50];
  double  *data, *mu, *R, *COV, *scal, *mixo;
  int     index;

  double invT;
  // Parse Arguments
  parse_args(argc,argv);

  // Open the File to Determine its Dimensions
  if ((fp = fopen(infile,"r"))==NULL)
    {
      fprintf(stderr,"%s can't open %s for input.\n",argv[0],infile);
      exit(0);
    }
  fgets(line,MAXLINE,fp);


  std::istringstream ist(line);

  D = 0;
  while (ist>>b2) D++;
  N = 1;
  while (fgets(line, MAXLINE, fp) != NULL) N++;
  fclose(fp);

  printf("Data size is %d and dimensionality is %d\n",N,D);
  // Allocate the Right Size Data Array
  data = (double *)  malloc(D*N*sizeof(double));
  index = 0;
  
  // Reread the File and Load it into Data Array
  fp = fopen(infile,"r");
  while (fgets(line, MAXLINE, fp) != NULL)
    {
      std::istringstream ist(line);
      while (ist>>b2)
	{
	  data[index] = atof(b2);
	  index++;
	}
    }
  fclose(fp);

  // Allocate and Initialize EM Model
  allocateGaussianMixtureModelEM(D, M, &mu, &R, &COV, &scal, &mixo);
  for (i=0; i<M; i++) mixo[i] = 1.0/((double) M);
  randomInitGaussianMixtureModelEM(data, N, D, M, mu, R, scal);

  double likeli = -9e9;
  double likeliOLD = -9e9;

  for (i=0; i<iters; i++)
    { 
      // Iterate EM 1 Step
      likeli = updateParamsGaussianMixtureModelEM(data, N, D, M, mu, R, COV, scal, mixo);
      printf("%f \n", likeli);
      if ((likeli-likeliOLD)<TOL)
	{
	  i = iters+10;
	  writeoften = 1;
	}
      likeliOLD=likeli;
      if (i==(iters-1)) writeoften = 1;
      if (writeoften>0)
	{
	  // Save Model to Text File
	  fp = fopen(outfile,"w");
	  for (m=0; m<M; m++)
	    {
	      for (d=0; d<D; d++) fprintf(fp,"%e ", mixo[m]);
	      fprintf(fp,"\n");
	      for (d=0; d<D; d++) fprintf(fp,"%e ",mu[m*D+d]);
	      fprintf(fp,"\n");
	      for (d=0; d<D; d++)
		{
		  for (d1=0; d1<D; d1++)
		    fprintf(fp,"%e ",COV[m*D*D + d*D + d1]);
		  fprintf(fp,"\n");
		}
	    }
	  fclose(fp);
	}
    }
  
  // Free Memory
  freeGaussianMixtureModelEM(D, M, &mu, &R, &COV, &scal, &mixo);
  free(data);
  return(1);
}