#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <math.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>

#include <nicklib.h>
#include <getpars.h>

#include "badpairs.h"
#include "admutils.h"
#include "mcio.h"  
#include "mcmcpars.h"  
#include "globals.h"  

#define WVERSION   "1000" 


#define MAXFL  50   
#define MAXSTR  512
#define MAXSIZE 8.0e9

typedef enum outputmodetype inputmodetype;  

extern int packmode ;
extern int malexhet ;  
extern int verbose ;  
extern int plotmode ;  

char *trashdir = "/var/tmp" ;
int qtmode = NO ;

/* major data structures */
Indiv **indivmarkers;
SNP **snpmarkers ;
int numsnps, numindivs ; 

char  *genotypename = NULL ;		/* name of genotype file */
char  *snpname = NULL ; 		/* name of SNP file */
char  *indivname = NULL ;		/* name of sample file */
char  *pcaname = NULL ;			/* name of pca file */
char *imode = "eigenstrat";		/* input mode */
char *outputname = NULL ;		/* name of output file */
int numpc = 10;				/* number of principal components
						to correct */

/* 
   If these are to be global, remove them from function parameter lists.
   If they're going to be local, put the rest in (chisq routines)    
*/

int NSAMPLES;
int *outlier;
int L;

inputmodetype inmode;
FILE *fpout; 				/* output file */

void readcommands(int argc, char **argv) ;
void setinmode(inputmodetype *inmode, char *imode);
void readpcafile(double **Vp, int **outlierp, int *kp, int L, int NSAMPLES);
void getphenos(int NSAMPLES, double **iscasep, int *outlier, 
  double **iscasecorrp, int L, double *V);
double compute_chisq(double *source, double *target);
double compute_chisqE(double *source, double *target);


int main(int argc, char **argv)
{
  double *V;
  double *xx;
  double *iscase;
  double *iscasecorr;
  int K;
  int k,m,n;
  int nignore;
  double rowsum, rowsum1;
  double chisq, Echisq, gamma, denom;

  readcommands(argc, argv) ;
  if (outputname != NULL) 
    openit(outputname, &fpout, "w") ;
  else 
    fpout = stdout;
  fprintf(fpout, "Chisq EIGENSTRAT\n");

  setinmode(&inmode, imode);
  packmode = YES;

  numsnps = 
    getsnps(snpname, &snpmarkers, 0.0,  NULL, &nignore, 1) ;

  NSAMPLES = getindivs(indivname, &indivmarkers) ;

  setstatus(indivmarkers, NSAMPLES, "Case") ;
  setgenotypename(&genotypename, indivname) ;
  if (genotypename != NULL)  {
   getgenos(genotypename, snpmarkers, indivmarkers, 
    numsnps, NSAMPLES, nignore) ;
  }

  /*******************************************************************/
  /*  Free memory:  Usually this is done in outfiles:                */
  /*                                                                 */
  /*  nind = rmindivs(&snpmarkers, numsnps, &indmarkers, NSAMPLES);  */
  /*                                                                 */
  /*  But where is the snpmarkers array released?                    */
  /*******************************************************************/

  L = numpc;
  readpcafile(&V, &outlier, &K, L, NSAMPLES);
  getphenos(NSAMPLES, &iscase, outlier, &iscasecorr, L, V);

  /* main eigenstrat loop here */

  if ((xx = (double *)malloc(NSAMPLES*sizeof(*xx))) == NULL)
  {  fprintf(stderr,"CM\n");  exit(1);  }

  for(m=0;m<numsnps;m++)  {

    SNP *cupt = snpmarkers[m];
    for(n=0; n<NSAMPLES; n++)
    {
      int j = getgtypes(cupt,n);

      if(j == 0)       { xx[n] = 0.0; }
      else if(j == 1)  { xx[n] = 0.5; }
      else if(j == 2 ) { xx[n] = 1.0; }
      else if(j == -1) { xx[n] = -100.0; }

      if(outlier[n] == 1) xx[n] = -100.0;

    }

    /* mean-adjust xx */
    rowsum = 0.0; rowsum1 = 0.0;
    for(n=0; n<NSAMPLES; n++)
    {
      if(qtmode == NO && ((outlier[n]) || (xx[n] < -99.0))) continue;
      if(qtmode == YES && ((outlier[n]) || (xx[n] == -100.0))) continue;
      rowsum += xx[n];
      rowsum1 += 1.0;
    }
    for(n=0; n<NSAMPLES; n++)
    {  
      if(outlier[n]) continue;
      if(qtmode == NO)  {
        if (xx[n] < -99.0) 
          xx[n] = -100.0; /* still keep track */
        else 
	  xx[n] -= rowsum/rowsum1;
      }
      else  {
        if (xx[n] == -100.0) 
          xx[n] = -100.0; /* still keep track */
        else 
	  xx[n] -= rowsum/rowsum1;
      }
    }

    /* Chisq */
    chisq = compute_chisq(xx,iscase);

    /* EIGENSTRAT */
    for(k=0; k<L; k++)
    {
      gamma = 0.0;
      denom = 0.0;
      for(n=0; n<NSAMPLES; n++) 
      {
        if(qtmode == NO && (outlier[n] || xx[n] < -99.0)) continue;
        if(qtmode == YES && (outlier[n] || xx[n] == -100.0)) continue;
        gamma += xx[n]*V[NSAMPLES*n+k];
        denom += V[NSAMPLES*n+k]*V[NSAMPLES*n+k];
      }
      gamma /= denom;
      for(n=0; n<NSAMPLES; n++) 
      {
        if(qtmode == NO && (outlier[n] || xx[n] < -99.0)) continue;
        if(qtmode == YES && (outlier[n] || xx[n] == -100.0)) continue;
        xx[n] -= gamma*V[NSAMPLES*n+k];
      }
    }
    Echisq = compute_chisqE(xx,iscasecorr);

    if(rowsum1 == 0.0)
    {
      chisq = -1.0; Echisq = -1.0;
    }

    if(chisq >= 0.0) fprintf(fpout,"%.04f",chisq);
    else fprintf(fpout,"NA");
    if(Echisq >= 0.0) fprintf(fpout," %.04f\n",Echisq);
    else fprintf(fpout," NA\n");

    if(NSAMPLES*m > MAXSIZE)
    {
      fprintf(stderr,"OOPS genotype file has > %g genotypes\n",MAXSIZE);
      fprintf(fpout,"OOPS genotype file has > %g genotypes\n",MAXSIZE);
      exit(1);
    }
  }
  return 0;
}


