// $Id: chainpack.cpp,v 1.36 2018/01/03 21:32:54 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 <cstdlib>
#include <iostream>

#include "chainpack.h"
#include "constants.h"
#include "region.h"
#include "registry.h"
#include "stringx.h"
#include "sumfilehandler.h"
#include "xmlsum_strings.h"             // for xml sumfile handling

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

using namespace std;

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

ChainOut ChainPack::GetChain(long int region, long int rep, long int chain) const
{
    return(m_chains[region][rep][chain]);
} // ChainPack::GetChain

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

ChainOut ChainPack::GetLastChain() const
{
    long int region = m_chains.size() - 1;
    if (region != -1)
    {
        long int rep = m_chains[region].size() - 1;
        if (rep != -1)
        {
            long int chain  = m_chains[region][rep].size() - 1;
            if (chain != -1)
            {
                return(m_chains[region][rep][chain]);
            }
        }
    }
    //There was no last chain
    ChainOut blankchout;
    if (region == -1) region = 0;
    ForceParameters fp(registry.GetForceSummary().GetStartParameters(), region);
    blankchout.SetEstimates(fp);
    return blankchout;

} // ChainPack::GetLastChain

ChainOut ChainPack::GetLastChain(long int region) const
{
    assert(region <= static_cast<long int>(m_chains.size()));
    assert(m_chains[region].size() > 0);
    unsigned long int rep = m_chains[region].size() - 1;
    assert(m_chains[region][rep].size() > 0);
    unsigned long int chain = m_chains[region][rep].size() - 1;
    return(m_chains[region][rep][chain]);
} // ChainPack::GetLastChain

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

ChainOut ChainPack::GetRegion(long int region) const
{
    if (m_regions.size() == 0)          // no sum over replicates
    {
        long int last = (m_chains[region][0].size())-1;
        return(this->GetChain(region, 0, last));
    }
    else
    {
        return(m_regions[region]);
    }
} // ChainPack::GetRegion

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

ChainOut ChainPack::GetOverall() const
{
    if (m_overall.size() == 0)          // no sum over regions
    {
        return(this->GetRegion(0));
    }
    else
    {
        return(m_overall[0]);
    }
} // ChainPack::GetOverall

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

// Two functions for the summary file reader to make sure the
//  region/replicate summaries are being set.

long int ChainPack::GetLenRegionsVec() const
{
    return (m_regions.size());
}

long int ChainPack::GetLenOverallVec() const
{
    return (m_overall.size());
}

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

void ChainPack::SetChain(ChainOut& chout, long int region, long int rep, long int chain)
{
    if (region > m_currentRegion)
    {
        EndRegion();
        VerifyPosition(region, rep, chain);
        SetChain(chout);
        return;
    }
    if (rep > m_currentRep)
    {
        EndReplicate();
        VerifyPosition(region, rep, chain);
        SetChain(chout);
        return;
    }
    VerifyPosition(region, rep, chain);
    SetChain(chout);
}

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

void ChainPack::ResetLastChain(ChainOut& chout, long int region)
{
    long int rep = m_chains[region].size()-1;
    long int chain = m_chains[region].size()-1;
    assert (rep >= 0);
    assert (chain >= 0);
    m_chains[region][rep][chain] = chout;
}

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

void ChainPack::ResetSummaryOverReps(ChainOut& chout, long int region)
{
    assert(region <= static_cast<long int>(m_regions.size()));
    m_regions[region] = chout;
}

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

void ChainPack::SetChain(ChainOut& chout)
{
    // The estimates in chout might not have known what region they were for.
    // If that's the case, this will set it straight.
    ForceParameters fp(chout.GetEstimates(), m_currentRegion);
    chout.SetEstimates(fp);

    if (m_currentRep == 0 && m_currentChain == 0) // first entry of its region
    {
        vector<ChainOut> v1;
        v1.push_back(chout);
        vector<vector<ChainOut> > v2;
        v2.push_back(v1);
        m_chains.push_back(v2);
        m_currentChain++;
        return;
    }
    if (m_currentChain == 0)            // first entry of its replicate
    {
        vector<ChainOut> v1;
        v1.push_back(chout);
        m_chains[m_currentRegion].push_back(v1);
        m_currentChain++;
        return;
    }

    m_chains[m_currentRegion][m_currentRep].push_back(chout);
    m_currentChain++;
} // ChainPack::SetChain

