// $Id: calculators.cpp,v 1.29 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>
#include <numeric>
#include <map>

#include "calculators.h"
#include "datapack.h"
#include "dlmodel.h"
#include "locus.h"
#include "mathx.h" //for ScaleToSumToOne
#include "region.h"
#include "registry.h"

using std::map;

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

DoubleVec1d FrequenciesFromData(long int regionId, long int locusId, model_type modelType)
{
    const Region& thisRegion = registry.GetDataPack().GetRegion(regionId);
    const Locus& thisLocus = thisRegion.GetLocus(locusId);
    return FrequenciesFromData(thisLocus,modelType);
}

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

DoubleVec1d FrequenciesFromData(const Locus& locus, model_type modelType)
{
    assert(modelType == F84 || modelType == GTR);
    DoubleVec1d basefreqs(BASES,0.0);
    double totalmarkers = 0.0;
    const vector<TipData>& tips = locus.GetTipData();

    vector<TipData>::const_iterator tit;
    for (tit = tips.begin(); tit != tips.end(); ++tit)
    {
        StringVec1d data = tit->data;
        StringVec1d::const_iterator sit;

        for (sit = data.begin(); sit != data.end(); ++sit)
        {
            double errorRate= locus.GetPerBaseErrorRate();
            DoubleVec1d site = NucModel::StaticDataToLikes(*sit,errorRate);
            double total = site[baseA]+site[baseC]+site[baseG]+site[baseT];
            basefreqs[baseA] += site[baseA]/total;
            basefreqs[baseC] += site[baseC]/total;
            basefreqs[baseG] += site[baseG]/total;
            basefreqs[baseT] += site[baseT]/total;
        }
        totalmarkers += locus.GetNmarkers();
    }

    basefreqs[baseA] /= totalmarkers;
    basefreqs[baseC] /= totalmarkers;
    basefreqs[baseG] /= totalmarkers;
    basefreqs[baseT] /= totalmarkers;
    ScaleToSumToOne(basefreqs);

    // check for any zero freqs, try global freqs, then use DBL_EPSILON set
    double zero(0.0);
    if (std::count(basefreqs.begin(),basefreqs.end(),zero))
    {
        basefreqs = registry.GetDataPack().CountOverallBaseFreqs();
        ScaleToSumToOne(basefreqs);
        if (std::count(basefreqs.begin(),basefreqs.end(),zero))
        {
            DoubleVec1d::iterator baseit;
            for (baseit = basefreqs.begin(); baseit != basefreqs.end(); ++baseit)
            {
                if (*baseit == 0.0) *baseit = defaults::minLegalFrequency;
            }
        }
    }

    return basefreqs;
}

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

// MDEBUG LOCUS Function not currently used but being saved for
// multi-locus code
// EWFIX.P3 TEST -- IMPORTANT -- this function may contain code that sets
// the estimate to a default value when we don't want to. need
// to see if we can avoid that and instead post-process unsuitable
// values to correct them when appropriate. for example, if using
// this function to get an initial estimate, we may want to default
// to a reasonable value, but if reporting the value we do not.

void MigFSTLocus(const DataPack&, const Locus& loc, long int numMigs,
                 DoubleVec1d& estimate, std::deque<bool>& isCalculated)
{
    estimate.clear();
    estimate = DoubleVec1d(numMigs,0.0);
    assert(isCalculated.size() == estimate.size());
    DoubleVec1d fw = loc.GetDataTypePtr()->CalcFw(loc);
    DoubleVec1d fb = loc.GetDataTypePtr()->CalcFb(loc);

    long int index1, index2;
    long int pop1, pop2, npops = fw.size();
    assert(numMigs ==  (npops * npops));

    for(pop1 = 0; pop1 < npops; ++pop1)
    {
        for(pop2 = pop1+1; pop2 < npops; ++pop2)
        {
            index1 = pop1 * npops + pop2;  // position in linear vector
            double fb12 = fb[index1];
            double denom = -2.0 * fb12 + fw[pop1] + fw[pop2];
            if (denom == 0.0)
            {
                isCalculated[index1] = false;
                estimate[index1] = defaults::migration;
            }
            else estimate[index1] = 2.0 * fb12 / denom;
            if (estimate[index1] < 0.0)
            {
                isCalculated[index1] = false;
                estimate[index1] = defaults::migration;
            }
            index2 = pop2 * npops + pop1;  // lower triangle
            estimate[index2] = estimate[index1];
        }
    }
} // MigFSTLocus

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

