#include "likelihood.h"
#include <stdlib.h> /* for size_t */
#include "pmt.h"
#include <gsl/gsl_integration.h>
#include "muon.h"
#include "misc.h"
#include <gsl/gsl_sf_gamma.h>
#include "sno.h"
#include "vector.h"
#include "event.h"
#include "optics.h"
#include "sno_charge.h"
#include "pdg.h"
#include "path.h"
#include <stddef.h> /* for size_t */
#include "scattering.h"
#include "solid_angle.h"
#include <gsl/gsl_roots.h>
#include <gsl/gsl_errno.h>
#include "pmt_response.h"

typedef struct intParams {
    path *p;
    int i;
} intParams;

double F(double t, double mu_noise, double mu_indirect, double *mu_direct, size_t n, double *ts, double tmean, double sigma)
{
    /* Returns the CDF for the time distribution of photons at time `t`. */
    size_t i;
    double p, mu_total;

    p = mu_noise*t/GTVALID + mu_indirect*(pow(sigma,2)*norm(tmean,t,sigma) + (t-tmean)*norm_cdf(t,tmean,sigma))/(GTVALID-tmean);

    mu_total = mu_noise + mu_indirect;
    for (i = 0; i < n; i++) {
        p += mu_direct[i]*norm_cdf(t,ts[i],sigma);
        mu_total += mu_direct[i];
    }

    return p/mu_total;
}

double f(double t, double mu_noise, double mu_indirect, double *mu_direct, size_t n, double *ts, double tmean, double sigma)
{
    /* Returns the probability that a photon is detected at time `t`.
     *
     * The probability distribution is the sum of three different components:
     * dark noise, indirect light, and direct light. The dark noise is assumed
     * to be constant in time. The direct light is assumed to be a delta
     * function around the times `ts`, where each element of `ts` comes from a
     * different particle.  This assumption is probably valid for particles
     * like muons which don't scatter much, and the hope is that it is *good
     * enough* for electrons too.  The probability distribution for indirect
     * light is assumed to be a step function past some time `tmean`.
     *
     * The probability returned is calculated by taking the sum of these three
     * components and convolving it with a gaussian with standard deviation
     * `sigma` which should typically be the PMT transit time spread. */
    size_t i;
    double p, mu_total;

    p = mu_noise/GTVALID + mu_indirect*norm_cdf(t,tmean,sigma)/(GTVALID-tmean);

    mu_total = mu_noise + mu_indirect;
    for (i = 0; i < n; i++) {
        p += mu_direct[i]*norm(t,ts[i],sigma);
        mu_total += mu_direct[i];
    }

    return p/mu_total;
}

double log_pt(double t, size_t n, double mu_noise, double mu_indirect, double *mu_direct, size_t n2, double *ts, double tmean, double sigma)
{
    /* Returns the first order statistic for observing a PMT hit at time `t`
     * given `n` hits.
     *
     * The first order statistic is computed from the probability distribution
     * above. It's not obvious whether one should take the first order
     * statistic before or after convolving with the PMT transit time spread.
     * Since at least some of the transit time spread in SNO comes from the
     * different transit times across the face of the PMT, it seems better to
     * convolve first which is what we do here. In addition, the problem is not
     * analytically tractable if you do things the other way around. */
    return ln(n) + (n-1)*log1p(-F(t,mu_noise,mu_indirect,mu_direct,n2,ts,tmean,sigma)) + log(f(t,mu_noise,mu_indirect,mu_direct,n2,ts,tmean,sigma));
}

static double gsl_muon_time(double x, void *params)
{
    intParams *pars = (intParams *) params;
    double dir[3], pos[3], n_d2o, n_h2o, wavelength0, T, t, theta0, l_d2o, l_h2o;

    path_eval(pars->p, x, pos, dir, &T, &t, &theta0);

    get_path_length(pos,pmts[pars->i].pos,AV_RADIUS,&l_d2o,&l_h2o);

    /* FIXME: I just calculate delta assuming 400 nm light. */
    wavelength0 = 400.0;
    n_d2o = get_index_snoman_d2o(wavelength0);
    n_h2o = get_index_snoman_h2o(wavelength0);

    t += l_d2o*n_d2o/SPEED_OF_LIGHT + l_h2o*n_h2o/SPEED_OF_LIGHT;

    return t*get_expected_charge(x, T, theta0, pos, dir, pmts[pars->i].pos, pmts[pars->i].normal, PMT_RADIUS);
}