//------------------------------------------------------------------------------------
// RemoveLastChain is needed by ChainManager in the unlikely event that
//  we have read in chain summary information for the last chain in a
//  region/replicate, but have no tree summary information, and thus
//  need to reproduce that tree, giving us a new chain.
//                 --Lucian

void ChainPack::RemoveLastChain()
{
    if (m_currentChain == 0)
    {
        //We called EndRegion and/or EndReplicate too early, and have to back up.
        if (m_currentRep == 0)
        {
            if (m_currentRegion == 0)
            {
                //There are no chains to remove--do nothing
                assert(false); //Why did this happen?
                return;
            }
        }
        m_currentRegion = m_chains.size()-1;
        m_currentRep = m_chains[m_currentRegion].size()-1;
        m_currentChain = m_chains[m_currentRegion][m_currentRep].size()-1;
    }
    m_chains[m_currentRegion][m_currentRep].pop_back();
    m_currentChain--;
}

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

void ChainPack::SetSummaryOverReps(ChainOut& chout)
{
    // a convenience variable
    vector<vector<ChainOut> >& region = m_chains[m_currentRegion];

    // set regional start/end times
    long int lastrep = (region.size())-1;
    long int lastchain = (region[0].size())-1;
    chout.SetStarttime(region[0][0].GetStarttime());
    chout.SetEndtime(region[lastrep][lastchain].GetEndtime());

    // This information is not helpful for the user so we'll supress it:
    chout.SetNumBadTrees(FLAGLONG);
    chout.SetNumTinyPopTrees(FLAGLONG);
    chout.SetNumStretchedTrees(FLAGLONG);
    chout.SetNumZeroDLTrees(FLAGLONG);
    chout.SetAccrate(FLAGDOUBLE);
    chout.SetLlikedata(FLAGDOUBLE);

    // The estimates in chout might not have known what region they were for.
    // If that's the case, this will set it straight.
    ForceParameters fp(chout.GetEstimates(), m_currentRegion);
    chout.SetEstimates(fp);
    m_regions.push_back(chout);

} // ChainPack::SetSummaryOverReps

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

void ChainPack::SetSummaryOverRegions(ChainOut& chout)
{
    assert(chout.GetEstimates().GetParamSpace() == global_region);
    // set overall start/end times
    long int lastregion = (m_chains.size()) - 1;
    long int lastrep = (m_chains[0].size()) - 1;
    long int lastchain = (m_chains[0][0].size()) - 1;

    chout.SetStarttime(m_chains[0][0][0].GetStarttime());
    chout.SetEndtime(m_chains[lastregion][lastrep][lastchain].GetEndtime());

    // This information is not helpful for the user so we'll supress it:
    chout.SetNumBadTrees(FLAGLONG);
    chout.SetNumTinyPopTrees(FLAGLONG);
    chout.SetNumStretchedTrees(FLAGLONG);
    chout.SetNumZeroDLTrees(FLAGLONG);
    chout.SetAccrate(FLAGDOUBLE);
    chout.SetLlikedata(FLAGDOUBLE);

    if (m_overall.size() > 0)
    {
        m_overall[0] = chout;
        //Useful for re-setting this when reading from a summary file.
    }
    else
    {
        m_overall.push_back(chout);
    }
} // ChainPack::SetSummaryOverRegions

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

void ChainPack::EndRegionIfNecessary(long int region)
{
    if (m_currentRegion > region) return;
    EndRegion();
}

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

void ChainPack::EndRegionsOrRepsAsNeeded(long int maxreps, long int maxchains)
{
    if (m_currentChain >= maxchains)
    {
        EndReplicate();
    }

    if (m_currentRep >= maxreps)
    {
        EndRegion();
    }
}

/**************************************************
 The following functions provide summaries of the
 parameter values from a ChainPack.
**************************************************/

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

DoubleVec1d ChainPack::RegionalMeanParams() const
{
    // summarizes the current region
    unsigned long int rep;
    DoubleVec2d values;
    long int lastchain = m_chains[m_currentRegion][0].size() - 1;

    for (rep = 0; rep < m_chains[m_currentRegion].size(); ++rep)
    {
        const ChainOut& chout = GetChain(m_currentRegion, rep, lastchain);
        values.push_back(chout.GetEstimates().GetRegionalParameters());
    }

    return CalcMean(values);

} // RegionalMeanParams

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

