/*
    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 <stdlib.h>
#include <math.h>
#include <float.h>
#include "statistics.h"
#include "DLTmath.h"
#include "weibull_dist.h"
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

/* The Weibull distribution has the form,

       p(x) dx = (b/a) (x/a)^(b-1) exp(-(x/a)^b) dx

       a > 0        scale parameter
       b > 0        shape parameter
       0 <= x < +inf

   http://mathworld.wolfram.com/WeibullDistribution.html
*/

double
weibull_dev(const double a, const double b, const gsl_rng *r2)
{
    return (a * pow(-log(gsl_rng_uniform(r2)), (1.0 / b)));
}


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

    if (x < 0.0)
    {
        return(0.0);
    }
    else if (x == 0.0)
    {
        if (b == 1.0)
            return(1.0 / a);
        else
            return(0.0);
    }
    else if (b == 1.0)
    {
        return(exp(-x / a) / a);
    }
    else
    {
        p = (b / a) * exp(-pow(x/a, b) + (b - 1.0) * log(x / a));
        return(p);
    }
}


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

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

    return(p);
}


double
weibull_cdf(const double x, const double a, const double b)
{
    return(1.0 - weibull_sdf(x, a, b));
}


double
weibull_sdf(const double x, const double a, const double b)
{
    if (x == 0.0)
        return(1.0);
    else
        return(exp(-pow((x / a), b)));
}


double
weibull_int(const double x, const double y, const double a, const double b)
{
    return(weibull_sdf(x, a, b) - weibull_sdf(y, a, b));
}


/* From Cover and Thomas (1991) _Elements of Information Theory_,
   pp. 486-487. But this appears to be wrong. Do not use. */
// double
// weibull_logL(const double a, const double b)
// {
//     return((b - 1.0) * MY_EULER / b + log(pow(a, 1.0 / b) / b) + 1.0);
// }


/* From wikipedia.  Is it right??  */
double
weibull_logL(const double a, const double b)
{
    return(MY_EULER*(1.0-1/b) + log(a/ b) + 1.0);
}


/* Maximum likelihood fit of data to a Weibull distribution
   Based on _Statistical Distributions_ 3rd ed.
   Evans, Hastings, and Peacock, p 193, 196.
   and on my own derivation of the Jacobian matrix (ugh!).
*/
double
weibull_fit(const double *x, const int nx, double *eta, double *beta, double *prob)
{
    double           nd = (double) nx;
    double           a, b;
    double           guesses[2];
    double          *array;
    int              i;
    double           tol = 1e-8;

    /* initial eta and beta guesses are rank estimates:
       eta = 63.21 percentile x value (1 - e^-1)
       beta = ln(ln(2) / ln(median / eta)
    */
    array = (double *) malloc(nx * sizeof(double));

    for (i = 0; i < nx; ++i)
    {
        if (x[i] < 0.0)
        {
            fprintf(stderr, "\n ERROR345: weibull distributed data must be >= 0.0 ");
            return(-1.0);
        }
        else
        {
            array[i] = x[i];
        }
    }

    qsort(array, nx, sizeof(double), dblcmp);

    a = guesses[0] = array[(int)floor(0.632120559 * nd)];
    b = guesses[1] = log(log(2)) / log(array[(int) floor(0.5 * nd)] / guesses[0]);

    if (guesses[0] < 0.0)
        guesses[0] = tol;

    if (guesses[1] < 0.0)
        guesses[1] = tol;

    if (mnewt_weibullML(x, nx, 100, guesses, tol, tol) == 1)
    {
        *eta  = guesses[0];
        *beta = guesses[1];
    }
    else /* newton-raphson did not converge, use rough rank estimates from above */
    {
        *eta  = a;
        *beta = b;
    }

    free(array);

    return(chi_sqr_adapt(x, nx, 0, prob, *eta, *beta, weibull_pdf, weibull_lnpdf, weibull_int));
}


