/* 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>

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.
 *
 * `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.
 *
 * 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)
{
    size_t i, j, k;
    static int index[MAX_PMTS];
    static int index2[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 wavelength0;
    double n_d2o;
    double t1, t2;
    int s;
    double pmt_dir[3];
    double expected;
    static double t[MAX_PMTS];
    static size_t p2[MAX_PMTS];
    double tmin;
    size_t max_tries;
    gsl_error_handler_t *old_handler;
    int status;

    wavelength0 = 400.0;
    n_d2o = get_index_snoman_d2o(wavelength0);

    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;
    }

    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;
            t[nhit] = ev->pmt_hits[i].t;
            nhit++;
        }
    }

    /* Only use the first 50% of the hits to avoid counting late/reflected
     * light. */
    k = nhit > 10 ? ceil(nhit/2.0) : nhit;
    gsl_sort_smallest_index(p2,k,t,1,nhit);

    /* Reorder the hits. */
    for (i = 0; i < k; i++) {
        index2[i] = index[p2[i]];
    }

    nhit = k;

    if (nhit < 5) {
        sprintf(quad_err, "only %i pmt hits. 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. */
        gsl_ran_choose(r,xs,4,index2,nhit,sizeof(int));

        /* 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)) {
            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)) {
                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;
    }

    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;
}