/* 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 "pmt_response.h"
#include "dict.h"
#include <stdint.h> /* for uint32_t */
#include <gsl/gsl_interp2d.h>
#include <gsl/gsl_spline2d.h>
#include <gsl/gsl_integration.h>
#include <stdio.h> /* for sprintf() */
#include "misc.h"
#include "quantum_efficiency.h"
#include "db.h"
#include <gsl/gsl_errno.h> /* for gsl_strerror() */

static int initialized = 0;

/* Global error string set when pmt_response_init() returns -1. */
char pmtr_err[256];

/* 2D lookup table for the PMT response as a function of angle and wavelength. */
static gsl_spline2d *spline_resp;
static gsl_spline2d *spline_reflec;
static gsl_interp_accel *xacc;
static gsl_interp_accel *yacc;

/* Angles (in degrees) that the PMT response lookup table is defined for.
 * Ideally these would have been stored with the table itself, but these are
 * hardcoded in SNOMAN. */
static double thetas[NUM_ANGLES] = {
    0.0,
    1.0,
    2.0,
    3.0,
    4.0,
    5.0,
    6.0,
    7.0,
    8.0,
    9.0,
    10.0,
    11.0,
    12.0,
    13.0,
    14.0,
    15.0,
    16.0,
    17.0,
    18.0,
    19.0,
    20.0,
    21.0,
    22.0,
    23.0,
    24.0,
    25.0,
    26.0,
    27.0,
    28.0,
    29.0,
    30.0,
    31.0,
    32.0,
    33.0,
    34.0,
    35.0,
    36.0,
    37.0,
    38.0,
    39.0,
    40.0,
    41.0,
    42.0,
    43.0,
    44.0,
    45.0,
    46.0,
    47.0,
    48.0,
    49.0,
    50.0,
    51.0,
    52.0,
    53.0,
    54.0,
    55.0,
    56.0,
    57.0,
    58.0,
    59.0,
    60.0,
    61.0,
    62.0,
    63.0,
    64.0,
    65.0,
    66.0,
    67.0,
    68.0,
    69.0,
    70.0,
    71.0,
    72.0,
    73.0,
    74.0,
    75.0,
    76.0,
    77.0,
    78.0,
    79.0,
    80.0,
    81.0,
    82.0,
    83.0,
    84.0,
    85.0,
    86.0,
    87.0,
    88.0,
    89.0,
};

/* Wavelengths (in nm) that the PMT response lookup table is defined for.
 * Ideally these would have been stored with the table itself, but these are
 * hardcoded in SNOMAN. */
static double wavelengths[NUM_WAVELENGTHS] = {
    220.0,
    230.0,
    240.0,
    250.0,
    260.0,
    270.0,
    280.0,
    290.0,
    300.0,
    310.0,
    320.0,
    330.0,
    340.0,
    350.0,
    360.0,
    370.0,
    380.0,
    390.0,
    400.0,
    410.0,
    420.0,
    430.0,
    440.0,
    450.0,
    460.0,
    470.0,
    480.0,
    490.0,
    500.0,
    510.0,
    520.0,
    530.0,
    540.0,
    550.0,
    560.0,
    570.0,
    580.0,
    590.0,
    600.0,
    610.0,
    620.0,
    630.0,
    640.0,
    650.0,
    660.0,
    670.0,
    680.0,
    690.0,
    700.0,
    710.0,
};

/* PMT response as a function of angle and wavelength. */
static double resp[NUM_ANGLES*NUM_WAVELENGTHS];
/* PMT reflectivity as a function of angle and wavelength. */
static double reflec[NUM_ANGLES*NUM_WAVELENGTHS];

#define N 1000

/* Lower bound for cos(theta). */
static double xlo = 0.0;
/* Upper bound for cos(theta). */
static double xhi = 1.0;

/* Array holding the values of cos(theta). */
static double cos_thetas[N];

/* PMT response as a function of cos(theta) weighted by the Cerenkov spectrum
 * and the quantum efficiency. */
static double weighted_resp[N];
/* PMT reflectivity as a function of cos(theta) weighted by the Cerenkov
 * spectrum and the quantum efficiency. */
static double weighted_reflec[N];

/* Returns the probability that a photon is reflected as a function of the
 * cosine of the photon angle with respect to the PMT normal.
 *
 * Note: The angle should be relative to the negative of the PMT normal, i.e. a
 * photon incident perpendicular to the front of the PMT should have
 *
 *     cos(theta) = 1.
 *
 */
