#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <math.h>
#include <unistd.h>
#include <limits.h>
#include <sys/time.h>
#include <zlib.h>

//#define FIBS_UNIT 1000
#define DEFAULT_NDIGITS 10
#define SZBUF 1024
#define SZBYTE 256
#define N_GENOTYPES 4
#define NA_GENO_CHAR (N_GENOTYPES-1)
#define DEFAULT_ROW_SIZE 100000
#define DEFAULT_SIZE_MATRIX 1000000
#define DEFAULT_SIZE_HEADER 100000
#define DEFAULT_DELIMS " \t\r\n"
#define SZ_LONG_BUF 1000000
#define DEFAULT_TPED_NUM_HEADER_COLS 4
#define DEFAULT_TFAM_NUM_HEADER_COLS 6
#define DEFAULT_TPED_SNPID_INDEX 1
#define DEFAULT_PHENO_NUM_HEADER_COLS 2

struct HFILE {
  int gzflag;       // 1 if gz if used
  int wflag;        // r(0)/w(1) for plain, rb(0)/wb(1) for gz
  int nheadercols;  // # of header columns (0 if nrows=0)
  int nvaluecols;   // # of value cols (0 if nrows=0)
  int nrows;        // # of rows
  FILE* fp;         // plain file handle
  gzFile gzfp;      // gzip file handle
};

// Input routines
void close_file (struct HFILE* fhp);
struct HFILE open_file(char* filename, int gzflag, int wflag);
struct HFILE open_file_with_suffix(char* prefix, char* suffix, int gzflag, int wflag);
void read_matrix_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, int symmetric, int* p_nmiss, unsigned char** matrix, char*** headers);
unsigned char* tokenize_tped_line_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, char* lbuf, unsigned char* values, char** headers, int* p_nvalues, int* p_nmiss );

void emmax_error( const char* format, ... );
void print_help(void);
FILE* readfile(char* filename);

void print_help(void) {
  fprintf(stderr,"Usage: emmax-kin [tpedf]\n");
  fprintf(stderr,"Required parameters\n");
  fprintf(stderr,"\t[tpedf]     : tped file\n");
  fprintf(stderr,"\t[outf]     : output file name\n");
  fprintf(stderr,"Optional parameters\n");
  fprintf(stderr,"\t-d [# digits]  : precision of the kinship values (default : 10)\n");
  //  fprintf(stderr,"\t-r [# snp samples] : randomly sample a subset of snps to reduced the time complexity\n");
  //fprintf(stderr,"\t-c : Compute EIGENSTRAT genotype covariance matrix, instead of IBS matrix\n");
  fprintf(stderr,"\t-v : turn on verbose mode\n");
}

