// $Id: newick.cpp,v 1.32 2013/11/07 22:52:33 jyamato Exp $

#include <cassert>
#include <utility>                     // for make_pair in recombinant LamarcTreeToNewickString

#include "newick.h"
#include "tree.h"
#include "stringx.h"
#include "constants.h"                  // for FLAGDOUBLE
#include "branch.h"
#include "timelist.h"                   // to build the tree
#include "registry.h"
#include "force.h"                      // for Force::GetMaximum()
#include "runreport.h"                  // for ReportUrgent() in NewickConverter::
                                        //   LamarcTreeToNewickString()

//------------------------------------------------------------------------------------

const double DEFAULTLENGTH = 0.01;

// DEBUG Comments, including migration nodes, are ignored
// DEBUG Nested comments do not work yet (need to fix ProcessComments)
// NB No provision is made for recombination

//------------------------------------------------------------------------------------

NewickTree::NewickTree(const string& tree)
    : m_newicktree(tree),
      m_curr_char(0),
      m_numbers("0123456789.-+eE"),     // things found in numbers
      m_terminators(":,)[")             // things terminating names
{
    // intentionally blank
} // NewickTree constructor

//------------------------------------------------------------------------------------

void NewickTree::ToLamarcTree(Tree& stump)
{
    // create a holder structure to store ongoing info about the tree
    NewickNode base(stump);

    // create a working pointer into the holder
    NewickNode* current = &base;

    // begin at start of tree string
    m_curr_char = 0;

    while (m_newicktree[m_curr_char] != ';') // semicolon is end of tree
    {
        switch(m_newicktree[m_curr_char])
        {
            case '(':
                ++m_curr_char;
                // create first daughter
                current = current->AddChild();
                break;
            case ',':
                ++m_curr_char;
                // create additional daughter
                current = current->GetParent()->AddChild();
                break;
            case ')':
                ++m_curr_char;
                // coalesce daughters
                current = current->GetParent()->Terminate();
                break;
            case ' ':
            case '\n':
            case '\t':
            case '\r':
                ++m_curr_char;
                // skip whitespace
                break;
            case ':':
                current->SetLength(ProcessLength());
                break;
            case '[':
                ProcessComment();
                break;
            default:
                // Anything unrecognized must be a tip name
                current->AddBranch(ProcessName(stump));
                break;
        }
    }

    // hookup the lamarc tree
    base.Coalesce();

    stump.AttachBase(base.GetBranch());

} // ToLamarcTree

//------------------------------------------------------------------------------------

double NewickTree::ProcessLength()
{
    ++m_curr_char;  // skip the colon
    unsigned long endpos = m_newicktree.find_first_not_of(m_numbers, m_curr_char);
    unsigned long ndigits = endpos - m_curr_char;
    string lengthstr = m_newicktree.substr(m_curr_char, ndigits);
    double length;
    if (FromString(lengthstr,length))
    {
        m_curr_char = endpos;
        return length;
    }
    else
    { // error handling
        assert(false); // bad newick tree, not handled yet
        return FLAGDOUBLE;
    }

} // ProcessLength

//------------------------------------------------------------------------------------

Branch_ptr NewickTree::ProcessName(const Tree& stump)
{
    unsigned long endpos = m_newicktree.find_first_of(m_terminators, m_curr_char);
    unsigned long nchars = endpos - m_curr_char;
    string name = m_newicktree.substr(m_curr_char, nchars);
    m_curr_char = endpos;
    return stump.GetTip(name);
} // ProcessName

//------------------------------------------------------------------------------------

void NewickTree::ProcessComment()
// NB: We don't do anything with the contents of the comment yet
{
    unsigned long endpos = m_newicktree.find(']', m_curr_char);
    // comment parsing would go in here
    m_curr_char = endpos;
} // ProcessComment

//------------------------------------------------------------------------------------
//------------------------------------------------------------------------------------

NewickNode::NewickNode(Tree& stump)
    : m_tree(stump),
      m_parent(NULL),
      m_length(0.0)
{
    // deliberately blank
} // NewickTree constructor

//------------------------------------------------------------------------------------

NewickNode::~NewickNode()
{
    vector<NewickNode*>::iterator it;

    for (it = m_children.begin(); it != m_children.end(); ++it)
    {
        delete *it;
    }
} // NewickNode dtor

//------------------------------------------------------------------------------------

NewickNode* NewickNode::AddChild()
{
    NewickNode* newnode = new NewickNode(m_tree);
    m_children.push_back(newnode);
    newnode->m_parent = this;
    return newnode;
} // AddChild

