#include <stdio.h>

#ifdef MEMDBG
#include "memdbg.h"
#endif

#include "treestruct.h"

int (*tdup)()=NULL;
int (*tinit)()=NULL;
int (*tfree)()=NULL;

tree maketree(name, rooted)
char *name;
int rooted;
{
	tree tmp;

	tmp=(tree)newspace(tr_tree);

	if(tmp==NULL)
		return(NULL);
	if(name!=NULL)
		tmp->name=(char *)strdup(name);
	else
		tmp->name=NULL;
	tmp->rooted=rooted;
	tmp->root=NULL;
	tmp->branches=newlist();
	if(tmp->branches==NULL)
	{
		free(tmp);
		return(NULL);
	}
	tmp->nodes=newlist();
	if(tmp->nodes==NULL)
	{
		freelist(tmp->branches);
		free(tmp);
		return(NULL);
	}
	tree_set_tree(tmp, tmp);
	tmp->unspecdist=1.0;
	return(tmp);
}

treenode makenode(name)
char *name;
{
	treenode tmp;

	tmp=(treenode)newspace(tr_node);
	if(tmp==NULL)
		return(NULL);
	if(name!=NULL)
		tmp->name=(char *)strdup(name);
	else
		tmp->name=NULL;
	tmp->branches=newlist();
	if(tmp->branches==NULL)
	{
		free(tmp);
		return(NULL);
	}
	tmp->bn=NULL;
	tmp->al=NULL;
	tmp->an=NULL;
	tmp->bd=NULL;
	tmp->ad=NULL;
	return(tmp);
}

node_set_name(n, name)
treenode n;
char *name;
{
	if(n->name!=NULL)
		free(n->name);
	if(name==NULL)
		n->name=NULL;
	else
		n->name=(char *)strdup(name);
}

addsubtreetotree(n, t, up)
/* sets the node n so all it's branches and subtrees belong to t */
/* won't recurse on branch to up */
treenode n, up;
tree t;
{
	treebranch branch;
    treenode tmp;

	/* add node to tree */
	addnode(t->nodes, n);
	tree_set_tree(n, t);

	/* add each subtree to the tree */
    startlist(n->branches);
    while((branch=listnext(n->branches))!=NULL)
    {
        tmp=othernode(n, branch);

		/* add the subtree, if it's not up */
        if(tmp!=up && tmp!=NULL)
			addsubtreetotree(tmp, t, n);
		else
		{
			/* add the branch */
			addnode(t->branches, branch);
			tree_set_tree(branch, t);
		}

    }
}

treebranch addsubtree(up, down, dst, specified)
/* up is in the tree, down is new, down will be descendant of up */
/* if up is a tree, down will be added as the root node */
/* if down has a branch that points to NULL, no new branch will be created */
/* otherwise, a new branch will be created with the specified distance */
treenode up, down;
double dst;
int specified;
{
	treebranch tmp;
	treenode p;
	treebranch pbranch;
	tree t;

	/* check types */
	if(notnode(down))
		return(NULL);
	if(notnode(up) && nottree(up))
		return(NULL);
	/* check that down has no parents */
	startlist(down->branches);

	/* if it has a branch that doesn't point to NULL, then we can't use it */
	pbranch=getparentandbranch(down, &p);
	if(pbranch!=NULL && p!=NULL)
		return(NULL);

	if(pbranch==NULL)
	{
		/* create a branch for down */
		tmp=(treebranch)newspace(tr_branch);
		if(tmp==NULL)
			return(NULL);
		tmp->distance=dst;
		tmp->specified=specified;
		tmp->down=down;
		if(tinit!=NULL)
			(*tinit)(tmp, tr_branch);
		addnode(down->branches, tmp);
	}
	else
		/* use branch that's there */
		tmp=pbranch;

	/* setup pointers to tree or parent node */
	if(tree_get_type(up)==tr_tree)
	{
		t=(tree)up;

		t->root=down;
		/* now the tree points to it as root */

		tmp->up=NULL;
		/* and it's parent pointer points to NULL */
	}
	else
	{
		t=tree_get_tree(up);

		tmp->up=up;
		/* now it's parent pointer points to up */

		endlist(up->branches);
		addnode(up->branches, tmp);
		/* and up has it in it's list of branches, at end */
	}

	/* now, just add the subtree to the tree, set all of it's tree pointers,
		and add the branches and nodes to the tree's lists */
	addsubtreetotree(down, t, up);

	return(tmp);
}

