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

/*
   The beta distribution has the form

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


/* Described in Knuth */
double
beta_dev(const double a, const double b, const gsl_rng *r2)
{
    double          x1, x2, bdev;

    x1 = gamma_dev(1.0, a, r2);
    x2 = gamma_dev(1.0, b, r2);

    bdev = x1 / (x1 + x2);

    if (bdev >= 1.0)
        bdev = 1.0 - DBL_MIN;

    if (bdev <= 0.0)
        bdev = DBL_MIN;

    return(bdev);
}


double
beta_pdf(const double x, const double a, const double b)
{
    double          p;

    if (x < 0.0 || x >= 1.0)
    {
        return(0.0);
    }
    else 
    {
/*         gab = lgamma(a + b); */
/*         ga = lgamma(a); */
/*         gb = lgamma(b); */

        p = pow (x, a - 1.0) * pow (1.0 - x, b - 1.0) / gsl_sf_lnbeta(a,b);

        return(p);
    }
}


double
beta_lnpdf(const double x, const double a, const double b)
{
    double          p;

    if (x < 0.0 || x >= 1.0)
    {
        return(-INFINITY);
    }
    else 
    {
/*         gab = lgamma(a + b); */
/*         ga = lgamma(a); */
/*         gb = lgamma(b); */

        p = -gsl_sf_lnbeta(a,b) + (a - 1.0) * log(x) + (b - 1.0) * log(1.0 - x);

        return(p);
    }
}


double
beta_cdf(const double x, const double a, const double b)
{
    return(gsl_sf_beta_inc(a, b, x));
}


double
beta_sdf(const double x, const double a, const double b)
{
    return(1.0 - beta_cdf(x, a, b));
}


double
beta_int(const double x, const double y, const double a, const double b)
{
    return(gsl_sf_beta_inc(a, b, y) - gsl_sf_beta_inc(a, b, x));
}


/* From Cover and Thomas (1991) _Elements of Information Theory_ */
double
beta_logL(const double a, const double b)
{
    double         logL = -DBL_MAX;

    logL = lgamma(a) + lgamma(b) - lgamma(a + b)
           - (a - 1.0) * (gsl_sf_psi(a) - gsl_sf_psi(a + b))
           - (b - 1.0) * (gsl_sf_psi(b) - gsl_sf_psi(a + b));

    return(-logL);
}


/* Maximum likelihood fit. */
/* Based on _Statistical Distributions_ 3rd ed. Evans, Hastings, and Peacock, p 41. 
   We must solve the simultaneous equations:

        F_0 = E(logx)     + digamma(a+b) - digamma(a) = 0
        F_1 = E(log(1-x)) + digamma(a+b) - digamma(b) = 0

   which imply:

        F_3 = digamma(a) - digamma(b) - E(logx) + E(log(1-x)) = 0

   when these are zero the ML soln is found.
   The third equation is a measure of how good the ML solution is.
   I use 2-Dimensional Newton-Raphson to find the ML a and b
   (see NR 2nd ed. pp. 379-382)
   2D Newton root finding requires the Jacobi matrix,
   which is the matrix of partial derivatives of a function
   that takes a vector of variables:

       J_ij = dF_i/dx_j

   where x_j is the jth variable (here there are two, a and b).
   The first derivative of the digamma fxn is, of course, the trigamma fxn.
   So, the first derivatives of the above two top equations are:

       J[0][0] = dF_0/da = trigamma(a+b) - trigamma(a)
       J[0][1] = dF_0/db = trigamma(a+b)
       J[1][0] = dF_1/da = trigamma(a+b)
       J[1][1] = dF_1/db = trigamma(a+b) - trigamma(b)

   I think (!).
*/
static void
evalbetaML(double *x, double Elogx, double Elog1mx, double *fvec, double **fjac)
{
    fvec[0] = Elogx   + gsl_sf_psi(x[0] + x[1]) - gsl_sf_psi(x[0]);
    fvec[1] = Elog1mx + gsl_sf_psi(x[0] + x[1]) - gsl_sf_psi(x[1]);
    fjac[0][1] = fjac[1][0] = gsl_sf_psi_1(x[0] + x[1]);
    fjac[0][0] = fjac[0][1] - gsl_sf_psi_1(x[0]);
    fjac[1][1] = fjac[0][1] - gsl_sf_psi_1(x[1]);
}