void MigFSTMultiregion(const DataPack & dpack, const DoubleVec2d & regionalmuratios,
                       DoubleVec1d & perRegionMigsAccumulator, std::deque<bool> & isCalculated)
{
    // MDEBUG this is disabled for force_DIVMIG currently.  It should be enabled
    // eventually, contemporary populations only.

    // get MIG npop and DIVMIG npop
    long int migpopCount  = dpack.GetNPartitionsByForceType(force_MIG);
    long int divmigpopCount  = dpack.GetNPartitionsByForceType(force_DIVMIG);

    if (migpopCount == 0 && divmigpopCount <= 1)
        throw data_error("Data error--zero populations?  There should be at least one.");

    if (migpopCount <= 1) return;

    long int popCount = migpopCount;

    isCalculated.clear();
    perRegionMigsAccumulator.clear();
    long int numMigs = popCount * popCount;
    perRegionMigsAccumulator = DoubleVec1d(numMigs,0.0);
    isCalculated = std::deque<bool>(numMigs,true);

    const long int numRegions = dpack.GetNRegions();

    for(long int regionIndex = 0; regionIndex < numRegions; regionIndex++)
    {
        const Region& region = dpack.GetRegion(regionIndex);
        long int numLoci = region.GetNloci();
        DoubleVec1d perLociMigsAccumulator(numMigs,0.0);

        for(long int locusIndex = 0; locusIndex < numLoci; locusIndex++)
        {
            const Locus& locus = region.GetLocus(locusIndex);
            DoubleVec1d migsForOneLocus;
            MigFSTLocus(dpack,locus,numMigs,migsForOneLocus,isCalculated);

            // rescale for relative mutation rate
            std::transform(migsForOneLocus.begin(),
                           migsForOneLocus.end(),
                           migsForOneLocus.begin(),
                           std::bind2nd(std::multiplies<double>(),
                                        regionalmuratios[regionIndex][locusIndex]));

            // and add to the totals by region
            std::transform(  migsForOneLocus.begin(),
                             migsForOneLocus.end(),
                             perLociMigsAccumulator.begin(),
                             perLociMigsAccumulator.begin(),
                             std::plus<double>());
        }
        // average over the per-locus data
        std::transform(  perLociMigsAccumulator.begin(),
                         perLociMigsAccumulator.end(),
                         perLociMigsAccumulator.begin(),
                         std::bind2nd(std::divides<double>(),static_cast<double>(numLoci)));

        // and add to the totals by region
        std::transform(  perLociMigsAccumulator.begin(),
                         perLociMigsAccumulator.end(),
                         perRegionMigsAccumulator.begin(),
                         perRegionMigsAccumulator.begin(),
                         std::plus<double>());
    }
    // finally, average over the per-region data
    std::transform(  perRegionMigsAccumulator.begin(),
                     perRegionMigsAccumulator.end(),
                     perRegionMigsAccumulator.begin(),
                     std::bind2nd(std::divides<double>(),static_cast<double>(numRegions)));
}

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

// EWFIX.P3 TEST -- IMPORTANT -- this function may contain code that sets
// the estimate to a default value when we don't want to. need
// to see if we can avoid that and instead post-process unsuitable
// values to correct them when appropriate. for example, if using
// this function to get an initial estimate, we may want to default
// to a reasonable value, but if reporting the value we do not.