treebranch replacesubtree(up, oldn, newn, dst, specified, useold)
/* oldn is in the tree, newn will be */
/* up must be node or tree */
/* useold is 1 if we are to use the old parent distance, not the one on newn */
treenode up, oldn, newn;
double dst;
int specified;
int useold;
{
	treebranch tmp;
	treenode p;
	treebranch pbranch, b;
	tree t;
	list before;

	/* check types */
	if(notnode(oldn) || notnode(newn))
		return(NULL);
	if(notnode(up) && nottree(up))
		return(NULL);

	/* check that newn has no parents */

	/* if it has a branch that doesn't point to NULL, then we can't use it */
	pbranch=getparentandbranch(newn, &p);
	if(pbranch!=NULL && p!=NULL)
		return(NULL);

	if(pbranch==NULL)
	{
		/* create a branch for newn */
		tmp=(treebranch)newspace(tr_branch);
		if(tmp==NULL)
			return(NULL);
		tmp->distance=dst;
		tmp->specified=specified;
		tmp->down=newn;
		if(tinit!=NULL)
			(*tinit)(tmp, tr_branch);
		startlist(newn->branches);
		addnode(newn->branches, tmp);
	}
	else
		/* use branch that's there */
		tmp=pbranch;
	
	if(useold)
	{
		treebranch oldbranch;

		/* must copy distance from old to new */
		oldbranch=getparentandbranch(oldn, &p);
		if(oldbranch!=NULL)
		{
			tmp->distance=oldbranch->distance;
			tmp->specified=oldbranch->specified;
		}
	}

	/* setup pointers to tree or parent node */
	if(nottree(up))
	{
		/* find point in up's branch list before oldn */
		startlist(up->branches);
		do
		{
			before=listnode(up->branches);
			b=listnext(up->branches);
			if(othernode(up, b)==oldn)
				break;
		} while(b!=NULL);
		/* now before = listnode before the one to be removed */
	}

	/* remove oldn from tree */
	removesubtree(up, oldn);

	if(nottree(up))
	{
		/* set node to accept new subtree in correct spot */
		setnode(up->branches, before);
	}

	/* now add to tree */
	if(tree_get_type(up)==tr_tree)
	{
		/* if up is a tree, add as root */
		t=(tree)up;

		t->root=newn;
		/* now the tree points to it as root */

		tmp->up=NULL;
		/* and it's parent pointer points to NULL */
	}
	else
	{
		/* if up is a node, add in same spot as oldn */
		t=tree_get_tree(up);

		tmp->up=up;
		/* now it's parent pointer points to up */

		addnode(up->branches, tmp);
		/* and up has it in it's list of branches, at end */
	}

	/* now, just add the subtree to the tree, set all of it's tree pointers,
		and add the branches and nodes to the tree's lists */
	addsubtreetotree(newn, t, up);

	return(tmp);
}


removesubtreefromtree(n, t, up)
/* removes the subtree of n from the tree,
	up is the parent, and the tree is t */
/* also adjusts the tree parameters on all the nodes and branches */
treenode n, up;
tree t;
{
    treenode tmp;
	treebranch branch;

	/* remove node from tree */
	findnode(t->nodes, n);
	rmcurr(t->nodes);
	tree_set_tree(n, NULL);

	/* update tree root pointer */
	if(t->root==n)
		t->root=NULL;

	/* remove each subtree from the tree */
    startlist(n->branches);
    while((branch=listnext(n->branches))!=NULL)
    {
        tmp=othernode(n, branch);

		/* remove the subtree, if it's not up */
        if(tmp!=up && tmp!=NULL)
			removesubtreefromtree(tmp, t, n);
		else
		{
			/* remove the branch */
			findnode(t->branches, branch);
			rmcurr(t->branches);
			tree_set_tree(branch, NULL);
		}

    }
}

int removesubtree(up, down)
/* removes down from the tree, and all nodes on it that are on the opposite
	side from up.  If up is a tree, removes all nodes in the tree */
treenode up, down;
{
	treebranch branch;
	treenode nextto;
	tree t;
	int isnode;
	list old;

    if((nottree(up) && notnode(up)) || notnode(down))
        return(0);

	if(tree_get_type(up)==tr_tree)
	{
		isnode=0;
		t=(tree)up;
	}
	else
	{
		isnode=1;
		t=tree_get_tree(up);
	}

    if(t != tree_get_tree(down))
        return(0);

	if(!isnode)
		nextto=NULL;
	else
	{
		/* get the branch that goes to up, and find the first node on it */
		branch=branchtonode(down, up);

		/* set the other end of the branch of down,
			now it links down with nothing, instead of with up */
		if(branch->up==down)
			branch->down=NULL;
		else
			branch->up=NULL;

		/* remove branch from node */
		findnode(up->branches, branch);
		rmcurr(up->branches);
	}
	/* remove all nodes and branches */
	removesubtreefromtree(down, t, nextto);

	/* may have removed the root node from the tree, so setup a new one */
	if(t->root==NULL)
	{
		if(isnode)
			/* if user passed a node, be friendly and use that one */
			t->root=up;
		else
		{
			/* no node passed, so use the first one left in the tree */
			startlist(t->nodes);
			t->root=listnext(t->nodes);
			/* t->root may now still be NULL, if there are no nodes left
				in the tree */
		}
	}

	return(1);
}

