// $Id: dlcalc.cpp,v 1.115 2018/01/03 21:32:57 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>
#include <fstream>

#include "local_build.h"

#include "dlcalc.h"
#include "datapack.h"
#include "locus.h"
#include "dlmodel.h"
#include "tree.h"
#include "timelist.h"
#include "region.h"
#include "registry.h"
#include "random.h"

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

using namespace std;

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

DLCalculator::DLCalculator(const Locus& locus)
    : m_datamodel(*(locus.GetDataModel())),
      m_markerpos(locus.GetMarkerLocations())
{
    // deliberately blank
} // DLCalculator::DLCalculator

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

DLCalculator::DLCalculator(const DLCalculator& src)
    : m_datamodel(src.m_datamodel),
      m_markerpos(src.m_markerpos)
{
    // intentionally blank
} // DLCalculator copy constructor

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

ChildInfo
DLCalculator::GetChildInfo(Branch_ptr branch, long locus, long childindex,
                           long cellindex, long posn, bool moving) const
{
    Branch_ptr child = branch->GetValidChild(branch->Child(childindex), posn);
    double length = branch->HowFarTo(*child);
    return ChildInfo(child->GetDLCell(locus, cellindex, moving), length);
} // GetChildInfo

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

ChildInfo DLCalculator::NullChildInfo() const
{
    Cell_ptr nullcell(new NullCell);
    double length = FLAGDOUBLE;
    return ChildInfo(nullcell, length);
} // NullChildInfo

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

rangepair DLCalculator::SitePairToMarkerPair(rangepair sites)
{
    long firstmrkr = FLAGLONG;
    long lastmrkr = FLAGLONG;
    long firstsite = sites.first;
    long lastsite = sites.second;
    long nmarkers = m_markerpos.size();

    assert (firstsite < lastsite);

    --lastsite;    // calling code assumes half-open intervals

    // if we are at the extreme right or left, outside the range
    // of markers, return immediately

    if (lastsite < m_markerpos.front() || firstsite > m_markerpos.back())
        return(rangepair(FLAGLONG, FLAGLONG));

    // find first marker that is within site range, inclusive
    long marker;
    for (marker = 0; marker < nmarkers; ++marker)
    {
        if (m_markerpos[marker] > lastsite) // we've passed the end
        {
            return(rangepair(FLAGLONG,FLAGLONG));
        }
        if (m_markerpos[marker] >= firstsite) // found it
        {
            firstmrkr = marker;
            break;
        }
    }

    // Find first marker past the end of site range (which may be the
    // non-existent marker after all markers; this is a half-open
    // interval).
    for (marker = nmarkers - 1; marker >= 0; --marker)
    {
        if (m_markerpos[marker] <= lastsite)
        {
            lastmrkr = marker + 1;                   // found it; we return the
            // next marker for half-open
            break;
        }
    }

    assert(firstmrkr != FLAGLONG && lastmrkr != FLAGLONG);
    assert(lastmrkr > firstmrkr);

    return(rangepair(firstmrkr, lastmrkr));

} // SitePairToMarkerPair

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

void DLCalculator::SimulateData(Tree& tree, Locus& locus)
{
    locus.ClearVariability();
    rangepair span = locus.GetSiteSpan();

    // Need vector instead of single rate for autocorrelation.
    DoubleVec1d rates = m_datamodel.ChooseRandomRates(span.second - span.first);

    LongVec1d markers = locus.GetMarkerLocations();
    for (unsigned long marker=0; marker<markers.size(); marker++)
    {
        long site = markers[marker];
        double rate = rates[site - span.first];

        // reverse iterate through the tree simulating data upwards
        TimeList& timelist = tree.GetTimeList();
        for (Branchconstriter brrit = timelist.RBeginBranch(); brrit != timelist.REndBranch(); brrit++)
        {
            if ((*brrit)->Event() == btypeCoal ||
                (*brrit)->Event() == btypeTip)
            {
                // try to find applicable parent
                Branch_ptr pParent = (*brrit)->GetValidParent(site);

                DoubleVec1d state;
                if (pParent == Branch::NONBRANCH)
                {
                    // if no parent found:
                    state = m_datamodel.ChooseAncestralState(marker);
                }
                else
                {
                    double length = pParent->HowFarTo(**brrit);
                    state = m_datamodel.SimulateMarker(rate * length,
                                                       marker,
                                                       pParent->GetDLCell(locus.GetIndex(), markerCell, locus.IsMoving())->GetStateFor(marker, 0));
                }
                //Find where to save the information and then save it.
                Cell_ptr childdl = (*brrit)->GetDLCell(locus.GetIndex(), markerCell, locus.IsMoving());
                childdl->SetAllCategoriesTo(state, marker);
                if ((*brrit)->Event() == btypeTip)
                {
                    //Keep track of the variability of the simulated data.
                    locus.SaveState(state, marker, boost::dynamic_pointer_cast<TBranch>(*brrit)->m_label);
                    //LS TEST
                    //cerr << childdl->DLsToString(marker, marker) << endl;
                }
            }
        }
    }
} // SimulateData

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