/* evalweibullML() */
/* this function evaluates six functions:
   the two simultaneous equations to be solved to satisfy ML 
   and the 2x2 jacobi matrix for these two functions.
   Based on _Statistical Distributions_ 3rd ed. Evans, Hastings, and Peacock, p 196. 
   To do this monstrous ML fit, we must solve these two equations simultaneously:

        F_0 = a - [(1/n) * \Sum{x^b}]^{1/b} = 0
        F_1 = N/b + \Sum{ln(x/a)} - \Sum{((x/a)^b) * ln(x/a)} = 0

   When these are zero the ML soln is found.
   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 
   derivatives of a vector of functions:

       J_ij = dF_i/dx_j

   where x_j is the jth variable (here there are two, a and b).
   The first derivatives of the above two top equations are these four ungodlies
   (horizontal dashes are divides/fractions):

       J[0][0] = dF_0/da = 1  easy!

                           [(1/n) * \Sum{x^b}]^{1/b} * ln((1/n) * \Sum{x^b}))
       J[0][1] = dF_0/db = --------------------------------------------------
                                                  b^2

                           - [((1/n) * \Sum{x^b})^{(1/b) - 1} * \Sum{ln(x)} * x^b]
                             -----------------------------------------------------
                                                    (b*N)

                           -N + [\Sum{x/a}^b] + [b * \Sum{(ln(x/a)) * ((x/a)^b)}]
       J[1][0] = dF_1/da = ------------------------------------------------------
                                                     a

       J[1][1] = dF_1/db = [-N/(b^2)] + [\Sum{ln(x/a) * ln(x/a) * ((x/a)^b)}]

   I derived and re-derived and re-re-derived these several times, so I believe they 
   are correct.  They also work pretty well, and give very good results relative to
   other methods. It seems to converge in less than half-a-dozen iterations.

       x            array of data to be fit
       nx           number of data points
       params       vector of the 2 Weibull parameters, the scale and shape paramenter, respectively
       fvec         vector for results of the two ML equations
       fjac         2x2 matrix of partial deriviative of the two ML equations, for each param
*/
void
evalweibullML(const double *x, const int nx,
              double *params, double *fvec, double **fjac)
{
    double          a, b, xb, logxa, xablogxa, logxxb, xab, logxaxab, logxa2xab;
    double          nd = (double) nx;
    int             i;

    a = params[0];
    b = params[1];

    xb = logxxb = 0.0;
    for (i = 0; i < nx; ++i)
    {
        if (x[i] == 0.0)
            continue;

        xb += pow(x[i], b);
        logxxb += log(x[i]) * pow(x[i], b);
    }
    xb /= nd;

    fvec[0] = a - pow(xb, 1.0 / b);
    fjac[0][0] = 1.0;
    fjac[0][1] = (pow(xb, 1.0 / b) * log(xb) / mysquare(b));
    fjac[0][1] -= (pow(xb, (1.0 / b) - 1.0) * logxxb) / (b * nd);

    logxa = xablogxa = 0.0;
    for (i = 0; i < nx; ++i)
    {
        if (x[i] == 0.0)
            continue;

        logxa += log(x[i] / a);
        xablogxa += pow((x[i] / a), b) * log(x[i] / a);
    }

    fvec[1] = (nd / b) + logxa - xablogxa;

    xab = logxaxab = logxa2xab = 0.0;
    for (i = 0; i < nx; ++i)
    {
        if (x[i] == 0.0)
            continue;
 
        xab += pow((x[i] / a), b);
        logxaxab += log(x[i] / a) * pow(x[i]/ a, b);
        logxa2xab += mysquare(log(x[i] / a)) * pow(x[i] / a, b);
    }

    fjac[1][0] = -(nd / a) + (xab / a) + (b * logxaxab / a);
    fjac[1][1] = -(nd / mysquare(b)) - logxa2xab;
}


int
mnewt_weibullML(const double *data, const int ndata, const int ntrial,
                double *params, const double tolx, const 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)
    {
        /* printf(" iter[%4d]  a=%12.6e  b=%12.6e\n", k, params[0], params[1]);
        fflush(NULL); */

        evalweibullML(data, ndata, params, 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];
        }

        /* weibull dist a and b cannot be non-positive */
        for (i = 0; i < 2; ++i)
        {
            if (params[i] <= 0.0)
                params[i] = tolx;
        }

        if (errx < tolx)
            break;
    }

    MatDestroy(&fjac);

    if (k == ntrial)
        return(0);
    else
        return(1);
}


