// $Id: summary.cpp,v 1.55 2018/01/03 21:33:04 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 <cmath>
#include <iostream>

#include "constants.h"
#include "force.h"
#include "intervaldata.h"
#include "range.h"                      // For Link-related typedefs and constants.
#include "region.h"
#include "summary.h"

using namespace std;

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

const DoubleVec1d & Summary::GetShortPoint() const
{
    return m_shortpoint;

} // Summary::ShortPoint

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

const DoubleVec1d & Summary::GetShortWait() const
{
    return m_shortwait;
} // Summary::ShortWait

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

const LongVec2d & Summary::GetShortPicks() const
{
    return m_shortpicks;
} // Summary::ShortPicks

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

Interval const * Summary::GetLongPoint() const
{
    assert(!m_shortness);  // tried to get long-form inappropriately?
    return m_front;
} // Summary::LongPoint

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

const list<Interval> & Summary::GetLongWait() const
{
    assert(!m_shortness);  // tried to get long-form inappropriately?
    return m_intervalData.m_intervals;
}  // Summary::LongWait

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

bool Summary::GetShortness() const
{
    return m_shortness;
}

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

void Summary::AddInterval(double time, const LongVec2d & pk, const LongVec1d & xk,
                          Linkweight recweight, xpart_t ostat, xpart_t nstat,
                          long int recpoint, const LongVec1d & picks, force_type type)
{
    // "recweight" is a Link recombination weight (Biglink weight or number of Littlelinks).
    // "recpoint" is a recombination Littlelink (middle of target Biglink) or FLAGLONG.
    // Add interval, retaining a pointer to it.
    Interval* thisinterval = m_intervalData.AddInterval(m_back, time, pk, xk, recweight, ostat, nstat, recpoint, picks, type);
    // Put that pointer in appropriate place(s).
    if (!m_front) m_front = thisinterval;
    m_back = thisinterval;
    // m_intervalData.PrintIntervalData(); // JRM debug

} // Summary::AddInterval

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

bool Summary::Compress()
{
    ComputeShortWait();
    ComputeShortPoint();
    ComputeShortPicks();
    return m_shortness;
} // Summary::Compress

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

Summary * CoalSummary::Clone(IntervalData& interval) const
{
    // NB:  This takes a reference to the IntervalData to which
    // the new Summary will belong.  It does *not* copy the
    // contents of the old Summary; only the type and m_shortness.

    return new CoalSummary(interval, m_shortness);
} // Clone

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

void CoalSummary::ComputeShortPoint()
{
    // Set up the recipient vector.
    m_shortpoint.assign(m_nparams, 0.0);

    Interval* in;

    for (in = m_front; in != NULL; in = in->m_next)
    {
        ++m_shortpoint[in->m_oldstatus];
    }

} // CoalSummary::ComputeShortPoint

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

void CoalSummary::ComputeShortWait()
{
    // Set up the recipient vector.
    m_shortwait.assign(m_nparams, 0.0);

    list<Interval>::const_iterator interval = m_intervalData.begin();
    list<Interval>::const_iterator end = m_intervalData.end();
    double starttime = 0.0;

    for ( ; interval != end; ++interval)
    {
        double deltaTime = interval->m_endtime - starttime;
        long int param;
        for (param = 0; param < m_nparams; ++param)
        {
            // OPT:  pre-calculate k(k-1) in storage routines....
            double k = interval->m_xpartlines[param];
            m_shortwait[param] += deltaTime * k * (k - 1);
        }
        starttime = interval->m_endtime;
    }

} // CoalSummary::ComputeShortWait

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

