/*
    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 <stdlib.h>
#include <gsl/gsl_sf_gamma.h>
#include <gsl/gsl_sf_psi.h>
#include "statistics.h"
//#include "DLTmath.h"
#include "gamma_dist.h" /* gamma_dev() */
#include "chisqr_dist.h"

/* The chi^2 distribution has the form

   p(x) dx = (1/(2^(n/2) * Gamma(nu/2))) * x^((nu - 2)/2) * exp(-x/2)) dx

   0 <= x < +inf
   nu > 0
*/

double
chisqr_dev(const double nu, const double nullval, const gsl_rng *r2)
{
    return(2.0 * gamma_dev(1.0, 0.5*nu, r2));
}


double
chisqr_pdf(const double x, const double nu, const double nullval)
{
    double          p, nu2;

    if (x <= 0.0)
    {
        return(0.0);
    }
    else
    {
        nu2 = 0.5*nu;

        p = (nu2-1.0)*log(x) - 0.5*x - nu2*log(2.0) - lgamma(nu2);
        return(exp(p));
    }
}


double
chisqr_lnpdf(const double x, const double nu, const double nullval)
{
    double          nu2 = 0.5*nu;

    return((nu2-1.0)*log(x) - 0.5*x - nu2*log(2.0) - lgamma(nu2));
}


double
chisqr_cdf(const double x, const double nu, const double nullval)
{
    if (x <= 0.0)
        return(0.0);
    else
        return(gsl_sf_gamma_inc_P(0.5*nu, 0.5*x));
}


double
chisqr_sdf(const double x, const double nu, const double nullval)
{
    if (x <= 0.0)
        return(1.0);
    else
        return(1.0 - chisqr_cdf(x, nu, nullval));
}


double
chisqr_int(const double x, const double y, const double nu, const double nullval)
{
    if (x <= 0.0)
        return(chisqr_cdf(y, nu, nullval));
    else
        return(chisqr_cdf(y, nu, nullval) - chisqr_cdf(x, nu, nullval));
}


/* This is from Cover and Thomas (_Elements of Information Theory_),
   but it is wrong. Something's wrong.
   I even did the integration myself and got the same answer. Weird. 
*/
/* double */
/* chisqr_logL(const double nu, const double nullval) */
/* { */
/*     double         logL, nu2; */
/*      */
/*     nu2 = nu / 2.0; */
/*  */
/*     logL = -log(2.0 * tgamma(nu2)) - (1.0 - nu2) * gsl_sf_psi(nu2) - nu2; */
/*  */
/*     return(logL); */
/* } */

double
chisqr_logL(const double nu, const double nullval)
{
    double          nu2 = 0.5*nu;

/*     printf("\nchisqr logL: %f %f %f %f %f\n", */
/*            (nu2 - 1.0)*gsl_sf_psi(nu2), - log(2.0), - nu2, - lgammav, */
/*            (nu2 - 1.0)*gsl_sf_psi(nu2) - log(2.0) - nu2 - lgammav); */

    return((nu2 - 1.0)*gsl_sf_psi(nu2) - log(2.0) - nu2 - lgamma(nu2));
}


/* For the maximum likelihood fit we nust find the root of:

       F1 = (1/N)\Sum{log(x)} - log(2) - digamma{nu/2} = 0

   where the first derivative with repect to nu (dF1/dnu) is:

       F1' = -trigamma(nu/2)/2 = 0
*/
void
evalchisqrML(const double logterm, const double nu, double *fx, double *dfx)
{
    *fx  = logterm - gsl_sf_psi(0.5*nu);
    *dfx = -0.5*gsl_sf_psi_1(0.5*nu);
}


/* fit a chisqr distribution by maxinum likelihood */
double
chisqr_fit(const double *data, const int num, double *nu, double *nullp, double *prob)
{
    double          ave, logterm, fx, dfx, guess_nu;
    int             i;
    double          iter = 100;
    double          tol = 1e-8;

    ave = 0.0;
    for (i = 0; i < num; ++i)
    {
        if (data[i] < 0.0)
        {
            fprintf(stderr, "\n ERROR345: chi^2 distributed data must be >= 0.0 ");
            return(-1.0);
        }
        else
        {
            ave += data[i];
        }
    }
    ave /= (double) num;

    guess_nu = *nu = ave;

    logterm = 0.0;
    for (i = 0; i < num; ++i)
    {
        if(data[i] == 0.0)
            continue;

        logterm += log(data[i]);
    }
    logterm /= (double) num;
    logterm -= log(2.0);

    for (i = 0; i < iter; ++i)
    {
        evalchisqrML(logterm, *nu, &fx, &dfx);

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

        *nu -= (fx / dfx); /* Newton-Raphson correction */

        if (*nu <= 0.0)
            *nu = tol;
    }

    if (i == iter)
        *nu = guess_nu;

    chisqr_logL(*nu, 0.0);

    return(chi_sqr_adapt(data, num, 0, prob, *nu, 0, chisqr_pdf, chisqr_lnpdf, chisqr_int));
}
