#include "qTools.h"

#define MIN(X, Y)  ((X) < (Y) ? (X) : (Y))
#define MAX(X, Y)  ((X) > (Y) ? (X) : (Y))


// Constructor
QTools::QTools(StructureAlignment* sa) 
  : alignment(sa), qScores(0), qPerRes(0), qPower(0.15) {

  return;
}


// Destructor
QTools::~QTools() {

  printf("=>QTools::~QTools()\n");
  int structCount = alignment->getStructureCount();

  if (qScores != 0) {
    for (int i=structCount-1; i>=0; i--) {
      printf("   Hey%d\n",i);
      if (qScores[i] != 0) {
	printf("   delete qScores[%d];\n",i);
	delete qScores[i];
      }
    }
    
    printf("   delete qScores;\n");
    delete qScores;
  }

  if (qPerRes != 0) {
    for (int i=structCount-1; i>=0; i--) {
      printf("   Hey%d\n",i);
      if (qPerRes[i] != 0) {
	printf("   delete qPerRes[%d];\n",i);
	delete qPerRes[i];
      }
    }
    
    printf("   delete qPerRes;\n");
    delete qPerRes;
  }

  printf("<=QTools::~QTools()\n");

  return;
}


// q
//   Calculates a q-value which takes into account the gapped regions
//   Originally written by Patrick O'Donoghue
//   Translated from fortran to C by Michael Januszyk
//   Completely rewritten in C++ by John Eargle
int QTools::q(int ends) {

  //printf("=>QTools::q\n");

  int structCount = alignment->getStructureCount();

  if (qScores == 0) {
    qScores = new float* [structCount];
    for (int i=0; i<structCount; i++) {
      qScores[i] = new float[structCount];
    }  
  }

  // Don't need to initialize qScores with 0s
  
  float** qAln  = new float* [structCount];
  float** qGap  = new float* [structCount];
  int**   qNorm = new int*   [structCount];

  for (int i=0; i<structCount; i++) {
    qAln[i]  = new float[structCount];
    qGap[i]  = new float[structCount];
    qNorm[i] = new int[structCount];
    for (int j=0; j<structCount; j++) {
      qNorm[i][j] = 0;
    }
  }

  getQAln(qAln,qNorm);
  getQGap(qGap,qNorm,ends);
  
  for (int i=0; i<structCount; i++) {
    for (int j=0; j<structCount; j++) {
      if (qNorm[i][j] > 0) {
	qScores[i][j] = (qAln[i][j] + qGap[i][j]) / qNorm[i][j];
      }
      else {
	qScores[i][j] = 0;
      }
      if (i==j) {
	qScores[i][j] = 1.0;
      }
    }
  }

  // XXX - DEBUGGING CODE
  /*
    FILE *out = fopen("q.txt","w");
    fprintf(out,"qAln:\n");
    printMatrix(out,qAln);
    fprintf(out,"\n");
    fprintf(out,"qGap:\n");
    printMatrix(out,qGap);
    fprintf(out,"\n");
    fprintf(out,"qNorm:\n");
    printMatrix(out,qdNorm);
    fprintf(out,"\n");
    fprintf(out,"q:\n");
    printQ(out);
    fclose(out);
  */
  // \XXX

  //printf("<=QTools::q\n");
  return 1;
}