void CoalSummary::ComputeShortPicks()
{
    const ForceSummary & fs = registry.GetForceSummary();
    long int numlocalforces = fs.GetNLocalPartitionForces();

    // If there are no local-partition-forces, do nothing.
    if (numlocalforces == 0) return;
    // If there is no recombination, do nothing.
    if (!fs.CheckForce(force_REC)) return;

    m_shortpicks.clear();
    const ForceVec& forces = fs.GetPartitionForces();
    unsigned long int i;
    long int index = 0;
    for (i = 0; i < forces.size(); ++i)
    {
        if (forces[i]->IsLocalPartitionForce())
        {
            const PartitionForce* partforce =
                dynamic_cast<const PartitionForce*>(forces[i]);
            long int nparts = partforce->GetNPartitions();
            LongVec1d counts(nparts, 0L);
            list<Interval>::const_iterator interval = m_intervalData.begin();
            list<Interval>::const_iterator end = m_intervalData.end();
            for ( ; interval != end; ++interval)
            {
                if (!interval->m_partnerpicks.empty())
                {
                    ++counts[interval->m_partnerpicks[index]];
                }
            }
            m_shortpicks.push_back(counts);
            ++index;
        }
    }
} // CoalSummary::ComputeShortPicks

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

void CoalSummary::AdjustSummary(const DoubleVec1d & totals, long int region)
{

    Interval* fakefront = NULL;
    Interval* fakeback = NULL;
    assert(totals.size()==m_shortpoint.size()); // inconsistent input?
    unsigned long int nparams(totals.size());
    bool didadjust = false;

    unsigned long int param;
    for(param = 0; param < nparams; ++param)
    {
        if (totals[param] != 0.0) continue;

        if (!m_shortness)
        {
            // Add fake interval to correct for "fatal attraction"; args specify "nonexistent" interval.
            Interval * newinterval = m_fakeIntervals.AddDummyInterval(fakeback, ZERO, param, FLAGLONG, FLAGLONG, force_COAL);

            if (!fakefront) fakefront = newinterval;
            fakeback = newinterval;

            didadjust = true;
        }

        // adjust short-form summary statistics
        m_shortpoint[param] += 1.0;
    }

    if (didadjust)
    {
        // Hook fake intervals onto front of real intervals.
        fakeback->m_next = m_front;
        m_front = fakefront;
    }

} // CoalSummary::AdjustSummary

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

Summary * MigSummary::Clone(IntervalData& interval) const
{
    // NB:  This takes a reference to the IntervalData to which
    // the new Summary will belong.  It does *not* copy the
    // contents of the old Summary; only the type and m_shortness.

    return new MigSummary(interval, m_shortness);
} // Clone

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

void MigSummary::ComputeShortPoint()
{
    // Set up the recipient vector.
    m_shortpoint.assign(m_npop * m_npop, 0.0);

    Interval* in;

    for (in = m_front; in != NULL; in = in->m_next)
    {
        // Migration rates vary newstatus on rows and oldstatus on columns.
        ++m_shortpoint[m_npop * in->m_newstatus + in->m_oldstatus];
    }

} // MigSummary::ComputeShortPoint

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

void MigSummary::ComputeShortWait()
{
    // Set up the recipient vector.
    m_shortwait.assign(m_npop, 0.0);

    list<Interval>::const_iterator interval = m_intervalData.begin();
    list<Interval>::const_iterator end = m_intervalData.end();
    long int pop;
    double starttime = 0.0;

    for ( ; interval != end; ++interval)
    {
        double deltaTime = interval->m_endtime - starttime;
        for (pop = 0; pop < m_npop; ++pop)
        {
            double k = interval->m_partlines[m_migpartindex][pop];
            m_shortwait[pop] += deltaTime * k;
        }
        starttime = interval->m_endtime;
    }

} // MigSummary::ComputeShortWait

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

