/* 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 <stdio.h>
#include <errno.h>
#include <string.h>
#include <stdlib.h>
#include <gsl/gsl_errno.h>
#include <gsl/gsl_spline.h>
#include <math.h>
#include "optics.h"
#include "quantum_efficiency.h"
#include "solid_angle.h"
#include "pdg.h"
#include "vector.h"
#include "electron.h"
#include "sno.h"
#include "scattering.h"
#include "pmt_response.h"
#include "misc.h"
#include <gsl/gsl_integration.h>
#include <math.h> /* for fmax() */
#include <gsl/gsl_sf_gamma.h>
#include "util.h"

static int initialized = 0;

static double *x, *dEdx_rad, *dEdx, *csda_range;
static size_t size;

static gsl_interp_accel *acc_dEdx_rad;
static gsl_spline *spline_dEdx_rad;

static gsl_interp_accel *acc_dEdx;
static gsl_spline *spline_dEdx;

static gsl_interp_accel *acc_range;
static gsl_spline *spline_range;

/* Electron critical energy in H2O and D2O.
 *
 * These values come from
 * http://pdgprod.lbl.gov/~deg/AtomicNuclearProperties/HTML/deuterium_oxide_liquid.html for D2O and
 * http://pdg.lbl.gov/2018/AtomicNuclearProperties/HTML/water_liquid.html for H2O.
 */
static const double ELECTRON_CRITICAL_ENERGY_H2O = 78.33;
static const double ELECTRON_CRITICAL_ENERGY_D2O = 78.33;

/* Returns the average number of Cerenkov photons in the range 200-800 nm
 * produced by secondary particles in an electron shower.
 *
 * This comes from fitting the ratio # shower photons/rad loss to the function:
 *
 *     c0 + c1/log(T0*c2 + c3)
 *
 * I don't really have any good theoretical motivation for this. My initial
 * thought was that the number of photons should be roughly proportional to the
 * energy lost due to radiation which is why I chose to fit the ratio. This
 * ratio turned out to vary from approximately 520 at low energies (10 MeV) to
 * a roughly constant value of 420 at higher energies (> 1 GeV).
 *
 * This functional form just happened to fit the ratio as a function of energy
 * pretty well.
 *
 * `T0` is the initial kinetic energy of the electron in MeV and `rad` is the
 * energy lost due to radiation in MeV. */
double electron_get_shower_photons(double T0, double rad)
{
    return rad*(406.745 + 0.271816/log(5.31309e-05*T0+1.00174));
}

void electron_get_position_distribution_parameters(double T0, double *a, double *b)
{
    /* Computes the gamma distribution parameters describing the longitudinal
     * profile of the photons generated from an electromagnetic shower.
     *
     * From the PDG "Passage of Particles" section 32.5:
     *
     * "The mean longitudinal profile of the energy deposition in an
     * electromagnetic cascade is reasonably well described by a gamma
     * distribution."
     *
     * Here we use a slightly different form of the gamma distribution:
     *
     *     f(x) = x**(a-1)*exp(-x/b)/(Gamma(a)*b**a)
     *
     * I determined the b parameter by simulating high energy electrons using
     * RAT-PAC and determined that it's roughly equal to the radiation length.
     * To calculate the a parameter we use the formula from the PDG, i.e.
     *
     *     tmax = (a-1)/b = ln(E/E_C) - 0.5
     *
     * Therefore, we calculate a as:
     *
     *     a = tmax*b+1.
     *
     * `T` should be in units of MeV.
     *
     * Example:
     *
     *     double a, b;
     *     electron_get_position_distribution_parameters(1000.0, &a, &b);
     *     double pdf = gamma_pdf(x, a, b);
     *
     * See http://pdg.lbl.gov/2014/reviews/rpp2014-rev-passage-particles-matter.pdf. */
    double tmax;

    *b = RADIATION_LENGTH;
    tmax = log(T0/ELECTRON_CRITICAL_ENERGY_D2O) - 0.5;
    *a = fmax(1.1,tmax*(*b)/RADIATION_LENGTH + 1);
}

double electron_get_angular_distribution_alpha(double T0)
{
    /* To describe the angular distribution of photons in an electromagnetic
     * shower I came up with the heuristic form:
     *
     *     f(cos_theta) ~ exp(-abs(cos_theta-mu)^alpha/beta)
     *
     * I simulated high energy electrons using RAT-PAC in heavy water to fit
     * for the alpha and beta parameters as a function of energy and determined
     * that a reasonably good fit to the data was:
     *
     *     alpha = c0 + c1/log(c2*T0 + c3)
     *
     * where T0 is the initial energy of the electron in MeV and c0, c1, c2,
     * and c3 are constants which I fit out. */
    return fmin(100.0,3.141318e-1 + 2.08198e-01/log(6.33331e-03*T0 + 1.19213e+00));
}