// qPerResidue
//   
int QTools::qPerResidue() {

  //printf("=>QTools::qPerResidue\n");

  int len = alignment->getLength();
  int structCount = alignment->getStructureCount();
  //int norm = 0;

  float numerator = 0.0;
  float denominator = 0.0;
  float score = 0.0;
  float* distances = new float[structCount];

  if (qPerRes == 0) {
    qPerRes = new float* [structCount];
    for (int i=0; i<structCount; i++) {
      //qPerRes[i] = new float[alignment->getStructure(i)->getLength()];
      qPerRes[i] = new float[len];
    }  
  }

  int **qNorm = new int* [structCount];
  for (int i=0; i<structCount; i++) {
    //qNorm[i] = new int[alignment->getStructure(i)->getLength()];
    qNorm[i] = new int[len];
    //for (int j=0; j<alignment->getStructure(i)->getLength(); j++) {
    for (int j=0; j<len; j++) {
      qPerRes[i][j] = 0.0;
      qNorm[i][j] = 0;
    }
  }

  // Traverse columns i and j summing up the contributions to
  //   qAln from differences between interior backbone distances;
  //   ignore adjacent columns
  for (int i=0; i<len-2; i++) {
    for (int j=i+2; j<len; j++) {
      // Get backbone distances for current pair of columns
      getBackboneDistances(distances, i, j);

      // Unset distances for structures if the columns correspond
      //   to adjacent residues
      for (int row=0; row<structCount; row++) {
	if ( alignment->getStructure(row)->alignedToUnalignedIndex(i) ==
	     alignment->getStructure(row)->alignedToUnalignedIndex(j)-1 ) {
	  distances[row] = -1;
	}
      }

      // Calculate scores for current pair of columns
      // Traverse rows struct1 and struct2 calculating qPerRes
      //   scores for each set of rows and columns
      for (int struct1=0; struct1<structCount-1; struct1++) {
	for (int struct2=struct1+1; struct2<structCount; struct2++) {
	  // if (neither structure is gapped) add new score to qAln
	  if ( distances[struct1] >= 0.0 &&
	       distances[struct2] >= 0.0 ) {
	    numerator = -1 * pow( distances[struct1] - distances[struct2], 2 );
	    denominator = 2 * pow( j-i, qPower );
	    score = exp( numerator/denominator );
	    //printf("numerator: %f, denominator: %f, score: %f\n",numerator,denominator,score);
	    qPerRes[struct1][i] += score;
	    qPerRes[struct2][i] += score;
	    qPerRes[struct1][j] += score;
	    qPerRes[struct2][j] += score;
	    //printf("qAln[%d][%d] = %f\n",struct1,struct2,qAln[struct1][struct2]);
	    qNorm[struct1][i]++;
	    qNorm[struct2][i]++;
	    qNorm[struct1][j]++;
	    qNorm[struct2][j]++;
	  }
	  else {
	    // Complicated process for getting appropriate norms
	    //   
	    Structure* s1 = alignment->getStructure(struct1);
	    Structure* s2 = alignment->getStructure(struct2);
	    
	    if ( distances[struct1] >= 0.0 &&
		 s2->getAlphabet()->isGap(s2->getSymbol(i)) && 
		 !s2->getAlphabet()->isGap(s2->getSymbol(j)) ) {
	    qNorm[struct1][i]++;
	    }
	    else if ( distances[struct1] >= 0.0 &&
		      !s2->getAlphabet()->isGap(s2->getSymbol(i)) &&
		      s2->getAlphabet()->isGap(s2->getSymbol(j)) ) {
	      qNorm[struct1][j]++;
	    }
	    else if (distances[struct2] >= 0.0 &&
		     s1->getAlphabet()->isGap(s1->getSymbol(i)) &&
		     !s1->getAlphabet()->isGap(s1->getSymbol(j)) ) {
	      qNorm[struct2][i]++;
	    }
	    else if (distances[struct2] >= 0.0 &&
		     !s1->getAlphabet()->isGap(s1->getSymbol(i)) &&
		     s1->getAlphabet()->isGap(s1->getSymbol(j))) {
	      qNorm[struct2][j]++;
	    }
	  }
	}
      }
    }
  }

  // Normalize qPerRes scores
  for (int row=0; row<structCount; row++) {
    for (int col=0; col<len; col++) {
      if (qNorm[row][col] != 0) {
	qPerRes[row][col] /= qNorm[row][col];
      }
      else {
	qPerRes[row][col] = 0.0;
      }
    }
  }

  // Delete qNorm
  for (int row=0; row<structCount; row++) {
    delete qNorm[row];
  }
  delete qNorm;

  /*
  int structCount = alignment->getStructureCount();
  int alignLength = alignment->getLength();

  if (qPerRes == 0) {
    qPerRes = new float* [structCount];
    for (int i=0; i<structCount; i++) {
      qPerResidue[i] = new float[alignment->getStructure(i)->getLength()];
    }  
  }

  int **qNorm = new int* [structCount];
  for (int i=0; i<structCount; i++) {
    qNorm[i] = new int[alignment->getStructure(i)->getLength()];
    for (int j=0; j<alignment->getStructure(i)->getLength(); j++) {
      qPerResidue[i][j] = 0.0;
      qNorm[i][j] = 0;
    }
  }
  
  float dist1 = 0.0;
  float dist2 = 0.0;
  float numerator = 0.0;
  float denominator = 0.0;
  float score = 0.0;
  for (int struct1=0; struct1<structCount; struct1++) {
    for (int struct2=0; struct2<structCount; struct2++) {

      if (struct1 != struct2) {
	for (int col1=0; col1<alignLength; col1++) {
	  for (int col2=0; col2<alignLength; col2++) {

	    // if (no gap in struct1 or struct2)
	    if ( !struct1->getAlphabet()->isGap(struct1->getSymbol(col1)) &&
		 !struct2->getAlphabet()->isGap(struct1->getSymbol(col1)) &&
		 !struct1->getAlphabet()->isGap(struct1->getSymbol(col2)) &&
		 !struct2->getAlphabet()->isGap(struct1->getSymbol(col2)) ) {
	      dist1 = getBackboneDistance(struct1,col1,col2);
	      dist2 = getBackboneDistance(struct2,col1,col2);
	      numerator = -1 * pow( dist1 - dist2, 2 );
	      denominator = 2 * pow( abs(j-i), qPower );
	      score = exp( numerator/denominator );
	      
	      qPerResidue[struct1][col1] += score;
	      qNorm[struct1][col1]++;
	    }
	    // else if (no gap in struct1, gap in struct2)
	    else if ( ) {
	      
	    }
	    else {
	      // Mark current qPerResidue as gap
	      
	    }
	    
	  }
	}
      }
    }
  }
  */

  //printf("<=QTools::qPerResidue\n");
  return 1;
}