DoubleVec1d ChainPack::OverallMeanParams() const
{
    // summarizes over all regions
    unsigned long int region;
    DoubleVec2d values;

    for (region = 0; region < m_chains.size(); ++region)
    {
        const ChainOut& chout = GetRegion(region);
        values.push_back(chout.GetEstimates().GetGlobalParameters());
    }

    return CalcMean(values);

} // OverallMeanParams

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

time_t ChainPack::GetStartTime() const
{
    ChainOut firstchain = GetChain(0, 0, 0);
    return(firstchain.GetStarttime());
} // GetStartTime

time_t ChainPack::GetStartTime(long int region) const
{
    ChainOut firstchain = GetChain(region, 0, 0);
    return(firstchain.GetStarttime());
} // GetStartTime

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

time_t ChainPack::GetEndTime() const
{
    ChainOut lastchain = GetOverall();
    return(lastchain.GetEndtime());
    // daniel 040903
    // why does lastchain.GetEndTime() sometimes return 0?  when the below endtime is correct?
    // see output file from reading in sumfile
    // bug in the overall vector?

#if 0
    long int i, j, k;
    i = m_chains.size() - 1;
    j = m_chains[i].size() - 1;
    k = m_chains[i][j].size() - 1;
    return m_chains[i][j][k].endtime;
#endif
} // GetEndTime

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

// CalcMean assumes that all subvectors are as long as the first one
DoubleVec1d ChainPack::CalcMean(const DoubleVec2d& src) const
{
    long int index;
    long int size = src[0].size();
    DoubleVec2d::const_iterator values;
    DoubleVec1d mean;

    for(index = 0; index < size; ++index)
    {
        double sum = 0.0;
        for(values = src.begin(); values != src.end(); ++values)
            sum += (*values)[index];
        mean.push_back(sum /= src.size());
    }
    return mean;

} // ChainPack::CalcMean

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

// following setter and getter for resuming with sumfiles
vector<vector<vector<ChainOut > > > ChainPack::GetAllChains() const
{ return m_chains; }

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