int main(int argc, char** argv) {
  int i, j, k, n, ac0, ac1,  nmiss, nelems, nex, *nexes;
  int verbose, ndigits, tped_nheadercols, tfam_nheadercols, flag_autosomal, rand_fill_flag, ibs_flag, hetero_division_flag;
  unsigned char *snprow;
  char buf[SZBUF];
  int c;
  double f, denom;
  //long *fibs_sums, *scores, mean_score;
  double *fibs_sums, *scores, mean_score;
  char *tpedf, *delims, *lbuf;
  char **tfam_headers, **tped_headers;
  struct HFILE tpedh, tfamh, kinsh;
  struct timeval tv;

  gettimeofday(&tv, NULL);
  srand((unsigned int)tv.tv_usec);
  delims = DEFAULT_DELIMS;
  tped_nheadercols = DEFAULT_TPED_NUM_HEADER_COLS;
  tfam_nheadercols = DEFAULT_TFAM_NUM_HEADER_COLS;
  tped_headers = tfam_headers = NULL;
  tpedf = lbuf = 0;
  flag_autosomal = 1;
  rand_fill_flag = 0;
  hetero_division_flag = 0;
  ibs_flag = 0;

  verbose = 0;
  ndigits = DEFAULT_NDIGITS;
  while ((c = getopt(argc, argv, "d:rsc:vxS:zh")) != -1 ) {
    switch(c) {
    case 'd': // precision of digits
      ndigits = atoi(optarg);
      break;
    case 'r':
      rand_fill_flag = 1;
      break;
    case 's':
      ibs_flag = 1;
      break;
    case 'h':
      hetero_division_flag = 1;
      break;
    case 'v':
      verbose = 1;
      break;
    case 'x':
      flag_autosomal = 0;
      break;
    case 'S':
      srand(atoi(optarg));
      break;
    default:
      fprintf(stderr,"Error : Unknown option unsigned character %c\n",c);
      abort();
    }
  }

  // With no options given, present help but do not abort.
  if (1 == argc) {
    print_help();
    exit(-1);
  }

  if ( ( rand_fill_flag == 1 ) && ( hetero_division_flag == 1 ) ) {
    emmax_error("-r and -h option cannot be combined");
  }

  // Sanity check for the number of required parameters
  if ( argc != optind + 1 ) {
    print_help();
    abort();
  }

  // Read required parameters
  tpedf = argv[optind++];

  if ( verbose) fprintf(stderr,"\nReading TFAM file %s.tfam ....\n",tpedf);

  tfamh = open_file_with_suffix(tpedf, "tfam", 0, 0);
  read_matrix_with_col_headers( &tfamh, tfam_nheadercols, delims, 0, &nmiss, NULL, &tfam_headers);
  n = tfamh.nrows;

  snprow = (unsigned char*)malloc(sizeof(unsigned char)*n);
  tped_headers = (char**)malloc(sizeof(char*)*n);
  lbuf = (char*) malloc(sizeof(char*) * SZ_LONG_BUF);
  //fibs_sums = (long*)calloc(n*(n+1)/2, sizeof(long));
  fibs_sums = (double*)calloc(n*(n+1)/2, sizeof(double));
  if ( hetero_division_flag == 1 ) {
    nexes = (int*)calloc(n*(n+1)/2,sizeof(int));
  }
  else {
    nexes = NULL;
  }
  //scores = (long*)malloc(sizeof(long)*N_GENOTYPES*N_GENOTYPES);
  scores = (double*)malloc(sizeof(double)*N_GENOTYPES*N_GENOTYPES);

  if ( verbose) fprintf(stderr,"\nReading TPED file %s.tped ....\n",tpedf);

  tpedh = open_file_with_suffix(tpedf, "tped", 0, 0);
  tpedh.nheadercols = tped_nheadercols;

  for ( i=0, nex = 0; tokenize_tped_line_with_col_headers( &tpedh, tped_nheadercols, delims, lbuf, snprow, tped_headers, &nelems, &nmiss) != NULL; ++i ) {
    if ( ( verbose ) && ( i % 1000 ) == 0 ) fprintf(stderr,"Reading %d SNPs\n",i);
    if ( ( flag_autosomal == 1 ) && ( ( atoi(tped_headers[0]) == 0 ) || ( atoi(tped_headers[0]) > 22 ) ) ) // if SNP is not in autosomal chromosomes
    {
      ++nex;
      continue;
    }

    if ( nelems != n ) {
      emmax_error("Number of values %d in line %d do not match to %d, the number of columns\n", nelems, tpedh.nvaluecols, n);
    }

    ac0 = 0;
    ac1 = 0;
    for(j=0; j < n; ++j) {
      if ( snprow[j] > 0 ) {
	snprow[j] = (unsigned char)(snprow[j]-2); // encoded as 0 or 2,3,4
	switch(snprow[j]) {
	case 0:
	  ac0 += 2;
	  break;
	case 1:
	  ++ac0;
	  ++ac1;
	  break;
	case 2:
	  ac1 += 2;
	  break;
	default:
	  emmax_error("Unknown allele %s, converted to %d\n",buf,(int)snprow[j]);
	  break;
	}
      }
      else {
	snprow[j] = (unsigned char)NA_GENO_CHAR;
      }
    }
    if ( rand_fill_flag == 1 ) {
      double maf = (double)ac1/(double)(ac0+ac1);

      for(j=0; j < n; ++j) {
	if ( snprow[j] == (unsigned char)NA_GENO_CHAR ) {
	  if ( (rand() / (double) RAND_MAX) > maf ) {
	    if ( (rand() / (double) RAND_MAX) > maf ) {
	      snprow[j] = (unsigned char)2;
	      ac1 += 2;
	    }
	    else {
	      snprow[j] = (unsigned char)1;
	      ++ac0;
	      ++ac1;
	    }
	  }
	  else {
	    if ( (rand() / (double) RAND_MAX) > maf ) {
	      snprow[j] = (unsigned char)1;
	      ++ac0;
	      ++ac1;
	    }
	    else {
	      snprow[j] = (unsigned char)0;
	      ac0 += 2;
	    }
	  }
	}
      }
    }
    
    if ( ibs_flag == 1 ) {
      mean_score = 2.*(double)ac1/(double)(ac0+ac1);

      for(j=0; j < NA_GENO_CHAR; ++j) {
	for(k=0; k < NA_GENO_CHAR; ++k) {
	  scores[j+k*N_GENOTYPES] = 2.-abs(j-k);
	}
      }
      if ( hetero_division_flag == 0 ) {
	scores[NA_GENO_CHAR+0*N_GENOTYPES] = scores[0+NA_GENO_CHAR*N_GENOTYPES] = 2-mean_score;
	scores[NA_GENO_CHAR+1*N_GENOTYPES] = scores[1+NA_GENO_CHAR*N_GENOTYPES] = 1+mean_score-0.5*mean_score*mean_score;
	scores[NA_GENO_CHAR+2*N_GENOTYPES] = scores[2+NA_GENO_CHAR*N_GENOTYPES] = mean_score;
	scores[N_GENOTYPES*N_GENOTYPES-1] = 2.-2.*mean_score*(1.-mean_score/2)*(mean_score*mean_score/4.-mean_score/2.+1);
	  /*for(k=0; k < NA_GENO_CHAR-1; ++k) {
	  scores[NA_GENO_CHAR+k*N_GENOTYPES] = scores[k+NA_GENO_CHAR*N_GENOTYPES] = 2.-fabs(mean_score-(double)k);
	}
	scores[N_GENOTYPES*N_GENOTYPES-1] = mean_score*mean_score-2.*mean_score+2.;*/
      }
      
      for(j=0; j < n; ++j) {
	for(k=0; k <= j; ++k) {
	  if ( ( hetero_division_flag == 1 ) && ( ( snprow[j] == NA_GENO_CHAR ) || ( snprow[k] == NA_GENO_CHAR ) ) ) {
	    ++nexes[j*(j+1)/2+k];
	  }
	  else {
	    fibs_sums[j*(j+1)/2+k] += scores[snprow[j]*N_GENOTYPES+snprow[k]]/2.;
	  }
	}
      }
    }
    else {
      if ( ( ac0 > 0 ) && ( ac1 > 0 ) ) {
	mean_score = 2.*(double)ac1/(double)(ac0+ac1);

	denom = 2.*mean_score*(1.-mean_score/2.);

	for(j=0; j < NA_GENO_CHAR; ++j) {
	  for(k=0; k < NA_GENO_CHAR; ++k) {
	    scores[j+k*N_GENOTYPES] = ((double)j-mean_score)*((double)k-mean_score)/denom;
	  }
	}
	for(k=0; k < NA_GENO_CHAR; ++k) {
	  scores[NA_GENO_CHAR+k*N_GENOTYPES] = scores[k+NA_GENO_CHAR*N_GENOTYPES] = 0;
	}
	scores[N_GENOTYPES*N_GENOTYPES-1] = 0;

	for(j=0; j < n; ++j) {
	  for(k=0; k <= j; ++k) {
	    if ( ( hetero_division_flag == 1 ) && ( ( snprow[j] == NA_GENO_CHAR ) || ( snprow[k] == NA_GENO_CHAR ) ) ) {
	      ++nexes[j*(j+1)/2+k];
	    }
	    else {
	      fibs_sums[j*(j+1)/2+k] += scores[snprow[j]*N_GENOTYPES+snprow[k]];
	    }
	  }
	}
      }
    }
  }
  close_file(&tpedh);

  if ( ibs_flag == 1 ) {
    if ( rand_fill_flag == 1 ) {
      kinsh = open_file_with_suffix( tpedf, "rIBS.kinf", 0, 1 );
    }
    else if ( hetero_division_flag == 1 ) {
      kinsh = open_file_with_suffix( tpedf, "hIBS.kinf", 0, 1 );
    }
    else {
      kinsh = open_file_with_suffix( tpedf, "IBS.kinf", 0, 1 );
    }
  }
  else {
    if ( rand_fill_flag == 1 ) {
      kinsh = open_file_with_suffix( tpedf, "rBN.kinf", 0, 1 );
    }
    else if ( hetero_division_flag == 1 ) {
      kinsh = open_file_with_suffix( tpedf, "hBN.kinf", 0, 1 );
    }
    else {
      kinsh = open_file_with_suffix( tpedf, "BN.kinf", 0, 1 );
    }
  }

  if ( verbose ) fprintf(stderr,"Printing the kinship matrix to file\n");

  for(i=0; i < n; ++i) {
    for(j=0; j < n; ++j) {
      if ( j > 0 ) fprintf(kinsh.fp,"\t");
      if ( i >= j ) {
	if ( hetero_division_flag == 1 ) {
	  f = (double)fibs_sums[i*(i+1)/2+j]/(double)(tpedh.nrows-nex-nexes[i*(i+1)/2+j]);
	}
	else {
	  f = (double)fibs_sums[i*(i+1)/2+j]/(double)(tpedh.nrows-nex);
	}
      }
      else {
	if ( hetero_division_flag == 1 ) {
	  f = (double)fibs_sums[j*(j+1)/2+i]/(double)(tpedh.nrows-nex-nexes[j*(j+1)/2+i]);
	}
	else {
	  f = (double)fibs_sums[j*(j+1)/2+i]/(double)(tpedh.nrows-nex);
	}
      }
      fprintf(kinsh.fp,"%-.*lf",ndigits,f);  
    }
    fprintf(kinsh.fp,"\n");
  }
  close_file(&kinsh);
  free(snprow);
  free(fibs_sums);
  free(lbuf);
  free(tped_headers);
  free(tfam_headers);
  if ( hetero_division_flag == 1 ) {
    free(nexes);
  }
  return 0;
}

