/* 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 "quad.h"
#include "event.h"
#include <gsl/gsl_randist.h>
#include <gsl/gsl_rng.h>
#include <unistd.h> /* for exit() */
#include <stdlib.h> /* for size_t */
#include <gsl/gsl_linalg.h>
#include "pdg.h"
#include "pmt.h"
#include <gsl/gsl_statistics_double.h>
#include "optics.h"
#include "vector.h"
#include "likelihood.h"
#include <gsl/gsl_sort.h>
#include <gsl/gsl_cblas.h>
#include "misc.h"
#include <gsl/gsl_errno.h>
#include "sno_charge.h"
#include "sno.h"
#include <gsl/gsl_permute.h>
#include "random.h"

char quad_err[256];

void my_handler(const char *reason, const char *file, int line, int gsl_errno)
{
    fprintf(stderr, "gsl: %s:%d: %s: %s\n", file, line, "ERROR", reason);
    return;
}

/* Returns an approximate vertex position and time using the QUAD fitter.
 *
 * The QUAD fitter was originally developed for SNO and estimates the event
 * vertex using a sort of Hough Transform[1]. It works by selecting 4 random
 * PMT hits and computing the position and time of the event which is
 * consistent with those hits. This procedure is repeated many times and the
 * median of the resulting "quad cloud" is found as an estimate of the event
 * position and time.
 *
 * Note: In the original formulation I believe instead of taking the median it
 * used a minimization routine to find the position and time with the highest
 * density of points, but since we are dealing with mostly high nhit events we
 * just take the median because it is easier and quicker.
 *
 * Update: This version of quad now also weights the PMT hits it selects by the
 * probability that they are more than 1 PE. The idea here is that we'd like to
 * ignore scattered and reflected light when selecting these points and most
 * scattered and reflected light is single photons.
 *
 * `ev` should be a pointer to the event.
 *
 * `pos` is a pointer to a double array of length 3.
 *
 * `t0` is a pointer to a double.
 *
 * `npoints` is the number of quad cloud points to compute. This function will
 * try to compute this many points, but it will stop trying after 10*npoints
 * times to avoid an infinite loop.
 *
 * `f` is the quantile of t0 to cut on. It should be between 0 and 1. The
 * reason to cut on a quantile of t0 is that for particles with tracks much
 * longer than a centimeter, the quad cloud of points generally follows the
 * whole track. Since we are interested in finding the position and time of the
 * start of the track, we'd like to only select the quad points at the start of
 * the track. To do this we only include quad cloud points in the first `f`
 * quantile when computing the position and time. For the default quad behavior
 * without cutting on t0, you should set `f` to 1.0.
 *
 * Returns 0 on success and `pos` and `t0` will be set to the approximate
 * vertex position and time respectively. On error, returns -1 and sets the
 * `quad_err` string.
 *
 * Example:
 *
 *     double pos[3], t0;
 *     if (quad(&ev, pos, &t0, 1000)) {
 *         fprintf(stderr, "error running the quad fitter: %s\n", quad_err);
 *         goto err;
 *     }
 *
 * [1] The only reference to the QUAD fitter that I can find from the SNO days
 * is Stephen Brice's PHD thesis which cites three SNO technical reports, but I
 * haven't been able to find these. The first two of these reports are by Bill
 * Frati who I assume wrote the initial implementation. */
