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

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

double get_qhs_avg(event *ev, double *pos, double *dir)
{
    /* Returns the average QHS value for all PMTs within the Cerenkov cone of a
     * particle at position `pos` travelling in direction `dir`. */
    size_t i;
    double qhs_sum, n_d2o;
    double pmt_dir[3];
    size_t n;

    n_d2o = get_index_snoman_d2o(400.0);

    n = 0;
    qhs_sum = 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);
        normalize(pmt_dir);
        if (DOT(pmt_dir,dir) > 1/n_d2o) {
            qhs_sum += ev->pmt_hits[i].q;
            n += 1;
        }
    }

    return qhs_sum/n;
}

void get_hough_transform(event *ev, double *pos, 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;
    double qhs[MAX_PMTS], qhs_avg;

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

    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;
        qhs[i] = ev->pmt_hits[i].q;
    }

    /* Subtract off previous rings. */
    for (i = 0; i < len; i++) {
        qhs_avg = get_qhs_avg(ev,pos,last+3*i);
        for (j = 0; j < MAX_PMTS; j++) {
            if (ev->pmt_hits[i].flags || pmts[i].pmt_type != PMT_NORMAL || !ev->pmt_hits[i].hit) continue;

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

            normalize(pmt_dir);

            cos_theta2 = DOT(pmt_dir,last+3*i);

            qhs[j] -= qhs_avg*exp(-fabs(cos_theta2-1.0/n_d2o)/0.1);

            if (qhs[j] < 0.0) qhs[j] = 0.0;
        }
    }

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

        normalize(pmt_dir);

        /* Ignore PMT hits within the Cerenkov cone of previously found peaks. */
        skip = 0;
        for (j = 0; j < len; j++) {
            if (DOT(pmt_dir,last+3*j) > 1/n_d2o) {
                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] += qhs[i]*exp(-fabs(cos_theta2-1.0/n_d2o)/0.1);
            }
        }
    }

    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, 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,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 < i; 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);
}