void read_matrix_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, int symmetric, int* p_nmiss, unsigned char** matrix, char*** headers) {
  char* lbuf = (char*) malloc(sizeof(char*) * SZ_LONG_BUF);
  int szmat = DEFAULT_SIZE_MATRIX;
  int szheader = DEFAULT_SIZE_HEADER;
  unsigned char* cmat = (unsigned char*) malloc(sizeof(unsigned char) * szmat );
  char** cheaders = (char**) malloc(sizeof(char*) * szheader );
  int nvalues, i, j, nmiss;

  fhp->nheadercols = nheadercols; 
  nmiss = 0;

  while( tokenize_tped_line_with_col_headers(fhp, nheadercols, delims, lbuf, &cmat[fhp->nrows*fhp->nvaluecols], &cheaders[fhp->nrows*fhp->nheadercols], &nvalues, &nmiss) != NULL ) {
    if ( fhp->nrows == 1 ) {
      fhp->nvaluecols = nvalues;
    }
    else if ( fhp->nvaluecols != nvalues ) {
      emmax_error("The column size %d do not match to %d at line %d\n",nvalues,fhp->nvaluecols,fhp->nrows);
    }

    if ( (fhp->nrows+1)*(fhp->nvaluecols) > szheader ) {
      szheader *= 2;
      fprintf(stderr,"Header size is doubled to %d\n",szheader);
      cheaders = (char**) realloc( cheaders, sizeof(char*) * szheader );
    }

    if ( (fhp->nrows+1)*(fhp->nvaluecols) > szheader ) {
      szmat *= 2;
      fprintf(stderr,"Matrix size is doubled to %d\n",szmat);
      cmat = (unsigned char*) realloc( cmat, sizeof(unsigned char) * szmat );
    }
  }
  free(lbuf);

  *p_nmiss = nmiss;
  
  unsigned char* fmat = (unsigned char*) malloc(sizeof(unsigned char)*fhp->nrows*fhp->nvaluecols);
  char** fheaders = (char**) malloc(sizeof(char*)*fhp->nrows*fhp->nheadercols);
  for(i=0; i < fhp->nrows; ++i) {
    for(j=0; j < fhp->nvaluecols; ++j) {
      fmat[i+j*fhp->nrows] = cmat[i*fhp->nvaluecols+j];
    }
    for(j=0; j < fhp->nheadercols; ++j) {
      fheaders[i+j*fhp->nrows] = cheaders[i*fhp->nheadercols+j];
    }
  }
  free(cmat);
  free(cheaders);
  
  if ( matrix != NULL ) {
    if ( *matrix != NULL ) {
      free(*matrix);
    }
    *matrix = fmat;
  }
  
  if ( headers != NULL ) {
    if ( *headers != NULL ) {
      free(*headers);
    }
    *headers = fheaders;
  }
}

