/*  Nucleotide Sequence Generator - seq-gen, version 1.2.5 */
/*  (c) Copyright 1996-2001, Andrew Rambaut & Nick Grassly */
/*  Department of Zoology, University of Oxford            */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <ctype.h>

#include "treefile.h"

char treeErrorMsg[256];
int treeError;

TNode *avail=NULL;
long usedAvail=0;
long usedMalloc=0;

/* prototypes */

TNode *NewNode(TTree *tree);
void InitTree(TTree *tree);
void DisposeNode(TNode *aNode);
void DisposeTreeNodes(TTree *tree);
void FreeNodes(void);

char ReadToNextChar(FILE *fv);
void ReadUntil(FILE *fv, char stopChar, char *what);
TNode *ReadTip(FILE *fv, char ch, TTree *tree, int numNames, char **names);
TNode *ReadNode(FILE *fv, TTree *tree, int numNames, char **names, int detectPolytomies);
TNode *ReadBranch(FILE *fv, TTree *tree, int numNames, char **names);
void WriteNode(FILE *fv, TTree *tree, TNode *node);


/* functions */

/*----------*/
TNode *NewNode(TTree *tree)
{
	TNode *node;

	if ( avail!=NULL ) {
		node=avail;
		avail=avail->next;
		
		usedAvail++;
	} else {
		if ( (node=malloc(sizeof(TNode)))==NULL ) {
			strcpy(treeErrorMsg, "Out of memory");
			return NULL;
		}
		usedMalloc++;
	}
	node->branch0=node->branch1=node->branch2=NULL;
	node->length0=node->length1=node->length2=0.0;
	node->param=0.0;
	node->tipNo=-1;

	node->sequence = NULL;

	node->effectiveLength0 = NULL;
	node->hiddenStates = NULL;
	
	node->next=tree->nodeList;	
	tree->nodeList=node;
	
	tree->numNodes++;

	return node;
} /* NewNode */


/*----------*/
void InitTree(TTree *tree)
{	
	tree->root=NULL;
	tree->nodeList=NULL;
	tree->numNodes=0;
	tree->numTips=0;
	tree->totalLength=0.0;
	tree->rooted=0;
	tree->lengths=-1;
} /* InitTree */


/*----------*/
TTree *NewTree()
{
	TTree *tree;
	
	if ( (tree=(TTree *)malloc(sizeof(TTree)))==NULL ) {
		strcpy(treeErrorMsg, "Out of memory creating tree.");
		return NULL;
	}
	InitTree(tree);
	memset(tree->names, 0, sizeof(char*)*MAX_TIPS);
	
	return tree;
} /* NewTree */


void DisposeNode(TNode *aNode)
{
	aNode->next=avail;
	avail=aNode;
} /* DisposeNode */
			

/*----------*/
void DisposeTreeNodes(TTree *tree)
{
	TNode *P, *O;
	
	if ( tree ) {
		P=tree->nodeList;
		while (P!=NULL) {
			O=P;
			P=P->next;
			DisposeNode(O);
		}
		tree->nodeList=NULL;
	}
} /* DisposeTreeNodes */


/*----------*/
void DisposeTree(TTree *tree)
{
	if (tree) {
		DisposeTreeNodes(tree);
		InitTree(tree);
	}
} /* DisposeTree */


/*----------*/
void FreeNodes(void)
{
	TNode *P, *O;
	
	P=avail;
	while (P!=NULL) {
		O=P;
		P=P->next;
		free(O);
	}
} /* FreeNodes */


/*----------*/
void FreeTree(TTree *tree)
{
	if (tree) {
		DisposeTreeNodes(tree);
		free(tree);
	}
	FreeNodes();
} /* FreeTree */


