#include <math.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>
#include <gsl/gsl_sf.h>
#include "fft.h"
#include "rl.h"

/* Program options */
struct options {
  int convolved;
  int deconvolved;
  int original;
  size_t num_events;
  size_t bins;
  double sigma;
  size_t rounds;
  const char* mode;
  double noise;
};


/* Intensity function parameters */
struct param {
  double a;  // aperture radius
  double k;  // wavenumber
  double e;  // wave amplitude
  double l;  // screen distance
};


/* Intensity of the EM field at a large distance
 * from a circular aperture diffraction of a plane
 * wave.
 *
 * The intensity is given as a function of θ,
 * the diffraction angle. All the other physical
 * parameters are controlled by the structure `param`.
 *
 * See Fraunhöfer diffraction on "Hect - Optics" for 
 * a derivation of the formula.
 */
double intensity(double theta, struct param p) {
  double x = p.k * p.a * sin(theta);
  double R = p.l / cos(theta);

  /* The intensity function has an eliminable
   * discontinuoty at the origin, which must
   * be handled separately.
   */
  if (fabs(x) < 1e-15) {
    return 0.5 * pow(p.e/R * M_PI * pow(p.a, 2), 2);
  }

  double y = 2*M_PI*pow(p.a, 2) * gsl_sf_bessel_J1(x) / x;
  return 0.5 * pow(p.e * y/R, 2);
}


/* `gaussian_for(hist, sigma)` generates a histogram
 * with the same bin width of `hist` but smaller,
 * with a gaussian PDF with μ at the central bin
 * and σ equal to `sigma`.
 */
gsl_histogram* gaussian_for(gsl_histogram *hist, double sigma) {
  /* Set the size to 6% of the histogram.
   * Always choose n even to correctly center the
   * gaussian in the middle.
   */
  size_t n = ((double) hist->n * 6) / 100;
  n = n % 2 ? n+1 : n;

  /* Calculate the (single) bin width assuming
   * the ranges are uniformely spaced
   */
  double min, max;
  gsl_histogram_get_range(hist, 0, &min, &max);
  double dx = max - min;
  gsl_histogram *res = gsl_histogram_alloc(n);
  gsl_histogram_set_ranges_uniform(res, 0, dx * n);

  /* The histogram will be such that
   * the maximum falls in the central bin.
   */
  long int offset = res->n/2;

  for (long int i = 0; i < (long int)res->n; i++)
    res->bin[i] = gsl_ran_gaussian_pdf(
        ((double)(i - offset) + 0.5) * dx, sigma * dx);

  return res;
}


/* Convolves two histograms while keeping the
 * original bin edges. This is a simple wrapper
 * function around `gsl_vector_convolve`.
 */
gsl_histogram* histogram_convolve(gsl_histogram *a, gsl_histogram *b) {
  gsl_histogram *res = gsl_histogram_calloc(a->n + b->n - 1);

  /* Set the same edges as `a`*/
  double max = gsl_histogram_max(a);
  double min = gsl_histogram_min(a);
  gsl_histogram_set_ranges_uniform(res, min, max);

  /* Create vector views for everybody */
  gsl_vector_view va  = gsl_vector_view_array(a->bin, a->n);
  gsl_vector_view vb  = gsl_vector_view_array(b->bin, b->n);
  gsl_vector_view vres = gsl_vector_view_array(res->bin, res->n);

  gsl_vector_convolve(&va.vector, &vb.vector, &vres.vector);

  return res;
}