void DLCalculator::CopyDataFrom(Locus& destloc, Locus& origloc, Tree& tree)
{
    destloc.ClearVariability();
    LongVec1d markers = destloc.GetMarkerLocations();

    for (unsigned long marker=0; marker<markers.size(); marker++)
    {
        long site = markers[marker];
        assert(origloc.SiteInLocus(site));

        TimeList& timelist = tree.GetTimeList();
        DoubleVec2d origstates;
        DoubleVec2d deststates;
        long n_destalleles = (*timelist.BeginBranch())->GetDLCell(destloc.GetIndex(), markerCell, destloc.IsMoving())->GetStateFor(marker, 0).size();
        DoubleVec1d zeroes(n_destalleles, 0.0);
        for (Branchconstiter brit = timelist.FirstTip(); brit != timelist.EndBranch(); brit = timelist.NextTip(brit))
        {
            origstates.push_back((*brit)->GetDLCell(origloc.GetIndex(), markerCell, origloc.IsMoving())->GetStateFor(origloc.SiteToMarker(site), 0));
            deststates.push_back(zeroes); //Just to be sure.
        }

        //Now.  The trick is that the two loci may have different numbers
        // of possible alleles.  So if the destination locus has an equal number
        // or more alleles than the original, we'll simply overwrite the first X
        // alleles with the new alleles, leaving any extra alleles at zero.
        //
        //When there are 'missing' alleles in the destination that are not
        // in the original, we must punt, and assign multiple original alleles
        // to the same destination allele.  We want to be able to preserve any
        // diversity, however:  it's no good if the only two alleles the original
        // visited are lumped together in the destination.  So.  We first check
        // to see what the most frequent allele is in the original, and assign
        // that to the first allele of the desination.  Then we assign the next
        // most frequent allele to the second allele of the destination.  We keep
        // doing this until we run out of destination alleles.  Then we add
        // any remaining original alleles to the final destination allele.
        //
        //Whew!

        //So, OK.  first:  Sum the originals to figure out who's the greatest.
        DoubleVec1d origsums(origstates[0].size(), 0.0);
        for (unsigned long i=0; i<origstates.size(); i++)
        {
            transform(origstates[i].begin(),
                      origstates[i].end(),
                      origsums.begin(),
                      origsums.begin(),
                      plus<double>());
        }

        //Now make a multimap so we can loop over it.
        multimap<double, unsigned long> mindexes;
        for (unsigned long i=0; i<origsums.size(); i++)
        {
            mindexes.insert(make_pair(origsums[i], i));
        }
        long destindex = 0;
        for (multimap<double, unsigned long>::reverse_iterator mindex=mindexes.rbegin();
             mindex != mindexes.rend(); mindex++)
        {
            long origindex = mindex->second;

            //We now have the index of the original allele and of the destination
            // allele, so we're good to go.
            if (mindex->first != 0)
            {
                for (unsigned long tip=0; tip<origstates.size(); tip++)
                {
                    deststates[tip][destindex] += origstates[tip][origindex];
                }
            }

            if (destindex < n_destalleles - 1)
            {
                destindex++;
            }
        }

        //Now copy everything back into the dlcells of the destination
        unsigned long tip=0;
        for (Branchconstiter brit = timelist.FirstTip(); brit != timelist.EndBranch(); brit = timelist.NextTip(brit), tip++)
        {
            Cell_ptr childdl = (*brit)->GetDLCell(destloc.GetIndex(), markerCell, destloc.IsMoving());
            //Copy the information
            childdl->SetAllCategoriesTo(deststates[tip], marker);
            destloc.SaveState(deststates[tip], marker, boost::dynamic_pointer_cast<TBranch>(*brit)->m_label);
        }
    }
} // DLCalculator::CopyDataFrom

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

void DLCalculator::Randomize(Locus& destloc, rangeset rset, Tree& tree)
{
    cerr << "Deleted range: " << ToString(rset) << endl;

    // Take all sites in rset and equilibrate their dlcells.
    LongVec1d markers = destloc.GetMarkerLocations();

    for (unsigned long marker=0; marker<markers.size(); marker++)
    {
        long site = markers[marker];
        if (rset == AddPairToRange(make_pair(site, site+1), rset))
        {
            TimeList& timelist = tree.GetTimeList();
            long n_destalleles = (*timelist.BeginBranch())->GetDLCell(destloc.GetIndex(), markerCell, destloc.IsMoving())->GetStateFor(marker, 0).size();
            DoubleVec1d ones(n_destalleles, 1.0);
            for (Branchconstiter brit = timelist.FirstTip(); brit != timelist.EndBranch(); brit = timelist.NextTip(brit))
            {
                Cell_ptr childdl = (*brit)->GetDLCell(destloc.GetIndex(), markerCell, destloc.IsMoving());
                //Copy the information
                childdl->SetAllCategoriesTo(ones, marker);
            }
        }
    }
} //DLCalculator::Randomize

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

void DLCalculator::MarkPanelBranches(Tree& tree, const Locus&) //locus unused.
{
    // Walk down the tree and mark each branch.
    // The tips are marked at create time as 0 = panel, 1 = sample.
    // Anything but a coalescence gets marked the same as its child.
    // A coalescence counts all the children that have value > 0.
    // Thus if both branches coming into a coalescence ultimate point to
    // samples or coalescences with value > 0, the coalescence has a score of 2.
    TimeList& timelist = tree.GetTimeList();
    Branchiter brit;
    Branch_ptr branch = Branch::NONBRANCH;
    int node = 0;
    for (brit = timelist.FirstBody(); brit != timelist.EndBranch();
         brit = timelist.NextBody(brit))
    {
        // propagate down tree from tips - which were marked at creation time
        branch = *brit;
        node++;
        // MDEBUG wasCoalCalced appears to be dead code--it is NEVER set true anywhere anymore
        branch->m_wasCoalCalced = false;
        branch->m_isSample = 0;
        // currently there are always 2 children, but this allows for more
        for (long nc=0; nc<branch->NChildren(); nc++)
        {
            if (branch->Child(nc)!= Branch::NONBRANCH)
            {
                if (branch->Child(nc)->Event()== btypeCoal)
                {
                    if (branch->Child(nc)->m_isSample > 1)
                    {
                        branch->m_isSample++;
                    }
                }
                else
                {
                    if (branch->Child(nc)->m_isSample > 0)
                    {
                        branch->m_isSample++;
                    }
                }
            }
        }
    }
} // DLCalculator::MarkPanelBranches

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