double get_weighted_pmt_reflectivity(double cos_theta)
{
    if (!initialized) {
        fprintf(stderr, "pmt response hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(cos_theta, cos_thetas, weighted_reflec, LEN(cos_thetas));
}

/* Returns the probability that a photon is detected as a function of the
 * cosine of the photon angle with respect to the PMT normal.
 *
 * Note: The angle should be relative to the negative of the PMT normal, i.e. a
 * photon incident perpendicular to the front of the PMT should have
 *
 *     cos(theta) = 1.
 *
 * Note: This does *not* include the PMT quantum efficiency. */
double get_weighted_pmt_response(double cos_theta)
{
    if (!initialized) {
        fprintf(stderr, "pmt response hasn't been initialized!\n");
        exit(1);
    }

    return interp1d(cos_theta, cos_thetas, weighted_resp, LEN(cos_thetas));
}

double get_pmt_reflectivity(double wavelength, double theta)
{
    /* Returns the probability that a photon is reflected as a function of the
     * photon angle with respect to the PMT normal (in radians) and the
     * wavelength (in nm). */
    double deg;

    if (!initialized) {
        fprintf(stderr, "pmt response hasn't been initialized!\n");
        exit(1);
    }

    /* Convert to degrees since the PMTR table is in degrees. */
    deg = theta*180.0/M_PI;

    deg = fmod(deg,180.0);

    if (deg < 0.0)
        deg += 180.0;

    if (deg > thetas[NUM_ANGLES-1])
        deg = thetas[NUM_ANGLES-1];

    if (wavelength < wavelengths[0])
        wavelength = wavelengths[0];

    if (wavelength > wavelengths[NUM_WAVELENGTHS-1])
        wavelength = wavelengths[NUM_WAVELENGTHS-1];

    return gsl_spline2d_eval(spline_reflec, deg, wavelength, xacc, yacc);
}

double get_pmt_response(double wavelength, double theta)
{
    /* Returns the probability that a photon is detected as a function of the
     * photon angle with respect to the PMT normal (in radians) and the
     * wavelength (in nm).
     *
     * Note: This does *not* include the PMT quantum efficiency. */
    double deg, qe;

    if (!initialized) {
        fprintf(stderr, "pmt response hasn't been initialized!\n");
        exit(1);
    }

    /* Convert to degrees since the PMTR table is in degrees. */
    deg = theta*180.0/M_PI;

    deg = fmod(deg,180.0);

    if (deg < 0.0)
        deg += 180.0;

    if (deg > thetas[NUM_ANGLES-1])
        deg = thetas[NUM_ANGLES-1];

    if (wavelength < wavelengths[0])
        wavelength = wavelengths[0];

    if (wavelength > wavelengths[NUM_WAVELENGTHS-1])
        wavelength = wavelengths[NUM_WAVELENGTHS-1];

    /* The PMTR bank in SNOMAN included the effect of the PMT quantum
     * efficiency since it was used for the "grey disk" model.  Therefore,
     * since we already account for the quantum efficiency it is necessary to
     * divide the response by the quantum efficiency. */

    qe = get_quantum_efficiency(wavelength);

    /* If the quantum efficiency is zero, we have no way of knowing what the
     * response should be since it was multiplied by zero. Since for these
     * wavelengths the photon won't be detected, it doesn't really matter. */
    if (qe == 0.0) return 0.0;

    return gsl_spline2d_eval(spline_resp, deg, wavelength, xacc, yacc)/qe;
}

static double gsl_pmt_reflectivity(double wavelength, void *params)
{
    /* Returns the probability that a photon is reflected as a function of
     * wavelength for a specific angle weighted by the quantum efficiency and
     * the Cerenkov spectrum. */
    double qe, cos_theta;

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

    qe = get_quantum_efficiency(wavelength);

    return qe*get_pmt_reflectivity(wavelength,acos(cos_theta))/pow(wavelength,2);
}

static double gsl_pmt_response(double wavelength, void *params)
{
    /* Returns the probability that a photon is detected as a function of
     * wavelength for a specific angle weighted by the quantum efficiency and
     * the Cerenkov spectrum. */
    double qe, cos_theta;

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

    qe = get_quantum_efficiency(wavelength);

    return qe*get_pmt_response(wavelength,acos(cos_theta))/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 pmt_response_init(dict *db)
{
    int i, j;
    float *pmtr;
    double norm;

    double result, error;
    size_t nevals;
    int status;
    gsl_integration_cquad_workspace *w;
    gsl_function F;
    double params[1];

    spline_resp = gsl_spline2d_alloc(gsl_interp2d_bilinear, NUM_ANGLES, NUM_WAVELENGTHS);
    spline_reflec = gsl_spline2d_alloc(gsl_interp2d_bilinear, NUM_ANGLES, NUM_WAVELENGTHS);
    xacc = gsl_interp_accel_alloc();
    yacc = gsl_interp_accel_alloc();

    pmtr = (float *) get_bank(db, "PMTR", 1);

    if (!pmtr) {
        sprintf(pmtr_err, "failed to load PMTR bank\n");
        return -1;
    }

    for (i = 0; i < NUM_ANGLES; i++) {
        for (j = 0; j < NUM_WAVELENGTHS; j++) {
            gsl_spline2d_set(spline_resp, resp, i, j, pmtr[30+KPMTR_RESP+i+j*NUM_ANGLES]);
            gsl_spline2d_set(spline_reflec, reflec, i, j, pmtr[30+KPMTR_REFLEC+i+j*NUM_ANGLES]);
        }
    }

    gsl_spline2d_init(spline_resp, thetas, wavelengths, resp, NUM_ANGLES, NUM_WAVELENGTHS);
    gsl_spline2d_init(spline_reflec, thetas, wavelengths, reflec, NUM_ANGLES, NUM_WAVELENGTHS);

    initialized = 1;

    w = gsl_integration_cquad_workspace_alloc(100);

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

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

    if (status) {
        fprintf(stderr, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
        exit(1);
    }

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

    F.function = &gsl_pmt_response;

    for (i = 0; i < LEN(cos_thetas); i++) {
        params[0] = cos_thetas[i];

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

        weighted_resp[i] = result/norm;

        if (status) {
            fprintf(stderr, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            exit(1);
        }
    }

    F.function = &gsl_pmt_reflectivity;

    for (i = 0; i < LEN(cos_thetas); i++) {
        params[0] = cos_thetas[i];

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

        weighted_reflec[i] = result/norm;

        if (status) {
            fprintf(stderr, "error integrating cerenkov distribution: %s\n", gsl_strerror(status));
            exit(1);
        }
    }

    gsl_integration_cquad_workspace_free(w);

    return 0;
}

void pmt_response_free(void)
{
    if (spline_resp) gsl_spline2d_free(spline_resp);
    if (spline_reflec) gsl_spline2d_free(spline_reflec);
    if (xacc) gsl_interp_accel_free(xacc);
    if (yacc) gsl_interp_accel_free(yacc);
}