/*----------*/
void WriteAvailInfo()
{
	TNode *P;
	int count;
	
	count=0;
	P=avail;
	while (P!=NULL) {
		P=P->next;
		count++;	
	}

	fprintf(stderr, "Avail: %d nodes - availed: %ld, malloced: %ld\n", count, usedAvail, usedMalloc);
} /* WriteAvailInfo */


int CountTrees(FILE *fv)
{
	int n;
		
	if (fv==NULL)
		return 0;
		
	n=0;
	while (!feof(fv)) {
		if (fgetc(fv)==';')
			n++;
	}
	rewind(fv);
	
	return n;
}

char ReadToNextChar(FILE *fv)
{
	char ch;
	
	ch=fgetc(fv);
	while (!feof(fv) && isspace(ch)) 
		ch=fgetc(fv);

	return ch;
}


void ReadUntil(FILE *fv, char stopChar, char *what)
{
	char ch;
	
	ch=fgetc(fv);
	while (!feof(fv) && ch!=stopChar && 
			ch!='(' && ch!=',' && ch!=':' && ch!=')' && ch!=';') 
		ch=fgetc(fv);

	if (feof(fv) || ch!=stopChar) {
		sprintf(treeErrorMsg, "%s missing", what);
		treeError=1;
	}
}


TNode *ReadTip(FILE *fv, char ch, TTree *tree, int numNames, char **names)
{
	int i;
	char *P;
	char name[256];
	
	TNode *node;
	
	node=NewNode(tree);
	i=0;
	
	P=name;
	while (!feof(fv) && ch!=':' && ch!=',' && ch!=')' && i<MAX_NAME_LEN) {
		*P=ch;
		i++; P++;
		ch=fgetc(fv);
	}
	*P='\0';

	if (numNames == 0) {
		node->tipNo=tree->numTips;
		if (tree->names[node->tipNo]==NULL) {
			if ( (tree->names[node->tipNo]=(char *)malloc(MAX_NAME_LEN+1))==NULL ) {
				strcpy(treeErrorMsg, "Out of memory creating name.");
				return NULL;
			}
		}
		strcpy(tree->names[node->tipNo], name);
	} else {
	/* we already have some names so just look it up...*/
		i = 0;
		while (i < numNames && strcmp(name, names[i]) != 0)
			i++;

		if (i == numNames) {
			sprintf(treeErrorMsg, "Taxon names in trees for different partitions do not match.");
			return NULL;
		}
		
		node->tipNo=i;
	}
	
	tree->tips[node->tipNo]=node;
	tree->numTips++;

	while (!feof(fv) && ch!=':' && ch!=',' && ch!=')') 
		ch=fgetc(fv);

	if (feof(fv)) {
		sprintf(treeErrorMsg, "Unexpected end of file");
		return NULL;
	}
	ungetc(ch, fv);
	
	return node;
}


TNode *ReadNode(FILE *fv, TTree *tree, int numNames, char **names, int detectPolytomies)
{
	TNode *node, *node2;
	char ch;
	
	if ((node=NewNode(tree))==NULL)
		return NULL;

	if ((node2=ReadBranch(fv, tree, numNames, names))==NULL)
		return NULL;
	node->branch1=node2;
	node2->branch0=node;
	node->length1=node2->length0;
	ReadUntil(fv, ',', "Comma");
	if (treeError)
		return NULL;
	
	if ((node2=ReadBranch(fv, tree, numNames, names))==NULL)
		return NULL;
	node->branch2=node2;
	node2->branch0=node;
	node->length2=node2->length0;
	
	ch=fgetc(fv);
	while (!feof(fv) && ch!=':' && ch!=',' && ch!=')' && ch!=';') 
		ch=fgetc(fv);
		
	if (detectPolytomies && ch==',') {
		fprintf(stderr, "This tree contains nodes which aren't bifurcations. Resolve the node\n");
		fprintf(stderr, "with zero branch lengths to obtain correct results. This can be done\n");
		fprintf(stderr, "with a program called TreeEdit: http://evolve.zoo.ox.ac.uk/software/TreeEdit\n");
		exit(0);
	}

	if (feof(fv)) {
		sprintf(treeErrorMsg, "Unexpected end of file");
		return NULL;
	}
	ungetc(ch, fv);
	
	return node;
}


