// $Id: collmanager.cpp,v 1.36 2018/01/03 21:33:03 mkkuhner Exp $

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

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

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

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

*/

/* Authors: Peter Beerli, Mary Kuhner, Jon Yamato and Joseph Felsenstein */


#include <cassert>
#include <iostream>

#include "local_build.h"

#include "collmanager.h"
#include "registry.h"
#include "stringx.h"
#include "newick.h"
#include "parameter.h"

#include "tinyxml.h"

using namespace std;

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

CollectionManager::CollectionManager(long numreg, long numrep)
    : m_region(0),
      m_replicate(0)
{
    //This used to be the EstablishSamplingStrategy function, but it works
    // semantically better here --LS

    // This function implements a very simple rule for which aspects
    // of the chain to sample.  It will need to get more complex
    // eventually.
    ChainParameters& chainparam = registry.GetChainParameters();

    bool bayesian = chainparam.IsBayesian();

    if (bayesian)
    {
        m_sample_params = true;
#ifdef STATIONARIES
        m_sample_trees = true;
#else // STATIONARIES
        m_sample_trees = false;
#endif // STATIONARIES
    }
    else
    {
        m_sample_params = false;
        m_sample_trees = true;
    }

    m_sample_maps = false;    //We'll set this to true in StartChain if needed.

    // set up internal storage
    long reg, rep;
    vector<TreeCollector*> temptree;
    vector<ParamCollector*> tempparam;
    vector<MapCollector*> tempmap;
    for (reg = 0; reg < numreg; ++reg)
    {
        for (rep = 0; rep < numrep; ++rep)
        {
            temptree.push_back(new TreeCollector);
            tempparam.push_back(new ParamCollector);
            tempmap.push_back(new MapCollector(reg));
            //Note:  These usually get destroyed by StartChain(), but not in the
            // case of reading in from a summary file (ChainManager::ReadInSumFile)
        }
        m_treecoll.push_back(temptree);
        m_paramcoll.push_back(tempparam);
        m_mapcoll.push_back(tempmap);
        temptree.clear();
        tempparam.clear();
        tempmap.clear();
    }

} // CollectionManager ctor

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

CollectionManager::~CollectionManager()
{
    // Clean up the Collectors
    unsigned long reg, rep;
    for (reg = 0; reg < m_treecoll.size(); ++reg)
    {
        for (rep = 0; rep < m_treecoll[reg].size(); ++rep)
        {
            delete m_treecoll[reg][rep];
            delete m_paramcoll[reg][rep];
            delete m_mapcoll[reg][rep];
        }
    }
} // CollectionManager dtor

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

void CollectionManager::Collect(ChainState& chstate, long initialOrFinal, bool lastchain)
{
    // stick sampling is automatically done as part of tree sampling
    if (m_sample_trees) m_treecoll[m_region][m_replicate]->Score(chstate);
    if (m_sample_params) m_paramcoll[m_region][m_replicate]->Score(chstate);
    if (m_sample_maps) m_mapcoll[m_region][m_replicate]->Score(chstate);
    chstate.ParametersSampled();             // Might as well make it consistent.
    WriteTraceFile(chstate, initialOrFinal); // EWFIX.CHAINTYPE
    WriteNewickTreeFile(chstate);
    WriteReclocsFile(chstate, lastchain);
#ifdef LAMARC_QA_TREE_DUMP
    WriteArgFile(chstate,lastchain);
#endif
} // Collect

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

void CollectionManager::WriteTraceFile(ChainState& chstate, long initialOrFinal)
{
    //Note:  The column headers for this file are written in
    // UserParameters::UpdateFileNamesAndSteps(...) in userparam.cpp
    UserParameters& userparams = registry.GetUserParameters();
    if (userparams.GetWriteTraceFiles())
    {
        string tracefilename = userparams.GetCurrentTraceFileName();
        long stepnumber(userparams.GetNextStep(initialOrFinal));  // EWFIX.CHAINTYPE
        DoubleVec1d values;
        values.push_back(chstate.GetTree()->GetDLValue());
        if (registry.GetChainParameters().IsBayesian())
        {
            DoubleVec1d params = m_paramcoll[m_region][m_replicate]->GetLastParameterVec();
            const ParamVector pvec(true);
            assert(pvec.size() == params.size());
            for (size_t p=0; p<pvec.size(); p++)
            {
                if (pvec[p].IsValidParameter())
                {
                    values.push_back(params[p]);
                }
            }
        }
        ofstream of;
        of.open(tracefilename.c_str(),ios::out | ios::app);
        of << ToString(stepnumber);
        for (size_t i=0; i<values.size(); i++)
        {
            of << "\t" << ToString(values[i]);
        }
        of << endl;
        of.close();
    }
}

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

