// Evolutionary tracing: identification of potential binding sites 
// or other functional sites on the surface of a protein structure

#include <iostream.h>
#include <fstream.h>
#include <new>
#include <string.h>
#include "/home/jpei/bin/ET/align_info.h"
#include "/home/jpei/bin/ET/TNode.h"
#include "/home/jpei/bin/ET/arraystack.h"
#include "/home/jpei/bin/ET/arrayqueue.h"

class seqnode
{
public:
        char *seq;
	char *seqname;
        int *iseq;
        int length;
        int numleaf;
        double branchlen;
};

void writetree (TNode<class seqnode> *a, FILE *fp);
void inordertraversal(TNode<class seqnode> *a);
void node2align(TNode<class seqnode> *t, align_info *a);

int debug = 0;

int main(int argc, char *argv[])
{
	int i, j, k, l;
	align_info fullali(argv[1]);
	double **distanceMatrix;
	int fnal = fullali.getNal();
	int falilen = fullali.getAlilen();
	int **falign = fullali.getAlignment();
	int count, idcount;
	TNode<class seqnode> *root;

	// calculate the distance matrix for the sequences
	// d(i,j) = 1 - seqID(i, j);
	distanceMatrix = dmatrix(0,fullali.getNal(), 0, fullali.getNal() );
	for(i=1;i<=fnal;i++) {
	   for(j=1;j<=fnal;j++) {
		count = idcount = 0;
		for(k=1;k<=falilen;k++) {
		    if( (falign[i][k]==0)||(falign[j][k]==0) ) {
			continue;
		    }
		    count++;
		    if(falign[i][k]==falign[j][k])  idcount++;
		} 
	        if(count==0) distanceMatrix[i][j] = distanceMatrix[j][i] = 0;
		else
		   distanceMatrix[i][j] = distanceMatrix[j][i] = 1 - 1.0*idcount/count;
	   }
	}
	/*for(i=1;i<=fnal;i++) {
		for(j=1;j<=fnal;j++) {
		  fprintf(stderr, "%3.2f ", distanceMatrix[i][j]);
		}
		fprintf(stderr, "\n");
	}*/


	// an array of tree nodes of the type seqnode: leaves, number from 
	// 1 to fnal 
	TNode<class seqnode>* *leafnodes = new (TNode<class seqnode>*) [fnal+1];
	for(i=1;i<=fnal;i++) {
	 	leafnodes[i] = new TNode<class seqnode>;
		leafnodes[i]->setLeft(NULL);
		leafnodes[i]->setRight(NULL);
		leafnodes[i]->data.numleaf = 1;
		leafnodes[i]->data.branchlen = 0;
		leafnodes[i]->data.seqname = cvector(0, 1000);
		leafnodes[i]->data.seq = cvector(0,fullali.getAlilen() );
		leafnodes[i]->data.iseq = ivector(0, fullali.getAlilen() );
		strcpy(leafnodes[i]->data.seqname, fullali.aname[i-1]);
		strcpy(leafnodes[i]->data.seq, fullali.aseq[i-1]);
		for(j=1;j<=fullali.getAlilen();j++) {
			leafnodes[i]->data.iseq[j] = fullali.alignment[i][j];
		}
		//fprintf(stderr,"%s\n", leafnodes[i]->data.seqname);
	}

	// another array of tree nodes of the type seqnode: internal nodes, number 
	// from 1 to fnal-1
	TNode<class seqnode>* *internodes = new (TNode<class seqnode>*) [fnal];

	//build the tree by UPGMA
	int sindi1, sindi2;
	double smad;
	double *tempdist = dvector(0,fnal);
	double *branchlength = dvector(0,fnal);
	TNode<class seqnode> *internode;
	for(i=1;i<fnal;i++) {
	   
	   // find the smallest distance
	   smad = 1.5;
	   for(j=1;j<=fnal;j++) { 
	       for(k=j+1;k<=fnal;k++) {
	   	  if(distanceMatrix[j][k]<smad) {
			smad = distanceMatrix[j][k];
			sindi1 = j;
			sindi2 = k;
		   }
		}
	   }
	   //fprintf(stderr, "sindi1: %d; sindi2: %d\n", sindi1, sindi2);
	
	   // assign a new TNode as the parent of the two clusters with
	   // the smallest distance
	   //TNode<class seqnode> *internode = new TNode<class seqnode>;
	   internode = new TNode<class seqnode>;
	   internode->data.numleaf = leafnodes[sindi1]->data.numleaf + leafnodes[sindi2]->data.numleaf;
	   internode->data.branchlen = smad;
	   branchlength[fnal-i] = smad;
	   //fprintf(stderr, "i: %d; branchlen: %f\n", i, internode->data.branchlen);
	   
	   // adjust all distances since the two clusters are now merged to
	   // one. The one with smaller index will be replaced by the one with
	   // its new parent cluster, the one with larger index will be blanked
	   // 100 is an arbitrary large number used to blank an index
	   for(j=1;j<=fnal;j++) {
		tempdist[j] = (leafnodes[sindi1]->data.numleaf*distanceMatrix[j][sindi1]+
				leafnodes[sindi2]->data.numleaf*distanceMatrix[j][sindi2])/
				(leafnodes[sindi1]->data.numleaf+leafnodes[sindi2]->data.numleaf);
	   }
	  
	   for(j=1;j<=fnal;j++) {
		if(distanceMatrix[j][sindi1] ==100) continue;  
		distanceMatrix[j][sindi1] = distanceMatrix[sindi1][j] = tempdist[j];
	   }

	   for(j=1;j<=fnal;j++) {
	       	distanceMatrix[j][sindi2] = distanceMatrix[sindi2][j] = 100;
	   }
	
          /* for(l=1;l<=fnal;l++) {
                for(j=1;j<=fnal;j++) {
                  fprintf(stderr, "%6.1f ", distanceMatrix[l][j]);
                }
                fprintf(stderr, "\n");
           }
	   fprintf(stderr, "\n");
	  */


	   // replace the nodes and set up their relations
	   internode->setLeft(leafnodes[sindi1]);
	   internode->setRight(leafnodes[sindi2]);
	   leafnodes[sindi1]->setParent(internode);
	   leafnodes[sindi2]->setParent(internode);

	   leafnodes[sindi1] = internode;
	   leafnodes[sindi2] = 0;

	   //internal nodes sorted by decreasing order of branchlength
	   internodes[fnal-i] = internode;
	   
	}
	/*fprintf(stderr, "*******\n\n");
	for(i=1;i<fnal;i++) {
		fprintf(stderr, "i: %d; branchlength: %f\n", i, branchlength[i]);
		fprintf(stderr, "\n");
	}*/

	/*
	// testing traversal and print the tree file: routine writetree
	root = internode;
	//inordertraversal(root);
	char *treefile = cvector(0,1000);
	strcpy(treefile, argv[1]);
	strcat(treefile, ".tre");
	FILE *treefp;
	if((treefp = fopen(treefile, "w"))==NULL) {
		fprintf(stderr, "the tree file is not for write\n");
		exit(0);
	}
	fprintf(treefp, "(");
	writetree(root, treefp);
	fprintf(treefp, ",root:0);");
	fclose(treefp);
	*/

	// the functional residue list

	int fcount5 = 0, fcount8 = 0;
	int f5 = 0, f8 = 0;
	int residue5[100], residue8[100];
	ifstream ffunc(argv[2]);
	char str[100];
	while(!ffunc.eof()) {
	    str[0] = '\0';
	    ffunc >> str;
	    if(ffunc.eof()) break;
	    if(!str) continue;
	    if(!strcmp(str, "cutoff5:")){
		f5 = 1; continue;
	    }
	    if(!strcmp(str, "cutoff8:")) {
                f8 = 1; f5 = 0; continue;
            }
	    if(f5==1) {
		/*sprintf(str, "%d", residue5[fcount5] );*/
		residue5[fcount5] = atoi(str);
		fcount5++;
	    }
	    if(f8==1) {
		//sprintf(str, "%d", residue8[fcount8]);
		residue8[fcount8] = atoi(str);
	    	fcount8++;
	    }
	}

	if(debug>0) {	
	fprintf(stdout, "there are %d functional residues with 5 cutoff\n", fcount5);
	for(i=0;i<fcount5;i++) fprintf(stdout, "%d ", residue5[i]);
	fprintf(stdout, "\n");
	fprintf(stdout, "there are %d functional residues with 8 cutoff\n", fcount8);
	for(i=0;i<fcount8;i++) fprintf(stdout, "%d ", residue8[i]);
	fprintf(stdout, "\n");
	}
	

	// clustering step: given a patition number n, the sequence group 
	// will be partitioned to n clusters according to the tree
	int clusnum;
	double medianlen;
	int cluscount;

	int flag5 = 1;
	int flag10 = 1;
	int flag15 = 1;
	int flag20 = 1;


	int tracecount, convtracecount;
        int prevf5 = 0, prevf8 = 0;
        int prevtracecount = 0;
	double a5, a8;

   strcpy(str, argv[2]);
   strcat(str, ".et");
   FILE *fp;
   if((fp=fopen(str, "a+") )==NULL ) {
	fprintf(stderr, "cannot open the file for append\n") ;
	exit(0);
   }

   fprintf(fp, "%c %c %c ", argv[1][strlen(argv[1])-6], argv[1][strlen(argv[1])-5], argv[3][0] );

   fprintf(stdout, "fnal: %d\n", fnal);

   for(clusnum=2;clusnum<fnal;clusnum++){

	// judge if internodes[clusnum-1] has the same branchlength as 
	// internodes[clusnum]
	// if it is true, then there will be degeneration of clusters

	if(internodes[clusnum]->data.branchlen==0) break;
	if(internodes[clusnum-1]->data.branchlen == internodes[clusnum]->data.branchlen )
	{
		fprintf(stderr, "degeneration occurs;\n");
		fprintf(stderr, "there will be more clusters than designated.\n"); fflush(stdout);
	}
	while(internodes[clusnum-1]->data.branchlen == internodes[clusnum]->data.branchlen )
	{
		clusnum = clusnum+1;
		if(clusnum>=fnal) break;
	}
	if(debug>0) fprintf(stdout, "clusnum: %d\n", clusnum);

	if(clusnum>=fnal) break;
	//fprintf(stderr, "There will be %d clusters\n", clusnum);

	//medianlen serves as the threshold to judge the candidates
	medianlen = 0.5*(internodes[clusnum-1]->data.branchlen+internodes[clusnum]->data.branchlen);

	// an array of pointers to the align_info type of objects, number from 1 
	// to clusnum
	align_info **clus_align = new (align_info*) [clusnum+1];

	// for each of the first clusnum-1 internal nodes, judge if its two children are 
	// candidates of the root of the cluster
	cluscount = 0;
	for(i=1;i<clusnum;i++) {
		if(internodes[i]->getLeft()->data.branchlen<medianlen) {
			// compute the cluster alignment
			// the alignment is made up the sequences in 
			// the leaves under the root
			cluscount++;
			//fprintf(stderr, "count is now %d\n", cluscount);
			clus_align[cluscount] = new align_info;
			clus_align[cluscount]->setNal(0);
			clus_align[cluscount]->setAlilen(fullali.getAlilen());
			clus_align[cluscount]->aname = cmatrix(0,internodes[i]->getLeft()->data.numleaf,0,1000);
			clus_align[cluscount]->aseq = cmatrix(0,internodes[i]->getLeft()->data.numleaf,0,fullali.getAlilen());
			clus_align[cluscount]->alignment = imatrix(0,internodes[i]->getLeft()->data.numleaf,0,fullali.getAlilen());
			node2align(internodes[i]->getLeft(), clus_align[cluscount]);
			if(debug>0)fprintf(stderr, "cluster number %d\n", cluscount);
			//clus_align[cluscount]->printali2(80);
			//fprintf(stderr, "*****\n\n");
		}
		if(internodes[i]->getRight()->data.branchlen<medianlen) {
			// compute the cluster alignment
			// the alignment is made up the sequences in
			// the leaves under the root
			cluscount++;
			clus_align[cluscount] = new align_info;
			clus_align[cluscount]->setNal(0);
			clus_align[cluscount]->setAlilen(fullali.getAlilen());
			clus_align[cluscount]->aname = cmatrix(0,internodes[i]->getRight()->data.numleaf,0,1000);
			clus_align[cluscount]->aseq = cmatrix(0,internodes[i]->getRight()->data.numleaf,0,fullali.getAlilen());
			clus_align[cluscount]->alignment = imatrix(0,internodes[i]->getRight()->data.numleaf+1,0,fullali.getAlilen());
			node2align(internodes[i]->getRight(), clus_align[cluscount]);
			if(debug>0)fprintf(stderr, "cluster number %d\n", cluscount);
			//clus_align[cluscount]->printali2(80);
			//fprintf(stderr, "*****\n\n");
		}
	}
	if(cluscount<clusnum) continue;

	// set gap threshold that suppresses the calculation of tracing in a position
	// default: 0.2
	double gap_threshold = 0.2;
	/*if(argv[4] != NULL) {
		gap_threshold = atof(argv[4]);
		if( (gap_threshold>=1)||(gap_threshold<=0) ) gap_threshold = 0.2;
	}
	fprintf(stderr, "Gap threshold to suppress calculation: %3.2f\n", gap_threshold); 
	*/
	fullali.freq(gap_threshold);
	
	// tracing
	int *traceindex;
	traceindex = ivector(0, fullali.getAlilen() );
	char *classspecific;
	classspecific = cvector(0, clusnum);
	for(i=1;i<=fullali.getAlilen();i++) traceindex[i] = 1;

	for(i=1;i<=fullali.getAlilen();i++) {
		if(debug>0) fprintf(stdout, "i %d ", i); fflush(stdout);
		if(fullali.fq[0][i] == INDI ) {
			if(debug>0)fprintf(stdout, "-1 ");fflush(stdout);
			traceindex[i] = -1;
			continue;
		}
		if(debug>0)fprintf(stdout, "*");fflush(stdout);
		for(j=1;j<=clusnum;j++) {
			//fprintf(stdout, "j %d ", j);fflush(stdout);
			classspecific[j] = clus_align[j]->alignment[1][i];
			for(k=1;k<=clus_align[j]->getNal();k++) {
				if(classspecific[j] != clus_align[j]->alignment[k][i]) {
					traceindex[i] = -1;
					break;
				}
			}
			if(traceindex[i] == -1) continue;
			if(j>1) {
				if(classspecific[j]!=classspecific[j-1]) 
				traceindex[i] = 0;
			}
		}
		//fprintf(stdout, "i %d %d ", i, traceindex[i]); fflush(stdout);
	}
	if(debug>0)fprintf(stdout, "\n"); fflush(stdout);

	for(i=1;i<cluscount;i++) {
	  delete clus_align[i];
	}
	delete [] clus_align;

	tracecount = 0; convtracecount = 0;
	f5=0, f8=0;
	for(i=1;i<=fullali.getAlilen();i++) {
	    if(traceindex[i] == 1) {tracecount++; convtracecount++;}
	    if(traceindex[i] == 0) tracecount++;
	}

	for(i=1;i<=fullali.getAlilen();i++) {
	    if( (traceindex[i] == 1)||(traceindex[i] == 0) ) {
		for(j=0;j<fcount5;j++) {
		  if(i==residue5[j]) {
			f5++; break;
		  }
		}
		for(j=0;j<fcount8;j++) {
		  if(i==residue8[j]) {
			f8++; break;
		  }
		}
	     }
	}

	fflush(stdout);
	if(debug>0)fprintf(stdout, "group %d: %d/%d  %d/%d  %d\n", clusnum, f5, tracecount, f8, tracecount, convtracecount);
	fflush(stdout);

	if( (tracecount==5)&&flag5)  {
	    a5 = f5;
	    a8 = f8;
	    fprintf(stdout, "5 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	    fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
	    flag5 = 0;
	}

	if( (tracecount>5) && (flag5==1) ) {
	    a5 = prevf5 + 1.0*(f5-prevf5)*(5-prevtracecount)/(tracecount-prevtracecount);
	    a8 = prevf8 + 1.0*(f8-prevf8)*(5-prevtracecount)/(tracecount-prevtracecount);
	    fprintf(stdout, "5 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	    fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
	    flag5 = 0;
	}

	if( (tracecount==10)&&flag10)  {
	  a5 = f5;
	  a8 = f8;
	  flag10 = 0;
	  fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
	  fprintf(stdout, "10 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	}
	
	if( (tracecount>10) && (flag10==1))  {
	    a5 = prevf5 + 1.0*(f5-prevf5)*(10-prevtracecount)/(tracecount-prevtracecount);
            a8 = prevf8 + 1.0*(f8-prevf8)*(10-prevtracecount)/(tracecount-prevtracecount);
            fprintf(stdout, "10 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	    fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
            flag10 = 0;
	}

	if( (tracecount==15)&&flag15 ) {
          a5 = f5;
          a8 = f8;
          flag15 = 0;
          fprintf(stdout, "15 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	  fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
        }

        if( (tracecount>15) && (flag15==1) ) {
            a5 = prevf5 + 1.0*(f5-prevf5)*(15-prevtracecount)/(tracecount-prevtracecount);
            a8 = prevf8 + 1.0*(f8-prevf8)*(15-prevtracecount)/(tracecount-prevtracecount);
            fprintf(stdout, "15 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	    fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
            flag15 = 0;
        }

	if( (tracecount==20)&&flag20 ) {
          a5 = f5;
          a8 = f8;
          flag20 = 0;
          fprintf(stdout, "20 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	  fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
	  break;
        }

        if( (tracecount>20) && (flag20==1) ) {
            a5 = prevf5 + 1.0*(f5-prevf5)*(20-prevtracecount)/(tracecount-prevtracecount);
            a8 = prevf8 + 1.0*(f8-prevf8)*(20-prevtracecount)/(tracecount-prevtracecount);
            fprintf(stdout, "20 residues: a5: %3.2f a8: %3.2f\n", a5, a8);
	    fprintf(fp, "%4.2f %4.2f %4.2f ", a5, a8, a8-a5);
            flag20 = 0;
        }

	if(tracecount!=prevtracecount) {
	   prevtracecount = tracecount;
	   prevf5 = f5;
	   prevf8 = f8;
	}

	if(debug>0) {
	for(i=1;i<=fullali.getAlilen();i++) {
            if((traceindex[i] == 1)||(traceindex[i]==0)) fprintf(stdout, "%d ", i);
        }
	fprintf(stdout, "\n");
	}

	fflush(fp);

    }

    if(flag5==1) fprintf(fp, "0 0 0 ");

    if(flag10==1) fprintf(fp, "0 0 0 ");
	
    if(flag15==1)  fprintf(fp, "0 0 0 ");
  
    if(flag20==1) fprintf(fp, "0 0 0 ");

    fprintf(fp, "%4d\n", convtracecount);

    fclose(fp);

	return 0;
	
}
	  
void writetree (TNode<class seqnode> *a, FILE *fp)
{
	if(a->getLeft()==NULL) return;
	fprintf(fp, "(");
	//fprintf(stderr, "%5.3f", a->data.branchlen);
	if(a->getLeft()->getLeft()!=NULL) {
		writetree( a->getLeft(),fp );
	} else {
	    	fprintf(fp, "%s", a->getLeft()->data.seqname);
	}
	fprintf(fp, ":%5.3f,",a->data.branchlen - a->getLeft()->data.branchlen );
	if(a->getRight()->getLeft()!=NULL) {
		writetree(a->getRight(),fp);
	} else {
		fprintf(fp, "%s", a->getRight()->data.seqname);
	}
	fprintf(fp, ":%5.3f)", a->data.branchlen - a->getRight()->data.branchlen );

}

void inordertraversal(TNode<class seqnode> *a)
{
	if(a->getLeft()==NULL) return;
	if(a->getRight()==NULL) return;
	fprintf(stderr, "%5.3f\n", a->data.branchlen);
	inordertraversal(a->getLeft());
	inordertraversal(a->getRight());
}
	 
void node2align(TNode<class seqnode> *t, align_info *a) 
{
	int i,j,k;
		

	if(t->getLeft()==NULL) {
		a->setNal( a->getNal()+1 );
		strcpy(a->aname[a->getNal()-1], t->data.seqname);
		strcpy(a->aseq[a->getNal()-1], t->data.seq);
		for(i=1;i<=a->getAlilen();i++) {
			a->alignment[a->getNal()][i] = t->data.iseq[i];
		}
		return;
	}

	node2align(t->getRight(), a);
	node2align(t->getLeft(),  a);
}
		
