/*
    Theseus - maximum likelihood superpositioning of macromolecular structures

    Copyright (C) 2004-2015 Douglas L. Theobald

    This program is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 2 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program; if not, write to the:

    Free Software Foundation, Inc.,
    59 Temple Place, Suite 330,
    Boston, MA  02111-1307  USA

    -/_|:|_|_\-
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <float.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include "statistics.h"
#include "DLTmath.h"
#include "negbinom_dist.h"

/*
   The negative binomial distribution has the form

   p(x) dx = (Gamma(a+x)/(Gamma(a) Gamma(x+1))) p^(a) (1-p)^(x) dx
*/


double
negbinom_dev(const double a, const double p, const gsl_rng *r2)
{
    int             nless, nmore;

    nless = nmore = 0;
    do
    {
        if(gsl_rng_uniform(r2) <= p)
            ++nless;
        else
            ++nmore;
    }
    while(nless < a);

    return((double) nmore);
}


double
negbinom_pdf(const double x, const double a, const double p)
{
    double          gax, ga, gx;

    if (p < 0.0 || p > 1.0)
    {
        return(0.0);
    }
    else 
    {
        gax = lgamma(a+x);
        ga = lgamma(a);
        gx = lgamma(x+1);

        return(exp(gax - ga - gx) * pow (p, a) * pow(1.0-p, x));
    }
}


double
negbinom_lnpdf(const double x, const double a, const double p)
{
    double          gax, ga, gx;

    if (p < 0.0 || p > 1.0)
    {
        return(-1.0);
    }
    else 
    {
        gax = lgamma(a+x);
        ga = lgamma(a);
        gx = lgamma(x+1);

        return(gax - ga - gx + a*log(p) + x*log(1.0-p));
    }
}


double
negbinom_cdf(const double x, const double a, const double p)
{
    return(InBeta(a, x+1, p));
}


double
negbinom_sdf(const double x, const double a, const double p)
{
    return(1.0 - negbinom_cdf(x, a, p));
}


double
negbinom_int(const double x, const double y, const double a, const double p)
{
    return(InBeta(a, y+1, p) - InBeta(a, x+1, p));
}


double
negbinom_logL(const double a, const double p)
{
    return(0.0); /* DLT debug */
}


static void
psigamma_sums(const int x, const double a, double *psigd, double *psigdf)
{
    int             j;
    double          tmp;

    if (x == 0.0)
    {
        *psigd = 1.0;
        *psigdf = 1.0;
        return;
    }
    else
    {
        *psigd = *psigdf = 0.0;
        for (j = 1; j <= x; ++j)
        {
            tmp = 1.0 / (a+x-j);
            *psigd += tmp;
            *psigdf -= tmp*tmp;
        }
    }

/*     *psigd = *psigdf = 0.0; */
/*     for (j = 0; j <= x-1; ++j) */
/*     { */
/*         tmp = 1.0 / (a+j); */
/*         *psigd += tmp; */
/*         *psigdf -= tmp*tmp; */
/*     } */
}


/* Maximum likelihood fit. */
/* Adapted from:

       Walter W. Piegorsch (1990)
       "Maximum likelihood estimation for the negative binomial dispersion
       parameter."
       Biometrics 46:863-867.

   which uses a mathematical trick to avoid calls to Digamma() or Polygamma().
*/
static void
evalnegbinomML(const double *x, const int num, const double a, const double p,
               const double ave, double *fx, double *dfx)
{
/*     int             i; */
/*     double          sum; */
/*  */
/*     sum = 0.0; */
/*     for (i = 0; i < num; ++i) */
/*         sum += gsl_sf_psi(a+x[i]); */
/*  */
/*     *fx = sum - num*gsl_sf_psi(a) + num*log(p); */
/* printf("\n% 10.6e % 10.6e ", sum - num*gsl_sf_psi(a), num*log(p)); */
/*     sum = 0.0; */
/*     for (i = 0; i < num; ++i) */
/*         sum += gsl_sf_psi(a+x[i]); */
/*  */
/*     *dfx = sum - num*gsl_sf_psi(a) + num*((1.0/a) - (1/(a+ave))); */

    int             i;
    double          psigd, psigdf, psigd_sum, psigdf_sum;

    psigd_sum = psigdf_sum = 0.0;
    for (i = 0; i < num; ++i)
    {
        psigamma_sums((int) x[i], a, &psigd, &psigdf);
        psigd_sum += psigd;
        psigdf_sum += psigdf;
    }

    *fx = psigd_sum + num*log(p);
    *dfx = psigdf_sum + num*((1.0/a) - (1/(a+ave)));
}


/* Maximum likelihood fit of data to a negbinom distribution */
double
negbinom_fit(const double *x, const int num, double *a, double *p, double *logL)
{
    int             i;
    double          ave, var, amme, pmme, chi2;
    double          fx = 0.0, dfx = 1.0, fxdfx = 0.0;
    double         *data = malloc(num * sizeof(double));
    int             maxit = 100;
    double          tol = 1e-8;

    for (i = 0; i < num; ++i)
        data[i] = round(x[i]);

    /* Method of moments initial guess at shape parameters.
       See _Statistical Distributions_ 3rd ed.
       Evans, Hastings, and Peacock, p 141.
       Mean and variance rearranged. */
    ave = average(data, num);
    var = variance(data, num, ave);

    /* MMEs */
    *p = pmme = ave / var;

    if (*p <= 0.0)
        *p = 0.1;
    else if (*p >= 1.0)
        *p = 0.9;

    *a = amme = ave * ave / (var - ave);

    /* printf("\n% 10.6e % 10.6e % 10.6e", *a, *p, (*a / (*a + ave))); */

    /* make sure they aren't non-positive */
    if (*a <= 0.0)
        *a = DBL_EPSILON;

    if (*p <= 0.0)
        *p = DBL_EPSILON;

    for (i = 0; i < maxit; ++i)
    {
/*         printf("\n% 10.6e % 10.6e % 10.6e % 10.6e % 10.6e", */
/*                *a, *p, fx, dfx, fx/dfx); */
/*         fflush(NULL); */

        evalnegbinomML(data, num, *a, *p, ave, &fx, &dfx);

        fxdfx = fx/dfx;

        if (fabs(fx) < tol && fabs(fxdfx) < tol)
            break; /* success */

        *a -= fxdfx; /* Newton-Raphson correction */

        if (*a <= 0.0)
            *a = 0.1;

        /* ML estimate of p */
        *p = (*a / (*a + ave));
        /* *a = round(*a); */
    }

    if (i == maxit)
    {
        *a = amme;
        *p = pmme;
    }

    *a = round(*a);
    *p = (*a / (*a + ave));

    /* printf("\n\nnegbinom logL %e\n", negbinom_logL(*a, *p)); */

    chi2 = chi_sqr_adapt(data, num, 0, logL, *a, *p, negbinom_pdf, negbinom_lnpdf, negbinom_int);

    free(data);

    return(chi2);
}
