/*
    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 <math.h>
#include <float.h>
#include <stdlib.h>
#include <gsl/gsl_sf_gamma.h>
#include <gsl/gsl_sf_psi.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include "statistics.h"
#include "DLTmath.h"
#include "binomial_dist.h"

/*
   The binomial distribution has the form

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

static double
geometric_dev(const double p, const double nullv, const gsl_rng *r2)
{
    return(log(gsl_rng_uniform(r2)) / log (1.0 - p) - 1.0);
}


double
binomial_dev(const double n, const double p, const gsl_rng *r2)
{
    int             i;

    if (p > 0.01)
    {
        int         nless = 0;
        for (i = 0; i < n; ++i)
            if(genrand_real2() < p)
                ++nless;

        return((double) nless);
    }
    else
    {
        double      sum = 0.0;

        i = 0;
        do
        {
            ++i;
            sum += geometric_dev(p, 0, r2);
        }
        while(sum <= n - i);

        return((double) i);
    }
}


double
binomial_pdf(const double x, const double n, const double p)
{
    if (p <= 0.0 || p > 1.0)
        return(0.0);
    else if (x == 0.0)
        return(pow(1.0 - p, n));
    else 
        /* return(Combination(n, x) * pow(p, x) * pow(1.0 - p, n - x)); */
        return(exp(binomial_lnpdf(x, n, p)));
}


double
binomial_lnpdf(const double x, const double n, const double p)
{
    if (p <= 0.0 || p > 1.0)
        return(INFINITY);
    else if (x > n)
        return(INFINITY);
    else if (x == 0.0)
        return(n * log(1.0 - p));
    else 
        return(lgamma(n+1) - lgamma(n-x+1) - lgamma(x+1) + x*log(p) + (n-x)*log(1.0-p));
}


double
binomial_cdf(const double x, const double n, const double p)
{
/*     double sum; */
/*     int i; */

    if (x == n)
        return(1.0);

/*     sum = 0.0; */
/*     if (x > n) */
/*     { */
/*         for (i = 0; i <= n; ++i) */
/*             sum += binomial_pdf(i, n, p); */
/*     } */
/*     else */
/*     { */
/*         for (i = 0; i <= (int)x; ++i) */
/*             sum += binomial_pdf(i, (int)n, p); */
/*     } */
/*  */
/* printf("\n####%f %f", sum, InBeta((int)n - (int)x, (int)x + 1, 1.0 - p)); */

    /* return(InBeta((int)n - (int)x, (int)x + 1, 1.0 - p)); */
/*     printf("\nBeta:::: %f %f %f\n", */
/*            InBeta((int)n - (int)x, (int)x + 1, 1.0 - p), */
/*            InBeta(n - x, x + 1.0, 1.0 - p), */
/*            integrate_qsimp(binomial_pdf, n, p, 0.0, x)); */
    return(gsl_sf_beta_inc(n - x, x + 1.0, 1.0 - p));
}


double
binomial_sdf(const double x, const double n, const double p)
{
    return(1.0 - binomial_cdf(x, n, p));
}


double
binomial_int(const double x, const double y, const double n, const double p)
{
    return(binomial_cdf(y, n, p) - binomial_cdf(x, n, p));
}


double
binomial_logL(const double n, const double p)
{
    return(0.0); /* DLT debug */
}


/* For integer x and n, this calculates:

        Digamma(n+1)/Digamma(n+1-x)
    and
        Trigamma(n+1)/Digamma(n+1-x)

    using a mathematical trick to avoid calls to Digamma() or Polygamma()

    Inspiration from:

        Walter W. Piegorsch (1990)
        "Maximum likelihood estimation for the negative binomial dispersion
        parameter."
        Biometrics 46:863-867.
*/
void
psigamma_sums(const int x, const double n, double *psigd, double *psigdf)
{
    int             i;
    double          tmp;

    if (x == 0.0)
    {
        *psigd = 1.0;
        *psigdf = 1.0;
        return;
    }
    else
    {
        *psigd = *psigdf = 0.0;
        for (i = 0; i < x; ++i)
        {
            tmp = 1.0 / (n - i);
            *psigd += tmp;
            *psigdf -= tmp*tmp;
        }
    }
}


/* Maximum likelihood fit. */
static void
evalbinomialML(const double *x, const int num, const double n, const double p,
               const double ave, double *fx, double *dfx)
{
    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], n, &psigd, &psigdf); */