//------------------------------------------------------------------------------------

NewickNode* NewickNode::Terminate()
{
    assert(m_children.size() > 1);

    return this;
} // Terminate

//------------------------------------------------------------------------------------

NewickNode* NewickNode::GetParent() const
{
    return m_parent;
} // GetParent

//------------------------------------------------------------------------------------

Branch_ptr NewickNode::GetBranch() const
{
    return m_branch;
} // GetBranch

//------------------------------------------------------------------------------------

void NewickNode::AddBranch(Branch_ptr br)
{
    m_branch = br;
} // AddBranch

//------------------------------------------------------------------------------------

void NewickNode::SetLength(double newlength)
{
    m_length = newlength;
} // SetLength

//------------------------------------------------------------------------------------

double NewickNode::Coalesce()
{
    if (m_children.empty())             // we're at a tip
    {
        return m_length;
    }

    double eventtime = FLAGDOUBLE;      // initialized to invalid value

    vector<NewickNode*>::const_iterator kid;
    for(kid = m_children.begin(); kid != m_children.end(); ++kid)
    {
        eventtime = (*kid)->Coalesce(); // These should all be the same.
                                        // There may be differences due to
                                        // rounding, but we ignore them.
    }

    assert(m_children[0]->m_branch && m_children[1]->m_branch);
    assert(eventtime != FLAGDOUBLE);

    rangeset fcsites;                   // No sites are fc, so this is empty.
    m_branch = m_tree.Coalesce(m_children[0]->m_branch, m_children[1]->m_branch, eventtime, fcsites);

    return eventtime + m_length;
} // Coalesce

//------------------------------------------------------------------------------------
//------------------------------------------------------------------------------------

string NewickConverter::LamarcTreeToMSNewickString(const Tree& tree) const
{
    string newick;

    if (registry.GetForceSummary().CheckForce(force_REC))
    {
        rangepair span(std::make_pair(0L, tree.GetNsites()));
        rangevector subtrees(tree.GetLocusSubtrees(span));
        long nsubtrees(subtrees.size());
        long tr;
        for(tr = 0; tr < nsubtrees; ++tr)
        {
            // print the range
            long length = subtrees[tr].second - subtrees[tr].first;
            newick += "[" + ToString(length) + "]";
            // print the rangetree
            long endval(subtrees[tr].second);
            --endval; // to deal with open ended ranges
            RecurseWriteIntervalTreeNewick(newick,endval,tree.GetTimeList().Root());
            newick += ":0.0;\n";
        }
    }
    else
    {
        if (tree.GetTimeList().ContainsOnlyTipsAndCoals())
        {
            RecurseWriteNewick(newick,tree.GetTimeList().Root());
            newick += ":0.0;\n";
        }
        else
        {
            string msg("WARNING--tried to make a newick tree from an ");
            msg += "enhanced coalescent tree.  Newick formatting failed.\n";
            registry.GetRunReport().ReportUrgent(msg);
        }
    }

    return newick;
} // LamarcTreeToMSNewickString

//------------------------------------------------------------------------------------

string NewickConverter::LamarcTreeToNewickString(const Tree& tree) const
{
    string newick;

    if (registry.GetForceSummary().CheckForce(force_REC))
    {
        rangepair span(std::make_pair(0L, tree.GetNsites()));
        rangevector subtrees(tree.GetLocusSubtrees(span));
        long nsubtrees(subtrees.size());
        long tr;
        for(tr = 0; tr < nsubtrees; ++tr)
        {
            // print the range
            newick += "[" + ToString(subtrees[tr].first);
            long endval(subtrees[tr].second);
            --endval; // to deal with open ended ranges
            newick += "," + ToString(endval) + "]\n";
            // print the rangetree
            RecurseWriteIntervalTreeNewick(newick,endval,tree.GetTimeList().Root());
            newick += ":0.0;\n";
        }
    }
    else
    {
        if (tree.GetTimeList().ContainsOnlyTipsAndCoals())
        {
            RecurseWriteNewick(newick,tree.GetTimeList().Root());
            newick += ":0.0;\n";
        }
        else
        {
            string msg("WARNING--tried to make a newick tree from an ");
            msg += "enhanced coalescent tree.  Newick formatting failed.\n";
            registry.GetRunReport().ReportUrgent(msg);
        }
    }

    return newick;
} // LamarcTreeToNewickString

//------------------------------------------------------------------------------------