void CollectionManager::WriteNewickTreeFile(ChainState& chstate)
{
    UserParameters& userparams = registry.GetUserParameters();
    if (userparams.GetWriteNewickTreeFiles())
    {
        const Tree* tree = chstate.GetTree();
        double dl = tree->GetDLValue();
        double bestlike = userparams.GetCurrentBestLike();
        string newickfilename = userparams.GetCurrentNewickTreeFileName();
        if (dl > bestlike)
        {
            userparams.SetCurrentBestLike(dl);
            NewickConverter nc;
            string newick = nc.LamarcTreeToNewickString(*tree);
            ofstream of;
            of.open(newickfilename.c_str(),ios::out | ios::trunc);
            of << dl << endl;
            of << newick << endl;
            of.close();
            userparams.AddNewickTreeFileName(newickfilename);
        }
    }
}

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

void CollectionManager::WriteReclocsFile(ChainState& chstate, bool lastchain)
{
    if(lastchain)
    {
        if (registry.GetForceSummary().CheckForce(force_REC) && registry.GetUserParameters().GetWriteReclocFiles())
        {
            const Tree* tree = chstate.GetTree();

            UserParameters& userparams = registry.GetUserParameters();
            string reclocfilename = userparams.GetCurrentReclocFileName();
            ofstream of;
            of.open(reclocfilename.c_str(),ios::out | ios::app);

            long offset = registry.GetCurrentReclocOffset();
            bool mustConvert = registry.GetConvertOutputToEliminateZeroes();

            const TimeList & timeList = tree->GetTimeList();
            Branchconstiter brit;
            std::set<long> recs;
            for (brit = timeList.FirstBody(); brit != timeList.EndBranch(); brit = timeList.NextBody(brit))
            {
                Branch_ptr pBranch = *brit;
                if (pBranch->Event() == btypeRec)
                {
                    long site = boost::dynamic_pointer_cast<RBranch>(*brit)->GetRecpoint();
                    recs.insert(site);
                }
            }

            std::set<long>::iterator it;
            for(it=recs.begin(); it != recs.end(); it++)
            {
                long newindex = *it + offset;
                if (mustConvert && (newindex <= 0))
                {
                    newindex = newindex - 1;
                }

                of << newindex << endl;
            }

            of.close();
        }
    }
}

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

#ifdef LAMARC_QA_TREE_DUMP

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

void CollectionManager::WriteArgFile(ChainState& chstate, bool lastchain)
{
    if(lastchain)
    {
        if (registry.GetUserParameters().GetWriteArgFiles())
        {
            const Tree* tree = chstate.GetTree();
            const TimeList & tl = tree->GetTimeList();

            TiXmlDocument * docP = tl.AssembleGraphML();

            UserParameters& userparams = registry.GetUserParameters();
            string argfilename = userparams.GetCurrentArgFileName();
            bool writeMany = userparams.GetWriteManyArgs();
            if (!writeMany)
            {
                ofstream of;
                of.open(argfilename.c_str(),ios::trunc);
                of.close();
            }

            FILE * argOutput = TiXmlFOpen(argfilename.c_str(),"a");

            docP->Print(argOutput,0);
            fclose(argOutput);

            // EWFIX.TREEDUMP -- clean up pointer storage
        }
    }
}

#endif

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

void CollectionManager::CorrectForFatalAttraction(long region)
{
    if (m_sample_trees)
        m_treecoll[m_region][m_replicate]->CorrectForFatalAttraction(region);
    if (m_sample_params)
        m_paramcoll[m_region][m_replicate]->CorrectForFatalAttraction(region);
    if (m_sample_maps)
        m_mapcoll[m_region][m_replicate]->CorrectForFatalAttraction(region);
} // CorrectForFatalAttraction

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