NucCalculator::NucCalculator(const Locus& locus)
    : DLCalculator(locus)
{
    // deliberately blank
} // NucCalculator ctor

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

NucCalculator::NucCalculator(const NucCalculator& src)
    : DLCalculator(src)
{
    // deliberately blank
} // NucCalculator copy constructor

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

LongVec1d NucCalculator::CalculateAliasesFromData(const vector<DoubleVec2d>& data) const
{
    // compute aliases, starting at position 1 because 0 can never be aliased.
    // If an alias partner is not found, the marker's alias will remain at the
    // -1 (no alias) with which it was initialized
    long nmarkers = data.size();
    LongVec1d alias(nmarkers, -1L);
#ifndef LAMARC_QA_NOALIAS
    long partner, marker;
    for (marker = 1; marker < nmarkers; ++marker)
    {
        for (partner = marker - 1; partner >= 0; --partner)
        {
            if (data[marker] == data[partner])
            {
                // found an alias
                alias[marker] = partner;
                break;
            }
        }
    }
#endif

    return alias;
} // CalculateAliasesFromData

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

LongVec1d NucCalculator::RecalculateAliases(const Tree& tree, const Locus& locus) const
{
    const TimeList& timelist = tree.GetTimeList();
    long cat = 0; // we alias based on first category

    long ntips = timelist.GetNTips();
    long nmarkers = locus.GetNmarkers();
    DoubleVec1d basepattern(4,0.0);
    DoubleVec2d tippattern(ntips, basepattern);
    vector<DoubleVec2d> markerpattern(nmarkers, tippattern);

    Branchconstiter brit = timelist.FirstTip();
    long tip, marker, base;
    for (tip = 0; brit != timelist.EndBranch(); brit = timelist.NextTip(brit), ++tip)
    {
        Cell_ptr dlcell = (*brit)->GetDLCell(locus.GetIndex(), markerCell, locus.IsMoving());
        //Note:  We never worry about aliasing for moving loci.
        for (marker = 0; marker < nmarkers; ++marker)
        {
            SiteDL siteDLs = dlcell->GetSiteDLs(marker);
            for (base = 0; base < 4; ++base)
            {
                markerpattern[marker][tip][base] = siteDLs[cat][base];
            }
        }
    }

    assert(tip == ntips);
    return CalculateAliasesFromData(markerpattern);

} // RecalculateAliases

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

LongVec1d NucCalculator::SetupAliases(const Locus& locus) const
{
    long nmarkers = locus.GetNmarkers();

    long cat = 0; // we alias based on the first category

    long ntips = locus.GetNTips();
    const vector<LocusCell>& tipcells = locus.GetTipCells();
    DoubleVec1d basepattern(4,0.0);
    DoubleVec2d tippattern(ntips, basepattern);
    vector<DoubleVec2d> markerpattern(nmarkers, tippattern);

    // dump the data into a usable form
    long marker, tip, base;
    for (tip = 0; tip < ntips; ++tip)
    {
        Cell_ptr dlcell = tipcells[tip][0]; // the primary Cell
        for (marker = 0; marker < nmarkers; ++marker)
        {
            SiteDL siteDLs = dlcell->GetSiteDLs(marker);
            for (base = 0; base < 4; ++base)
            {
                markerpattern[marker][tip][base] = siteDLs[cat][base];
            }
        }
    }

    return CalculateAliasesFromData(markerpattern);

} // SetupAliases

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

void NucCalculator::CalculateSite(Cell_ptr child1, Cell_ptr child2, Cell_ptr newcell, long pos, long alias)
{
    SiteDL newsiteDLs;

    SiteDL child1DLs = child1->GetSiteDLs(pos);
    SiteDL child2DLs = child2->GetSiteDLs(pos);

    // if the alias is invalid, calculate; else use the alias
    if (alias == -1)
    {
        newsiteDLs = dynamic_cast<NucModel&>(m_datamodel).
            ComputeSiteDLs(child1DLs,child2DLs);
        if (m_datamodel.ShouldNormalize())
        {
            double newnorm = newcell->Normalize(newsiteDLs) +
                child1->GetNorms(pos) + child2->GetNorms(pos);
            newcell->SetNorms(newnorm, pos);
        }
    }
    else
    {
        // cerr << " " << pos;
        newsiteDLs = newcell->GetSiteDLs(alias);
        newcell->SetNorms(newcell->GetNorms(alias),pos);
    }

    newcell->SetSiteDLs(pos,newsiteDLs);

} // NucCalculator::CalculateSite

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

void NucCalculator::Breakalias(LongVec1d& aliases, const rangevector& subtrees)
{
    long tr;
    long nsubtrees(subtrees.size());

    // skip the first subtree because it can't break any aliases;
    // for all other subtrees, break all aliases that cross the
    // subtree boundary
    for (tr = 1; tr < nsubtrees; ++tr)
    {
        rangepair marker = SitePairToMarkerPair(subtrees[tr]);
        long pos;
        for (pos = marker.first; pos < marker.second; ++pos)
        {
            assert(static_cast<unsigned long>(pos) < aliases.size());
            if (aliases[pos] < marker.first) aliases[pos] = -1;
        }
    }

} // NucCalculator::Breakalias

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

