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

#include "common.h"
#include "fft.h"
#include "rl.h"
#include "dist.h"


/* Program options */
struct options {
  size_t num_events;
  size_t num_exps;
  size_t bins;
  double sigma;
  size_t rounds;
  const char* mode;
  double noise;
};

int show_help(char **argv) {
  fprintf(stderr, "Usage: %s -[hexbsrn]\n", argv[0]);
  fprintf(stderr, "  -h\t\tShow this message.\n");
  fprintf(stderr, "  -e N\t\tThe number of events. (default: 50000)\n");
  fprintf(stderr, "  -x N\t\tThe number of experiments to run. (default: 1000)\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, "  -n SIGMA\tThe σ of the gaussian noise to add to "
                  "the convolution. (default: 0)\n");
  return EXIT_SUCCESS;
}

/* Performs an experiment consisting in
 *
 *  1. Measuring the distribution I(θ) sampling from
 *     an RNG;
 *  2. Convolving the I(θ) sample with a kernel
 *     to simulate the instrumentation response;
 *  3. Applying a gaussian noise with σ=opts.noise
 *  4. Deconvolving the result with both the
 *     Richardson-Lucy (RL) algorithm and the Fourier
 *     transform (FFT) method;
 *  5. Computing the Earth Mover's Distance (EMD) between
 *     the deconvolution results and the original,
 *     uncorrupted sample.
 *
 *  The function returns a pair of the FFT and RL
 *  distances, in this order.
 */
double *experiment(
  struct options opts,
  gsl_rng *r,
  gsl_histogram *kernel) {

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

  /* Sample events following the intensity
   * of a circular aperture diffraction I(θ)
   * while producing a histogram.
   */
  gsl_histogram* hist = gsl_histogram_alloc(opts.bins);
  gsl_histogram_set_ranges_uniform(hist, 0, M_PI/2);

  for (size_t i = 0; i < opts.num_events; i++){
    double theta;
    do {
      // uniform sphere polar angle
      theta = acos(1 - gsl_rng_uniform(r));
    } while(intensity(0, p) * gsl_rng_uniform(r) > intensity(theta, p));
    gsl_histogram_increment(hist, theta);
  }

  /* Convolve the hisogram with the kernel */
  gsl_histogram *conv = histogram_convolve(hist, kernel);

  /* Add gaussian noise with σ=opts.noise */
  if (opts.noise > 0) {
    for (size_t i = 0; i < conv->n; i++)
      conv->bin[i] += conv->bin[i] * gsl_ran_gaussian(r, opts.noise);
  }

  /* Deconvolve the histogram with both methods: RL and FFT.
   * The FFT is used as reference for the distance.
   */
  gsl_histogram *fft_clean = fft_deconvolve(conv, kernel);
  gsl_histogram *rl_clean  = rl_deconvolve(conv, kernel, opts.rounds);

  /* Compute the earth mover's distances from the original
   * histogram and store add each one to the respective
   * distance histogram.
   */
  static double dist[3];
  dist[0] = emd_between(hist, fft_clean);
  dist[1] = emd_between(hist, rl_clean);
  dist[2] = emd_between(hist, conv);

  // free memory
  gsl_histogram_free(hist);
  gsl_histogram_free(conv);
  gsl_histogram_free(fft_clean);
  gsl_histogram_free(rl_clean);

  return dist;
}