void readcommands(int argc, char **argv) 
{
  int i;
  char *parname = NULL ;
  phandle *ph ;

  while ((i = getopt (argc, argv, "p:vV")) != -1) {

    switch (i)
      {

      case 'p':
	parname = strdup(optarg) ;
	break;

      case 'v':
	printf("version: %s\n", WVERSION) ; 
	break; 

      case 'V':
	verbose = YES ;
	break; 

      case '?':
	printf ("Usage: bad params.... \n") ;
	fatalx("bad params\n") ;
      }
  }

         
   pcheck(parname,'p') ;
   printf("parameter file: %s\n", parname) ;
   ph = openpars(parname) ;
   dostrsub(ph) ;

   getint(ph, "packmode:", &packmode) ; // controls internals 
   getstring(ph, "genotypename:", &genotypename) ;
   getstring(ph, "snpname:", &snpname) ;
   getstring(ph, "indivname:", &indivname) ;
   getstring(ph, "pcaname:", &pcaname) ;
   getint(ph, "numpc:", &numpc) ;
   getint(ph, "numeigs:", &numpc) ;
   getint(ph, "qtmode:", &qtmode) ;
   getint(ph, "hashcheck:", &hashcheck) ;

   getstring(ph, "outputname:", &outputname);

   writepars(ph) ;
   closepars(ph) ;

}


void setinmode(inputmodetype *inmode, char *imode)  {
  char *ss = strdup(imode);
  int len = strlen(ss);

  int i;
  for(i=0;i<len;i++)  {
    ss[i] = tolower(ss[i]);
  }

  *inmode = EIGENSTRAT;			/* default */
  if ( strcmp(ss, "eigenstrat") == 0 )  
    *inmode = EIGENSTRAT;
  if ( strcmp(ss, "ped") == 0 )  
    *inmode = PED;
  if ( strcmp(ss, "packedped") == 0 )  
    *inmode = PACKEDPED;
  if ( strcmp(ss, "ancestrymap") == 0 )  
    *inmode = ANCESTRYMAP;
  if ( strcmp(ss, "packedancestrymap") == 0 )  
    *inmode = PACKEDANCESTRYMAP;

}

