#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);
}