TNode *ReadBranch(FILE *fv, TTree *tree, int numNames, char **names)
{
	char ch;
	double len, param=0.0;
	TNode *node;
	
	ch=ReadToNextChar(fv);
	if (ch=='(') {	/* is a node */
		node=ReadNode(fv, tree, numNames, names, 1);
		ReadUntil(fv, ')', "Closing bracket");
		if (treeError)
			return NULL;
	} else {		/* is a tip */
		node=ReadTip(fv, ch, tree, numNames, names);
	}
	
	ch=ReadToNextChar(fv);
	if (ch==':') {
		if (tree->lengths==0) {
			sprintf(treeErrorMsg, "Some branches don't have branch lengths");
			return NULL;
		} else 
			tree->lengths=1;
			
		if (fscanf(fv, "%lf", &len)!=1) {
			sprintf(treeErrorMsg, "Unable to read branch length");
			return NULL;
		}

		ch=ReadToNextChar(fv);
		if (ch=='[') {
			if (fscanf(fv, "%lf", &param)!=1) {
				sprintf(treeErrorMsg, "Unable to read branch parameter");
				return NULL;
			}
			ReadUntil(fv, ']', "Close square bracket");
		} else
			ungetc(ch, fv);
	} else {
		if (tree->lengths==1) {
			sprintf(treeErrorMsg, "Some branches don't have branch lengths");
			return NULL;
		} else 
			tree->lengths=0;
	
		len=0.0;
		ungetc(ch, fv);
	}
	node->length0=len;
	node->param=param;
	
	return node;
}	


void ReadTree(FILE *fv, TTree *tree, int treeNum, int numNames, char **names, 
				int *outNumSites, double *outRelRate)
{
	char ch;
	TNode *P, *node2;
		
	treeError=0;
	tree->numNodes=0;
	tree->numTips=0;
	tree->rooted=1;
	tree->lengths=-1;
	
	(*outRelRate) = 1.0;
	
	ch=fgetc(fv);
	while (!feof(fv) && ch!='(' && ch!='[') 
		ch=fgetc(fv);

	if (ch == '[') {
		if (fscanf(fv, "%d", outNumSites)!=1) {
			sprintf(treeErrorMsg, "Unable to read partition length");
			exit(0);
		}

		ch=fgetc(fv);
		while (!feof(fv) && ch!=',' && ch!='(') 
			ch=fgetc(fv);
			
		if (ch == ',') {
			if (fscanf(fv, "%lf", outRelRate)!=1) {
				sprintf(treeErrorMsg, "Unable to read partition relative rate");
				exit(0);
			}

			ch=fgetc(fv);
			while (!feof(fv) && ch!='(') 
				ch=fgetc(fv);
		}
	}
	
	if (ch!='(' || (tree->root=ReadNode(fv, tree, numNames, names, 0))==NULL) {
		fprintf(stderr, "Error reading tree number %d: %s.\n", treeNum, treeErrorMsg);
		exit(0);
	}
	
	ch=fgetc(fv);
	while (!feof(fv) && ch!=',' && ch!=')' && ch!=';') 
		ch=fgetc(fv);

	if (ch==',') {		/* the tree is unrooted */
		tree->rooted=0;
		if ((tree->root->branch0=ReadBranch(fv, tree, numNames, names))==NULL) {
			fprintf(stderr, "Error reading tree number %d: %s.\n", treeNum, treeErrorMsg);
			exit(0);
		}
		tree->root->branch0->branch0=tree->root;
		tree->root->length0=tree->root->branch0->length0;
	}
	
	tree->totalLength=0.0;
	
	if (tree->rooted) {
		P=tree->root;
		while (P!=NULL) {
			tree->totalLength+=P->length0;
			P=P->branch1;
		}
	}
}


