#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <search.h>
#include <math.h>
#include <malloc.h>
#include <sys/timeb.h>
#include <time.h>


#define max(x,y) (x>y?x:y)
#define min(x,y) (x<y?x:y)

#define MAXSEQ 20000


char*  get_parameter(int argc, char** argv, char* param);
void   chomp(char* s); 
char*  complement(char* s) ;


int    nucl(char c); 

double lcumhyper(int i, int s1, int s2, int N); 
double cumhyper(int i, int s1, int s2, int N); 
double bico(int n, int k);
double factln(int n);
double gammln(double xx);
double factrl(int n) ;
void   nrerror(char str[]);
double hypergeom(int i, int s1, int s2, int N);
int    CmpFunc(const void* _a, const void* _b);

char** allocate_more_kmers(char** kmers, int n, int max_nb_kmers, int inc_nb_kmers); 


int nucl(char c); 

typedef struct _prefixNode {
  int next[4];    // A, C, G, T
 } prefixNode;

typedef struct _Kmer {
  int   index;      // where is it found
  float score;
} Kmer;



int addWord(char* s, int m, prefixNode* pnode, int* nextfree); 
int existWord(char* s, int m, prefixNode* pnode); 


char* nextSequence(FILE* fp, char** name, int* size, int* nextSequence_started, int* nextSequence_ended, char* nextSequence_currentLine); 
char* nextSequence_currentLine; 
int   nextSequence_started = 0;
int   nextSequence_ended = 0;