// printQ
//   
int QTools::printQ(FILE* outfile) {

  if (qScores == 0) {
    return 0;
  }

  int structCount = alignment->getStructureCount();

  for (int i=0; i<structCount; i++) {
    for (int j=0; j<structCount; j++)
      fprintf(outfile,"%6.4f ",(float)(1-qScores[i][j]));
    fprintf(outfile,"\n");
  }

  return 1;
}


// printMatrix - DEBUGGING METHOD
//
int QTools::printMatrix(FILE* outfile, float** mat) {

  int structCount = alignment->getStructureCount();

  for (int i=0; i<structCount; i++) {
    for (int j=0; j<structCount; j++)
      fprintf(outfile,"%6.4f ",mat[i][j]);
    fprintf(outfile,"\n");
  }

  return 1;
}


// printMatrix - DEBUGGING METHOD
//
int QTools::printMatrix(FILE* outfile, int** mat) {

  int structCount = alignment->getStructureCount();

  for (int i=0; i<structCount; i++) {
    for (int j=0; j<structCount; j++)
      fprintf(outfile,"%d ",mat[i][j]);
    fprintf(outfile,"\n");
  }

  return 1;
}


// printQPerResidue
//   
int QTools::printQPerResidue(FILE* outfile) {

  int structCount = alignment->getStructureCount();
  int alignedRes = 0;

  for (int i=0; i<structCount; i++) {
    for (int j=0; j<alignment->getStructure(i)->getUnalignedLength(); j++) {
      alignedRes = alignment->getStructure(i)->unalignedToAlignedIndex(j);
      fprintf(outfile,"%f ",qPerRes[i][alignedRes]);
    }
    fprintf(outfile,"\n");
  }

  return 0;
}