static double gsl_muon_reflected_charge(double x, void *params)
{
    intParams *pars = (intParams *) params;
    double dir[3], pos[3], T, t, theta0;

    path_eval(pars->p, x, pos, dir, &T, &t, &theta0);

    return get_expected_reflected_charge(x, T, theta0, pos, dir, pmts[pars->i].pos, pmts[pars->i].normal, PMT_RADIUS);
}

static double gsl_muon_charge(double x, void *params)
{
    intParams *pars = (intParams *) params;
    double dir[3], pos[3], T, t, theta0;

    path_eval(pars->p, x, pos, dir, &T, &t, &theta0);

    return get_expected_charge(x, T, theta0, pos, dir, pmts[pars->i].pos, pmts[pars->i].normal, PMT_RADIUS);
}

double get_total_charge_approx(double T0, double *pos, double *dir, muon_energy *m, int i, double smax, double theta0, double *t, double *mu_reflected)
{
    /* Returns the approximate expected number of photons seen by PMT `i` using
     * an analytic formula.
     *
     * To come up with an analytic formula for the expected number of photons,
     * it was necessary to make the following approximations:
     *
     * - the index of refraction is constant
     * - the particle track is a straight line
     * - the integral along the particle track is dominated by the gaussian
     *   term describing the angular distribution of the light
     *
     * With these approximations and a few other ones (like using a Taylor
     * expansion for the distance to the PMT), it is possible to pull
     * everything out of the track integral and assume it's equal to it's value
     * along the track where the exponent of the gaussian dominates.
     *
     * The point along the track where the exponent dominates is calculated by
     * finding the point along the track where the angle between the track
     * direction and the PMT is equal to the Cerenkov angle. If this point is
     * before the start of the track, we use the start of the track and if it's
     * past the end of `smax` we use `smax`.
     *
     * Since the integral over the track also contains a term like
     * (1-1/(beta**2*n**2)) which is not constant near the end of the track, it
     * is necessary to define `smax` as the point along the track where the
     * particle velocity drops below some threshold.
     *
     * `smax` is currently calculated as the point where the particle velocity
     * drops to 0.8 times the speed of light. */
    double pmt_dir[3], tmp[3], R, cos_theta, x, z, s, a, b, beta, E, p, T, omega, cos_theta_cerenkov, sin_theta_cerenkov, n_d2o, n_h2o, sin_theta, E0, p0, beta0, f, cos_theta_pmt, absorption_length_h2o, absorption_length_d2o, absorption_length_acrylic, l_h2o, l_d2o, l_acrylic, wavelength0, f_reflected;

    /* Calculate beta at the start of the track. */
    E0 = T0 + MUON_MASS;
    p0 = sqrt(E0*E0 - MUON_MASS*MUON_MASS);
    beta0 = p0/E0;

    /* First, we find the point along the track where the PMT is at the
     * Cerenkov angle. */
    SUB(pmt_dir,pmts[i].pos,pos);
    /* Compute the distance to the PMT. */
    R = NORM(pmt_dir);
    normalize(pmt_dir);

    /* Calculate the cosine of the angle between the track direction and the
     * vector to the PMT at the start of the track. */
    cos_theta = DOT(dir,pmt_dir);
    sin_theta = sqrt(1-pow(cos_theta,2));

    /* Compute the Cerenkov angle at the start of the track. */
    wavelength0 = 400.0;
    n_d2o = get_index_snoman_d2o(wavelength0);
    n_h2o = get_index_snoman_h2o(wavelength0);
    cos_theta_cerenkov = 1/(n_d2o*beta0);
    sin_theta_cerenkov = sqrt(1-pow(cos_theta_cerenkov,2));

    /* Now, we compute the distance along the track where the PMT is at the
     * Cerenkov angle.
     *
     * Note: This formula comes from using the "Law of sines" where the three
     * vertices of the triangle are the starting position of the track, the
     * point along the track that we want to find, and the PMT position. */
    s = R*(sin_theta_cerenkov*cos_theta - cos_theta_cerenkov*sin_theta)/sin_theta_cerenkov;

    /* Make sure that the point is somewhere along the track between 0 and
     * `smax`. */
    if (s < 0) s = 0.0;
    else if (s > smax) s = smax;

    /* Compute the vector from the point `s` along the track to the PMT. */
    tmp[0] = pmts[i].pos[0] - (pos[0] + s*dir[0]);
    tmp[1] = pmts[i].pos[1] - (pos[1] + s*dir[1]);
    tmp[2] = pmts[i].pos[2] - (pos[2] + s*dir[2]);

    /* To do the integral analytically, we expand the distance to the PMT along
     * the track in a Taylor series around `s0`, i.e.
     *
     *     r(s) = a + b*(s-s0)
     *
     * Here, we calculate `a` which is the distance to the PMT at the point
     * `s`. */
    a = NORM(tmp);

    /* `z` is the distance to the PMT projected onto the track direction. */
    z = R*cos_theta;

    /* `x` is the perpendicular distance from the PMT position to the track. */
    x = R*sqrt(1-pow(cos_theta,2));

    /* `b` is the second coefficient in the Taylor expansion. */
    b = (s-z)/a;

    /* Compute the kinetic energy at the point `s` along the track. */
    T = muon_get_energy(s,m);

    /* Calculate the particle velocity at the point `s`. */
    E = T + MUON_MASS;
    p = sqrt(E*E - MUON_MASS*MUON_MASS);
    beta = p/E;

    if (beta < 1/n_d2o) return 0.0;

    /* `prob` is the number of photons emitted per cm by the particle at a
     * distance `s` along the track. */
    double prob = get_probability2(beta);

    /* Compute the position of the particle at a distance `s` along the track. */
    tmp[0] = pos[0] + s*dir[0];
    tmp[1] = pos[1] + s*dir[1];
    tmp[2] = pos[2] + s*dir[2];

    SUB(pmt_dir,pmts[i].pos,tmp);

    cos_theta_pmt = DOT(pmt_dir,pmts[i].normal)/NORM(pmt_dir);

    /* Calculate the sine of the angle between the track direction and the PMT
     * at the position `s` along the track. */
    cos_theta = DOT(dir,pmt_dir)/NORM(pmt_dir);
    sin_theta = sqrt(1-pow(cos_theta,2));

    /* Get the solid angle of the PMT at the position `s` along the track. */
    omega = get_solid_angle_approx(tmp,pmts[i].pos,pmts[i].normal,PMT_RADIUS);

    theta0 = fmax(theta0*sqrt(s),MIN_THETA0);

    double frac = sqrt(2)*n_d2o*x*beta0*theta0;

    f = get_weighted_pmt_response(acos(-cos_theta_pmt));
    f_reflected = get_weighted_pmt_reflectivity(acos(-cos_theta_pmt));

    absorption_length_d2o = get_absorption_length_snoman_d2o(wavelength0);
    absorption_length_h2o = get_absorption_length_snoman_h2o(wavelength0);
    absorption_length_acrylic = get_weighted_absorption_length_snoman_acrylic();

    get_path_length(tmp,pmts[i].pos,AV_RADIUS,&l_d2o,&l_h2o);

    l_acrylic = AV_RADIUS_OUTER - AV_RADIUS_INNER;

    /* Assume the particle is travelling at the speed of light. */
    *t = s/SPEED_OF_LIGHT + l_d2o*n_d2o/SPEED_OF_LIGHT + l_h2o*n_h2o/SPEED_OF_LIGHT;

    double charge = f*exp(-l_d2o/absorption_length_d2o-l_h2o/absorption_length_h2o-l_acrylic/absorption_length_acrylic)*n_d2o*x*beta0*prob*(1/sin_theta)*omega*(erf((a+b*(smax-s)+n_d2o*(smax-z)*beta0)/frac) + erf((-a+b*s+n_d2o*z*beta0)/frac))/(b+n_d2o*beta0)/(4*M_PI);

    *mu_reflected = f_reflected*charge;

    return f*charge;
}

