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

/* The von Mises (circular normal) distribution has the form

   p(x) dx = 1/(2 pi I_0(b)) exp(b cos[x-a])) dx

   where I_0(b) is the modified Bessel function of the first kind.

   variate 0 < x <= 2 pi
   scale param b > 0
   location param 0 < a < 2 pi
*/


/*
static double
pseudo_vonmises_target(const double a, const double b, const double x)
{
    return(a * cos(x) + b * cos(x));
}
 */



/*
static double
mardia_gadsden_target(const double a, const double b, const double x)
{
    return((a * cos(x) + b * cos(x)) * cos(x));
}
 */



/*
double
pseudo_vonmises_met(const double a, const double b, const double x)
{
    double          r, y, u, width = 0.1;
    unsigned long   seed = (unsigned long) time(NULL);

    init_genrand(seed);

    y = x + (2.0 * width * genrand_real3() - width);
    r = pseudo_vonmises_target(a, b, y) / pseudo_vonmises_target(a, b, x);
    u = genrand_real3();

    if (r > u)
        return(y);
    else
        return(x);
}
 */



/*
double
mardia_gadsden_met(const double a, const double b, const double x)
{
    double          r, y, u, width = 0.1;
    unsigned long   seed = (unsigned long) time(NULL);

    init_genrand(seed);

    y = x + (2.0 * width * genrand_real3() - width);
    r = mardia_gadsden_target(a, b, y) / mardia_gadsden_target(a, b, x);
    u = genrand_real3();

    if (r > u)
        return(y);
    else
        return(x);
}
 */



static double
wrap_0_2PI(double x)
{
    while(x < 0.0)
        x += 2.0*MY_PI;

    while(x > 2.0*MY_PI)
        x -= 2.0*MY_PI;

    return(x);
}


/*
double
vonmises_dev2(const double k)
{
    double rho, theta, r, s, u, v, z, w, y;
    int accept;
    unsigned long   seed = (unsigned long) time(NULL);

    init_genrand(seed);

    r = 1.0 + sqrt(1.0 + 4.0 * k*k);
    rho = (r - sqrt(2.0 *r) / (2.0 * k));
    s = (1.0 + rho*rho) / (2.0 * rho);

    accept = 0;
    while (accept == 0)
    {
        u = 2.0 * genrand_real3() - 1.0;
        v = 2.0 * genrand_real3() - 1.0;

        z = cos(MY_PI * u);
        w = (1.0 + s*z) / (s + z);
        y = k * (s - w);

        if (w * (2.0 - w) - v > DBL_MIN)
            accept = 1;
        else if (log(w/v) + 1.0 > DBL_MIN)
            accept = 1;
    }

    if (u < 0)
        theta = -1.0 / cos(w);
    else
        theta = 1.0 / cos(w);

    return(theta);
}
 */


double
vonmises_dev2(const double mean, const double k, const gsl_rng *r2)
{
    double result = 0.0;

    double a = 1.0 + sqrt(1 + 4.0 * (k * k));
    double b = (a - sqrt(2.0 * a))/(2.0 * k);
    double r = (1.0 + b * b)/(2.0 * b);

    while (1)
    {
        double U1 = gsl_ran_flat(r2, 0.0, 1.0);
        double z = cos(MY_PI * U1);
        double f = (1.0 + r * z)/(r + z);
        double c = k * (r - f);
        double U2 = gsl_ran_flat(r2, 0.0, 1.0);

        if (c * (2.0 - c) - U2 > 0.0)
        {
            double U3 = gsl_ran_flat(r2, 0.0, 1.0);
            double sign = 0.0;

            if (U3 - 0.5 < 0.0)
                sign = -1.0;

            if (U3 - 0.5 > 0.0)
                sign = 1.0;

            result = sign * acos(f) + mean;

            while (result >= 2.0 * MY_PI)
                result -= 2.0 * MY_PI;

            break;
        }
        else
        {
            if(log(c/U2) + 1.0 - c >= 0.0)
            {
                double U3 = gsl_ran_flat(r2, 0.0, 1.0);
                double sign = 0.0;

                if (U3 - 0.5 < 0.0)
                    sign = -1.0;

                if (U3 - 0.5 > 0.0)
                    sign = 1.0;

                result = sign * acos(f) + mean;

                while (result >= 2.0 * MY_PI)
                    result -= 2.0 * MY_PI;

                break;
            }
        }
    }

    return(result);
}


/* D. J. Best and N. I. Fisher (1979)
   "Efficient simulation of the von Mises distribution."
   Applied Statistics 28:152Ð157. */
