#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <math.h>
#include <stdlib.h>

#define MAXSIZE 2000000000
int *outlier, noutlier;
int K, L, NSAMPLES, nSNP;

int
main (int argc, char **argv)
{
  int k, x, m, n, i, nflags;
  double *xx, *iscase;
  double *iscasecorr, gamma, rowsum, rowsum1, *V, denom;
  double tempdouble;
  double chisq, Echisq;
  char Xchar;
  FILE *fp, *fppheno, *fppca, *fpout;
  char *INFILE = NULL;
  char *OUTFILE = NULL;
  char *PCAFILE = NULL;
  char *PHENOFILE = NULL;

  double compute_chisq (double *source, double *target);
  double compute_chisqE (double *source, double *target);

  /* set default values */
  L = 10;

  /* process flags */
  nflags = 0;
  while ((i = getopt (argc, argv, "i:j:p:l:o:")) != -1) {
    switch (i) {
    case 'i':			/* input genotype file */
      INFILE = (char *) strdup (optarg);
      nflags++;
      break;
    case 'j':			/* input phenotype file */
      PHENOFILE = (char *) strdup (optarg);
      nflags++;
      break;
    case 'p':			/* input phenotype file */
      PCAFILE = (char *) strdup (optarg);
      nflags++;
      break;
    case 'l':
      L = atoi (optarg);	/* number of principal components to correct */
      break;
    case 'o':			/* output file */
      OUTFILE = (char *) strdup (optarg);
      nflags++;
      break;
    }
  }

  if (nflags != 4) {
    fprintf (stderr, "Usage: -i -j -p -o flags must all be specified\n");
    exit (1);
  }

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

  /* print parameters */
  fprintf (fpout, "eigenstrat program run using parameters\n");
  fprintf (fpout, " -i %s\n", INFILE);
  fprintf (fpout, " -j %s\n", PHENOFILE);
  fprintf (fpout, " -p %s\n", PCAFILE);
  fprintf (fpout, " -l %d\n", L);
  fprintf (fpout, " -o %s\n", OUTFILE);
  fprintf (fpout, "\n");
  fprintf (fpout, "Chisq EIGENSTRAT\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);

  /* malloc */
  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 ((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);
  }
  if ((xx = (double *) malloc (NSAMPLES * sizeof (*xx))) == NULL) {
    fprintf (stderr, "CM\n");
    exit (1);
  }

  /* FIRST, do PCAFILE */
  /* Build V[] and K and outlier[] */
  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);
  }

  /* SECOND, do PHENOFILE */
  /* get phenotypes */
  if ((fppheno = fopen (PHENOFILE, "r")) == NULL) {
    fprintf (stderr, "Could not open input file %s\n", PHENOFILE);
    exit (1);
  }
  for (n = 0; n < NSAMPLES; n++) {
    fscanf (fppheno, "%c", &Xchar);
    if (Xchar == '0') {
      iscase[n] = 0.0;
    }
    else if (Xchar == '1') {
      iscase[n] = 1.0;
    }
    else if (Xchar == '9') {
      iscase[n] = -100.0;
    }
    else {
      fprintf (stderr, "OOPS bad phenotype %c\n", Xchar);
      fprintf (fpout, "OOPS bad phenotype %c\n", Xchar);
      exit (1);
    }
    if (feof (fppheno)) {
      fprintf (stderr, "OOPS: %s contains less than %d entries\n",
	       PHENOFILE, NSAMPLES);
      fprintf (fpout, "OOPS: %s contains less than %d entries\n",
	       PHENOFILE, NSAMPLES);
      exit (1);
    }
  }
  fscanf (fppheno, "%c", &Xchar);	/* should be end of file */
  if (((Xchar == '0') || (Xchar == '1') || (Xchar == '9'))
      && (!(feof (fppheno)))) {
    fprintf (stderr, "OOPS: %s contains more than %d entries\n", PHENOFILE,
	     NSAMPLES);
    fprintf (fpout, "OOPS: %s contains more than %d entries\n", PHENOFILE,
	     NSAMPLES);
    exit (1);
  }
  /* mean-adjust iscase */
  rowsum = 0.0;
  rowsum1 = 0.0;
  for (n = 0; n < NSAMPLES; n++) {
    if ((outlier[n] == 1) || (iscase[n] == -100.0))
      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 ((outlier[n]) || (iscase[n] < -99.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 ((outlier[n]) || (iscase[n] < -99.0))
	continue;
      iscasecorr[n] -= gamma * V[NSAMPLES * n + k];
    }
  }

  /* THIRD, do INFILE */
  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') {
	xx[n] = 0.0;
      }
      else if (Xchar == '1') {
	xx[n] = 0.5;
      }
      else if (Xchar == '2') {
	xx[n] = 1.0;
      }
      else if (Xchar == '9') {
	xx[n] = -100.0;
      }
      else if (!(feof (fp))) {
	fprintf (stderr, "OOPS bad char %c at m=%d n=%d\n", Xchar, m, n);
	fprintf (fpout, "OOPS bad char %c at m=%d n=%d\n", Xchar, m, n);
	exit (1);
      }
      if (outlier[n] == 1)
	xx[n] = -100.0;
    }
    if (feof (fp))
      break;
    fscanf (fp, "%c", &Xchar);	/* should be \n character */

    /* mean-adjust xx */
    rowsum = 0.0;
    rowsum1 = 0.0;
    for (n = 0; n < NSAMPLES; n++) {
      if ((outlier[n]) || (xx[n] < -99.0))
	continue;
      rowsum += xx[n];
      rowsum1 += 1.0;
    }
    for (n = 0; n < NSAMPLES; n++) {
      if (outlier[n])
	continue;
      if (xx[n] < -99.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 ((outlier[n]) || (xx[n] < -99.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 ((outlier[n]) || (xx[n] < -99.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");

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

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 (source[n] < -99.0)
      continue;
    if (target[n] < -99.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 (source[n] < -99.0)
      continue;
    if (target[n] < -99.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);
}
