/* 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 <math.h>
#include "optics.h"
#include "misc.h"
#include "quantum_efficiency.h"
#include <stdio.h> /* for fprintf() */
#include <gsl/gsl_integration.h>
#include <gsl/gsl_errno.h> /* for gsl_strerror() */
#include <gsl/gsl_spline.h>
#include "misc.h"
#include "db.h"
#include "sno.h"

char optics_err[256];

/* Index of refraction in water and heavy water averaged over the Cerenkov
 * spectrum and the PMT quantum efficiency.
 *
 * Set these to 0.0 so that it's obvious if someone forgets to call
 * optics_init() before using them. */
double avg_index_h2o = 0.0, avg_index_d2o = 0.0;

/* Absorption coefficients for H2O, D2O, and acrylic as a function of
 * wavelength from SNOMAN.
 *
 * Note: These numbers come from mcprod/media_qoca_d2o_20060110.cmd. */
static double absorption_coefficient_d2o_wavelengths[6] = {337.0, 365.0, 386.0, 420.0, 500.0, 620.0};
static double absorption_coefficient_d2o[6] = {6.057e-05, 2.680e-05, 3.333e-05, 3.006e-05, 2.773e-05, 2.736e-05};
static double rayleigh_scl_fac_d2o = 1.000;
static double isothermal_comp_d2o = 4.92e-10;

static double absorption_coefficient_h2o_wavelengths[6] = {337.0, 365.0, 386.0, 420.0, 500.0, 620.0};
static double absorption_coefficient_h2o[6] = {8.410e-05, 2.880e-06, 1.431e-04, 1.034e-04, 5.239e-04, 2.482e-03};
static double rayleigh_scl_fac_acrylic = 0.950;
static double isothermal_comp_acrylic = 3.55e-10;

static double absorption_coefficient_acrylic_wavelengths[6] = {337.0, 365.0, 386.0, 420.0, 500.0, 620.0};
static double absorption_coefficient_acrylic[6] = {5.610e-02, 2.279e-02, 1.204e-02, 7.587e-03, 7.036e-03, 7.068e-03};
static double rayleigh_scl_fac_h2o = 0.870;
static double isothermal_comp_h2o = 4.78e-10;

static int initialized = 0;

/* Number of points to sample when computing the average absorption and
 * scattering probabilities.
 *
 * Note that we use this array to sample for very small distances when
 * computing the absorption in the acrylic so (xhi-xlo)/N should be less than
 * 10 cm. */
#define N 10000

/* Lower and upper bounds for the distances used to calculate the average
 * absorption and scattering probabilities. */
static const double xlo = 0.0;
static const double xhi = 2*PSUP_RADIUS;

/* Array of lengths used in the lookup tables for the average absorption and
 * scattering probabilities. */
static double fabs_x[N];

/* Average fraction of light which is not absorbed by D2O as a function of
 * length weighted by the Cerenkov spectrum and the quantum efficiency.
 *
 * We calculate the fraction not absorbed instead of absorbed */
static double fabs_d2o[N];
static double fabs_h2o[N];
static double fabs_acrylic[N];
static double fsct_d2o[N];
static double fsct_h2o[N];
static double fsct_acrylic[N];

/* From Table 4 in the paper. */
static double A0 = 0.243905091;
static double A1 = 9.53518094e-3;
static double A2 = -3.64358110e-3;
static double A3 = 2.65666426e-4;
static double A4 = 1.59189325e-3;
static double A5 = 2.45733798e-3;
static double A6 = 0.897478251;
static double A7 = -1.63066183e-2;
static double UV = 0.2292020;
static double IR = 5.432937;

static double RIND_H2O_C1 = 1.302;
static double RIND_H2O_C2 = 0.01562;
static double RIND_H2O_C3 = 0.32;

static double RIND_D2O_C1 = 1.302;
static double RIND_D2O_C2 = 0.01333;
static double RIND_D2O_C3 = 0.32;

static double RIND_ACRYLIC_C1 = 1.452;
static double RIND_ACRYLIC_C2 = 0.02;
static double RIND_ACRYLIC_C3 = 0.32;

/* Wavelength of 1 eV particle in nm.
 *
 * From http://pdg.lbl.gov/2018/reviews/rpp2018-rev-phys-constants.pdf. */
static double HC = 1.239841973976e3;