void readpcafile(double **Vp, int **outlierp, int *kp, int L, 
  int NSAMPLES)  {

  /* Build V[] and K and outlier[] */
  char *PCAFILE = strdup(pcaname);
  double tempdouble;
  int K;
  FILE *fppca;
  double *V;
  int *outlier;
  int noutlier;
  int x,n,k;

  if ((V = (double *) malloc(NSAMPLES*NSAMPLES*sizeof(*V))) == NULL)  
  {  fprintf(stderr, "CM\n");  exit(1);  }
  if ((outlier = (int *)malloc(NSAMPLES*sizeof(*outlier))) == NULL)
  {  fprintf(stderr,"CM\n");  exit(1);  }


  if( (fppca = fopen(PCAFILE, "r")) == NULL)
  {
    fprintf(stderr,"Could not open input file %s\n", PCAFILE);  
    exit(1);
  }

  fscanf(fppca,"%d",&K);
  if(L > K)
  {
    fprintf(stderr,"OOPS l=%d is larger than k=%d in %s\n",L,K,PCAFILE);
    fprintf(fpout,"OOPS l=%d is larger than k=%d in %s\n",L,K,PCAFILE);
    exit(1);
  }
  for(x=0; x<K; x++) fscanf(fppca,"%lf",&tempdouble); /* eigenvalues */
  for(n=0; n<NSAMPLES; n++)
  {
    for(k=0; k<K; k++) fscanf(fppca,"%lf",&V[NSAMPLES*n+k]);
    if(feof(fppca))
    {
      fprintf(stderr,"OOPS: %s contains less than %d times %d entries\n",
        PCAFILE,NSAMPLES,K);
      fprintf(fpout,"OOPS: %s contains less than %d times %d entries\n",
        PCAFILE,NSAMPLES,K);
      exit(1);
    }
    /* check for outliers */
    outlier[n] = 1;
    for(k=0; k<K; k++) { if(V[NSAMPLES*n+k] != 0.0) outlier[n] = 0; }
    if(outlier[n] == 1) noutlier++;
  }
  fscanf(fppca,"%lf",&tempdouble);
  if(!(feof(fppca)))
  {
    fprintf(stderr,"OOPS: %s contains too many entries\n",PCAFILE);
    fprintf(fpout,"OOPS: %s contains too many entries\n",PCAFILE);
    exit(1);
  }
  fclose(fppca);

  *kp = K;
  *Vp = V;
  *outlierp = outlier;
}

void getphenos(int NSAMPLES, double **iscasep, int *outlier, 
  double **iscasecorrp, int L, double *V)  {

  int k,n;
  double *iscase, *iscasecorr;
  double rowsum, rowsum1, gamma, denom ;

  /* allocate iscase */
  if ((iscase = (double *)malloc(NSAMPLES*sizeof(*iscase))) == NULL)  
  {  fprintf(stderr,"CM\n");  exit(1);  }
  if ((iscasecorr = (double *)malloc(NSAMPLES*sizeof(*iscasecorr))) == NULL)  
  {  fprintf(stderr,"CM\n");  exit(1);  }

  /* get phenotypes */
  for(n=0; n<NSAMPLES; n++)
  {

    if ( qtmode == NO )  {

      char *grp = strdup(indivmarkers[n]->egroup);
      int i;
      for(i=0;i<strlen(grp);i++)  
        grp[i] = tolower(grp[i]);

      if ( !strcmp(grp,"case") )  {
        iscase[n] = 1.0;
      }
      else if ( !strcmp(grp,"control") )  {
        iscase[n] = 0.0;
      }
      else if ( indivmarkers[n]->ignore == YES )  {
        iscase[n] = -100.0;
      }
      else  {
        fprintf(stderr,"OOPS bad phenotype %s\n",grp);
        fprintf(fpout,"OOPS bad phenotype %s\n",grp);
        exit(1);
      }
      free(grp);

    }
    else  {
      iscase[n] = indivmarkers[n]->qval;
    }
  }

  /* mean-adjust iscase */
  rowsum = 0.0; rowsum1 = 0.0;
  for(n=0; n<NSAMPLES; n++)
  {
    if((outlier[n] == 1) || (iscase[n] == -100) )
      continue;
    rowsum += iscase[n];
    rowsum1 += 1.0;
  }
  for(n=0; n<NSAMPLES; n++)
  {
    if(outlier[n]) continue;
    if(iscase[n] == -100.0) iscase[n] = -100.0; /* still keep track */
    else iscase[n] -= rowsum/rowsum1;
  }

  /* make iscasecorr */
  for(n=0; n<NSAMPLES; n++) 
  {
    if(outlier[n] == 0) iscasecorr[n] = iscase[n];
  }
  for(k=0; k<L; k++)
  {
    gamma = 0.0;
    denom = 0.0;
    for(n=0; n<NSAMPLES; n++) 
    {
      if(qtmode == NO   && ((outlier[n]) || (iscase[n] < -99.0))) continue;
      if(qtmode == YES  && ((outlier[n]) || (iscase[n] == -100.0))) continue;
      gamma += iscasecorr[n]*V[NSAMPLES*n+k];
      denom += V[NSAMPLES*n+k]*V[NSAMPLES*n+k];
    }
    gamma /= denom;
    for(n=0; n<NSAMPLES; n++) 
    {
      if(qtmode == NO   && ((outlier[n]) || (iscase[n] < -99.0))) continue;
      if(qtmode == YES  && ((outlier[n]) || (iscase[n] == -100.0))) continue;
      iscasecorr[n] -= gamma*V[NSAMPLES*n+k];
    }
  }
  *iscasep = iscase;
  *iscasecorrp = iscasecorr;
}

