/*
    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_rng.h>
#include <gsl/gsl_randist.h>
#include "DLTmath.h"
#include "statistics.h"
#include "logistic_dist.h"

/* The Logistic probability distribution is

   p(x) dx = 
   
   a = location param
   b = shape param
*/

double
logistic_dev(const double a, const double b, const gsl_rng *r2)
{
    double         x;

    x = gsl_rng_uniform(r2);

    return(a + (b * log(x / (1.0 - x))));
}


double
logistic_pdf(const double x, const double a, const double b)
{
    double         u, p;

    u = exp((x - a) / b);

    p = u / (b * (1.0 + u)*(1.0 + u));

    return(p);
}


double
logistic_lnpdf(const double x, const double a, const double b)
{
    return(log(logistic_pdf(x, a, b)));
}


double
logistic_cdf(const double x, const double a, const double b)
{
    return(1.0 - 1.0 / (1.0 + exp((a - x) / b)));
}


double
logistic_sdf(const double x, const double a, const double b)
{
    return(1.0 / (1.0 + exp((a - x) / b)));
}


double
logistic_int(const double x, const double y, const double a, const double b)
{
    return(logistic_cdf(x, a, b) - logistic_cdf(y, a, b));
}


double
logistic_logL(const double a, const double b)
{
    return(-log(b) - 2.0);
}


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

    /* guess at parameters by method of moments */
    a = 0.0;
    for (i = 0; i < n; ++i)
        a += x[i];
    a /= (double) n;

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

    b = sqrt(3.0 * b / ((double) n * MY_PI * MY_PI));

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

    //printf(" Logistic [a,b]: %e %e\n", a, b);

    mnewt_logisticML(100, x, n, guesses, 1e-9, 1e-9);

    *ra = guesses[0];
    *rb = guesses[1];
    
    /* printf(" Logistic logL: %f\n", logistic_logL(*ra, *rb)); */

    return(chi_sqr_adapt(x, n, 0, logL, *ra, *rb, logistic_pdf, logistic_lnpdf, logistic_int));
}


/* Maximum likelihood fit for the Logistic.
   Based on my own calculations of the derivatives.
   We must solve the simultaneous equations:

       F_0 = \Sum{tanh((a-x)/(2b))} = 0
       F_1 = \Sum{(a-x) * tanh((a-x)/(2b))} - Nb = 0

   when these are zero the 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 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:

       J[0][0] = dF_0/da = (1/b) * \Sum[ 1/(1 + cosh{(a-x)/b}) ]

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

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

                           -1               (a-x)^2
       J[1][1] = dF_1/db = --- * \Sum ------------------- - N
                           b^2        (1 + cosh((x-a)/b))
*/
void
evallogisticML(const double *x, const int n,
               double *params, double *fvec, double **fjac)
{
    int            i;
    double         a, b, amxb, /* exp_xmab,*/ cosh_amxb, nd, xi, diff;

    a = params[0];
    b = params[1];
    nd = (double) n;

    fvec[0] = fvec[1] = fjac[0][0] = fjac[1][0] = fjac[0][1] = fjac[1][1] = 0.0;
    for (i = 0; i < n; ++i)
    {
        xi = x[i];
        /* xmab = (x[i] - a) / b; */
        diff = a - xi;
        amxb = diff / b;
        cosh_amxb = cosh(amxb);

        fvec[0] += tanh(amxb/2.0); /* good */
        fvec[1] += diff * tanh(amxb/2.0); /* good */

        fjac[0][0] += (1.0 / (1.0 + cosh(amxb))); /* good */
        fjac[0][1] += (diff / (1.0 + cosh_amxb)); /* good */
        fjac[1][0] += (diff + b * sinh(amxb)) / (1.0 + cosh_amxb); /* good */
        fjac[1][1] += diff * diff / (1.0 + cosh_amxb); /* good */
    }

    fvec[0] = fvec[0]; /* good */
    fvec[1] -= (b * nd); /* good */

    fjac[0][0] /= b; /* good */
    fjac[0][1] /= b*b; /* good */
    fjac[1][0] /= b; /* good */
    fjac[1][1] = -fjac[1][1]/(b*b) - nd; /* good */
}


void
mnewt_logisticML(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], p[2];
    double          **fjac = NULL;

    /* 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)
    {
        evallogisticML(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 Logistic: iter[%4d]  a=%12.6e b=%12.6e", k, params[0], params[1]);
        //fflush(NULL);

        if (errx < tolx)
            break;
    }

    MatDestroy(&fjac);
}


double
sech(const double d)
{
    return(1.0 / cosh(d));
}