int main(int argc, char **argv) {
  struct options opts;
  /* Set default options */
  opts.num_events = 50000;
  opts.num_exps   = 1000;
  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], "-e")) opts.num_events = atol(argv[++i]);
    else if (!strcmp(argv[i], "-x")) opts.num_exps   = 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], "-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);

  /* Generate the gaussian kernel */
  gsl_histogram *kernel = gaussian_kernel(opts.bins, opts.sigma);

  /* Performs experiments and record the deconvolution
   * EMD distances to the original sample.
   */
  double* fft_dist  = calloc(opts.num_exps, sizeof(double));
  double* rl_dist   = calloc(opts.num_exps, sizeof(double));
  double* conv_dist = calloc(opts.num_exps, sizeof(double));

  for (size_t i = 0; i < opts.num_exps; i++) {
    double *dist = experiment(opts, r, kernel);
    fft_dist[i]  = dist[0];
    rl_dist[i]   = dist[1];
    conv_dist[i] = dist[2];
  }

  /* Compute some statistics of the EMD */
  // 1 is the array stride.
  double fft_mean  = gsl_stats_mean(fft_dist, 1, opts.num_exps);
  double fft_stdev = gsl_stats_sd(fft_dist, 1, opts.num_exps);
  double fft_skew  = gsl_stats_skew_m_sd(fft_dist, 1, opts.num_exps, fft_mean, fft_stdev);
  double fft_min, fft_max; gsl_stats_minmax(&fft_min, &fft_max, fft_dist, 1, opts.num_exps);

  double rl_mean  = gsl_stats_mean(rl_dist, 1, opts.num_exps);
  double rl_stdev = gsl_stats_sd(rl_dist, 1, opts.num_exps);
  double rl_skew  = gsl_stats_skew_m_sd(rl_dist, 1, opts.num_exps, rl_mean, rl_stdev);
  double rl_min, rl_max; gsl_stats_minmax(&rl_min, &rl_max, rl_dist, 1, opts.num_exps);

  double conv_mean  = gsl_stats_mean(conv_dist, 1, opts.num_exps);
  double conv_stdev = gsl_stats_sd(conv_dist, 1, opts.num_exps);
  double conv_skew  = gsl_stats_skew_m_sd(conv_dist, 1, opts.num_exps, conv_mean, conv_stdev);
  double conv_min, conv_max; gsl_stats_minmax(&conv_min, &conv_max, conv_dist, 1, opts.num_exps);

  /* Create EDM distance histograms.
   * Since the distance depends wildly on the noise we can't
   * set a fixed range and therefore use the above values.
   */
  gsl_histogram *fft_hist  = gsl_histogram_alloc(sqrt(opts.num_exps));
  gsl_histogram *rl_hist   = gsl_histogram_alloc(sqrt(opts.num_exps));
  gsl_histogram *conv_hist = gsl_histogram_alloc(sqrt(opts.num_exps));
  gsl_histogram_set_ranges_uniform(  fft_hist, fft_min,  fft_max);
  gsl_histogram_set_ranges_uniform(  rl_hist,   rl_min,   rl_max);
  gsl_histogram_set_ranges_uniform(conv_hist, conv_min, conv_max);

  for (size_t i = 0; i < opts.num_exps; i++) {
    gsl_histogram_increment( fft_hist,  fft_dist[i]);
    gsl_histogram_increment(  rl_hist,   rl_dist[i]);
    gsl_histogram_increment(conv_hist, conv_dist[i]);
  }

  // print results to stderr
  fprintf(stderr, "# EMD distance\n\n");
  fprintf(stderr, "## FFT deconvolution\n");
  fprintf(stderr, "- mean:  %.2e\n"
                  "- stdev: %.1e\n"
                  "- skew:  %.2f\n", fft_mean, fft_stdev, fft_skew);
  fprintf(stderr, "\n## RL deconvolution\n");
  fprintf(stderr, "- mean:  %.2e\n"
                  "- stdev: %.1e\n"
                  "- skew:  %.2f\n", rl_mean, rl_stdev, rl_skew);
  fprintf(stderr, "\n## convolution\n");
  fprintf(stderr, "- mean:  %.2e\n"
                  "- stdev: %.1e\n"
                  "- skew:  %.2f\n", conv_mean, conv_stdev, conv_skew);

  // print histograms to stdout
  fprintf(stderr, "\n# EMD histogram\n");
  fprintf(stdout, "%d\n", (int)sqrt(opts.num_exps));
  gsl_histogram_fprintf(stdout,  fft_hist, "%.10g", "%.10g");
  gsl_histogram_fprintf(stdout,   rl_hist, "%.10g", "%.10g");
  gsl_histogram_fprintf(stdout, conv_hist, "%.10g", "%.10g");

  // free memory
  gsl_rng_free(r);
  gsl_histogram_free(kernel);
  gsl_histogram_free( fft_hist);
  gsl_histogram_free(  rl_hist);
  gsl_histogram_free(conv_hist);
  free( fft_dist);
  free(  rl_dist);
  free(conv_dist);

  return EXIT_SUCCESS;
}