double electron_get_angular_distribution_beta(double T0)
{
    /* To describe the angular distribution of photons in an electromagnetic
     * shower I came up with the heuristic form:
     *
     * f(cos_theta) ~ exp(-abs(cos_theta-mu)^alpha/beta)
     *
     * I simulated high energy electrons using RAT-PAC in heavy water to fit
     * for the alpha and beta parameters as a function of energy and determined
     * that a reasonably good fit to the data was:
     *
     * beta = c0 + c1/log(c2*T0 + c3)
     *
     * where T0 is the initial energy of the electron in MeV and c0, c1, c2,
     * and c3 are constants which I fit out. */
    return 1.35372e-01 + 2.22344e-01/log(1.96249e-02*T0 + 1.23912e+00);
}

double electron_get_angular_pdf_no_norm(double cos_theta, double alpha, double beta, double mu)
{
    return exp(-pow(fabs(cos_theta-mu),alpha)/beta);
}

double electron_get_angular_pdf_norm(double alpha, double beta, double mu)
{
    return pow(beta,1/alpha)*gsl_sf_gamma(1/alpha)*(gsl_sf_gamma_inc_P(1/alpha,pow(1-mu,alpha)/beta)+gsl_sf_gamma_inc_P(1/alpha,pow(1+mu,alpha)/beta))/alpha;
}

double electron_get_angular_pdf_delta_ray(double cos_theta, double alpha, double beta, double mu)
{
    /* Returns the probability density that a photon in the wavelength range
     * 200 nm - 800 nm is emitted at an angle cos_theta.
     *
     * The angular distribution is modelled by the pdf:
     *
     *     f(cos_theta) ~ exp(-abs(cos_theta-mu)^alpha/beta)
     *
     * where alpha and beta are constants which depend on the initial energy of
     * the particle.
     *
     * There is no nice closed form solution for the integral of this function,
     * so we just compute it on the fly. To make things faster, we keep track
     * of the last alpha, beta, and mu parameters that were passed in so we
     * avoid recomputing the normalization factor. */
    size_t i;
    static double last_alpha, last_beta, last_mu, norm;
    static int first = 1;
    static double x[10000], y[10000];

    if (first || alpha != last_alpha || beta != last_beta || mu != last_mu) {
        norm = electron_get_angular_pdf_norm(alpha, beta, mu);

        last_alpha = alpha;
        last_beta = beta;
        last_mu = mu;

        for (i = 0; i < LEN(x); i++) {
            x[i] = -1.0 + 2.0*i/(LEN(x)-1);
            y[i] = electron_get_angular_pdf_no_norm(x[i], alpha, beta, mu)/norm;
        }

        first = 0;
    }

    return interp1d(cos_theta,x,y,LEN(x));
}

double electron_get_angular_pdf(double cos_theta, double alpha, double beta, double mu)
{
    /* Returns the probability density that a photon in the wavelength range
     * 200 nm - 800 nm is emitted at an angle cos_theta.
     *
     * The angular distribution is modelled by the pdf:
     *
     *     f(cos_theta) ~ exp(-abs(cos_theta-mu)^alpha/beta)
     *
     * where alpha and beta are constants which depend on the initial energy of
     * the particle.
     *
     * There is no nice closed form solution for the integral of this function,
     * so we just compute it on the fly. To make things faster, we keep track
     * of the last alpha, beta, and mu parameters that were passed in so we
     * avoid recomputing the normalization factor. */
    size_t i;
    static double last_alpha, last_beta, last_mu, norm;
    static int first = 1;
    static double x[10000], y[10000];

    if (first || alpha != last_alpha || beta != last_beta || mu != last_mu) {
        norm = electron_get_angular_pdf_norm(alpha, beta, mu);

        last_alpha = alpha;
        last_beta = beta;
        last_mu = mu;

        for (i = 0; i < LEN(x); i++) {
            x[i] = -1.0 + 2.0*i/(LEN(x)-1);
            y[i] = electron_get_angular_pdf_no_norm(x[i], alpha, beta, mu)/norm;
        }

        first = 0;
    }

    return interp1d(cos_theta,x,y,LEN(x));
}