double compute_chisq(double *source, double *target)
{
  int n;
  double sum1, sumx, sumxx, sumy, sumyy, sumxy, numer, denom1, denom2;
  double corr;

  sum1 = 0.0; sumx = 0.0; sumxx = 0.0; sumy = 0.0; sumyy = 0.0; sumxy = 0.0;
  for(n=0; n<NSAMPLES; n++)
  {
    if(outlier[n]) continue;
    if(qtmode == NO && source[n] < -99.0) continue;
    if(qtmode == YES && source[n] == -100.0) continue;

    if(qtmode == NO && target[n] < -99.0) continue;
    if(qtmode == YES && target[n] == -100.0) continue;

    sumx += source[n];
    sumxx += source[n]*source[n];
    sumy += target[n];
    sumyy += target[n]*target[n];
    sumxy += source[n]*target[n];
    sum1 += 1.0;
  }
  if(sumxx == 0.0) return -1.0;
  if(sumyy == 0.0) return -1.0;
  numer = sumxy/sum1 - (sumx/sum1)*(sumy/sum1);
  denom1 = (sumxx/sum1 - (sumx/sum1)*(sumx/sum1));
  denom2 = (sumyy/sum1 - (sumy/sum1)*(sumy/sum1));
  if(denom1 <= 0.0) return -1.0;
  if(denom2 <= 0.0) return -1.0;

  corr = (numer/sqrt(denom1*denom2));
  return (sum1*corr*corr);
}

double compute_chisqE(double *source, double *target)
{
  int n;
  double sum1, sumx, sumxx, sumy, sumyy, sumxy, numer, denom1, denom2;
  double corr;

  sum1 = 0.0; sumx = 0.0; sumxx = 0.0; sumy = 0.0; sumyy = 0.0; sumxy = 0.0;
  for(n=0; n<NSAMPLES; n++)
  {
    if(outlier[n]) continue;
    if(qtmode == NO && source[n] < -99.0) continue;
    if(qtmode == YES && source[n] == -100.0) continue;

    if(qtmode == NO && target[n] < -99.0) continue;
    if(qtmode == YES && target[n] == -100.0) continue;

    sumx += source[n];
    sumxx += source[n]*source[n];
    sumy += target[n];
    sumyy += target[n]*target[n];
    sumxy += source[n]*target[n];
    sum1 += 1.0;
  }
  if(sumxx == 0.0) return -1.0;
  if(sumyy == 0.0) return -1.0;
  numer = sumxy/sum1 - (sumx/sum1)*(sumy/sum1);
  denom1 = (sumxx/sum1 - (sumx/sum1)*(sumx/sum1));
  denom2 = (sumyy/sum1 - (sumy/sum1)*(sumy/sum1));
  if(denom1 <= 0.0) return -1.0;
  if(denom2 <= 0.0) return -1.0;

  corr = (numer/sqrt(denom1*denom2));
  sum1 = sum1 - ((double)(L+1));
  return (sum1*corr*corr);
}