/*         psigd_sum += psigd; */
/*         psigdf_sum += psigdf; */
/*     } */
/*  */
/*     *fx = psigd_sum + num * log(1.0 - p); */
/*     *dfx = psigdf_sum; */

    psigd_sum = psigdf_sum = 0.0;
    for (i = 0; i < num; ++i)
    {
        psigd_sum += gsl_sf_psi(n-x[i]+1);
        psigdf_sum += gsl_sf_psi_1(n-x[i]+1);
    }

    *fx = num*gsl_sf_psi(n+1) - psigd_sum + num * log(1.0 - p);
    *dfx = num*gsl_sf_psi_1(n+1) - psigdf_sum;
}


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

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

    max = findmax(data, num);

    /* 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) / ave;

    if (pmme < 0.0)
        *p = pmme = 0.1;

    *n = nmme = ave / *p;

    if (*n < max)
        *n = max+1;

    printf("    % 10.6e % 10.6e % 10.6e % 10.6e % 10.6e\n",
           *n, *p, fx, dfx, fx/dfx);

    for (i = 0; i < maxit; ++i)
    {
        evalbinomialML(data, num, *n, *p, ave, &fx, &dfx);

        fxdfx = fx / dfx;

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

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

        if (*n < 1.0)
            *n = 1.0;
        else if (*n < max)
            *n = max;

        /* ML estimate of p */
        *p = ave / *n;

        printf("%3d % 10.6e % 10.6e % 10.6e % 10.6e % 10.6e\n",
               i, *n, *p, fx, dfx, fx/dfx);
    }

    *n = round(*n);
    *p = ave / *n;

/*     if (i == maxit) */
/*     { */
/*         *n = nmme; */
/*         *p = pmme; */
/*     } */

    printf("\n\nbinomial logL %e\n", binomial_logL(*n, *p));

/*     *n = 100; */
/*     *p = 0.5; */

    chi2 = chi_sqr_adapt(x, num, 0, logL, *n, *p, binomial_pdf, binomial_lnpdf, binomial_int);

    free(data);

    return(chi2);
}



/* Maximum likelihood fit of data to a binomial distribution */
double
binomial_fit(const double *x, const int num, double *n, double *p, double *logL)
{
    int             i, start;
    double          ave, var, pmme, nmme, chi2, max;
    double          llogL, plogL, lp, pp;
    int             ln, pn, lslope, pslope, dir;
    double         *data = malloc(num * sizeof(double));
    int             maxit = 100;

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

    max = data[findmax(data, num)];

    /* 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) / ave;

    if (pmme < 0.0)
        *p = pmme = 0.01;

    *n = nmme = ave / *p;
    *n = nmme = ave*ave/(ave-var);

    if (*n < max)
        *n = max;

    *p = ave / *n;

/*     printf("\n    % 10.6e % 10.6e % 10.6e % 10.6e", *n, *p, ave, max); */
/*     fflush(NULL); */

    start = (int) *n;

/*     ln = start-1; */
/*     lp = ave / ln; */
/*  */
/*     llogL = dist_logL(binomial_lnpdf, ln, lp, data, num); */
/*  */
/*     printf("\n    logL:% e  n:%-3d  p:% e", */
/*             llogL, ln, lp); */

    ln = start;
    pn = ln + 1;
    lp = ave / ln;
    pp = ave / pn;
    llogL = dist_logL(binomial_lnpdf, ln, lp, data, num);
    plogL = dist_logL(binomial_lnpdf, pn, pp, data, num);

/*     printf("\n    logL:% e  n:%-3d  p:% e", */
/*             llogL, ln, lp); */
/*     printf("\n    logL:% e  n:%-3d  p:% e", */
/*             plogL, pn, pp); */
/*     fflush(NULL); */

    if (plogL > llogL)
    {
        dir = 1;
        pslope = 1;
    }
    else if (plogL < llogL)
    {
        dir = -1;
        pslope = -1;
    }
    else
    {
        dir = 0;
        pslope = 0;
    }

    for (i = 0; i < maxit; ++i)
    {
        llogL = plogL;
        ln = pn;
        lp = pp;
        lslope = pslope;

        if (pn == max && dir == -1)
        {
            *n = pn;
            *p = pp;
            break;
        }

        pn += dir;

        /* ML estimate of p */
        pp = ave / pn;
        plogL = dist_logL(binomial_lnpdf, pn, pp, data, num);

        if (plogL > llogL)
            pslope = 1 * dir;
        else if (plogL < llogL)
            pslope = -1 * dir;
        else
        {
            pslope = 0;
        }

/*         printf("\n%3d logL:% e  n:%-3d  p:% e  slope:% d -- logL:% e  n:%-3d  p:% e  slope:% d", */
/*                 i, plogL, pn, pp, pslope, llogL, ln, lp, lslope); */
/*         fflush(NULL); */

        if (pslope == 0)
        {
            *n = pn;
            *p = pp;
            break;
        }

        if (pslope != lslope)
        {
            *n = ln;
            *p = lp;
            break;
        }
    }

    chi2 = chi_sqr_adapt(x, num, 0, logL, *n, *p, binomial_pdf, binomial_lnpdf, binomial_int);

    free(data);

    return(chi2);
}
