#include <stdio.h>
#include <nicklib.h>
#include <stdlib.h>
#include "eigsubs.h"
#include <unistd.h>

int MAXITER, K, TOPK, NSAMPLES, nSNP;
double SIGMATHRESH;

int main(int argc, char **argv)
{
  int k, n, m, nn, rowvalid, *outlier, i;
  int iter, nonewoutliers, nflags;
  char Xchar;
  double *X, *XTX, rowsum, rowmean = 0, rowmeanbayes = 0, *sigmaoutlier;
  double *eval, *evec, sum, summ, sum1, mean, sdev, sigma;
  FILE *fp, *fpout, *fplog, *fpeval;
  char *INFILE = NULL;
  char *OUTFILE = NULL;
  char *EVALFILE = NULL;
  char *LOGFILE = NULL;

  /* set default values */
  K=10; MAXITER=5; TOPK=10; SIGMATHRESH=6.0;

  /* process flags */
  nflags = 0;
  while((i = getopt(argc,argv,"i:k:o:e:l:m:t:s:")) != -1)
  {
    switch(i)
    {
      case 'i':          /* input file */
        INFILE = (char *) strdup(optarg);
        nflags++; break;
      case 'k':
        K = atoi(optarg); /* number of principal components to output */
        break;
      case 'o':          /* output file */
        OUTFILE = (char *) strdup(optarg);
        nflags++; break;
      case 'e':          /* output eval file */
        EVALFILE = (char *) strdup(optarg);
        nflags++; break;
      case 'l':          /* log file */
        LOGFILE = (char *) strdup(optarg);
        nflags++; break;
      case 'm':
        MAXITER = atoi(optarg); /* max # of outlier removal iterations */
        break;
      case 't':
        TOPK = atoi(optarg); /* # of PCs along which to remove outliers */
        break;
      case 's':
        SIGMATHRESH = atof(optarg); /* # sdev to declare as outlier */
        break;
    }
  }
  if(nflags != 4) 
  { 
    fprintf(stderr,"Usage: -i -o -e -l flags must all be specified\n");
    exit(1);
  }

  /* open output files */
  if( (fpout = fopen(OUTFILE, "w")) == NULL)
  {
    fprintf(stderr,"Could not open output file %s\n", OUTFILE);  exit(1);
  }
  if( (fpeval = fopen(EVALFILE, "w")) == NULL)
  {
    fprintf(stderr,"Could not open output file %s\n", OUTFILE);  exit(1);
  }
  if( (fplog = fopen(LOGFILE, "w")) == NULL)
  {
    fprintf(stderr,"Could not open input file %s\n", LOGFILE);  exit(1);
  }

  /* print parameters */
  fprintf(fplog,"pca program run using parameters\n");
  fprintf(fplog," -i %s\n",INFILE);
  fprintf(fplog," -k %d\n",K);
  fprintf(fplog," -o %s\n",OUTFILE);
  fprintf(fplog," -e %s\n",EVALFILE);
  fprintf(fplog," -l %s\n",LOGFILE);
  fprintf(fplog," -m %d\n",MAXITER);
  fprintf(fplog," -t %d\n",TOPK);
  fprintf(fplog," -s %.03f\n",SIGMATHRESH);
  fprintf(fplog,"\n");

  /* Determine NSAMPLES */
  if( (fp = fopen(INFILE, "r")) == NULL)
  {
    fprintf(stderr,"Could not open input file %s\n", INFILE);  exit(1);
  }
  n = 0;
  while(1)
  {
    fscanf(fp,"%c",&Xchar);
    if(Xchar == '\n') break;
    n++;
  }
  NSAMPLES = n;
  fclose(fp);
  if(K > NSAMPLES-1)
  {
    fprintf(stderr,"OOPS k=%d is too large for only %d samples\n",K,NSAMPLES);
    fprintf(fplog,"OOPS k=%d is too large for only %d samples\n",K,NSAMPLES);
    exit(1);
  }

  /* malloc */
  if((eval = (double *) malloc(NSAMPLES*sizeof(*eval))) == NULL)
    { fprintf(stderr,"CM\n");  exit(1); }
  if((evec = (double *) malloc(NSAMPLES*NSAMPLES*sizeof(*evec))) == NULL)
    { fprintf(stderr,"CM\n");  exit(1); }
  if((outlier = (int *) malloc(NSAMPLES*sizeof(*outlier))) == NULL)
    { fprintf(stderr,"CM\n");  exit(1); }
  if((sigmaoutlier = (double *) malloc(NSAMPLES*sizeof(*sigmaoutlier))) == NULL)
    { fprintf(stderr,"CM\n");  exit(1); }
  if((X = (double *) malloc(NSAMPLES*sizeof(*X))) == NULL)
    { fprintf(stderr,"CM\n");  exit(1); }
  if((XTX = (double *) malloc(NSAMPLES*NSAMPLES*sizeof(*X))) == NULL)
    { fprintf(stderr,"CM\n");  exit(1); }

  nonewoutliers = 0;
  for(n=0; n<NSAMPLES; n++) outlier[n] = 0;
  iter = 0;
  while(nonewoutliers == 0)
  {
    for(n=0; n<NSAMPLES; n++) sigmaoutlier[n] = 0.0;
    /* initialize XTX */
    for(n=0; n<NSAMPLES; n++)
    {
      for(nn=0; nn<NSAMPLES; nn++) 
        XTX[NSAMPLES*n+nn] = 0.0;
    }

    nonewoutliers = 1;
    /* get data */
    if( (fp = fopen(INFILE, "r")) == NULL)
    {
      fprintf(stderr,"Could not open input file %s\n", INFILE);  exit(1);
    }
    m = 0;
    while(1) /* do EVERYTHING for SNP m */
    {
      for(n=0; n<NSAMPLES; n++)
      {
        fscanf(fp,"%c",&Xchar);
        if(Xchar == '0') { X[n] = 0.0; }
        else if(Xchar == '1') { X[n] = 0.5; }
        else if(Xchar == '2') { X[n] = 1.0; }
        else if(Xchar == '9') { X[n] = -100.0; }
        else if(!(feof(fp)))
	{
          fprintf(stderr,"OOPS bad char %c at m=%d n=%d\n",Xchar,m,n);
          fprintf(fplog,"OOPS bad char %c at m=%d n=%d\n",Xchar,m,n);
          exit(1);
	}
        if(outlier[n] == 1) X[n] = -100.0;
      }
      if(feof(fp)) break;
      fscanf(fp,"%c",&Xchar); /* should be \n character */

      /* mean-adjust this SNP */
      rowvalid = 0;
      rowsum = 0.0;
      for(n=0; n<NSAMPLES; n++)
      {
        if(X[n] >= -99.0)
        {
          rowvalid++;
          rowsum += X[n];
        }
      }
      if(rowvalid > 0) 
      {
        rowmean = (rowsum)/((double)(rowvalid));
        rowmeanbayes = (rowsum+0.5)/((double)(1+rowvalid));
      }
      for(n=0; n<NSAMPLES; n++)
      {
        if(X[n] >= -99.0)
        {
	  X[n] -= rowmean;
  	  X[n] /= sqrt(rowmeanbayes*(1.0-rowmeanbayes));
        }
        else
          X[n] = 0.0;
      }
  
      /* update XTX */
      for(n=0; n<NSAMPLES; n++)
      {
        for(nn=n; nn<NSAMPLES; nn++) 
          XTX[NSAMPLES*n+nn] += X[n]*X[nn];
      }
      m++;
    }
    nSNP = m;
    if(K > nSNP-1)
    {
      fprintf(stderr,"OOPS k=%d is too large for only %d SNPs\n",K,nSNP);
      fprintf(fplog,"OOPS k=%d is too large for only %d SNPs\n",K,nSNP);
      exit(1);
    }
    if(iter == 0)
    {
      fprintf(fplog,"nSNP=%d NSAMPLES=%d\n",nSNP,NSAMPLES);
    }

    /* complete XTX */
    for(n=0; n<NSAMPLES; n++)
    {
      for(nn=n; nn<NSAMPLES; nn++) 
        XTX[NSAMPLES*n+nn] /= ((double)nSNP);
    }
    for(n=0; n<NSAMPLES; n++)
    {
      for(nn=0; nn<n; nn++) 
        XTX[NSAMPLES*n+nn] = XTX[NSAMPLES*nn+n];
    }

    /* do eigenanalysis */
    eigvecs(XTX, eval, evec, NSAMPLES); /* eigenvector k is evec[k*NSAMPLES+n] */

    if(iter == MAXITER) break; /* no need to look for outliers */

    /* find outliers */
    for(k=0; k<TOPK; k++)
    {
      sum=0.0; summ=0.0; sum1=0.0;
      for(n=0; n<NSAMPLES; n++)
      {
        if(outlier[n] == 1) continue;
        sum += evec[k*NSAMPLES+n];
        summ += evec[k*NSAMPLES+n]*evec[k*NSAMPLES+n];
        sum1 += 1.0;
      }
      mean = sum/sum1;
      sdev = sqrt(summ/sum1 - mean*mean);
      for(n=0; n<NSAMPLES; n++)
      {
        if(outlier[n] == 1) continue;
        sigma = (evec[k*NSAMPLES+n]-mean)/sdev;
        if(sigma < 0) sigma = -sigma;
        if(sigma > SIGMATHRESH) 
        { 
          if(sigma > sigmaoutlier[n]) sigmaoutlier[n] = sigma;
          nonewoutliers = 0; 
        }
      }
    }
    fprintf(fplog,"Outlier removal iteration %d:\n",iter);
    if(nonewoutliers) fprintf(fplog,"  no outliers detected\n");
    for(n=0; n<NSAMPLES; n++)
    {
      if(sigmaoutlier[n] > 0.0) 
      {
        fprintf(fplog,"  removed outlier individual %d (%.02f sigma)\n",n,sigmaoutlier[n]);
        outlier[n] = 1;
      }
    }
    iter++;
    fclose(fp);
  }

  /* print eval and evec */
  for(k=0; k<NSAMPLES; k++) fprintf(fpeval,"%.06f\n",eval[k]);
  fprintf(fpout,"%d\n",K);
  for(k=0; k<K; k++) fprintf(fpout,"%.04f\n",eval[k]);
  for(n=0; n<NSAMPLES; n++)
  {
    for(k=0; k<K; k++)
    {
      fprintf(fpout," ");
      if(evec[k*NSAMPLES+n] > 0) fprintf(fpout," ");
      fprintf(fpout,"%.04f",evec[k*NSAMPLES+n]);
    }
    fprintf(fpout,"\n");
  }
  return 0;
}
