// $Id: chain.cpp,v 1.100 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 <algorithm>
#include <cassert>
#include <fstream>
#include <iostream>                     // debugging
#include <map>
#include <vector>

#include "local_build.h"

#include "arrangervec.h"
#include "chain.h"
#include "chainparam.h"
#include "constants.h"
#include "defaults.h"
#include "errhandling.h"
#include "forcesummary.h"
#include "likelihood.h"
#include "newick.h"                     // for TREETRACK writing out of trees
#include "random.h"
#include "region.h"                     // for Region::MakeUserTree() and Region::HasUserTree()
#include "registry.h"
#include "runreport.h"
#include "timemanager.h"                // for StartReplicate's use of MakeStickTilTime()
                                        // when dealing with a user defined tree.
#include "tree.h"
#include "treesum.h"
#include "types.h"

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

using namespace std;

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

Chain::Chain (Random & rnd, RunReport & runrep,
              const ChainParameters & chainparam,
              CollectionManager& collmanager, double temper)
    : m_tempident(temper, chainparam.GetAllStringsForActiveArrangers()),
      m_ndiscard(0),
      m_chaintype(0),
      m_nsamples(0),
      m_sampleinterval(0),
      m_realstep(0),
      m_arrangers(chainparam.CloneArrangers()),
      m_randomsource (rnd),
      m_runreport (runrep),
      m_collmanager(collmanager)
{
    // deliberately blank
} // Chain::Chain

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

Chain::Chain (const Chain & src)
    :
    m_tempident(src.m_tempident),
    m_arrangers(src.m_arrangers),
    m_randomsource(src.m_randomsource),
    m_runreport(src.m_runreport),
    m_collmanager(src.m_collmanager)
{
    CopyMembers(src);
} // Chain::Chain

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

Chain::~Chain ()
{
    m_arrangers.ClearAll();
} // Chain::~Chain

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

Chain & Chain::operator=(const Chain & src)
{
    if (&src != this)                   // self assignment test
    {
        m_arrangers.ClearAll();

        CopyMembers(src);
    }

    return *this;

} // Chain::operator=

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

// Does not set the reference members: Chain::m_randomsource,
// Chain::m_runreport and Chain::m_collmanager.

void Chain::CopyMembers(const Chain & src)
{
    // temperature identity variables
    m_tempident = src.m_tempident;

    // non-temperature-identity variables
    m_nsamples = src.m_nsamples;
    m_ndiscard = src.m_ndiscard;
    m_chaintype = src.m_chaintype;
    m_sampleinterval = src.m_sampleinterval;
    m_realstep = src.m_realstep;
    m_chainstate = src.m_chainstate;

    m_arrangers = src.m_arrangers;

} // Chain::CopyMembers

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

Arranger * Chain::ChooseArranger (double rnd)
{
    return m_arrangers.GetRandomArranger(rnd);
} // Chain::ChooseArranger

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

void Chain::SetChainType (long chtype, const ChainParameters & chparam)
{
    m_chaintype = chtype;
    m_nsamples = chparam.GetNSamples(chtype);
    m_ndiscard = chparam.GetNDiscard(chtype);
    m_sampleinterval = chparam.GetInterval(chtype);

} // Chain::SetChainType

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

double Chain::GetCurrentDataLlike()
{
    return m_chainstate.GetTree()->GetDLValue();
} // GetCurrentDataLike

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

void Chain::StartRegion (Tree * regiontree)
{
#ifdef DENOVO
    m_runreport.ReportNormal("Running denovo case--no rearrangement!");
#endif // DENOVO
    m_chainstate.SetTree(regiontree);

    if (regiontree->NoPhaseUnknownSites())
    {
        m_arrangers.ZeroHapArranger();
    }
    if (registry.GetForceSummary().CheckForce(force_REC))
    {
        if (dynamic_cast<RecTree*>(regiontree)->GetNumMovingLoci() == 0)
        {
            m_arrangers.ZeroLocusArranger();
        }
        if (dynamic_cast<RecTree*>(regiontree)->AnyRelativeHaplotypes() == false)
        {
            m_arrangers.ZeroProbHapArranger();
        }
    }
} // Chain::StartRegion

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