double get_index(double p, double wavelength, double T)
{
    /* Returns the index of refraction of pure water for a given density,
     * wavelength, and temperature. The density should be in units of g/cm^3,
     * the wavelength in nm, and the temperature in Celsius.
     *
     * See "Refractive Index of Water and Steam as a function of Wavelength,
     * Temperature, and Density" by Schiebener et al. 1990. */
    /* normalize the temperature and pressure */
    wavelength = wavelength/589.0;
    T = (T+273.15)/273.15;

    /* first we compute the right hand side of Equation 7 */
    double c = A0 + A1*p + A2*T + A3*pow(wavelength,2)*T + A4/pow(wavelength,2) + A5/(pow(wavelength,2)-pow(UV,2)) + A6/(pow(wavelength,2)-pow(IR,2)) + A7*pow(p,2);
    c *= p;

    return sqrt((2*c+1)/(1-c));
}

double get_index_snoman_h2o(double wavelength)
{
    /* Returns the index of refraction of light water that SNOMAN uses for a
     * given wavelength.  The wavelength should be in units of nm.
     *
     * The coefficients come from media.dat in SNOMAN version 5.0294 and the
     * formula used to compute the index of refraction comes from the SNOMAN
     * companion page for titles MEDA. */
    double E;

    /* Calculate the energy of the photon in eV. */
    E = HC/wavelength;

    return RIND_H2O_C1 + RIND_H2O_C2*exp(RIND_H2O_C3*E);
}

double get_index_snoman_d2o(double wavelength)
{
    /* Returns the index of refraction of heavy water that SNOMAN uses for a
     * given wavelength.  The wavelength should be in units of nm.
     *
     * The coefficients come from media.dat in SNOMAN version 5.0294 and the
     * formula used to compute the index of refraction comes from the SNOMAN
     * companion page for titles MEDA. */
    double E;

    /* Calculate the energy of the photon in eV. */
    E = HC/wavelength;

    return RIND_D2O_C1 + RIND_D2O_C2*exp(RIND_D2O_C3*E);
}

double get_index_snoman_acrylic(double wavelength)
{
    /* Returns the index of refraction of SNO acrylic for a given wavelength.
     * The wavelength should be in units of nm.
     *
     * The coefficients come from media.dat in SNOMAN version 5.0294 and the
     * formula used to compute the index of refraction comes from the SNOMAN
     * companion page for titles MEDA. */
    double E;

    /* Calculate the energy of the photon in eV. */
    E = HC/wavelength;

    return RIND_ACRYLIC_C1 + RIND_ACRYLIC_C2*exp(RIND_ACRYLIC_C3*E);
}

/* Returns the probability per unit length (1/cm) for Rayleigh scattering in a
 * material with an isothermal compressibility of `isothermal_comp` at a given
 * wavelength in nm.
 *
 * The isothermal compressibility should be in units of N^-1 m^2.
 *
 * This function is mostly copied from snoman's rayint_prob.for. */
double rayint_prob(double wavelength, double n, double isothermal_comp)
{
    double E;
    double confac = 1.53E26;

    /* Calculate the energy of the photon in MeV. */
    E = 1e-6*HC/wavelength;

    return isothermal_comp*confac*pow(E,4)*pow(n*n - 1.0,2)*pow(n*n + 2.0,2);
}

/* Returns the probability per unit length (1/cm) for Rayleigh scattering in
 * D2O at a given wavelength in nm. */
double rayint_prob_d2o(double wavelength)
{
    return rayleigh_scl_fac_d2o*rayint_prob(wavelength, get_index_snoman_d2o(wavelength), isothermal_comp_d2o);
}

/* Returns the probability per unit length (1/cm) for Rayleigh scattering in
 * H2O at a given wavelength in nm. */
double rayint_prob_h2o(double wavelength)
{
    return rayleigh_scl_fac_h2o*rayint_prob(wavelength, get_index_snoman_h2o(wavelength), isothermal_comp_h2o);
}

/* Returns the probability per unit length (1/cm) for Rayleigh scattering in
 * acrylic at a given wavelength in nm. */
double rayint_prob_acrylic(double wavelength)
{
    return rayleigh_scl_fac_acrylic*rayint_prob(wavelength, get_index_snoman_acrylic(wavelength), isothermal_comp_acrylic);
}

