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

/* The Cauchy probability distribution is

   p(x) dx = (1/(PI*b)) (1 + ((x - a))^2 / b)^(-1) dx
   
   a = location param
   b = shape param (roughly a SD)

   http://mathworld.wolfram.com/CauchyDistribution.html
   http://en.wikipedia.org/wiki/Lorentzian_function
*/

double
cauchy_dev(const double a, const double b, const gsl_rng *r2)
{
    return (a + b * tan(MY_PI * (gsl_rng_uniform(r2) - 0.5)));
}


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

    p = (MY_PI * b) * (1.0 + mysquare((x - a) / b));

    p = 1.0 / p;

    return(p);
}


double
cauchy_lnpdf(const double x, const double a, const double b)
{
    return(- log(MY_PI * b) - log(1.0 + mysquare((x - a) / b)));
}


double
cauchy_cdf(const double x, const double a, const double b)
{
    return(0.5 - atan((x - a) / b) / MY_PI);
}


double
cauchy_sdf(const double x, const double a, const double b)
{
    return(0.5 + atan((x - a) / b) / MY_PI);
}


double
cauchy_int(const double x, const double y, const double a, const double b)
{
    return(cauchy_sdf(y, a, b) - cauchy_sdf(x, a, b));
}


double
cauchy_logL(const double a, const double b)
{
    return(-4.0 * MY_PI * b);
}


/* Maximum likelihood fit of data to a cauchy distribution */
double
cauchy_fit(const double *x, const int n, double *ra, double *rb, double *logL)
{
    double           a, b;
    double          *array;
    double           guesses[2];
    int              i, status;

    /* guess at parameters -- use median for a, the location param
       we could use the median absolute deviation (MAD) for the
       shape param, b (much better than the variance or SD,
       which are infinite for the Cauchy), but the interquartile
       range provides a very good estimate.
       (2*b is also the full-width at half maximum height.)
    */
    array = (double *) malloc(n * sizeof(double));
    for (i = 0; i < n; ++i)
        array[i] = x[i];
    qsort(array, n, sizeof(double), dblcmp);
    a = array[(int) ((n+1) / 2)];

    b = (array[(int)(0.75 * n)] - array[(int)(0.25 * n)]) / 2.0;

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

    status = mnewt_cauchyML(150, x, n, guesses, 1e-8, 1e-8);

    /* if it bonks, refit with MAD guess for shape param */
    if (status == 0)
    {
        b = 0.0;
        for (i = 0; i < n; ++i)
            b += fabs(x[i] - a);
        b = sqrt(b / n);

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

        status = mnewt_cauchyML(150, x, n, guesses, 1e-8, 1e-8);

        /* if it bonks, refit with mean for location param */
        if (status == 0)
        {
            a = 0.0;
            for (i = 0; i < n; ++i)
                a += x[i];
            a /= (double) n;
    
            guesses[0] = a;
            guesses[1] = (array[(int)(0.75 * n)] - array[(int)(0.25 * n)]) / 2.0;;
    
            status = mnewt_cauchyML(150, x, n, guesses, 1e-8, 1e-8);
        }
    }

    *ra = guesses[0];
    *rb = guesses[1];

    free(array);
    /* printf(" Cauchy logL: %f\n", cauchy_logL(*ra, *rb)); */

    return(chi_sqr_adapt(x, n, 0, logL, *ra, *rb, cauchy_pdf, cauchy_lnpdf, cauchy_int));
}


/* Maximum likelihood fit for the Cauchy.

   Based on _Johnson and Kotz_ and my own derivation of 
   the Jacobian. A Cauchy distribution is very difficult
   to fit, as the likelihood function is well-known for 
   being multinodal, esp. when N is small. This may
   diverge. Alternatively, this may
   converge on a local maximum, but it seems to work
   very well in practice. According to J&K, the two-
   variable ML fits appear to be more robust than the
   one-variable, strangely enough. I use the median as a 
   starting guess for the location parameter,
   and the interquartile range as an initial
   scale parameter.
   We must solve the simultaneous equations:

        F_0 = 2 \SUM{(x - a) / (b^2 + (x - a)^2)} = 0
        F_1 = n/b - 2b \Sum{1 / (b^2 + (x - a)^2)} = 0

   when these are zero a local ML solution 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 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 derivatives of the above two equations are:

       J[0][0] = dF_0/da = -2 \Sum {(b^2 - (x - a)^2) / ((b^2 + (x - a)^2))^2}
       J[0][1] = dF_0/db = -4b \Sum {(x - a) / ((b^2 + (x - a)^2))^2}
       J[1][0] = dF_1/da = -4b \Sum {(x - a) / ((b^2 + (x - a)^2))^2}
       J[1][1] = dF_1/db = -n/b^2 - 2\Sum {(b^2 - (x - a)^2) / ((b^2 + (x - a)^2))^2}
*/
void
evalcauchyML(const double *x, const int n,
             double *params, double *fvec, double **fjac)
{
    int            i;
    double         a, b, b2, xma, xma2, b2mxma2, b2pxma2, ib2pxma2, sqr_b2pxma2;

    a = params[0];
    b = params[1];
    b2 = mysquare(b);

    fvec[0] = fvec[1] = fjac[0][0] = fjac[1][0] = fjac[0][1] = fjac[1][1] = 0.0;
    for (i = 0; i < n; ++i)
    {
        xma = x[i] - a;
        xma2 = mysquare(xma);
        b2mxma2 = b2 - xma2;
        b2pxma2 = b2 + xma2;
        ib2pxma2 = 1.0 / b2pxma2;
        sqr_b2pxma2 = mysquare(b2pxma2);

        fvec[0] += (xma * ib2pxma2);
        fvec[1] += ib2pxma2;

        fjac[0][0] += (b2mxma2 / sqr_b2pxma2);
        fjac[0][1] += (xma / sqr_b2pxma2);
        fjac[1][1] += (b2mxma2 / sqr_b2pxma2);
    }

    fvec[0] *= 2.0;
    fvec[1] = (double) n / b - (2.0 * b * fvec[1]);

    fjac[0][0] *= -2.0;
    fjac[0][1] *= (-4.0 * b);
    fjac[1][0] = fjac[0][1];
    fjac[1][1] = -2.0 * fjac[1][1] - (n / mysquare(b));
}


int
samesign(double a, double b)
{
    if ((a > 0.0 && b > 0.0) || (a < 0.0 && b < 0.0))
        return (1);
    else
        return (0);
}


int
mnewt_cauchyML(const int ntrials, const double *x, const int n,
               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 < ntrials; ++k)
    {
        evalcauchyML(x, n, 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];
        }

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

        if (errx < tolx)
            break;
    }

    MatDestroy(&fjac);

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