#include "likelihood.h"
#include <stdio.h>
#include "zebra.h"
#include "Record_Info.h"
#include "event.h"
#include "zdab_utils.h"
#include "scattering.h"
#include "pmt.h"
#include "sno_charge.h"
#include "db.h"
#include "dqxx.h"
#include <nlopt.h>
#include <math.h> /* for sin(), cos(), etc. */
#include <sys/time.h> /* for gettimeofday() */
#include <inttypes.h> /* for PRIu32 macro */
#include <string.h> /* for memcpy() */
#include <errno.h> /* for errno */
#include "pdg.h"
#include "optics.h"

#define EV_RECORD    0x45562020 // 'EV  '  (as written to ZDAB file)

static size_t iter;

typedef struct fitParams {
    event *ev;
    double epsrel;
} fitParams;

/* In order to start the fitter close to the minimum, we first do a series of
 * "quick" minimizations starting at the following points. We keep track of the
 * parameters with the best likelihood value and then start the "real"
 * minimization from those parameters. */

static struct startingParameters {
    double x;
    double y;
    double z;
    double theta;
    double phi;
} startingParameters[] = {
    {0.0, 0.0, 0.0,    0.0, 0.0},
    {0.0, 0.0, 0.0, M_PI/2, 0.0},
    {0.0, 0.0, 0.0, M_PI/2, M_PI/2},
    {0.0, 0.0, 0.0, M_PI/2, M_PI},
    {0.0, 0.0, 0.0, M_PI/2, 3*M_PI/2},
    {0.0, 0.0, 0.0,   M_PI, 0.0},
};

double nll(unsigned int n, const double *x, double *grad, void *params)
{
    fitParams *fpars = (fitParams *) params;
    double T, theta, phi, t0;
    double pos[3], dir[3];
    double fval;
    double z1[1], z2[1];
    struct timeval tv_start, tv_stop;

    T = x[0];

    pos[0] = x[1];
    pos[1] = x[2];
    pos[2] = x[3];

    theta = x[4];
    phi = x[5];
    dir[0] = sin(theta)*cos(phi);
    dir[1] = sin(theta)*sin(phi);
    dir[2] = cos(theta);

    t0 = x[6];

    z1[0] = x[7];
    z2[0] = x[8];

    gettimeofday(&tv_start, NULL);
    fval = nll_muon(fpars->ev, T, pos, dir, t0, z1, z2, 1, fpars->epsrel);
    gettimeofday(&tv_stop, NULL);

    long long elapsed = (tv_stop.tv_sec - tv_start.tv_sec)*1000 + (tv_stop.tv_usec - tv_start.tv_usec)/1000;

    printf("%5zu %10.2f %7.2f %7.2f %7.2f %5.2f %5.2f %5.2f %5.2f %5.2f f() = %7.3e took %lld ms\n",
           iter++,
           x[0],
           x[1],
           x[2],
           x[3],
           x[4],
           x[5],
           x[6],
           x[7],
           x[8],
           fval,
           elapsed);

    return fval;
}

int fit_event(event *ev, double *xopt, double *fmin)
{
    size_t i;
    fitParams fpars;
    double x[9], ss[9], lb[9], ub[9], fval, n;
    int rv;

    nlopt_opt opt = nlopt_create(NLOPT_LN_BOBYQA, 9);
    nlopt_set_min_objective(opt,nll,&fpars);

    x[0] = 1200.0;
    x[1] = 0.0;
    x[2] = 0.0;
    x[3] = 0.0;
    x[4] = 1.57;
    x[5] = 0.0;
    x[6] = 130.0;
    x[7] = 0.0;
    x[8] = 0.0;

    ss[0] = 10.0;
    ss[1] = 10.0;
    ss[2] = 10.0;
    ss[3] = 10.0;
    ss[4] = 0.1;
    ss[5] = 0.1;
    ss[6] = 1.0;
    ss[7] = 0.1;
    ss[8] = 0.1;

    n = get_index_snoman_d2o(400.0);

    lb[0] = MUON_MASS/sqrt(1.0-1.0/(n*n));
    lb[1] = -1000.0;
    lb[2] = -1000.0;
    lb[3] = -1000.0;
    lb[4] = -INFINITY;
    lb[5] = -INFINITY;
    lb[6] = 0.0;
    lb[7] = -10.0;
    lb[8] = -10.0;

    ub[0] = 10000.0;
    ub[1] = 1000.0;
    ub[2] = 1000.0;
    ub[3] = 1000.0;
    ub[4] = INFINITY;
    ub[5] = INFINITY;
    ub[6] = 400.0;
    ub[7] = 10.0;
    ub[8] = 10.0;

    nlopt_set_lower_bounds(opt, lb);
    nlopt_set_upper_bounds(opt, ub);

    nlopt_set_initial_step(opt, ss);

    fpars.ev = ev;

    iter = 0;

    /* First we do a set of "quick" minimizations to try and start the
     * minimizer close to the minimum. To make these function evaluations
     * faster, we set the absolute tolerance on the likelihood to 1.0, the
     * maximum number of function evaluations to 100, and the relative
     * tolerance on the numerical integration to 10%. */
    fpars.epsrel = 1e-1;
    nlopt_set_ftol_abs(opt, 1.0);
    nlopt_set_maxeval(opt, 100);

    for (i = 0; i < sizeof(startingParameters)/sizeof(startingParameters[0]); i++) {
        x[0] = 1200.0;
        x[1] = startingParameters[i].x;
        x[2] = startingParameters[i].y;
        x[3] = startingParameters[i].z;
        x[4] = startingParameters[i].theta;
        x[5] = startingParameters[i].phi;
        x[6] = 130.0;
        x[7] = 0.0;
        x[8] = 0.0;

        rv = nlopt_optimize(opt,x,&fval);

        if (fval < *fmin || i == 0) {
            *fmin = fval;
            memcpy(xopt,x,sizeof(x));
        }
    }

    memcpy(x,xopt,sizeof(x));

    /* Now, we do the "real" minimization. */
    fpars.epsrel = 1e-4;
    nlopt_set_ftol_abs(opt, 1e-5);
    nlopt_set_maxeval(opt, 1000);


    do {
        *fmin = fval;
        memcpy(xopt,x,sizeof(x));
        rv = nlopt_optimize(opt,x,&fval);
    } while (fval < *fmin && fabs(fval-*fmin) > 1e-5);

    if (fval < *fmin) {
        *fmin = fval;
        memcpy(xopt,x,sizeof(x));
    }

    nlopt_destroy(opt);

    return rv;
}