typedef struct betaRootParams {
    muon_energy *m;
    double beta_min;
} betaRootParams;

static double beta_root(double x, void *params)
{
    /* Function used to find at what point along a track a particle crosses
     * some threshold in beta. */
    double T, E, p, beta;
    betaRootParams *pars;

    pars = (betaRootParams *) params;

    T = muon_get_energy(x, pars->m);

    /* Calculate total energy */
    E = T + MUON_MASS;
    p = sqrt(E*E - MUON_MASS*MUON_MASS);
    beta = p/E;

    return beta - pars->beta_min;
}

static int get_smax(muon_energy *m, double beta_min, double range, double *smax)
{
    /* Find the point along the track at which the particle's velocity drops to
     * `beta_min`. */
    int status;
    betaRootParams pars;
    gsl_root_fsolver *s;
    gsl_function F;
    int iter = 0, max_iter = 100;
    double r, x_lo, x_hi;

    s = gsl_root_fsolver_alloc(gsl_root_fsolver_brent);

    pars.m = m;
    pars.beta_min = beta_min;

    F.function = &beta_root;
    F.params = &pars;

    gsl_root_fsolver_set(s, &F, 0.0, range);

    do {
        iter++;
        status = gsl_root_fsolver_iterate(s);
        r = gsl_root_fsolver_root(s);
        x_lo = gsl_root_fsolver_x_lower(s);
        x_hi = gsl_root_fsolver_x_upper(s);

        /* Find the root to within 1 mm. */
        status = gsl_root_test_interval(x_lo, x_hi, 1e-1, 0);

        if (status == GSL_SUCCESS) break;
    } while (status == GSL_CONTINUE && iter < max_iter);

    gsl_root_fsolver_free(s);

    *smax = r;

    return status;
}

