// $Id: derivatives.cpp,v 1.56 2018/01/03 21:33:02 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 */


// derivatives of parameter likelihood
// call sequence is in maximizer.cpp

#include <cassert>
#include <iostream>                     // debugging

#include "collector.h"
#include "likelihood.h"
#include "mathx.h"
#include "plforces.h"

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

#ifndef NDEBUG
//#define TEST
#endif // NDEBUG

using namespace std;

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

// Default implementation of a function to compute and return the following:
//
// sum_over_all_trees( (Prob(G|P)/Prob(G|Po)) * (d/dp)log(Prob(G|P)) )
// == sum_over_unique_trees(nG*(Prob(G|P)/Prob(G|Po))*(d/dp)log(Prob(G|P)))
//
// where d/dp   means partial derivative with respect to parameter p,
//       P      is the full vector of current parameter values ("param"),
//       Po     is the starting point ("driving values") vector ("param0"),
//       G      is one of the genealogies over which we're summing,
//  and  nG     is the number of identical copies of genealogy G.
//
// Input paramters:  param, ln_ProbGP_over_ProbGPo, pGenealogies, whichparam.
//
// This function can be used with or without Geyer replicates.  If Geyer replicates
// are being used, then input parameter ln_ProbGP_over_ProbGPo must have been
// computed using an effective "lnProbGPo" == "lnProb(G|param0),"
// where "lnProb(G|param0)" = log(sum_over_replicates((NtrR/wR)*Prob(G|param0R))),
//       NtrR    is the total number of (non-unique) trees in replicate R,
//       wR      is the Geyer weight assigned to replicate R,
//  and  param0R is the starting point ("driving value") for replicate R.
//
// This function returns the "sum_over_..." quantity spelled out at the top
// of this set of comments.
//
double PostLike::DCalc_sumG_BasicNumerator(const vector<double>& param, const vector<double>& ln_ProbGP_over_ProbGPo,
                                           const vector<TreeSummary*> *pGenealogies, const long& whichparam)
{
    double DlnPoint(0.0), DlnWait(0.0), termG(0.0), result(0.0);
    vector<PLForces*>::const_iterator theForce;
    long whichforce = FindForce(whichparam);
    theForce = m_forces.begin() + whichforce;

    for (unsigned long G = 0; G < pGenealogies->size(); G++)
    {
        DlnPoint = (*theForce)->DlnPoint(param, (*pGenealogies)[G], whichparam);
        DlnWait = (*theForce)->DlnWait(param, (*pGenealogies)[G], whichparam);
        termG = SafeProductWithExp(DlnPoint + DlnWait, ln_ProbGP_over_ProbGPo[G]);
        result += ((*pGenealogies)[G])->GetNCopies() * termG;

        // These should never be not-a-number.
        assert(!systemSpecificIsnan(ln_ProbGP_over_ProbGPo[G]));
        assert(!systemSpecificIsnan(DlnPoint));
        assert(!systemSpecificIsnan(DlnWait));
        //assert(!systemSpecificIsnan(termG)); // termG is OK if the others are OK
    }

    return result;
}

