// $Id: collector.cpp,v 1.33 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 <fstream>

#include "local_build.h"

#include "arranger.h"
#include "arrangervec.h"
#include "collector.h"
#include "registry.h"
#include "region.h"
#include "sumfilehandler.h"
#include "treesum.h"
#include "xmlsum_strings.h"

using namespace std;

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

Collector::Collector()
    : m_areWriting(false), m_sumout()
{
    // intentionally blank
}

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

void Collector::WriteCollectionWhenScore(ofstream* out)
{
    m_areWriting = true;
    m_sumout = out;
}

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

TreeCollector::TreeCollector()
    : forceParameters(unknown_region),
      m_findbesttree(true /* false */), // JDEBUG--set from menu/xml/compileflag?
      m_besttreedatallike(NEGMAX)
{
    // intentionally blank
}

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

TreeCollector::~TreeCollector()
{
    vector<TreeSummary*>::iterator treeiter;
    for(treeiter = treeSummaries.begin(); treeiter != treeSummaries.end(); treeiter++)
    {
        delete(*treeiter);
    }
}

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

void TreeCollector::Score(ChainState& chstate)
{
    if (chstate.TreeNeedsSampling() || chstate.StickNeedsSampling())
    {
        // a novel tree; summarize it
        if (m_areWriting)
        {
            WriteLastSummary();
            // Write the last tree summary.  Note that this will fail to write
            //  the one we're collecting--do that later.
        }
        treeSummaries.push_back(chstate.GetTree()->SummarizeTree());
    }
    else
    {
        // a previously seen tree; increment its count
        treeSummaries.back()->AddCopy();
    }
    chstate.TreeSampled();

} // Score

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

void TreeCollector::WriteLastSummary()
{
    if (m_areWriting && treeSummaries.size())
    {
        treeSummaries[treeSummaries.size()-1]->WriteTreeSummary(*m_sumout);
    }
}

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

void TreeCollector::WriteAllSummaries(ofstream& out)
{
    for (unsigned long index=0; index<treeSummaries.size(); index++)
    {
        treeSummaries[index]->WriteTreeSummary(out);
    }
}

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

void TreeCollector::CorrectForFatalAttraction(long region)
{
    vector<TreeSummary*>::iterator treesum = treeSummaries.begin();
    map<force_type, DoubleVec1d> totals;
    map<force_type, DoubleVec1d>::iterator totit;
    map<force_type, DoubleVec1d>::iterator newit;

    if (treeSummaries.empty()) return;  // nothing to do--possible in the
    // case of a Bayesian run

    // accumulate, over all treesummaries, total events for each force
    totals = (*treesum)->InspectSummary();   // first summary establishes number of params
    ++treesum;

    for (; treesum != treeSummaries.end(); ++treesum)
    {
        map<force_type, DoubleVec1d> newtot = (*treesum)->InspectSummary();
        assert(totals.size() == newtot.size());  // different number of forces?!
        for (totit = totals.begin(), newit = newtot.begin();
             totit != totals.end();
             ++totit, ++newit)
        {
            assert(totit->second.size() == newit->second.size());  // different n. of parameters?!
            // this line assumes that forces are in same order everywhere!
            transform(totit->second.begin(), totit->second.end(),
                      newit->second.begin(),
                      totit->second.begin(),
                      plus<double>());
        }
    }

    // We adjust the bin counts in the first tree summary.  This is
    // an arbitrary choice; as far as I know any tree summary would do.
    // The first tree might be atypical anyway, and therefore messing it
    // up is relatively harmless.

    // This call will do nothing if all parameters have at least one
    // event, but that condition cannot be detected here because of
    // legitimate zero entries on the diagonal for some forces.

#ifndef LAMARC_QA_SINGLE_DENOVOS
    treeSummaries.front()->AdjustSummary(totals, region);
#endif // LAMARC_QA_SINGLE_DENOVOS

} // CorrectForFatalAttraction

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

void TreeCollector::AddTreeSummary(TreeSummary* ts)
{
    assert(ts);  // bad tree summary!
    treeSummaries.push_back(ts);

} // AddTreeSummary

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

long TreeCollector::TreeCount() const
{
    long ntrees(0L);

    vector<TreeSummary*>::const_iterator treesum = treeSummaries.begin();
    for (; treesum != treeSummaries.end(); ++treesum)
        ntrees += (*treesum)->GetNCopies();

    return ntrees;
} // TreeCount

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