double
vonmises_dev(const double a, const double b, const gsl_rng *r2)
{
    double          z, f, c, t, p, r;

    if (b < DBL_EPSILON)
        return(a);

    t = 1.0 + sqrt(1.0 + 4.0*b*b);
    p = 0.5*(t - sqrt(2.0*t))/b;
    r = 0.5*(1.0 + p*p)/p;

    do
    {
        z = cos(MY_PI*gsl_rng_uniform(r2));
        f = (1.0 + r*z)/(r + z);
        c = b*(r-f);
    }
    while(log(c/gsl_rng_uniform(r2)) + 1.0 < c);

    if (gsl_rng_uniform(r2) > 0.5)
        return(wrap_0_2PI(a + acos(f)));
    else
        return(wrap_0_2PI(a - acos(f)));
}


double
vonmises_pdf(const double x, const double a, const double b)
{
    if (b <= 1e-8)
        return(1.0/(2.0*MY_PI));
    else
        return((1.0/(2.0*MY_PI*BesselI0(b))) * exp(b*cos(x-a)));
}


double
vonmises_lnpdf(const double x, const double a, const double b)
{
    if (b <= 1e-8)
        return(-log(2.0*MY_PI));
    else
        return(b*cos(x-a) - log(2.0*MY_PI) - log(BesselI0(b)));
}


/* DLT debug - something's wrong */
double
vonmises_cdf(const double x, const double a, const double b)
{
    double          tol = FLT_EPSILON;
    double          sum, dterm;
    int             j;

    sum = 0.0;
    j = 1;
    do
    {
        dterm = BesselI(j, b) * sin(j*(x-a)) / (double) j;
        sum += dterm;
        printf("% f % f % f % f % f % f, %3d\n", a, x, BesselI(j,b), sin(j*(x-a)), dterm, sum, j);
        fflush(NULL);
        j++;
    }
    while(fabs(dterm) > tol && j < 100);
    printf("vonMises iters=%d\n", j);

    /* return((x*BesselI0(b) + 2.0 * dterm)/(2.0 * MY_PI * BesselI0(b))); */
    return((x + 2.0 * dterm / BesselI0(b)) / (2.0 * MY_PI));
}


double
vonmises_sdf(const double x, const double a, const double b)
{
    return(1.0 - vonmises_cdf(x, a, b));
}


double
vonmises_int(const double x, const double y, const double a, const double b)
{
    /* return(vonmises_cdf(y, a, b) - vonmises_cdf(x, a, b)); */
    return(integrate_romberg(vonmises_pdf, a, b, x, y));
}


/* This is correct. */
double
vonmises_logL(const double a, const double b)
{
    double          i0b = BesselI0(b);

    return(b*BesselI1(b)/i0b - log(2.0*MY_PI*i0b));
}


/* For the maximum likelihood fit we nust find the root of:

       F1 = I_0(b) \Sum{cos(x-a)} - N I_1(b) = 0

   where the first derivative with repect to b (dF1/db) is:

       F1' = I_0(b) [\Sum{cos(x-a)} - N] - (N/2) I_1(b)
*/
void
evalvonmisesML(const double cosdif, const double b, const int num,
                   double *fx, double *dfx)
{
    double          i0b = BesselI0(b);
    double          i1b = BesselI1(b);

    *fx  = i0b*cosdif - num*i1b;
    *dfx = i1b*(cosdif - 0.5*num) - num*i0b;
}


void
evalvonmisesML_EHP(const double sincosterm, const double b, const int num,
               double *fx, double *dfx)
{
    double          i0b = BesselI0(b);
    double          i1b = BesselI1(b);

    *fx  = num*i1b - sincosterm*i0b;
    *dfx = num*i0b - (num/b + sincosterm)*i1b;
}


static void
evalvonmisesMLmu(const double *data, const double a, const int num,
                 double *fx, double *dfx)
{
    double          sinsum, cossum;
    int             i;

    sinsum = 0.0;
    for (i = 0; i < num; ++i)
        sinsum += sin(data[i] - a);

    cossum = 0.0;
    for (i = 0; i < num; ++i)
        cossum += cos(data[i] - a);

    *fx  = sinsum;
    *dfx = -cossum;
}