DLCalculator* DNACalculator::Clone() const
{
    DNACalculator* pnewcalc = new DNACalculator(*this);

    return(pnewcalc);

} // DNACalculator::Clone

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

double DNACalculator::Calculate(Tree& tree, const Locus& locus, bool moving)
{
    assert(tree.IsValidTree());

    Branch_ptr branch = Branch::NONBRANCH;
    Cell_ptr dlcell;
    double totallike = 0.0;
    long nsubtrees, tr;
    NucModel& datmodel = dynamic_cast<NucModel&>(m_datamodel);
    ChildInfo child0, child1;
    long curmarker;
    long loc = locus.GetIndex();
    LongVec1d aliases = tree.GetAliases(locus.GetIndex());

    // Find the subtrees
    rangepair span = locus.GetSiteSpan();
    rangevector subtrees = tree.GetLocusSubtrees(span);
    nsubtrees = subtrees.size();

    // Update the aliases.
    Breakalias(aliases, subtrees);

    // Initialize the model categories storage
    datmodel.ResetCatCells();

    // Step through the subtrees to compute the data likelhoods
    TimeList& timelist = tree.GetTimeList();
    for (tr = 0; tr < nsubtrees; tr++)
    {
        rangepair marker = SitePairToMarkerPair(subtrees[tr]);
        // bail out if there are no markers in this subtree
        if (marker.first == FLAGLONG) continue;

        long firstsite = subtrees[tr].first;

        Branchiter brit;
        Branch_ptr last_active_coal_branch = Branch::NONBRANCH;
        for (brit = timelist.FirstCoal(); brit != timelist.EndBranch();
             brit = timelist.NextCoal(brit))
        {
            branch = *brit;
            if (branch->CanCalcDL(firstsite))
            {
                // we want the last such branch saved after fall-through, since
                // it will be the correct place to calculate the likelihood
                last_active_coal_branch = branch;

                if (branch->ShouldCalcDL(firstsite))
                {
                    // Find "real" children and appropriate branch lengths.
                    child0 = GetChildInfo(branch, loc, 0, markerCell, firstsite, moving);
                    child1 = GetChildInfo(branch, loc, 1, markerCell, firstsite, moving);

                    // Precalculate exponential terms
                    datmodel.RescaleLengths(child0.m_length, child1.m_length);

                    // Calculate the data likelihood at each position.
                    dlcell = branch->GetDLCell(loc, markerCell, moving);
                    for(curmarker = marker.first; curmarker < marker.second;++curmarker)
                    {
                        CalculateSite(child0.m_cell, child1.m_cell, dlcell, curmarker,aliases[curmarker]);
                    }
                }
            }
        }

        assert(last_active_coal_branch != Branch::NONBRANCH);
        // We want to calculate likelihood at last_active_coal_branch,
        // so we set the dlcell to point there--whether it does
        // already or not.  This makes it legal to compute the
        // data likelihood of a tree with no flags set, which is
        // useless but hard to prevent and thus needs to work.
        dlcell = last_active_coal_branch->GetDLCell(loc, markerCell, moving);

        totallike += datmodel.ComputeSubtreeDL(*dlcell,
                                               dlcell->GetSiteDLs(marker.first),
                                               dlcell->GetSiteDLs(marker.second),
                                               marker.first);
    }

    return totallike;
} // DNACalculator::Calculate

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

SNPCalculator::SNPCalculator(const Locus& locus)
    : NucCalculator(locus),
      m_invarmodel(dynamic_cast<NucModel*>(m_datamodel.Clone()))
{
    assert(m_invarmodel);
    m_invarmodel->SetNmarkers(INVARIANTS+1); // the invariant model should have
    // 1 "marker" for each base-type plus one extra for the conjoint number

} // SNPCalculator::SNPCalculator

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

SNPCalculator::SNPCalculator(const SNPCalculator& src)
    : NucCalculator(src),
      m_invarmodel(dynamic_cast<NucModel*>(src.m_invarmodel->Clone()))
{
    assert(m_invarmodel);
} // SNPCalculator::copy ctor

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

SNPCalculator::~SNPCalculator()
{
    delete m_invarmodel;
} // SNPCalculator::dtor

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

DLCalculator* SNPCalculator::Clone() const
{
    SNPCalculator* pnewcalc = new SNPCalculator(*this);
    return(pnewcalc);
} // SNPCalculator::Clone

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

void SNPCalculator::CalculateInvarSite(Cell_ptr child1, Cell_ptr child2, Cell_ptr newcell, long pos)
{
    SiteDL child1DLs = child1->GetSiteDLs(pos);
    SiteDL child2DLs = child2->GetSiteDLs(pos);

    // no alias for invariant sites

    SiteDL newsiteDLs = m_invarmodel->ComputeSiteDLs(child1DLs,child2DLs);

    if (m_invarmodel->ShouldNormalize())
    {
        double norm1 = child1->GetNorms(pos);
        double norm2 = child2->GetNorms(pos);
        double newnorm = newcell->Normalize(newsiteDLs) + norm1 + norm2;
        newcell->SetNorms(newnorm,pos);
    }

    newcell->SetSiteDLs(pos,newsiteDLs);

} // SNPCalculator::CalculateInvarSite

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