list dupcommentlist(cm)
list cm;
{
	char *comment;
	list newcm;

	if(cm==NULL)
		return(NULL);
	newcm=newlist();
	if(newcm==NULL)
		return(NULL);

	startlist(cm);
	while((comment=listnext(cm))!=NULL)
		addnode(newcm, strdup(comment));
}

treenode dupnode(n)
treenode n;
{
	treenode newnode;

	newnode=makenode(n->name);
	newnode->bn=dupcommentlist(n->bn);
	newnode->al=dupcommentlist(n->al);
	newnode->an=dupcommentlist(n->an);
	newnode->bd=dupcommentlist(n->bd);
	newnode->ad=dupcommentlist(n->ad);
	n->duped=newnode;
	return(newnode);
}

treenode duplicatenode(n)
treenode n;
{
	treenode newnode;

	newnode=dupnode(n);
	if(newnode==NULL)
		return(NULL);
	if(tdup!=NULL)
		(*tdup)(n, newnode, tr_node, tree_get_data(n));
}

treebranch dupbranch(b)
treebranch b;
{
	treebranch newbranch;

	newbranch=(treebranch)newspace(tr_branch);
	if(newbranch==NULL)
		return(NULL);
	newbranch->distance=b->distance;
	newbranch->specified=b->specified;
	return(newbranch);
}

treebranch duplicatebranch(n)
treebranch n;
{
	treebranch newbranch;

	newbranch=dupbranch(n);
	if(newbranch==NULL)
		return(NULL);
	if(tdup!=NULL)
		(*tdup)(n, newbranch, tr_branch, tree_get_data(n));
}

treenode duplicatesubtree1(n, up, newup)
/* duplicates node n and it's branch that points to up, makes the new branch
	point to newup.  Duplicates n's subtrees */
treenode n, up, newup;
{
	treebranch br, newbranch;
	treenode newnode;

	newnode=dupnode(n);
	/* only thing missing on our new node is the branches.  The tree for our
		new node is correct, it points to NULL since this node doesn't belong
		to a tree yet */

	startlist(n->branches);
	while((br=listnext(n->branches))!=NULL)
	{
		if(br->up==n)
		{
			if(br->down==NULL)
			{
				/* just duplicate the branch */
				newbranch=dupbranch(br);
				addnode(newnode->branches, newbranch);
				newbranch->up=newnode;
				newbranch->down=NULL;
				if(tdup!=NULL)
					(*tdup)(br, newbranch, tr_branch, tree_get_data(br));
			}
			else if(br->down==up)
			{
				/* duplicate the branch and add to parent */
				newbranch=dupbranch(br);
				addnode(newnode->branches, newbranch);
				newbranch->up=newnode;
				newbranch->down=newup;
				if(newup!=NULL)
					addnode(newup->branches, newbranch);
				if(tdup!=NULL)
					(*tdup)(br, newbranch, tr_branch, tree_get_data(br));
			}
			else
				/* duplicate the subtree */
				/* already adds itself to newnode */
				duplicatesubtree1(br->down, n, newnode);
		}
		else /* br->down==n */
		{
			if(br->up==NULL)
			{
				/* just duplicate the branch */
				newbranch=dupbranch(br);
				addnode(newnode->branches, newbranch);
				newbranch->down=newnode;
				newbranch->up=NULL;
				if(tdup!=NULL)
					(*tdup)(br, newbranch, tr_branch, tree_get_data(br));
			}
			else if(br->up==up)
			{
				/* duplicate the branch and add to parent */
				newbranch=dupbranch(br);
				addnode(newnode->branches, newbranch);
				newbranch->down=newnode;
				newbranch->up=newup;
				if(newup!=NULL)
					addnode(newup->branches, newbranch);
				if(tdup!=NULL)
					(*tdup)(br, newbranch, tr_branch, tree_get_data(br));
			}
			else
				/* duplicate the subtree */
				/* already adds itself to newnode */
				duplicatesubtree1(br->up, n, newnode);
		}
		/* now the branch is duplicated, with it's subtree if any */
	}
	
	if(tdup!=NULL)
		(*tdup)(n, newnode, tr_node, tree_get_data(n));

	return(newnode);
}