void CollectionManager::StartChain(long reg, long rep, bool lastchain)
{
    //LS Note:  If you add code in this function, make sure that this doesn't
    // break reading in partially-written summary files.
    // See ChainManager::ReadInSumFile and ChainManager::ReadInRecover.

    // this is another chain for the same region and replicate.
    // We store only one chain per region/replicate pair, so throw
    // away the old one and store this one.
    delete m_treecoll[reg][rep];
    delete m_paramcoll[reg][rep];
    delete m_mapcoll[reg][rep];

    m_region = reg;
    m_replicate = rep;

    m_treecoll[reg][rep]  = new TreeCollector;
    m_paramcoll[reg][rep] = new ParamCollector;
    m_mapcoll[reg][rep] = new MapCollector(reg);

    m_sample_maps = m_mapcoll[reg][rep]->DoWeCollectFor(lastchain);
} // StartChain

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

void CollectionManager::WriteThisChainsCollections(ofstream* out) const
{
    if (m_sample_trees)
        m_treecoll[m_region][m_replicate]->WriteCollectionWhenScore(out);
    if (m_sample_params)
        m_paramcoll[m_region][m_replicate]->WriteCollectionWhenScore(out);
    if (m_sample_maps)
        m_mapcoll[m_region][m_replicate]->WriteCollectionWhenScore(out);
}

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

void CollectionManager::WriteLastSummaries() const
{
    if (m_sample_trees)
        m_treecoll[m_region][m_replicate]->WriteLastSummary();
    if (m_sample_params)
        m_paramcoll[m_region][m_replicate]->WriteLastSummary();
    if (m_sample_maps)
        m_mapcoll[m_region][m_replicate]->WriteLastSummary();
}

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

void CollectionManager::WriteAllSummaries(long reg, long rep, ofstream& out) const
{
    if (m_sample_trees)
        m_treecoll[reg][rep]->WriteAllSummaries(out);
    if (m_sample_params)
        m_paramcoll[reg][rep]->WriteAllSummaries(out);
    if (m_sample_maps)
        m_mapcoll[reg][rep]->WriteAllSummaries(out);
}

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

TreeCollector* CollectionManager::GetTreeColl(long region, long rep) const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(region < static_cast<long>(m_treecoll.size()));
    assert(rep < static_cast<long>(m_treecoll[region].size()));
    assert(m_sample_trees);

    return m_treecoll[region][rep];
} // GetTreeColl

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

vector<TreeCollector*> CollectionManager::GetTreeColl(long region) const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(region < static_cast<long>(m_treecoll.size()));
    assert(m_sample_trees);

    return m_treecoll[region];
} // GetTreeColl

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

vector<vector<TreeCollector*> > CollectionManager::GetTreeColl() const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(m_sample_trees);

    return m_treecoll;
} // GetTreeColl

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

ParamCollector* CollectionManager::GetParamColl(long region, long rep) const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(region < static_cast<long>(m_paramcoll.size()));
    assert(rep < static_cast<long>(m_paramcoll[region].size()));
    assert(m_sample_params);

    return m_paramcoll[region][rep];
} // GetParamColl

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

vector<ParamCollector*> CollectionManager::GetParamColl(long region) const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(region < static_cast<long>(m_paramcoll.size()));
    assert(m_sample_params);

    return m_paramcoll[region];
} // GetParamColl

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

vector<vector<ParamCollector*> > CollectionManager::GetParamColl() const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(m_sample_params);

    return m_paramcoll;

} // GetParamColl

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

MapCollector* CollectionManager::GetMapColl(long region, long rep) const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(region < static_cast<long>(m_mapcoll.size()));
    assert(rep < static_cast<long>(m_mapcoll[region].size()));
    assert(m_mapcoll[region][rep]->DoWeCollectFor(true));

    return m_mapcoll[region][rep];
} // GetMapColl

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

vector<MapCollector*> CollectionManager::GetMapColl(long region) const
{
    // are we trying to retrieve a chain we didn't expect?  That's bad.
    assert(region < static_cast<long>(m_mapcoll.size()));
    assert(m_mapcoll[region][0]->DoWeCollectFor(true));

    return m_mapcoll[region];
} // GetMapColl

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

vector<vector<MapCollector*> > CollectionManager::GetMapColl() const
{
    // We only collect for some regions, and can't tell (here) which ones.
    // assert(m_sample_maps);

    return m_mapcoll;
} // GetMapColl

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

void CollectionManager::UpdateArrangers(ArrangerVec& arrangers, ChainState& chstate) const
{
    // Currently, only parameter changes trigger arranger updates.
    // This can easily be changed here.
    m_paramcoll[m_region][m_replicate]->UpdateArrangers(arrangers, chstate);
} // UpdateArrangers

//____________________________________________________________________________________