//------------------------------------------------------------------------------------
// Single-region derivatives
//
bool SinglePostLike::DCalculate(const vector<double>& param, vector<double>& gradient)
{
    assert(param.size() == gradient.size());
    double denominator = m_totalNumTrees * exp(m_last_lnL);
    // "denominator" is equal to
    // sum_over_all_trees(Prob(G|P)/Prob(G|Po)).

    if (0.0 == denominator)
    {
#ifdef TEST
        cerr << "SinglePostLike::DCalculate -- zero likelihood.  "
             << "Returning...." << endl;
#endif // TEST
        return false; // otherwise we'd divide by zero
    }

    vector<double>::const_iterator p; // param
    vector<ParamStatus>::const_iterator guideIt;
    vector<double>::iterator g; // gradient
    long whichparam = 0;
    for (p = param.begin(), g = gradient.begin(),
             guideIt = m_working_pstatusguides.begin();
         p != param.end(); ++p, ++g, ++whichparam, ++guideIt)
    {
        ParamStatus mystatus = *guideIt;
        if (!mystatus.Varies())
        {
            //guide is 0 if the gradient does not need to be calculated
            (*g) = 0.0;
        }
        else
        {
            (*g) = PostLike::DCalc_sumG_BasicNumerator(param, m_ln_ProbGP_over_ProbGPo,
                                                       &m_data->treeSummaries, whichparam);
            (*g) /= denominator; // defined at the top of this method
        }
        // 2004/03/19 erynes -- For every parameter "p" except linear
        // parameters like growth, we maximize with respect to log(p)
        // instead of with respect to p.
        // Not only does this make the math easier, it also keeps our
        // positive-definite parameters (e.g., theta) positive.
        // The "external" effect of maximizing with respect to log(p)
        // is visible only in Maximizer::SetParam(), where the step is applied
        // to the log parameters instead of the linear parameters,
        // and right here, where we multiply the gradient by the parameter.
        // This latter step comes from the chain rule:
        // if f(x1, x2, x3, ..., xN) = g(y1, y2, y3, ..., xN),
        // where y1 is a function of x1 alone, y2 = y2(x2), etc.
        // (note g can contain both y's and x's),
        // then if we use "d" to stand for "partial derivative," we have:
        //
        // g = f
        // dg/dy2 = df/dy2  = df/dx2 * dx2/dy2,
        // and similarly for y1, y3, etc.
        //
        // For yN = log(xN), we have xN = exp(yN),
        // and dxN/dyN = d(exp(yN))/dyN = exp(yN) = xN.
        // Substituting xN = dxN/dyN in the equation above, we have
        // dg/dy2 = df/dx2 * x2,
        // and similarly for y1, y3, and all log parameters generally.
        if (!isLinearParam(whichparam))
            (*g) *= (*p);
    }

    return true;
}

//------------------------------------------------------------------------------------
// Multi-replicate, single-region derivatives
//
bool ReplicatePostLike::DCalculate(const vector<double>& param, vector<double>& gradient)
{
    assert(param.size() == gradient.size());
    double denominator = exp(m_last_lnL); // reduce num. of calls to exp()

    if (0.0 == denominator)
    {
#ifdef TEST
        cerr << "ReplicatePostLike::DCalculate -- zero likelihood.  "
             << "Returning...." << endl;
#endif // TEST
        return false; // otherwise we'd divide by zero
    }

    vector<double>::const_iterator p; // param
    vector<ParamStatus>::const_iterator guideIt;
    vector<double>::iterator g; // gradient
    long whichparam = 0;

    for (p = param.begin(), g = gradient.begin(),
             guideIt = m_working_pstatusguides.begin();
         p != param.end(); ++p, ++g, ++whichparam, ++guideIt)
    {
        ParamStatus mystatus = *guideIt;
        (*g) = 0.0;
        // non-varying parameters are not computed
        if (mystatus.Varies())
        {
            for (unsigned long rep = 0; rep < m_nReplicate; rep++)
            {
                (*g) += PostLike::DCalc_sumG_BasicNumerator(param, m_ln_ProbGP_over_ProbGPo[rep],
                                                            &(m_data[rep]->treeSummaries), whichparam);
            }
            (*g) /= denominator; // defined at the top of this method
            if (!isLinearParam(whichparam))
                (*g) *= (*p); // see comment for log params in SinglePostLike::DCalculate()
        }
    }
    return true;
}