// EWFIX -- add comment and/or change name of "long i" to be more clear
void SNPCalculator::MaskPanelTips(TimeList& timeList, long loc, long first, long last, long i)
{
    // determine the mask
    DoubleVec1d mask(INVARIANTS,0.0);
    if (i<INVARIANTS)
    {
        mask = m_invarmodel->DataToLikes(SINGLEBASES[i]);

    }
    else
    {
        // unknown so all equally probable
        mask[0] = 1.0;
        mask[1] = 1.0;
        mask[2] = 1.0;
        mask[3] = 1.0;
    }

    // apply the mask to all panel tips
    Branchconstiter brit;
    Branch_ptr branch;
    for (brit = timeList.FirstTip(); brit != timeList.EndBranch(); brit = timeList.NextTip(brit))
    {
        branch = *brit;
        if (branch->m_isSample == 0)
        {
            // found a panel tip, set its mask
            Cell_ptr dlcell = branch->GetDLCell(loc, markerCell, false);
            long curmarker;
            for (curmarker = first; curmarker < last; ++curmarker)
            {
                dlcell->SetAllCategoriesTo(mask, curmarker);
            }
        }
    }

}  // MaskPanelTips

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

// A horrible truth to remember about this code:  the storage used
// for invariant site calculations is *re-used*.  You can never assume
// that that information is still around; the next subtree will wipe
// it out.  This is why we recalculate invariant sites throughout the
// whole tree every time, even though, in theory, this is not necessary.
// There is a possible but very difficult OPTIMIZATION opportunity here.
// Mary June 11 2002
//
// Preconditions (non-panel):
//    Sample tips have correct markercells and invarcells
//    Tree is correctly marked with updateDL flags
// Postconditions (non-panel):
//    Sample tips have correct markercells and invarcells
//    markercells (ONLY!) contain correct likelihood values in all
//       branches where the site was live and not FC
//    INVARCELLS DO NOT CONTAIN VALID LIKELIHOODS EXCEPT AT TIPS
//    Value returned is correct data likelihood
//    UpdateDL status has not changed
//
// Preconditions (panel):
//    Sample tips have correct markercells and invarcells
//    Panel tips have correct invarcells BUT NOT MARKERCELLS
// Postconditions (panel):
//    Sample tips have correct markercells and invarcells
//    Panel tips have correct invarcells BUT NOT MARKERCELLS
//    NO CELLS CONTAIN VALID LIKELIHOODS EXCEPT AT TIPS
//    Value returned is correct data likelihood
//    All branches are marked as needing updating

double SNPCalculator::Calculate(Tree& tree, const Locus& locus, bool moving)
{
    if (moving == true)  // This is not an appropriate Calculator for a floating locus.
    {
        assert(false);
        throw implementation_error ("Floating loci may not use the SNP model.");
    }

    // Determine subtrees
    double totallike = 0.0;
    rangepair span = locus.GetSiteSpan();
    rangevector subtrees = tree.GetLocusSubtrees(span);
    long nsubtrees = subtrees.size();
    long loc = locus.GetIndex();

    // Update the aliases.
    LongVec1d aliases = tree.GetAliases(locus.GetIndex());
    Breakalias(aliases, subtrees);

    // Initialize the model categories storage
    // we reference m_datamodel with a local variable since we need and know it to
    // be NucModel as this is SNP-only code. m_invarmodel doesn't need it 'cause
    // it is ours alone
    NucModel& varmodel = dynamic_cast<NucModel&>(m_datamodel);
    varmodel.ResetCatCells();
    m_invarmodel->ResetCatCells();

    // check if there are panels
    if(!tree.GetSnpPanelFlag())
    {
        // no-panel case
        for (long ntree = 0; ntree < nsubtrees; ++ntree)
        {
            // loop over the subtrees
            totallike += CalcNoPanel(tree, loc, subtrees[ntree], aliases);
        }
    }
    else
    {
        // panel case
        long curmarker;
        TimeList& timeList = tree.GetTimeList();
        Branch_ptr root = timeList.Root();
        Cell_ptr rootcell_m = root->GetDLCell(loc, markerCell, false);
        Cell_ptr rootcell_i = root->GetDLCell(loc, invariantCell, false);

        // Temporary DLCells for scratch use
        Cell_ptr basecell(rootcell_m->Clone());  // sample markers only
        Cell_ptr sumcell(rootcell_m->Clone());   // sum of sample+panel markers
        Cell_ptr invarcell(rootcell_i->Clone()); // panel invariant sites only

        // mark the coalescences  // MDEBUG dead code?
        MarkPanelBranches(tree, locus);

        // Regrettably, we have to recalculate every node every time, so the updateDL
        // information is set accordingly here.
        // Calculating a node does NOT unset its updateDL flag (for subtree reasons) so
        // once these are set they stay set until data likelihood calculation is complete.
        //
        // We should cache these for benefit of further loci -- MDEBUG
        timeList.SetAllUpdateDLs();

        // loop over the subtrees
        for (long ntree = 0; ntree < nsubtrees; ++ntree)
        {
            // clear working storage
            sumcell->EmptyCell();

            rangepair marker = SitePairToMarkerPair(subtrees[ntree]);
            DoubleVec2d markerstates;

            // work on the variable sites, if they exist
            if (marker.first != FLAGLONG)
            {
                // sum up the single base panel corrections
                for (int i=0; i<INVARIANTS; i++)
                {
                    MaskPanelTips(timeList, loc, marker.first, marker.second, i);
                    Branch_ptr topBranch = CalcPanel(tree, loc, subtrees[ntree], aliases);
                    sumcell->AddTo(topBranch->GetDLCell(loc, markerCell, false));
                }

                // subtract the summed single base panel corrections from the unknown panel values
                MaskPanelTips(timeList, loc, marker.first, marker.second, INVARIANTS);
                Branch_ptr topBranch = CalcPanel(tree, loc, subtrees[ntree], aliases);
                basecell->CopyInitialize(*(topBranch->GetDLCell(loc, markerCell, false)));
                basecell->SubtractFrom(sumcell);
            }

            // now calculate unobserved sites invariant correction,
            // (masks set when tips created in timelist.cpp)
            Branch_ptr topBranch = CalcPanelInvariants(tree, loc, subtrees[ntree]);
            invarcell->CopyInitialize(*(topBranch->GetDLCell(loc, invariantCell, false)));

            totallike += SumPanel(basecell, subtrees[ntree], invarcell);
        }
    }

    return totallike;

} // SNPCalculator::Calculate

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

