/* Copyright (c) 2019, Anthony Latorre <tlatorre at uchicago>
 *
 * 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 3 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, see <https://www.gnu.org/licenses/>.
 */

#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"
#include <nlopt.h>

/* Maximum number of PE to compute charge PDFs for.
 *
 * For PE greater than this we compute the PDF using a Gaussian distribution. */
#define MAX_PE 100

/* Maximum number of PE to consider when computing P(q|mu) in nlopt_log_pq().
 * This should be at least as large as the maximum charge we expect to record
 * in the detector which should be approximately (4095-600)/30 ~ 100 for QHS.
 */
#define MAX_PE_LOOKUP 1000

/* Number of points to compute the table for when calculating the most likely
 * mumber of mean PE given an observed charge. */
#define N_CHARGE_MEAN_PE 1000

/* Minimum and maximum charges to compute the table for when calculating the
 * most likely mumber of mean PE given an observed charge. */
static double xlo_mean_pe = -1.0;
static double xhi_mean_pe = 100.0;

/* Minimum charge to compute the PDF at (in units of PE). */
double qlo = -1.0;
/* Maximum charge to compute the PDF at (in units of PE). */
double qhi = 100.0;
/* Number of points in the charge PDF. */
#define nq 10000

/* Charge PDFs. */
static double x[nq];
static double pq[MAX_PE][nq];
/* Log of the charge PDFs used for calculations in the likelihood. */
static double log_pq[MAX_PE][nq];

/* Mean and standard deviation of the single PE charge distribution.
 *
 * This is used to evaluate the probability of a high charge where we don't
 * have a precomputed PDF. We assume the central limit theorem and just compute
 * the probability density as a Gaussian with mean n*qmean and standard
 * deviation sqrt(n)*qstd where `n` is the number of PE. */
static double qmean, qstd;

/* Lookup table for the log of the probability of the discriminator not firing
 * as a function of the number of PE which hit the tube. We compute this in
 * init_charge() by integrating the charge PDF up to the mean threshold for all
 * channels. */
static double log_pmiss[MAX_PE];
static double log_phit[MAX_PE];

/* Keep track of whether init_charge() has been called so we can exit if a
 * function is called before being initialized. */
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 get_qlo(void)
{
    return qlo;
}

double get_qhi(void)
{
    return qhi;
}

static double nlopt_log_pq(unsigned int n, const double *x, double *grad, void *params)
{
    int i;
    double logp[MAX_PE_LOOKUP];
    double q;
    double mu = *x;

    q = ((double *) params)[0];

    for (i = 1; i < MAX_PE_LOOKUP; i++)
        logp[i] = get_log_pq(q,i) + get_log_phit(i) - mu + i*log(mu) - lnfact(i);

    return -logsumexp(logp+1,MAX_PE_LOOKUP-1);
}

/* Returns the most likely number of PE to produce a given charge `q`, i.e. it
 * returns the `mu` which maximizes
 *
 *     P(q|mu) = \sum_n P(q|n) P(n|mu)
 *
 * Note: You must call init_charge() before calling this function!. */
double get_most_likely_mean_pe(double q)
{
    size_t i;
    static int initialized = 0;
    static double xs[N_CHARGE_MEAN_PE];
    static double ys[N_CHARGE_MEAN_PE];
    double lb[1];
    double ub[1];
    double ss[1];
    double fval;
    nlopt_opt opt;

    if (!initialized) {
        /* Create the minimizer object. */
        opt = nlopt_create(NLOPT_LN_BOBYQA, 1);

        lb[0] = xlo_mean_pe;
        ub[0] = xhi_mean_pe/qmean + 10.0;
        ss[0] = 0.1;

        nlopt_set_lower_bounds(opt, lb);
        nlopt_set_upper_bounds(opt, ub);
        nlopt_set_initial_step(opt, ss);
        nlopt_set_xtol_abs1(opt, 1e-10);

        for (i = 0; i < LEN(xs); i++) {
            xs[i] = xlo_mean_pe + (xhi_mean_pe-xlo_mean_pe)*i/(LEN(xs)-1);
            nlopt_set_min_objective(opt,nlopt_log_pq,&xs[i]);
            ys[i] = fmax(xs[i]/qmean,0.1);
            nlopt_optimize(opt,&ys[i],&fval);
            /* This is a bit of a hack, but there is a discontinuity in
             * get_log_pq() when it switches over from using the Polya
             * distribution to a gaussian which I think is causing the
             * optimization to occasionally stop too early. By minimizing
             * twice, it seems to fix itself and the most likely number of PE
             * is a smooth monotonic function of the charge. */
            nlopt_optimize(opt,&ys[i],&fval);
        }

        nlopt_destroy(opt);

        initialized = 1;
    }

    if (q > xhi_mean_pe) return q/qmean;

    return interp1d(q,xs,ys,LEN(xs));
}

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 get_log_pq(double q, int n)
{
    /* Returns the log of the probability of observing a charge `q` given that
     * `n` PE were generated.
     *
     * `q` should be in units of PE. */
    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 log_norm(q,n*qmean,sqrt(n)*qstd);
    }

    if (q < qlo || q > qhi) return -INFINITY;

    return interp1d(q,x,log_pq[n-1],nq);
}

