// $Id: datapack.cpp,v 1.71 2018/01/03 21:32:57 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 <functional>     // for plus<> used in CountOverallBaseFreqs
#include <iostream>
#include <numeric>        // for accumulate() used in CountOverallBaseFreqs
#include <fstream>        // for WriteFlucFile()

#include "datapack.h"
#include "region.h"
#include "registry.h"
#include "constants.h"
#include "individual.h"
#include "xml_strings.h"  // for class xmlstr in DataPack::ToXML()
#include "stringx.h"      // for MakeTag(), MakeCloseTag(), etc in
                          //    DataPack::ToXML()

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

using namespace std;

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

DataPack::~DataPack()
{
    /* EWFIX.P3 -- causes seg fault in static destruction on mac
       EWFIX.P3 -- LEAK
       // once and for all, clean out those DLCell arrays
       DLCell::ClearStore();
    */

    // clean out the regions
    vector<Region*>::iterator rit;
    for(rit = m_regions.begin(); rit != m_regions.end(); ++rit)
        delete *rit;
} // DataPack::~DataPack

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

long DataPack::GetNTips() const
{
    long ntips = 0;
    vector<Region*>::const_iterator rit;
    for(rit = m_regions.begin(); rit != m_regions.end(); ++rit)
        ntips += (*rit)->GetNTips();

    return ntips;
} // DataPack::GetNTips

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

long DataPack::GetNumTipPopulations() const
{
    set<string> popnames;
    vector<Region*>::const_iterator rit;
    for(rit = m_regions.begin(); rit != m_regions.end(); ++rit)
        (*rit)->AddUniqueNamesTo(popnames);

    long ntippops(popnames.size());
    return ntippops;
} // DataPack::GetNumTipPopulations

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

long DataPack::AddPartition(force_type forceType, const string& partname)
{
    // create the new force entry (if it already doesn't exist) and
    // make sure the new partition entry doesn't already exist.
    // NB:  the code doesn't look like it adds a new force, but due to
    // std::map behavior, it does!
    StringVec1d& partitions(m_partnames[forceType]);
    long part, nparts = partitions.size();
    for(part = 0; part < nparts; ++part)
        if (partitions[part] == partname)
            return part;

    // new partition to add
    partitions.push_back(partname);
    return partitions.size()-1;
} // DataPack::AddPartition

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

long DataPack::GetNPartitionsByForceType(force_type forceType) const
{
    PartitionNames::const_iterator pnames = m_partnames.find(forceType);
    if (pnames != m_partnames.end()) return pnames->second.size();

    // NB:  This is a hack, but what to do?  The function is widely used
    // in both phase 1 and phase 2, and its appropriate answer for
    // MIG depends on whether DIVMIG exists.
    // (The following if statement is "If there's MIG and no DIVMIG".)

    if (forceType == force_MIG && m_partnames.find(force_DIVMIG) == m_partnames.end())
        return 1L;
    else return 0L;
} // DataPack::GetNPartitionsByForceType

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

string DataPack::GetPartitionName(force_type forceType, long index) const
{
    assert(static_cast<unsigned long>(index) < m_partnames.find(forceType)->second.size());
    return m_partnames.find(forceType)->second[index];
} // DataPack::GetPartitionName

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

long DataPack::GetPartitionNumber(force_type forceType, const string& partname) const
{
    const StringVec1d& partitions(m_partnames.find(forceType)->second);
    long part, nparts = partitions.size();
    for(part = 0; part < nparts; ++part)
        if (partitions[part] == partname)
            return part;

    assert(false);  // failed to find anything!
    return FLAGLONG;
} // DataPack::GetPartitionNumber

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

StringVec1d DataPack::GetAllPartitionNames(force_type forceType) const
{
    StringVec1d emptyvec;
    PartitionNames::const_iterator pnames = m_partnames.find(forceType);

    if (pnames != m_partnames.end())
    {
        return pnames->second;
    }
    else if (forceType == force_MIG || forceType == force_DIVMIG)
    {
        return m_outputpopnames;
    }

    return emptyvec;
} // DataPack::GetAllPartitionNames

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

StringVec1d DataPack::GetAllCrossPartitionNames() const
{
    vector<map<force_type,string> > ids = GetCrossPartitionIds(false);
    StringVec1d xpnames;

    vector<map<force_type,string> >::iterator idit;
    for(idit = ids.begin(); idit != ids.end(); ++idit)
    {
        string xpname = idit->begin()->second;
        if (idit->size() > 1) xpname += " (";
        map<force_type,string>::iterator nameit;
        map<force_type,string>::iterator startit(idit->begin());
        for(nameit = ++startit; nameit != idit->end(); ++nameit)
        {
            if (nameit != startit) xpname += ", ";
            xpname += nameit->second;
        }
        if (idit->size() > 1) xpname += ")";
        xpnames.push_back(xpname);
    }

    return xpnames;
} // DataPack::GetAllCrossPartitionNames

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

