/*********************************************************************/
/*                                                                   */
/*                             idc.c                                 */
/*            Program for  the calculation of the of                 */
/*              independent contrasts from a phylogeny.              */
/*           written by Liam J. Revell (ljrevell@wustl.edu)          */
/*               much code adapted from Luke J. Harmon               */
/*                                                                   */
/*********************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#define MAXCHAR 10000
#define MAXNAME 50

// structures
struct node {
	struct node *left_ptr; // Left daughter
	struct node *right_ptr; // Right daughter
	struct node *parent_ptr; // Parent node 
	double branch_length;	 // Branch length leading to parent branch
	double corrected_bl;  // corrected branch length following (Felsenstein 1985) 
	double depth;		 // Depth in tree from root
	int tot_species;		 // Total number of species as of this point in tree 
	char name[40];		// Name of tip or interior node (not used)
	int tip_number;
	char node_name[101];
	int side;
	double *characters; // vector of character values 
	double *contrasts;
};

// end structures

// function prototypes
double **malloc_double_array(int rows, int cols);
int count_species(struct node *top);
int read_file_line(FILE *input_file, char *file_ch);
int read_tree(char *filechar, int n, struct node *newnode, struct node *currnode, int, double **, double **, FILE *cont_fp, struct node *list);
void free_tree(struct node *top);
void init_node(struct node *newnode, int numtraits);
void measure_tree(struct node *top);
void calc_mcp(int, int, double **, double **);
void print_tree(struct node *top, FILE *fp);
void cov_to_corr(int, double **, double **);
double *mem_1d(int length);
void sleep_seconds ( long seconds );
int max_of_two(int integer1, int integer2);

// functions

// finds larger of two integers

int max_of_two(int integer1, int integer2){

	int biggest;

	biggest=integer1;
		
	if(integer2>biggest) biggest=integer2;
	
	return biggest;
	
}	

// waits

void sleep_seconds (long seconds){
  clock_t limit, now = clock();
  limit = now + seconds * CLOCKS_PER_SEC;
  while ( limit > now )
    now = clock();
}
// memory allocation for double vector

double *mem_1d(int length){
  double *result;

  result=malloc(length*sizeof(double));
  if(!result){ printf("Memory allocation failure.\n"); fflush(stdout); exit(1);}
  return result;
}

// this function converts a covariance matrix to a correlation matrix

void cov_to_corr(int vars, double **covmat, double **corrmat){
  int i,j;
  for(i=0;i<vars;i++){
    for(j=0;j<vars;j++){
      corrmat[i][j]=covmat[i][j]/sqrt(covmat[i][i]*covmat[j][j]);
    }
  }
}

// this function measures and assigns tree depth to nodes (LJH)

void measure_tree(struct node *top){
	struct node *parent;
	if (top==NULL)
		return;
	if(top->parent_ptr==NULL)
		top->depth=0;
	else
	{
		parent=top->parent_ptr;
		top->depth = top->branch_length + parent->depth;
	}
	measure_tree(top->left_ptr);
	measure_tree(top->right_ptr);
}

// this function prints the tree from a given node (LJR)

void print_tree(struct node *top, FILE *fp){
	if(top->left_ptr==NULL){
		fprintf(fp,"%d",top->tip_number);
		return;
	}
	if(top->left_ptr!=NULL){
		fprintf(fp,"(");
		print_tree(top->left_ptr,fp);
	}
	if(top->right_ptr!=NULL){
		fprintf(fp,",");
		print_tree(top->right_ptr,fp);
	}
	if(top->right_ptr!=NULL){
		fprintf(fp,")");
	}
}


// this function calculates a mean-cross-product (covariance) matrix from data (LJR)

void calc_mcp(int N, int chars, double **data, double **matrix){
	int i,j,k;
  	double MCP; // M[2];
  	for(i=0;i<chars;i++){
    	for(j=0;j<chars;j++){
      		MCP=0.0; // M[0]=M[1]=0.0;
      		for(k=0;k<N;k++){
	  			MCP+=data[k][i]*data[k][j]/(double)N;
	  			// M[0]+=data[k][i]/(double)N;
	  			// M[1]+=data[k][j]/(double)N;
      		}
			matrix[i][j]=MCP;
      		// matrix[i][j]=(double)N*(MCP-M[0]*M[1])/(double)(N-1);
    	}
  	}
}

// this function calculate average tree depth to a given node (LJH/LJR)

double tree_depth(struct node *top){

	double tot_depth_left, tot_depth_right;
	if (top->left_ptr==NULL && top->right_ptr==NULL)
		return(top->depth);

	tot_depth_left=tree_depth(top->left_ptr);
	tot_depth_right=tree_depth(top->right_ptr);
	return((tot_depth_left+tot_depth_right)/2);
}

// this function allocates memory to an array of doubles (LJH/LJR)

double **malloc_double_array(int rows, int cols){
	
	int i, j;
	double **array;

	array=(double **)malloc(rows*sizeof(double *));
	if(array==NULL){
		printf("Out of memory.\n"); fflush(stdout);
		return(NULL);
	}
	for (i=0; i<rows; i++){
      	array[i]= (double *)malloc(cols*sizeof(double));
		if(array[i]==NULL){
			printf("Out of memory.\n"); fflush(stdout);
			return(NULL);
		}
	}
	
	for(i=0; i<rows; i++)
		for(j=0; j<cols; j++)
			array[i][j]=0;

	return(array);
}

// this function calculates the number of species descended from node *top (LJH)

int count_species(struct node *top){

	if (top==NULL)
		return(0);
	if(top->left_ptr==NULL && top->right_ptr==NULL)
		return(1);
	else return(count_species(top->left_ptr)+count_species(top->right_ptr));
}

// this function reads a line of text from the input file (LJH)

int read_file_line(FILE *input_file, char *file_ch){
	
	int i=0;
	while(1){
		file_ch[i]=fgetc(input_file);
		if (file_ch[i]=='\n' || file_ch[i]==EOF) break;
		i++;
	}
	return(i-1);
}

// this function reads in the tree, calculates the node character states, and calculates all contrasts (LJH/LJR)

int read_tree(char *filechar, int n, struct node *newnode, struct node *currnode, int traits, double **char_array, double **contrasts_array, FILE *cont_fp, struct node *list){
	int i=0, j, k, scinote,sign,exponent, numcontrasts=0;
	char temp[20],temp_pre_e[20],temp_post_e[2];
	float floater;	

	while(1){
	
		if (filechar[i]=='('){
 	   		// Create new daughter node 
			newnode=malloc(sizeof(struct node));
			if (newnode==NULL) {printf ("memory error\n"); fflush(stdout);}
			init_node(newnode,traits);
			newnode->parent_ptr=currnode;
			newnode->side=0;
			currnode->left_ptr=newnode;
			currnode=newnode;
			i++;
        }
	   	else if (filechar[i]==','){
			// Create new right node 
        	newnode=malloc(sizeof(struct node));
			if (newnode==NULL) {printf ("memory error\n"); fflush(stdout);}
			init_node(newnode,traits);
			newnode->parent_ptr=currnode->parent_ptr;
			newnode->side=1;
			currnode->parent_ptr->right_ptr=newnode;
			currnode=newnode;
			i++;
    	}
      	else if (filechar[i]==':'){
			// Read branch length 
			j=0;
			i++;
			scinote=0;
			while((filechar[i]>='0' && filechar[i]<='9')||(filechar[i]=='.' || filechar[i]=='e' || filechar[i]=='-')){
		  		if(filechar[i]=='e') scinote=1;
		        temp[j]=filechar[i];
				i++;
				j++;
			}
			temp[j]='\0';
			if(scinote==0){
		  		sscanf(temp, "%f", &floater);
			}
			if(scinote==1){
		  		j=0;
		  		while(j<5){
		    		temp_pre_e[j]=temp[j];
		    		j++;
		  		}
		  		temp_pre_e[6]='\0';
		  		sscanf(temp_pre_e,"%f",&floater);
		  		if(temp[6]=='-') sign=-1;
		  		if(temp[6]=='+') sign=1;
		  		k=0;
		  		for(j=7;j<10;j++){
		    		if(temp[j]=='0');
		    		if(temp[j]>='1' && temp[j]<='9'){ 
		      			temp_post_e[k]=temp[j];
		      			k++;
		    		}
		  		}
		  		k++;
		  		temp_post_e[k]='\0';
		  		sscanf(temp_post_e,"%d",&exponent);
		  		exponent*=sign;
		  		floater=floater*pow(10,exponent);
			} 
			// assign estimated branch length
     		currnode->branch_length = floater;
			// also calculate and assign corrected branch length [following Felsenstein (1985)]
			currnode->corrected_bl=currnode->branch_length;
			if(currnode->left_ptr!=NULL){
				// apply correction
				currnode->corrected_bl+=currnode->left_ptr->corrected_bl*currnode->right_ptr->corrected_bl/(currnode->left_ptr->corrected_bl+currnode->right_ptr->corrected_bl);
			}
			// already at next character 		
	   	}
	   	else if (filechar[i]==')'){
			// Back down into tree 
			currnode=currnode->parent_ptr;
			
			// calculate ancestral state for node

			for(j=0;j<traits;j++){
				// calculate ancestral state as the weighted mean
				currnode->characters[j]=((1.0/currnode->left_ptr->corrected_bl)*currnode->left_ptr->characters[j]+(1.0/currnode->right_ptr->corrected_bl)*currnode->right_ptr->characters[j])/(1.0/currnode->left_ptr->corrected_bl+1.0/currnode->right_ptr->corrected_bl);
				// calculate contrasts
				currnode->contrasts[j]=(currnode->left_ptr->characters[j]-currnode->right_ptr->characters[j])/sqrt(currnode->left_ptr->corrected_bl+currnode->right_ptr->corrected_bl);
				contrasts_array[numcontrasts][j]=currnode->contrasts[j];
			}
			list[numcontrasts]=*currnode;
			numcontrasts++;
			i++;
	   	}
	   	else if (filechar[i]==';'){
			break;
			// End tree reading
	   	}
      	else{
	   		// Read species name 
			j=0;
			while(filechar[i]!=':'){
				temp[j]=filechar[i];						
				// currnode->name[j]=filechar[i];
				i++;
				j++;
			}
			temp[j]='\0';
			sscanf(temp,"%d",&currnode->tip_number);
			// currnode->name[j]='\0';
			for(j=0;j<traits;j++)
				currnode->characters[j]=char_array[currnode->tip_number][j];		
	   	}
	}
}

// this function frees the memory allocated to the tree (LJH)

void free_tree(struct node *top){

	if (top==NULL)
		return;
	free_tree(top->left_ptr);
	free_tree(top->right_ptr);
	free(top);

}

// this function initiliazes a new node with 0s and NULLs (LJH/LJR)

void init_node(struct node *newnode, int numtraits){

	int i;
	
	newnode->left_ptr=NULL;
	newnode->right_ptr=NULL;
	newnode->parent_ptr=NULL;
	newnode->name[0]='\0';
	newnode->branch_length=0;
	newnode->depth=0;
	newnode->tot_species=0;
	newnode->tip_number=0;
	newnode->characters=mem_1d(numtraits);
	newnode->contrasts=mem_1d(numtraits);
	for(i=0;i<numtraits;i++){
		newnode->characters[i]=0.0;
		newnode->contrasts[i]=0.0;
	}
}

// this is the main program (LJR)

int main(){
		
	char filename[50], c, d;
	int *tip_names, numtraits, numspecies;
	double **data_array, **contrasts_array;
	double **cov_matrix, **corr_matrix;
	FILE *data_file, *tree_file, *out_file, *corr_file, *cov_file;
	char temp[20];
	struct node *curr_node;
	struct node *root = NULL;
	struct node *new_node;
	struct node *node_list;
	int i, j, nchars, nspecies=0;	char file_char[MAXCHAR], line[100];
	float floater;
	double avg_tree_depth;
	int done=0;
	int single_many_data=0;
	int single_many_tree=0;
	int print_vcv_corr=0;
	int unit_variance=0;
	int minimal=0;
	int bad_option=0;
	int numdata=1, numtrees=1, line_num;
	int loops=0, numloops;
	char *credits[]=
			{
  		"  --------------------------------------------------------------- ",
  		" |                          ++ IDC ++                            |",
  		" |                          ~~~~~~~~~                            |",
  		" |                                                               |",  
  		" |                                                               |",  
  		" |       A program for the calculation of phylogenetically       |",
  		" |                    independent contrasts                      |",  		
  		" |                                                               |",  
  		" |                                                               |",  
  		" |               Written by Liam J. Revell (2006)                |",
  		" |               Code adapted from Luke J. Harmon                |",  
  		" |                                                               |",  
  		" |                                                               |",
  		" |                                                               |",
  		" |                                                               |",
  		"  --------------------------------------------------------------- "
		};	
	
	char *option1[]=
		{
  		" |   (S) Single dataset or many (Single)                         |",
  		" |   (S) Single dataset or many (many)                           |",
		};
	char *option2[]=
		{
  		" |   (T) Single tree or many (Single)                            |",
  		" |   (T) Single tree or many (many)                              |",
		};
	char *option3[]=
		{
  		" |   (P) Print VCV & correlation matrix to file (Print)          |",
  		" |   (P) Print VCV & correlation matrix to file (do not print)   |",
		};
	char *option4[]=
		{
  		" |   (U) Standardize contrasts to unit var. (Don't standardize)  |",
  		" |   (U) Standardize contrasts to unit var. (standardize)        |",
  	};
  char *option5[]=
		{
  		" |   (M) Minimalist output (Full output)                         |",
  		" |   (M) Minimalist output (minimal output)                      |",
  	};
	char *menu[] =
		{
  		"  --------------------------------------------------------------- ",
  		" |                         ++ OPTIONS ++                         |",
  		" |                         ~~~~~~~~~~~~~                         |",
  		" |   (S) Single dataset or many (Single)                         |",
  		" |   (T) Single tree or many (Single)                            |",
  		" |   (P) Print VCV & correlation matrix to file (Print)          |",
  		" |   (U) Standardize contrasts to unit var. (Don't standardize)  |",
  		" |   (M) Minimalist output (Full output)                         |",
  		" |                                                               |",  		
  		" |   (D) Done entering options (not-done)                        |",
 			" |   (Q/q) Quit                                                  |",
  		" |                                                               |",
  		" |   Selected option in ().                                      |",
  		" |   Hit (key) Enter to toggle options.                          |",
  		" |                                                               |",
  		"  --------------------------------------------------------------- "
		};
		
	// print credits
	
	
	for(line_num=0;line_num<16;line_num++){
      printf ("%s\n", credits[line_num]);
  }
	printf("  Press a key to continue      > "); fflush(stdout);
	c=fgetc(stdin);
	
	// print menu
	
	while(done==0){
	
//		system("clear");

  	for (line_num=0;line_num<16;line_num++){
      if(line_num!=3&&line_num!=4&&line_num!=5&&line_num!=6&&line_num!=7) printf ("%s\n", menu[line_num]);
      else if(line_num==3) printf("%s\n",option1[single_many_data]);
      else if(line_num==4) printf("%s\n",option2[single_many_tree]);
      else if(line_num==5) printf("%s\n",option3[print_vcv_corr]);
      else if(line_num==6) printf("%s\n",option4[unit_variance]);
      else if(line_num==7) printf("%s\n",option5[minimal]);
    }
    if(bad_option==1){
    	printf("  Bad option!\n");
    	bad_option=0;
    }
 fflush(stdout);
   	printf("  Enter option                 > ");
 fflush(stdout);
		c=fgetc(stdin); d=fgetc(stdin);
		if(c=='S'||c=='s'){ 
			if(single_many_data==0){ 
				single_many_data=1;
				printf("  Enter number of data sets    > "); fflush(stdout);
				scanf("%d",&numdata);
				d=fgetc(stdin);
			}		
			else { 
				single_many_data=0;
				numdata=1;
			}
		}
		else if(c=='T'||c=='t'){
			if(single_many_tree==0){ 
				single_many_tree=1;
				printf("  Enter number of trees        > "); fflush(stdout);
				scanf("%d",&numtrees);
				d=fgetc(stdin);
				}
			else { 
				single_many_tree=0;
				numtrees=1;
			}
		}
		else if(c=='P'||c=='p'){
		 	if(print_vcv_corr==0) print_vcv_corr=1;
		 	else print_vcv_corr=0;
		 }
		else if(c=='U'||c=='u'){
		 	if(unit_variance==0) unit_variance=1;
		 	else unit_variance=0;
		 }
		else if(c=='M'||c=='m'){
		 	if(minimal==0) minimal=1;
		 	else minimal=0;
		 }
		else if(c=='D'||c=='d'){
			if((numdata!=numtrees)&&(single_many_data==1)&&(single_many_tree==1)){
				printf("  Error! For multiple data sets and multiple trees,\n  number of trees must = number of data sets.\n");
				printf("  Resetting options.\n"); fflush(stdout);
				single_many_data=single_many_tree=0;
				sleep_seconds(4);
			}
			else done=1;
		}
		else if(c=='Q'||c=='q'){
			printf("  Quitting!\n"); fflush(stdout);
			sleep_seconds(1);
			system("clear");
			exit(1);
		}
		else{ 
			bad_option=1;
		}
	}
	
	// open I/O files

	printf("  Filename with character data > "); fflush(stdout);
	fgets(filename,sizeof(filename),stdin);
	filename[strlen(filename)-1]='\0';
	data_file=fopen(filename,"r");
	if(!data_file){ 
		printf("Couldn't open %s\n",filename); fflush(stdout);
		exit(1);
	}
	
	printf("  Filename with Newick tree    > "); fflush(stdout);
	fgets(filename,sizeof(filename),stdin);
	filename[strlen(filename)-1]='\0';
	tree_file=fopen(filename,"r");
	if(!tree_file){
		printf("Couldn't open %s\n",filename); fflush(stdout);
		exit(1);
	}

	printf("  Output filename              > "); fflush(stdout);
	fgets(filename,sizeof(filename),stdin);
	filename[strlen(filename)-1]='\0';
	out_file=fopen(filename,"w");
	if(!out_file){
		printf("couldn't create %s\n",filename); fflush(stdout);
		exit(1);
	}
	
	numloops=max_of_two(numdata,numtrees);
	
	if(((numloops>1)||(minimal==1))&&(print_vcv_corr==0)){
		printf("  Filename for vcv matrices    > "); fflush(stdout);
		fgets(filename,sizeof(filename),stdin);
		filename[strlen(filename)-1]='\0';
		cov_file=fopen(filename,"w");
		if(!cov_file){
			printf("couldn't create %s\n",filename); fflush(stdout);
			exit(1);
		}
		printf("  Filename for corr. matrices  > "); fflush(stdout);
		fgets(filename,sizeof(filename),stdin);
		filename[strlen(filename)-1]='\0';
		corr_file=fopen(filename,"w");
		if(!corr_file){
			printf("couldn't create %s\n", filename); fflush(stdout);
			exit(1);
		}
	}
	else {
		cov_file=out_file;
		corr_file=out_file;
	}
	
	for(loops=1;loops<=numloops;loops++){
	
		// read character data

		if(loops<=numdata){ 
			fscanf(data_file,"%d %d",&numtraits,&numspecies);

			// allocate for tip and character arrays

			tip_names=malloc((numspecies+1)*sizeof(int));
			if(!tip_names){
				printf("malloc failure\n"); fflush(stdout);
				exit(1);
			}
	
			data_array=malloc_double_array(numspecies+1,numtraits);

			// get data
			for(i=0;i<numspecies;i++){
				j=0;
				while(((temp[j]=fgetc(data_file))!='\t')&&temp[j]!=' ')
					j++;
				temp[j]='\0';
				sscanf(temp,"%d",&tip_names[i]);
				for(j=0;j<numtraits;j++){
					fscanf(data_file,"%lf",&data_array[i+1][j]);
				}
			}

			// allocate memory for contrasts array

			contrasts_array=malloc_double_array(numspecies-1,numtraits);
	
			// allocate memory for node list
	
			node_list=malloc((numspecies-1)*sizeof(struct node));
			if(!node_list){
				printf("malloc failure.\n"); fflush(stdout);
				exit(1);
			}
		}

		// get tree

		if(loops<=numtrees){
		
			// free old tree
			
			free_tree(root);
		
			root=(struct node *)malloc(sizeof(struct node));

			init_node(root,numtraits);

 			nchars=read_file_line(tree_file, file_char);

		}
		
		read_tree(file_char, nchars, new_node, root, numtraits, data_array, contrasts_array, out_file, node_list);

		nspecies=count_species(root);

		measure_tree(root);
		avg_tree_depth=tree_depth(root);

		// check number of species

		if(numspecies!=nspecies){
			printf("Oops-> the number of species specified in the character file differs from the number in your tree.\n"); fflush(stdout);
			exit(1);
		}	

		// multiply all contrasts by total tree length
	
		for(i=0;i<(numspecies-1);i++){
			for(j=0;j<numtraits;j++){
				contrasts_array[i][j]*=sqrt(avg_tree_depth);
			}
		}

		// print to ouput file
		if(minimal==0){
			if(numloops<=1)
				fprintf(out_file,"Phylogenetically independent contrasts:\n");
			else
				fprintf(out_file,"Data set %d\n",loops);
			fprintf(out_file,"number\tlength  \tcorrected");
			for(j=0;j<numtraits;j++) fprintf(out_file,"\tz_%d      ",j+1);
			fprintf(out_file,"\tcontrast\n");
		}
		else
			fprintf(out_file,"%d %d %d\n",loops,numspecies-1,numtraits);
			
		// calculate a covariance matrix for the contrasts
	
		cov_matrix=malloc_double_array(numtraits,numtraits);

		calc_mcp(numspecies-1,numtraits,contrasts_array,cov_matrix);
		
		for(i=0;i<numspecies-1;i++){
			if(minimal==0){
				fprintf(out_file,"%d    \t%lf",i+1,node_list[i].left_ptr->branch_length+node_list[i].right_ptr->branch_length);
				fprintf(out_file,"\t%lf\t",node_list[i].left_ptr->corrected_bl+node_list[i].right_ptr->corrected_bl);
			}
			for(j=0;j<numtraits;j++){
				if(unit_variance==0)
					fprintf(out_file,"%lf\t",contrasts_array[i][j]);
				else if(unit_variance==1)
					fprintf(out_file,"%lf\t",contrasts_array[i][j]/sqrt(cov_matrix[j][j]));
			}
			if(minimal==0){
				print_tree(node_list[i].left_ptr,out_file);
				fprintf(out_file," vs. ");
				print_tree(node_list[i].right_ptr,out_file);
			}
			fprintf(out_file,"\n");			
		}

		if(print_vcv_corr==0){

			// calculate a covariance matrix for the contrasts
	
			cov_matrix=malloc_double_array(numtraits,numtraits);

			calc_mcp(numspecies-1,numtraits,contrasts_array,cov_matrix);
			if((numloops<=1)&&(minimal==0))
				fprintf(out_file,"\nVCV matrix of independent contrasts:\n");
			else
				fprintf(cov_file,"%d\n",loops);
			for(i=0;i<numtraits;i++){
				for(j=0;j<numtraits;j++){
					fprintf(cov_file,"%.10lf\t",cov_matrix[i][j]);
				}
				fprintf(cov_file,"\n");
			}
			if((numloops>1)||(minimal==1))
				fprintf(cov_file,"\n");

			// calculate a correlation matrix for the contrasts

			corr_matrix=malloc_double_array(numtraits,numtraits);

			cov_to_corr(numtraits,cov_matrix,corr_matrix);

			if((numloops<=1)&&(minimal==0))
				fprintf(out_file,"\nCorrelation matrix of independent contrasts:\n");
			else
				fprintf(corr_file,"%d\n",loops);
			for(i=0;i<numtraits;i++){
				for(j=0;j<numtraits;j++){
					fprintf(corr_file,"%.10lf\t",corr_matrix[i][j]);
				}
				fprintf(corr_file,"\n");
			}
			fprintf(corr_file,"\n");
		}
		fprintf(out_file,"\n");
	}

	// free tree from memory

	free_tree(root);

	// free file pointers

	fclose(data_file);
	fclose(tree_file);
	fclose(out_file);

	printf("  Analysis complete.  Data printed to file.\n");
	printf("  Exiting. ");
	fflush(stdout);
	return 0;
}