unsigned char* tokenize_tped_line_with_col_headers( struct HFILE* fhp, int nheadercols, char* delims, char* lbuf, unsigned char* values, char** headers, int* p_nvalues, int* p_nmiss ) {
  int j;
  char *token;
  unsigned char ctoken;

  char *ret = (fhp->gzflag == 1) ? gzgets(fhp->gzfp, lbuf, SZ_LONG_BUF) : fgets( lbuf, SZ_LONG_BUF, fhp->fp );
  int nmiss = 0;

  if ( ret == NULL ) {
    return NULL;
  }

  if ( fhp->nheadercols != nheadercols ) {
    emmax_error("# of header columns mismatch (%d vs %d) at line %d",fhp->nheadercols,nheadercols,fhp->nrows);
  }

  //fprintf(stderr,"tokenize-line called %s\n",lbuf);

  token = strtok(lbuf, delims);
  for( j=0; token != NULL; ++j ) {
    if ( j < nheadercols ) {
      headers[j] = strdup(token);
    }
    // if zero_miss_flag is set, assume the genotypes are encoded 0,1,2
    // Additively encodes the two genotypes in the following way
    // when (j-nheadercols) is even, 0->MISSING, add 1->0, 2->1
    // when (j-nheadercols) is odd, check 0-0 consistency, and add 1->0, 2->1
    else {
      ctoken = (unsigned char)(token[0]-'0');
      
      if ( ctoken > 2 ) {
	fprintf(stderr,"Unrecognized token %s\n",token);
	abort();
      }
      
      if ( (j-nheadercols) % 2 == 0 ) {
	values[(j-nheadercols)/2] = ctoken;
      }
      else {
	if ( ( ctoken > 0 ) && ( values[(j-nheadercols)/2] == 0 ) ) {
	  fprintf(stderr,"Unmatched token pair 0 %s\n",token);
	  abort();
	}
	else if ( ( ctoken == 0 ) && ( values[(j-nheadercols)/2] > 0 ) ) {
	  fprintf(stderr,"Unmatched token pair - %d 0\n",(int)values[(j-nheadercols)/2]);
	  abort();
	}
	values[(j-nheadercols)/2] += ctoken;
      }
    }
    token = strtok(NULL, delims);
  }
  //fprintf(stderr,"tokenize-line ended %d %d\n",j,nheadercols);
  if ( (j-nheadercols) % 2 != 0 ) {
    fprintf(stderr,"Number of value tokens are not even %d\n",j-nheadercols);
    abort();
  }

  *p_nvalues = (j-nheadercols)/2;
  *p_nmiss += nmiss;
  ++(fhp->nrows);

  if ( j < nheadercols ) {
    fprintf(stderr,"Number of header columns are %d, but only %d columns were observed\n", nheadercols, j);
    abort();
  }

  return values;
}