vector<map<force_type,string> >
DataPack::GetCrossPartitionIds(bool ignoresinglepop) const
{
    if (m_partnames.empty())   // there are no partition forces
    {
        vector<map<force_type,string> > ids(1L);
        if (!ignoresinglepop)
        {
            assert(m_outputpopnames.size() == 1);
            // force_type ft = registry.GetForceSummary().GetNonLocalPartitionForceTag();
            //
            // ids[0].insert(make_pair(ft,m_outputpopnames[0]));
            //
            // cannot use the above lines as the ForceSummary doesn't yet exist in at least
            // one of the function callers
            ids[0].insert(make_pair(force_MIG,m_outputpopnames[0]));
        }

        return ids;   // return the saved population names
    }

    // setup separate collections of forcenames, first partition name and
    // iterator-last partition name, all dimensioned by the number of forces
    vector<force_type> fname;
    vector<StringVec1d::const_iterator> startname;
    vector<StringVec1d::const_iterator> endname;
    PartitionNames::const_iterator forceit;
    for(forceit = m_partnames.begin(); forceit != m_partnames.end(); ++forceit)
    {
        fname.push_back(forceit->first);
        startname.push_back(forceit->second.begin());
        endname.push_back(forceit->second.end());
    }

    // initialize the "current name" to the first partition name of each
    // force
    vector<StringVec1d::const_iterator> curname(startname);

    // loop over each cross partition, building up and adding the appropriate
    // map of names for each cross partition and then storing the map away
    // in the return collection "ids".
    long nxparts = GetNCrossPartitions();
    vector<map<force_type,string> > ids(nxparts);
    long xpart;
    for(xpart = 0; xpart < nxparts; ++xpart)
    {
        map<force_type,string> id;
        long force, nforces = m_partnames.size();
        for(force = 0; force < nforces; ++force)
            id.insert( make_pair( fname[force],*(curname[force])));
        ids[xpart] = id;

        // advance the string iterators for the next cross partition
        for(force = nforces - 1; force > -1; --force)
        {
            if (++curname[force] != endname[force]) break;
            curname[force] = startname[force];
        }
    }

    return ids;
} // DataPack::GetCrossPartitionIds

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

long DataPack::GetNPartitionForces() const
{
    return m_partnames.size();
} // DataPack::GetNPartitionForces

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

long DataPack::GetNCrossPartitions() const
{
    long count = 1;
    PartitionNames::const_iterator pnames;
    for(pnames = m_partnames.begin(); pnames != m_partnames.end(); ++pnames)
        count *= pnames->second.size();

    return count;
} // DataPack::GetNCrossPartitions()

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

long DataPack::GetNPartitions(unsigned long index) const
{
    assert(index < m_finalcount.size());
    return m_finalcount[index];
} // DataPack::GetNPartitions

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

LongVec1d DataPack::GetAllNPartitions() const
{
    return m_finalcount;
} // DataPack::GetAllNPartitions

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

long DataPack::GetCrossPartitionIndex(const LongVec1d& membership) const
{
    if (membership.empty()) return 0L;

    // The following code computes a cross-partition index based on
    // the number of states in each partition force and the given
    // membership vector.  The index is equivalent to writing the
    // forces out in a table which varies the last force first
    // (like counting in binary, except with arbitrary bases for each
    // force) and then numbering the table entries starting at zero.

    // How the code actually works:  It initializes the index with the
    // state of the last force.  Then, for each of the remaining forces,
    // it computes a multiplier based on all later forces and
    // multiplies the force's state by this before adding it to the
    // index.  (For example, if the last force has eight
    // states, each entry of the next-to-last force is equivalent to
    // a table of eight entries.)

    long lastforce = membership.size() - 1;
    long index = membership[lastforce], force, totalcount = 1;
    for (force = lastforce-1; force > -1; --force)
    {
        totalcount *= GetNPartitions(force+1);
        index += totalcount * membership[force];
    }

    return index;
} // DataPack::GetCrossPartitionIndex

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

map<force_type,string> DataPack::GetTipId(long xpartindex) const
{
    map<force_type,string> tipid;
    LongVec1d membership(GetBranchMembership(xpartindex));

    PartitionNames::const_iterator force;
    for(force = m_partnames.begin(); force != m_partnames.end(); ++force)
    {
        force_type forceType = force->first;
        long partitionindex = membership[registry.GetForceSummary().GetPartIndex(forceType)];
        string partitionname = GetPartitionName(forceType,partitionindex);
        tipid.insert(make_pair(forceType,partitionname));
    }

    return tipid;
} // DataPack::GetTipId

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

