#include "sno_charge.h"
#include <gsl/gsl_sf_gamma.h>
#include <math.h>
#include <gsl/gsl_spline.h>
#include <gsl/gsl_integration.h>
#include <unistd.h>
#include <stdio.h>
#include "misc.h"

#define MAX_PE 10

double qlo = -1.0;
double qhi = 100.0;
size_t nq = 10000;
static gsl_spline *splines[MAX_PE];
static gsl_interp_accel *acc;

static double qmean, qstd;

static double pmiss[MAX_PE];

static int initialized = 0;

/* Parameters for the single PE charge distribution. These numbers come from
 * mc_generator.dat in SNOMAN 5.0294. */

/* Polya parameter. */
static double M = 8.1497;
/* Amplitude of 1st Polya. */
static double NG1 = 0.27235e-1;
/* Slope of Polya Gain 1 vs. high-half point. */
static double SLPG1 = 0.71726;
/* Offset of Polya Gain 1 vs. high-half point. */
static double OFFG1 = 1.2525;
/* Amplitude of 2nd Polya. */
static double NG2 = 0.7944e-3;
/* Slope of Polya Gain 2 vs. high-half point. */
static double SLPG2 = 0.71428;
/* Offset of Polya Gain 2 vs. high-half point. */
static double OFFG2 = 118.21;
/* Amplitude of exponential noise peak. */
static double NEXP = 0.81228e-1;
/* Falloff parameter of exponential noise peak. */
static double Q0 = 20.116;
/* Mean of high-half point distribution. */
static double MEAN_HIPT = 46.0;
/* Width of noise before discriminator in ADC counts. */
static double QNOISE = 0.61;
/* Width of smearing after discriminator in ADC counts. */
static double QSMEAR_ADC = 3.61;
/* Mean of simulated threshold distribution. */
static double MEAN_THRESH = 8.5;

double spe_pol2exp(double q)
{
    /* Returns the probability distribution to get a charge `q` on a PMT with a
     * high-half point `hhp` from a single photoelectron.
     *
     * We assume that the charge has been normalized by the high-half point.
     *
     * This function models the single PE charge distribution as a double Polya
     * with an exponential noise term. */
    double qpe1, qpe2, funpol1, funpol2, gmratio, g1, g2, norm, hhp;

    if (q < 0) return 0.0;

    hhp = MEAN_HIPT;

    q *= MEAN_HIPT;

    g1 = OFFG1 + SLPG1*hhp;
    g2 = OFFG2 + SLPG2*hhp;

    qpe1 = q/g1;
    qpe2 = q/g2;

    funpol1 = pow(M*qpe1,M-1)*exp(-M*qpe1);
    funpol2 = pow(M*qpe2,M-1)*exp(-M*qpe2);

    gmratio = M/gsl_sf_gamma(M);

    norm = (NG1*g1 + NG2*g2 + NEXP*Q0)/MEAN_HIPT;

    return (gmratio*(NG1*funpol1 + NG2*funpol2) + NEXP*exp(-q/Q0))/norm;
}

double pq(double q, int n)
{
    if (!initialized) {
        fprintf(stderr, "charge interpolation hasn't been initialized!\n");
        exit(1);
    }

    if (n > MAX_PE) {
        /* Assume the distribution is gaussian by the central limit theorem. */
        return norm(q,n*qmean,qstd);
    }

    if (q < qlo || q > qhi) return 0.0;

    return gsl_spline_eval(splines[n-1], q, acc);
}

double get_pmiss(int n)
{
    if (!initialized) {
        fprintf(stderr, "charge interpolation hasn't been initialized!\n");
        exit(1);
    }

    if (n == 0) {
        return 1.0;
    } else if (n > MAX_PE) {
        /* Assume the distribution is gaussian by the central limit theorem. */
        return 0.0;
    }

    return pmiss[n-1];
}

static double gsl_charge(double x, void *params)
{
    double q = ((double *) params)[0];
    int n = (int) ((double *) params)[1];

    return pq(x,1)*pq(q-x,n-1);
}

static double gsl_charge2(double x, void *params)
{
    int n = (int) ((double *) params)[0];

    return pq(x,n);
}

static double gsl_smear(double x, void *params)
{
    double q = ((double *) params)[0];
    int n = (int) ((double *) params)[1];

    return pq(x,n)*norm(q-x,0.0,QSMEAR_ADC/MEAN_HIPT);
}

static double sno_charge1(double q, void *params)
{
    return q*spe_pol2exp(q);
}

static double sno_charge2(double q, void *params)
{
    return q*q*spe_pol2exp(q);
}

void init_charge(void)
{
    int i, j;
    gsl_integration_cquad_workspace *w;
    double result, error;
    size_t nevals;
    double *x, *y;
    double params[2];
    gsl_function F;
    double q, q2;

    initialized = 1;

    x = malloc(nq*sizeof(double));
    y = malloc(nq*sizeof(double));

    for (i = 0; i < nq; i++) {
        x[i] = qlo + (qhi-qlo)*i/(nq-1);
    }

    w = gsl_integration_cquad_workspace_alloc(100);

    F.function = &sno_charge1;
    F.params = NULL;

    gsl_integration_cquad(&F, 0, qhi, 0, 1e-9, w, &result, &error, &nevals);
    q = result;

    F.function = &sno_charge2;
    F.params = NULL;

    gsl_integration_cquad(&F, 0, qhi, 0, 1e-9, w, &result, &error, &nevals);
    q2 = result;

    qmean = q;
    qstd = sqrt(q2 - q*q);

    F.function = &gsl_charge;
    F.params = &params;

    acc = gsl_interp_accel_alloc();

    splines[0] = gsl_spline_alloc(gsl_interp_linear,nq);
    for (j = 0; j < nq; j++) {
        y[j] = spe_pol2exp(x[j]);
    }
    gsl_spline_init(splines[0],x,y,nq);

    for (i = 2; i <= MAX_PE; i++) {
        splines[i-1] = gsl_spline_alloc(gsl_interp_linear,nq);
        for (j = 0; j < nq; j++) {
            params[0] = x[j];
            params[1] = i;
            if (x[j] < 0) {
                y[j] = 0.0;
                continue;
            }
            gsl_integration_cquad(&F, 0, x[j], 0, 1e-4, w, &result, &error, &nevals);
            y[j] = result;
        }
        gsl_spline_init(splines[i-1],x,y,nq);
    }

    /* Integrate the charge distribution before smearing to figure out the
     * probability that the discriminator didn't fire.
     *
     * Note: Technically there is some smearing before the discriminator, but
     * it is very small, so we ignore it. */
    F.function = &gsl_charge2;
    for (i = 1; i <= MAX_PE; i++) {
        params[0] = i;
        gsl_integration_cquad(&F, 0, MEAN_THRESH/MEAN_HIPT, 0, 1e-9, w, &result, &error, &nevals);
        pmiss[i-1] = result;
    }

    F.function = &gsl_smear;

    for (i = 1; i <= MAX_PE; i++) {
        for (j = 0; j < nq; j++) {
            params[0] = x[j];
            params[1] = i;
            gsl_integration_cquad(&F, x[j]-QSMEAR_ADC*10/MEAN_HIPT, x[j]+QSMEAR_ADC*10/MEAN_HIPT, 0, 1e-4, w, &result, &error, &nevals);
            y[j] = result;
        }
        gsl_spline_free(splines[i-1]);
        splines[i-1] = gsl_spline_alloc(gsl_interp_linear,nq);
        gsl_spline_init(splines[i-1],x,y,nq);
    }

    gsl_integration_cquad_workspace_free(w);
    free(x);
    free(y);
}