int IsTreeAvail(FILE *fv)
{
	char ch;

	ch=fgetc(fv);
	while (!feof(fv) && ch!='(' && ch!='[') 
		ch=fgetc(fv);
		
	if (ch=='(' || ch=='[')
		ungetc(ch, fv);
		
	return (!feof(fv));
}


void WriteNode(FILE *fv, TTree *tree, TNode *node)
{
	if (node->tipNo==-1) {
		fputc('(', fv);				
		WriteNode(fv, tree, node->branch1);
		fputc(',', fv);			
		WriteNode(fv, tree, node->branch2);
		fputc(')', fv);				
	} else {
		fprintf(fv, "%s", tree->names[node->tipNo]);
	}
	if (tree->lengths)
		fprintf(fv, ":%.6f", node->length0);
}


void WriteTree(FILE *fv, TTree *tree)
{
	fputc('(', fv);	
	if (tree->rooted) {
		WriteNode(fv, tree, tree->root->branch1);
		fputc(',', fv);				
		WriteNode(fv, tree, tree->root->branch2);
	} else {
		WriteNode(fv, tree, tree->root->branch1);
		fputc(',', fv);				
		WriteNode(fv, tree, tree->root->branch2);
		fputc(',', fv);	
		WriteNode(fv, tree, tree->root->branch0);
	}
	fprintf(fv, ");\n");
}


void UnrootRTree(TTree *tree) 
/* Used to unroot a rooted tree */
{
	TNode *P, *Q, *R, *newRoot;
	double len, len2;
		
	if (!tree->rooted || tree->numTips<3)
		return;
	
	P=tree->tips[0];
	Q=P->branch0; newRoot=Q;
	while (Q!=tree->root) {
		R=Q->branch0;
		len=Q->length0;
		if (Q->branch1==P) {
			len2=Q->length1;
			Q->branch1=R;
			Q->length1=len;
		} else {
			len2=Q->length2;
			Q->branch2=R;
			Q->length2=len;
		}
		Q->branch0=P;
		Q->length0=len2;

		P=Q;
		Q=R;
	}

	len=R->length1+R->length2;
	if (R->branch1==P)
		Q=R->branch2;
	else
		Q=R->branch1;

	Q->branch0=P;
	Q->length0=len;
	
	if (P->branch1==R) {
		P->branch1=Q;
		P->length1=len;
	} else {
		P->branch2=Q;
		P->length2=len;
	}
	
	tree->root=newRoot;
	
	DisposeNode(R);
	
	tree->rooted=0;
}


void RerootUTree(TTree *tree, int tip) 
/* Used to reroot an unrooted tree */
/* This may sound strange but all it does is move
   the root trichotomy  */
{
	TNode *P, *Q, *R, *newRoot;
	double len, len2;
		
	if (tree->rooted)
		return;
	
	P=tree->tips[tip];
	Q=P->branch0; newRoot=Q;
	while (P!=tree->root) {
		R=Q->branch0;
		len=Q->length0;
		if (Q->branch1==P) {
			len2=Q->length1;
			Q->branch1=R;
			Q->length1=len;
		} else {
			len2=Q->length2;
			Q->branch2=R;
			Q->length2=len;
		}
		Q->branch0=P;
		Q->length0=len2;

		P=Q;
		Q=R;
	}
	 	
	tree->root=newRoot;
}


void RootUTree(TTree *tree, TNode *node1, TNode *node2, double pos) 
/* Used to root an unrooted tree between node1 and node2 */
/* This is done at pos (0-1.0) proportion of branch length */
{
}


void RerootRTree(TTree *tree, TNode *node1, TNode *node2, double pos) 
/* Used to reroot a rooted tree between node1 and node2 */
/* This is done at pos (0-1.0) proportion of branch length */
{
}