int main(int argc, char** argv) {
  
  char* fastafile1;
  char* fastafile2;
  char* kmersfile;
  
  char* mce, *mcb, *c_mce;
  
  int a;
  
  char** kmers;
    
  int    i, s;

  FILE*  fp1, *fp2, *fp;
  char*  name1;
  char*  name2;
  int    size1;
  int    size2;
  char*  seq1;
  char*  seq2;
  
  int    n=0, l, ii, j;

  char*  buff;

  int    nbgenes;


  float* scores;
  int*   a_s1;
  int*   a_s2;
  int*   a_i;
  int    i_limit     = 0;

  int    gap         = 0;
  int    motifsize; 
  char   stmp[50];

  int   nextSequence_started1 = 0;
  int   nextSequence_ended1   = 0;
  char* nextSequence_currentLine1; 
 
  
  int   nextSequence_started2 = 0;
  int   nextSequence_ended2   = 0;
  char* nextSequence_currentLine2; 

  char* local_sites1 = 0;
  char* local_sites2 = 0;

  Kmer* kmers_array;


  // PREFIX TREE VARS
  prefixNode* pnode;
  int         nextfree;
  int         lastnode;
  int*        node2index;
  int         nb1, nb2;

  int singlestrand = 0;

  int k1=-1, k2=-1;
  int stop;


  int   max_nb_kmers = 1000;
  int   inc_nb_kmers = 1000;


  if (argc < 2) {
    printf("Usage : fastcompare -kmers FILE -fasta1 FILE -fasta2 FILE [ -gap INT -limit INT -singlestrand 1 ]\n");
    exit(0);
  }

  nextSequence_currentLine1 = (char*)malloc(50000 * sizeof(char));
  nextSequence_currentLine2 = (char*)malloc(50000 * sizeof(char));

  if (strcmp(get_parameter(argc, argv, "-singlestrand"), "1") == 0) 
    singlestrand = 1;

  kmersfile  = get_parameter(argc, argv, "-kmers");
  
  fastafile1 = get_parameter(argc, argv, "-fasta1");
  fastafile2 = get_parameter(argc, argv, "-fasta2");

  if (strlen(get_parameter(argc, argv, "-limit")) == 0) 
    i_limit = 0;
  else
    i_limit = atoi(get_parameter(argc, argv, "-limit"));


  if (strlen(get_parameter(argc, argv, "-gap")) == 0) 
    gap = 0;
  else
    gap = atoi(get_parameter(argc, argv, "-gap"));

  if (strlen(get_parameter(argc, argv, "-l")) == 0) 
    k1 = -1;
  else {
    k1 = atoi(get_parameter(argc, argv, "-l"));
  }

  
  //
  // START READING THE LIST OF KMERS
  //
  kmers = (char**)malloc(max_nb_kmers * sizeof(char*));
  // create a line buffer
  buff = (char*)malloc(10000 * sizeof(char));
  fp = fopen(kmersfile, "r");
  if (!fp) {
    printf("could not open set data %s\n", kmersfile);
  }
  n = 0;
  motifsize = 0;
  while (!feof(fp)) {
    fscanf(fp, "%s\n", buff);
    kmers[n] = (char*)calloc(100, sizeof(char));   
    strcat(kmers[n], buff);    
    n++;
    // allocate more memory if we need it
    if (n == max_nb_kmers) {
      kmers = allocate_more_kmers(kmers, n, max_nb_kmers, inc_nb_kmers);
      max_nb_kmers += inc_nb_kmers;
    }
  }
  
  // motif size becomes the length of the first k-mer in the list
  motifsize = strlen(kmers[0]);    
  fclose(fp);

  if (k1 > 0) {
    k2 = motifsize - k1;
  }
  
  //
  // END READING THE LIST OF KMERS
  //

  // alloc memory for local storage
  local_sites1    = (char*)malloc(n * sizeof(char));
  local_sites2    = (char*)malloc(n * sizeof(char));
  

  // alloc more memory for calculating hypergeometric p-values
  scores         = (float*)malloc(n * sizeof(float));
  a_s1           = (int*)calloc(n, sizeof(int));
  a_s2           = (int*)calloc(n, sizeof(int));
  a_i            = (int*)calloc(n, sizeof(int));


    
  
  //
  //  START BUILDING A PREFIX TREE  KMER => INDEX
  //
  
  // associates a node to an index
  node2index = (int*)malloc(n * motifsize * sizeof(int));
  
  // we are going to need many nodes ...
  pnode = (prefixNode*)malloc(motifsize * n * (sizeof(prefixNode)));
  pnode[0].next[0] = -1;
  pnode[0].next[1] = -1;
  pnode[0].next[2] = -1;
  pnode[0].next[3] = -1;

  //
  //  add all the words in the list
  //
  nextfree = 1;
  for (i=0; i<n; i++) {
    lastnode = addWord(kmers[i], motifsize, pnode, &nextfree);
    //
    //   at this point, lastnode contains the last node inserted in the tree
    //
    node2index[lastnode] = i;
  }
  
  //
  //  END BUILDING A PREFIX TREE  KMER => INDEX
  //
  
    
  //
  //  open sequence file 1
  //
  fp1 = fopen(fastafile1, "r");
  if (!fp1) {
    printf("cannot open %s ..\n", fastafile1);
    exit(0);
  }

  //
  //  open sequence file 2
  //
  fp2 = fopen(fastafile2, "r");
  if (!fp2) {
    printf("cannot open %s ..\n", fastafile2);
    exit(0);
  }
    
  
  mce   = (char*)calloc(motifsize+2, sizeof(char));
  mcb   = (char*)calloc(motifsize+2, sizeof(char));
  c_mce = (char*)calloc(motifsize+2, sizeof(char));



  //
  //  START SCANNING SPECIES 1 AND 2
  //

  s = 0;
  nbgenes = 0;
  while ( (seq1 = nextSequence(fp1, &name1, &size1, &nextSequence_started1, &nextSequence_ended1, nextSequence_currentLine1)) ) {   

    //printf("Read SEQ 1 %s, %s ..\n", name1, seq1); 

        
    l = strlen(seq1);

    //  initilialize the local store
    memset(local_sites1, 0, n);
    stop = strlen(seq1) - motifsize - gap + 1;
    if (stop < 0)
      stop = 0;
    
    //printf("stop1 = %d\n", stop);
    for (j=0; j<stop; j++) {
      //printf("j=%d\n", j);
      if (gap > 0) {
	
	if (k1 > 0) {
	  strncpy(mce, seq1+j, k1);
	  strncpy(mce + k1, seq1 + j + k1 + gap, k2);

	} else {
	  strncpy(mce, seq1+j,  (motifsize / 2));
	  strncpy(mce+(motifsize / 2), seq1+j + (motifsize / 2) + gap, motifsize - (motifsize / 2));
	}
      } else {
	strncpy(mce, seq1+j, motifsize);
      }


      mce[motifsize] = '\0';
      c_mce          = complement(mce);
      

      nb1 = existWord(mce, motifsize, pnode);

      nb2 = -1;
      if (singlestrand == 0)
	nb2 = existWord(c_mce, motifsize, pnode);

      
      // if any of them has been found, no need to go on
      if ((nb1 > 0) || (nb2 > 0)) {
	
	if (nb1 > 0)
	  i = node2index[nb1];
	else
	  i = node2index[nb2];
	
	// increment s1 for the ith kmer if it has not already been incremented
	if (local_sites1[i] == 0) {
	  a_s1[i] ++;
	  
	  // the ith kmer is present in this upstream region
	  local_sites1[i] = 1; 	  
	}
	

	

      }

      // free memory
      free(c_mce);
    }
    
    
    free(name1);
    free(seq1);

    //
    // SCAN SEQUENCE 2
    //

    seq2 = nextSequence(fp2, &name2, &size2, &nextSequence_started2, &nextSequence_ended2, nextSequence_currentLine2);

    // printf("Read SEQ 2 %s, %s ..\n", name2, seq2); 
    //  initilialize the local store
    memset(local_sites2, 0, n);
 
    
    l = strlen(seq2);
    stop = strlen(seq2) - motifsize - gap + 1;
    if (stop < 0)
      stop = 0;
    //printf("stop2 = %d\n", stop);
    for (j=0; j<stop; j++) {
      
      if (gap > 0) {
	
	if (k1 > 0) {
	  strncpy(mce, seq1+j, k1);
	  strncpy(mce + k1, seq1 + j + k1 + gap, k2);
	} else {
	  strncpy(mce, seq2+j,  (motifsize / 2));
	  strncpy(mce+(motifsize/2), seq2+j + (motifsize / 2) + gap, motifsize - (motifsize / 2));
	}

      } else {
	strncpy(mce, seq2+j, motifsize);
      }
      
      mce[motifsize] = '\0';
      c_mce = complement(mce);
      

      nb1 = existWord(mce, motifsize, pnode);
      
      nb2 = -1;
      if (singlestrand == 0)
	nb2 = existWord(c_mce, motifsize, pnode);
      
      
      // if any of them has been found, no need to go on
      if ((nb1 > 0) || (nb2 > 0)) {
	
	if (nb1 > 0)
	  i = node2index[nb1];
	else
	  i = node2index[nb2];
	
	// increment s2 for this kmer 
	if (local_sites2[i] == 0) {
	  a_s2[i] ++;
	  
	  // the ith kmer is present in this upstream region
	  local_sites2[i] = 1; 	  

	  // the same kmer was present in the orth upstream sequence
	  
	  if (local_sites1[i] > 0) {
	    a_i[i] ++;
	  } 
	  

	}
	
	

      }
      
      free(c_mce);
    }
    
    
    free(name2);
    free(seq2);


    s++;

    nbgenes ++;
    
  }

  fclose(fp1);
  fclose(fp2);

  //
  //  END SCANNING SPECIES 1 AND 2
  //



  //
  // traverse every kmers and calculate conservation score
  //
  kmers_array = (Kmer*)malloc(n * sizeof(Kmer));
  for (i=0; i<n; i++) {
    kmers_array[i].index = i;
    kmers_array[i].score = lcumhyper(a_i[i], a_s1[i], a_s2[i], nbgenes);
  }

  //
  // sort the kmers by their score 
  //
  qsort((void*)kmers_array, n, sizeof(Kmer), CmpFunc);
  
 
  a = 0;
  for (i=0; i<n; i++) {
        
    // get index in initial array
    ii = kmers_array[i].index;
  

    // output the kmer itself (with gap if needed)
    if (gap > 0) {
      
      if (k1 > 0) {

	strncpy(stmp, kmers[ii], k1);  
	stmp[k1] = '\0';
	printf("%s", stmp);	    
	for (j=0; j<gap; j++) 
	  printf("N");
	printf("%s\t", kmers[ii] + k1);

      } else {
	strncpy(stmp, kmers[ii], motifsize/2);  
	stmp[motifsize/2] = '\0';
	printf("%s", stmp);	    
	for (j=0; j<gap; j++) 
	  printf("N");
	printf("%s\t", kmers[ii] + motifsize/2);	    
      }

    } else {
      printf("%s\t", kmers[ii]);
    }	



    // output s1, s2, i and score for this kmer
    printf("%d\t%d\t%d\t%5.4f\n", a_s1[ii], a_s2[ii],  a_i[ii], - kmers_array[i].score);
    
    
    
    a++;
    
    // abort the output if the user asked for the highest scoring kmers
    if ( (i_limit > 0) && (i_limit == a)) {
	break;
    }
  }



  

  for (i=0; i<n; i++) {
    free(kmers[i]);
  }
  free(kmers);
  free(buff);
  free(mce);

  
  free( nextSequence_currentLine1);
  free( nextSequence_currentLine2);
    
  return 0;
}





