#include <stdlib.h>
#include <string.h>
#include "common.h"
#include "likelihood.h"
#include "chisquared.h"
#include <gsl/gsl_rng.h>
#include <gsl/gsl_histogram2d.h>

/* Options for the program */
struct options {
  int show_hist;      // whether to print the 2D histogram
  int only_hist;      // whether to skip the minimisations
  size_t num_events;  // number of generated events
  size_t bins_theta;  // number of bins for θ
  size_t bins_phi;    // number of bins for φ
};


int main(int argc, char **argv) {
  /* Set default options */
  struct options opts;
  opts.num_events = 50000;
  opts.show_hist  = 0;
  opts.only_hist  = 0;
  opts.bins_theta = 30;
  opts.bins_phi   = 60;

  /* Process CLI arguments */
  for (size_t i = 1; i < argc; i++) {
         if (!strcmp(argv[i], "-i")) opts.show_hist  = 1;
    else if (!strcmp(argv[i], "-I")) opts.only_hist  = 1;
    else if (!strcmp(argv[i], "-n")) opts.num_events = atol(argv[++i]);
    else if (!strcmp(argv[i], "-t")) opts.bins_theta = atol(argv[++i]);
    else if (!strcmp(argv[i], "-p")) opts.bins_phi   = atol(argv[++i]);
    else {
      fprintf(stderr, "Usage: %s -[hiIntp]\n", argv[0]);
      fprintf(stderr, "\t-h\tShow this message.\n");
      fprintf(stderr, "\t-i\tPrint the histogram to stdout.\n");
      fprintf(stderr, "\t-I\tPrint *only* the histogram.\n");
      fprintf(stderr, "\t-n N\tThe number of events to generate. (default: 50000)\n");
      fprintf(stderr, "\t-t N\tThe number of θ bins. (default: 30)\n");
      fprintf(stderr, "\t-p N\tThe number of φ bins. (default: 60)\n");
      return EXIT_FAILURE;
    }
  }

  /* Initialize an RNG. */
  gsl_rng_env_setup();
  gsl_rng *r = gsl_rng_alloc(gsl_rng_default);

  /* Initialise the sample data structure with
   * the number of events and allocate an array
   * for the events.
   */
  struct sample s = { opts.num_events, NULL };
  s.events = calloc(s.size, sizeof(struct event));

  /* Prepare a 2D histogram of the sample
   * with 30x60 bins.
   */
  gsl_histogram2d* hist =
    gsl_histogram2d_alloc(opts.bins_theta, opts.bins_phi);
  // set the ranges θ∈[0, π] and φ∈[0, 2π]
  gsl_histogram2d_set_ranges_uniform(hist, 0,M_PI, 0,2*M_PI);

  /* F(θ.φ) parameters we will generate
   * a sample with, and later try to recover.
   * Note: for this choice max F(θ,φ) ≈ 0.179.
   */
  gsl_vector *def_par = gsl_vector_alloc(3);
  gsl_vector_set(def_par, 0,  0.65);  // α
  gsl_vector_set(def_par, 1,  0.06);  // β
  gsl_vector_set(def_par, 2, -0.18);  // γ

  /* Generate the points (θ, φ) on the unit sphere
   * using the inverse transform:
   *   θ = acos(1 - 2X)
   *   φ = 2πY
   * where X,Y are random uniform variables in [0,1).
   * This ensures that θ,φ are uniformly distributed, then we
   * proceed to reject the events as follow:
   *   1. generate a y ∈ [0, max F(θ,φ)] uniformly
   *   2. reject the sample if y > F(θ,φ) else save it.
   */
  fputs("# event sampling\n\n", stderr);
  fprintf(stderr, "generating %ld events...", opts.num_events);
  struct event *e;
  for (size_t i = 0; i < s.size; i++){
    e = &s.events[i];
    do {
      e->th = acos(1 - 2*gsl_rng_uniform(r));
      e->ph = 2 * M_PI * gsl_rng_uniform(r);
    } while(0.2 * gsl_rng_uniform(r) > distr(def_par, e));

    // update the histogram
    gsl_histogram2d_increment(hist, e->th, e->ph);
  }
  fputs("done\n", stderr);

  // print histogram to stdout
  if (opts.show_hist || opts.only_hist) {
    fputs("\n# histogram\n", stderr);
    /* Print the bins on the first line:
     * this is needed by plot.py to make a
     * 2D histogram plot.
     */
    fprintf(stdout, "%ld %ld\n", opts.bins_theta, opts.bins_phi);
    gsl_histogram2d_fprintf(stdout, hist, "%g", "%g");
  }

  if (!opts.only_hist) {
    /* Estimate the parameters α,β,γ by
     * the maximum likelihood method.
     */
    min_result like = maxLikelihood(&s);

    /* Estimate the parameters α,β,γ by 
     * a minimum χ² method.
     */
    struct hist data;
    data.h = hist;
    data.n = s.size;
    min_result chi2 = minChiSquared(data);

    /* Compare the results with the original ones
     * and check whether they are compatible.
     */
    fputs("\n# vector decay compatibility\n", stderr);

    fputs("\n## likelihood results\n\n", stderr);
    compatibility(def_par, like);

    fputs("\n## χ² results\n\n", stderr);
    compatibility(def_par, chi2);

    /* Test the isotropic decay hypothesys
     * with the χ² results
     */
    fputs("\n# isotropic hypotesys\n\n", stderr);

    /* If the decay is isotropic then F = 1/4π.
     * Solving for α,β,γ yields α=1/3,β=γ=0.
     */
    gsl_vector *iso_par = gsl_vector_alloc(3);
    gsl_vector_set(iso_par, 0, 0.33);  // α
    gsl_vector_set(iso_par, 1, 0.00);  // β
    gsl_vector_set(iso_par, 2, 0.00);  // γ

    compatibility(iso_par, chi2);

    // free memory
    min_result_free(like);
    min_result_free(chi2);
  };

  // free memory
  gsl_rng_free(r);
  free(s.events);
  gsl_histogram2d_free(hist);
  gsl_vector_free(def_par);

  return EXIT_SUCCESS;
}