LongVec1d DataPack::GetBranchMembership(long xpartindex) const
{
    long npartforces = GetNPartitionForces();

    // first build up the divisor
    long force, divisor = 1;
    for(force = 0; force < npartforces; ++force)
        divisor *= GetNPartitions(force);

    // now build the vector
    LongVec1d membership(npartforces);
    long remainder = xpartindex;
    for(force = 0; force < npartforces; ++force)
    {
        divisor /= GetNPartitions(force);
        membership[force] = remainder / divisor;
        remainder %= divisor;
    }

    return membership;
} // DataPack::GetBranchMembership

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

void DataPack::SetFinalPartitionCounts(LongVec1d pcounts)
{
    m_finalcount = pcounts;
} // DataPack::SetFinalPartitionCounts

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

void DataPack::SetRegion(Region* r)
{
    r->SetID(m_regions.size());
    m_regions.push_back(r);
} // DataPack::SetRegion

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

long DataPack::GetNMaxLoci() const
{
    long maxLoci = 0;
    vector<Region*>::const_iterator iter;
    for(iter=m_regions.begin(); iter != m_regions.end(); iter++)
    {
        long numLoci = (*iter)->GetNloci();
        if(numLoci > maxLoci)
        {
            maxLoci = numLoci;
        }
    }
    return maxLoci;
}

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

long DataPack::GetNumMovingLoci(long reg) const
{
    return GetRegion(reg).GetNumMovingLoci();
}

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

StringVec1d DataPack::GetAllRegionNames() const
{
    StringVec1d rnames;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
        rnames.push_back((*region)->GetRegionName());

    return rnames;
} // DataPack::GetAllRegionNames

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

StringVec2d DataPack::GetAllLociNames() const
{
    StringVec2d lnames;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
        lnames.push_back((*region)->GetAllLociNames());

    return lnames;
} // DataPack::GetAllLociNames

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

StringVec2d DataPack::GetAllLociDataTypes() const
{
    StringVec2d lnames;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
        lnames.push_back((*region)->GetAllLociDataTypes());

    return lnames;
} // DataPack::GetAllLociDataTypes

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

StringVec2d DataPack::GetAllLociMuRates() const
{
    StringVec2d lnames;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
        lnames.push_back((*region)->GetAllLociMuRates());

    return lnames;
} // DataPack::GetAllLociMuRates

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

StringVec1d DataPack::ToXML(unsigned long nspaces) const
{
    StringVec1d xmllines;
    string line = MakeIndent(MakeTag(xmlstr::XML_TAG_DATA),nspaces);
    xmllines.push_back(line);

    nspaces += INDENT_DEPTH;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        StringVec1d regxml((*region)->ToXML(nspaces));
        xmllines.insert(xmllines.end(),regxml.begin(),regxml.end());
    }
    nspaces -= INDENT_DEPTH;

    line = MakeIndent(MakeCloseTag(xmlstr::XML_TAG_DATA),nspaces);
    xmllines.push_back(line);

    return xmllines;
} // DataPack::ToXML

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

bool DataPack::CanHapArrange() const
{
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        if ((*region)->CanHaplotype()) return true;
    }
    return false;
} // DataPack::CanHapArrange

//------------------------------------------------------------------------------------
// True if you have at least two linked markers.

bool DataPack::CanMeasureRecombination() const
{
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        size_t nMarkersThisRegion = 0;

        long nLoci = (*region)->GetNloci();
        for(long i=0; i < nLoci ; i++)
        {
            const Locus & locusRef = (*region)->GetLocus(i);
            if(! locusRef.IsMovable())
            {
                long nMarkersThisLocus = locusRef.GetNmarkers();
                if(nMarkersThisLocus > 1)
                {
                    return true;
                }
                nMarkersThisRegion += nMarkersThisLocus;
            }
        }

        if (nMarkersThisRegion > 1 )
        {
            return true;
        }
    }
    return false;
} // DataPack::CanMeasureRecombination

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

bool DataPack::AnyMapping() const
{
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        if ((*region)->AnyMapping()) return true;
    }
    return false;
}

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

void DataPack::RemoveUneededPartitions()
{
    // use in Phase 1 only
    // of all partition forces, only force_MIG might have mentioned partitions
    // in the XML and then discovered there was only one partition and the
    // force was not, in fact, active.

    force_type ft = force_MIG;
    if (GetNPartitionsByForceType(ft) == 1)
    {
        vector<Region*>::const_iterator region;
        for(region = m_regions.begin(); region != m_regions.end(); ++region)
            (*region)->RemovePartitionFromLoci(ft);
        m_outputpopnames = m_partnames[ft];
        m_partnames.erase(ft);
    }
} // DataPack::RemoveUneededPartitions

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

