/* 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 <stdio.h>
#include "find_peaks.h"
#include "misc.h"
#include "zebra.h"
#include "zdab_utils.h"
#include <stddef.h> /* for size_t */
#include <unistd.h> /* for exit() */
#include "pmt.h"
#include <math.h> /* for M_PI */
#include <errno.h> /* for errno */
#include <string.h> /* for strerror() */
#include "vector.h"
#include "util.h"
#include "sno.h"
#include "quad.h"
#include "likelihood.h"
#include "optics.h"
#include "sno_charge.h"
#include "db.h"
#include "pmt_response.h"
#include "dqxx.h"

#define EV_RECORD      0x45562020
#define MCTK_RECORD    0x4d43544b
#define MCVX_RECORD    0x4d435658

#define TEST_FIND_PEAKS_NUM_ANGLES 1000

void plot_hough_transform(double *x, double *y, double *z, size_t n, size_t m)
{
    size_t i, j;

    FILE *pipe = popen("gnuplot -p", "w");

    if (!pipe) {
        fprintf(stderr, "error running gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    fprintf(pipe, "set title 'Hough Transform'\n");
    /* Not entirely sure what these do, but following the instructions from
     * http://lowrank.net/gnuplot/plot3d2-e.html. */
    fprintf(pipe, "set dgrid3d %zu,%zu\n",n,m);
    fprintf(pipe, "set hidden3d\n");
    fprintf(pipe, "set contour\n");
    fprintf(pipe, "splot '-' u 1:2:3 with lines\n");

    for (i = 0; i < n; i++) {
        for (j = 0; j < m; j++) {
            fprintf(pipe, "%.10g %.10g %.10g\n", x[i], y[j], z[i*m+j]);
        }
    }
    fprintf(pipe,"e\n");

    /* Pause so that you can rotate the 3D graph. */
    fprintf(pipe,"pause mouse keypress\n");

    if (pclose(pipe)) {
        fprintf(stderr, "error closing gnuplot command: %s\n", strerror(errno));
        exit(1);
    }
}

void plot_find_peaks(event *ev, double *pos, double *peak_theta, double *peak_phi, size_t n)
{
    size_t i, j;
    double r, theta, phi, last_phi;
    double dir[3], k[3], tmp[3], tmp2[3], hit[3];
    double l;
    double wavelength0, n_d2o;

    wavelength0 = 400.0;
    n_d2o = get_index_snoman_d2o(wavelength0);

    FILE *pipe = popen("gnuplot -p", "w");

    if (!pipe) {
        fprintf(stderr, "error running gnuplot command: %s\n", strerror(errno));
        exit(1);
    }

    fprintf(pipe, "set macros\n");
    fprintf(pipe, "load 'viridis.pal'\n");
    fprintf(pipe, "set title 'Hough Transform'\n");
    /* Not entirely sure what these do, but following the instructions from
     * http://lowrank.net/gnuplot/plot3d2-e.html. */
    fprintf(pipe, "set xrange [0:6.28]\n");
    fprintf(pipe, "set yrange [3.14:0]\n");
    fprintf(pipe, "plot '-' u 1:2:3:4 with circles palette fillstyle solid, '-' u 1:2 with lines lc rgb \"red\" lw 2\n");

    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) {
            r = NORM(pmts[i].pos);
            theta = acos(pmts[i].pos[2]/r);
            phi = atan2(pmts[i].pos[1],pmts[i].pos[0]);
            phi = (phi < 0) ? phi + 2*M_PI: phi;
            fprintf(pipe, "%.10g %.10g %.10g %.10g\n", phi, theta, 0.01, ev->pmt_hits[i].q);
        }
    }
    fprintf(pipe,"e\n");

    /* Now we draw the Cerenkov rings for each peak. We do this by computing
     * rays which lie along the Cerenkov cone, and intersecting each of these
     * with the PSUP. */
    for (i = 0; i < n; i++) {
        /* Compute the direction of the peak. */
        dir[0] = sin(peak_theta[i])*cos(peak_phi[i]);
        dir[1] = sin(peak_theta[i])*sin(peak_phi[i]);
        dir[2] = cos(peak_theta[i]);

        k[0] = 0;
        k[1] = 0;
        k[2] = 1;

        CROSS(tmp,dir,k);

        normalize(tmp);

        /* Rotate the direction by the Cerenkov angle. */
        rotate(tmp2,dir,tmp,acos(1/n_d2o));

        for (j = 0; j < TEST_FIND_PEAKS_NUM_ANGLES; j++) {
            /* Rotate the new direction around the peak by 2 pi. */
            rotate(tmp,tmp2,dir,j*2*M_PI/(TEST_FIND_PEAKS_NUM_ANGLES-1));

            /* Calculate the intersection of the ray with the sphere. */
            if (!intersect_sphere(pos,tmp,PSUP_RADIUS,&l)) continue;

            hit[0] = pos[0] + l*tmp[0];
            hit[1] = pos[1] + l*tmp[1];
            hit[2] = pos[2] + l*tmp[2];

            /* Calculate the theta and phi angles for the intersection. */
            theta = acos(hit[2]/NORM(hit));
            phi = atan2(hit[1],hit[0]);

            /* Since the phi angle is periodic, avoid drawing a horizontal line
             * when we switch from the left side of the plot to the right side
             * or vice versa. */
            if ((j > 0) && ((last_phi < 0 && phi > 0) || (last_phi > 0 && phi < 0)))
                fprintf(pipe,"\n");

            last_phi = phi;

            phi = (phi < 0) ? phi + 2*M_PI: phi;

            fprintf(pipe, "%.10g %.10g\n", phi, theta);
        }
        fprintf(pipe,"\n");
    }
    fprintf(pipe,"e\n");

    /* Pause so that you can rotate the 3D graph. */
    fprintf(pipe,"pause mouse keypress\n");

    if (pclose(pipe)) {
        fprintf(stderr, "error closing gnuplot command: %s\n", strerror(errno));
        exit(1);
    }
}