treenode duplicatesubtree(up, down)
treenode up, down;
{
	treebranch branch;
	treenode nextto;
	treenode newnode;
	tree t;
	int isnode;

    if(notnode(down))
        return(0);

	if(up==NULL)
	{
		t=tree_get_tree(down);
		isnode=0;
	}
	else if(tree_get_type(up)==tr_tree)
	{
		isnode=0;
		t=(tree)up;
	}
	else if(tree_get_type(up)==tr_node)
	{
		isnode=1;
		t=tree_get_tree(up);
	}
	else
		return(0);

    if(t != tree_get_tree(down))
        return(0);

	if(!isnode)
		nextto=NULL;
	else
	{
		/* get the branch that goes to up, and find the first node on it */
		branch=branchtonode(down, up);
		nextto=othernode(down, branch);
	}
	/* duplicate the subtree */
	newnode=duplicatesubtree1(down, nextto, NULL);

	return(newnode);
}

freecommentlist(cm)
list cm;
{
	char *comment;

	if(cm==NULL)
		return;
	startlist(cm);
	while((comment=listnext(cm))!=NULL)
		free(comment);
	freelist(cm);
}

freenode(n)
treenode n;
{
	if(n->name!=NULL)
		free(n->name);
	freelist(n->branches);
	freecommentlist(n->bn);
	freecommentlist(n->al);
	freecommentlist(n->an);
	freecommentlist(n->bd);
	freecommentlist(n->ad);
	free(n);
}

freebranch(b)
treebranch b;
{
	free(b);
}

freesubtree1(n, up)
/* frees the node, with all branches that lead out from it, and all nodes
	on those branches */
treenode n, up;
{
    treenode tmp;
    treebranch branch;

    /* notify the user */
    if(tfree!=NULL)
		(*tfree)(n, tr_node, tree_get_data(n));

    startlist(n->branches);
    while((branch=listnext(n->branches))!=NULL)
    {
        tmp=othernode(n, branch);

        if(tmp==up || tmp==NULL)
	    {
		/* just free the branch */
		if(tfree!=NULL)
		    (*tfree)(branch, tr_branch, tree_get_data(branch));
		freebranch(branch);
	    }
	    else
	    {
		/* this will free the branch too */
		freesubtree1(tmp, n);
	    }
    }

    /* free this node */
    freenode(n);
}

int freesubtree(n)
/* frees a subtree.  Must have removed from tree first */
treenode n;
{
	if(notnode(n))
		return(0);

	freesubtree1(n, NULL);
	return(1);
}

int freetree(t)
/* frees a tree. */
tree t;
{
	if(nottree(t))
		return(0);
	if(t->root!=NULL)
		return(0);
	
	freelist(t->branches);
	freelist(t->nodes);
	if(t->name!=NULL)
		free(t->name);
	free(t);
	return(1);
}

tree duplicatetree(t)
tree t;
{
	tree tmp;
	treenode newnode;

	if(nottree(t))
		return(NULL);

	tmp=maketree(t->name, t->rooted);
	if(tmp==NULL)
		return(NULL);
	newnode=duplicatesubtree(NULL, t->root);
	(void) addsubtree(tmp, newnode, 1.0, 0);
	if(tdup!=NULL)
		(*tdup)(t, tmp, tr_tree, tree_get_data(t));

	return(tmp);
}

int treereroot(n)
/* makes n the new root of the tree it's in */
/* works if it's not in a tree, too */
treenode n;
{
	treenode p;  /* parent of current node */
	treenode c;  /* current node */
	treebranch l;/* last treebranch we worked with */
	treebranch b;
	tree t;
	int more;

	if(notnode(n))
		return(0);

	more=1;
	c=n;
	l=NULL;
	/* work to the root, flipping branches */
	while(more)
	{
		startlist(c->branches);
		/* find branch to parent */
		while((b=listnext(c->branches))!=NULL)
			if(b->up!=c && b!=l)
				/* if this is the parent branch */
				break;
		if(b!=NULL)
		{
		    /* found a parent branch, must flip it or move it */
		    p=b->up;
		    if(p!=NULL)
		    {
			/* there's a node above it, so flip direction */
			b->up=c;
			b->down=p;
			/* go to parent node and do it again */
			c=p;
			l=b;
		    }
		    else
		    {
			/* there's no node above it, this is the
			    dangling distance */
			/* so have it point to n now */
			b->down=n;
			/* remove it from the current node */
			rmcurr(c->branches);
			/* add it to our new root */
			addnode(n->branches, b);
			more=0;
		    }
		}
		else
		    more=0;
	}
	t=tree_get_tree(n);
	if(t!=NULL)
		t->root=n;
	return(1);
}