// getQAln
//   Calculate QAln scores for each pair of structures
//   QAln = sum((i<j-2), exp [ - ((r(ij) - r(i'j'))^2 / (2 * sigma(ij)^2)) ]
int QTools::getQAln(float** qAln, int** qNorm) {

  //printf("=>QTools::getQAln\n");

  int len = alignment->getLength();
  int structCount = alignment->getStructureCount();
  //int norm = 0;

  float numerator = 0.0;
  float denominator = 0.0;
  float score = 0.0;
  float* distances = new float[structCount];

  // Zero out qAln
  for (int struct1=0; struct1<structCount; struct1++) {
    for (int struct2=0; struct2<structCount; struct2++) {
      qAln[struct1][struct2] = 0.0;
    }
  }

  // Traverse columns i and j summing up the contributions to
  //   qAln from differences between interior backbone distances;
  //   ignore adjacent columns
  for (int i=0; i<len-2; i++) {
    for (int j=i+2; j<len; j++) {
      // Get backbone distances for current pair of columns
      getBackboneDistances(distances, i, j);

      // Unset distances for structures if the columns correspond
      //   to adjacent residues
      for (int row=0; row<structCount; row++) {
	if ( alignment->getStructure(row)->alignedToUnalignedIndex(i) ==
	     alignment->getStructure(row)->alignedToUnalignedIndex(j)-1 ) {
	  distances[row] = -1;
	}
      }

      // Calculate scores for current pair of columns
      // Traverse rows struct1 and struct2 calculating qAln
      //   scores for each pair of rows
      for (int struct1=0; struct1<structCount-1; struct1++) {
	for (int struct2=struct1+1; struct2<structCount; struct2++) {
	  // if (neither structure is gapped) add new score to qAln
	  if ( distances[struct1] >= 0.0 &&
	       distances[struct2] >= 0.0 ) {
	    numerator = -1 * pow( distances[struct1] - distances[struct2], 2 );
	    denominator = 2 * pow( j-i, qPower );
	    score = exp( numerator/denominator );
	    //printf("numerator: %f, denominator: %f, score: %f\n",numerator,denominator,score);
	    qAln[struct1][struct2] += score;
	    qAln[struct2][struct1] += score;
	    //printf("qAln[%d][%d] = %f\n",struct1,struct2,qAln[struct1][struct2]);
	    qNorm[struct1][struct2]++;
	    qNorm[struct2][struct1]++;
	  }
	}
      }
    }
  }

  //printf("<=QTools::getQAln\n");
  return 1;
}


// getQGap
//   Calculate and return QGap scores for each pair of structures
//   in the alignment
int QTools::getQGap(float** qGap, int** qNorm, int ends) {

  //printf("=>QTools::getQGap\n");

  int structCount = alignment->getStructureCount();

  float qGapPair1 = 0.0;   // QGap score for a pair of structures
  float qGapPair2 = 0.0;   // QGap score for a pair of structures

  for (int struct1=0; struct1<structCount; struct1++) {
    for (int struct2=0; struct2<structCount; struct2++) {
      qGap[struct1][struct2] = 0.0;
    }
  }
  
  for (int struct1=0; struct1<structCount-1; struct1++) {
    for (int struct2=struct1+1; struct2<structCount; struct2++) {
      // getQGap(x,y,e) gets score for x with respect to y
      //   e is whether the ends should be counted or not
      qGapPair1 = getQGap(struct1, struct2, qNorm, ends);
      qGapPair2 = getQGap(struct2, struct1, qNorm, ends);
      if (qGapPair1 >= 0.0 && qGapPair2 >= 0.0) {
	//printf("qGapPair1: %f, qGapPair2: %f\n",qGapPair1,qGapPair2);
	qGap[struct1][struct2] = qGapPair1 + qGapPair2;
	qGap[struct2][struct1] = qGapPair1 + qGapPair2;
	//qNorm[struct1][struct2]++;
	//qNorm[struct2][struct1]++;
      }
    }
  }

  //printf("<=QTools::getQGap\n");
  return 1;
}