void usage(void)
{
    fprintf(stderr,"Usage: ./fit [options] FILENAME\n");
    fprintf(stderr,"  -o      output file\n");
    fprintf(stderr,"  -h      display this help message\n");
    exit(1);
}

int main(int argc, char **argv)
{
    int i;
    zebraFile *f;
    bank b;
    int rv;
    PMTBank bpmt;
    EVBank bev;
    event ev = {0};
    int crate, card, channel;
    int id;
    double fmin;
    double xopt[9];
    char *filename = NULL;
    char *output = NULL;
    FILE *fout = NULL;

    for (i = 1; i < argc; i++) {
        if (argv[i][0] == '-') {
            switch (argv[i][1]) {
            case 'o':
                output = argv[++i];
                break;
            case 'h':
                usage();
            default:
                fprintf(stderr, "unrecognized option '%s'\n", argv[i]);
                exit(1);
            }
        } else {
            filename = argv[i];
        }
    }

    if (!filename)
        usage();

    f = zebra_open(filename);

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

    if (output) {
        fout = fopen(output, "w");

        if (!fout) {
            fprintf(stderr, "failed to open '%s': %s\n", output, strerror(errno));
            return 1;
        }
    }

    load_pmt_info();

    for (i = 0; i < MAX_PMTS; i++) {
        ev.pmt_hits[i].hit = 0;
        ev.pmt_hits[i].flags = 0;
    }

    ev.run = -1;

    init_interpolation();
    init_charge();
    dict *db = db_init();

    if (load_file(db, "DQXX_0000010000.dat")) {
        fprintf(stderr, "failed to load DQXX_0000010000.dat: %s\n", db_err);
        exit(1);
    }

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

    while (1) {
        rv = next_bank(f, &b);

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

        if (b.name == PMT_RECORD) {
            unpack_pmt(b.data, &bpmt);
            card = bpmt.pin/1024;
            crate = (bpmt.pin % 1024)/32;
            channel = bpmt.pin % 32;
            id = crate*512 + card*32 + channel;
            ev.pmt_hits[id].hit = 1;
            ev.pmt_hits[id].t = bpmt.pt;
            ev.pmt_hits[id].qhl = bpmt.phl;
            ev.pmt_hits[id].qhs = bpmt.phs;
            ev.pmt_hits[id].qlx = bpmt.plx;
            //ev.pmt_hits[id].flags |= bpmt.pn & KPF_DIS;
        } else if (b.name == EV_RECORD) {
            unpack_ev(b.data, &bev);
            if (ev.run != -1) {
                if (ev.run != bev.run || ev.gtid != bev.gtr_id) {
                    /* New event, so we need to fit the old event. */
                    fit_event(&ev,xopt,&fmin);

                    if (fout) {
                        fprintf(fout, "%" PRIu32 " %10.2f %7.2f %7.2f %7.2f %5.2f %5.2f %5.2f %5.2f %5.2f f() = %7.3e\n",
                                ev.gtid,
                                xopt[0],
                                xopt[1],
                                xopt[2],
                                xopt[3],
                                xopt[4],
                                xopt[5],
                                xopt[6],
                                xopt[7],
                                xopt[8],
                                fmin);
                        fflush(fout);
                    }
                }
            }
            ev.run = bev.run;
            ev.gtid = bev.gtr_id;

            for (i = 0; i < MAX_PMTS; i++) {
                ev.pmt_hits[i].hit = 0;
            }
        }
    }

    free_interpolation();

    db_free(db);

    if (fout) fclose(fout);

    zebra_close(f);

    return 0;

err:

    zebra_close(f);
    return 1;
}