double SNPCalculator::CalcNoPanel(Tree& tree, long loc, pair<long,long> sitespan, LongVec1d aliases)
{
    Branch_ptr branch = Branch::NONBRANCH;
    Branch_ptr child1 = Branch::NONBRANCH;
    Branch_ptr child2 = Branch::NONBRANCH;

    // set up some necessary values
    NucModel& varmodel = dynamic_cast<NucModel&>(m_datamodel);
    TimeList& timelist = tree.GetTimeList();

    // no markers in this subtree is handled by the rangepair,
    // marker, having the same value in both first and second.
    // Currently that value is FLAGLONG.
    rangepair marker = SitePairToMarkerPair(sitespan);
    long firstsite = sitespan.first;

    bool snpsPresent = (marker.first != FLAGLONG);

    Cell_ptr dlcell;
    Cell_ptr dlcell_invar;
    Cell_ptr dlcell1;
    Cell_ptr dlcell2;
    double length1 = 0.0;
    double length2 = 0.0;

    // Data likelihood for ONE subtree

    Branchiter brit;
    for (brit = timelist.FirstCoal(); brit != timelist.EndBranch();
         brit = timelist.NextCoal(brit))
    {
        branch = *brit;
        if (branch->CanCalcDL(firstsite))
        {
            // obtain the children of this coalescence and their branch lengths
            child1 = branch->GetValidChild(branch->Child(0),firstsite);
            length1 = branch->HowFarTo(*child1);
            child2 = branch->GetValidChild(branch->Child(1),firstsite);
            length2 = branch->HowFarTo(*child2);

            // note the marker DLCell of this coalescence--needed later even if not updated now
            dlcell  = branch->GetDLCell(loc, markerCell, false);

            // Calculate the data likelihood at each position for variable sites,
            // if appropriate; we skip this if the branch is marked "doesn't need updating"
            if (branch->ShouldCalcDL(firstsite))
            {
                dlcell1 = child1->GetDLCell(loc, markerCell, false);
                dlcell2 = child2->GetDLCell(loc, markerCell, false);
                varmodel.RescaleLengths(length1, length2);
                long curmarker;
                for (curmarker = marker.first; curmarker < marker.second; ++curmarker)
                {
                    CalculateSite(dlcell1,dlcell2,dlcell,curmarker,aliases[curmarker]);
                }
            }

            // Calculate the invariant data likelihood for this branch,
            // using baseA for allAs, baseC for allCs, etc.  Must be done even if "doesn't need updating"
            dlcell_invar  = branch->GetDLCell(loc, invariantCell, false);
            dlcell1 = child1->GetDLCell(loc, invariantCell, false);
            dlcell2 = child2->GetDLCell(loc, invariantCell, false);
            m_invarmodel->RescaleLengths(length1, length2);
            long curbase;
            for (curbase = baseA; curbase <= baseT; ++curbase)
            {
                CalculateInvarSite(dlcell1,dlcell2,dlcell_invar,curbase);
            }
        }
    }
    // Code above sets dlcell and dlcell_invar only in a two-legged coalescence.  Therefore,
    // the last time it sets them, it is at the subtree root.  The last setting of dlcell
    // and dlcell_invar thus contains the subtree root values and will be used for summing.

    return SumNoPanel(dlcell, dlcell_invar, sitespan);
} //SNPCalculator::CalcNoPanel

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

double SNPCalculator::SumNoPanel(Cell_ptr varcell, Cell_ptr invarcell, pair<long,long> sitespan)
{
    double totallike = 0.0;
    NucModel& varmodel = dynamic_cast<NucModel&>(m_datamodel);

    rangepair marker = SitePairToMarkerPair(sitespan);
    // if the subtree contains no markers skip computing marker datalike
    if (marker.first != FLAGLONG)
    {
        totallike += varmodel.ComputeSubtreeDL(*varcell,
                                               varcell->GetSiteDLs(marker.first),
                                               varcell->GetSiteDLs(marker.second),
                                               marker.first);
    }

    // Calculate the invariant marker likelihood for the subtree,
    // using baseA for allAs, baseC for allCs, etc.
    // number invariant sites = total sites - markers
    long ninvarsites = sitespan.second - sitespan.first;
    if (marker.first != FLAGLONG)
        ninvarsites -= marker.second - marker.first;

    if (ninvarsites > 0)
    {
        // We are going to use a trick here, and add together the four
        // invariant site likelihoods into the next (fifth) bin.
        // This is ugly, but necessary to keep
        // SNP-specific code out of the DataModel class.  The DataModel
        // needs to treat the four invariant markers as one big
        // supermarker, and cannot determine this itself; so we help it
        // out.

        invarcell->SumMarkers(baseA, baseEnd, m_invarmodel->ShouldNormalize());

        //We multiply the log likelihood by the number of invariant markers
        // because this is equivalent to raising the likelihood to the power
        // of the number of invariants.
        totallike += ninvarsites * m_invarmodel->
            ComputeSubtreeDL(*invarcell,
                             invarcell->GetSiteDLs(baseEnd),
                             invarcell->GetSiteDLs(baseEnd+1),    // one past the end
                             baseEnd);
    }

    return totallike;
} // SNPCalculator::SumNoPanel

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

