#include <math.h>
#include <string.h>
#include <gsl/gsl_rng.h>
#include <gsl/gsl_randist.h>

#include "common.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;
};


int show_help(char **argv) {
  fprintf(stderr, "Usage: %s -[hcdoebsrmn]\n", argv[0]);
  fprintf(stderr, "  -h\t\tShow this message.\n");
  fprintf(stderr, "  -c\t\tPrint the convolved histogram to stdout.\n");
  fprintf(stderr, "  -d\t\tAttempt and print the deconvolved histogram.\n");
  fprintf(stderr, "  -o\t\tPrint the original histogram to stdout.\n");
  fprintf(stderr, "  -e N\t\tThe number of events. (default: 50000)\n");
  fprintf(stderr, "  -b N\t\tThe number of θ bins. (default: 150)\n");
  fprintf(stderr, "  -s SIGMA\tThe sigma of gaussian kernel. "
                  "(default: 0.8)\n");
  fprintf(stderr, "  -r N\t\tThe number of RL deconvolution rounds."
                  "(default: 3)\n");
  fprintf(stderr, "  -m MODE\tThe deconvolution mode: 'fft' or 'rl'."
                  "(default: 'fft')\n");
  fprintf(stderr, "  -n SIGMA\tThe σ of the gaussian noise to add to "
                  "the convolution. (default: 0)\n");
  return EXIT_SUCCESS;
}


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 return show_help(argv);
  }

  /* 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_kernel(hist->n, opts.sigma);
  fputs("done\n", stderr);

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

  /* Add gaussian noise with σ=opts.noise to
   * the convolution.
   */
  if (opts.noise > 0) {
    fputs("2.1 adding noise...", stderr);
    for (size_t i = 0; i < conv->n; i++)
      conv->bin[i] += conv->bin[i] * gsl_ran_gaussian(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;
}