#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() */

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 *x, double *y, size_t n, size_t m, double *result)
{
    /* 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_theta;

    double wavelength0 = 400.0;
    double n_d2o = get_index_snoman_d2o(wavelength0);

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

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

        normalize(pmt_dir);

        for (j = 0; j < n; j++) {
            for (k = 0; k < m; k++) {
                double r = 1+x[j]*x[j]+y[k]*y[k];
                dir[0] = 2*x[j]/r;
                dir[1] = 2*y[k]/r;
                dir[2] = (-1+x[j]*x[j]+y[k]*y[k])/r;

                cos_theta = DOT(pmt_dir,dir);

                result[j*m + k] += ev->pmt_hits[i].qhs*exp(-fabs(cos_theta-1.0/n_d2o)/0.1);
            }
        }
    }
}

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 threshold)
{
    size_t i;

    double *x = calloc(n,sizeof(double));
    double *y = calloc(m,sizeof(double));
    double *result = calloc(n*m,sizeof(double));

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

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

    get_hough_transform(ev,pos,x,y,n,m,result);

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

    find_peaks_array(result,n,m,imax,jmax,npeaks,max_peaks,threshold);

    for (i = 0; i < *npeaks; i++) {
        double R, theta;
        R = sqrt(x[imax[i]]*x[imax[i]]+y[jmax[i]]*y[jmax[i]]);
        theta = atan2(y[jmax[i]],x[imax[i]]);
        peak_theta[i] = 2*atan(1/R);
        peak_phi[i] = theta;
    }

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