// getQGap
//   Return the QGap score for a structure with respect
//   to a second structure in the alignment
float QTools::getQGap(int s1, int s2, int** qNorm, int ends) {

  //printf("  =>QTools::getQGap(%d,%d,%d)\n",s1,s2,ends);

  int len = alignment->getLength();
  int gapHead = -1;
  int gapTail = getGapTail(gapHead, s1, s2);
  //int norm = 0;
  float dist1 = 0.0;
  float dist2 = 0.0;
  float numerator = 0.0;
  float denominator = 0.0;
  float score1 = 0.0;
  float score2 = 0.0;
  float qGap = 0.0;
  AlignedStructure* struct1 = alignment->getStructure(s1);
  AlignedStructure* struct2 = alignment->getStructure(s2);

  if (struct1 == 0 || struct2 == 0) {
    printf("    dying\n");
    printf("  <=QTools::getQGap\n");
    return -1.0;
  }

  gapHead = -1;
  gapTail = getGapTail(gapHead, s1, s2);
  for (int i=0; i<len; i++) {
    //gapHead = -1;
    //gapTail = getGapTail(gapHead, s1, s2);
    if (i == gapTail) {
      gapHead = getGapHead(gapTail, s1, s2);
      gapTail = getGapTail(gapHead, s1, s2);
    }
    // if (struct1 has insertion (gap in struct2))
    if ( !struct1->getAlphabet()->isGap(struct1->getSymbol(i)) &&
	  struct2->getAlphabet()->isGap(struct2->getSymbol(i)) ) {
      for (int j=0; j<len; j++) {
	// if (the structure pair have aligned elements at one end &&
	//     measuring non-adjacent residues)
	if ( !struct1->getAlphabet()->isGap(struct1->getSymbol(j)) &&
	     !struct2->getAlphabet()->isGap(struct2->getSymbol(j)) &&
	     abs( alignment->getStructure(s1)->alignedToUnalignedIndex(i) -
		  alignment->getStructure(s1)->alignedToUnalignedIndex(j) ) > 1 &&
	     abs( alignment->getStructure(s2)->alignedToUnalignedIndex(i) -
		  alignment->getStructure(s2)->alignedToUnalignedIndex(j) ) > 1) {
	  
	  // Head case - haven't reached aligned elements yet
	  if (gapHead == -1 && ends != 0) {
	    //printf("  Headcase\n");
	    dist1 = getBackboneDistance(s1,i,j);
	    dist2 = getBackboneDistance(s2,gapTail,j);
	    numerator = -1 * pow( dist1 - dist2, 2 );
	    denominator = 2 * pow( abs(j-i), qPower );
	    score1 = exp( numerator/denominator );
	    //printf("  dist1: %f, dist2: %f, score1: %f\n",dist1,dist2,score1);
	    qGap +=  score1;
	    qNorm[s1][s2]++;
	    qNorm[s2][s1]++;
	  }
	  // Tail case - no more aligned elements
	  else if (gapTail == -1 && ends != 0) {
	    //printf("  Tailcase\n");
	    dist1 = getBackboneDistance(s1,i,j);
	    dist2 = getBackboneDistance(s2,gapHead,j);
	    numerator = -1 * pow( dist1 - dist2, 2 );
	    denominator = 2 * pow( abs(j-i), qPower );
	    score1 = exp( numerator/denominator );
	    //printf("  dist1: %f, dist2: %f, score1: %f\n",dist1,dist2,score1);
	    qGap +=  score1;
	    qNorm[s1][s2]++;
	    qNorm[s2][s1]++;
	  }
	  // Body case - aligned elements abound
	  else if (gapHead >= 0 && gapTail >= 0) {
	    //printf("  Bodycase\n");
	    dist1 = getBackboneDistance(s1,i,j);
	    dist2 = getBackboneDistance(s2,gapHead,j);
	    numerator = -1 * pow( dist1 - dist2, 2 );
	    denominator = 2 * pow( abs(j-i), qPower );
	    score1 = exp( numerator/denominator );
	    
	    dist2 = getBackboneDistance(s2,gapTail,j);
	    numerator = -1 * pow( dist1 - dist2, 2 );
	    //denominator = 2 * pow( j-i, qPower );
	    score2 = exp( numerator/denominator );
	    
	    //printf("  dist1: %f, dist2: %f, score1: %f, score2: %f\n",dist1,dist2,score1,score2);
	    qGap +=  MAX(score1,score2);
	    qNorm[s1][s2]++;
	    qNorm[s2][s1]++;
	  }
	  //norm++;
	}
      }
    }
  }

  //printf("  <=QTools::getQGap\n");
  return qGap;
}


