#include <u.h>
#include <libn.h>
#include <libbio.h>

// -----------------------------------------------------------------------
// subtree manipulation methods
// NOTE: As of now these don't update nnode & nleaf stats.
//       It is the caller's responsibility to refresh counts.

error
phylo·addchild(bio·Node* parent, bio·Node* child)
{
    bio·Node *it, *sibling;
    if (!parent->nchild) {
        parent->child = child;
        goto SUCCESS;
    }

    for (it = parent->child, sibling = it; it != nil; it = it->sibling) {
        sibling = it;
    }
    sibling->sibling = child;

SUCCESS:
    child->parent = parent;
    parent->nchild++;
    return 0;
}

error
phylo·rmchild(bio·Node *parent, bio·Node *child)
{
    bio·Node *it, *prev;
    enum {
        error·nil,
        error·notfound,
        error·nochildren,
    };

    prev = nil;
    for (it = parent->child; it != nil; it = it->sibling) {
        if (it == child) goto FOUND;
        prev = it;
    }
    return error·notfound;

FOUND:
    if (prev == nil) {
        parent->child = child->sibling;
    } else {
        prev->sibling = child->sibling;
    }
    return error·nil;
}

// -----------------------------------------------------------------------
// subtree statistics

error
phylo·countnodes(bio·Node *node, int *n)
{
    int       m;
    error     err;
    bio·Node *child;
    
    m = *n;
    for (child = node->child; child != nil; child = child->sibling) {
        if (err = phylo·countnodes(child, n), err) {
            errorf("node count: failure at '%s'", child->name);
            return 1;
        }
    }
    node->nnode = *n - m;
    *n += 1;

    return 0;
}

error
phylo·countleafs(bio·Node *node, int *n)
{
    error     err;
    bio·Node *child;
    
    if (!node->nchild) {
        *n += 1;
    }

    for (child = node->child; child != nil; child = child->sibling) {
        if (err = phylo·countleafs(child, n), err) {
            errorf("leaf count: failure at '%s'", child->name);
            return 1;
        }
    }

    return 0;
}

// -----------------------------------------------------------------------
// tree editing

static
void
sortnodelist(bio·Node **head, bio·Node *next)
{
    bio·Node tmp, *it;

    it          = &tmp;
    tmp.sibling = *head;

    while (it->sibling != nil && it->sibling->nnode < next->nnode) {
        it = it->sibling;
    }

    next->sibling = it->sibling;
    it->sibling = next;
    *head = tmp.sibling;
}

error
phylo·ladderize(bio·Node *root)
{
    int       i;
    error     err;
    bio·Node *child, *sorted, *sibling;

    if (!root->nchild) return 0;

    // ladderize below 
    for (child = root->child; child != nil; child = child->sibling) {
        if (err = phylo·ladderize(child), err) {
            errorf("ladderize: failure at '%s'", child->name);
            return 1;
        }
    }

    // ladderize yourself
    sorted = nil;
    child  = root->child;
    while (child != nil) {
        sibling = child->sibling;
        sortnodelist(&sorted, child);
        child   = sibling;
    }
    root->child = sorted;

    return 0;
}

static
error
rotateparent(bio·Node *node, bio·Node *to)
{
    error err;
    bio·Node *it, *prev;

    if (node->parent == to) {
        return 0;
    }

    prev = nil;
    for (it = node->child; it != nil; it = it->sibling) {
        if (it == to) goto FOUND;
        prev = it;
    }
    errorf("inconsistent topology: attempting to rotate to non-child");
    return 1;

FOUND:
    if (err = rotateparent(node->parent, node), err) {
        errorf("failure: broken tree");
        return err;
    }

    if (!prev) {
        node->child   = node->parent;
    } else {
        prev->sibling = node->parent;
    }
    node->parent->sibling = it->sibling;
    node->parent          = to;

    return 0;
}

#define FLOAT_PREC .00000001
error
phylo·reroot(bio·Tree *tree, bio·Node *node, double d)
{
    bio·Node *old;
    bio·Node *new;

    // TODO: should check that node is part of this tree?
    
    old = tree->root;
    if (fabs(d) < FLOAT_PREC)            new = node;
    if (fabs(d-node->dist) < FLOAT_PREC) new = node->parent;
    else {
        new = tree->heap.alloc(tree->h, 1, sizeof(*new));
        memset(new, 0, sizeof(*new));
    }

    // TODO: Use rotateparent() to perform necessary flips

    tree->root = new;
    return 0;
}

// -----------------------------------------------------------------------
// ancestral inference

struct phylo·InferOpts
{
    int     nstates;
    double  *Q;
};

error
phylo·ancestralinfer(bio·Tree *tree, struct phylo·InferOpts opts)
{
    return 0;
}