void Chain::EndRegion()
{
    // In case we turned haplotyping off in a previous region, turn it back on.
    m_arrangers.RestoreTiming();

} // Chain::EndRegion

//------------------------------------------------------------------------------------
// StartRegion() must be called before StartReplicate().

void Chain::StartReplicate (const ForceSummary & forcesum, Region& region)
{
    bool tree_is_bad = false;
    long bad_denovo = 0;
    Arranger* denovoarranger = m_arrangers.GetDenovoArranger();
    ForceParameters fp(forcesum.GetStartParameters(), region.GetID());
    m_chainstate.SetParameters(fp);

    if (region.HasUserTree())
    {
        // JDEBUG--this creates a tree that ignores the stick,
        // if the user hasn't given us branchlengths, I don't know
        // what this will do to the rest of the program
        region.MakeUserTree(m_chainstate.GetTree());
        if (forcesum.UsingStick())
            m_chainstate.GetTree()->GetTimeManager()->
                MakeStickTilTime(fp, m_chainstate.GetTree()->RootTime());
    }
    else
    {
        // no user tree is provided, so we make a de novo tree
        denovoarranger->SetParameters (m_chainstate);
        DestroyStick(forcesum);
        do {
            tree_is_bad = false;
            try
            {
                denovoarranger->Rearrange(m_chainstate);
                m_chainstate.SimulateDataIfNeeded();
                if (region.AnySimulatedLoci())
                {
                    string msg = "Number of recombinations in this tree:  "
                        + ToString(m_chainstate.GetTree()->NumberOfRecombinations());
                    registry.GetRunReport().ReportNormal(msg);
                }
#ifndef STATIONARIES
                m_chainstate.GetTree()->CalculateDataLikes();
                //LS NOTE:  It's possible to reject a tree during dl calculation.
#endif // STATIONARIES
            }
            catch (rejecttree_error& e)
            {
                tree_is_bad = true;
                ++bad_denovo;
                // LS TEST:  First-tree rejection messages.
#if 1
                switch (e.GetType())
                {
                    case OVERRUN:
                        cerr << "error = OVERRUN" << endl;
                        break;
                    case TINYPOP:
                        cerr << "error = TINYPOP" << endl;
                        break;
                    case ZERODL:
                        cerr << "error = ZERODL" << endl;
                        break;
                    case STRETCHED:
                        cerr << "error = STRETCHED" << endl;
                        break;
                    case EPOCHBOUNDARYCROSS:
                        cerr << "error = EPOCHBOUNDS" << endl;
                        break;
                }
#endif

                if (bad_denovo > defaults::tooManyDenovo)
                {
                    string errorMessage = "Unable to create initial tree.  Starting parameter ";
                    errorMessage += "\n values may be too extreme; try using more conservative ones.";
                    denovo_failure e(errorMessage);
                    throw e;
                }
                // reject the unsatisfactory tree and prepare to try again
                denovoarranger->AcceptAndSynchronize(m_chainstate, m_tempident.GetTemperature(), tree_is_bad);
            }
        } while(tree_is_bad);
    }

    assert(m_chainstate.GetTree()->IsValidTree());

    // accept the tree and make the oldtree identical to it
    denovoarranger->AcceptAndSynchronize(m_chainstate, m_tempident.GetTemperature(), tree_is_bad);
    assert(m_chainstate.GetOldTree()->IsValidTree());

} // Chain::StartReplicate

//------------------------------------------------------------------------------------
// SetChainType() must be called before StartChain()

void Chain::StartChain (long chainnumber, long chaintype, const ForceSummary & forces, const ForceParameters & starts)
{
    if (IsCold())
    {
        long nsteps = m_nsamples * m_sampleinterval + m_ndiscard;
        m_runreport.SetBarParameters(nsteps, m_ndiscard, chainnumber, chaintype);
        m_runreport.PrintBar(0L);       // initialize printed bar
    }

    m_tempident.StartChain();  // set up temperature identity tracking info
    m_chainstate.AllChanged();

    m_realstep = 0;

    m_chainstate.SetParameters(starts);
    m_chainstate.OverwriteOldParameters();
    m_chainstate.UpdateOldStickParams();
    m_chainstate.UpdateNewStickParams();

    arrangerit arr;
    for (arr = m_arrangers.begin (); arr != m_arrangers.end (); ++arr)
        (*arr).second->SetParameters(m_chainstate);

#ifdef TREETRACK
    if (m_tempident.IsCold())
    {
        ofstream of;
        of.open("treedls", ios::app);
        of << "chain" << chainnumber << "={" << endl;
        of.close();
    }
#endif

} // Chain::StartChain