// getGapHead
//   Return the column before the next gapped region,
//   -1 if end (gaps at head before the first aligned pair)
int QTools::getGapHead(int gapTail, int s1, int s2) {

  //printf("=>QTools::getGapHead\n");

  if (gapTail < 0) {
    //printf("   Nope1\n");
    //printf("<=QTools::getGapHead\n");
    return -1;
  }

  int len = alignment->getLength();
  AlignedStructure* struct1 = alignment->getStructure(s1);
  AlignedStructure* struct2 = alignment->getStructure(s2);

  int i = gapTail;
  for (; i<len-1; i++) {
    if ( struct1->getAlphabet()->isGap(struct1->getSymbol(i+1)) ||
	 struct2->getAlphabet()->isGap(struct2->getSymbol(i+1)) ) {
      //printf("   Yep1\n");
      //printf("<=QTools::getGapHead\n");
      return i;
    }
  }

  //printf("   Yep2\n");
  //printf("<=QTools::getGapHead\n");
  return i;
}


// getGapTail
//   Return the next column where struct1 and struct2 are aligned,
//   -1 if end (gaps at tail after the last aligned pair)
int QTools::getGapTail(int gapHead, int s1, int s2) {
  
  //printf("=>QTools::getGapTail\n");

  if (gapHead < -1) {
    //printf("   Nope1\n");
    //printf("<=QTools::getGapTail\n");
    return -1;
  }

  int len = alignment->getLength();
  AlignedStructure* struct1 = alignment->getStructure(s1);
  AlignedStructure* struct2 = alignment->getStructure(s2);

  for (int i=gapHead+1; i<len; i++) {
    if ( !struct1->getAlphabet()->isGap(struct1->getSymbol(i)) &&
	 !struct2->getAlphabet()->isGap(struct2->getSymbol(i)) ) {
      //printf("   Yep1\n");
      //printf("<=QTools::getGapTail\n");
      return i;
    }
  }

  //printf("   Nope2\n");
  //printf("<=QTools::getGapTail\n");
  return -1;
}


// getBackboneDistances
//   Calculate distance between the two backbone atoms that
//   correspond to the aligned residues in two columns;  Return
//   the distances in an array
int QTools::getBackboneDistances(float* distances, int bb1, int bb2) {

  //printf("=>QTools::getBackboneDistances\n");

  for (int i=0; i<alignment->getStructureCount(); i++) {
    //printf("%d",i);
    distances[i] = getBackboneDistance(i, bb1, bb2);
  }
  //printf("\n");

  //printf("<=QTools::getBackboneDistances\n");
  return 1;
}


// getBackboneDistance
//   Calculate distance between two backbone atoms in a structure;
//   Returns -1 on error
float QTools::getBackboneDistance(int structure, int bb1, int bb2) {

  //printf("=>QTools::getBackboneDistance\n");
  //printf("   bb1: %d, bb2: %d\n",bb1,bb2);

  float dist;
  AlignedStructure* tempStruct = alignment->getStructure(structure);
  if (tempStruct == 0) {
    //printf("   Nope1\n");
    //printf("<=QTools::getBackboneDistance\n");
    return -1;
  }

  Coordinate3D* tempCoord1 = tempStruct->getCoordinate(bb1);
  Coordinate3D* tempCoord2 = tempStruct->getCoordinate(bb2);
  if (tempCoord1 == 0 || tempCoord2 == 0) {
    //printf("   Nope2\n");
    //printf("<=QTools::getBackboneDistance\n");
    return -1;
  }

  dist = sqrt( pow( tempCoord1->getX() -
		    tempCoord2->getX(),2 ) +
	       pow( tempCoord1->getY() -
		    tempCoord2->getY(),2 ) +
	       pow( tempCoord1->getZ() -
		    tempCoord2->getZ(),2 ));
  
  //printf("<=QTools::getBackboneDistance\n");
  return dist;
}