void ParamCollector::Score(ChainState& chstate)
{
    if (chstate.ParametersNeedSampling())
    {
        // a novel set of parameters; summarize them
        if (m_areWriting)
        {
            WriteLastSummary();
        }
        ForceParameters& fp = chstate.GetParameters();
        AddParamSummary(fp, 1L);
    }
    else
    {
        // a previously seen set; increment their count
        pair<ForceParameters, long>& previous = m_paramsum.back();
        ++previous.second;
    }
    chstate.ParametersSampled();

} // Score

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

void ParamCollector::AddParamSummary(ForceParameters fp, long ncopy)
{
    m_paramsum.push_back(make_pair(fp, ncopy));
}

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

void ParamCollector::WriteLastSummary()
{
    long index = static_cast<long>(m_paramsum.size())-1;
    if (index >= 0)
    {
        WriteParamSumm(index);
    }
}

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

void ParamCollector::WriteAllSummaries(ofstream& out)
{
    m_sumout = &out;
    m_areWriting = true;
    for (unsigned long index = 0; index<m_paramsum.size(); index++)
    {
        WriteParamSumm(index);
    }
}

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

void ParamCollector::WriteParamSumm(unsigned long index)
{
    assert (m_paramsum.size() > index);
    assert (m_areWriting);

    ForceParameters fp = m_paramsum[index].first;
    long ncopy         = m_paramsum[index].second;
    if ( m_sumout->is_open() )
    {
        *m_sumout << "\t" << xmlsum::PARAM_SUMMARY_START << endl;
        *m_sumout << "\t\t" << xmlsum::NCOPY_START << " " << ncopy
                  << " "    << xmlsum::NCOPY_END   << endl;
        fp.WriteForceParameters(*m_sumout, 2);
        *m_sumout << "\t" << xmlsum::PARAM_SUMMARY_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("WriteParamSumm");
}

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

void ParamCollector::UpdateArrangers(ArrangerVec& arrangers, ChainState& chstate) const
{
    // If the parameters have changed, inform the Arrangers
    if (chstate.ParametersNeedExport())
    {
        arrangerit arr = arrangers.begin();
        for ( ; arr != arrangers.end(); ++arr)
        {
            (*arr).second->SetParameters(chstate);
        }
        chstate.UpdateOldStickParams();
        chstate.UpdateNewStickParams();
        chstate.ParametersExported();
    }
} // UpdateArrangers

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

const DoubleVec1d ParamCollector::GetLastParameterVec() const
{
    assert (m_paramsum.size() > 0);
    return m_paramsum[m_paramsum.size()-1].first.GetGlobalParameters();
}

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

MapCollector::MapCollector(long region)
    : Collector(),
      m_region(region),
      m_nsamples(registry.GetChainParameters().GetNSamples(defaults::final)),
      m_mapsum(),
      m_lastmapsum(),
      m_lastcount(0)
{
    DoubleVec1d zeroes(registry.GetDataPack().GetRegion(region).GetNumSites(), 0.0);
    DoubleVec2d mapsum(registry.GetDataPack().GetNumMovingLoci(region), zeroes);
    m_mapsum = mapsum;
}

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

bool MapCollector::DoWeCollectFor(bool lastchain)
{
    if (m_mapsum.size() == 0)
    {
        return false;
    }
    if ((registry.GetDataPack().GetRegion(m_region).AnyJumpingAnalyses()) || (lastchain))
    {
        return true;
    }
    return false;
}

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

void MapCollector::Score(ChainState& chstate)
{
    if (chstate.MapNeedsSampling())
    {
        // a novel tree; summarize it
        if (m_areWriting)
        {
            WriteLastSummary();
            // Write the last tree summary.  Note that this will fail to write
            //  the one we're collecting--do that later.
        }
        AccumulateMap();
        m_lastmapsum = dynamic_cast<RecTree*>(chstate.GetTree())->GetMapSummary();
        m_lastcount = 1;
    }
    else
    {
        // a previously seen tree; increment its count
        m_lastcount++;
    }
    chstate.MapSampled();

} // ScoreMap

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

void MapCollector::AddMapSummary(DoubleVec2d& maplikes, long ncopy)
{
    AccumulateMap();
    m_lastmapsum = maplikes;
    m_lastcount = ncopy;
}

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

void MapCollector::WriteLastSummary()
{
    if (m_lastmapsum.size() == 0) return;
    if ( m_sumout->is_open() )
    {
        *m_sumout << "\t" << xmlsum::MAP_SUMMARIES_START << endl;
        *m_sumout << "\t\t" << xmlsum::NCOPY_START << " " << m_lastcount
                  << " "    << xmlsum::NCOPY_END   << endl;
        for (unsigned long summary=0; summary < m_lastmapsum.size(); summary++)
        {
            *m_sumout << "\t\t" << xmlsum::MAP_SUMMARY_START;
            rangeset range = registry.GetDataPack().GetRegion(m_region).GetMovingLocus(summary).GetAllowedRange();
            WriteVec1D(*m_sumout, m_lastmapsum[summary], range);
            *m_sumout << " " << xmlsum::MAP_SUMMARY_END << endl;
        }
        *m_sumout << "\t" << xmlsum::MAP_SUMMARIES_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("WriteLastSummary");
}

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

void MapCollector::WriteAllSummaries(ofstream& out)
{
    AccumulateMap();
    if ( out.is_open() )
    {
        out << "\t" << xmlsum::MAP_SUMMARIES_START << endl;
        out << "\t\t" << xmlsum::NCOPY_START << " " << m_nsamples
            << " "    << xmlsum::NCOPY_END   << endl;
        for (unsigned long mlocus=0; mlocus < m_mapsum.size(); mlocus++)
        {
            out << "\t\t" << xmlsum::MAP_SUMMARY_START;
            const Locus& mloc = registry.GetDataPack().GetRegion(m_region).GetMovingLocus(mlocus);
            rangeset range = mloc.GetAllowedRange();
            if (mloc.GetAnalysisType() == mloc_mapfloat)
            {
                DoubleVec1d logs = SafeLog(m_mapsum[mlocus]);
                WriteVec1D(out, logs, range);
            }
            else
            {
                WriteVec1D(out, m_mapsum[mlocus], range);
            }
            out << " " << xmlsum::MAP_SUMMARY_END << endl;
        }
        out << "\t" << xmlsum::MAP_SUMMARIES_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("WriteAllSummaries");
}

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

DoubleVec2d MapCollector::GetMapSummary()
{
    AccumulateMap();
    return m_mapsum;
}

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

void MapCollector::AccumulateMap()
{
    if (m_lastmapsum.size() == 0)
    {
        //Nothing to accumulate
        return;
    }
    for (unsigned long locus=0; locus<m_lastmapsum.size(); locus++)
    {
        DoubleVec1d reallikes = m_lastmapsum[locus];
        if (registry.GetDataPack().GetRegion(m_region).GetMovingLocus(locus).GetAnalysisType() == mloc_mapfloat)
        {
            reallikes = SafeExp(m_lastmapsum[locus]);
        }
        //Multiply m_lastmapsum by m_lastcount/m_nsamples
        double weight = static_cast<double>(m_lastcount) / static_cast<double>(m_nsamples);
        transform(reallikes.begin(),
                  reallikes.end(),
                  reallikes.begin(),
                  bind2nd(multiplies<double>(),weight));
        assert(reallikes.size() == m_mapsum[locus].size());

        transform(reallikes.begin(),
                  reallikes.end(),
                  m_mapsum[locus].begin(),
                  m_mapsum[locus].begin(),
                  plus<double>());
    }
    m_lastmapsum.clear();
    m_lastcount = 0;
}

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

void MapCollector::WriteVec1D(ofstream& sumout, DoubleVec1d& vec, rangeset range)
{
    if (sumout.is_open())
    {
        StringVec1d strings(vec.size(), "-");
        for(rangeset::iterator currRange=range.begin(); currRange != range.end(); currRange++)
        {
            for (long site = currRange->first; site < currRange->second; ++site)
            {
                strings[site] = ToString(vec[site]);
            }
        }
        for (StringVec1d::iterator num = strings.begin(); num != strings.end(); ++num)
        {
            sumout << *num << " ";
        }
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("WriteVec1D");
} // WriteVec1D

//____________________________________________________________________________________