char* get_parameter(int argc, char** argv, char* param)
{
  int i = 0;
  while ((i < argc) && (strcmp(param, argv[i])))
    i++;
  if (i<argc)
    return (argv[i+1]);
  else
    return "";
}



void chomp(char* s) 
{

  int j;

  j = strlen(s);
  while (j >= 0) {
    if (s[j] == '\n') {
       s[j] = '\0';
       return;
    }
    j--;
  }


}


char* complement(char* s) 
{

  int l = strlen(s);
  char* c;
  char  d;
  int i;
  
  c = (char*)calloc(l+1, sizeof(char));

  for (i=l-1; i>=0; i--) {
    if (s[i] == 'A')
      d = 'T';
    else if (s[i] == 'T')
      d = 'A';
    else if (s[i] == 'G')
      d = 'C';
    else if (s[i] == 'C') 
      d = 'G';
    else 
      d = 'N';
    
    strncat(c, &d, 1);
  }
	 
  return c;
  
}

//
//  returns the next sequence in the file fp
//
char* nextSequence(FILE* fp, char** name, int* size, int* nextSequence_started, int* nextSequence_ended, char* nextSequence_currentLine) 
{
  
  int len = 20000;
  char* seq;
  

  // not yet started, get the first line
  if (*nextSequence_started == 0) {
    fgets(nextSequence_currentLine, len, fp);  
    chomp(nextSequence_currentLine);      
    *nextSequence_started = 1;
 
    
 }
  
  //printf("%s\n", nextSequence_currentLine);

  // if file is finished, return 0
  if (*nextSequence_ended == 1) {
    return 0;
  }

  // if started, line should be filled with ">..."
  if (nextSequence_currentLine[0] == '>') {
    

    *name = strdup(nextSequence_currentLine + 1);
    
    // create a new line
    seq = (char*)calloc(100000, sizeof(char) ); 

    while (1) {
      
      fgets(nextSequence_currentLine, len, fp);    
      
      if (feof(fp)) {
	*nextSequence_ended  = 1;
	return seq;
      }

      chomp(nextSequence_currentLine);      
      if (strlen(nextSequence_currentLine) == 0) { 
	continue;
      }
      
      if (nextSequence_currentLine[0] == '>') {	
	return seq;	
      } else {
	strcat(seq, nextSequence_currentLine);
      }
      
    }

  } else return 0;
  
}