int quad(event *ev, double *pos, double *t0, size_t npoints, double f)
{
    size_t i, j, k;
    static int index[MAX_PMTS];
    int nhit;
    const gsl_rng_type *T;
    gsl_rng *r;
    int xs[4];
    size_t nresults;
    double M[3][3], Minv[3][3];
    double K[3], g[3], h[3], n[3], tmp[3];
    double a, b, c;
    double pos1[3], pos2[3];
    static double results[QUAD_MAX][4];
    double c2;
    double n_d2o;
    double t1, t2;
    int s;
    double pmt_dir[3];
    double expected;
    static double w[MAX_PMTS];
    double tmin;
    size_t max_tries;
    gsl_error_handler_t *old_handler;
    int status;
    size_t reorder[QUAD_MAX];
    double pq = 0.0;

    n_d2o = get_avg_index_d2o();

    c2 = SPEED_OF_LIGHT*SPEED_OF_LIGHT/pow(n_d2o,2);

    if (npoints > QUAD_MAX) {
        sprintf(quad_err, "npoints must be less than %i", QUAD_MAX);
        return 1;
    }

    double expected_pe = 0.0;
    nhit = 0;
    for (i = 0; i < MAX_PMTS; i++) {
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL) continue;

        if (ev->pmt_hits[i].hit) {
            expected_pe += ev->pmt_hits[i].q;
            nhit += 1;
        }
    }
    expected_pe /= nhit;

    nhit = 0;
    for (i = 0; i < MAX_PMTS; i++) {
        if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL) continue;

        if (ev->pmt_hits[i].hit) {
            index[nhit] = i;

            /* Weights are equal to the probability that the hit is > 1 photon.
             * This is supposed to be a rough proxy for the probability that
             * it's not reflected and/or scattered light (which is typically
             * single photons). */
            if (ev->pmt_hits[i].q > QUAD_MAX_PE*get_qmean()/2) {
                /* If the charge is greater than QUAD_MAX_PE, it's almost
                 * certainly multiple photons so we don't waste time calculating P(1 PE|q). */
                w[nhit] = 1.0;
            } else {
                /* We want to calculate P(multiple photons|q) which we calculate as:
                 *
                 *     P(multiple photons|q) = 1 - P(1 PE|q) = 1 - P(q|1 PE)P(1 PE)/P(q)
                 *
                 * To calculate the second two quantities we assume the number
                 * of PE is Poisson distributed with a mean equal to
                 * expected_photons. */
                pq = 0.0;
                for (j = 1; j < QUAD_MAX_PE; j++) {
                    pq += get_pq(ev->pmt_hits[i].q,j)*gsl_ran_poisson_pdf(j,expected_pe);
                }

                w[nhit] = 1-get_pq(ev->pmt_hits[i].q,1)*gsl_ran_poisson_pdf(1,expected_pe)/pq;
            }
            nhit++;
        }
    }

    if (nhit < 5) {
        sprintf(quad_err, "only %i pmt hit(s). quad needs at least 5 points!", nhit);
        return 1;
    }

    T = gsl_rng_default;
    r = gsl_rng_alloc(T);

    gsl_permutation *p = gsl_permutation_alloc(3);

    i = 0;
    max_tries = npoints*10;
    nresults = 0;
    while (nresults < npoints && i++ < max_tries) {
        /* Choose 4 random hits. */
        ran_choose_weighted(xs,w,4,index,nhit);

        /* Shuffle them since GSL always returns the random choices in order.
         *
         * I'm not actually sure if that affects the quad calculation. */
        gsl_ran_shuffle(r,xs,4,sizeof(int));

        for (j = 1; j < 4; j++) {
            for (k = 0; k < 3; k++) {
                M[j-1][k] = pmts[xs[j]].pos[k] - pmts[xs[0]].pos[k];
            }
            n[j-1] = ev->pmt_hits[xs[j]].t - ev->pmt_hits[xs[0]].t;

            K[j-1] = (c2*(pow(ev->pmt_hits[xs[0]].t,2) - pow(ev->pmt_hits[xs[j]].t,2)) - DOT(pmts[xs[0]].pos,pmts[xs[0]].pos) + DOT(pmts[xs[j]].pos,pmts[xs[j]].pos))/2;
        }

        gsl_matrix_view m = gsl_matrix_view_array(&M[0][0],3,3);
        gsl_matrix_view minv = gsl_matrix_view_array(&Minv[0][0],3,3);
        gsl_vector_view k_view = gsl_vector_view_array(K,3);
        gsl_vector_view g_view = gsl_vector_view_array(g,3);
        gsl_vector_view h_view = gsl_vector_view_array(h,3);
        gsl_vector_view n_view = gsl_vector_view_array(n,3);

        gsl_linalg_LU_decomp(&m.matrix, p, &s);

        /* Ocassionaly the matrix is singular and we can't invert it. Since we
         * don't want our program to quit when this happens, we install a new
         * gsl error handler, do the matrix inversion, and then restore the old
         * handler. If the matrix inversion failed, then we just continue with
         * the next quad point. */

        /* save original handler, install new handler */
        old_handler = gsl_set_error_handler(&my_handler);

        status = gsl_linalg_LU_invert(&m.matrix, p, &minv.matrix);

        /* restore original handler */
        gsl_set_error_handler(old_handler);

        if (status) continue;

        gsl_blas_dgemv(CblasNoTrans,1.0,&minv.matrix,&k_view.vector,0.0,&g_view.vector);
        gsl_blas_dgemv(CblasNoTrans,1.0,&minv.matrix,&n_view.vector,0.0,&h_view.vector);

        SUB(tmp,pmts[xs[0]].pos,g);
        a = c2*(c2*DOT(h,h) - 1.0);
        b = -2*c2*(DOT(tmp,h) - ev->pmt_hits[xs[0]].t);
        c = DOT(tmp,tmp) - c2*pow(ev->pmt_hits[xs[0]].t,2);

        if (b*b - 4*a*c < 0) continue;

        t1 = (-b + sqrt(b*b - 4*a*c))/(2*a);
        COPY(tmp,h);
        MUL(tmp,c2*t1);
        ADD(pos1,g,tmp);

        /* Check the first result. */
        SUB(pmt_dir,pmts[xs[0]].pos,pos1);
        expected = t1 + NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT;

        tmin = ev->pmt_hits[xs[0]].t;
        for (j = 1; j < 4; j++) {
            if (ev->pmt_hits[xs[j]].t < tmin)
                tmin = ev->pmt_hits[xs[j]].t;
        }

        if (t1 < tmin && isclose(ev->pmt_hits[xs[0]].t,expected,0,1e-2)) {
            if (NORM(pos1) < PSUP_RADIUS) {
                results[nresults][0] = pos1[0];
                results[nresults][1] = pos1[1];
                results[nresults][2] = pos1[2];
                results[nresults][3] = t1;
                nresults++;
            }
        } else {
            /* Check the second solution. */
            t2 = (-b - sqrt(b*b - 4*a*c))/(2*a);
            COPY(tmp,h);
            MUL(tmp,c2*t2);
            ADD(pos2,g,tmp);

            SUB(pmt_dir,pmts[xs[0]].pos,pos2);
            expected = t2 + NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT;

            if (t2 < tmin && isclose(ev->pmt_hits[xs[0]].t,expected,0,1e-2)) {
                if (NORM(pos2) < PSUP_RADIUS) {
                    results[nresults][0] = pos2[0];
                    results[nresults][1] = pos2[1];
                    results[nresults][2] = pos2[2];
                    results[nresults][3] = t2;
                    nresults++;
                }
            }
        }
    }

    if (nresults < 1) {
        sprintf(quad_err, "no valid quad points found!");
        goto err;
    }

    /* Compute the permutation required to sort the results by t0. */
    gsl_sort_index(reorder,&results[0][3],4,nresults);

    /* Sort the results by t0. */
    gsl_permute(reorder,&results[0][0],4,nresults);
    gsl_permute(reorder,&results[0][1],4,nresults);
    gsl_permute(reorder,&results[0][2],4,nresults);
    gsl_permute(reorder,&results[0][3],4,nresults);

    if (f > 0.0 && f < 1.0) {
        /* Now, we filter only the results with t0 less than the quantile given
         * by `f`. The idea here is that for high energy particles which travel
         * macroscopic distances in the detector we want to only sample the
         * quad points near the start of the track, i.e. the points with the
         * earliest times. */
        nresults = (int) (f*nresults);
    }

    /* Sort the x, y, z, and t0 columns so we can calculate the median. Note:
     * The rows of the results array don't represent the quad cloud anymore
     * since x, y, z, and t0 are mixed up! */
    gsl_sort(&results[0][0],4,nresults);
    gsl_sort(&results[0][1],4,nresults);
    gsl_sort(&results[0][2],4,nresults);
    gsl_sort(&results[0][3],4,nresults);

    pos[0] = gsl_stats_median_from_sorted_data(&results[0][0],4,nresults);
    pos[1] = gsl_stats_median_from_sorted_data(&results[0][1],4,nresults);
    pos[2] = gsl_stats_median_from_sorted_data(&results[0][2],4,nresults);
    *t0 = gsl_stats_median_from_sorted_data(&results[0][3],4,nresults);

    gsl_permutation_free(p);
    gsl_rng_free(r);

    return 0;

err:
    gsl_permutation_free(p);
    gsl_rng_free(r);

    return 1;
}