void ThetaFSTLocus(const DataPack&, const Locus& loc, const long int numThetas,
                   vector<map<force_type,string> > tipids, DoubleVec1d& estimates,
                   std::deque<bool>& isCalculated)
{
    assert(numThetas > 1);

    StringVec3d data;
    for(long int xpart = 0; xpart < numThetas; xpart++)
    {
        data.push_back(loc.GetCrossPartitionGeneticData(tipids[xpart]));
    }

    DoubleVec1d fw = loc.GetDataTypePtr()->CalcXFw(loc,data);
    DoubleVec1d fb = loc.GetDataTypePtr()->CalcXFb(loc,data);

    long int index1;
    assert(fw.size() == static_cast<unsigned long int>(numThetas));

    DoubleVec1d thetaEstimates(numThetas,defaults::theta);

    for(long int xpart = 0; xpart < numThetas; xpart++)
    {
        double estimate = 0.0;
        for(long int otherpart = 0; otherpart < numThetas; ++otherpart)
        {
            if (otherpart == xpart) continue;
            index1 = xpart * numThetas + otherpart;  // position in linear vector
            double fb12 = fb[index1];
            double numer = -2.0 * fb12 + fw[xpart] + fw[otherpart];
            double denom = -2.0 * fb12 *fb12 + fw[xpart] + fw[otherpart] + fw[xpart] * fw[xpart];

            if (denom != 0.0)
            {
                estimate += (numer * (1.0 - fw[xpart])) / denom;
            }
        }

        if (estimate > 0.0)
        {
            thetaEstimates[xpart] = estimate / numThetas;
        }
        else
        {
            isCalculated[xpart] = false;
        }

    }
    estimates = thetaEstimates;
} // ThetaFSTLocus

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

void ThetaFSTMultiregion(const DataPack& dpack, const DoubleVec1d & regionalthetascalars,
                         const DoubleVec2d & regionalmuratios, DoubleVec1d & estimates,
                         std::deque<bool> & isCalculated)
{
    const long int numRegions = dpack.GetNRegions();
    const long int numThetas  = dpack.GetNCrossPartitions();
    vector<map<force_type,string> > tipids = dpack.GetCrossPartitionIds(true);
    isCalculated.clear();
    isCalculated = std::deque<bool>(numThetas,true);

    DoubleVec1d OverallThetasAccumulator(numThetas,0.0);

    for(long int regionIndex = 0; regionIndex < numRegions; regionIndex++)
    {
        const Region& region = dpack.GetRegion(regionIndex);
        long int numLoci = region.GetNloci();
        DoubleVec1d RegionalThetasAccumulator(numThetas,0.0);

        for(long int locusIndex = 0; locusIndex < numLoci; locusIndex++)
        {
            const Locus& locus = region.GetLocus(locusIndex);
            DoubleVec1d thetasForOneLocus;
            ThetaFSTLocus(dpack,locus,numThetas,tipids,thetasForOneLocus,isCalculated);
            // correct for varying mutation rates
            std::transform(thetasForOneLocus.begin(),
                           thetasForOneLocus.end(),
                           thetasForOneLocus.begin(),
                           std::bind2nd(std::divides<double>(),
                                        regionalmuratios[regionIndex][locusIndex]));

            // and add to the totals by locus
            std::transform(  thetasForOneLocus.begin(),
                             thetasForOneLocus.end(),
                             RegionalThetasAccumulator.begin(),
                             RegionalThetasAccumulator.begin(),
                             std::plus<double>());
        }

        // average over the per-locus data; while scaling appropiately
        // for the current region
        std::transform(  RegionalThetasAccumulator.begin(),
                         RegionalThetasAccumulator.end(),
                         RegionalThetasAccumulator.begin(),
                         std::bind2nd(std::divides<double>(),regionalthetascalars[regionIndex]*numLoci));

        // and add to the totals by region
        std::transform(  RegionalThetasAccumulator.begin(),
                         RegionalThetasAccumulator.end(),
                         OverallThetasAccumulator.begin(),
                         OverallThetasAccumulator.begin(),
                         std::plus<double>());
    }
    // finally, average over the per-region data
    std::transform(  OverallThetasAccumulator.begin(),
                     OverallThetasAccumulator.end(),
                     OverallThetasAccumulator.begin(),
                     std::bind2nd(std::divides<double>(),static_cast<double>(numRegions)));
    estimates = OverallThetasAccumulator;

} // ThetaFSTMultiregion

//------------------------------------------------------------------------------------
// EWFIX.P3 TEST -- IMPORTANT -- this function may contain code that sets the estimate to a default value
// when we don't want to.  Need to see if we can avoid that and instead post-process unsuitable values
// to correct them when appropriate. for example, if using this function to get an initial estimate,
// we may want to default to a reasonable value, but if reporting the value we do not.