double lcumhyper(int i, int s1, int s2, int N) {
  
  return log(cumhyper(i, s1, s2, N));

}


double cumhyper(int i, int s1, int s2, int N) {
  
  int min = (s1<s2?s1:s2);
  double prod = 0.0;
  int a;
  double tmp = 0.0;

  for (a=i; a<=min; a++) {
    tmp = (double)hypergeom(a, s1, s2, N);
    prod += (double)hypergeom(a, s1, s2, N);
  }
  
  return prod;
}



void nrerror(char str[]) {
  printf("%s\n", str);
}


double hypergeom(int i, int s1, int s2, int N) {
  double factln(int n);
    
  return exp(factln(s1) + factln(N - s1) + factln(s2) + factln(N - s2) - 
	     factln(i) - factln(N) - factln(s1 - i) - factln(s2 - i) 
	     - factln(N - s1 - s2 + i));
    
}



double factln(int n) {
  double gammln(double xx);
  void nrerror(char error_text[]);
  static double a[101];                                        
  if (n < 0) nrerror("Negative factorial in routine factln");
  if (n <= 1) return 0.0;
  if (n <= 100) return a[n] ? a[n] : (a[n]=gammln(n+1.0));                                    
  else return gammln(n+1.0);                                  
}


double gammln(double xx) {
  double x,y,tmp,ser;
  static double cof[6]={76.18009172947146,-86.50532032941677,
			24.01409824083091,-1.231739572450155,
			0.1208650973866179e-2,-0.5395239384953e-5};
  int j;
  y=x=xx;
  tmp=x+5.5;
  tmp -= (x+0.5)*log(tmp);
  ser=1.000000000190015;
  for (j=0;j<=5;j++) ser += cof[j]/++y;
  return -tmp+log(2.5066282746310005*ser/x);
}