/* Returns the average probability that a photon is not absorbed after a
 * distance `x` in cm in D2O.
 *
 * This average probability is found by averaging
 *
 *     e^(-x/l(wavelength))
 *
 * over the wavelength range 200-800 nm weighted by the Cerenkov wavelenght
 * distribution and the quantum efficiency. */
double get_fabs_d2o(double x)
{
    if (!initialized) {
        fprintf(stderr, "weighted absorption length hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(x,fabs_x,fabs_d2o,LEN(fabs_x));
}

/* Returns the average probability that a photon is not absorbed after a
 * distance `x` in cm in H2O. See get_fabs_d2o() for how this is calculated. */
double get_fabs_h2o(double x)
{
    if (!initialized) {
        fprintf(stderr, "average probability hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(x,fabs_x,fabs_h2o,LEN(fabs_x));
}

/* Returns the average probability that a photon is not absorbed after a
 * distance `x` in cm in the SNO acrylic. See get_fabs_d2o() for how this is
 * calculated.
 *
 * FIXME: Should include a term for the "dark" acrylic. */
double get_fabs_acrylic(double x)
{
    if (!initialized) {
        fprintf(stderr, "weighted absorption length hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(x,fabs_x,fabs_acrylic,LEN(fabs_x));
}

/* Returns the average probability that a photon is not scattered after a
 * distance `x` in cm in D2O. See get_fabs_d2o() for how this is calculated. */
double get_fsct_d2o(double x)
{
    if (!initialized) {
        fprintf(stderr, "weighted rayleigh scattering length hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(x,fabs_x,fsct_d2o,LEN(fabs_x));
}

/* Returns the average probability that a photon is not scattered after a
 * distance `x` in cm in H2O. See get_fabs_d2o() for how this is calculated. */
double get_fsct_h2o(double x)
{
    if (!initialized) {
        fprintf(stderr, "weighted rayleigh scattering length hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(x,fabs_x,fsct_h2o,LEN(fabs_x));
}

/* Returns the average probability that a photon is not scattered after a
 * distance `x` in cm in SNO acrylic. See get_fabs_d2o() for how this is
 * calculated. */
double get_fsct_acrylic(double x)
{
    if (!initialized) {
        fprintf(stderr, "weighted rayleigh scattering length hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(x,fabs_x,fsct_acrylic,LEN(fabs_x));
}

/* Returns the absorption length as a function of wavelength in D2O.
 *
 * The wavelength should be in nm and the absorption length is in cm. */
double get_absorption_length_snoman_d2o(double wavelength)
{
    /* Note: We use GSL splines here instead of interp1d() since we aren't
     * guaranteed that the values in the lookup table are evenly spaced. */
    static gsl_spline *spline;
    static gsl_interp_accel *acc;

    if (!spline) {
        spline = gsl_spline_alloc(gsl_interp_linear, LEN(absorption_coefficient_d2o_wavelengths));
        gsl_spline_init(spline, absorption_coefficient_d2o_wavelengths, absorption_coefficient_d2o, LEN(absorption_coefficient_d2o_wavelengths));
        acc = gsl_interp_accel_alloc();
    }

    /* If the wavelength is outside the range of the lookup table we just
     * return the value at the beginning or end. We do this because by default
     * gsl_spline_eval will exit if the wavelength is out of bounds. */
    if (wavelength < absorption_coefficient_d2o_wavelengths[0])
        wavelength = absorption_coefficient_d2o_wavelengths[0];

    if (wavelength > absorption_coefficient_d2o_wavelengths[LEN(absorption_coefficient_d2o_wavelengths)-1])
        wavelength = absorption_coefficient_d2o_wavelengths[LEN(absorption_coefficient_d2o_wavelengths)-1];

    return 1.0/gsl_spline_eval(spline, wavelength, acc);
}

/* Returns the absorption length as a function of wavelength in H2O.
 *
 * The wavelength should be in nm and the absorption length is in cm. */
double get_absorption_length_snoman_h2o(double wavelength)
{
    /* Note: We use GSL splines here instead of interp1d() since we aren't
     * guaranteed that the values in the lookup table are evenly spaced. */
    static gsl_spline *spline;
    static gsl_interp_accel *acc;

    if (!spline) {
        spline = gsl_spline_alloc(gsl_interp_linear, LEN(absorption_coefficient_h2o_wavelengths));
        gsl_spline_init(spline, absorption_coefficient_h2o_wavelengths, absorption_coefficient_h2o, LEN(absorption_coefficient_h2o_wavelengths));
        acc = gsl_interp_accel_alloc();
    }

    /* If the wavelength is outside the range of the lookup table we just
     * return the value at the beginning or end. We do this because by default
     * gsl_spline_eval will exit if the wavelength is out of bounds. */
    if (wavelength < absorption_coefficient_h2o_wavelengths[0])
        wavelength = absorption_coefficient_h2o_wavelengths[0];

    if (wavelength > absorption_coefficient_h2o_wavelengths[LEN(absorption_coefficient_h2o_wavelengths)-1])
        wavelength = absorption_coefficient_h2o_wavelengths[LEN(absorption_coefficient_h2o_wavelengths)-1];

    return 1.0/gsl_spline_eval(spline, wavelength, acc);
}

/* Returns the absorption length as a function of wavelength in acrylic.
 *
 * The wavelength should be in nm and the absorption length is in cm. */
double get_absorption_length_snoman_acrylic(double wavelength)
{
    /* Note: We use GSL splines here instead of interp1d() since we aren't
     * guaranteed that the values in the lookup table are evenly spaced. */
    static gsl_spline *spline;
    static gsl_interp_accel *acc;

    if (!spline) {
        spline = gsl_spline_alloc(gsl_interp_linear, LEN(absorption_coefficient_acrylic_wavelengths));
        gsl_spline_init(spline, absorption_coefficient_acrylic_wavelengths, absorption_coefficient_acrylic, LEN(absorption_coefficient_acrylic_wavelengths));
        acc = gsl_interp_accel_alloc();
    }

    /* If the wavelength is outside the range of the lookup table we just
     * return the value at the beginning or end. We do this because by default
     * gsl_spline_eval will exit if the wavelength is out of bounds. */
    if (wavelength < absorption_coefficient_acrylic_wavelengths[0])
        wavelength = absorption_coefficient_acrylic_wavelengths[0];

    if (wavelength > absorption_coefficient_acrylic_wavelengths[LEN(absorption_coefficient_acrylic_wavelengths)-1])
        wavelength = absorption_coefficient_acrylic_wavelengths[LEN(absorption_coefficient_acrylic_wavelengths)-1];

    return 1.0/gsl_spline_eval(spline, wavelength, acc);
}

static double gsl_rayleigh_scattering_length_snoman_d2o(double wavelength, void *params)
{
    /* Returns the Rayleigh scattering length in D2O weighted by the quantum
     * efficiency and the Cerenkov spectrum. */
    double qe, x;

    x = *((double *) params);

    qe = get_quantum_efficiency(wavelength);

    return qe*exp(-x*rayint_prob_d2o(wavelength))/pow(wavelength,2);
}

static double gsl_rayleigh_scattering_length_snoman_h2o(double wavelength, void *params)
{
    /* Returns the Rayleigh scattering length in H2O weighted by the quantum
     * efficiency and the Cerenkov spectrum. */
    double qe, x;

    x = *((double *) params);

    qe = get_quantum_efficiency(wavelength);

    return qe*exp(-x*rayint_prob_h2o(wavelength))/pow(wavelength,2);
}

static double gsl_rayleigh_scattering_length_snoman_acrylic(double wavelength, void *params)
{
    /* Returns the Rayleigh scattering length in acrylic weighted by the
     * quantum efficiency and the Cerenkov spectrum. */
    double qe, x;

    x = *((double *) params);

    qe = get_quantum_efficiency(wavelength);

    return qe*exp(-x*rayint_prob_acrylic(wavelength))/pow(wavelength,2);
}

static double gsl_absorption_length_snoman_d2o(double wavelength, void *params)
{
    /* Returns the absorption length in D2O weighted by the quantum efficiency
     * and the Cerenkov spectrum. */
    double qe, x;

    x = *((double *) params);

    qe = get_quantum_efficiency(wavelength);

    return qe*exp(-x/get_absorption_length_snoman_d2o(wavelength))/pow(wavelength,2);
}

static double gsl_absorption_length_snoman_h2o(double wavelength, void *params)
{
    /* Returns the absorption length in H2O weighted by the quantum efficiency
     * and the Cerenkov spectrum. */
    double qe, x;

    x = *((double *) params);

    qe = get_quantum_efficiency(wavelength);

    return qe*exp(-x/get_absorption_length_snoman_h2o(wavelength))/pow(wavelength,2);
}

static double gsl_absorption_length_snoman_acrylic(double wavelength, void *params)
{
    /* Returns the absorption length in acrylic weighted by the quantum
     * efficiency and the Cerenkov spectrum. */
    double qe, x;

    x = *((double *) params);

    qe = get_quantum_efficiency(wavelength);

    return qe*exp(-x/get_absorption_length_snoman_acrylic(wavelength))/pow(wavelength,2);
}

/* Returns the product of the index of refraction in D2O multiplied by the
 * quantum efficiency and the Cerenkov spectrum. */
static double gsl_index_d2o(double wavelength, void *params)
{
    double qe;

    qe = get_quantum_efficiency(wavelength);

    return qe*get_index_snoman_d2o(wavelength)/pow(wavelength,2);
}

/* Returns the product of the index of refraction in H2O multiplied by the
 * quantum efficiency and the Cerenkov spectrum. */
static double gsl_index_h2o(double wavelength, void *params)
{
    double qe;

    qe = get_quantum_efficiency(wavelength);

    return qe*get_index_snoman_h2o(wavelength)/pow(wavelength,2);
}

static double gsl_cerenkov(double wavelength, void *params)
{
    /* Returns the quantum efficiency multiplied by the Cerenkov spectrum. */
    double qe;

    qe = get_quantum_efficiency(wavelength);

    return qe/pow(wavelength,2);
}

int optics_init(void)
{
    int i;
    double norm;
    double result, error;
    size_t nevals;
    int status;
    gsl_integration_cquad_workspace *w;
    gsl_function F;

    w = gsl_integration_cquad_workspace_alloc(100);

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

    status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &norm, &error, &nevals);

    if (status) {
        sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
        return -1;
    }

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

    status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

    if (status) {
        sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
        return -1;
    }

    avg_index_d2o = result/norm;

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

    status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

    if (status) {
        sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
        return -1;
    }

    avg_index_h2o = result/norm;

    for (i = 0; i < LEN(fabs_x); i++) {
        fabs_x[i] = xlo + (xhi-xlo)*i/(LEN(fabs_x)-1);
    }

    F.function = &gsl_absorption_length_snoman_d2o;

    for (i = 0; i < LEN(fabs_x); i++) {
        F.params = fabs_x+i;

        status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

        if (status) {
            sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            return -1;
        }

        fabs_d2o[i] = result/norm;
    }

    F.function = &gsl_absorption_length_snoman_h2o;

    for (i = 0; i < LEN(fabs_x); i++) {
        F.params = fabs_x+i;

        status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

        if (status) {
            sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            return -1;
        }

        fabs_h2o[i] = result/norm;
    }

    F.function = &gsl_absorption_length_snoman_acrylic;

    for (i = 0; i < LEN(fabs_x); i++) {
        F.params = fabs_x+i;

        status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

        if (status) {
            sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            return -1;
        }

        fabs_acrylic[i] = result/norm;
    }

    F.function = &gsl_rayleigh_scattering_length_snoman_d2o;

    for (i = 0; i < LEN(fabs_x); i++) {
        F.params = fabs_x+i;

        status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

        if (status) {
            sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            return -1;
        }

        fsct_d2o[i] = result/norm;
    }

    F.function = &gsl_rayleigh_scattering_length_snoman_h2o;

    for (i = 0; i < LEN(fabs_x); i++) {
        F.params = fabs_x+i;

        status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

        if (status) {
            sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            return -1;
        }

        fsct_h2o[i] = result/norm;
    }

    F.function = &gsl_rayleigh_scattering_length_snoman_acrylic;

    for (i = 0; i < LEN(fabs_x); i++) {
        F.params = fabs_x+i;

        status = gsl_integration_cquad(&F, 200, 800, 0, 1e-2, w, &result, &error, &nevals);

        if (status) {
            sprintf(optics_err, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            return -1;
        }

        fsct_acrylic[i] = result/norm;
    }

    gsl_integration_cquad_workspace_free(w);

    initialized = 1;

    return 0;
}