Branch_ptr SNPCalculator::CalcPanel(Tree& tree, long loc, pair<long,long> sitespan, LongVec1d aliases)
{
    Branch_ptr branch = Branch::NONBRANCH;
    Branch_ptr child1 = Branch::NONBRANCH;
    Branch_ptr child2 = Branch::NONBRANCH;

    // set up some necessary values
    NucModel& varmodel = dynamic_cast<NucModel&>(m_datamodel);
    TimeList& timelist = tree.GetTimeList();

    // no markers in this subtree is handled by the rangepair,
    // marker, having the same value in both first and second.
    // Currently that value is FLAGLONG.
    rangepair marker = SitePairToMarkerPair(sitespan);
    long firstsite = sitespan.first;

    Cell_ptr dlcell;
    Cell_ptr dlcell1;
    Cell_ptr dlcell2;
    double length1 = 0.0;
    double length2 = 0.0;
    long curmarker;

    // Compute marker likelihood for one subtree
    Branch_ptr last_calc_branch = Branch::NONBRANCH;
    Branchiter brit;
    for (brit = timelist.FirstCoal(); brit != timelist.EndBranch();
         brit = timelist.NextCoal(brit))
    {
        branch = *brit;
        if (branch->CanCalcDL(firstsite))
        {
            // store which node we did; last one is subtree root and is returned
            last_calc_branch = branch;
            dlcell  = branch->GetDLCell(loc, markerCell, false);

            // obtain the children of this coalescence and their branch lengths
            child1 = branch->GetValidChild(branch->Child(0),firstsite);
            length1 = branch->HowFarTo(*child1);
            dlcell1 = child1->GetDLCell(loc, markerCell, false);
            child2 = branch->GetValidChild(branch->Child(1),firstsite);
            length2 = branch->HowFarTo(*child2);
            dlcell2 = child2->GetDLCell(loc, markerCell, false);

            // Precalculate the exponential values based on branch length.
            varmodel.RescaleLengths(length1, length2);
            for (curmarker = marker.first; curmarker < marker.second; ++curmarker)
            {
                CalculateSite(dlcell1,dlcell2,dlcell,curmarker,aliases[curmarker]);
            }
        }
    }

    return last_calc_branch;
} // SNPCalculator::CalcPanel
//------------------------------------------------------------------------------------

Branch_ptr SNPCalculator::CalcPanelInvariants(Tree& tree, long loc, pair<long,long> sitespan)
{
    Branch_ptr branch = Branch::NONBRANCH;
    Branch_ptr child1 = Branch::NONBRANCH;
    Branch_ptr child2 = Branch::NONBRANCH;

    // set up some necessary values
    TimeList& timelist = tree.GetTimeList();

    // no markers in this subtree is handled by the rangepair,
    // marker, having the same value in both first and second.
    // Currently that value is FLAGLONG.
    rangepair marker = SitePairToMarkerPair(sitespan);
    long firstsite = sitespan.first;

    Cell_ptr dlcell;
    Cell_ptr dlcell1;
    Cell_ptr dlcell2;
    double length1 = 0.0;
    double length2 = 0.0;
    long curbase;

    // Compute marker likelihood for one subtree
    Branch_ptr last_calc_branch = Branch::NONBRANCH;
    Branchiter brit;
    for (brit = timelist.FirstCoal(); brit != timelist.EndBranch();
         brit = timelist.NextCoal(brit))
    {
        branch = *brit;
        if (branch->CanCalcDL(firstsite))
        {
            // this is how we remember the final coalescence point
            // for this subtree
            last_calc_branch = branch;

            // obtain the children of this coalescence and their branch lengths
            child1 = branch->GetValidChild(branch->Child(0),firstsite);
            length1 = branch->HowFarTo(*child1);
            child2 = branch->GetValidChild(branch->Child(1),firstsite);
            length2 = branch->HowFarTo(*child2);
            dlcell1 = child1->GetDLCell(loc, invariantCell, false);
            dlcell2 = child2->GetDLCell(loc, invariantCell, false);
            dlcell  = branch->GetDLCell(loc, invariantCell, false);
            m_invarmodel->RescaleLengths(length1, length2);
            for (curbase = 0; curbase < INVARIANTS; ++curbase)
            {
                CalculateInvarSite(dlcell1,dlcell2,dlcell,curbase);
            }
        }
    }
    return last_calc_branch;
}

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