//------------------------------------------------------------------------------------
// SetChainType() must be called before StartBayesianChain()
// This is a version of StartChain which does not change the
// parameters.

void Chain::StartBayesianChain(long chainnumber, long chaintype, const ForceSummary& forces)
{
    if (IsCold())
    {
        long nsteps = m_nsamples * m_sampleinterval + m_ndiscard;
        m_runreport.SetBarParameters(nsteps, m_ndiscard, chainnumber, chaintype);
        m_runreport.PrintBar(0L);       // initialize printed bar
    }
    m_tempident.StartChain();  // set up temperature identity tracking info
    m_chainstate.AllChanged();

    m_realstep = 0;

    m_chainstate.OverwriteOldParameters();
    m_chainstate.UpdateOldStickParams();
    m_chainstate.UpdateNewStickParams();

    // The following is still needed in case this is the first use of
    // this Chain object
    arrangerit arr;
    for (arr = m_arrangers.begin (); arr != m_arrangers.end (); ++arr)
        (*arr).second->SetParameters(m_chainstate);

#ifdef TREETRACK
    if (m_tempident.IsCold())
    {
        ofstream of;
        of.open("treedls", ios::app);
        of << "chain" << chainnumber << "={" << endl;
        of.close();
    }
#endif
} // Chain::StartBayesianChain

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

