// $Id: branchbuffer.cpp,v 1.28 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: Peter Beerli, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>
#include <iostream>                     // debugging

#include "local_build.h"

#include "branchbuffer.h"
#include "datapack.h"                   // for BranchBuffer ctor
#include "errhandling.h"                // for exceptions
#include "force.h"                      // for force ctor in BranchBuffer::RemoveBranch()
#include "fc_status.h"                  // for BranchBuffer::SetupFC()
#include "range.h"                      // For Link-related typedefs and constants.
#include "registry.h"                   // for access to global datapack for handling
                                        // partition logic in BranchBuffer class

#ifdef DMALLOC_FUNC_CHECK
#include "/usr/local/include/dmalloc.h"
#endif

using namespace std;

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

extern Registry registry;

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

BranchBuffer::BranchBuffer(const DataPack& dpack)
    :
    m_branchxparts(dpack.GetNCrossPartitions(), 0L),
    m_branchparts(dpack.GetNPartitionForces())
{
    long force, nforces = dpack.GetNPartitionForces();
    for(force = 0; force < nforces; ++force)
    {
        LongVec1d parts(dpack.GetNPartitions(force), 0L);
        m_branchparts[force] = parts;
    }
} // BranchBuffer::BranchBuffer

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

void BranchBuffer::Clear()
{
    m_branches.clear();

    m_branchxparts.assign(m_branchxparts.size(), 0L);
    LongVec2d::iterator force;

    for(force = m_branchparts.begin(); force != m_branchparts.end(); ++force)
        force->assign(force->size(), 0L);
} // BranchBuffer::Clear

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

long BranchBuffer::GetNBranches(force_type force, const LongVec1d & membership) const
{
    long count = m_branchxparts[registry.GetDataPack().GetCrossPartitionIndex(membership)];

    return count;
} // BranchBuffer::GetNBranches

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

long BranchBuffer::GetNPartBranches(force_type force, long part) const
{
    long count = m_branchparts[registry.GetForceSummary().GetPartIndex(force)][part];

    return count;
} // BranchBuffer::GetNPartBranches

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

void BranchBuffer::UpdateBranchCounts(const LongVec1d & newpartitions, bool addbranch)
{
    unsigned long xpartindex = registry.GetDataPack().GetCrossPartitionIndex(newpartitions);

    assert(xpartindex < m_branchxparts.size());

    if (addbranch)
    {
        ++m_branchxparts[xpartindex];
        long force, nforces = newpartitions.size();
        for(force = 0; force < nforces; ++force)
            ++m_branchparts[force][newpartitions[force]];
    }
    else
    {
        --m_branchxparts[xpartindex];
        long force, nforces = newpartitions.size();
        for(force = 0; force < nforces; ++force)
            --m_branchparts[force][newpartitions[force]];
    }
} // BranchBuffer::UpdateBranchCounts

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

LongVec1d BranchBuffer::GetBranchParts(force_type force) const
{
    return m_branchparts[registry.GetForceSummary().GetPartIndex(force)];
} // BranchBuffer::GetBranchParts

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

LongVec1d BranchBuffer::GetBranchXParts() const
{
    return m_branchxparts;
} // BranchBuffer::GetBranchXParts

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

void BranchBuffer::Append(Branch_ptr newbranch)
{
    // OPTIMIZE -- should we change to list push_back?
    m_branches.push_back(newbranch);

    UpdateBranchCounts(newbranch->m_partitions);
} // BranchBuffer::Append

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

void BranchBuffer::Append(vector<Branch_ptr> newbranches)
{
    vector<Branch_ptr>::iterator branch;

    for(branch = newbranches.begin(); branch != newbranches.end(); ++branch)
        Append(*branch);
} // BranchBuffer::Append

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

void BranchBuffer::AddAfter(Branchiter here, Branch_ptr newbranch)
{
    // debug DEBUG warning WARNING--assume that we can keep time sorted
    //    ordering by simple append after here.
    ++here;
    m_branches.insert(here, newbranch);

    UpdateBranchCounts(newbranch->m_partitions);
} // BranchBuffer::AddAfter

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

void BranchBuffer::Collate(Branch_ptr newbranch)
{
    // Degenerate collate into an empty list
    if (m_branches.empty())
    {
        m_branches.push_front(newbranch);
        UpdateBranchCounts(newbranch->m_partitions);
        return;
    }

    // Regular collate
    Branchiter br = m_branches.begin();

    // Catch insertion before the first entry
    if ((*br)->Parent(0)->m_eventTime > newbranch->Parent(0)->m_eventTime)
    {
        m_branches.push_front(newbranch);
        UpdateBranchCounts(newbranch->m_partitions);
        return;
    }

    // Insertion after an entry
    Branchiter end = m_branches.end();
    Branchiter previous = br;
    br++;                               // First entry has been taken care of already.

    for ( ; br != end; ++br)
    {
        if ((*br)->Parent(0)->m_eventTime > newbranch->Parent(0)->m_eventTime) break;
        previous = br;
    }

    // We rely on flowthrough for the case where this will be the new bottommost branch.
    AddAfter(previous, newbranch);

} // BranchBuffer::Collate

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

void BranchBuffer::Remove(Branchiter here)
{
    UpdateBranchCounts((*here)->m_partitions, false);
    m_branches.erase(here);

} // BranchBuffer::Remove