static void
mnewt_betaML(int ntrial, double *params, double Elogx, double Elog1mx, double tolx, double tolf)
{
    int             k, i, indx[2];
    double          errx, errf, d, fvec[2], **fjac, p[2];
    /* fjac is the jacobian of the function vector, i.e. the matrix of partial deriviatives */

    fjac = MatAlloc(2, 2);

    for (k = 0; k < ntrial; ++k)
    {
        evalbetaML(params, Elogx, Elog1mx, fvec, fjac);

        errf = 0.0;
        for (i = 0; i < 2; ++i)
            errf += fabs(fvec[i]);

        if (errf < tolf)
            break;

        for (i = 0; i < 2; ++i)
            p[i] = -fvec[i];

        ludcmp(fjac, 2, indx, &d);
        lubksb(fjac, 2, indx, p);

        errx = 0.0;
        for (i = 0; i < 2; ++i)
        {
            errx += fabs(p[i]);
            params[i] += p[i];
        }

        /* beta dist a and b cannot be less than zero */
        for (i = 0; i < 2; ++i)
        {
            if (params[i] <= 0.0)
                params[i] = tolx;
        }

        //printf(" iter[%4d]  a=%12.6e b=%12.6e\n", k, params[0], params[1]);
        //fflush(NULL);

        if (errx < tolx)
            break;
    }

    MatDestroy(&fjac);

    return;
}


/* Maximum likelihood fit of data to a beta distribution */
double
beta_fit(const double *x, const int n, double *ra, double *rb, double *logL)
{
    double           nd = (double) n;
    double           a, b, ave, var, xi;
    double           guesses[2];
    int              i;
    double           Elogx, Elog1mx;

    /* find constants for ML equations */
    Elogx = Elog1mx = 0.0;
    for (i = 0; i < n; ++i)
    {
        xi = x[i];

        if (xi < 0.0 || xi > 1.0)
        {
            fprintf(stderr, "\n ERROR345: beta distributed data must be >= 0.0 && <= 1.0");
            return(-1.0);
        }
        else if (xi == 1.0)
        {
            Elogx += log(1.0 - DBL_EPSILON);
            Elog1mx += log(DBL_EPSILON);
        }
        else if (xi == 0.0)
        {
            Elogx += log(DBL_EPSILON);
            /* Elog1mx += 0.0 */
        }
        else
        {
            Elogx += log(xi);
            Elog1mx += log(1.0 - xi);
        }
    }

    Elogx /= nd;
    Elog1mx /= nd;

    /* Method of moments initial guess at shape parameters.
       Based on _Statistical Distributions_ 3rd ed.
       Evans, Hastings, and Peacock, p 40. */
    ave = 0.0;
    for (i = 0; i < n; ++i)
        ave += x[i];
    ave /= nd;

    var = 0.0;
    for (i = 0; i < n; ++i)
        var += mysquare(x[i] - ave);
    var /= (nd - 1.0);

    a = guesses[0] =  ave * ((ave*(1.0 - ave)/var) - 1.0);
    b = guesses[1] = (1.0 - ave) * ((ave*(1.0 - ave)/var) - 1.0);

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

    if (b <= 0.0)
        b = DBL_EPSILON;

    mnewt_betaML(100, guesses, Elogx, Elog1mx, 1e-9, 1e-9);

    *ra = guesses[0];
    *rb = guesses[1];
    /* printf("\n\nbeta logL %e\n", beta_logL(*ra, *rb)); */

    return(chi_sqr_adapt(x, n, 0, logL, *ra, *rb, beta_pdf, beta_lnpdf, beta_int));
}