/* Maximum likelihood fit of histogram data to a Weibull distribution */
double
weibull_histfit(double *x, double *freq, int nx, double *eta, double *beta, double *logL)
{
    double           nd = (double) nx;
    double           a, b;
    double           guesses[2];
    double          *array;
    int              i;
    double           tol = 1e-6;

    /* estimate eta and beta with:
       eta = 63.21 percentile x value (1 - e^-1)
       beta = ln(ln(2) / ln(median / eta)
    */
    array = (double *) malloc(nx * sizeof(double));

    for (i = 0; i < nx; ++i)
        array[i] = x[i];

    qsort(array, nx, sizeof(double), dblcmp);

    a = guesses[0] = array[(int)floor(0.632120559 * nd)];
    b = guesses[1] = log(log(2)) / log(array[(int)floor(0.5 * nd)] / guesses[0]);

    if (guesses[0] < 0.0)
        guesses[0] = tol;

    if (guesses[1] < 0.0)
        guesses[1] = tol;

/* SCREAME(guesses[0]); */
/* SCREAME(guesses[1]); */

    /* guesses[0] = guesses[1] = 1.0; */
    if (mnewt_hist_weibullML(x, freq, nx, 30, guesses, tol, tol) == 1)
    {
        *eta  = guesses[0];
        *beta = guesses[1];
    }
    else /* newton-raphson did not converge, use rough rank estimates from above */
    {
        *eta  = a;
        *beta = b;
    }

    free(array);

    /* SCREAMF(chi_sqr_adapt(x, nx, 0, prob, a, b, weibull_pdf)); */
    return(chi_sqr_hist(x, freq, nx, logL, *eta, *beta, weibull_pdf, weibull_lnpdf, weibull_int));
}


int
mnewt_hist_weibullML(const double *data, double *freq,
                     const int ndata, const int ntrial,
                     double *params, const double tolx, const 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)
    {
        eval_hist_weibullML(data, freq, ndata, params, fvec, fjac);

/*         SCREAME(fvec[0]); */
/*         SCREAME(fvec[1]); */

        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];
        }

/* SCREAME(p[0]); */
/* SCREAME(p[1]); */

        /* weibull dist a and b cannot be non-positive */
        for (i = 0; i < 2; ++i)
        {
            if (params[i] <= 0.0)
                params[i] = tolx;
        }

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

        if (errx < tolx)
            break;
    }

/* SCREAMD(k); */

    MatDestroy(&fjac);

    if (k == ntrial)
        return(0);
    else
        return(1);
}


void
eval_hist_weibullML(const double *x, double *freq, const int nx,
                    double *params, double *fvec, double **fjac)
{
    double          a, b, norm;
    double          xb, logxa, xablogxa, logxxb, xab, logxaxab, logxa2xab;
    int             i;

    norm = 0.0;
    for (i = 0; i < nx; ++i)
        norm += freq[i];

    a = params[0];
    b = params[1];

    xb = logxxb = 0.0;
    for (i = 0; i < nx; ++i)
    {
        if (x[i] == 0.0)
            continue;

        xb += (freq[i] * pow(x[i], b));
        logxxb += (freq[i] * log(x[i]) * pow(x[i], b));
    }
    xb /= norm;

    fvec[0] = a - pow(xb, 1.0 / b);
    fjac[0][0] = 1.0;
    fjac[0][1] = (pow(xb, 1.0 / b) * log(xb) / mysquare(b));
    fjac[0][1] -= (pow(xb, (1.0 / b) - 1.0) * logxxb) / (b * norm);

    logxa = xablogxa = 0.0;
    for (i = 0; i < nx; ++i)
    {
        if (x[i] == 0.0)
            continue;

        logxa += (freq[i] * log(x[i] / a));
        xablogxa += (freq[i] * pow((x[i] / a), b) * log(x[i] / a));
    }

    fvec[1] = (norm / b) + logxa - xablogxa;

    xab = logxaxab = logxa2xab = 0.0;
    for (i = 0; i < nx; ++i)
    {
        if (x[i] == 0.0)
            continue;
 
        xab += freq[i] * pow((x[i] / a), b);
        logxaxab += freq[i] * log(x[i] / a) * pow(x[i]/ a, b);
        logxa2xab += freq[i] * mysquare(log(x[i] / a)) * pow(x[i] / a, b);
    }

    fjac[1][0] = -(norm / a) + (xab / a) + (b * logxaxab / a);
    fjac[1][1] = -(norm / mysquare(b)) - logxa2xab;
}
