#include "lib.h"
#include <math.h>
#include <stdlib.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_min.h>
#include <gsl/gsl_deriv.h>


/* Process CLI arguments and print usage */
int parser(size_t *N, size_t *n, double *p_max,
           size_t argc, char **argv, size_t *go) {
  for (size_t i = 1; i < argc; i++) {
         if (!strcmp(argv[i], "-n")) *N     = atol(argv[++i]);
    else if (!strcmp(argv[i], "-b")) *n     = atol(argv[++i]);
    else if (!strcmp(argv[i], "-p")) *p_max = atof(argv[++i]);
    else if (!strcmp(argv[i], "-o")) *go = 1;
    else {
      fprintf(stderr, "Usage: %s -[hnbpo]\n", argv[0]);
      fprintf(stderr, "\t-h\tShow this message.\n");
      fprintf(stderr, "\t-n N\tThe number of particles to generate. (default: 50000)\n");
      fprintf(stderr, "\t-b N\tThe number of bins of the histogram. (default: 50)\n");
      fprintf(stderr, "\t-p PMAX\tThe maximum value of momentum. (default: 10)\n");
      fprintf(stderr, "\t-o\tPrint histogram to stdout.\n");
      return 0;
    }
  }
  return 1;
}


int main(int argc, char **argv) {
  // Set default options.
  size_t N = 50000;
  size_t n = 50;
  double p_max = 10;
  size_t go = 0;

  // Get eventual CLI arguments.
  int res = parser(&N, &n, &p_max, argc, argv, &go);
  if (go == 0 && res == 1) {
    printf("Generating histogram with:\n"
           "%ld points\n"
           "%ld bins\n"
           "p_max = %.3f\n\n", N, n, p_max);
  }
  else if (res == 0)
    return EXIT_FAILURE;

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

  /* Generate the angle θ uniformly distributed on a sphere using the
   * inverse transform:
   *
   *   θ = acos(1 - 2X)
   *
   * where X is a random uniform variable in [0,1), and the module p of
   * the vector:
   *
   *   p² = p_v² + p_h²
   * 
   * uniformly distributed between 0 and p_max. The two components are
   * then computed as:
   *
   *   p_v = p⋅cos(θ)
   *   p_h = p⋅sin(θ)
   *
   * The histogram is updated this way.
   * The j-th bin where p_h goes in is given by:
   *
   *   step = p_max / n
   *   j = floor(p_h / step)
   *
   * Thus an histogram was created and a structure containing the number of
   * entries in each bin and the sum of |p_v| in each of them is created and
   * filled while generating the events (struct bin).
   */
  struct bin *histo = calloc(n, sizeof(struct bin));

  // Some useful variables.
  double step = p_max / n;
  struct bin *b;
  double theta;
  double p;
  double p_v;
  double p_h;
  size_t j;

  for (size_t i = 0; i < N; i++) {
    // Generate the event
    theta = acos(1 - 2*gsl_rng_uniform(r));
    p = p_max * gsl_rng_uniform(r);

    // Compute the components
    p_v = p * cos(theta);
    p_h = p * sin(theta);

    // Update the histogram
    j = floor(p_h / step);
    b = &histo[j];
    b -> amo++;
    b -> sum += fabs(p_v);
  }

  /* Compute the mean value of each bin and print it to stodut
   * together with other useful things to make the histogram.
   */
  if (go == 1) {
    printf("bins: \t%ld\n", n);
    printf("step: \t%.5f\n", step);
  }
  for (size_t i = 0; i < n; i++) {
    histo[i].sum = histo[i].sum / histo[i].amo;    // Average P_v
    if (go == 1) printf("\n%.5f", histo[i].sum);
  };
  
  /* Compare the histigram with the expected function:
   *
   *   x * log(p_max/x)/arctan(sqrt(p_max^2/x^2 - 1))
   *
   * using the χ² test.
   */
  struct parameters params;
  params.histo = histo;
  params.n     = n;
  params.step  = step;

  gsl_function func;
  func.function = &chi2;
  func.params = &params;

  double min_p = p_max - 5;
  double max_p = p_max + 5;

  // Initialize minimization
  double x = p_max;
  int    max_iter = 100;
  double prec = 1e-7;
  int    status = GSL_CONTINUE; 
  const  gsl_min_fminimizer_type *T = gsl_min_fminimizer_brent;
  gsl_min_fminimizer *s = gsl_min_fminimizer_alloc(T);
  gsl_min_fminimizer_set(s, &func, x, min_p, max_p);

  // Minimization
  for (int iter = 0; status == GSL_CONTINUE && iter < max_iter; iter++) {  
    status = gsl_min_fminimizer_iterate(s);
    x      = gsl_min_fminimizer_x_minimum(s);
    min_p  = gsl_min_fminimizer_x_lower(s);
    max_p  = gsl_min_fminimizer_x_upper(s);
    status = gsl_min_test_interval(min_p, max_p, 0, prec);
  }

  double result = x;
  double res_chi = chi2(result, &params);
  if (go == 0) {
    printf("Results:\n");
    printf("χ² = %.3f\n", res_chi);
    printf("p_max = %.3f\n", result);
  }

  /* Compute the second derivative of χ² in its minimum for the result error.
   *   
   *   p_max = α
   *                    2
   *           (Eᵢ - Oᵢ) 
   *   χ² = Σi ──────────
   *               Eᵢ    
   *               
   *                ⎛      2⎞
   *                ⎜    Oᵢ ⎟
   *   ∂αχ² = Σi Eᵢ⋅⎜1 - ───⎟⋅∂αE
   *                ⎜      2⎟
   *                ⎝    Eᵢ ⎠
   *
   *              ⎛           2         ⎛      2⎞ ⎞
   *              ⎜         Oᵢ          ⎜    Oᵢ ⎟ ⎟
   *   ∂²αχ² = Σi ⎜(∂αE)²⋅2⋅─── + ∂²αE ⋅⎜1 - ───⎟ ⎟
   *              ⎜           2         ⎜      2⎟ ⎟
   *              ⎝         Eᵢ          ⎝    Eᵢ ⎠ ⎠
   */
  double expecto, A, B;
  double error = 0;
  for (size_t i = 0; i < n; i++) {
    x = (i + 0.5) * step;
    expecto = expected(x, result);
    A = 2 * pow(exp1d(x, result) * histo[i].sum / expecto, 2);
    B = exp2d(x, result) * (1 - pow((histo[i].sum / expecto), 2));
    error = error + A + B;
  };
  error = 1/error;
  if (go == 0)
    printf("ΔP_max = %.3f\n", error);

  // Check compatibility
  double t = fabs(result - p_max)/error;
  double alpha = 1 - erf(t/sqrt(2));
  if (go == 0) {
    printf("\nCompatibility:\n");
    printf("t = %.3f\n", t);
    printf("α = %.3f\n\n", alpha);
  }

  // Free memory
  gsl_min_fminimizer_free(s);
  gsl_rng_free(r);
  free(histo);

  return EXIT_SUCCESS;
}