//------------------------------------------------------------------------------------
// Multi-region, multi-replicate derivatives
//
bool RegionPostLike::DCalculate(const vector<double>& param, vector<double>& gradient)
{
    assert(param.size() == gradient.size());

    DoubleVec2d scaledparams = CreateVec2d(m_nRegions,param.size(),0.0);
    for(unsigned long region = 0; region < m_nRegions; ++region)
    {
        transform(param.begin(),param.end(),
                  m_paramscalars[region].begin(),scaledparams[region].begin(),
                  multiplies<double>());
    }
    assert(scaledparams.size() == m_data.size());

    vector<ParamStatus>::const_iterator guideIt;
    vector<double>::iterator g;
    long whichparam = 0;
    double numeratorForThisRegion(0.0);

    for (g = gradient.begin(),
             guideIt = m_working_pstatusguides.begin();
         g != gradient.end(); ++g, ++whichparam, ++guideIt)
    {
        ParamStatus mystatus = *guideIt;
        (*g) = 0.0;
        //guide is 0 if the gradient does not need to be calculated
        if (mystatus.Varies())
        {
            for (unsigned long reg = 0; reg < m_lnProbGPo.size(); reg++)
            {
                numeratorForThisRegion = 0.0;
                for (unsigned long rep = 0; rep < m_lnProbGPo[reg].size(); rep++)
                    numeratorForThisRegion +=
                        PostLike::DCalc_sumG_BasicNumerator(scaledparams[reg],
                                                            m_ln_ProbGP_over_ProbGPo[reg][rep],
                                                            &(m_data[reg][rep]->treeSummaries),
                                                            whichparam);
                if (!isLinearParam(whichparam))
                    numeratorForThisRegion *= scaledparams[reg][whichparam];
                // see comment for log params in SinglePostLike::DCalculate()

                if (0.0 == m_sumG_ProbGP_over_ProbGPo[reg])
                {
#ifdef TEST
                    cerr << "RegionPostLike::DCalculate -- zero likelihood for region "
                         << reg << ".  Returning...." << endl;
#endif // TEST
                    return false; // otherwise we'd divide by zero
                }
                (*g) += numeratorForThisRegion / m_sumG_ProbGP_over_ProbGPo[reg];
            }
        }
    }

    return true;
}