void MigSummary::AdjustSummary(const DoubleVec1d & totals, long int region)
{
    assert(static_cast<xpart_t>(totals.size()) == m_npop*m_npop);
    Interval* fakefront = NULL;
    Interval* fakeback = NULL;

    //LS NOTE: It's possible that this will add a migration event for a
    // migration that's been constrained to be zero.  However, the maximizer
    // ignores all such events in this case, so adding one is not fatal.  Though
    // one should be wary.

    bool didadjust = false;

    xpart_t frompop, topop;

    for (frompop = 0; frompop != m_npop; ++frompop)
    {
        for (topop = 0; topop != m_npop; ++topop)
        {

            // no adjustment to diagonal entries
            if (frompop == topop) continue;

            long int param = frompop * m_npop + topop; // index into linearized vector
            if (totals[param] == 0.0)                  // no events of this type
            {
                if (!m_shortness)
                {
                    // Add fake interval to correct for "fatal attraction"; args specify "nonexistent" interval.
                    Interval* newinterval = m_fakeIntervals.AddDummyInterval(fakeback, ZERO, frompop, topop, FLAGLONG, force_MIG);

                    if (!fakefront) fakefront = newinterval;
                    fakeback = newinterval;

                    didadjust = true;
                }

                // adjust short-form summary statistics
                m_shortpoint[param] += 1.0;
            }
        }
    }

    if (didadjust)
    {
        // Hook fake intervals onto front of real intervals.
        fakeback->m_next = m_front;
        m_front = fakefront;
    }
} // MigSummary::AdjustSummary

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

Summary * DiseaseSummary::Clone(IntervalData& interval) const
{
    // NB:  This takes a reference to the IntervalData to which
    // the new Summary will belong.  It does *not* copy the
    // contents of the old Summary; only the type and m_shortness.

    return new DiseaseSummary(interval, m_shortness);
} // Clone

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

void DiseaseSummary::ComputeShortPoint()
{
    // Set up the recipient vector.
    m_shortpoint.assign(m_nstati * m_nstati, 0.0);

    Interval* in;

    for (in = m_front; in != NULL; in = in->m_next)
    {
        // Disease rates vary newstatus on rows and oldstatus on columns.
        ++m_shortpoint[m_nstati * in->m_newstatus + in->m_oldstatus];
    }

} // DiseaseSummary::ComputeShortPoint

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

void DiseaseSummary::ComputeShortWait()
{
    // Set up the recipient vector.
    m_shortwait.assign(m_nstati, 0.0);

    list<Interval>::const_iterator interval = m_intervalData.begin();
    list<Interval>::const_iterator end = m_intervalData.end();
    long int status;
    double starttime = 0.0;

    for ( ; interval != end; ++interval)
    {
        double deltaTime = interval->m_endtime - starttime;
        for (status = 0; status < m_nstati; ++status)
        {
            double k = interval->m_partlines[m_dispartindex][status];
            m_shortwait[status] += deltaTime * k;
        }
        starttime = interval->m_endtime;
    }

} // DiseaseSummary::ComputeShortWait

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

void DiseaseSummary::AdjustSummary(const DoubleVec1d & totals, long int region)
{
    assert(static_cast<xpart_t>(totals.size()) == m_nstati*m_nstati);
    Interval* fakefront = NULL;
    Interval* fakeback = NULL;

    bool didadjust = false;

    xpart_t oldstatus, newstatus;

    for (oldstatus = 0; oldstatus != m_nstati; ++oldstatus)
    {
        for (newstatus = 0; newstatus != m_nstati; ++newstatus)
        {

            // no adjustment to diagonal entries
            if (oldstatus == newstatus) continue;

            // index into linearized vector
            long int param = oldstatus * m_nstati + newstatus;

            if (totals[param] == 0.0)   // no events of this type
            {
                if (!m_shortness)
                {
                    // Add fake interval to correct for "fatal attraction"; args specify "nonexistent" interval.
                    Interval* newinterval = m_fakeIntervals.AddDummyInterval(fakeback, ZERO, oldstatus, newstatus, FLAGLONG, force_DISEASE);

                    if (!fakefront) fakefront = newinterval;
                    fakeback = newinterval;

                    didadjust = true;
                }

                // adjust short-form summary statistics
                m_shortpoint[param] += 1.0;
            }
        }
    }

    if (didadjust)
    {
        // Hook fake intervals onto front of real intervals.
        fakeback->m_next = m_front;
        m_front = fakefront;
    }
} // DiseaseSummary::AdjustSummary

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