int main(int argc, char **argv) {
  struct options opts;
  /* Set default options */
  opts.convolved   = 0;
  opts.deconvolved = 0;
  opts.original    = 1;
  opts.num_events  = 50000;
  opts.bins        = 150;
  opts.sigma       = 0.8;
  opts.rounds      = 3;
  opts.mode        = "fft";
  opts.noise       = 0;
    
  /* Process CLI arguments */
  for (int i = 1; i < argc; i++) {
         if (!strcmp(argv[i], "-c")) opts.convolved   = 1;
    else if (!strcmp(argv[i], "-d")) opts.deconvolved = 1;
    else if (!strcmp(argv[i], "-o")) opts.original    = 1;
    else if (!strcmp(argv[i], "-e")) opts.num_events  = atol(argv[++i]);
    else if (!strcmp(argv[i], "-b")) opts.bins        = atol(argv[++i]);
    else if (!strcmp(argv[i], "-s")) opts.sigma       = atof(argv[++i]);
    else if (!strcmp(argv[i], "-r")) opts.rounds      = atol(argv[++i]);
    else if (!strcmp(argv[i], "-m")) opts.mode        = argv[++i];
    else if (!strcmp(argv[i], "-n")) opts.noise       = atof(argv[++i]);
    else {
      fprintf(stderr, "Usage: %s -[cdoebsrmh]\n", argv[0]);
      fprintf(stderr, "\t-h\t\tShow this message.\n");
      fprintf(stderr, "\t-c\t\tPrint the convolved histogram to stdout.\n");
      fprintf(stderr, "\t-d\t\tAttempt and print the deconvolved histogram.\n");
      fprintf(stderr, "\t-o\t\tPrint the original histogram to stdout.\n");
      fprintf(stderr, "\t-e N\t\tThe number of events. (default: 50000)\n");
      fprintf(stderr, "\t-b N\t\tThe number of θ bins. (default: 150)\n");
      fprintf(stderr, "\t-s SIGMA\tThe sigma of gaussian kernel.(default: 0.8)\n");
      fprintf(stderr, "\t-r N\t\tThe number of RL deconvolution rounds. (default: 3)\n");
      fprintf(stderr, "\t-m MODE\t\tThe deconvolution mode: 'fft' or 'rl'. (default: 'fft')\n");
      fprintf(stderr, "\t-n MU\t\tThe mean (μ) of Poisson noise to add to the convolution. (default: 0)\n");
      return EXIT_FAILURE;
    }
  }

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

  gsl_histogram* hist = gsl_histogram_alloc(opts.bins);
  gsl_histogram_set_ranges_uniform(hist, 0, M_PI/2);

  double theta;
  struct param p = { 0.01, 0.0001, 1e4, 1 };

  /* Sample events following the intensity
   * of a circular aperture diffraction I(θ)
   * while producing a histogram.
   */
  fputs("# event sampling\n", stderr);
  fprintf(stderr, "max itensity: %.2f\n", intensity(0, p));
  fprintf(stderr, "1. generating %ld events...", opts.num_events);
  for (size_t i = 0; i < opts.num_events; i++){
    do {
      theta = acos(1 - gsl_rng_uniform(r));
    } while(intensity(0, p) * gsl_rng_uniform(r) > intensity(theta, p));
    gsl_histogram_increment(hist, theta);
  }
  fputs("done\n", stderr);

  /* Generate the gaussian kernel, convolve the
   * sample histogram with it and try to deconvolve it.
   */
  fputs("\n# convolution\n", stderr);
  fputs("1. generating gaussian kernel...", stderr);
  gsl_histogram *kernel = gaussian_for(hist, opts.sigma);
  fputs("done\n", stderr);

  fputs("2. convolving...", stderr);
  gsl_histogram *conv = histogram_convolve(hist, kernel);
  fputs("done\n", stderr);

  /* Add Poisson noise with μ=opts.noise to
   * the convolution.
   */
  if (opts.noise > 0) {
    fputs("2.1 adding poisson noise...", stderr);
    for (size_t i = 0; i < conv->n; i++)
      conv->bin[i] += gsl_ran_poisson(r, opts.noise);
    fputs("done\n", stderr);
  }

  fprintf(stderr, "3. %s deconvolution...", opts.mode);
  gsl_histogram *clean;
  if (!strcmp(opts.mode, "fft"))
    clean = fft_deconvolve(conv, kernel);
  else if (!strcmp(opts.mode, "rl"))
    clean = rl_deconvolve(conv, kernel, opts.rounds);
  else {
    fputs("\n\nerror: invalid mode. select either 'fft' or 'rl'\n", stderr);
    return EXIT_FAILURE;
  }
  fputs("done\n", stderr);

  /* Print the selected histogram*/
  fputs("\n# histogram \n", stderr);
  gsl_histogram *showing;
       if (opts.convolved)   showing = conv;
  else if (opts.deconvolved) showing = clean;
  else if (opts.original)    showing = hist;
  gsl_histogram_fprintf(stdout, showing, "%g", "%g");

  // free memory
  gsl_rng_free(r);
  gsl_histogram_free(hist);
  gsl_histogram_free(kernel);
  gsl_histogram_free(conv);
  gsl_histogram_free(clean);

  return EXIT_SUCCESS;
}