/* 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 "muon.h"
#include "random.h"
#include "optics.h"
#include "quantum_efficiency.h"
#include <math.h>
#include <gsl/gsl_histogram.h>
#include "sno.h"
#include "pdg.h"
#include "vector.h"
#include "solid_angle.h"
#include <stdlib.h> /* for atoi() and strtod() */
#include <unistd.h> /* for exit() */
#include "scattering.h"
#include <errno.h> /* for errno */
#include <string.h> /* for strerror() */

void simulate_cos_theta_distribution(int N, gsl_histogram *h, double T, double theta0)
{
    /* Simulate the cos(theta) distribution around the original track direction
     * for a muon with kinetic energy T. The angle from the original track
     * distribution is simulated as a gaussian distribution with standard
     * deviation `theta0`. */
    int i;
    double theta, phi, wavelength, u, qe, index, cerenkov_angle, dir[3], n[3], dest[3], E, p, beta, cos_theta, thetax, thetay;

    i = 0;
    while (i < N) {
        /* Generate a random wavelength in the range 300-600 nm from the
         * distribution of Cerenkov light. */
        u = genrand_real2();
        wavelength = 300.0*600.0/(u*(300.0-600.0) + 600.0);

        qe = get_quantum_efficiency(wavelength);

        /* Check to see if the photon was detected. */
        if (genrand_real2() > qe) continue;

        index = get_index_snoman_d2o(wavelength);

        /* Calculate total energy */
        E = T + MUON_MASS;
        p = sqrt(E*E - MUON_MASS*MUON_MASS);
        beta = p/E;

        cerenkov_angle = acos(1/(index*beta));

        /* Assuming the muon track is dominated by small angle scattering, the
         * angular distribution looks like the product of two uncorrelated
         * Gaussian distributions with a standard deviation of `theta0` in the
         * plane perpendicular to the track direction. Here, we draw two random
         * angles and then compute the polar and azimuthal angle for the track
         * direction. */
        thetax = randn()*theta0;
        thetay = randn()*theta0;

        theta = sqrt(thetax*thetax + thetay*thetay);
        phi = atan2(thetay,thetax);

        n[0] = sin(theta)*cos(phi);
        n[1] = sin(theta)*sin(phi);
        n[2] = cos(theta);

        /* To compute the direction of the photon, we start with a vector which
         * has the same azimuthal angle as the track direction but is offset
         * from the track direction in the polar angle by the Cerenkov angle
         * and then rotate it around the track direction by a random angle
         * `phi`. */
        dir[0] = sin(cerenkov_angle + theta)*cos(phi);
        dir[1] = sin(cerenkov_angle + theta)*sin(phi);
        dir[2] = cos(cerenkov_angle + theta);

        phi = genrand_real2()*2*M_PI;

        rotate(dest,dir,n,phi);

        cos_theta = dest[2];

        gsl_histogram_increment(h, cos_theta);

        i += 1;
    }
}

void usage(void)
{
    fprintf(stderr,"Usage: ./test-likelihood [options]\n");
    fprintf(stderr,"  -n     number of events\n");
    fprintf(stderr,"  -T     kinetic energy of muon (MeV)\n");
    fprintf(stderr,"  -t     standard deviation of angular distribution\n");
    fprintf(stderr,"  -b     number of bins\n");
    fprintf(stderr,"  --xmin lowest value of cos(theta)\n");
    fprintf(stderr,"  --xmax highest value of cos(theta)\n");
    fprintf(stderr,"  -h     display this help message\n");
    exit(1);
}

int main(int argc, char **argv)
{
    size_t i, N, bins;
    double T, theta0;
    double E, p, beta;
    double xmin, xmax;

    N = 100000;
    bins = 1000;
    T = 1000.0;
    theta0 = 0.1;
    xmin = -1.0;
    xmax = 1.0;

    for (i = 1; i < argc; i++) {
        if (!strncmp(argv[i], "--", 2)) {
            if (!strcmp(argv[i]+2,"xmin")) {
                xmin = strtod(argv[++i],NULL);
                continue;
            } else if (!strcmp(argv[i]+2,"xmax")) {
                xmax = strtod(argv[++i],NULL);
                continue;
            }
        } else if (argv[i][0] == '-') {
            switch (argv[i][1]) {
            case 'n':
                N = atoi(argv[++i]);
                break;
            case 'b':
                bins = atoi(argv[++i]);
                break;
            case 'T':
                T = strtod(argv[++i],NULL);
                break;
            case 't':
                theta0 = strtod(argv[++i],NULL);
                break;
            case 'h':
                usage();
            default:
                fprintf(stderr, "unrecognized option '%s'\n", argv[i]);
                exit(1);
            }
        }
    }

    gsl_histogram *h = gsl_histogram_alloc(bins);
    gsl_histogram_set_ranges_uniform(h,xmin,xmax);

    simulate_cos_theta_distribution(N, h, T, theta0);

    gsl_histogram_scale(h, 1.0/gsl_histogram_sum(h));

    FILE *pipe = popen("graph -T X --bitmap-size 2000x2000 -X 'Cos(theta)' -Y Probability", "w");

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

    for (i = 0; i < h->n; i++) {
        fprintf(pipe, "%g %g\n", h->range[i], h->bin[i]);
        fprintf(pipe, "%g %g\n", h->range[i+1], h->bin[i]);
    }
    fprintf(pipe, "\n\n");

    gsl_histogram_reset(h);

    init_interpolation();

    /* Calculate total energy */
    E = T + MUON_MASS;
    p = sqrt(E*E - MUON_MASS*MUON_MASS);
    beta = p/E;

    for (i = 0; i < bins; i++) {
        double lo, hi;
        gsl_histogram_get_range(h, i, &lo, &hi);
        double cos_theta = (lo+hi)/2.0;
        double sin_theta = sqrt(1-cos_theta*cos_theta);
        h->bin[i] = get_probability(beta, cos_theta, sin_theta, theta0);
    }

    free_interpolation();

    printf("\n\n");

    gsl_histogram_scale(h, 1.0/gsl_histogram_sum(h));

    for (i = 0; i < h->n; i++) {
        fprintf(pipe, "%g %g\n", h->range[i], h->bin[i]);
        fprintf(pipe, "%g %g\n", h->range[i+1], h->bin[i]);
    }
    fprintf(pipe, "\n\n");

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

    gsl_histogram_free(h);

    return 0;
}