Summary * EpochSummary::Clone(IntervalData& interval) const
{
    // NB:  This takes a reference to the IntervalData to which
    // the new Summary will belong.  It does *not* copy the
    // contents of the old Summary; only the type and m_shortness.

    return new EpochSummary(interval, m_shortness);
} // Clone

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

RecSummary::RecSummary(IntervalData& interval, bool shortform)
    :
    Summary(interval,registry.GetForceSummary().GetNParameters(force_REC), shortform),
    m_nrecsbyxpart(registry.GetForceSummary().GetNParameters(force_COAL), 0L),
    m_nrecs(0L)
{
    long int i(0L);
    const ForceVec forces(registry.GetForceSummary().GetPartitionForces());
    ForceVec::const_iterator force;

    (void)i;  // Silence compiler warning about unused variable.

    for(force = forces.begin(); force != forces.end(); ++force)
    {
        if ((*force)->IsLocalPartitionForce())
        {
            LongVec1d fparts((*force)->GetNPartitions(), 0L);
            m_nrecsbydispart = fparts;
        }
    }

} // RecSummary::ctor

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

Summary * RecSummary::Clone(IntervalData& interval) const
{
    // NB:  This takes a reference to the IntervalData to which
    // the new Summary will belong.  It does *not* copy the
    // contents of the old Summary; only the type and m_shortness.

    return new RecSummary(interval, m_shortness);
} // Clone

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

void RecSummary::ComputeShortPoint()
{
    // Set up the recipient vector.
    m_shortpoint.assign(m_nparams, 0.0);

    Interval* in;
    for (in = m_front; in != NULL; in = in->m_next)
    {
        ++m_shortpoint[0];
    }

} // RecSummary::ComputeShortPoint

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

void RecSummary::ComputeShortWait()
{
    // Set up the recipient vector.
    m_shortwait.assign(m_nparams, 0.0);

    list<Interval>::const_iterator interval = m_intervalData.begin();
    list<Interval>::const_iterator end = m_intervalData.end();
    double starttime = 0.0;
    //  long int count = 1;
    for ( ; interval != end; ++interval)
    {
        double deltaTime = interval->m_endtime - starttime;
        m_shortwait[0] += deltaTime * interval->m_recweight;
        starttime = interval->m_endtime;
#if 0
        cerr << "ShortWait "
             << count++
             << ": ts="
             << starttime
             << " te="
             << interval->m_endtime
             << " recweight="
             << interval->m_recweight
             << " r="
             << m_shortwait[0]
             << endl;
#endif
    }

} // RecSummary::ComputeShortWait

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

void RecSummary::AdjustSummary(const DoubleVec1d & totals, long int region)
{
    assert(totals.size() == 1);  // inconsistent forces?!
    if (!registry.GetDataPack().GetRegion(region).RecombinationCanBeEstimated())
    {
        return;
    }
    Interval* fakefront = NULL;
    Interval* fakeback = NULL;
    bool didadjust = false;

    long int param = 0;
    if (totals[param] == 0.0)  // no events of this type
    {
        if (!m_shortness)
        {
            // Add fake interval to correct for "fatal attraction"; args specify "nonexistent" interval.
            Interval* newinterval = m_fakeIntervals.AddDummyInterval(fakeback, ZERO, FLAGLONG, FLAGLONG, 0L, force_REC);

            if (!fakefront) fakefront = newinterval;
            fakeback = newinterval;

            didadjust = true;
        }

        // adjust short-form summary statistics
        m_shortpoint[param] += 1.0;
    }
    if (didadjust)
    {
        // Hook fake intervals onto front of real intervals.
        fakeback->m_next = m_front;
        m_front = fakefront;
    }
} // RecSummary::AdjustSummary

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

