/* 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 "find_peaks.h"
#include "vector.h"
#include "event.h"
#include "pmt.h"
#include <stddef.h> /* for size_t */
#include <stdlib.h> /* for malloc() */
#include "optics.h"
#include <math.h> /* for exp() */
#include "misc.h"
#include "pmt.h"
#include "sno_charge.h"
#include <gsl/gsl_randist.h>
#include "pdg.h"

typedef struct peak {
    size_t i;
    size_t j;
    double value;
} peak;

/* Compare two different peaks.
 *
 * Note: We return 1 if peak b is greater than peak b and -1 if peak a is
 * greater than peak b. This is backwards from what you would normally expect,
 * but it's because we want to keep the peaks sorted in *descending* order. */
static int peak_compare(const void *a, const void *b)
{
    const peak *pa = (peak *) a;
    const peak *pb = (peak *) b;

    if (pa->value > pb->value)
        return -1;
    else if (pa->value < pb->value)
        return 1;

    return 0;
}

void find_peaks_array(double *x, size_t n, size_t m, size_t *imax, size_t *jmax, size_t *npeaks, size_t max_peaks, double threshold)
{
    /* Find a maximum of `max_peaks` in the 2D array `x` indexed as:
     *
     *     x[i,j] = x[i*m + j]
     *
     *  and store the indices in the two arrays `imax` and `jmax`.
     *
     *  Peaks are defined as any array element which is greater than all 8 of
     *  it's neighbors. Peaks are also required to be at least max*threshold
     *  times as high as the highest peak.
     *
     *  The returned peaks will always be the *highest* peaks and they will be
     *  returned in sorted order from highest to lowest. */
    size_t i, j;
    double max;
    peak *p;

    if (n*m == 0) {
        *npeaks = 0;
        return;
    }

    p = malloc(sizeof(peak)*(max_peaks+1));

    /* First, find the highest value in the array (which is guaranteed to be a peak). */
    max = x[0];
    p[0].i = 0;
    p[0].j = 0;
    p[0].value = max;
    for (i = 0; i < n; i++) {
        for (j = 0; j < m; j++) {
            if (x[i*m + j] > max) {
                max = x[i*m + j];
                p[0].i = i;
                p[0].j = j;
                p[0].value = max;
            }
        }
    }

    *npeaks = 1;

    if (*npeaks >= max_peaks) goto end;

    /* Now we look for other peaks which are at least max*threshold high. */
    for (i = 0; i < n; i++) {
        for (j = 0; j < m; j++) {
            if (x[i*m + j] <= max*threshold) continue;

            if (i == p[0].i && j == p[0].j) continue;

            /* Check to see if it is actually a peak. */
            if (x[i*m + (j+1) % m] < x[i*m + j]               && /* 0 +1 */
                x[i*m + (j+m-1) % m] < x[i*m + j]             && /* 0 -1 */
                x[((i+1) % n)*m + j] < x[i*m + j]             && /* +1 0 */
                x[((i+1) % n)*m + (j+1) % m] < x[i*m + j]     && /* +1 +1 */
                x[((i+1) % n)*m + (j+m-1) % m] < x[i*m + j]   && /* +1 -1 */
                x[((i+n-1) % n)*m + j] < x[i*m + j]           && /* -1 0 */
                x[((i+n-1) % n)*m + (j+1) % m] < x[i*m + j]   && /* -1 +1 */
                x[((i+n-1) % n)*m + (j+m-1) % m] < x[i*m + j]) { /* -1 -1 */
                p[*npeaks].i = i;
                p[*npeaks].j = j;
                p[*npeaks].value = x[i*m+j];
                *npeaks += 1;

                qsort(p,*npeaks,sizeof(peak),peak_compare);

                if (*npeaks >= max_peaks) *npeaks = max_peaks;
            }
        }
    }

end:
    for (i = 0; i < *npeaks; i++) {
        imax[i] = p[i].i;
        jmax[i] = p[i].j;
    }

    free(p);

    return;
}