void ThetaWattersonLocus(const DataPack&, const Locus& locus, const long int numThetas,
                         const vector<map<force_type,string> > tipids, DoubleVec1d& estimates,
                         std::deque<bool>& isCalculated)
{
    const long int nsites = locus.GetNsites();
    DoubleVec1d thetaEstimates(numThetas,defaults::theta);

    for(long int thetaIndex = 0; thetaIndex < numThetas; thetaIndex++)
    {
        long int nvarmarkers = 0;
        const StringVec2d data = locus.GetCrossPartitionGeneticData(tipids[thetaIndex]);
        if (!data.empty())
        {
            nvarmarkers = locus.GetDataTypePtr()->CalcNVarMarkers(data);
        }

        const long int nseqs = locus.GetNTips(tipids[thetaIndex]);

        if(nseqs > 1)
        {
            double harmonic = 0.0;
            for(long int seq = 1; seq < nseqs; seq++)
            {
                harmonic += 1.0/seq;
            }
            thetaEstimates[thetaIndex] = nvarmarkers / (nsites * harmonic);
        }
        else
        {
            isCalculated[thetaIndex] = false;
        }
    }
    estimates = thetaEstimates;
} // ThetaWattersonLocus

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

void ThetaWattersonMultiregion(const DataPack& dpack,
                               const DoubleVec1d& regionalthetascalars,
                               const DoubleVec2d& regionalmuratios,
                               DoubleVec1d& estimates,
                               std::deque<bool>& isCalculated)
{
    const long int numRegions = dpack.GetNRegions();
    const long int numThetas  = dpack.GetNCrossPartitions();
    isCalculated.assign(numThetas,true);

    DoubleVec1d OverallThetasAccumulator(numThetas,0.0);

    for(long int regionIndex = 0; regionIndex < numRegions; regionIndex++)
    {
        const Region& region = dpack.GetRegion(regionIndex);
        DoubleVec1d RegionalThetasAccumulator;
        ThetaWattersonRegion(dpack,region,regionalthetascalars[regionIndex],
                             regionalmuratios[regionIndex],RegionalThetasAccumulator,isCalculated);

        // and add to the totals by region
        std::transform(  RegionalThetasAccumulator.begin(),
                         RegionalThetasAccumulator.end(),
                         OverallThetasAccumulator.begin(),
                         OverallThetasAccumulator.begin(),
                         std::plus<double>());
    }
    // finally, average over the per-region data
    std::transform(  OverallThetasAccumulator.begin(),
                     OverallThetasAccumulator.end(),
                     OverallThetasAccumulator.begin(),
                     std::bind2nd(std::divides<double>(),static_cast<double>(numRegions)));
    estimates = OverallThetasAccumulator;

} // ThetaWattersonMultiregion

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

void ThetaWattersonRegion(const DataPack& dpack,
                          const Region& region,
                          double thetascalar,
                          const DoubleVec1d& muratios,
                          DoubleVec1d& estimates,
                          std::deque<bool>& isCalculated)
{
    long int numLoci = region.GetNloci();
    long int numThetas = dpack.GetNCrossPartitions();
    DoubleVec1d ThetasAccumulator(numThetas,0.0);
    vector<map<force_type,string> > tipids = dpack.GetCrossPartitionIds(true);

    for(long int locusIndex = 0; locusIndex < numLoci; locusIndex++)
    {
        const Locus& locus = region.GetLocus(locusIndex);
        DoubleVec1d thetasForOneLocus;
        ThetaWattersonLocus(dpack,locus,numThetas,tipids,thetasForOneLocus,isCalculated);

        // rescale the locus values for their relative mutation rates
        std::transform(thetasForOneLocus.begin(),
                       thetasForOneLocus.end(),
                       thetasForOneLocus.begin(),
                       std::bind2nd(std::divides<double>(),muratios[locusIndex]));

        // and add to the totals by region
        std::transform(  thetasForOneLocus.begin(),
                         thetasForOneLocus.end(),
                         ThetasAccumulator.begin(),
                         ThetasAccumulator.begin(),
                         std::plus<double>());
    }
    // average over the per-locus data; while scaling appropriately
    // for the current region
    std::transform(  ThetasAccumulator.begin(),
                     ThetasAccumulator.end(),
                     ThetasAccumulator.begin(),
                     std::bind2nd(std::divides<double>(),thetascalar*numLoci));

    estimates = ThetasAccumulator;
} // ThetaWattersonRegion

//____________________________________________________________________________________