void BranchBuffer::Remove(Branch_ptr branch)
{
    UpdateBranchCounts(branch->m_partitions, false);
    m_branches.remove(branch);
}

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

Branch_ptr BranchBuffer::GetFirst()
{
    return m_branches.front();
} // BranchBuffer::GetFirst

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

Branch_ptr BranchBuffer::RemoveFirst()
{
    Branch_ptr pbranch = GetFirst();
    UpdateBranchCounts(pbranch->m_partitions, false);
    m_branches.pop_front();

    return pbranch;
} // BranchBuffer::RemoveFirst

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

Branch_ptr BranchBuffer::RemovePartitionBranch(force_type force, long part, double rnd)
{
    long count = static_cast<long>(rnd * GetNPartBranches(force, part));
    Branchiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
        if ((*brit)->GetPartition(force) == part)
        {
            if (!count) break;
            --count;
        }

    assert(count == 0);                 // Didn't find a branch.
    Branch_ptr retbranch = *brit;
    Remove(brit);

    return retbranch;
} // BranchBuffer::RemovePartitionBranch

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

Branch_ptr BranchBuffer::RemoveBranch(force_type forceid, long xpartindex, double rnd)
{
    const LongVec1d & membership = registry.GetDataPack().GetBranchMembership(xpartindex);
    long count = static_cast<long>(rnd * GetNBranches(forceid, membership));
    const Force & force(**registry.GetForceSummary().GetForceByTag(forceid));
    Branchiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
    {
        if ((*brit)->IsAMember(force, membership))
        {
            if (!count) break;
            --count;
        }
    }

    assert(count == 0);                 // Didn't find a branch.
    assert(brit != m_branches.end());

    Branch_ptr retbranch = *brit;
    Remove(brit);

    return retbranch;
} // BranchBuffer::RemoveBranch

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

Branch_ptr BranchBuffer::RemoveBranch(force_type forceid, const LongVec1d & membership, double rnd)
{
    long count = static_cast<long>(rnd * GetNBranches(forceid, membership));
    const Force & force(**registry.GetForceSummary().GetForceByTag(forceid));
    Branchiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
        if ((*brit)->IsAMember(force, membership))
        {
            if (!count) break;
            --count;
        }

    assert(count == 0); // didn't find a branch

    Branch_ptr retbranch = *brit;
    Remove(brit);

    return retbranch;
} // BranchBuffer::RemoveBranch

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

double BranchBuffer::IntervalBottom()
{
    return GetFirst()->Parent(0)->m_eventTime;
} // BranchBuffer::IntervalBottom()

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

vector<Branch_ptr> BranchBuffer::ExtractConstBranches() const
{
    vector<Branch_ptr> br;
    Branchconstiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
        br.push_back(*brit);

    return br;
} // BranchBuffer::ExtractConstBranches

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

long BranchBuffer::Nsites() const
{
    if (m_branches.empty()) return 0L;

    return (*BeginBranch())->Nsites();
} // BranchBuffer::Nsites

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

void BranchBuffer::IncreaseFCCount(FC_Status & fcstatus) const
{
#if FINAL_COALESCENCE_ON
    Branchconstiter brit;
    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
    {
        fcstatus.Increment_FC_Counts((*brit)->GetLiveSites());
    }
#endif
} // BranchBuffer::IncreaseFCCount

//------------------------------------------------------------------------------------
// Debugging function.

void BranchBuffer::PrintEvents() const
{
    Branchconstiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
        cerr << ToString((*brit)->Event()) << ", ";

    cerr << endl;
} // BranchBuffer::PrintEvents

//------------------------------------------------------------------------------------
// Debugging function.

void BranchBuffer::PrintRanges() const
{
    cerr << "Ranges of branches in this list:" << endl;
    Branchconstiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
    {
        (*brit)->GetRangePtr()->PrintInfo();
    }
} // BranchBuffer::PrintRanges

//------------------------------------------------------------------------------------
// Debugging function.

void BranchBuffer::PrintXParts() const
{
    cerr << "Xpartition values for this buffer:" << endl;
    LongVec1d::size_type i;

    for(i = 0; i < m_branchxparts.size(); i++)
    {
        cerr << " pop: " << i
             << " xpart: " << m_branchxparts[i]
             << endl;
    }
} // BranchBuffer::PrintXParts

//------------------------------------------------------------------------------------
// Debugging function (used only in ASSERT); call for recombinant branches only.
// Returns accumulated weight of Current Target Links over all branches in BranchBuffer object.

Linkweight BranchBuffer::AccumulatedCurTargetLinkweight() const
{
    Linkweight recweight(ZERO);
    Branchconstiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
    {
        recweight += (*brit)->GetRangePtr()->GetCurTargetLinkweight();
    }

    return recweight;
} // BranchBuffer::AccumulatedCurTargetLinkweight

//------------------------------------------------------------------------------------
// Debugging function (used only in ASSERT); call for recombinant branches only.
// Returns accumulated weight of New Target Links over all branches in BranchBuffer object.

Linkweight BranchBuffer::AccumulatedNewTargetLinkweight() const
{
    Linkweight recweight(ZERO);
    Branchconstiter brit;

    for(brit = m_branches.begin(); brit != m_branches.end(); ++brit)
    {
        recweight += (*brit)->GetRangePtr()->GetNewTargetLinkweight();
    }

    return recweight;
} // BranchBuffer::AccumulatedNewTargetLinkweight

//____________________________________________________________________________________
