#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 log(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], pmt_dir[3], R, n, wavelength0, T, t, theta0, distance;

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

    R = NORM(pos);

    SUB(pmt_dir,pmts[pars->i].pos,pos);

    distance = NORM(pmt_dir);

    /* FIXME: I just calculate delta assuming 400 nm light. */
    wavelength0 = 400.0;
    if (R <= AV_RADIUS) {
        n = get_index_snoman_d2o(wavelength0);
    } else {
        n = get_index_snoman_h2o(wavelength0);
    }

    t += distance*n/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_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, int i, double smax, double theta0, double *t)
{
    /* 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, theta, x, z, s, a, b, beta, E, p, T, omega, theta_cerenkov, n, sin_theta, E0, p0, beta0, f, cos_theta_pmt, absorption_length, wavelength0;

    /* 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);
    /* 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 = get_index_snoman_d2o(wavelength0);
    theta_cerenkov = acos(1/(n*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 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);

    /* Assume the particle is travelling at the speed of light. */
    *t = s/SPEED_OF_LIGHT + a*n/SPEED_OF_LIGHT;

    /* `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*fabs(sin(theta));

    /* `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 = get_T(T0,s,HEAVY_WATER_DENSITY);

    /* 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) 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. */
    sin_theta = fabs(sin(acos(DOT(dir,pmt_dir)/NORM(pmt_dir))));

    /* 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*x*beta0*theta0;

    f = get_weighted_pmt_response(acos(-cos_theta_pmt));

    absorption_length = get_absorption_length_snoman_d2o(wavelength0);

    return f*exp(-a/absorption_length)*n*x*beta0*prob*(1/sin_theta)*omega*(erf((a+b*(smax-s)+n*(smax-z)*beta0)/frac) + erf((-a+b*s+n*z*beta0)/frac))/(b+n*beta0)/(4*M_PI);
}

double getKineticEnergy(double x, double T0)
{
    return get_T(T0, x, HEAVY_WATER_DENSITY);
}

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, T0, beta_min;

    T0 = ((double *) params)[0];
    beta_min = ((double *) params)[1];

    T = get_T(T0, x, HEAVY_WATER_DENSITY);

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

    return beta - beta_min;
}

static int get_smax(double T0, 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;
    double params[2];
    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);

    params[0] = T0;
    params[1] = beta_min;

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

    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 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 total_charge;
    double logp[MAX_PE], nll[MAX_PMTS], range, theta0, E0, p0, beta0, smax, log_mu, max_logp;
    double tmean = 0.0;

    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);

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

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

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

        params.i = i;

        if (fast) {
            mu_direct[i] = get_total_charge_approx(T0, pos, dir, i, smax, theta0, &ts[i]);
            ts[i] += t0;
        } else {
            F.function = &gsl_muon_charge;
            gsl_integration_cquad(&F, 0, range, 0, epsrel, w, &result, &error, &nevals);
            mu_direct[i] = result;

            ts[i] = t0;
            if (mu_direct[i] > 1e-9) {
                F.function = &gsl_muon_time;
                gsl_integration_cquad(&F, 0, range, 0, epsrel, w, &result, &error, &nevals);
                ts[i] += result/mu_direct[i];
            }
        }

        tmean += ts[i]*mu_direct[i];

        total_charge += mu_direct[i];
    }

    path_free(params.p);

    if (total_charge > 0)
        tmean /= total_charge;

    gsl_integration_cquad_workspace_free(w);

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

    for (i = 0; i < MAX_PMTS; i++) {
        if (ev->pmt_hits[i].flags || (pmts[i].pmt_type != PMT_NORMAL && pmts[i].pmt_type != PMT_OWL)) continue;
        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 && pmts[i].pmt_type != PMT_OWL)) 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], tmean, 1.5);
                if (j == 1 || logp[j] > max_logp) max_logp = logp[j];

                if (logp[j] - max_logp < MIN_RATIO*log(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);
        }
    }

    return kahan_sum(nll,nhit);
}