/*
    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 <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 "betasym_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
betasym_dev(const double a, const double nullv, const gsl_rng *r2)
{
    double          x1, x2;

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

    return(x1 / (x1 + x2));
}


double
betasym_pdf(const double x, const double a, const double nullv)
{
    double          p, gab, g2a;

    if (x < 0.0 || x > 1.0)
    {
        return(0.0);
    }
    else 
    {
        gab = lgamma(2*a);
        g2a = lgamma(a);

        p = exp(gab - 2*g2a) * pow (x - x*x, a - 1.0);

        return(p);
    }
}


double
betasym_lnpdf(const double x, const double a, const double nullv)
{
    double          p, gab, g2a;

    if (x < 0.0 || x > 1.0)
    {
        return(INFINITY);
    }
    else 
    {
        gab = lgamma(2*a);
        g2a = lgamma(a);

        p = gab - 2*g2a + (a - 1.0) * log(x - x*x);

        return(p);
    }
}


double
betasym_cdf(const double x, const double a, const double nullv)
{
    return(gsl_sf_beta_inc(a, a, x));
}


double
betasym_sdf(const double x, const double a, const double nullv)
{
    return(1.0 - betasym_cdf(a, a, x));
}


double
betasym_int(const double x, const double y, const double a, const double nullv)
{
    return(gsl_sf_beta_inc(a, a, y) - gsl_sf_beta_inc(a, a, x));
}


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

    logL = 2.0 * lgamma(a) - lgamma(2.0 * a)
           - 2.0 * (a - 1.0) * (gsl_sf_psi(a) - gsl_sf_psi(2.0 * a));

    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(const double logterm, const double a, double *fx, double *dfx)
{
    *fx = gsl_sf_psi(2.0 * a) - gsl_sf_psi(a) + logterm;
    *dfx = 2.0 * gsl_sf_psi_1(2.0 * a) - gsl_sf_psi_1(a);
}


/* Maximum likelihood fit of data to a beta distribution */
double
betasym_fit(const double *x, const int n, double *ra, double *nullv, double *logL)
{
    double           nd = (double) n;
    double           a, ave, var, fx, dfx, tmp;
    double           guess, logterm, tol = 1e-8;
    int              i, maxiter = 100;

    /* find constants for ML equations */
    logterm = 0.0;
    for (i = 0; i < n; ++i)
    {
        if (x[i] < 0.0 || x[i] > 1.0)
        {
            fprintf(stderr, "\n ERROR345: beta distributed data must be >= 0.0 && <= 1.0");
            return(-1.0);
        }
        else if (x[i] == 0.0)
        {
            continue;
        }
        else
        {
            logterm += log(x[i] - x[i]*x[i]);
        }
    }
    logterm /= (2.0 * 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)
    {
        tmp = (x[i] - 0.5);
        var += tmp*tmp;
    }
    var /= nd;

    a = guess = (1.0 - 4.0 * var) / (4.0 * var);

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

    for (i = 0; i < maxiter; ++i)
    {
        evalbetaML(logterm, a, &fx, &dfx);

        a -= (fx / dfx); /* Newton-Raphson correction */
        /* *b = ave / *c; */

        if (a < DBL_EPSILON)
        {
            a = DBL_EPSILON;
            break;
        }

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

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

    *ra = a;
    /* printf("\n\nbeta logL %e\n", betasym_logL(*ra, *rb)); */

    return(chi_sqr_adapt(x, n, 0, logL, *ra, *nullv, betasym_pdf, betasym_lnpdf, betasym_int));
}