double getKineticEnergy(double x, void *params)
{
    return muon_get_energy(x,(muon_energy *) params);
}

double nll_muon(event *ev, double T0, double *pos, double *dir, double t0, double *z1, double *z2, size_t n, double epsrel, int fast)
{
    size_t i, j, nhit;
    intParams params;
    double logp[MAX_PE], nll[MAX_PMTS], range, theta0, E0, p0, beta0, smax, log_mu, max_logp;
    muon_energy *m;
    double pmt_dir[3], R, cos_theta, theta, wavelength0, n_d2o, n_h2o, theta_cerenkov, s, l_h2o, l_d2o;
    double logp_path;
    double a, b;
    double tmp[3];
    double mu_reflected;

    double mu_direct[MAX_PMTS];
    double ts[MAX_PMTS];
    double mu[MAX_PMTS];
    double mu_noise, mu_indirect;

    gsl_integration_cquad_workspace *w = gsl_integration_cquad_workspace_alloc(100);
    double result, error;

    size_t nevals;

    gsl_function F;
    F.params = &params;

    range = get_range(T0, HEAVY_WATER_DENSITY);

    /* Calculate total energy */
    E0 = T0 + MUON_MASS;
    p0 = sqrt(E0*E0 - MUON_MASS*MUON_MASS);
    beta0 = p0/E0;

    /* FIXME: is this formula valid for muons? */
    theta0 = get_scattering_rms(range/2,p0,beta0,1.0,HEAVY_WATER_DENSITY)/sqrt(range/2);

    m = muon_init_energy(T0,HEAVY_WATER_DENSITY,10000);

    params.p = path_init(pos, dir, T0, range, theta0, getKineticEnergy, m, z1, z2, n, MUON_MASS);

    if (beta0 > BETA_MIN)
        get_smax(m, BETA_MIN, range, &smax);
    else
        smax = 0.0;

    mu_indirect = 0.0;
    for (i = 0; i < MAX_PMTS; i++) {
        if (pmts[i].pmt_type != PMT_NORMAL) continue;

        /* Skip PMTs which weren't hit when doing the "fast" likelihood
         * calculation. */
        if (fast && !ev->pmt_hits[i].hit) continue;

        params.i = i;

        if (fast) {
            mu_direct[i] = get_total_charge_approx(T0, pos, dir, m, i, smax, theta0, &ts[i], &mu_reflected);
            ts[i] += t0;
            mu_indirect += mu_reflected;
        } else {
            /* First, we try to compute the distance along the track where the
             * PMT is at the Cerenkov angle. The reason for this is because for
             * heavy particles like muons which don't scatter much, the
             * probability distribution for getting a photon hit along the
             * track looks kind of like a delta function, i.e. the PMT is only
             * hit over a very narrow window when the angle between the track
             * direction and the PMT is *very* close to the Cerenkov angle
             * (it's not a perfect delta function since there is some width due
             * to dispersion). In this case, it's possible that the numerical
             * integration completely skips over the delta function and so
             * predicts an expected charge of 0.
             *
             * To fix this, we only integrate 1 meter on both sides of the
             * point along the track where the PMT is at the Cerenkov angle. */

            /* First, we find the point along the track where the PMT is at the
             * Cerenkov angle. */
            SUB(pmt_dir,pmts[i].pos,pos);
            /* Compute the distance to the PMT. */
            R = NORM(pmt_dir);
            normalize(pmt_dir);

            /* Calculate the cosine of the angle between the track direction and the
             * vector to the PMT at the start of the track. */
            cos_theta = DOT(dir,pmt_dir);
            /* Compute the angle between the track direction and the PMT. */
            theta = acos(cos_theta);

            /* Compute the Cerenkov angle at the start of the track. */
            wavelength0 = 400.0;
            n_d2o = get_index_snoman_d2o(wavelength0);
            theta_cerenkov = acos(1/(n_d2o*beta0));

            /* Now, we compute the distance along the track where the PMT is at the
             * Cerenkov angle.
             *
             * Note: This formula comes from using the "Law of sines" where the three
             * vertices of the triangle are the starting position of the track, the
             * point along the track that we want to find, and the PMT position. */
            s = R*sin(theta_cerenkov-theta)/sin(theta_cerenkov);

            /* Make sure that the point is somewhere along the track between 0 and
             * `smax`. */
            if (s < 0) s = 0.0;
            else if (s > smax) s = smax;

            /* Compute the endpoints of the integration. */
            a = s - 100.0;
            b = s + 100.0;

            if (a < 0.0) a = 0.0;
            if (b > range) b = range;

            F.function = &gsl_muon_charge;

            gsl_integration_cquad(&F, a, b, 0, epsrel, w, &result, &error, &nevals);
            mu_direct[i] = result;

            F.function = &gsl_muon_reflected_charge;

            gsl_integration_cquad(&F, a, b, 0, epsrel, w, &result, &error, &nevals);
            mu_indirect += result;

            F.function = &gsl_muon_time;

            if (mu_direct[i] > 1e-9) {
                gsl_integration_cquad(&F, a, b, 0, epsrel, w, &result, &error, &nevals);
                ts[i] = t0 + result/mu_direct[i];
            } else {
                /* If the expected number of PE is very small, our estimate of
                 * the time can be crazy since the error in the total charge is
                 * in the denominator. Therefore, we estimate the time using
                 * the same technique as in get_total_charge_approx(). See that
                 * function for more details. */
                n_h2o = get_index_snoman_h2o(wavelength0);

                /* Compute the position of the particle at a distance `s` along the track. */
                tmp[0] = pos[0] + s*dir[0];
                tmp[1] = pos[1] + s*dir[1];
                tmp[2] = pos[2] + s*dir[2];

                SUB(pmt_dir,pmts[i].pos,tmp);

                get_path_length(tmp,pmts[i].pos,AV_RADIUS,&l_d2o,&l_h2o);

                /* Assume the particle is travelling at the speed of light. */
                ts[i] = t0 + s/SPEED_OF_LIGHT + l_d2o*n_d2o/SPEED_OF_LIGHT + l_h2o*n_h2o/SPEED_OF_LIGHT;
            }
        }
    }

    path_free(params.p);

    muon_free_energy(m);

    gsl_integration_cquad_workspace_free(w);

    mu_noise = DARK_RATE*GTVALID*1e-9;
    mu_indirect *= CHARGE_FRACTION/10000.0;

    for (i = 0; i < MAX_PMTS; i++) {
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL) continue;

        /* Skip PMTs which weren't hit when doing the "fast" likelihood
         * calculation. */
        if (fast && !ev->pmt_hits[i].hit) continue;

        if (fast)
            mu[i] = mu_direct[i] + mu_noise;
        else
            mu[i] = mu_direct[i] + mu_indirect + mu_noise;
    }

    nhit = 0;
    for (i = 0; i < MAX_PMTS; i++) {
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL) continue;

        /* Skip PMTs which weren't hit when doing the "fast" likelihood
         * calculation. */
        if (fast && !ev->pmt_hits[i].hit) continue;

        log_mu = log(mu[i]);

        if (ev->pmt_hits[i].hit) {
            for (j = 1; j < MAX_PE; j++) {
                logp[j] = log(pq(ev->pmt_hits[i].qhs,j)) - mu[i] + j*log_mu - lnfact(j) + log_pt(ev->pmt_hits[i].t, j, mu_noise, mu_indirect, &mu_direct[i], 1, &ts[i], ts[i], PMT_TTS);
                if (j == 1 || logp[j] > max_logp) max_logp = logp[j];

                if (logp[j] - max_logp < MIN_RATIO*ln(10)) {
                    j++;
                    break;
                }
            }

            nll[nhit++] = -logsumexp(logp+1, j-1);
        } else {
            logp[0] = -mu[i];
            if (fast) {
                nll[nhit++] = -logp[0];
                continue;
            }
            for (j = 1; j < MAX_PE_NO_HIT; j++) {
                logp[j] = get_log_pmiss(j) - mu[i] + j*log_mu - lnfact(j);
            }
            nll[nhit++] = -logsumexp(logp, MAX_PE_NO_HIT);
        }
    }

    logp_path = 0.0;
    for (i = 0; i < n; i++)
        logp_path += log(norm(z1[i],0,1)) + log(norm(z2[i],0,1));

    return kahan_sum(nll,nhit) + logp_path;
}