double NewickConverter::RecurseWriteNewick(string& newick, Branch_ptr pBranch) const
{
    // NB This code makes some simplifying assumptions:
    // No recombination; mutation and disease are ignored.
    // No more than 2 children per branch; no "one-legged coalescence nodes".

    double time, newtime;

    // Skip a disease, epoch, or migration branch.
    if (pBranch->Event() == btypeMig || pBranch->Event() == btypeDisease ||
        pBranch->Event() == btypeEpoch || pBranch->Event() == btypeDivMig)
    {
        time = RecurseWriteNewick(newick, pBranch->Child(0));
        return time + (pBranch->m_eventTime - pBranch->Child(0)->m_eventTime);
    }

    // Terminate on a tip branch.
    if (pBranch->BranchGroup() == bgroupTip)
    {
        // write tip name
        newick += boost::dynamic_pointer_cast<TBranch>(pBranch)->m_label;
        return 0.0;
    }

    // Process a coalescence branch and recurse.
    assert(pBranch->Event() == btypeCoal);
    assert(pBranch->Child(0));
    assert(pBranch->Child(1));

    newick += '(';

    // Recurse left.
    time = RecurseWriteNewick(newick, pBranch->Child(0));
    newick += ':';
    newtime = time + pBranch->m_eventTime - pBranch->Child(0)->m_eventTime;
    newick += ToString(newtime);
    newick += ',';

    // Recurse right.
    time = RecurseWriteNewick(newick, pBranch->Child(1));
    newick += ':';
    newtime = time + pBranch->m_eventTime - pBranch->Child(1)->m_eventTime;
    newick += ToString(newtime) + ')';

    // Return newtime; this return value will lead to node times being printed.
    return 0.0;  // return zero if not trying to accumulate lengths

} // RecurseWriteNewick

//____________________________________________________________________________________

double NewickConverter::RecurseWriteIntervalTreeNewick(string& newick, long int site,
                                                       Branch_ptr pBranch) const
{
    // NB This code makes some simplifying assumptions:
    // Mutation and disease are legal but ignored.
    // No more than 2 children per branch.

    double time, newtime;

    // Skip a migration, disease, recombination, epoch, or div-mig branch.
    if (pBranch->Event() == btypeMig || pBranch->Event() == btypeDisease ||
        pBranch->Event() == btypeRec || pBranch->Event() == btypeEpoch ||
        pBranch->Event() == btypeDivMig)
    {
        time = RecurseWriteIntervalTreeNewick(newick, site, pBranch->Child(0));
        return time + (pBranch->m_eventTime - pBranch->Child(0)->m_eventTime);
    }

    // Terminate on a tip branch.
    if (pBranch->BranchGroup() == bgroupTip)
    {
        // write tip name
        newick += boost::dynamic_pointer_cast<TBranch>(pBranch)->m_label;
        return 0.0;
    }

    // Process a coalescence branch and recurse.
    assert(pBranch->Event() == btypeCoal);
    assert(pBranch->Child(0));
    assert(pBranch->Child(1));

    // Deal with a "one-legged" coalescence for site "site".
    if (!pBranch->Child(0)->GetRangePtr()->IsSiteLive(site))
    {
        time = RecurseWriteIntervalTreeNewick(newick, site, pBranch->Child(1));
        return time + (pBranch->m_eventTime - pBranch->Child(1)->m_eventTime);
    }

    if (!pBranch->Child(1)->GetRangePtr()->IsSiteLive(site))
    {
        time = RecurseWriteIntervalTreeNewick(newick, site, pBranch->Child(0));
        return time + (pBranch->m_eventTime - pBranch->Child(0)->m_eventTime);
    }

    // We must have a normal "two-legged" coalescence here.
    newick += '(';

    time = RecurseWriteIntervalTreeNewick(newick, site, pBranch->Child(0));
    newick += ':';
    newtime = time + pBranch->m_eventTime - pBranch->Child(0)->m_eventTime;
    newick += ToDecimalString(newtime);
    newick += ',';

    time = RecurseWriteIntervalTreeNewick(newick, site, pBranch->Child(1));
    newick += ':';
    newtime = time + pBranch->m_eventTime - pBranch->Child(1)->m_eventTime;
    newick += ToDecimalString(newtime) + ')';

    // Return newtime; this return value will lead to node times being printed.
    return 0.0;  // return zero if not trying to accumulate lengths

} // RecurseWriteIntervalTreeNewick

//____________________________________________________________________________________