void Chain::DoOneChain (long nsteps, bool lastchain)
{
    long step;
    bool tree_is_bad;

#ifdef LAMARC_QA_SINGLE_DENOVOS
    long denovo_tree_reject_count = 0;
#endif // LAMARC_QA_SINGLE_DENOVOS

    for (step = 0; step < nsteps; ++step)
    {
        // The eventloop call used to be here when we were supporting
        // Mac OS 9 eventloop();

        ++m_realstep;  // running total of all steps

        // Make sure all Arrangers are up to date.
        m_collmanager.UpdateArrangers(m_arrangers, m_chainstate);

        // If we're doing purely DENOVO arrangement, don't let
        // ourselves choose any other arranger.

#ifndef DENOVO
#ifndef LAMARC_QA_SINGLE_DENOVOS
        Arranger * arranger = ChooseArranger(m_randomsource.Float());
#else
        Arranger * arranger = m_arrangers.GetDenovoArranger();
#endif // LAMARC_QA_SINGLE_DENOVOS
#else
        Arranger * arranger = m_arrangers.GetDenovoArranger();
#endif // DENOVO

        assert(m_chainstate.GetTree()->IsValidTree());
        assert(m_chainstate.GetOldTree()->IsValidTree());

#ifdef LAMARC_QA_SINGLE_DENOVOS
        tree_is_bad = true;
        while(tree_is_bad)
        {
#endif // LAMARC_QA_SINGLE_DENOVOS

            tree_is_bad = false;
            try
            {
                arranger->Rearrange(m_chainstate);
#ifndef STATIONARIES
                arranger->ScoreRearrangement(m_chainstate);
#endif // STATIONARIES
                // In the Stationaries code, the tree retains its initial likelihood.
                // We must nevertheless call Accept because of the Hastings ratio issue.
                // Similiarly, we call Accept with a "bad" tree to clean it up correctly.
            }
            catch (rejecttree_error& e)
            {
                tree_is_bad = true;
#ifdef LAMARC_QA_SINGLE_DENOVOS
                denovo_tree_reject_count++;
#endif // LAMARC_QA_SINGLE_DENOVOS
                switch (e.GetType())
                {
                    case OVERRUN:
                        m_tempident.NoteBadTree();
                        break;
                    case TINYPOP:
                        m_tempident.NoteTinyPopTree();
                        break;
                    case ZERODL:
                        m_tempident.NoteZeroDLTree();
                        break;
                    case STRETCHED:
                        m_tempident.NoteStretchedTree();
                        break;
                    case EPOCHBOUNDARYCROSS:
                        // We don't bookkeep the number of such events at this time.
                        break;
                }
            }
#ifdef LAMARC_QA_SINGLE_DENOVOS
        }
#endif // LAMARC_QA_SINGLE_DENOVOS

        bool accepted = arranger->AcceptAndSynchronize(m_chainstate, m_tempident.GetTemperature(), tree_is_bad);

// MREMOVE special-case divergence stationaries code
        if (m_tempident.IsCold())
        {
           ofstream of;
           of.open("divstat", ios::app);
           TimeList& tl = m_chainstate.GetTree()->GetTimeList();
           Branchconstiter br;
           for (br = tl.FirstBody(); br != tl.EndBranch(); br = tl.NextBody(br)) {
              if ((*br)->Event() == btypeCoal) {
                of << "coal" << endl;
                break;
              }
              if ((*br)->Event() == btypeEpoch) {
                of << "epoch" << endl;
                break;
              }
           }
           of.close();
        }

// end MREMOVE

#ifdef TREETRACK
        if (m_tempident.IsCold())
        {
            ofstream of;
            of.open("treedls", ios::app);
            //   of << m_chainstate.GetTree()->GetDLValue() << "," << endl;
            of << m_chainstate.GetTree()->GetDLValue() << ",";
            of << "step=" << step << ",";

#if 0
            TreeSummary* trsum = m_chainstate.GetTree()->SummarizeTree();

            double pscore = registry.GetSinglePostLike().Calc_lnProbGP
                (m_chainstate.GetParameters().GetRegionalParameters(),
                 m_chainstate.GetParameters().GetRegionalLogParameters(),
                 trsum);

            delete trsum;
            of << pscore << endl;
#endif

            NewickConverter nc;
            of << nc.LamarcTreeToNewickString(*(m_chainstate.GetTree())) << endl;

            of.close();

            of.open("mstrees", ios::app);
            of << "//" << endl;
            of << nc.LamarcTreeToMSNewickString(*(m_chainstate.GetTree()));
            of << endl;
            of.close();
        }
#endif // TREETRACK

#ifndef LAMARC_QA_SINGLE_DENOVOS
        // this fouls up LAMARC_QA_SINGLE_DENOVOS, not sure why

        // We don't count acceptances during burn-in.
        if (m_realstep > m_ndiscard)
        {
            string arrangername = arranger->GetName();
            m_tempident.ScoreRearrangement(arrangername, accepted);
        }
#endif  // ! LAMARC_QA_SINGLE_DENOVOS

#if LIKETRACK
        ofstream of;
        of.open("likes1", ios::app);
        of << newtree << endl;
        of.close();
#endif

        if ((m_realstep > m_ndiscard) &&
            ((m_realstep - m_ndiscard) % m_sampleinterval == 0))
        {
            m_tempident.Sample(m_collmanager, m_chainstate, m_chaintype, lastchain);   // EWFIX.CHAINTYPE
        }

        if (IsCold())
        {
            m_runreport.PrintBar(m_realstep);
        }
    } // end of "for" loop over steps

#ifdef LAMARC_QA_SINGLE_DENOVOS
    registry.AddDenovoTreeRejectCount(denovo_tree_reject_count);
#endif // LAMARC_QA_SINGLE_DENOVOS

} // Chain::DoOneChain

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

ChainOut Chain::EndChain()
{
    ChainOut& chout = m_tempident.EndChain();
    chout.SetLlikedata(m_chainstate.GetTree()->GetDLValue());
    chout.SetEndtime();

#ifdef TREETRACK
    if (m_tempident.IsCold())
    {
        ofstream of;
        of.open("treedls", ios::app);
        of << "}" << endl;
        of.close();
    }
#endif

    return chout;

} // Chain::EndChain

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