int fliplist(n)
/* flips the list of subtrees for the given node */
treenode n;
{
	treebranch b;
	list tmp;

	if(notnode(n))
		return(0);
	
	tmp=newlist();
	startlist(n->branches);
	b=listnext(n->branches);
	while(b!=NULL)
	{
		startlist(tmp);
		addnode(tmp, b);
		rmcurr(n->branches);
		startlist(n->branches);
		b=listnext(n->branches);
	}
	/* now, all added in opposite order to tmp, and removed from n->branches */
	/* now, let's swap lists */
	freelist(n->branches);
	n->branches=tmp;
	n->currentbranch=NULL;

	return(1);
}

treenode break_branch(b, p)
/* breaks a branch at p% from up node, inserting a new node */
treebranch b;
double p;
{
	treenode newnode;
	treebranch newbranch;
	tree t;

	t=tree_get_tree(b);
	/* make new node */
	newnode=makenode(NULL);
	/* setup branch completely */
	newbranch=(treebranch)newspace(tr_branch);
	if(b->specified)
	{
		newbranch->distance=p*b->distance;
		newbranch->specified=1;
	}
	else
	{
		newbranch->specified=0;
		newbranch->distance=t->unspecdist;
	}
	newbranch->up=b->up;
	newbranch->down=newnode;
	/* add both to tree */
	addnode(t->branches, newbranch);
	addnode(t->nodes, newnode);
	if(t->root==newbranch->down)
		t->root=newnode;
	addnode(newnode->branches, newbranch);
	addnode(newnode->branches, b);
	b->up=newnode;
	if(b->specified)
		b->distance=(1-p)*b->distance;
	if(newbranch->up!=NULL)
	{
		/* replace the old branch pointer with the new one */
		findnode(newbranch->up->branches, b);
		setcurr(newbranch->up->branches, newbranch);
	}
	tree_set_tree(newbranch, t);
	tree_set_tree(newnode, t);
	if(tinit!=NULL)
		(*tinit)(newbranch, tr_branch);
	if(tinit!=NULL)
		(*tinit)(newnode, tr_node);
	return(newnode);
}

treenode merge_branch(n)
/* if the treenode has only 1 child, removes the treenode and concats the
	branch */
treenode n;
{
	treebranch b, c, d;
	treenode p;
	tree t;

	/* find if there is more than one child */
	startlist(n->branches);
	c=listnext(n->branches);
	d=listnext(n->branches);
	/* find first real child */
	if(c->up!=n)
		c=d;
	/* if there is another, too many */
	b=listnext(n->branches);
	if(b!=NULL)
		return(NULL);
	/* if there is none, too few */
	if(c==NULL)
		return(NULL);

	t=tree_get_tree(n);
	/* get parent and branch to parent */
	/* c already = branch to child */
	b=getparentandbranch(n, &p);
	/* update distance on branch to child */
	if(b->specified)
		c->distance+=b->distance;
	/* remove node from tree list */
	findnode(t->nodes, n);
	rmcurr(t->nodes);
	/* remove branch to parent from tree list */
	findnode(t->branches, b);
	rmcurr(t->branches);
	/* have branch to child now point to parent */
	c->up=p;
	if(p!=NULL)
	{
		/* not merging root node, so change branch list of parent to
			include the branch to child, not the old branch */
		findnode(p->branches, b);
		setcurr(p->branches, c);
	}
	else
		/* adjust root node */
		t->root=c->down;

    if(tfree!=NULL)
	{
		/* call free routines */
		(*tfree)(n, tr_node, tree_get_data(n));
		(*tfree)(b, tr_branch, tree_get_data(b));
	}
	/* free data */
	freenode(n);
	freebranch(b);
	/* return child, the node that has taken this node's place */
	return(c->down);
}

treenode change_outgroup(n, percent)
/* makes n the new outgroup of the tree it's in */
/* cuts n's parent branch at percent*dist from the parent,
	adds the new root node there */
/* reroots the tree at the new root node */
/* checks old root node, if only one child, removes and merges branch */
/* works only if n is in a tree */
treenode n;
double percent;
{
	treebranch b;
	treenode parent, newroot, oldroot;
	tree t;

	if(notnode(n))
		return(NULL);

	t=tree_get_tree(n);
	b=getparentandbranch(n, &parent);
	newroot=break_branch(b, percent);
	oldroot=treeroot(t);
	treereroot(newroot);
	/*remove old root if necessary*/
	merge_branch(oldroot);
	return(newroot);
}