void RecSummary::AddToRecombinationCounts(const LongVec1d & membership)
{
    // NB: THIS CODE ASSUMES A DISEASE FORCE EXISTS!
    // m_nrecsbyxpart
    m_nrecsbyxpart[registry.GetDataPack().GetCrossPartitionIndex(membership)]++;

    // m_nrecsbydispart
    // it also assumes that Disease is the first (and only) local partition force
    long int lpforce = registry.GetForceSummary().GetLocalPartitionIndexes()[0];
    m_nrecsbydispart[membership[lpforce]]++;
    m_nrecs++;

} // RecSummary::AddToRecombinationCounts

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

Summary * DivMigSummary::Clone(IntervalData& interval) const
{
    // NB:  This takes a reference to the IntervalData to which
    // the new Summary will belong.  It does *not* copy the
    // contents of the old Summary; only the type and m_shortness.

    return new DivMigSummary(interval, m_shortness);
} // Clone

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

void DivMigSummary::ComputeShortPoint()
{

    // Set up the recipient vector.
    m_shortpoint.assign(m_npop * m_npop, 0.0);

    Interval* in;

    for (in = m_front; in != NULL; in = in->m_next)
    {
        // Migration rates vary newstatus on rows and oldstatus on columns.
        ++m_shortpoint[m_npop * in->m_newstatus + in->m_oldstatus];
    }

} // DivMigSummary::ComputeShortPoint

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

void DivMigSummary::ComputeShortWait()
{
    // Set up the recipient vector.
    m_shortwait.assign(m_npop * m_npop, 0.0);

    list<Interval>::const_iterator interval = m_intervalData.begin();
    list<Interval>::const_iterator end = m_intervalData.end();
    long int pop;
    vector<Epoch>::const_iterator curepoch(m_epochs->begin());
    LongVec1d epochpops(curepoch->PopulationsHere());
    double starttime = 0.0;

    for ( ; interval != end; ++interval)
    {
        double deltaTime = interval->m_endtime - starttime;

        for (pop = 0; pop < m_npop; ++pop)
        {
            double k = interval->m_partlines[m_divmigpartindex][pop];

            // now add k * deltaTime to all populations extant in this epoch
            LongVec1d::iterator epop;
            for(epop = epochpops.begin(); epop != epochpops.end(); ++epop)
            {
                if (pop == (*epop)) continue;
                // Migration rates vary newstatus on rows and oldstatus on columns.
                m_shortwait[pop * m_npop + (*epop)] += deltaTime * k;
            }
        }

        if (interval->m_type == force_DIVERGENCE)
        {
            ++curepoch;
            assert(curepoch != m_epochs->end());
            epochpops = curepoch->PopulationsHere();
        }

        starttime = interval->m_endtime;
    }

} // DivMigSummary::ComputeShortWait

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

void DivMigSummary::AdjustSummary(const DoubleVec1d & totals, long int region)
// NB:  AdjustSummary prevents fatal attraction in likelihood runs.  DivMig is
//  not compatible with likelihood runs, therefore we do nothing here.
{
} // DivMigSummary::AdjustSummary

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

void Summary::AddShortPoint(string s)
{
    double a = static_cast<long int>(atoi(s.c_str()));
    m_shortpoint.push_back(a);
}

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

void Summary::AddShortWait(string s)
{
    double a = static_cast<double>(atof(s.c_str()));
    m_shortwait.push_back(a);
}

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

void Summary::AddShortPicks(const StringVec1d & svec)
{
    LongVec1d newpicks;
    StringVec1d::const_iterator values;

    for(values = svec.begin(); values != svec.end(); ++values)
    {
        long int a = static_cast<long int>(atoi(values->c_str()));
        newpicks.push_back(a);
    }

    m_shortpicks.push_back(newpicks);
}

//____________________________________________________________________________________