double factrl(int n)  { 
  double gammln(double xx); 
  void nrerror(char error_text[]); 
  static int ntop=4; 
  static double a[33]={1.0,1.0,2.0,6.0,24.0}; 
  int j; 
  
  if (n < 0) 
    nrerror("Negative factorial in routine factrl"); 
  if (n > 32) 
    return exp(gammln(n+1.0)); 
  while (ntop<n) { 
    j=ntop++; 
    a[ntop]=a[j]*ntop; 
  } 
  
  return a[n]; 

}



int CmpFunc(const void* _a, const void* _b)
{
  const Kmer* a = (const Kmer*) _a;
  const Kmer* b = (const Kmer*) _b;
  
  if (a->score > b->score) 
    return 1; 
  else if(a->score == b->score) 
    return  0;
  else         
    return -1;
}




/*** PREFIX TREE OPTIONS ***/

int nucl(char c) 
{
  
  if (c == 'A')
    return 0;
  else if (c == 'C')
    return 1;
  else if (c == 'G')
    return 2;
  else if (c == 'T')
    return 3;
  else return -1;
  
  
}






// add word to the prefix tree (if it is not there)
int addWord(char* s, int m, prefixNode* pnode, int* nextfree) 
{

  //printf("Adding %s to the prefix tree ..\n", s);
  int i, k;

  int j = 0;

  for (i=0; i<m; i++) {
    k = pnode[j].next[nucl(s[i])];
    if (k == -1) {
      pnode[j].next[nucl(s[i])] = *nextfree;

      pnode[*nextfree].next[0] = -1;
      pnode[*nextfree].next[1] = -1;
      pnode[*nextfree].next[3] = -1;
      pnode[*nextfree].next[2] = -1;

      (*nextfree)++;
    }
    j = pnode[j].next[nucl(s[i])];
  }

  return j;
}


// add word to the prefix tree (if it is not there)
int existWord(char* s, int m, prefixNode* pnode) 
{

  //printf("Adding %s to the prefix tree ..\n", s);
  int i, k;
  int t;
  int j = 0;

  for (i=0; i<m; i++) {
    t = nucl(s[i]);
    
    if (t == -1)
      return -2;

    k = pnode[j].next[nucl(s[i])];
    if (k == -1)
      return -1;
  
    j = pnode[j].next[nucl(s[i])];

  }

  return j;
}


//
//  function to allocate more kmers
//
char** allocate_more_kmers(char** kmers, int n, int max_nb_kmers, int inc_nb_kmers) 
{

  char** more_kmers;
  int    i;

  more_kmers = (char**)malloc( (max_nb_kmers + inc_nb_kmers) * sizeof(kmers));
  for (i=0; i<n; i++) {
    more_kmers[i] = strdup(kmers[i]); 
  }
  
  return more_kmers;

}