// $Id: argtree.cpp,v 1.8 2018/01/03 21:33:03 mkkuhner Exp $

/*
Copyright 2002-2005 Mary K. Kuhner, Peter Beerli, and Joseph Felsenstein

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

*/

/* Authors: Jim McGill, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>
#include <iostream>                     // for debug cerr
#include <map>
#include <stdio.h>                      // for use of printf in assorted functions
#include <utility>                      // for make_pair in recombinant

#include "local_build.h"

#include "argtree.h"
#include "constants.h"                  // for FLAGDOUBLE
#include "rangex.h"
#include "registry.h"
#include "stringx.h"
#include "tinyxml.h"
#include "tree.h"
#include "ui_strings.h"
#include "xml_strings.h"

using namespace std;

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

ARGEdge::ARGEdge(long target, long source)
    : m_target(target),
      m_source(source),
      m_targetptr(),
      m_partitions(""),
      m_livesites(""),
      m_transmittedsites(""),
      m_targetid(FLAGLONG),
      m_type(""),
      m_time(FLAGDOUBLE),
      m_label(""),
      m_recloc(FLAGLONG)
{
    // intentionally blank
} // ARGEdge constructor

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

FILE* TiXmlFOpen( const char* filename, const char* mode );  // pick this up from tinyxml

void ARGTree::ToLamarcTree(Tree& stump, vector<ARGEdge> argedges)
{
    // This is an ugly piece of code, that takes the edge centric translation
    // of an input ARG tree and grafts it onto an existing Lamarc tree. This is something
    // Lamarc was never designed to expect, so the process is arcane and convoluted.

#ifndef NDEBUG
    // debug print of edge list
    printf("\n****Initial edge list passed into ARGTree::ToLamarcTree****\n");
    for(size_t edge=0; edge<argedges.size(); edge++)
    {
        printf("\nedge: %li  target: %li source: %li\n", edge, (argedges[edge].GetTarget()), argedges[edge].GetSource());
        printf("Partitions: %s\n", argedges[edge].GetPartitions().c_str());
        printf("Live sites: %s\n",  argedges[edge].GetLiveSites().c_str());
        printf("Transmitted sites: %s\n",  argedges[edge].GetTransmittedSites().c_str());
        printf("Label: %s\n", argedges[edge].GetLabel().c_str());
        printf("Type: %s  Time: %f\n", argedges[edge].GetType().c_str(), argedges[edge].GetTime());
        printf("Rec Loc: %li\n", (argedges[edge].GetRecLoc()));
    }
#endif

    // target/source connections
    typedef pair<Branch_ptr, Branch_ptr> targetpair;
    map<long, targetpair> targetmap;
    branchpair recbranches;

    // find out what kind of tree is being built
    bool isRecTree = true;
    //bool firstRecFound = true;  //JRM Debug
    Tree* testTree = &stump;
    RecTree* recTree = dynamic_cast<RecTree*>(testTree);
    if (recTree == NULL)
    {
        isRecTree = false;
    }
    m_totalSites = 0;

    // The following is tricky because it depends on the argedge vector being time ordered
    // It will blow up if the ARG tree reading in parsetreetodata is not done right
    // All the tips are first because they are at time 0
    // The tree is then built rootward in time order so that at every
    // step the targetmap will contain the branch being made
    // It's elegant but obscure.
    Branch_ptr newbranch;
    bool forcefound;
    bool recfound;
    size_t recpar1 = 0;
    size_t recpar2 = 0;
    FC_Status fc_status;
    for(size_t edge=0; edge<argedges.size(); edge++)
    {
        forcefound = false;
        recfound   = false;
        if (argedges[edge].GetType() == xmlstr::XML_BRANCHTYPE_TIP)
        {
            forcefound = true;
            // find the tip branch pointer in the current tree using the name
            string tipname = argedges[edge].GetLabel();
            newbranch = stump.GetTip(tipname);
            fc_status.Increment_FC_Counts(newbranch->GetLiveSites());

            //cerr << "tip: " << newbranch->GetID() << endl;
            argedges[edge].SetTargetPtr(newbranch);
            TransferEdgeData(newbranch, argedges[edge]);

#if 0
            if (m_totalSites == 0)
            {
                // set this the first time through, from the tips which have the total number of sites
                m_totalSites = newbranch->GetRangePtr()->NumRegionSites();
            }
#endif
        }

        else if (argedges[edge].GetType() == xmlstr::XML_BRANCHTYPE_COAL)
        {
            forcefound = true;
            map<long, targetpair>::iterator pos = targetmap.find(argedges[edge].GetTarget());
            if (pos != targetmap.end())
            {
                targetpair p = pos->second;
                if ((p.first != NULL) && (p.second != NULL))
                {
                    // make coalesence
                    rangeset fcsites;
                    fcsites = Intersection(p.first->GetLiveSites(), p.second->GetLiveSites());
                    fc_status.Decrement_FC_Counts(fcsites);
                    fcsites = fc_status.Coalesced_Sites();

                    newbranch = stump.CoalesceActive(argedges[edge].GetTime(), p.first, p.second, fcsites);
#if 0
                    if (firstRecFound)
                    {
                        cerr  << endl << "coalescence: " << newbranch->GetID() << endl;
                        cerr << "     side 1: " << endl;
                        p.first->GetRangePtr()->PrintInfo();
                        cerr << "     side 2: " << endl;
                        p.second->GetRangePtr()->PrintInfo();
                        cerr << "     result: " << endl;
                        newbranch->GetRangePtr()->PrintInfo();
                    }
#endif
                    TransferEdgeData(newbranch, argedges[edge]);
                }
                else
                {
                    // error - child missing
                    if (p.first == NULL)
                    {
                        printf("edge: %li coalesence: %li  child 1 not defined\n", edge, argedges[edge].GetTarget());
                    }
                    if (p.second == NULL)
                    {
                        printf("edge: %li coalesence: %li  child 2 not defined\n", edge, argedges[edge].GetTarget());
                    }
                }
            }
            else
            {
                // error - coalescence missing
                printf("edge: %li coalesence: %li not found in target map\n", edge, argedges[edge].GetTarget());
            }
        }

        else if (argedges[edge].GetType() == xmlstr::XML_BRANCHTYPE_DISEASE)
        {
            forcefound = false;
        }

        else if (argedges[edge].GetType() == xmlstr::XML_BRANCHTYPE_EPOCH)
        {
            forcefound = true;

            map<long, targetpair>::iterator pos = targetmap.find(argedges[edge].GetTarget());
            if (pos != targetmap.end())
            {
                targetpair p = pos->second;
                if (p.first != NULL)
                {
                    // make migration
                    long maxEvents = registry.GetForceSummary().GetMaxEvents(force_DIVMIG);
                    long topop = p.first->GetPartition(force_DIVMIG);
                    newbranch = stump.Migrate(argedges[edge].GetTime(), topop, maxEvents, p.first);
                    TransferEdgeData(newbranch, argedges[edge]);
                }
                else
                {
                    // error - child missing
                    printf("edge: %li div-mig: %li child not defined\n", edge, argedges[edge].GetTarget());
                }
            }
            else
            {
                // error - div-mig missing
                printf("edge: %li div-mig: %li not found in target map\n",
                       edge, argedges[edge].GetTarget());
            }
        }

        else if (argedges[edge].GetType() == xmlstr::XML_BRANCHTYPE_MIG)
        {
            forcefound = true;

            map<long, targetpair>::iterator pos = targetmap.find(argedges[edge].GetTarget());
            if (pos != targetmap.end())
            {
                targetpair p = pos->second;
                if (p.first != NULL)
                {
                    // make migration
                    long maxEvents = registry.GetForceSummary().GetMaxEvents(force_MIG);
                    long topop = p.first->GetPartition(force_MIG);
                    newbranch = stump.Migrate(argedges[edge].GetTime(), topop, maxEvents, p.first);
                    TransferEdgeData(newbranch, argedges[edge]);
                }
                else
                {
                    // error - child missing
                    printf("edge: %li migration: %li child not defined\n", edge, argedges[edge].GetTarget());
                }
            }
            else
            {
                // error - migration missing
                printf("edge: %li migration: %li not found in target map\n",
                       edge, argedges[edge].GetTarget());
            }
        }

        else if (argedges[edge].GetType() == xmlstr::XML_BRANCHTYPE_REC)
        {
            if (isRecTree)
            {
                forcefound = true;
                recfound   = true;
                //firstRecFound = true; //JRM debug

                map<long, targetpair>::iterator pos = targetmap.find(argedges[edge].GetTarget());
                if (pos != targetmap.end())
                {
                    targetpair p = pos->second;
                    if (p.first != NULL)
                    {
                        if (argedges[edge+1].GetType() == xmlstr::XML_BRANCHTYPE_REC)
                        {
                            // make sure next edge is associated with same recombination
                            if (argedges[edge].GetTarget() == argedges[edge+1].GetTarget())
                            {
                                // find parents
                                for(size_t paredge=edge+2; paredge<argedges.size(); paredge++)
                                {
                                    if(argedges[paredge].GetTarget() == argedges[edge].GetSource())
                                    {
                                        recpar1 = paredge;
                                    }
                                    if(argedges[paredge].GetTarget() == argedges[edge+1].GetSource())
                                    {
                                        recpar2 = paredge;
                                    }
                                }

                                if ((recpar1 > 0) && (recpar2 > 0))
                                {
                                    // make recombination
                                    long maxEvents = registry.GetForceSummary().GetMaxEvents(force_REC);
                                    FPartMap fparts; // not used, we'll override whatever RecombineActive
                                                     // does with what is in the parent branches

                                    // decide which way the lower sites go from parent A
                                    bool lowleft = true;
                                    string transmittedsites =  argedges[edge].GetTransmittedSites();

                                    // Old format: "MIN:MAX" where both endpoints are INCLUDED.
                                    //unsigned int idx = transmittedsites.find_last_of(':');
                                    //string endsite = transmittedsites.substr(idx+1, transmittedsites.length());
                                    //long endval = atol(endsite.c_str());

                                    // New format: "[MIN,MAX)" where MIN is closed lower endpoint and MAX is open upper endpoint.
                                    unsigned int idx = transmittedsites.find_last_of(',');
                                    string endsite = transmittedsites.substr(idx+1, transmittedsites.length()-1);
                                    long endval = atol(endsite.c_str()) - 1; // correct for open interval

                                    if (endval > argedges[edge].GetRecLoc())
                                    {
                                        lowleft = false;
                                    }

                                    // build the rec branches
                                    rangeset fcsites;
                                    fcsites = fc_status.Coalesced_Sites();
                                    recbranches = recTree->RecombineActive(argedges[edge].GetTime(), maxEvents, fparts,
                                                                           p.first,argedges[edge].GetRecLoc(), fcsites, lowleft);
                                    TransferEdgeData(recbranches.first, argedges[edge]);
                                    TransferEdgeData(recbranches.second, argedges[edge+1]);

                                    //cerr  << endl << "recbranch.first : " << recbranches.first->GetID() << endl;
                                    //recbranches.first->GetRangePtr()->PrintInfo();

                                    //cerr << endl << "recbranch.second : " << recbranches.second->GetID() << endl;
                                    //recbranches.second->GetRangePtr()->PrintInfo();

                                    //recTree->GetTimeList().PrintTreeList();           // JRM debug

                                    // override the generated partitions with the parent partitions from the ARG tree
                                    string par1part = argedges[recpar1].GetPartitions();
                                    if (!par1part.empty())
                                    {
                                        ProcessForce(recbranches.first, par1part);
                                    }

                                    string par2part = argedges[recpar2].GetPartitions();
                                    if (!par2part.empty())
                                    {
                                        ProcessForce(recbranches.second, par2part);
                                    }
                                }
                                else if (recpar1 == 0)
                                {
                                    printf("edge: %li source: %li not defined\n", edge, argedges[edge].GetSource());
                                }
                                else if (recpar2 == 0)
                                {
                                    printf("edge: %li source: %li not defined\n", edge+1, argedges[edge+1].GetSource());
                                }
                            }
                            else
                            {
                                // error - rec numbers don't match
                                printf("edge: %li rec number: %li does not match edge: %li rec number: %li\n",
                                       edge, argedges[edge].GetTarget(), edge+1, argedges[edge+1].GetTarget());
                            }
                        }
                        else
                        {
                            // error - recombinations not paired
                            printf("edge: %li is not a recombination, recombinations are not properly paired\n", edge+1);
                        }
                    }
                    else
                    {
                        // error - child missing
                        printf("edge: %li rec: %li child not defined\n", edge, argedges[edge].GetTarget());
                    }
                }
                else
                {
                    // error - rec missing
                    printf("edge: %li rec: %li not found in target map\n",
                           edge, argedges[edge].GetTarget());
                }
            }
            else
            {
                // error - nor a rec tree
                printf("edge: %li contains a recombination and this is not a recombinant tree\n", edge);
            }
        }

        if (forcefound)
        {
            // add to target map array
            Branch_ptr null_ptr;
            if (recfound)
            {
                // recombination case - returns a pair of branches

                // branch 1
                map<long, targetpair>::iterator pos1 =  targetmap.find(argedges[edge].GetSource());
                if (pos1 != targetmap.end())
                {
                    // second branch in coalescence case
                    targetpair p = pos1->second;
                    p.second = recbranches.first;
                    pos1->second =  p;
                }
                else
                {
                    // everything else
                    targetpair p = make_pair(recbranches.first, null_ptr);
                    targetmap.insert(make_pair(argedges[edge].GetSource(), p));
                }

                // branch 2
                map<long, targetpair>::iterator pos2 =  targetmap.find(argedges[edge+1].GetSource());
                if (pos2 != targetmap.end())
                {
                    // second branch in coalescence case
                    targetpair p = pos2->second;
                    p.second = recbranches.second;
                    pos2->second =  p;
                }
                else
                {
                    // everything else
                    targetpair q = make_pair(recbranches.second, null_ptr);
                    targetmap.insert(make_pair(argedges[edge+1].GetSource(), q));
                }
            }
            else
            {
                map<long, targetpair>::iterator pos =  targetmap.find(argedges[edge].GetSource());
                if (pos != targetmap.end())
                {
                    // second branch in coalescence case
                    targetpair p = pos->second;
                    p.second = newbranch;
                    pos->second =  p;
                }
                else
                {
                    // everything else
                    targetpair p = make_pair(newbranch, null_ptr);
                    targetmap.insert(make_pair(argedges[edge].GetSource(), p));
                }
            }

            if (recfound)
            {
                edge += 1; // move forward one extra because recombination uses two edges up
            }
        }
        else
        {
            // unknown force
            printf("edge: %li specifies unknown force: %s\n", edge, argedges[edge].GetType().c_str());
        }
    }

#ifndef NDEBUG
    recTree->GetTimeList().PrintTreeList();           // JRM debug
    TiXmlDocument * docP = recTree->GetTimeList().AssembleGraphML();
    FILE * argOutput = TiXmlFOpen("checktree.xml","a");
    docP->Print(argOutput,0);
    fclose(argOutput);
#endif
    // add root
    stump.AttachBase(newbranch);

    assert(stump.IsValidTree());

} // ARGTree::ToLamarcTree

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

void ARGTree::TransferEdgeData(Branch_ptr bptr, ARGEdge edge)
{
    // transfer the other data to the new branch

    //  partition forces
    string partitions =  edge.GetPartitions();
    if (partitions.length() >0)
    {
        unsigned int stidx = 0;
        unsigned int idx = partitions.find(',');
        string psubstr;
        while (idx<partitions.length())
        {
            // pull off individual forces
            psubstr = partitions.substr(stidx, idx);
            ProcessForce(bptr, psubstr);
            stidx = idx + 1;
            idx = partitions.find(',', stidx);
        }
        psubstr = partitions.substr(stidx, partitions.length());
        ProcessForce(bptr, psubstr);
    }

} // ARGTree::TransferEdgeData

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

bool ARGTree::ProcessForce(Branch_ptr bptr, string forcestr)
{
    // process force
    unsigned int idx = forcestr.find(':');
    if (idx == forcestr.length())
    {
        printf("ARGTree::ProcessForce: forcestr: %s has no : so it is not a valid definition\n",forcestr.c_str());
        return false;
    }
    else
    {
        string forcekind = forcestr.substr(0, idx);
        string forcevalue = forcestr.substr(idx+1, forcestr.length());

        if (forcekind == uistr::disease)
        {
            string error_msg = "ERROR found in ARGTree::ProcessForce: force: ";
            error_msg += forcekind;
            error_msg +=" not implemented.\n";
            printf("%s", error_msg.c_str());
            return false;
        }
        else if (forcekind == uistr::divmigration)
        {
            long partNum = registry.GetDataPack().GetPartitionNumber(force_DIVMIG, forcevalue);
            bptr->SetPartition(force_DIVMIG, partNum);
        }
        else if(forcekind == uistr::migration)
        {
            long partNum = registry.GetDataPack().GetPartitionNumber(force_MIG, forcevalue);
            bptr->SetPartition(force_MIG, partNum);
        }
        else
        {
            string error_msg = "ERROR found in ARGTree::ProcessForce: force: ";
            error_msg += forcekind;
            error_msg +=" not a known force.\n";
            printf("%s", error_msg.c_str());
            return false;
        }
    }
    return true;
}

//____________________________________________________________________________________