void usage(void)
{
    fprintf(stderr,"Usage: ./test-find-peaks [options] FILENAME\n");
    fprintf(stderr,"  -n     number of peaks to plot\n");
    fprintf(stderr,"  --plot plot hough transform\n");
    fprintf(stderr,"  --gtid               only fit a single GTID\n");
    fprintf(stderr,"  -h     display this help message\n");
    exit(1);
}

int main(int argc, char **argv)
{
    int i;
    char *filename = NULL;
    zebraFile *f;
    zebraBank bmast, b;
    EVBank bev;
    event ev = {0};
    int rv;
    double pos[3];
    double peak_theta[10];
    double peak_phi[10];
    size_t npeaks, max_npeaks = 10;
    int plot = 0;
    int status;
    double t0;
    int last_run;
    char dqxx_file[256];
    int32_t gtid = -1;

    for (i = 1; i < argc; i++) {
        if (strlen(argv[i]) >= 2 && !strncmp(argv[i], "--", 2)) {
            if (!strcmp(argv[i]+2,"plot")) {
                plot = 1;
                continue;
            } else if (!strcmp(argv[i]+2,"gtid")) {
                gtid = atoi(argv[++i]);
                continue;
            }
        } else if (argv[i][0] == '-') {
            switch (argv[i][1]) {
            case 'n':
                max_npeaks = atoi(argv[++i]);
                if (max_npeaks > LEN(peak_theta)) {
                    fprintf(stderr, "number of peaks must be less than %lu\n", LEN(peak_theta));
                    exit(1);
                }
                break;
            case 'h':
                usage();
            default:
                fprintf(stderr, "unrecognized option '%s'\n", argv[i]);
                exit(1);
            }
        } else {
            filename = argv[i];
        }
    }

    if (!filename)
        usage();

    load_pmt_info();

    init_charge();

    f = zebra_open(filename);

    if (!f) {
        fprintf(stderr, "%s\n", zebra_err);
        return 1;
    }

    dict *db = db_init();

    last_run = -1;

    if (load_file(db, "pmt_response_qoca_d2o_20060216.dat", 0)) {
        fprintf(stderr, "failed to load pmt_response_qoca_d2o_20060216.dat: %s\n", db_err);
        goto err;
    }

    if (load_file(db, "rsp_rayleigh.dat", 0)) {
        fprintf(stderr, "failed to load rsp_rayleigh.dat: %s\n", db_err);
        goto err;
    }

    if (pmt_response_init(db)) {
        fprintf(stderr, "failed to initialize PMTR bank: %s\n", pmtr_err);
        goto err;
    }

    if (optics_init()) {
        fprintf(stderr, "failed to initialize optics: %s\n", optics_err);
        goto err;
    }

    while (1) {
        rv = zebra_read_next_logical_record(f);

        if (rv == -1) {
            fprintf(stderr, "error getting MAST bank: %s\n", zebra_err);
            goto err;
        } else if (rv == 1) {
            /* EOF */
            break;
        }

        if (f->mast_bank == -1) {
            fprintf(stderr, "no MAST bank in logical record! Skipping...\n");
            continue;
        }

        rv = zebra_get_bank(f, &bmast, f->mast_bank);

        if (rv) {
            fprintf(stderr, "error getting MAST bank: %s\n", zebra_err);
            goto err;
        }

        if (bmast.links[KMAST_EV-1] == 0) {
            /* First logical record in SNOCR files don't have an EV bank. */
            continue;
        }

        rv = zebra_get_bank(f,&b,bmast.links[KMAST_EV-1]);

        if (rv) {
            fprintf(stderr, "error getting EV bank: %s\n", zebra_err);
            goto err;
        }

        /* Skip to the last event so we can traverse them in reverse order. The
         * reason for this is that for some reason SNOMAN puts the events in
         * reverse order within each logical record. */
        while (b.next) {
            rv = zebra_get_bank(f,&b,b.next);

            if (rv) {
                fprintf(stderr, "error getting EV bank: %s\n", zebra_err);
                goto err;
            }
        }

        unpack_ev(b.data, &bev);
        ev.run = bev.run;
        ev.gtid = bev.gtr_id;

        if (gtid > 0 && ev.gtid != gtid) continue;

        if (bev.run != last_run) {
            printf("loading DQXX file for run %010i\n", bev.run);

            sprintf(dqxx_file, "DQXX_%010i.dat", bev.run);
            if (load_file(db, dqxx_file, 1)) {
                fprintf(stderr, "failed to load %s: %s\n", dqxx_file, db_err);
                goto err;
            }

            if (dqxx_init(db, &ev)) {
                fprintf(stderr, "failed to initialize DQXX bank: %s\n", dqxx_err);
                goto err;
            }

            last_run = bev.run;
        }

        rv = get_event(f,&ev,&b);

        size_t n = 100;
        size_t m = 100;

        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] = i*M_PI/(n-1);
        }

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

        /* Guess the position and t0 of the event using the QUAD fitter. */
        status = quad(&ev,pos,&t0,10000,0.1);

        if (status || t0 < 0 || t0 > GTVALID || NORM(pos) > PSUP_RADIUS) {
            /* If the QUAD fitter fails or returns something outside the PSUP or
             * with an invalid time we just assume it's at the center. */
            fprintf(stderr, "quad returned pos = %.2f, %.2f, %.2f t0 = %.2f. Assuming vertex is at the center!\n",
                    pos[0], pos[1], pos[2], t0);
            pos[0] = 0.0;
            pos[1] = 0.0;
            pos[2] = 0.0;
        }

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

        find_peaks(&ev,pos,n,m,peak_theta,peak_phi,&npeaks,max_npeaks,0.1);

        printf("gtid %i\n", ev.gtid);
        for (i = 0; i < npeaks; i++) {
            printf("%.2f %.2f\n", peak_theta[i], peak_phi[i]);
        }

        if (plot)
            plot_find_peaks(&ev,pos,peak_theta,peak_phi,npeaks);
    }

    zebra_close(f);

    return 0;

err:
    zebra_close(f);

    return 1;
}