double get_pq(double q, int n)
{
    /* Returns the probability of observing a charge `q` given that `n` PE were
     * generated.
     *
     * `q` should be in units of PE. */
    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,sqrt(n)*qstd);
    }

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

    return interp1d(q,x,pq[n-1],nq);
}

double get_log_phit(int n)
{
    /* Returns the log of the probability that the channel discriminator fired
     * given `n` PE were generated. */
    if (!initialized) {
        fprintf(stderr, "charge interpolation hasn't been initialized!\n");
        exit(1);
    }

    if (n == 0) {
        /* If no PE were generated, there is a 100% chance that the
         * discriminator didn't fire. */
        return -INFINITY;
    } else if (n > MAX_PE) {
        /* For n > MAX_PE we assume there is a 100% chance that the
         * discriminator fires. */
        return 0.0;
    }

    return log_phit[n-1];
}

double get_log_pmiss(int n)
{
    /* Returns the log of the probability that the channel discriminator didn't
     * fire given `n` PE were generated. */
    if (!initialized) {
        fprintf(stderr, "charge interpolation hasn't been initialized!\n");
        exit(1);
    }

    if (n == 0) {
        /* If no PE were generated, there is a 100% chance that the
         * discriminator didn't fire, so the log is zero. */
        return 0.0;
    } else if (n > MAX_PE) {
        /* For n > MAX_PE we assume there is a 100% chance that the
         * discriminator fires. */
        return -INFINITY;
    }

    return log_pmiss[n-1];
}

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

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

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

    return get_pq(x,n);
}

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

    return get_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)
{
    /* Compute the charge PDFs for up to MAX_PE PE. */
    int i, j;
    gsl_integration_cquad_workspace *w;
    double result, error;
    size_t nevals;
    double params[2];
    gsl_function F;
    double q, q2;
    double y[nq];

    initialized = 1;

    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;

    /* Compute the average single PE charge. */
    gsl_integration_cquad(&F, 0, qhi, 0, 1e-9, w, &result, &error, &nevals);
    q = result;

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

    /* Compute the expected value of the single PE charge squared. */
    gsl_integration_cquad(&F, 0, qhi, 0, 1e-9, w, &result, &error, &nevals);
    q2 = result;

    /* Compute the mean and standard deviation of the single PE charge PDF. */
    qmean = q;
    qstd = sqrt(q2 - q*q);

    /* Compute the single PE charge PDF. */
    for (j = 0; j < nq; j++) {
        pq[0][j] = spe_pol2exp(x[j]);
    }

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

    /* Compute the charge PDFs for `n` PE by convolving the `n-1` charge PDF
     * with the single PE charge PDF. */
    for (i = 2; i <= MAX_PE; i++) {
        for (j = 0; j < nq; j++) {
            params[0] = x[j];
            params[1] = i;
            if (x[j] < 0) {
                pq[i-1][j] = 0.0;
                continue;
            }
            gsl_integration_cquad(&F, 0, x[j], 0, 1e-4, w, &result, &error, &nevals);
            pq[i-1][j] = result;
        }
    }

    /* 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);
        log_pmiss[i-1] = log(result);
        log_phit[i-1] = log(1.0-result);
    }

    /* Set the PDF for the charge to 0 before the discriminator level. */
    for (i = 1; i <= MAX_PE; i++) {
        for (j = 0; j < nq; j++) {
            if (x[j] < MEAN_THRESH/MEAN_HIPT) pq[i-1][j] = 0.0;
        }
        double norm = trapz(&pq[i-1][0],x[1]-x[0],nq);
        for (j = 0; j < nq; j++) {
            pq[i-1][j] /= norm;
        }
    }

    F.function = &gsl_smear;

    /* Convolve the charge PDFs with a gaussian distribution to represent
     * smearing from the electronics. */
    for (i = 1; i <= MAX_PE; i++) {
        for (j = 0; j < nq; j++) {
            params[0] = x[j];
            params[1] = i;
            gsl_integration_cquad(&F, fmax(0.0,x[j]-QSMEAR_ADC*100/MEAN_HIPT), x[j]+QSMEAR_ADC*100/MEAN_HIPT, 0, 1e-4, w, &result, &error, &nevals);
            /* Store the smeared charge PDF in a temporary array since we need
             * to keep the original PDF for the next iteration of the loop. */
            y[j] = result;
        }
        for (j = 0; j < nq; j++) {
            pq[i-1][j] = y[j];
        }
    }

    gsl_integration_cquad_workspace_free(w);

    /* Compute the log of the charge distributions to speed up the likelihood
     * calculation. */
    for (i = 1; i <= MAX_PE; i++) {
        for (j = 0; j < nq; j++) {
            log_pq[i-1][j] = log(pq[i-1][j]);
        }
    }
}