//------------------------------------------------------------------------------------
// Multi-region, multi-replicate derivatives,
// with the background mutation rates gamma-distributed over regions.
//
bool GammaRegionPostLike::DCalculate(const vector<double>& param, vector<double>& gradient)
{
    if (param.size() != gradient.size())
    {
        assert(false);
        string msg = "GammaRegionPostLike::DCalculate(), received two ";
        msg += "vectors of unequal size:  param.size() = ";
        msg += ToString(param.size()) + " and gradient.size() = ";
        msg += ToString(gradient.size()) + ".";
        throw implementation_error(msg);
    }

    DoubleVec2d scaledparams = CreateVec2d(m_nRegions,
                                           m_paramscalars[0].size(),0.0);
    for(unsigned long region = 0; region < m_nRegions; ++region)
    {
        transform(param.begin(), param.end(), m_paramscalars[region].begin(),
                  scaledparams[region].begin(), multiplies<double>());
    }
    assert(scaledparams.size() == m_data.size());

    vector<ParamStatus>::const_iterator guideIt;
    vector<double>::iterator g;
    bool retval(true);
    double DlnPoint(0.0), DlnWait(0.0), termRegRepG(0.0), numeratorForThisRegion(0.0);
    vector<PLForces*>::const_iterator theForce;
    unsigned long whichparam(0), whichforce(0), the_alpha_param(param.size() - 1);
    double alpha = param[the_alpha_param], SqrtAlpha = sqrt(alpha);

    for (g = gradient.begin(),
             guideIt = m_working_pstatusguides.begin();
         g != gradient.end(); ++g, ++whichparam, ++guideIt)
    {
        ParamStatus mystatus = *guideIt;
        (*g) = 0.0;
        //guide is 0 if the gradient does not need to be calculated
        if (!mystatus.Varies()) continue;

        if (whichparam != the_alpha_param)
        {
            whichforce = FindForce(whichparam);
            theForce = m_forces.begin() + whichforce;
            for (unsigned long reg = 0; reg < m_lnProbGPo.size(); reg++)
            {
                numeratorForThisRegion = 0.0;
                for (unsigned long rep = 0; rep < m_lnProbGPo[reg].size(); rep++)
                {
                    const vector<TreeSummary*> *pGenealogies =
                        &(m_data[reg][rep]->treeSummaries);
                    for (unsigned long G = 0; G < pGenealogies->size(); G++)
                    {
                        DlnPoint = (*theForce)->DlnPoint(scaledparams[reg],
                                                         (*pGenealogies)[G], whichparam);
                        DlnWait = (*theForce)->DlnWait(scaledparams[reg],
                                                       (*pGenealogies)[G], whichparam);
                        termRegRepG = m_K_alpha_minus_nevents_plus_1[reg][rep][G] /
                            m_K_alpha_minus_nevents[reg][rep][G];
                        termRegRepG *= sqrt(alpha/(-m_lnWait[reg][rep][G]));
                        termRegRepG += (alpha - m_nevents_G[reg][rep][G]) /
                            m_lnWait[reg][rep][G]; // note:  not -m_lnWait
                        termRegRepG *= DlnWait;
                        termRegRepG += DlnPoint;
                        termRegRepG *= m_C[reg][rep][G];

                        numeratorForThisRegion += termRegRepG
                            * ((*pGenealogies)[G])->GetNCopies();
                    }
                }
                if (0.0 == m_sumG_C[reg])
                {
#ifdef TEST
                    cerr << "GammaRegionPostLike::DCalculate(), sum over reps and trees of region "
                         << reg << " is " << m_sumG_C[reg] << "; this must be a positive number.  "
                         << "Returning false and a gradient component of DBL_BIG....";
#endif // TEST
                    (*g) = DBL_BIG;
                    retval = false;
                    continue;
                }

                if (isLinearParam(whichparam))
                {
                    // Growth is the lone linear parameter at present.
                    // But at present, growth and gamma can't mix.
                    // If the user asked for both growth and gamma,
                    // this error was almost surely handled earlier,
                    // but we need to throw here, just to be safe.
                    throw implementation_error("Unable to co-estimate growth and alpha.");
                }
                numeratorForThisRegion *= scaledparams[reg][whichparam];
                // see comment for log params in SinglePostLike::DCalculate()
                (*g) += numeratorForThisRegion / m_sumG_C[reg];
            }
        }

        else // partial derivative with respect to alpha
        {
            for (unsigned long reg = 0; reg < m_lnProbGPo.size(); reg++)
            {
                numeratorForThisRegion = 0.0;
                for (unsigned long rep = 0; rep < m_lnProbGPo[reg].size(); rep++)
                {
                    const vector<TreeSummary*> *pGenealogies =
                        &(m_data[reg][rep]->treeSummaries);

                    for (unsigned long G = 0; G < pGenealogies->size(); G++)
                    {
                        double SqrtNegLogWait = sqrt(-m_lnWait[reg][rep][G]);
                        termRegRepG = m_K_alpha_minus_nevents_plus_1[reg][rep][G];
                        termRegRepG *= -SqrtNegLogWait / SqrtAlpha;
                        termRegRepG += DvBesselK(alpha - m_nevents_G[reg][rep][G],
                                                 2.0*SqrtNegLogWait*SqrtAlpha,
                                                 m_K_alpha_minus_nevents[reg][rep][G]);
                        termRegRepG /= m_K_alpha_minus_nevents[reg][rep][G];
                        termRegRepG += 1.0 + log(SqrtNegLogWait*SqrtAlpha);
                        termRegRepG *= m_C[reg][rep][G];

                        numeratorForThisRegion += termRegRepG
                            * ((*pGenealogies)[G])->GetNCopies();
                    }
                }
                if (0.0 == m_sumG_C[reg])
                {
#ifdef TEST
                    cerr << "GammaRegionPostLike::DCalculate(), sum over reps and trees of region "
                         << reg << " is " << m_sumG_C[reg] << "; this must be a positive number.  "
                         << "Returning false and a gradient component of DBL_BIG....";
#endif // TEST
                    (*g) = DBL_BIG;
                    retval = false;
                    continue;
                }

                (*g) += numeratorForThisRegion / m_sumG_C[reg]; // psi() added later
            }

            (*g) -= m_nRegions * psi(alpha);
            (*g) *= alpha; // Alpha is treated as a log parameter to ensure it's kept
            // positive.  See comment in SinglePostLike::DCalculate().
        }
    }

    return retval; // true unless m_sumG_C[reg] is zero
}

//____________________________________________________________________________________
