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


/* The rice distribution has the form
*/
double
rice_dev(const double mean, const double var, const gsl_rng *r2)
{
    double mean2 = M_SQRT1_2 * mean;
    double sigma = sqrt(var);
    double var1 = gsl_ran_gaussian_ziggurat(r2, sigma) + mean2;
    double var2 = gsl_ran_gaussian_ziggurat(r2, sigma) + mean2;

    return (sqrt(var1*var1 + var2*var2));
}


///////////////////////////////////////
///////////////////////////////////////
//  NB: THE REST OF THIS BELOW IS BUNK!
//  it is currently just filler (based on gaussian)
///////////////////////////////////////
///////////////////////////////////////


double
rice_pdf(const double x, const double mean, const double var)
{
    double          p, term;

    if (var == 0.0)
    {
        if (x == mean)
            return(1.0);
        else
            return(0.0);
    }

    term = x - mean;
    p = (1.0 / sqrt(2.0 * MY_PI * var)) * exp(-term * term / (2.0 * var));

    return (p);
}


double
rice_lnpdf(const double x, const double mean, const double var)
{
    double          p;

    p = (-0.5 * log(2.0 * MY_PI * var)) - (mysquare(x - mean) / (2.0 * var));

    return (p);
}


double
rice_cdf(const double x, const double mean, const double var)
{
    double          p;

    p = 0.5 * (1.0 + erf((x - mean) / (sqrt(2.0 * var))));

    return (p);
}


double
rice_sdf(const double x, const double mean, const double var)
{
    double          p;

    p = 0.5 * erfc((x - mean) / (sqrt(2.0 * var)));

    return (p);
}


double
rice_int(const double x, const double y, const double mean, const double var)
{
    double          p, xn, yn;

    xn = (x - mean) / (sqrt(2.0 * var));
    yn = (y - mean) / (sqrt(2.0 * var));

    if ((x + y) < 2.0 * mean)
        p = 0.5 * (erf(yn) - erf(xn));
    else
        p = 0.5 * (erfc(xn) - erfc(yn));

    return (p);
}


double
rice_logL(const double mean, const double var)
{
    return(-log(sqrt(var * 2.0 * MY_PI * MY_E)));
}


/* maximum likelihood fit of data to a rice distribution */
/* screw bias, it's overblown */
double
rice_fit(const double *x, const int n, double *mean, double *var, double *logL)
{
    double           nd = (double) n;
    double           ave, variance, verror;
    int              i;

    ave = 0.0;
    for (i = 0; i < n; ++i)
        ave += x[i];
    ave /= nd;

    variance = verror = 0.0;
    for (i = 0; i < n; ++i)
    {
        variance += mysquare(x[i] - ave);
        verror += (x[i] - ave);
    }

    verror = mysquare(verror) / nd;
    variance = (variance - verror) / nd; /* more accurate corrected two-pass algorithm */
                                         /* with no rounding error, verror = 0 */

    *mean = ave;
    *var = variance;
    /* printf("\n\nrice logL %e\n", rice_logL(*mean, *var)); */

    return(chi_sqr_adapt(x, n, 0, logL, *mean, *var, rice_pdf, rice_lnpdf, rice_int));
}


double
rice_fit_w(const double *x, const int n, const double *wts, double *mean,
             double *var, double *logL)
{
    double           ave, variance, tmp, wtsum;
    int              i;

    ave = wtsum = 0.0;
    for (i = 0; i < n; ++i)
    {
        ave += wts[i] * x[i];
        wtsum += wts[i];
    }
    ave /= wtsum;

    variance = 0.0;
    for (i = 0; i < n; ++i)
    {
        tmp = (x[i] - ave);
        variance += wts[i] * tmp * tmp;
    }

    variance /= wtsum;

    *mean = ave;
    *var = variance;

    return(1.0);
}


void
rice_init_mix_params(const double *x, const int n, const int mixn, double *mean,
                       double *var)
{
    int             j;
    double          imean, ivar, ilogL;

    rice_fit(x, n, &imean, &ivar, &ilogL);

    for (j = 0; j < mixn; ++j)
    {
        mean[j] = imean - (2.0 * sqrt(ivar)) + (4.0 * sqrt(ivar) * (j + 1.0) / (1.0 + mixn));
        var[j] = ivar / mixn;
    }
}