double SNPCalculator::SumPanel(Cell_ptr varcell, pair<long,long> sitespan, Cell_ptr invarcell)
{
    double totallike = 0.0;
    NucModel& varmodel = dynamic_cast<NucModel&>(m_datamodel);
    rangepair marker = SitePairToMarkerPair(sitespan);

    // if the subtree contains no markers skip computing marker datalike
    if (marker.first != FLAGLONG)
    {
        // calculate variant panel likelihood for this subtree
        totallike = varmodel.ComputeSubtreeDL(*varcell,
                                              varcell->GetSiteDLs(marker.first),
                                              varcell->GetSiteDLs(marker.second),
                                              marker.first);


    }

    // Calculate the invariant marker likelihood for the subtree,
    // using baseA for allAs, baseC for allCs, etc.
    long ninvarmarkers = sitespan.second - sitespan.first;
    if (marker.first != FLAGLONG)
        ninvarmarkers -= marker.second - marker.first;

    // don't bother if there aren't any
    if (ninvarmarkers > 0)
    {
        double invarlike = 0;
        // We are going to use a trick here, and add together the four
        // invariant site likelihoods into the bin normally occupied
        // by the first one.  This is ugly, but necessary to keep
        // SNP-specific code out of the DataModel class.  The DataModel
        // needs to treat the four invariant markers as one big
        // supermarker, and cannot determine this itself; so we help it
        // out.

        // Changed to sum to "last" bin

        invarcell->SumMarkers(baseA, baseEnd, m_invarmodel->ShouldNormalize());

        // We multiply the log likelihood by the number of invariant markers
        // because this is equivalent to raising the likelihood to the power
        // of the number of invariants.
        invarlike = m_invarmodel->ComputeSubtreeDL(*invarcell,
                                                   invarcell->GetSiteDLs(baseEnd),
                                                   invarcell->GetSiteDLs(baseEnd+1),    // one past the end
                                                   baseEnd);

        invarlike *= ninvarmarkers;
        totallike += invarlike;
    }

    return totallike;
} // SNPCalculator::SumPanel

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

// A horrible truth to remember about this code:  the storage used
// for invariant site calculations is *re-used*.  You can never assume
// that that information is still around; the next subtree will wipe
// it out.  This is why we recalculate invariant sites throughout the
// whole tree every time, even though, in theory, this is not necessary.
// There is a possible but very difficult OPTIMIZATION opportunity here.
// Mary June 11 2002

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

void SNPCalculator::SimulateData(Tree&, Locus&)
{
    assert(false);
    //The default data simulator probably doesn't work with SNPs.

} // SimulateData

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

AlleleCalculator::AlleleCalculator(const Locus& locus)
    : DLCalculator(locus)
{
    // deliberately blank
} // AlleleCalculator::AlleleCalculator

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

AlleleCalculator::AlleleCalculator(const AlleleCalculator& src)
    : DLCalculator(src)
{
    // deliberately blank
} // AlleleCalculator::copy ctor, for internal use only

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

// Similar to DNA case except:
// no aliasing (identical microsats are too rare)
// uses a unique approach to non-marker sites (different from SNPs)

double AlleleCalculator::Calculate(Tree& tree, const Locus& locus, bool moving)
{
    Branch_ptr branch = Branch::NONBRANCH;
    Cell_ptr dlcell;
    double totallike = 0.0;
    long posn1, posn2, firstsite, nsubtrees, tr, pos;
    AlleleModel& datmodel = dynamic_cast<AlleleModel&>(m_datamodel);
    DoubleVec1d scaled0, scaled1;
    ChildInfo child0, child1;
    if (moving)
    {
        //We need to recalculate the marker positions
        m_markerpos = locus.GetMarkerLocations();
        //LS DEBUG MAPPING:  we might need a similar thing for DNACalc:Calc,
        // but only if/when we allow mapping something with a DNA model (we
        // currently are restricted to K-Allele models).
    }

    long loc = locus.GetIndex();

    // Initialize the model categories storage
    datmodel.ResetCatCells();

    rangepair span = locus.GetSiteSpan();
    rangevector subtrees = tree.GetLocusSubtrees(span);
    nsubtrees = subtrees.size();

    // Step through the subtrees to compute the data likelhoods
    TimeList& timelist = tree.GetTimeList();
    for (tr = 0; tr < nsubtrees; tr++)
    {
        rangepair markers = SitePairToMarkerPair(subtrees[tr]);
        if (markers.first == FLAGLONG) continue;  // no markers in this subtree
        posn1 = markers.first;
        posn2 = markers.second;
        firstsite = subtrees[tr].first;

        Branchiter brit;
        Branch_ptr last_calc_branch = Branch::NONBRANCH;
        for (brit = timelist.FirstCoal(); brit != timelist.EndBranch();
             brit = timelist.NextCoal(brit))
        {
            branch = *brit;
            if (branch->CanCalcDL(firstsite))
            {
                last_calc_branch = branch;

                if (branch->ShouldCalcDL(firstsite))
                {
                    dlcell = branch->GetDLCell(loc, markerCell, moving);
                    // Find "real" children and appropriate branch lengths.
                    child0 = GetChildInfo(branch, loc, 0, markerCell, firstsite, moving);
                    child1 = GetChildInfo(branch, loc, 1, markerCell, firstsite, moving);

                    // Precalculate branchlength terms
                    scaled0 = datmodel.RescaleLength(child0.m_length);
                    scaled1 = datmodel.RescaleLength(child1.m_length);

                    // Calculate the data likelihood at each position.
                    for(pos = posn1; pos < posn2; ++pos)
                    {
                        datmodel.ComputeSiteDLs(child0.m_cell, child1.m_cell, dlcell, scaled0, scaled1, pos);
                    }
                }
            }
        }

        // we want to calculate likelihood at last_calc_branch so set dlcell accordingly
        dlcell = last_calc_branch->GetDLCell(loc, markerCell, moving);

        totallike += datmodel.ComputeSubtreeDLs(*dlcell, dlcell->GetSiteDLs(posn1),
                                                dlcell->GetSiteDLs(posn2), posn1);
    }

    // oddly enough, the microsat likelihood can be 1, and sometimes is, so
    // this assert is inappropriate.
    // assert(totallike != 0);  // this would be a likelihood=1, which is *too* likely

    return totallike;

} // AlleleCalculator::Calculate

//____________________________________________________________________________________