DoubleVec1d DataPack::CountOverallBaseFreqs() const
{
    DoubleVec1d basecnts(BASES,0.0);
    double nbases(0.0);
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        DoubleVec1d regbasecnts((*region)->CountNBases());
        transform(basecnts.begin(),basecnts.end(),regbasecnts.begin(),
                  basecnts.begin(), plus<double>());
        double initial(0.0);
        nbases += accumulate(regbasecnts.begin(),regbasecnts.end(),initial);
    }

    DoubleVec1d basefreqs(BASES);
    transform(basecnts.begin(),basecnts.end(),basefreqs.begin(),
              bind2nd(divides<double>(),nbases));

    return basefreqs;
} // DataPack::CountOverallBaseFreqs

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

DoubleVec2d DataPack::CreateParamScalarVector() const
{
    DoubleVec2d pscalars;
    for (unsigned long reg=0; reg<m_regions.size(); reg++)
    {
        ForceParameters fp(reg);
        pscalars.push_back(fp.GetRegionalScalars());
    }

    return pscalars;
} // DataPack::CreateParamScalarVector

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

DoubleVec1d DataPack::GetRegionalPopSizeScalars() const
{
    DoubleVec1d pscalars;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
        pscalars.push_back((*region)->GetEffectivePopSize());

    return pscalars;
} // DataPack::GetRegionalPopSizeScalars

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

LongVec1d DataPack::GetNumAllLociPerRegion() const
{
    LongVec1d nloci;
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
        nloci.push_back((*region)->GetNumAllLoci());

    return nloci;
}

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

bool DataPack::HasMultiLocus() const
{
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        if ((*region)->GetNumAllLoci() > 1)
        {
            return true;
        }
    }

    return false;
}

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

bool DataPack::IsDuplicateRegionName(const string& newname) const
{
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        if ((*region)->GetRegionName() == newname) return true;
    }

    return false;
} // DataPack::IsDuplicateRegionName

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

Region* DataPack::GetRegionByName(const string& newname)
{
    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        if ((*region)->GetRegionName() == newname) return *region;
    }

    return NULL;
} // DataPack::IsDuplicateRegionName

//------------------------------------------------------------------------------------
// WriteFlucFile is only used when JSIM is defined (as of 3/30/06 --LS)

void DataPack::WriteFlucFile(const string& outputfilename) const
{
    ofstream ofile;
    ofile.open(outputfilename.c_str(),ios::out | ios::trunc);
    ofile << m_regions.size() << endl;
    ofile.close();

    vector<Region*>::const_iterator region;
    for(region = m_regions.begin(); region != m_regions.end(); ++region)
    {
        (*region)->WriteFlucFile(outputfilename,false);
    }
} // DataPack::WriteFlucFile

//------------------------------------------------------------------------------------
// doesn't handle spacing or loci
// WritePopulationXMLFiles used only when JSIM defined (3/30/06 --LS)

void DataPack::WritePopulationXMLFiles(bool separate_regions) const
{
    if (separate_regions)
    {
        vector<Region*>::const_iterator region;
        for(region = m_regions.begin(); region != m_regions.end(); ++region)
        {
            (*region)->WritePopulationXMLFiles();
        }
    }
    else
    {
        long xmlindent(INDENT_DEPTH);
        StringVec3d xmldata; // region X pop X data
        vector<Region*>::const_iterator region;
        for(region = m_regions.begin(); region != m_regions.end(); ++region)
        {
            xmldata.push_back((*region)->MakeByPopXML(xmlindent));
        }

        // write the all-populations file
        ofstream afile;
        afile.open("popall", ios::out | ios::trunc);
        afile << "<lamarc>" << endl;
        afile << "<data>" << endl;
        long reg, nregions(GetNRegions());
        for(reg = 0; reg < nregions; ++reg)
        {
            m_regions[reg]->WriteToXMLFileUsing(afile,xmldata[reg],true);
        }
        afile << "</data>" << endl;
#ifndef JSIM
        afile << "</lamarc>" << endl;
#endif
        afile.close();

        // write the single population files
        force_type ft = registry.GetForceSummary().GetNonLocalPartitionForceTag();
        long pop, npops = GetNPartitionsByForceType(ft);
        for(pop = 0; pop < npops; ++pop)
        {
            string fname("pop"+ToString(pop));
            ofstream ofile;
            ofile.open(fname.c_str(), ios::out | ios::trunc);

            ofile << "<lamarc>" << endl;
            ofile << "<data>" << endl;
            for(reg = 0; reg < nregions; ++reg)
            {
                m_regions[reg]->WriteToXMLFileUsing(ofile,xmldata[reg][pop]);
            }
            ofile << "</data>" << endl;
#ifndef JSIM
            ofile << "</lamarc>" << endl;
#endif
            ofile.close();
        }
    }
} // DataPack::WritePopulationXMLFiles

//____________________________________________________________________________________