void ChainPack::WriteLastChain ( ofstream& sumout) const
{
    if ( sumout.is_open() )
    {
        sumout << xmlsum::CHAINPACK_START << endl;
        sumout << "\t" << xmlsum::NUMBER_START
               << " " << m_currentRegion << " " << m_currentRep << " " << (m_currentChain -1)
               << " " << xmlsum::NUMBER_END << endl;
        WriteChainOut( sumout, m_chains[m_currentRegion][m_currentRep][(m_currentChain-1)] );
        WriteAlphas(sumout, m_currentRegion, m_currentRep, m_currentChain-1);
        sumout << xmlsum::CHAINPACK_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("ChainPack::WriteLastChain");
} // ChainPack::WriteLastChain

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

void ChainPack::WriteChain(ofstream& sumout, long int region, long int rep, long int chain) const
{
    if ( sumout.is_open() )
    {
        sumout << xmlsum::CHAINPACK_START << endl;
        sumout << "\t" << xmlsum::NUMBER_START
               << " " << region << " " << rep << " " << chain
               << " " << xmlsum::NUMBER_END << endl;
        WriteChainOut( sumout, m_chains[region][rep][chain] );
        WriteAlphas(sumout, region, rep, chain);
        sumout << xmlsum::CHAINPACK_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("ChainPack::WriteChain");
} // ChainPack::WriteLastChain

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

void ChainPack::WriteChainOut( ofstream& sumout, const ChainOut& chout) const
{
    if ( sumout.is_open() )
    {
        long int badtrees       = chout.GetNumBadTrees();
        long int tinytrees      = chout.GetTinyPopTrees();
        long int stretchedtrees = chout.GetStretchedTrees();
        long int zerodltrees    = chout.GetZeroDLTrees();
        double accrate      = chout.GetAccrate();
        double llikemle     = chout.GetLlikemle();
        double llikedata    = chout.GetLlikedata();
        time_t starttime    = chout.GetStarttime();
        time_t endtime      = chout.GetEndtime();
        ratemap rates       = chout.GetAllAccrates();
        ForceParameters fp  = chout.GetEstimates();
        DoubleVec1d temperatures = chout.GetTemperatures();
        DoubleVec1d swaprates    = chout.GetSwaprates();
        LongVec1d   bayesunique  = chout.GetBayesUnique();

        sumout << "\t" << xmlsum::CHAINOUT_START << endl;
        sumout << "\t\t" << xmlsum::BADTREES_START  << " " << badtrees
               << " " << xmlsum::BADTREES_END << endl;
        sumout << "\t\t" << xmlsum::TINYTREES_START << " " << tinytrees
               << " " << xmlsum::TINYTREES_END << endl;
        sumout << "\t\t" << xmlsum::STRETCHEDTREES_START << " " << stretchedtrees
               << " " << xmlsum::STRETCHEDTREES_END << endl;
        sumout << "\t\t" << xmlsum::ZERODLTREES_START << " " << zerodltrees
               << " " << xmlsum::ZERODLTREES_END << endl;
        sumout << "\t\t" << xmlsum::ACCRATE_START   << " " << accrate
               << " " << xmlsum::ACCRATE_END << endl;
        sumout << "\t\t" << xmlsum::LLIKEMLE_START  << " " << llikemle
               << " " << xmlsum::LLIKEMLE_END << endl;
        sumout << "\t\t" << xmlsum::LLIKEDATA_START << " " << llikedata
               << " " << xmlsum::LLIKEDATA_END << endl;
        sumout << "\t\t" << xmlsum::STARTTIME_START << " " << starttime
               << " " << xmlsum::STARTTIME_END << endl;
        sumout << "\t\t" << xmlsum::ENDTIME_START   << " " << endtime
               << " " << xmlsum::ENDTIME_END << endl;

        sumout << "\t\t" << xmlsum::RATES_START << " ";
        map< string, pair<long int, long int> >::iterator rit;
        for (rit = rates.begin(); rit != rates.end(); ++rit)
        {
            sumout << xmlsum::MAP_START << " "; // map<string, pair<long int, long int> >
            sumout << rit->first << " "
                   << rit->second.first << " " << rit->second.second << " ";
            sumout << xmlsum::MAP_END << " ";
        }
        sumout << xmlsum::RATES_END << endl;

        long int numtemps = chout.GetNumtemps();
        if (numtemps > 1)
        {
            sumout << "\t\t" << xmlsum::TEMPERATURES_START << " ";
            SumFileHandler::WriteVec1D(sumout, temperatures);
            sumout << xmlsum::TEMPERATURES_END << endl;

            sumout << "\t\t" << xmlsum::SWAPRATES_START << " ";
            SumFileHandler::WriteVec1D(sumout, swaprates);
            sumout << xmlsum::SWAPRATES_END << endl;
        }
        if (bayesunique.size() > 0)
        {
            sumout << "\t\t" << xmlsum::BAYESUNIQUE_START << " ";
            SumFileHandler::WriteVec1D(sumout, bayesunique);
            sumout << xmlsum::BAYESUNIQUE_END << endl;
        }
        fp.WriteForceParameters(sumout, 2);
        sumout << "\t" << xmlsum::CHAINOUT_END << endl;
    }
    else
        SumFileHandler::HandleSumOutFileFormatError("ChainPack::WriteChainOut");
}

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

void ChainPack::WriteAlphas(ofstream& sumout, long int reg, long int rep, long int chain) const
{
    const Region& region = registry.GetDataPack().GetRegion(reg);
    for (long int loc = 0; loc < region.GetNloci(); loc++)
    {
        const Locus & locus = region.GetLocus(loc);
        locus.GetDataModel()->WriteAlpha(sumout, loc, rep, chain);
    }
}

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

void ChainPack::VerifyPosition( long int test_reg, long int test_rep, long int test_chain)
{
    if ((test_reg != m_currentRegion) || (test_rep != m_currentRep) || (test_chain != m_currentChain))
    {
        string msg = "Tried to set the chainpack for region " + ToString(test_reg)
            + ", replicate " + ToString(test_rep) + ", and chain "
            + ToString(test_chain) + ", but we needed the chainpack for region "
            + ToString(m_currentRegion) + ", replicate " + ToString(m_currentRep)
            + ", and chain " + ToString(m_currentChain) + ".";
        throw data_error(msg);
    }
}

//____________________________________________________________________________________