void get_hough_transform(event *ev, double *pos, double t0, double *x, double *y, size_t n, size_t m, double *result, double *last, size_t len)
{
    /* Computes the "Hough transform" of the event `ev` and stores it in `result`. */
    size_t i, j, k;
    double dir[3], pmt_dir[3], cos_theta2;
    int skip;
    double *sin_theta, *cos_theta, *sin_phi, *cos_phi;
    double wavelength0, n_d2o;
    static double w[MAX_PMTS];
    double expected_pe;
    int nhit;
    double pq;
    double dt;

    sin_theta = malloc(sizeof(double)*n);
    cos_theta = malloc(sizeof(double)*n);
    sin_phi = malloc(sizeof(double)*n);
    cos_phi = malloc(sizeof(double)*n);

    wavelength0 = 400.0;
    n_d2o = get_index_snoman_d2o(wavelength0);

    /* Precompute sin(theta), cos(theta), sin(phi), and cos(phi) to speed
     * things up. */
    for (i = 0; i < n; i++) {
        sin_theta[i] = sin(x[i]);
        cos_theta[i] = cos(x[i]);
    }

    for (i = 0; i < m; i++) {
        sin_phi[i] = sin(y[i]);
        cos_phi[i] = cos(y[i]);
    }

    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;

    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) {
            /* 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 > FIND_PEAKS_MAX_PE*get_qmean()/2) {
                /* If the charge is greater than FIND_PEAKS_MAX_PE, it's almost
                 * certainly multiple photons so we don't waste time
                 * calculating P(1 PE|q). */
                w[i] = 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 < FIND_PEAKS_MAX_PE; j++) {
                    pq += get_pq(ev->pmt_hits[i].q,j)*gsl_ran_poisson_pdf(j,expected_pe);
                }

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

    /* Zero out the result array. */
    for (i = 0; i < n*m; i++) result[i] = 0.0;

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

        SUB(pmt_dir,pmts[i].pos,pos);

        dt = ev->pmt_hits[i].t - t0 - NORM(pmt_dir)*n_d2o/SPEED_OF_LIGHT;

        /* Skip PMT hits with a time residual of more than 10 nanoseconds since
         * they are most likely to be scattered and/or reflected light. */
        if (dt > 10) continue;

        normalize(pmt_dir);

        /* Ignore PMT hits within the Cerenkov cone of previously found peaks. */
        skip = 0;
        for (j = 0; j < len; j++) {
            if (fabs(DOT(pmt_dir,last+3*j) - 1/n_d2o) < 0.1) {
                skip = 1;
                break;
            }
        }

        if (skip) continue;

        for (j = 0; j < n; j++) {
            for (k = 0; k < m; k++) {
                dir[0] = sin_theta[j]*cos_phi[k];
                dir[1] = sin_theta[j]*sin_phi[k];
                dir[2] = cos_theta[j];

                cos_theta2 = DOT(pmt_dir,dir);

                result[j*n + k] += w[i]*exp(-pow(cos_theta2-1.0/n_d2o,2)/0.01);
            }
        }
    }

    free(sin_theta);
    free(cos_theta);
    free(sin_phi);
    free(cos_phi);
}

/* Finds rings in the event `ev` by looking for peaks in the Hough transform.
 * The directions of potential particle rings are stored in the arrays
 * `peak_theta` and `peak_phi`. The number of rings found are stored in the
 * variable `npeaks`.
 *
 * This function first computes the Hough transform of the event `ev` and looks
 * for the highest peak. The next peak is then found by computing the Hough
 * transform after subtracting off the previous ring and ignoring any PMTs
 * within the Cerenkov cone of the first peak. This process is repeated until
 * `max_peaks` peaks are found.
 *
 * In addition, only peaks which are `delta` radians away from other peaks are
 * returned. So, for example if three peaks are found but the first two are
 * very close to each other, this function will only return the first and the
 * third peaks. */
void find_peaks(event *ev, double *pos, double t0, size_t n, size_t m, double *peak_theta, double *peak_phi, size_t *npeaks, size_t max_peaks, double delta)
{
    size_t i, j;
    double *x, *y, *result, *dir;
    size_t *imax, *jmax;
    size_t max;
    double theta, phi;
    int unique;

    x = calloc(n,sizeof(double));
    y = calloc(m,sizeof(double));
    result = malloc(n*m*sizeof(double));
    dir = calloc(max_peaks*3,sizeof(double));

    for (i = 0; i < n; i++) {
        x[i] = i*M_PI/(n-1);
    }

    for (i = 0; i < m; i++) {
        y[i] = i*2*M_PI/(m-1);
    }

    imax = calloc(max_peaks,sizeof(size_t));
    jmax = calloc(max_peaks,sizeof(size_t));

    *npeaks = 0;

    for (i = 0; i < max_peaks; i++) {
        get_hough_transform(ev,pos,t0,x,y,n,m,result,dir,i);
        max = argmax(result,n*m);
        theta = x[max/m];
        phi = y[max % m];
        get_dir(dir+i*3,theta,phi);

        /* Only add the latest peak to the results if it's more than `delta`
         * radians away from all other results so far. */
        unique = 1;
        for (j = 0; j < *npeaks; j++) {
            /* If the angle between the latest peak and a previous peak is
             * within `delta` radians, it's not unique. */
            if (acos(DOT(dir+i*3,dir+j*3)) < delta) unique = 0;
        }

        /* Add it to the results if it's unique. */
        if (unique) {
            peak_theta[*npeaks] = theta;
            peak_phi[*npeaks] = phi;
            *npeaks += 1;
        }
    }

    free(imax);
    free(jmax);
    free(x);
    free(y);
    free(result);
    free(dir);
}