/*
    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 <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include "DLTmath.h"
#include "statistics.h"
#include "gamma_dist.h" /* gamma_dev() */
#include "chi_dist.h"

/* The chi 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
chi_dev(const double nu, const double nullval, const gsl_rng *r2)
{
    return(sqrt(2.0 * gamma_dev(1.0, 0.5*nu, r2)));
}


double
chi_pdf(const double x, const double nu, const double nullval)
{
    double          p, lgammav;

    if (x <= 0.0)
    {
        return(0.0);
    }
    else
    {
        lgammav = lgamma(0.5*nu);
      
        p = exp((nu - 1.0)*log(x) - (0.5*x*x) - (0.5*nu - 1.0)*log(2.0) - lgammav);
        return(p);
    }
}


double
chi_lnpdf(const double x, const double nu, const double nullval)
{
    double          p, lgammav;

    lgammav = lgamma(nu / 2.0);
      
    p = (nu - 1.0)*log(x) - (0.5*x*x) - (0.5*nu - 1.0)*log(2.0) - lgammav;

    return(p);
}


double
chi_cdf(const double x, const double nu, const double nullval)
{
    return(gsl_sf_gamma_inc_P(0.5*nu, 0.5*x*x));
}


double
chi_sdf(const double x, const double nu, const double nullval)
{
    return(1.0 - chi_cdf(x, nu, nullval));
}


double
chi_int(const double x, const double y, const double nu, const double nullval)
{
    return(chi_cdf(y, nu, nullval) - chi_cdf(x, nu, nullval));
}


/* This is from Cover and Thomas (_Elements of Information Theory_),
   but it is wrong. Something's wrong. */
double
chi_logL(const double nu, const double nullval)
{
    double         logL, nu2;

    nu2 = 0.5*nu;
    logL = lgamma(nu2) - 0.5 * log(2.0) - (nu2 - 0.5) * gsl_sf_psi(nu2) + nu2;

    return(-logL);
}


/* For Maximum likelihood fit we must find the root of:

       F1 = (2/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
evalchiML(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);
}


double
chi_fit(const double *data, const int num, double *nu, double *nullp, double *prob)
{
    double          rawmom2; /* second raw moments */
    double          logterm, fx, dfx, guess_nu;
    int             i;
    double          iter = 100;
    double          tol = 1e-10;

    /* nu = (second raw moment)
       the second raw moment (second moment about the origin)
       is just the average squared values.
       Based on _Statistical Distributions_ 3rd ed. Evans, Hastings, and Peacock, p 57.
       Gamma(n + 1)/Gamma(n) = n
    */
    rawmom2 = 0.0;
    for (i = 0; i < num; ++i)
    {
        if (data[i] < 0.0)
        {
            fprintf(stderr, "\n ERROR345: chi distributed data must be >= 0.0 ");
            return(-1.0);
        }
        else
        {
            rawmom2 += mysquare(data[i]);
        }
    }
    rawmom2 /= (double) num;

    guess_nu = *nu = rawmom2;

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

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

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

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

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

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

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

    /* printf(" Chi logL: %f\n", chi_logL(*nu, *nullp)); */

    return(chi_sqr_adapt(data, num, 0, prob, *nu, 0, chi_pdf, chi_lnpdf, chi_int));
}