static int init()
{
    int i, j;
    char line[256];
    char *str;
    double value;
    int n;

    FILE *f = open_file("e_water_liquid.txt", "r");

    if (!f) {
        fprintf(stderr, "failed to open e_water_liquid.txt: %s\n", strerror(errno));
        return -1;
    }

    i = 0;
    n = 0;
    /* For the first pass, we just count how many values there are. */
    while (fgets(line, sizeof(line), f)) {
        size_t len = strlen(line);
        if (len && (line[len-1] != '\n')) {
            fprintf(stderr, "got incomplete line on line %i: '%s'\n", i, line);
            goto err;
        }

        i += 1;

        /* Skip the first 8 lines since it's just a header. */
        if (i <= 8) continue;

        if (!len) continue;
        else if (line[0] == '#') continue;

        str = strtok(line," \n");

        while (str) {
            value = strtod(str, NULL);
            str = strtok(NULL," \n");
        }

        n += 1;
    }

    x = malloc(sizeof(double)*n);
    dEdx_rad = malloc(sizeof(double)*n);
    dEdx = malloc(sizeof(double)*n);
    csda_range = malloc(sizeof(double)*n);
    size = n;

    i = 0;
    n = 0;
    /* Now, we actually store the values. */
    rewind(f);
    while (fgets(line, sizeof(line), f)) {
        size_t len = strlen(line);
        if (len && (line[len-1] != '\n')) {
            fprintf(stderr, "got incomplete line on line %i: '%s'\n", i, line);
            goto err;
        }

        i += 1;

        /* Skip the first 8 lines since it's just a header. */
        if (i <= 8) continue;

        if (!len) continue;
        else if (line[0] == '#') continue;

        str = strtok(line," \n");

        j = 0;
        while (str) {
            value = strtod(str, NULL);
            switch (j) {
            case 0:
                x[n] = value;
                break;
            case 2:
                dEdx_rad[n] = value;
                break;
            case 3:
                dEdx[n] = value;
                break;
            case 4:
                csda_range[n] = value;
                break;
            }
            j += 1;
            str = strtok(NULL," \n");
        }

        n += 1;
    }

    fclose(f);

    acc_dEdx_rad = gsl_interp_accel_alloc();
    spline_dEdx_rad = gsl_spline_alloc(gsl_interp_linear, size);
    gsl_spline_init(spline_dEdx_rad, x, dEdx_rad, size);

    acc_dEdx = gsl_interp_accel_alloc();
    spline_dEdx = gsl_spline_alloc(gsl_interp_linear, size);
    gsl_spline_init(spline_dEdx, x, dEdx, size);

    acc_range = gsl_interp_accel_alloc();
    spline_range = gsl_spline_alloc(gsl_interp_linear, size);
    gsl_spline_init(spline_range, x, csda_range, size);

    initialized = 1;

    return 0;

err:
    fclose(f);

    return -1;
}

/* Returns the maximum kinetic energy for an electron in the range tables.
 *
 * If you call electron_get_range() or electron_get_dEdx() with a kinetic
 * energy higher you will get a GSL interpolation error. */
double electron_get_max_energy(void)
{
    if (!initialized) {
        if (init()) {
            exit(1);
        }
    }

    return x[size-1];
}

double electron_get_range(double T, double rho)
{
    /* Returns the approximate range a electron with kinetic energy `T` will travel
     * in water before losing all of its energy. This range is interpolated
     * based on data from the PDG which uses the continuous slowing down
     * approximation.
     *
     * `T` should be in MeV, and `rho` should be in g/cm^3.
     *
     * Return value is in cm.
     *
     * See http://pdg.lbl.gov/2018/AtomicNuclearProperties/adndt.pdf. */
    if (!initialized) {
        if (init()) {
            exit(1);
        }
    }

    return gsl_spline_eval(spline_range, T, acc_range)/rho;
}

double electron_get_dEdx_rad(double T, double rho)
{
    /* Returns the approximate radiative dE/dx for a electron in water with
     * kinetic energy `T`.
     *
     * `T` should be in MeV and `rho` in g/cm^3.
     *
     * Return value is in MeV/cm.
     *
     * See http://pdg.lbl.gov/2018/AtomicNuclearProperties/adndt.pdf. */
    if (!initialized) {
        if (init()) {
            exit(1);
        }
    }

    if (T < spline_dEdx_rad->x[0]) return spline_dEdx_rad->y[0];

    return gsl_spline_eval(spline_dEdx_rad, T, acc_dEdx_rad)*rho;
}

double electron_get_dEdx(double T, double rho)
{
    /* Returns the approximate dE/dx for a electron in water with kinetic energy
     * `T`.
     *
     * `T` should be in MeV and `rho` in g/cm^3.
     *
     * Return value is in MeV/cm.
     *
     * See http://pdg.lbl.gov/2018/AtomicNuclearProperties/adndt.pdf. */
    if (!initialized) {
        if (init()) {
            exit(1);
        }
    }

    if (T < spline_dEdx->x[0]) return spline_dEdx->y[0];

    return gsl_spline_eval(spline_dEdx, T, acc_dEdx)*rho;
}