/* Swap trees between two different-temperature Chains, if
   appropriate.  The equation is:

   hotlike ^ (1/coldtemp) * coldlike ^ (1/hottemp)
   -----------------------------------------------
   hotlike ^ (1/hottemp) * coldlike ^ (1/coldtemp)

   but we use log units for flow-protection reasons.
   This is a Chain function in order to use private parts of
   Chain; it is actually symmetrical in its two Chain arguments.

   We actually swap the temperatures (and associated variables)
   between the Chains, not the Trees, as it's easier to
   parallelize this way.

   NB:  Since we DO NOT swap the Tree, we also DO NOT swap anything
   that should stay with the Tree, such as (in a Bayesian run) the
   ForceParameters.  Watch out for this!  Only things that go with
   the temperature, which should all be in the TempIdent package,
   are swapped.

   NB:  The equation above uses likelihoods; the code below
   uses log likelihoods.
*/

bool Chain::SwapTemperatureIdentities(Chain& hot)
{
    bool bayesian = registry.GetChainParameters().IsBayesian();

    Chain& cold = *this;
    double hottemp = hot.m_tempident.GetTemperature();
    double coldtemp = cold.m_tempident.GetTemperature();
    cold.m_tempident.SwapTried();

    // these are log likelihoods
    double hotlike = hot.m_chainstate.GetTree()->GetDLValue();
    double coldlike = cold.m_chainstate.GetTree()->GetDLValue();

    double test;

    if (bayesian)
    {
        SinglePostLike& postlike = registry.GetSinglePostLike();

        ForceParameters& hotparams = hot.m_chainstate.GetParameters();
        ForceParameters& coldparams = cold.m_chainstate.GetParameters();
        vector<double> hotlogs = hotparams.GetRegionalLogParameters();
        vector<double> coldlogs = coldparams.GetRegionalLogParameters();

        // WARNING program is not exception safe while these pointers
        // are held!
        TreeSummary* hotsum = hot.m_chainstate.GetTree()->SummarizeTree();
        TreeSummary* coldsum = cold.m_chainstate.GetTree()->SummarizeTree();

        double hotprior = postlike.Calc_lnProbGP(hotparams.GetRegionalParameters(),
                                                 hotlogs, hotsum);
        double coldprior = postlike.Calc_lnProbGP(coldparams.GetRegionalParameters(),
                                                  coldlogs, coldsum);

        delete hotsum;
        delete coldsum;

        test = ((hotlike + hotprior) / coldtemp) + ((coldlike + coldprior) / hottemp)
            - (((hotlike + hotprior) / hottemp) + ((coldlike + coldprior) / coldtemp));

    }
    else
    {   // non-bayesian
        test = (hotlike/coldtemp) + (coldlike/hottemp) -
            ((hotlike/hottemp) + (coldlike/coldtemp));
    }

    if (test >= 0.0)
    {
        // always accept improvements
        cold.SwapTemperatures(hot);
        // DO NOT DO THIS: if (bayesian) cold.SwapParameters(hot);
        return true;
    }

    // conditionally accept non-improvements
    double comparison = log(m_randomsource.Float());
    if (test < comparison)
    {
        return false;
    }
    else
    {
        cold.SwapTemperatures(hot);
        // DO NOT DO THIS:  if (bayesian) cold.SwapParameters(hot);
        return true;
    }
} // SwapTemperatureIdentities

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

void Chain::SwapTemperatures(Chain& hot)
{
    // We swap only the "temperature identity" variables.
    // NB:  this code assumes that it's been called on the colder
    // of the two arguments!
    m_tempident.NoteSwap(hot.m_tempident, hot.m_chainstate, m_chainstate);
    swap(m_tempident, hot.m_tempident);
} // SwapTemperatures

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

void Chain::SwapParameters(Chain& hot)
{
    // swaps the working parameters of two chains; needed in
    // Bayesian runs only.
    m_chainstate.SwapParameters(hot.m_chainstate);
} // SwapParameters

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

void Chain::DestroyStick(const ForceSummary & forces)
{
    if (forces.UsingStick())
    {
        m_chainstate.GetTree()->DestroyStick();
        if (m_chainstate.GetOldTree())
            m_chainstate.GetOldTree()->DestroyStick();
    }
} // DestroyStick

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

void Chain::RecalculateDataLikes()
{
    m_chainstate.GetTree()->GetTimeList().SetAllUpdateDLs();
    m_chainstate.GetTree()->CalculateDataLikes();
    m_chainstate.OverwriteOldTree();
}

//____________________________________________________________________________________