// open_file_with_suffix()
// - [prefix].[suffix] : file name to open
// - gzflag : gzip flag (use gzfp if gzflag=1, otherwise use fp)
// - wflag : write flag (1 if write mode otherwise read mode
struct HFILE open_file_with_suffix(char* prefix, char* suffix, int gzflag, int wflag) {
  char filename[SZBUF];
  sprintf(filename,"%s.%s",prefix,suffix);
  return open_file(filename,gzflag,wflag);
}

// open_file()
// - filename : file name to open
// - gzflag : gzip flag (use gzfp if gzflag=1, otherwise use fp)
// - wflag : write flag (1 if write mode otherwise read mode)
struct HFILE open_file(char* filename, int gzflag, int wflag) {
  struct HFILE fh;
  fh.gzflag = gzflag;
  fh.wflag = wflag;
  fh.nheadercols = 0;
  fh.nvaluecols = 0;
  fh.nrows = 0;
  if ( gzflag == 1 ) {
    char* mode = (wflag == 1) ? "wb" : "rb";
    fh.gzfp = gzopen(filename,mode);
    fh.fp = NULL;

    if ( fh.gzfp == NULL ) {
      emmax_error("Cannot open file %s for writing",filename);
    }
  }
  else {
    char* mode = (wflag == 1) ? "w" : "r";
    fh.gzfp = (gzFile) NULL;
    fh.fp = fopen(filename,mode);

    if ( fh.fp == NULL ) {
      emmax_error("Cannot open file %s for writing",filename);
    }
  }
  return fh;
}

void emmax_error( const char* format, ... ) {
  va_list args;
  fprintf(stderr, "ERROR: ");
  va_start (args, format);
  vfprintf(stderr, format, args);
  va_end (args);
  fprintf(stderr,"\n");
  abort();
}

void close_file(struct HFILE* fhp) {
  if ( fhp->gzflag == 1 ) {
    gzclose(fhp->gzfp);
    fhp->gzfp = NULL;
  }
  else {
    fclose(fhp->fp);
    fhp->fp = NULL;
  }
}