/*
    Theseus - maximum likelihood superpositioning of macromolecular structures

    Copyright (C) 2004-2009 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 <stdlib.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include "DLTmath.h"
#include "statistics.h"
#include "EVD_dist.h"

/* generates a random deviate from an extreme value distribution */
double
EVD_dev(const double mu, const double lambda, const gsl_rng *r2)
{
    return(mu - log(-log(gsl_rng_uniform(r2))) / lambda);
}


double
EVD_pdf(const double x, const double mu, const double lambda)
{
    double          explxa;

    explxa = exp(-lambda * (x - mu));

    return(lambda * explxa * exp(-explxa));
}


double
EVD_lnpdf(const double x, const double mu, const double lambda)
{
    double          lxa;

    lxa = -lambda * (x - mu);

    return(log(lambda) + lxa  - exp(lxa));
}


double
EVD_cdf(const double x, const double mu, const double lambda)
{
    return(EVD_sdf(x, mu, lambda));
}


double
EVD_sdf(const double x, const double mu, const double lambda)
{
    return(1.0 - exp(-exp(-lambda * (x - mu))));
}


double
EVD_int(const double x, const double y, const double mu, const double lambda)
{
    return(EVD_sdf(x, mu, lambda) - EVD_sdf(y, mu, lambda));
}


/* Function: Lawless416()
 * Date:     SRE, Thu Nov 13 11:48:50 1997 [St. Louis]
 *
 * Purpose:  Equation 4.1.6 from [Lawless82], pg. 143, and
 *           its first derivative with respect to lambda,
 *           for finding the ML fit to EVD lambda parameter.
 *           This equation gives a result of zero for the maximum
 *           likelihood lambda.
 *
 *           Can either deal with a histogram or an array.
 *
 *           Warning: beware overflow/underflow issues! not bulletproof.
 *
 * Args:     x      - array of sample values
 *           n      - number of samples
 *           lambda - a lambda to test
 *           ret_f  - RETURN: 4.1.6 evaluated at lambda
 *           ret_df - RETURN: first derivative of 4.1.6 evaluated at lambda
 *
 * Return:   (void)
 */
void
Lawless416(const double *x, const int n, const double lambda,
           double *ret_f, double *ret_df)
{
    double          esum;        /* \sum e^(-lambda xi)      */
    double          xesum;       /* \sum xi e^(-lambda xi)   */
    double          xxesum;      /* \sum xi^2 e^(-lambda xi) */
    double          xsum;        /* \sum xi                  */
    double          xi, elxi, xelxi;
    int             i;

    esum = xesum = xsum = xxesum = 0.0;
    for (i = 0; i < n; i++)
    {
        xi = x[i];
        xsum += xi;
        elxi = exp(-lambda * xi);
        xelxi = xi * elxi;
        xesum += xelxi;
        xxesum += xi * xelxi;
        esum += elxi;
    }

    *ret_f = (1.0 / lambda) -
             (xsum / (double) n) +
             (xesum / esum);

    *ret_df = ((xesum * xesum) / (esum * esum)) -
              (xxesum / esum) -
              (1.0 / (lambda * lambda));
}


/* Function: EVD_fit()
 *
 * Purpose:  Given a list of EVD-distributed samples,
 *           find maximum likelihood parameters lambda and
 *           mu.
 *
 * Algorithm: Uses approach described in [Lawless82]. Solves
 *           for lambda using Newton/Raphson iterations;
 *           then substitutes lambda into Lawless' equation 4.1.5
 *           to get mu.
 *
 *           Newton/Raphson algorithm developed from description in
 *           Numerical Recipes in C [Press88].
 *
 * Args:     x          - list of EVD distributed samples or x-axis of histogram
 *           n          - number of samples
 *           rmu       : RETURN: ML estimate of mu
 *           rlambda   : RETURN: ML estimate of lambda
 *
 * Return:   reduced chi^2 for the fit
 */
double
EVD_fit(const double *x, const int n, double *rmu, double *rlambda, double *prob)
{
    double           nd = (double) n;
    double           lambda, mu;
    double           fx;         /* f(x)  */
    double           dfx;        /* f'(x) */
    double           esum;       /* \sum e^(-lambda xi) */
    double           std_dev, ave;
    double           tol = 1e-8;
    double           guess, diff;
    int              i;

    /* Initial guess at lambda, based on method of moments.
       Based on _Statistical Distributions_ 3rd ed. Evans, Hastings, and Peacock, p 86.
       variance = pi^2 / (6 * lambda^2) */
    ave = 0.0;
    for (i = 0; i < n; ++i)
        ave += x[i];

    ave /= nd;

    std_dev = 0.0;
    for (i = 0; i < n; ++i)
    {
        diff = x[i] - ave;
        std_dev += diff * diff;
    }

    std_dev /= (nd - 1.0);
    std_dev = sqrt(std_dev);

    lambda = guess = MY_PI / (std_dev * sqrt(6.0));

    /* Use Newton/Raphson to solve Lawless 4.1.6 and find ML lambda */
    for (i = 0; i < 200; ++i)
    {
        Lawless416(x, n, lambda, &fx, &dfx);

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

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

        if (lambda <= 0.0)
            lambda = tol * tol; /* but be a little careful  */
    }

    /* If we did 200 iterations Newton/Raphson failed.
       Resort to a bisection search. Worse convergence speed but
       guaranteed to converge (unlike Newton/Raphson). We assume (!?)
       that fx is a monotonically decreasing function of x;
       i.e. fx > 0 if we are left of the root, fx < 0
       if we are right of the root.
    */
    if (i == 200)
    {
        double           left, right, mid;

        fprintf(stderr, "\n EVD_fit(): Newton/Raphson failed; switching to bisection ");

        /* First we need to bracket the root */
        lambda = right = left = guess;
        Lawless416(x, n, lambda, &fx, &dfx);

        if (fx < 0.0)
        { /* fix right; search left. */
            do
            {
                left -= 0.1;

                if (left < 0.0)
                {
                    fprintf(stderr, "\n ERROR124: EVD_fit(): failed to bracket root ");
                    return(-1.0);
                }

                Lawless416(x, n, left, &fx, &dfx);
            }
            while (fx < 0.0);
        }
        else
        { /* fix left; search right. */
            do
            {
                right += 0.1;
                Lawless416(x, n, right, &fx, &dfx);

                if (right > 100.0)
                {
                    fprintf(stderr, "\n ERROR124: EVD_fit(): failed to bracket root ");
                    return(-1.0);
                }
            }
            while (fx > 0.0);
        }

        /* now we bisection search in left/right interval */
        for (i = 0; i < 200; i++)
        {
            mid = (left + right) / 2.0;
            Lawless416(x, n, mid, &fx, &dfx);

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

            if (fx > 0.0)
                left = mid;
            else
                right = mid;
        }

        if (i == 200)
        {
            fprintf(stderr, "\n ERROR124: EVD_fit(): bisection bit it too ");
            return(-1.0);
        }

        lambda = mid;
    }

    /* 3. Substitute into Lawless 4.1.5 to find mu */
    esum = 0.0;
    for (i = 0; i < n; ++i)
        esum += exp(-lambda * x[i]);

    mu = log(nd / esum) / lambda;

    *rlambda = lambda;
    *rmu = mu;

    return(chi_sqr_adapt(x, n, 0, prob, mu, lambda, EVD_pdf, EVD_lnpdf, EVD_int));
}