/* Fit a vonmises distribution by maximum likelihood. */
double
vonmises_fit(const double *data, const int num, double *a, double *b, double *prob)
{
    double          ave, cosdif, cossum, sinsum, fx, dfx, fxdfx, chi2, bguess;
    double          kappa, sincosterm;
    int             i;
    double          iter = 1000;
    double          tol = 1e-11;
    double         *x = malloc(num * sizeof(double));

    cossum = sinsum = ave = 0.0;
    for (i = 0; i < num; ++i)
    {
        /* x[i] = fmod(data[i], 2.0*MY_PI); */
/*         if (x[i] < 0.0) */
/*             x[i] += 2.0*MY_PI; */
        /* printf("\n %e", x[i]); */
        x[i] = data[i];

        ave += x[i];
        cossum += cos(x[i]);
        sinsum += sin(x[i]);
    }

    ave /= num;

    *a = atan(sinsum/cossum);
    /* printf("\natan:  % f % f % f % f % f", sinsum, cossum, sinsum/cossum, *a, ave); */
    *a = atan2(sinsum,cossum);
    /* printf("\natan2: % f % f % f % f % f", sinsum, cossum, sinsum/cossum, *a, ave); */

    if (*a < 0.0)
        *a += 2.0*MY_PI;

    if (fabs((ave - *a)/ave) > 50.0)
    {
        printf("% f % f % f % f\n", cossum, sinsum, *a, ave);
        *a = ave;
    }

    /* For a (mu), well, we use my own derivation that I've never seen
       anywhere but it seems straight-forward to me.
       Evans, Hastings, and Peacock 2000 (_Statistical Distributions_) give
       the MLE as:

           \hat{mu} = arctan[\Sum{sin x_i}/\Sum{cos x_i}]

       but this has problems, it gives a negative answer for real
       mu = 3, 4 etc. Part of the problem is fixed by using atan2,
       which computes the correct quadrant, but it still doesn't
       always work. So, I simply did the standard MLE derivation --
       take the log of the PDF, take the first derivative, set it
       to zero, and you get:

            \Sum{sin(x - mu)} = 0

        which can't be solved analytically, but can easily be solved
        with Newton-Raphson using the second derivative correction:

            -\Sum{cos(x - mu)}

        It seems to work well. */

    fx = fxdfx = 0.0;
    dfx = 1.0;
    *a = ave;
    for (i = 0; i < iter; ++i)
    {
        /* printf("\n% e % e % e % e", *a, fx, dfx, fx/dfx); */
        evalvonmisesMLmu(x, *a, num, &fx, &dfx);

        fxdfx = fx/dfx;

        if (i > 3 && (fabs(fx) < tol || fabs(fxdfx) < tol))
            break; /* success */

        *a -= fxdfx; /* Newton-Raphson correction */
    }

    while (*a < 0.0)
        *a += 2.0*MY_PI;

    while (*a > 2.0*MY_PI)
        *a -= 2.0*MY_PI;

    /* cossin = sqrt(sinsum*sinsum + cossum*cossum); */

    /* For b (kappa), we first use the estimate recommended by
       Best and Fisher (1981),
       then refine it nicely with Newton-Raphson ML.
       The ML estimator is actually more biased. */

    /* Best, D. and Fisher N. (1981).
       "The bias of the maximum likelihood estimators of the von Mises-Fisher
       concentration parameters."
       Communications in Statistics - Simulation and Computation
       B10(5), 493Ð502.
    */
    sincosterm = sqrt(cossum*cossum + sinsum*sinsum);

    cosdif = 0.0;
    for (i = 0; i < num; ++i)
        cosdif += cos(x[i] - *a);

    if (cosdif < 0.0)
        cosdif += MY_PI;
    kappa = cosdif / num;

    if (kappa < 0.53)
        *b = bguess = 2.0*kappa + kappa*kappa*kappa
             + (5.0*mysquare(kappa*kappa)*kappa)/6.0;
    else if(kappa < 0.85)
        *b = bguess = -0.4 +1.39*kappa + 0.43/(1.0 - kappa);
    else
        *b = bguess = 1.0 / (mycube(kappa) - 4.0*kappa*kappa) + 3.0*kappa;

    fx = fxdfx = 0.0;
    dfx = 1.0;
    for (i = 0; i < iter; ++i)
    {
        /* printf("\n% e % e % e % e", *b, fx, dfx, fx/dfx); */
        evalvonmisesML_EHP(/* cosdif */sincosterm, *b, num, &fx, &dfx);

        fxdfx = fx/dfx;

        if (i > 3 && (fabs(fx) < tol || fabs(fxdfx) < tol))
            break; /* success */

        *b -= fxdfx; /* Newton-Raphson correction */

        if (*b < DBL_EPSILON)
            *b = DBL_EPSILON;
    }

    if (i == iter)
        *b = bguess;

    /* printf("\n von Mises logL:% e", vonmises_logL(*a, *b)); */

    chi2 = chi_sqr_adapt(x, num, 0, prob, *a, *b, vonmises_pdf, vonmises_lnpdf, vonmises_int);

    free(x);

    return(chi2);
}
