/*
 *  targets.cpp
 *  MetaQuant
 *
 *  Kevin McLoughlin
 *  Based on code from eXpress, created by Adam Roberts in 2011.
 *  Copyright 2014 Kevin McLoughlin, Adam Roberts. All rights reserved.
 */

#include "main.h"
#include "targets.h"
#include "lengthdistribution.h"
#include "fragments.h"
#include "biascorrection.h"
#include "mismatchmodel.h"
#include "mapparser.h"
#include "sequence.h"
#include "library.h"
#include <iostream>
#include <fstream>
#include <cassert>
#include <stdio.h>
#include <limits>
#include <float.h>
#include <boost/thread.hpp>
#include <boost/chrono.hpp>
#include <boost/atomic.hpp>
// =-= For memory debugging only
#include "memdebug.h"

using namespace std;

Target::Target(TargID id, const std::string& name, const std::string& description,
               const std::string& seq,
               bool prob_seq, double alpha, double beta, const Librarian* libs,
               const BiasBoss* known_bias_boss, const LengthDistribution* known_fld)
   : _libs(libs),
     _id(id),
     _name(name),
     _description(description),
     _seq_f(seq, 0, prob_seq),
     _seq_r(_seq_f),
     _alpha_prior(alpha),
     _beta_prior(beta),
     _ret_params(&_curr_params),
     _uniq_counts(0),
     _tot_counts(0),
     _avg_bias(0),
     _avg_bias_buffer(0),
     _solvable(false) {
  if ((_libs->curr_lib()).bias_table) {
    _start_bias.reset(new std::map<const size_t, float>);
    _start_bias_buffer.reset(new std::map<const size_t, float>);
    _end_bias.reset(new std::map<const size_t, float>);
    _end_bias_buffer.reset(new std::map<const size_t, float>);
  }
  update_target_bias_buffer(known_bias_boss, known_fld);
  swap_bias_parameters();
}

void Target::round_reset() {
  // =-= TODO: Note that this only makes sense for batch VBayes. For online VB _ret_params
  // =-= should always point to _curr_params. Reimplement this when we implement batch
  // =-= VBayes. Will need to implement a copy constructor for RoundParams, to handle
  // =-= the atomic spin_lock variable.
  //_last_params = _curr_params;
  //_curr_params = RoundParams();
  //_ret_params = &_last_params;
}

double Target::stick_breaking_frac() const {
  return digamma(_ret_params->alpha) - digamma(_ret_params->alpha + _ret_params->beta);
}

double Target::stick_breaking_complement() const {
  return digamma(_ret_params->beta) - digamma(_ret_params->alpha + _ret_params->beta);
}

double Target::stick_breaking_frac_var() const {
  int ierr = 0;
  return trigamma(_ret_params->alpha, &ierr) - trigamma(_ret_params->alpha + _ret_params->beta, &ierr);
}

double Target::stick_breaking_complement_var() const {
  int ierr = 0;
  return trigamma(_ret_params->beta, &ierr) - trigamma(_ret_params->alpha + _ret_params->beta, &ierr);
}

void Target::update_alpha_beta(double frag_posterior, double posterior_sum,
       double learning_rate, size_t total_frags) {
  while (_curr_params.spin_lock.exchange(1, boost::memory_order_acquire) == 1) {
    // spin until free
  }
  _curr_params.alpha = _curr_params.alpha * (1. - learning_rate)
                       + learning_rate*(_alpha_prior + total_frags*sexp(frag_posterior));
  _curr_params.beta = _curr_params.beta * (1. - learning_rate)
                + learning_rate*(_beta_prior + total_frags*sexp(posterior_sum));
  _curr_params.spin_lock.store(0, boost::memory_order_release);
}

double Target::rho() const {
  return _relative_abundance;
}

void Target::add_hit(double p, double m) {
  _curr_params.mass = log_add(_curr_params.mass, p+m);
  if (p != LOG_1) {
    if (p != LOG_0) {
      _curr_params.ambig_mass = log_add(_curr_params.ambig_mass, p+m);
      _curr_params.tot_ambig_mass = log_add(_curr_params.tot_ambig_mass, m);
    }
  }
  (_libs->curr_lib()).targ_table->update_total_fpb(m - _cached_eff_len);
}

double Target::est_effective_length(const LengthDistribution* fld,
                                    bool with_bias) const {
  if (!fld) {
    fld = (_libs->curr_lib()).fld.get();
  }

  double eff_len = LOG_0;

  double log_length = log((double)length());
  if (log_length < fld->mean()) {
    eff_len = log_length;
  } else {
    for(size_t l = fld->min_val(); l <= min(length(), fld->max_val()); l++) {
      eff_len = log_add(eff_len, fld->pmf(l)+log((double)length()-l+1));
    }
  }

  if (with_bias) {
    eff_len += _avg_bias;
  }

  return eff_len;
}

double Target::start_bias(const BiasBoss* bias_table, const size_t start_pos) {
  double bias;
  try {
    bias = _start_bias->at(start_pos);
  } catch (out_of_range err) {
    bias = bias_table->get_target_start_bias(*this, start_pos);
    (*_start_bias)[start_pos] = bias;
  }
  return bias;
}

double Target::end_bias(const BiasBoss* bias_table, const size_t end_pos) {
  double bias;
  try {
    bias = _end_bias->at(end_pos);
  } catch (out_of_range err) {
    bias = bias_table->get_target_end_bias(*this, end_pos);
    (*_end_bias)[end_pos] = bias;
  }
  return bias;
}

double Target::cached_effective_length(bool with_bias) const {
  if (with_bias) {
    return _cached_eff_len + _avg_bias;
  }
  return _cached_eff_len;
}

void Target::update_target_bias_buffer(const BiasBoss* bias_table,
                                       const LengthDistribution* fld) {
  long update_time = 0;
  long efflen_time = 0;

  if (bias_table) {
    boost::chrono::high_resolution_clock::time_point update_start =
       boost::chrono::high_resolution_clock::now();

    bias_table->update_target_bias(*_start_bias_buffer, *_end_bias_buffer, *this);

    boost::chrono::high_resolution_clock::time_point update_end =
       boost::chrono::high_resolution_clock::now();
    update_time = (long)((update_end-update_start).count());

    // =-= REMOVE: avg_bias never gets used in calculating effective length
    //_avg_bias_buffer = bias_table->get_average_target_bias(*_start_bias_buffer,
    //   *_end_bias_buffer, *this);

  }
  //assert(!isnan(_avg_bias_buffer) && !isinf(_avg_bias_buffer));
  boost::chrono::high_resolution_clock::time_point efflen_start =
       boost::chrono::high_resolution_clock::now();
  _cached_eff_len_buffer = est_effective_length(fld, false);
  boost::chrono::high_resolution_clock::time_point update_end =
       boost::chrono::high_resolution_clock::now();

  efflen_time = (long)((update_end-efflen_start).count());

  //logger.info("Updated bias for target %s with hits at %d positions, rho = %f, update_time=%ld efflen_time=%ld",
  //    _name.c_str(), (*_start_bias_buffer).size(), sexp(_relative_abundance),
  //    update_time, efflen_time);
}

void Target::swap_bias_parameters() {
  _cached_eff_len = _cached_eff_len_buffer;
  _avg_bias = _avg_bias_buffer;
  _start_bias.swap(_start_bias_buffer);
  _end_bias.swap(_end_bias_buffer);
}

TargetTable::TargetTable(string targ_db_name,
                         bool prob_seqs, bool known_aux_params,
                         const Librarian* libs)
    :  _libs(libs) {
  string info_msg = "Loading target sequences";
  const Library& lib = _libs->curr_lib();
  const TargIndex& targ_index = lib.map_parser->targ_index();
  const TargIndex& targ_lengths = lib.map_parser->targ_lengths();
  const TargIDSet& hit_target_ids = lib.map_parser->hit_target_ids();
  if (lib.bias_table && !known_aux_params) {
    info_msg += " and measuring bias background";
  }
  info_msg += "...";
  logger.info(info_msg.c_str());

  boost::unordered_set<string> target_names;

  // Copy target index entries from the parser for targets with hits only into the
  // TargID set and vector objects
  size_t index = 0;
  for(TargIndex::const_iterator it = targ_index.begin(); it != targ_index.end(); ++it) {
    TargID targ_id = it->second;
    if (hit_target_ids.find(targ_id) != hit_target_ids.end()) {
      _hit_target_vec.push_back(targ_id);
      _targ_index_map[targ_id] = index++;
    }
  }

  // Find the length of the longest target sequence we might have to deal with
  size_t max_targ_length = 0;
  for(TargIndex::const_iterator it = targ_lengths.begin(); it != targ_lengths.end(); ++it) {
    size_t len = it->second;
    if (len > max_targ_length) {
      max_targ_length = len;
    }
  }

  // Instead of reading the whole FASTA file, we preprocess the reference FASTA
  // using code from refseq_db.py to generate a catalog, and then use fseek and fread
  // to only grab the sequences that have hits.

  string targ_catalog_file = targ_db_name + "_catalog.txt";
  string targ_data_file = targ_db_name + "_seq.dat";

  _target_db_fp = fopen(targ_data_file.c_str(), "rb");

  ifstream infile (targ_catalog_file.c_str());
  const size_t BUFF_SIZE = 2000;
  char line_buff[BUFF_SIZE];
  string name = "";
  long int offset = 0;
  size_t targ_len = 0;
  string targ_desc = "";
  size_t ntargs = 0;

  if (infile.is_open()) {
    char* seq_buff = (char*) malloc(max_targ_length*sizeof(char));
    while (infile.good()) {
      infile.getline(line_buff, BUFF_SIZE, '\n');
      // Split line into tab separated fields
      char *p = strtok(line_buff, "\t");
      int nf = 0;
      while (p && nf < 4) {
        switch(nf++) {
          case 0: {
            name = p;
            break;
          }
          case 1: {
            offset = (long int)atol(p);
            break;
          }
          case 2: {
            targ_len = (size_t)atoi(p);
            break;
          }
          case 3: {
            targ_desc = p;
            break;
          }
        }
        p = strtok(NULL, "\t");
      }
      long targ_size = add_targ(name, offset, targ_len, targ_desc, prob_seqs, known_aux_params,
         alpha_prior, beta_prior, targ_index, targ_lengths, hit_target_ids,
         seq_buff);
      if (targ_size > 0) {
        ntargs++;
        //total_targ_data += targ_size;
        if (ntargs % 1000 == 1) {
          logger.info("Loaded target %d (%s) %s", ntargs, name.c_str(),
              targ_desc.c_str());
        }
      }
    }
    infile.close();
    free(seq_buff);

    if (lib.bias_table && !known_aux_params) {
      lib.bias_table->normalize_expectations();
    }
  } else {
    logger.severe("Unable to open target DB catalog file '%s'.",
                  targ_catalog_file.c_str());
  }

  if (size() == 0) {
    logger.severe("No targets found in target DB catalog file '%s'.",
                  targ_catalog_file.c_str());
  }

  for (size_t i=0; i < size(); i++) {
    TargID targ_id = _hit_target_vec[i];
    if (_targ_map.find(targ_id) == _targ_map.end()) {
      logger.severe("Target %d not found in target DB catalog file "
                    "'%s'.", targ_id, targ_catalog_file.c_str());
    }
  }
  fclose(_target_db_fp);
  logger.info("Initialized %d targets.", ntargs);

}

TargetTable::~TargetTable() {
  foreach(TargPair targ_pair, _targ_map) {
    Target* targ = targ_pair.second;
    delete targ;
  }
}

long TargetTable::add_targ(const string& name, long int offset, size_t targ_len,
                           const string& description, bool prob_seq,
                           bool known_aux_params, double alpha, double beta,
                           const TargIndex& targ_index,
                           const TargIndex& targ_lengths,
                           const TargIDSet& hit_target_set,
                           char* seq_buffer) {

  long targ_size = 0;
  TargIndex::const_iterator it = targ_index.find(name);
  if (it == targ_index.end()) {
    logger.warn("Target '%s' exists in target DB but not alignment "
                   "(SAM/BAM) file.", name.c_str());
    return targ_size;
  }
  TargID targ_id = it->second;

  if (hit_target_set.find(targ_id) == hit_target_set.end()) {
    // Do nothing; we only add table entries for targets with hits
    return targ_size;
  }

  if (targ_lengths.find(name)->second != targ_len) {
    logger.severe("Target '%s' differs in length between target DB and "
                  "alignment (SAM/BAM) files (%d  vs. %d).", name.c_str(),
                  targ_len, targ_lengths.find(name)->second);
  }

  // Load the target sequence from the target DB data file
  int ferr = fseek(_target_db_fp, offset, SEEK_SET);
  if (ferr != 0) {
    logger.severe("Error trying to seek in target DB for target %s at offset %d",
       name.c_str(), offset);
  }
  size_t nread = fread((void *)seq_buffer, sizeof(char), targ_len, _target_db_fp);
  if (nread != targ_len) {
    logger.severe("Attempt to read sequence of length %d for target %s returned only %d chars",
       targ_len, name.c_str(), nread);
  }
  string seq(seq_buffer, targ_len);

  const Library& lib = _libs->curr_lib();
  const BiasBoss* known_bias_boss = (known_aux_params) ? lib.bias_table.get()
                                                       : NULL;
  const LengthDistribution* known_fld = (known_aux_params) ? lib.fld.get()
                                                           : NULL;

  Target* targ = new Target(targ_id, name, description, seq, prob_seq, alpha, beta, _libs,
                            known_bias_boss, known_fld);

  targ_size = sizeof(Target) + sizeof(SequenceFwd) + sizeof(SequenceRev) +
     2*sizeof(RoundParams) + name.size() + description.size() + targ_len;
  if (lib.bias_table && !known_aux_params) {
    (lib.bias_table)->update_expectations(*targ);
  }
  _targ_map[targ_id] = targ;
  targ->bundle(_bundle_table.create_bundle(targ));

  return targ_size;
}

Target* TargetTable::get_targ(TargID id) {
    try {
      Target* t = _targ_map.at(id);
      return t;
    } catch (out_of_range err) {
      return NULL;
    }
}

Bundle* TargetTable::merge_bundles(Bundle* b1, Bundle* b2) {
  if (b1 != b2) {
    return _bundle_table.merge(b1, b2);
  }
  return b1;
}

void TargetTable::round_reset() {
  foreach(TargPair targ_pair, _targ_map) {
    Target* targ = targ_pair.second;
    targ->round_reset();
  }
}

void project_to_polytope(vector<Target*> bundle_targ,
                         vector<double>& targ_counts, double bundle_counts) {
  vector<bool> polytope_bound(bundle_targ.size(), false);
  while (true) {
    double unbound_counts = 0;
    double bound_counts = 0;
    for (size_t i = 0; i < bundle_targ.size(); ++i) {
      Target& targ = *bundle_targ[i];

      if (targ_counts[i] > targ.tot_counts()) {
        targ_counts[i] = targ.tot_counts();
        polytope_bound[i] = true;
      } else if (targ_counts[i] < targ.uniq_counts()) {
        targ_counts[i] = targ.uniq_counts();
        polytope_bound[i] = true;
      }

      if (polytope_bound[i]) {
        bound_counts += targ_counts[i];
      } else {
        unbound_counts += targ_counts[i];
      }
    }

    if (approx_eq(unbound_counts + bound_counts, bundle_counts)) {
      return;
    }

	  if (unbound_counts == 0) {
      polytope_bound = vector<bool>(bundle_targ.size(), false);
	    unbound_counts = bound_counts;
	    bound_counts = 0;
	  }

    double normalizer = (bundle_counts - bound_counts)/unbound_counts;
    for (size_t i = 0; i < bundle_targ.size(); ++i) {
      if (!polytope_bound[i]) {
        targ_counts[i] *= normalizer;
      }
    }
  }
}

void TargetTable::estimate_abundances() {
  double sum_rho = 0.0;
  for (size_t i = 0; i < size(); i++) {
    TargID targ_i = _hit_target_vec[i];
    Target* ti = get_targ(targ_i);
    double log_theta = ti->stick_breaking_frac();
    double var_log_theta = ti->stick_breaking_frac_var();
    for (size_t j = 0; j < i; j++) {
      TargID targ_j = _hit_target_vec[j];
      Target* tj = get_targ(targ_j);
      log_theta += tj->stick_breaking_complement();
      var_log_theta += tj->stick_breaking_complement_var();
    }
    ti->_sampling_prob = log_theta;
    // Estimate unnormalized log relative abundances and the standard error of the estimate
    ti->_relative_abundance = ti->_sampling_prob - ti->_cached_eff_len;
    // =-= Debugging code
    if (isnan(ti->_relative_abundance)) {
      logger.info("Target %s rho is NaN; sampling_prob = %f, cached_eff_len = %f, stick breaking frac = %f", ti->name().c_str(), ti->_sampling_prob, ti->_cached_eff_len, ti->stick_breaking_frac());
    }
    // =-= End debugging code

    sum_rho += sexp(ti->_relative_abundance);
    ti->_relative_abundance_sd = sqrt(var_log_theta);
  }
  // Normalize the relative abundances and confidence limits
  for (size_t i = 0; i < size(); i++) {
    TargID targ_i = _hit_target_vec[i];
    Target* ti = get_targ(targ_i);
    ti->_relative_abundance -= log(sum_rho);
    double log_ci = 1.959964 * ti->_relative_abundance_sd;
    ti->_relative_abundance_lo_cl = ti->_relative_abundance - log_ci;
    ti->_relative_abundance_hi_cl = ti->_relative_abundance + log_ci;
  }
}

void TargetTable::output_results(string output_dir, size_t tot_counts) {
  FILE * expr_file = fopen((output_dir + "/metaquant_results.txt").c_str(), "w");
  fprintf(expr_file, "bundle_id\ttarget_name\tlength\teff_length\t"
                     "rel_abundance\trel_abundance_conf_low\trel_abundance_conf_high\t"
                     "tot_counts\tuniq_counts\t"
                     "sampling_prob\tassigned_fragments\t"
                     "alpha\tbeta\tfpkm\tsolvable\tdescription\n");

  const double l_bil = log(1000000000.);

  size_t bundle_id = 0;

  bundle_id = 0;
  foreach (Bundle* bundle, _bundle_table.bundles()) {
    ++bundle_id;

    const vector<Target*>& bundle_targ = *(bundle->targets());

    for (size_t i = 0; i < bundle_targ.size(); ++i) {

      Target& targ = *bundle_targ[i];

      double fpkm = sexp(l_bil + targ._sampling_prob - targ._cached_eff_len);
      double eff_len = sexp(targ._cached_eff_len);
      double rho = sexp(targ._relative_abundance);
      double rho_lo_cl = sexp(targ._relative_abundance_lo_cl);
      double rho_hi_cl = sexp(targ._relative_abundance_hi_cl);
      double sample_prob = sexp(targ._sampling_prob);
      double assigned_frags = sample_prob * tot_counts;

      fprintf(expr_file, "" SIZE_T_FMT "\t%s\t" SIZE_T_FMT "\t%f\t"
              "%f\t%f\t%f\t"
              SIZE_T_FMT "\t" SIZE_T_FMT "\t"
              "%f\t%.1f\t"
              "%f\t%f\t%f\t"
              "%c\t%s\n",
              bundle_id, targ.name().c_str(), targ.length(), eff_len,
              rho, rho_lo_cl, rho_hi_cl,
              targ.tot_counts(), targ.uniq_counts(),
              sample_prob, assigned_frags,
              targ.alpha(), targ.beta(), fpkm,
              (targ.solvable())?'T':'F', targ.description().c_str());
    }
  }
  fclose(expr_file);
}

double TargetTable::total_fpb() const {
  boost::unique_lock<boost::mutex>(_fpb_mut);
  return _total_fpb;
}

void TargetTable::update_total_fpb(double incr_amt) {
  boost::unique_lock<boost::mutex>(_fpb_mut);
  _total_fpb = log_add(_total_fpb, incr_amt);
}

void TargetTable::asynch_bias_update(boost::mutex* mutex, int sleep_time) {
  BiasBoss* bg_table = NULL;
  boost::scoped_ptr<BiasBoss> bias_table;
  boost::scoped_ptr<LengthDistribution> fld;

  bool burned_out_before = false;

  const Library& lib = _libs->curr_lib();

  while(running) {
    if (bg_table) {
      bg_table->normalize_expectations();
    }
    {
      boost::unique_lock<boost::mutex> lock(*mutex);
      logger.info("Synchronizing auxiliary parameter tables...");
      if(!fld) {
        fld.reset(new LengthDistribution(*(lib.fld)));
      } else {
        *fld = *(lib.fld);
      }
      if (lib.bias_table) {
        BiasBoss& lib_bias_table = *(lib.bias_table);
        if (!bias_table) {
          bias_table.reset(new BiasBoss(lib_bias_table));
        } else {
          lib_bias_table.copy_expectations(*bg_table);
          bg_table->copy_observations(lib_bias_table);
          bias_table.reset(bg_table);
        }
        bg_table = new BiasBoss(lib_bias_table.order(), 0);
      }
      // Update target abundance estimates. This is safe to do now since
      // process_fragment() is blocked by mutex for this code block.
      estimate_abundances();

      logger.info("Synchronized auxiliary parameter tables.");
    }

    if (!edit_detect && burned_out && burned_out_before) {
      break;
    }

    burned_out_before = burned_out;

    vector<double> fl_cdf = fld->cmf();

    // Buffer results of long computations.
    logger.info("Updating target bias buffers.");
    foreach(TargPair targ_pair, _targ_map) {
      Target* targ = targ_pair.second;
      // =-= Don't think we need to lock targets while calculating biases
      //targ->lock();
      targ->update_target_bias_buffer(bias_table.get(), fld.get());
      if (bg_table) {
        bg_table->update_expectations(*targ, targ->rho(), fl_cdf);
      }
      //targ->unlock();
    }
    logger.info("Updated target bias buffers.");
    {
      boost::unique_lock<boost::mutex> lock(*mutex);
      // Do quick atomic swap
      foreach(TargPair targ_pair, _targ_map) {
        Target* targ = targ_pair.second;
        targ->lock();
      }
      foreach(TargPair targ_pair, _targ_map) {
        Target* targ = targ_pair.second;
        targ->swap_bias_parameters();
        targ->unlock();
      }
    }
    logger.info("Swapped target bias parameters.");

    // Go to sleep for a while
    for (int remaining_sleep = sleep_time; remaining_sleep > 0; remaining_sleep -= 5) {
      if (!running) {
        break;
      }
      boost::this_thread::sleep_for(boost::chrono::seconds(5));
    }
  }

  if (bg_table) {
    delete bg_table;
  }
}


/******************************************************
Time critical functions for the Deep Learning-based
reconstruction (DNN_reco) are written in C++ and
wrapped with boost python.
******************************************************/
#include <iostream>
#include <string>
#include <math.h>

#include "icetray/I3Frame.h"
#include "icetray/OMKey.h"
#include "dataclasses/physics/I3RecoPulse.h"
#include "dataclasses/physics/I3Particle.h"
#include "dataclasses/physics/I3MCTree.h"
#include "dataclasses/geometry/I3Geometry.h"
#include "dataclasses/geometry/I3OMGeo.h"
#include "dataclasses/I3Map.h"
#include "dataclasses/I3TimeWindow.h"
#include "phys-services/I3Calculator.h"
#include <boost/version.hpp>
#include <boost/python.hpp>

#include "utils.cpp"

#if BOOST_VERSION < 106500
    #include "reduced_summary_statistics__old_boost_version.cpp"
#else
    #include "reduced_summary_statistics.cpp"
#endif

/*
Depending on the boost version, we need to use numpy differently.
Prior to Boost 1.63 (BOOST_VERSION == 106300) numpy was not directly
included in boost/python

See answers and discussion provided here:
https://stackoverflow.com/questions/10701514/how-to-return-numpy-array-from-boostpython
*/
#if BOOST_VERSION < 106500
    #include "numpy/npy_3kcompat.h"
    typedef typename boost::python::numeric::array pyndarray;
    namespace arrayFunc = boost::python::numeric;
#else
    #include <boost/python/numpy.hpp>
    typedef typename boost::python::numpy::ndarray pyndarray;
    namespace arrayFunc = boost::python::numpy;
    namespace bn = boost::python::numpy;
#endif

namespace bp = boost::python;

/******************************************************
 Helper-Functions to return numpy arrays
******************************************************/

// https://stackoverflow.com/questions/10701514/how-to-return-numpy-array-from-boostpython
#if BOOST_VERSION < 106500

    template <typename T>
    struct select_npy_type
    {};

    template <>
    struct select_npy_type<double>
    {
        const static NPY_TYPES type = NPY_DOUBLE;
    };

    template <>
    struct select_npy_type<float>
    {
        const static NPY_TYPES type = NPY_FLOAT;
    };

    template <>
    struct select_npy_type<int>
    {
        const static NPY_TYPES type = NPY_INT;
    };

    template <typename T>
    boost::python::object stdVecToNumpyArray( std::vector<T> const& vec )
    {
          npy_intp size = vec.size();

         /* const_cast is rather horrible but we need a writable pointer
            in C++11, vec.data() will do the trick
            but you will still need to const_cast
          */

          T * data = size ? const_cast<T *>(&vec[0])
            : static_cast<T *>(NULL);

        // create a PyObject * from pointer and data
          PyObject * pyObj = PyArray_SimpleNewFromData( 1, &size, select_npy_type<T>::type, data );
          boost::python::handle<> handle( pyObj );
          pyndarray arr( handle );

        /* The problem of returning arr is twofold: firstly the user can modify
          the data which will betray the const-correctness
          Secondly the lifetime of the data is managed by the C++ API and not the
          lifetime of the numpy array whatsoever. But we have a simple solution..
         */

           return arr.copy(); // copy the object. numpy owns the copy now.
      }
    // todo: change this to templated function and remove duplicate

    template <std::size_t N, typename T>
    boost::python::object stdArrayToNumpyArray( std::array<T, N> const& vec )
    {
          npy_intp size = vec.size();

         /* const_cast is rather horrible but we need a writable pointer
            in C++11, vec.data() will do the trick
            but you will still need to const_cast
          */

          T * data = size ? const_cast<T *>(&vec[0])
            : static_cast<T *>(NULL);

        // create a PyObject * from pointer and data
          PyObject * pyObj = PyArray_SimpleNewFromData( 1, &size, select_npy_type<T>::type, data );
          boost::python::handle<> handle( pyObj );
          pyndarray arr( handle );

        /* The problem of returning arr is twofold: firstly the user can modify
          the data which will betray the const-correctness
          Secondly the lifetime of the data is managed by the C++ API and not the
          lifetime of the numpy array whatsoever. But we have a simple solution..
         */

           return arr.copy(); // copy the object. numpy owns the copy now.
      }
#else
    template <typename T>
    boost::python::object stdVecToNumpyArray( std::vector<T> const& vec )
    {
        Py_intptr_t shape[1] = { vec.size() };
        bn::ndarray result = bn::zeros(1, shape, bn::dtype::get_builtin<double>());
        std::copy(vec.begin(), vec.end(),
                  reinterpret_cast<double*>(result.get_data()));
        return result;
    }
    // todo: change this to templated function and remove duplicate
    template <std::size_t N, typename T>
    boost::python::object stdArrayToNumpyArray( std::array<T, N> const& vec )
    {
        Py_intptr_t shape[1] = { vec.size() };
        bn::ndarray result = bn::zeros(1, shape, bn::dtype::get_builtin<double>());
        std::copy(vec.begin(), vec.end(),
                  reinterpret_cast<double*>(result.get_data()));
        return result;
    }
#endif

/******************************************************
 Functions with pybinding for python-based usage
******************************************************/

template <typename T>
inline boost::python::dict get_cascade_classification_data(
                                  boost::python::object& frame_obj,
                                  const boost::python::object& pulse_map_obj,
                                  const boost::python::object& cascade_key_obj,
                                  boost::python::object& time_quantiles_obj,
                                  const bool add_dom_info
                                ) {

    // extract objects
    I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);
    const std::string cascade_key =
        boost::python::extract<std::string>(cascade_key_obj);

    const unsigned int n_bins =  len(time_quantiles_obj) - 1;

    std::vector<double> time_quantiles_abs;
    std::vector<double> time_quantiles_rel;
    for (int i = 0; i < n_bins + 1; ++i){
        time_quantiles_rel.push_back(
                boost::python::extract<double>(time_quantiles_obj[i][0]));
        time_quantiles_abs.push_back(
                boost::python::extract<double>(time_quantiles_obj[i][1]));
    }

    // Get geometry map
    const I3OMGeoMap& omgeo = frame.Get<I3Geometry>("I3Geometry").omgeo;
    // Get pulse map
    I3RecoPulseSeriesMap& pulse_map = boost::python::extract<I3RecoPulseSeriesMap&>(pulse_map_obj);

    // get parameters
    const I3MapStringDouble& params = frame.Get<I3MapStringDouble>(cascade_key);

    double x = params.at("VertexX");
    double y = params.at("VertexY");
    double z = params.at("VertexZ");
    double t = params.at("VertexTime");
    double unc_x = params.at("VertexX_unc");
    double unc_y = params.at("VertexY_unc");
    double unc_z = params.at("VertexZ_unc");
    double unc_t = params.at("VertexTime_unc");
    double c = 0.299792458; // m/ns from dataclasses.I3Constants.c
    double c_ice = 0.221030462863; // m/ns from dataclasses.I3Constants.c_ice
    double sqrt2 = std::sqrt(2);

    double unc_pos = std::sqrt(unc_x*unc_x + unc_y*unc_y + unc_z*unc_z);
    double unc_time = std::sqrt(std::pow(unc_pos / c, 2) + unc_t*unc_t);

    boost::python::dict data_dict;

    for (auto const& dom_pulses : pulse_map){

        // check if pulses are present
        if (dom_pulses.second.size() == 0){
            continue;
        }

        const I3Position& pos = omgeo.at(dom_pulses.first).position;

        // distance to DOM
        double distance = std::sqrt(std::pow(x - pos.GetX(), 2) +
                                    std::pow(y - pos.GetY(), 2) +
                                    std::pow(z - pos.GetZ(), 2));

        double delta_t = distance / c;

        std::vector<double> time_edges;

        double max_edge = 0;
        for (unsigned int i=0; i < n_bins + 1; i++){
            max_edge =
                t + time_quantiles_rel[i] * delta_t + time_quantiles_abs[i];
            time_edges.push_back(max_edge);
        }

        std::vector<double> bin_values_list;
        for (unsigned int i=0; i < n_bins; i++){
            bin_values_list.push_back(0);
        }

        // loop through pulses
        for (auto const& pulse : dom_pulses.second){
            //dom_charges.push_back(pulse.GetCharge());
            //dom_times.push_back(pulse.GetTime());

            // check if we need to keep going
            if (pulse.GetTime() > max_edge + 4*unc_time){
                break;
            }

            // get cumulative pdf at first time edge
            double sigma = (time_edges[0] - pulse.GetTime()) / unc_time;
            double current_cdf = 0.5*(1 + std::erf(sigma/sqrt2));

            // get iterator for vector
            std::vector<double>::iterator values_iterator =
                bin_values_list.begin();

            // loop through time windows
            for (unsigned int i=1; i<n_bins+1; i++){

                // calculate cdf at this time edge
                double sigma = (time_edges[i] - pulse.GetTime()) / unc_time;
                double new_cdf = 0.5*(1 + std::erf(sigma/sqrt2));

                // accumulate charge in bin
                /*bin_values_list[i-1] +=
                    (new_cdf - current_cdf) * pulse.GetCharge();*/
                *values_iterator +=
                    (new_cdf - current_cdf) * pulse.GetCharge();
                values_iterator += 1;

                current_cdf = new_cdf;
            }

        }

        unsigned int n_total = n_bins;
        if (add_dom_info){
            n_total += 2;
            bin_values_list.push_back(distance);
            bin_values_list.push_back(unc_time);
        }

        // compress output: throw out extremely small values
        boost::python::list bin_exclusions_list; // empty dummy exclusions
        boost::python::list bin_indices_list;
        boost::python::list bin_values_list_compressed;
        std::vector<double>::iterator values_iterator =
                bin_values_list.begin();
        bool found_at_least_one_element = false;
        for (unsigned int i=0; i < n_total; i++){
            if (*values_iterator > 1e-7){
                bin_values_list_compressed.append(*values_iterator);
                bin_indices_list.append(i);
                found_at_least_one_element = true;
            }
            values_iterator += 1;
        }

        if (found_at_least_one_element){
            data_dict[dom_pulses.first] = boost::python::make_tuple(
                    bin_values_list_compressed, bin_indices_list,
                    bin_exclusions_list);
        }
    }

    return  data_dict;
}

// equivalent to numpy searchsorted
template <typename T>
unsigned int find_index(std::vector<T> array, T value){
    unsigned int index = 0;
    while (value > array[index] && index < array.size()){
        index++;
    }
    return index;
}

template <typename T>
inline boost::python::dict get_mc_tree_input_data_dict(
                                  boost::python::object& frame_obj,
                                  boost::python::object& angle_bins_obj,
                                  boost::python::object& distance_bins_obj,
                                  const double distance_cutoff,
                                  const double energy_cutoff,
                                  const bool add_distance
                                ) {

    // extract objects
    I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);

    const unsigned int num_dist_bins =  len(distance_bins_obj) - 1;
    const unsigned int num_angle_bins =  len(angle_bins_obj) - 1;
    const unsigned int num_bins =
        1 + add_distance + num_dist_bins * num_angle_bins;

    std::vector<double> angle_bins;
    std::vector<double> distance_bins;
    for (int i = 0; i < num_angle_bins + 1; ++i){
        angle_bins.push_back(
                boost::python::extract<double>(angle_bins_obj[i]));
    }
    for (int i = 0; i < num_dist_bins + 1; ++i){
        distance_bins.push_back(
                boost::python::extract<double>(distance_bins_obj[i]));
    }

    // Get geometry map
    const I3OMGeoMap& omgeo = frame.Get<I3Geometry>("I3Geometry").omgeo;

    // Get I3MCTree map
    const I3MCTree& mctree = frame.Get<I3MCTree>("I3MCTree");

    // create data array for DOMs [shape: (86, 60, num_bins)]
    double init_dist_value = 0.;
    if (add_distance){
        init_dist_value = std::numeric_limits<double>::max();
    }
    std::vector<double>  dom_data[86][60];
    for(size_t string = 1; string <= 86; string++){
        for(size_t dom = 1; dom <= 60; dom++){
            dom_data[string-1][dom-1] = std::vector<double>();

            // intialize values
            dom_data[string-1][dom-1].push_back(init_dist_value);
            for(size_t i=1; i < num_bins; i++){
                dom_data[string-1][dom-1].push_back(0);
            }
        }
    }

    // walk through energy losses and calculate data
    for (auto const& loss : mctree){

        // skip energy loss if it is not one of the allowed types
        bool is_allowed_type = false;
        switch(loss.GetType()){
            case I3Particle::ParticleType::NuclInt:
            case I3Particle::ParticleType::PairProd:
            case I3Particle::ParticleType::Brems:
            case I3Particle::ParticleType::DeltaE:
            case I3Particle::ParticleType::EMinus:
            case I3Particle::ParticleType::EPlus:
            case I3Particle::ParticleType::Hadrons: is_allowed_type=true;
                                                    break;
            default: is_allowed_type=false;
        }
        if (!is_allowed_type || loss.GetEnergy() < energy_cutoff ||
            loss.GetPos().Magnitude() > 2000){
            continue;
        }

        double min_distance = std::numeric_limits<double>::max();
        OMKey min_omkey = OMKey();

        // walk through DOMs and calculate data
        for(size_t string = 1; string <= 86; string++){
            for(size_t dom = 1; dom <= 60; dom++){

                const OMKey& om_key = OMKey(string, dom);

                // get position of DOM
                const I3Position& pos = omgeo.at(om_key).position;

                // calculate distance and opening angle to DOM
                const I3Position& diff = pos - loss.GetPos();
                I3Particle diff_p = I3Particle();
                diff_p.SetDir(diff.GetX(), diff.GetY(), diff.GetZ());
                const double angle = I3Calculator::Angle(diff_p, loss);
                const double distance = diff.Magnitude();

                // sort loss energy to correct bin index
                unsigned int index_angle =
                    find_index<double>(angle_bins, angle) - 1;
                unsigned int index_dist =
                    find_index<double>(distance_bins, distance) - 1;

                // check if it is out of bounds
                bool out_of_bounds = false;
                if (index_angle < 0 || index_angle >= num_angle_bins){
                    out_of_bounds = true;
                }
                if (index_dist < 0 || index_dist >= num_dist_bins){
                    out_of_bounds = true;
                }

                if (!out_of_bounds){
                    // calculate index
                    unsigned int index = add_distance +
                        index_dist * num_angle_bins + index_angle;

                    // accumulate energy
                    dom_data[string - 1][dom - 1][index] += loss.GetEnergy();
                }

                // check if distance is closest so far
                if (add_distance && distance < dom_data[string - 1][dom - 1][0]){
                    dom_data[string - 1][dom - 1][0] = distance;
                }

                if (distance < min_distance){
                    min_distance = distance;
                    min_omkey = om_key;
                }
            }
        }

        // add energy loss to closest DOM within distance_cutoff
        if (loss.GetEnergy() > energy_cutoff){
            if (min_distance < distance_cutoff){
                dom_data[min_omkey.GetString() - 1]
                         [min_omkey.GetOM() - 1]
                         [num_bins-1] += loss.GetEnergy();
            }
        }
    }

    // empty dummy exclusions
    boost::python::list bin_exclusions_list;

    boost::python::dict data_dict;
    for(size_t string = 1; string <= 86; string++){
        for(size_t dom = 1; dom <= 60; dom++){
            const OMKey& om_key = OMKey(string, dom);

            // create value and bin indices lists
            boost::python::list bin_values_list;
            boost::python::list bin_indices_list;
            boost::python::list bin_values_list_compressed;

            std::vector<double>::iterator values_iterator =
                dom_data[string - 1][dom - 1].begin();
            bool found_at_least_one_element = false;

            for (unsigned int i=0; i < num_bins; i++){
                if (*values_iterator > 1e-7){
                    bin_values_list_compressed.append(*values_iterator);
                    bin_indices_list.append(i);
                    found_at_least_one_element = true;
                }
                values_iterator += 1;
            }

            if (found_at_least_one_element){
                data_dict[om_key] = boost::python::make_tuple(
                        bin_values_list_compressed, bin_indices_list,
                        bin_exclusions_list);
            }
        }
    }
    return  data_dict;
}

boost::python::list get_time_bin_exclusions(
                            boost::python::object& frame_obj,
                            boost::python::object& om_key_obj,
                            const boost::python::list& time_bins,
                            const boost::python::list& excluded_doms_obj,
                            const boost::python::object& partial_exclusion_obj,
                            const double time_offset
                            ){
    /*
    This method creates a list of bin indices of a given DOM that are excluded
    as specified with the excluded_doms and partial_exclusion.
    */

    // extract c++ data types from python objects
    I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);
    const OMKey& om_key = boost::python::extract<OMKey&>(om_key_obj);
    const bool partial_exclusion =
        boost::python::extract<bool>(partial_exclusion_obj);

    std::vector<std::string> excluded_doms;
    for (int i = 0; i < len(excluded_doms_obj); ++i){
        excluded_doms.push_back(
                boost::python::extract<std::string>(excluded_doms_obj[i]));
    }

    // create a list for the bin exclusions
    boost::python::list bin_exclusions_list;
    std::set<int> bin_exclusions_set;

    // ------------------------------------------------
    // remove all excluded DOMs (!= I3TimeWindowSeries)
    // ------------------------------------------------
    I3TimeWindowSeriesMap exclusions;
    for (const std::string& mapname: excluded_doms) {

            I3TimeWindowSeriesMapConstPtr exclusions_segment =
                frame.Get<I3TimeWindowSeriesMapConstPtr>(mapname);

            I3VectorOMKeyConstPtr excludedoms =
                frame.Get<I3VectorOMKeyConstPtr>(mapname);

            if (exclusions_segment && partial_exclusion) {
                /* These are TimeWindowSeries from which we have to remove
                pulses that lie within, since partial exclusion is true and
                we want to keep the rest of the pulses of a DOM.

                For now we will combine all of the provided TimeWindowSeries
                into a single I3TimeWindowSeriesMap: exclusions.
                */
                for (I3TimeWindowSeriesMap::const_iterator i =
                    exclusions_segment->begin(); i !=
                    exclusions_segment->end(); i++){
                        if (i->first == om_key){
                            exclusions[i->first] = exclusions[i->first] |
                                i->second;
                        }
                    }
            } else if (exclusions_segment && !partial_exclusion) {
                /* These are TimeWindowSeries, but since partial exclusion
                is false, we will remove all DOMs that have exclusion time
                windows defined.
                */
                for (I3TimeWindowSeriesMap::const_iterator i =
                    exclusions_segment->begin(); i !=
                    exclusions_segment->end(); i++){
                        if (i->first == om_key) {
                            // remove complete DOM
                            bin_exclusions_list.append(-1);
                            return bin_exclusions_list;
                        }
                    }

            } else if (excludedoms) {
                /* These are whole DOMs to be ommited.
                Examples can be BrightDOMs which is a list of DOM OMKeys that
                will be completely removed.
                */
                for (const OMKey& key: *excludedoms){
                        if (key == om_key) {
                            // remove complete DOM
                            bin_exclusions_list.append(-1);
                            return bin_exclusions_list;
                        }
                    }
            }
    }

    for (I3TimeWindowSeriesMap::const_iterator tws = exclusions.begin();
                tws != exclusions.end(); tws++){

        for (auto &readout : tws->second) {
            for (int index = 0; index < len(time_bins) - 1; ++index){
                if (readout.GetStart() - time_offset < time_bins[index+1] &&
                    readout.GetStop() - time_offset > time_bins[index] ){
                    bin_exclusions_set.insert(index);
                }
            }
        }

    }

    for(auto const index : bin_exclusions_set){
        bin_exclusions_list.append(index);
    }

    return bin_exclusions_list;
}

template <typename T>
inline boost::python::tuple restructure_pulses(
                                  const boost::python::object& pulse_map_obj
                                ) {

    // Get pulse map
    const I3RecoPulseSeriesMap& pulse_map = boost::python::extract<I3RecoPulseSeriesMap&>(pulse_map_obj);

    boost::python::dict dom_times_dict;
    boost::python::dict dom_charges_dict;

    std::vector<double> charges;
    std::vector<double> times;

    for (auto const& dom_pulses : pulse_map){
        std::vector<double> dom_charges;
        std::vector<double> dom_times;

        for (auto const& pulse : dom_pulses.second){
            dom_charges.push_back(pulse.GetCharge());
            dom_times.push_back(pulse.GetTime());
        }
        dom_times_dict[dom_pulses.first] = stdVecToNumpyArray<double>(dom_times);
        dom_charges_dict[dom_pulses.first] = stdVecToNumpyArray<double>(dom_charges);

        // extend total charges and times
        // reserve() is optional - just to improve performance
        charges.reserve(charges.size() + dom_charges.size());
        charges.insert(charges.end(), dom_charges.begin(), dom_charges.end());

        times.reserve(times.size() + dom_times.size());
        times.insert(times.end(), dom_times.begin(), dom_times.end());

    }
    boost::python::object charges_numpy = stdVecToNumpyArray<double>(charges);
    boost::python::object times_numpy = stdVecToNumpyArray<double>(times);

    return  boost::python::make_tuple(
                charges_numpy, times_numpy, dom_times_dict, dom_charges_dict );
}

/* Combine DOM exclusions into a single vector of DOMs and a single
TimeWindowsSeriesMap
Returns:
    I3VectorOMKey: the combined DOM exclusions
    I3TimeWindowSeriesMap: the combined time window exclusions
*/
boost::python::tuple combine_exclusions(
                    boost::python::object& frame_obj,
                    const boost::python::list& excluded_doms_obj,
                    const boost::python::object& partial_exclusion_obj){

    // extract c++ data types from python objects
    I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);
    const bool partial_exclusion =
        boost::python::extract<bool>(partial_exclusion_obj);

    std::vector<std::string> exclusion_series;
    for (int i = 0; i < len(excluded_doms_obj); ++i){
        exclusion_series.push_back(
                boost::python::extract<std::string>(excluded_doms_obj[i]));
    }

    I3VectorOMKey exclusion_doms;
    I3TimeWindowSeriesMap exclusion_tws;
    for (const std::string& mapname: exclusion_series) {

        I3TimeWindowSeriesMapConstPtr exclusions_segment =
            frame.Get<I3TimeWindowSeriesMapConstPtr>(mapname);

        I3VectorOMKeyConstPtr excludedoms =
            frame.Get<I3VectorOMKeyConstPtr>(mapname);

        if (exclusions_segment && partial_exclusion) {
            /* These are TimeWindowSeries from which we have to remove
            pulses that lie within, since partial exclusion is true and
            we want to keep the rest of the pulses of a DOM.

            For now we will combine all of the provided TimeWindowSeries
            into a single I3TimeWindowSeriesMap: exclusion_tws.
            */
            for (I3TimeWindowSeriesMap::const_iterator i =
                exclusions_segment->begin(); i !=
                exclusions_segment->end(); i++){
                    exclusion_tws[i->first] = exclusion_tws[i->first] |
                        i->second;
            }
        } else if (exclusions_segment && !partial_exclusion) {
            /* These are TimeWindowSeries, but since partial exclusion
            is false, we will remove all DOMs that have exclusion time
            windows defined.
            */
            for (I3TimeWindowSeriesMap::const_iterator i =
                exclusions_segment->begin(); i !=
                exclusions_segment->end(); i++){
                    exclusion_doms.push_back(i->first);
            }

        } else if (excludedoms) {
            /* These are whole DOMs to be ommited.
            Examples can be BrightDOMs which is a list of DOM OMKeys that
            will be completely removed.
            */
            for (const OMKey& key: *excludedoms){
                    exclusion_doms.push_back(key);
            }
        }
    }
    return  boost::python::make_tuple(exclusion_doms, exclusion_tws);
}

I3RecoPulseSeriesMapPtr get_valid_pulse_map(
                            boost::python::object& frame_obj,
                            const boost::python::object& pulse_key_obj,
                            const boost::python::list& excluded_doms_obj,
                            const boost::python::object& partial_exclusion_obj,
                            const boost::python::object& verbose_obj){
    /*
    This method creates a new pulse series based on the given pulse series
    name, but excludes pulses and DOMs as specified with the excluded_doms
    and partial_exclusion.
    */

    // extract c++ data types from python objects
    I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);
    const std::string pulse_key =
        boost::python::extract<std::string>(pulse_key_obj);
    const bool partial_exclusion =
        boost::python::extract<bool>(partial_exclusion_obj);
    const bool verbose = boost::python::extract<bool>(verbose_obj);

    std::vector<std::string> excluded_doms;
    for (int i = 0; i < len(excluded_doms_obj); ++i){
        excluded_doms.push_back(
                boost::python::extract<std::string>(excluded_doms_obj[i]));
    }

    // keep track of removed number of pulses and DOMs
    unsigned int removed_doms = 0;
    unsigned int removed_pulses = 0;

    // get pulses
    const I3RecoPulseSeriesMap& pulses =
        frame.Get<I3RecoPulseSeriesMap>(pulse_key);

    // copy pulses so that we can modify them
    I3RecoPulseSeriesMap pulses_masked = I3RecoPulseSeriesMap(pulses);

    // ------------------------------------------------
    // remove all excluded DOMs (!= I3TimeWindowSeries)
    // ------------------------------------------------
    I3TimeWindowSeriesMap exclusions;
    for (const std::string& mapname: excluded_doms) {

            I3TimeWindowSeriesMapConstPtr exclusions_segment =
                frame.Get<I3TimeWindowSeriesMapConstPtr>(mapname);

            I3VectorOMKeyConstPtr excludedoms =
                frame.Get<I3VectorOMKeyConstPtr>(mapname);

            if (exclusions_segment && partial_exclusion) {
                /* These are TimeWindowSeries from which we have to remove
                pulses that lie within, since partial exclusion is true and
                we want to keep the rest of the pulses of a DOM.

                For now we will combine all of the provided TimeWindowSeries
                into a single I3TimeWindowSeriesMap: exclusions.
                */
                    for (I3TimeWindowSeriesMap::const_iterator i =
                        exclusions_segment->begin(); i !=
                        exclusions_segment->end(); i++){
                            exclusions[i->first] = exclusions[i->first] |
                                i->second;
                        }
            } else if (exclusions_segment && !partial_exclusion) {
                /* These are TimeWindowSeries, but since partial exclusion
                is false, we will remove all DOMs that have exclusion time
                windows defined.
                */
                    for (I3TimeWindowSeriesMap::const_iterator i =
                        exclusions_segment->begin(); i !=
                        exclusions_segment->end(); i++){
                            //exclusions[i->first].push_back(I3TimeWindow());
                            removed_doms += pulses_masked.erase(i->first);
                        }

            } else if (excludedoms) {
                /* These are whole DOMs to be ommited.
                Examples can be BrightDOMs which is a list of DOM OMKeys that
                will be completely removed.
                */
                    for (const OMKey& key: *excludedoms){
                            //exclusions[key].push_back(I3TimeWindow());
                            removed_doms += pulses_masked.erase(key);
                        }
            }
    }

    for (I3TimeWindowSeriesMap::const_iterator tws = exclusions.begin();
                tws != exclusions.end(); tws++){

        // ----------------------------
        // Get effective readout window
        // ----------------------------
        I3TimeWindowSeries effective_readouts;
        effective_readouts.push_back(I3TimeWindow());
        {
            // The effective readout windows are the set difference of the
            // global readout window and exclusion windows
            effective_readouts = effective_readouts & (~(tws->second));
        }
        // ----------------------------

        I3RecoPulseSeriesMap::const_iterator dom_pulse_ptr =
            pulses_masked.find(tws->first);

        // make sure key exists in pulses
        if (dom_pulse_ptr != pulses_masked.end()){

            // make iterator point at pulse series for the given OMKey
            I3RecoPulseSeries rps = dom_pulse_ptr->second;
            I3RecoPulseSeries masked_pulse_series = I3RecoPulseSeries();

            // reserve memory (Note: this reserves more than necessary)
            masked_pulse_series.reserve(rps.size());

            // define pointer to begin and end of pulse series
            auto begin(rps.cbegin()), end(rps.cend());
            for (auto &readout : effective_readouts) {
                auto pulse_time_sort = [](const I3RecoPulse &p, double t) {
                                                     return p.GetTime() < t; };
                begin = std::lower_bound(begin, end, readout.GetStart(),
                                         pulse_time_sort);
                end = std::lower_bound(begin, end, readout.GetStop(),
                                       pulse_time_sort);

                // insert valid pulses
                masked_pulse_series.insert(masked_pulse_series.end(),
                                           begin, end);

                // update new begin and end position
                begin = end;
                end = rps.cend();
            }

            // delete old pulse series and update with masked pulse series
            //pulses_masked.erase(tws->first);
            removed_pulses += rps.size() - masked_pulse_series.size();
            pulses_masked[tws->first] = masked_pulse_series;
        }
    }
    // ------------------------------------------------

    // write pulses to frame
    I3RecoPulseSeriesMapPtr fr_pulses =
        boost::make_shared<I3RecoPulseSeriesMap>(pulses_masked);
    //frame.Put(pulse_key + "_masked", fr_pulses);

    if (verbose){
        /*log_info(
            "[MaskPulses] Removed %d DOMs and %d additional pulses from %s",
            removed_doms, removed_pulses, pulse_key.c_str());*/
        std::cout << "[MaskPulses] Removed " << removed_doms <<" DOMs and "
                  << removed_pulses << " additional pulses from "
                  << pulse_key << std::endl;
    }

    return fr_pulses;
}

I3RecoPulseSeriesMapPtr merge_pulses(
                            boost::python::object& frame_obj,
                            const boost::python::object& pulse_key_obj,
                            const boost::python::object& time_threshold_obj){
    /*
    This method creates a new pulse series by combining pulses that are nearby
    in time. The threshold time under which pulses are combined to one is
    defined by `time_threshold`.
    */

    // extract c++ data types from python objects
    I3Frame& frame = boost::python::extract<I3Frame&>(frame_obj);
    const std::string pulse_key =
        boost::python::extract<std::string>(pulse_key_obj);
    const double time_threshold =
        boost::python::extract<double>(time_threshold_obj);

    // get pulses
    const I3RecoPulseSeriesMap& pulses =
        frame.Get<I3RecoPulseSeriesMap>(pulse_key);

    // create a new and empty pulse series mask
    I3RecoPulseSeriesMap pulses_merged = I3RecoPulseSeriesMap();

    // now iterate through DOMs
    for (I3RecoPulseSeriesMap::const_iterator dom_pulse_ptr = pulses.begin();
                dom_pulse_ptr != pulses.end(); dom_pulse_ptr++){

        // create new pulse series for this DOM
        I3RecoPulseSeries merged_pulse_series = I3RecoPulseSeries();

        // create iterator for pulse series
        I3RecoPulseSeries::const_iterator pulse_ptr =
            dom_pulse_ptr->second.begin();

        I3RecoPulse last_pulse = I3RecoPulse(*pulse_ptr);
        pulse_ptr += 1;

        // iterate through pulses at DOM
        while (pulse_ptr != dom_pulse_ptr->second.end()) {

            // check if this pulse is within the specified threshold
            if (pulse_ptr->GetTime() - last_pulse.GetTime() <= time_threshold){
                // merge pulse
                last_pulse.SetCharge(
                    pulse_ptr->GetCharge() + last_pulse.GetCharge());
            } else{
                // found new pulse: so append previous pulse
                merged_pulse_series.push_back(last_pulse);

                // set new last pulse
                last_pulse = I3RecoPulse(*pulse_ptr);
            }

            pulse_ptr += 1;
        }

        // add last pulse
        merged_pulse_series.push_back(last_pulse);

        // add to pulse series map
        pulses_merged[dom_pulse_ptr->first] = merged_pulse_series;

    }

    // create shared pulse series map
    I3RecoPulseSeriesMapPtr shared_pulse_series_map =
        boost::make_shared<I3RecoPulseSeriesMap>(pulses_merged);

    return shared_pulse_series_map;
}

BOOST_PYTHON_MODULE(ext_boost)
{
    #if BOOST_VERSION < 106500
        // Specify that py::numeric::array should refer to the Python type
        // numpy.ndarray (rather than the older Numeric.array).
        boost::python::numeric::array::set_module_and_type("numpy", "ndarray");
        // numpy requires this
        import_array();

    #else
        bn::initialize();

    #endif

    boost::python::def("restructure_pulses",
                       &restructure_pulses<double>);

    boost::python::def("get_valid_pulse_map",
                       &get_valid_pulse_map);

    boost::python::def("merge_pulses",
                       &merge_pulses);

    boost::python::def("combine_exclusions",
                       &combine_exclusions);

    boost::python::def("get_time_bin_exclusions",
                       &get_time_bin_exclusions);

    boost::python::def("get_reduced_summary_statistics_data",
                       &get_reduced_summary_statistics_data<double>);

    boost::python::def("fill_reduced_summary_statistics_data",
                       &fill_reduced_summary_statistics_data<double>,
                       (bp::arg("container"), bp::arg("pulse_key"),
                        bp::arg("add_total_charge"), bp::arg("add_t_first"),
                        bp::arg("add_t_std"), bp::arg("write_to_frame"),
                        bp::arg("batch_index"))
                       );

    boost::python::def("get_cascade_classification_data",
                       &get_cascade_classification_data<double>);

    boost::python::def("get_mc_tree_input_data_dict",
                       &get_mc_tree_input_data_dict<double>);
}


// Copyright 2020 Dow, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

#include <QtWidgets>

#include <array>
#include <limits>

#include <boost/log/trivial.hpp>
#include <boost/log/attributes/scoped_attribute.hpp>
#include <boost/optional.hpp>

#define BOOST_MULTI_INDEX_DISABLE_SERIALIZATION
#include <boost/multi_index_container.hpp>
#include <boost/multi_index/random_access_index.hpp>
#include <boost/multi_index/hashed_index.hpp>
#include <boost/multi_index/member.hpp>

#include "GenericDistributionDialog.h"
#include "AppStyle.h"

enum class DistributionType
{
    Discrete,
    Continuous,
    Step
};

enum class DistributionCategory
{
    Uniform,
    Bernoulli,
    Poisson,
    Normal,
    Sampling,
    Miscellaneous
};

using DT = DistributionType;
using DC = DistributionCategory;

// FIXME: refactor lookup tables

const QMap<DistributionCategory, QString> categoryMap {
    { DC::Uniform,       "Uniform Distributions"       },
    { DC::Bernoulli,     "Bernoulli Distributions"     },
    { DC::Poisson,       "Poisson Distributions"       },
    { DC::Normal,        "Normal Distributions"        },
    { DC::Sampling,      "Sampling Distributions"      },
    { DC::Miscellaneous, "Miscellaneous Distributions" }
};

// Reverse Sort
const QMultiMap<DistributionCategory, GenericDistribution::DistributionID> categoryToIDs {
    { DC::Uniform,        GenericDistribution::ID_UniformReal       },
    { DC::Uniform,        GenericDistribution::ID_UniformInt        },
    { DC::Bernoulli,      GenericDistribution::ID_NegativeBinomial  },
    { DC::Bernoulli,      GenericDistribution::ID_Geometric         },
    { DC::Bernoulli,      GenericDistribution::ID_Binomial          },
    { DC::Poisson,        GenericDistribution::ID_Laplace           },
    { DC::Poisson,        GenericDistribution::ID_Beta              },
    { DC::Poisson,        GenericDistribution::ID_ExtremeValue      },
    { DC::Poisson,        GenericDistribution::ID_Weibull           },
    { DC::Poisson,        GenericDistribution::ID_Gamma             },
    { DC::Poisson,        GenericDistribution::ID_Exponential       },
    { DC::Poisson,        GenericDistribution::ID_Poisson           },
    { DC::Normal,         GenericDistribution::ID_StudentT          },
    { DC::Normal,         GenericDistribution::ID_FisherF           },
    { DC::Normal,         GenericDistribution::ID_Cauchy            },
    { DC::Normal,         GenericDistribution::ID_NCChiSquared      },
    { DC::Normal,         GenericDistribution::ID_ChiSquared        },
    { DC::Normal,         GenericDistribution::ID_Lognormal         },
    { DC::Normal,         GenericDistribution::ID_Normal            },
    { DC::Sampling,       GenericDistribution::ID_PiecewiseLinear   },
    { DC::Sampling,       GenericDistribution::ID_PiecewiseConstant },
    { DC::Sampling,       GenericDistribution::ID_Discrete          },
    { DC::Miscellaneous,  GenericDistribution::ID_Triangle          }
};

struct DistributionInfo
{
    DistributionCategory category;
    DistributionType type;
    QString fullName;
};

const QMap<GenericDistribution::DistributionID, DistributionInfo> distributionMap {
    { GenericDistribution::ID_UniformReal,        { DC::Uniform,        DT::Continuous,  "Uniform Distribution (Real)"         }},
    { GenericDistribution::ID_UniformInt,         { DC::Uniform,        DT::Discrete,    "Uniform Distribution (Int)"          }},
    { GenericDistribution::ID_NegativeBinomial,   { DC::Bernoulli,      DT::Discrete,    "Negative Binomial Distribution"      }},
    { GenericDistribution::ID_Geometric,          { DC::Bernoulli,      DT::Discrete,    "Geometric Distribution"              }},
    { GenericDistribution::ID_Binomial,           { DC::Bernoulli,      DT::Discrete,    "Binomial Distribution"               }},
    { GenericDistribution::ID_Laplace,            { DC::Poisson,        DT::Continuous,  "Laplace Distribution"                }},
    { GenericDistribution::ID_Beta,               { DC::Poisson,        DT::Continuous,  "Beta Distribution"                   }},
    { GenericDistribution::ID_ExtremeValue,       { DC::Poisson,        DT::Continuous,  "Extreme Value Distribution"          }},
    { GenericDistribution::ID_Weibull,            { DC::Poisson,        DT::Continuous,  "Weibull Distribution"                }},
    { GenericDistribution::ID_Gamma,              { DC::Poisson,        DT::Continuous,  "Gamma (Erlang) Distribution"         }},
    { GenericDistribution::ID_Exponential,        { DC::Poisson,        DT::Continuous,  "Exponential Distribution"            }},
    { GenericDistribution::ID_Poisson,            { DC::Poisson,        DT::Discrete,    "Poisson Distribution"                }},
    { GenericDistribution::ID_StudentT,           { DC::Normal,         DT::Continuous,  "Student's t Distribution"            }},
    { GenericDistribution::ID_FisherF,            { DC::Normal,         DT::Continuous,  "F Distribution"                      }},
    { GenericDistribution::ID_Cauchy,             { DC::Normal,         DT::Continuous,  "Cauchy-Lorentz Distribution"         }},
    { GenericDistribution::ID_NCChiSquared,       { DC::Normal,         DT::Continuous,  "Noncentral Chi-Squared Distribution" }},
    { GenericDistribution::ID_ChiSquared,         { DC::Normal,         DT::Continuous,  "Chi Squared Distribution"            }},
    { GenericDistribution::ID_Lognormal,          { DC::Normal,         DT::Continuous,  "Log Normal Distribution"             }},
    { GenericDistribution::ID_Normal,             { DC::Normal,         DT::Continuous,  "Normal (Gaussian) Distribution"      }},
    { GenericDistribution::ID_PiecewiseLinear,    { DC::Sampling,       DT::Continuous,  "Piecewise Linear Distribution"       }},
    { GenericDistribution::ID_PiecewiseConstant,  { DC::Sampling,       DT::Step,        "Piecewise Constant Distribution"     }},
    { GenericDistribution::ID_Discrete,           { DC::Sampling,       DT::Discrete,    "Discrete Distribution"               }},
    { GenericDistribution::ID_Triangle,           { DC::Miscellaneous,  DT::Continuous,  "Triangular Distribution"             }}
};

struct ParameterInfo
{
    QString name;
    boost::optional<double> value;
    boost::optional<double> min;
    boost::optional<double> max;
    int precision;
    double step;
};

// Reverse Sort
const QMultiMap<GenericDistribution::DistributionID, ParameterInfo> parameterMap = {
//    Distribution                                 Name                               Value       Minimum       Maximum      Decimals    Step
//    -----------------------------------------    ------------------------     -----------   -----------   -----------   -----------    ----
    { GenericDistribution::ID_UniformInt,        { "Upper: ",                   boost::none,  boost::none,  boost::none,            0,      1 }},
    { GenericDistribution::ID_UniformInt,        { "Lower: ",                   boost::none,  boost::none,  boost::none,            0,      1 }},
    { GenericDistribution::ID_UniformReal,       { "Upper: ",                   boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_UniformReal,       { "Lower: ",                   boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Binomial,          { "Probability of success: ",  boost::none,     0.000001,        1.000,            6,   0.01 }},
    { GenericDistribution::ID_Binomial,          { "Trials: ",                  boost::none,            1,  boost::none,            0,      1 }},
    { GenericDistribution::ID_Geometric,         { "Probability: ",             boost::none,     0.000001,        0.999,            6,   0.01 }},
    { GenericDistribution::ID_NegativeBinomial,  { "Probability of success: ",  boost::none,     0.000001,        0.999,            6,   0.01 }},
    { GenericDistribution::ID_NegativeBinomial,  { "Successes: ",               boost::none,            1,  boost::none,            0,      1 }},
    { GenericDistribution::ID_Beta,              { "Beta: ",                    boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Beta,              { "Alpha: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Exponential,       { "Rate: ",                    boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_ExtremeValue,      { "Scale: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_ExtremeValue,      { "Location: ",                boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Gamma,             { "Scale: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Gamma,             { "Shape: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Laplace,           { "Scale: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Laplace,           { "Location: ",                boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Poisson,           { "Mean: ",                    boost::none,     0.000001,  boost::none,            6,      1 }},
    { GenericDistribution::ID_Weibull,           { "Scale: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Weibull,           { "Shape: ",                   boost::none,            1,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Cauchy,            { "Scale: ",                   boost::none,            0,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Cauchy,            { "Location: ",                boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_ChiSquared,        { "Degrees of freedom: ",      boost::none,            2,         1000,            0,      1 }},
    { GenericDistribution::ID_NCChiSquared,      { "Non-centrality: ",          boost::none,     0.000001,         1000,            6,   0.01 }},
    { GenericDistribution::ID_NCChiSquared,      { "Degrees of freedom: ",      boost::none,            1,         1000,            0,      1 }},
    { GenericDistribution::ID_FisherF,           { "Degrees of freedom 2: ",    boost::none,            1,          100,            0,      1 }},
    { GenericDistribution::ID_FisherF,           { "Degrees of freedom 1: ",    boost::none,            2,           20,            0,      1 }},
    { GenericDistribution::ID_Lognormal,         { "Scale: ",                   boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Lognormal,         { "Location: ",                boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Normal,            { "Standard deviation: ",      boost::none,     0.000001,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Normal,            { "Mean: ",                    boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_StudentT,          { "Degrees of freedom: ",      boost::none,            1,          100,            0,      1 }},
    { GenericDistribution::ID_Triangle,          { "Upper: ",                   boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Triangle,          { "Mode: ",                    boost::none,  boost::none,  boost::none,            6,   0.01 }},
    { GenericDistribution::ID_Triangle,          { "Lower: ",                   boost::none,  boost::none,  boost::none,            6,   0.01 }}
};

// Remaps parameters of the Boost.Random distribution to ParameterInfo.
// TODO: store vector of functions? e.g.
// std::vector<boost::function<void>> functions;
// functions.push_back(&a));
// functions.push_back(&b));

struct ParamVisitor : public boost::static_visitor<>
{
    void operator()(const Distribution::UniformReal& d) {
        id = GenericDistribution::ID_UniformReal;
        paramList = parameterMap.values(GenericDistribution::ID_UniformReal);
        paramList[0].value = d.a();
        paramList[1].value = d.b();
    }
    void operator()(const Distribution::UniformInt& d) {
        id = GenericDistribution::ID_UniformInt;
        paramList = parameterMap.values(GenericDistribution::ID_UniformInt);
        paramList[0].value = d.a();
        paramList[1].value = d.b();
    }
    void operator()(const Distribution::Binomial& d) {
        id = GenericDistribution::ID_Binomial;
        paramList = parameterMap.values(GenericDistribution::ID_Binomial);
        paramList[0].value = d.t();
        paramList[1].value = d.p();
    }
    void operator()(const Distribution::Geometric& d) {
        id = GenericDistribution::ID_Geometric;
        paramList = parameterMap.values(GenericDistribution::ID_Geometric);
        paramList[0].value = d.p();
    }
    void operator()(const Distribution::NegativeBinomial& d) {
        id = GenericDistribution::ID_NegativeBinomial;
        paramList = parameterMap.values(GenericDistribution::ID_NegativeBinomial);
        paramList[0].value = d.k();
        paramList[1].value = d.p();
    }
    void operator()(const Distribution::Poisson& d) {
        id = GenericDistribution::ID_Poisson;
        paramList = parameterMap.values(GenericDistribution::ID_Poisson);
        paramList[0].value = d.mean();
    }
    void operator()(const Distribution::Exponential& d) {
        id = GenericDistribution::ID_Exponential;
        paramList = parameterMap.values(GenericDistribution::ID_Exponential);
        paramList[0].value = d.lambda();
    }
    void operator()(const Distribution::Gamma& d) {
        id = GenericDistribution::ID_Gamma;
        paramList = parameterMap.values(GenericDistribution::ID_Gamma);
        paramList[0].value = d.alpha();
        paramList[1].value = d.beta();
    }
    void operator()(const Distribution::Weibull& d) {
        id = GenericDistribution::ID_Weibull;
        paramList = parameterMap.values(GenericDistribution::ID_Weibull);
        paramList[0].value = d.a();
        paramList[1].value = d.b();
    }
    void operator()(const Distribution::ExtremeValue& d) {
        id = GenericDistribution::ID_ExtremeValue;
        paramList = parameterMap.values(GenericDistribution::ID_ExtremeValue);
        paramList[0].value = d.a();
        paramList[1].value = d.b();
    }
    void operator()(const Distribution::Beta& d) {
        id = GenericDistribution::ID_Beta;
        paramList = parameterMap.values(GenericDistribution::ID_Beta);
        paramList[0].value = d.alpha();
        paramList[1].value = d.beta();
    }
    void operator()(const Distribution::Laplace& d) {
        id = GenericDistribution::ID_Laplace;
        paramList = parameterMap.values(GenericDistribution::ID_Laplace);
        paramList[0].value = d.mean();
        paramList[1].value = d.beta();
    }
    void operator()(const Distribution::Normal& d) {
        id = GenericDistribution::ID_Normal;
        paramList = parameterMap.values(GenericDistribution::ID_Normal);
        paramList[0].value = d.mean();
        paramList[1].value = d.sigma();
    }
    void operator()(const Distribution::Lognormal& d) {
        id = GenericDistribution::ID_Lognormal;
        paramList = parameterMap.values(GenericDistribution::ID_Lognormal);
        paramList[0].value = d.m();
        paramList[1].value = d.s();
    }
    void operator()(const Distribution::ChiSquared& d) {
        id = GenericDistribution::ID_ChiSquared;
        paramList = parameterMap.values(GenericDistribution::ID_ChiSquared);
        paramList[0].value = d.n();
    }
    void operator()(const Distribution::NCChiSquared& d) {
        id = GenericDistribution::ID_NCChiSquared;
        paramList = parameterMap.values(GenericDistribution::ID_NCChiSquared);
        paramList[0].value = d.k();
        paramList[1].value = d.lambda();
    }
    void operator()(const Distribution::Cauchy& d) {
        id = GenericDistribution::ID_Cauchy;
        paramList = parameterMap.values(GenericDistribution::ID_Cauchy);
        paramList[0].value = d.a();
        paramList[1].value = d.b();
    }
    void operator()(const Distribution::FisherF& d) {
        id = GenericDistribution::ID_FisherF;
        paramList = parameterMap.values(GenericDistribution::ID_FisherF);
        paramList[0].value = d.m();
        paramList[1].value = d.n();
    }
    void operator()(const Distribution::StudentT& d) {
        id = GenericDistribution::ID_StudentT;
        paramList = parameterMap.values(GenericDistribution::ID_StudentT);
        paramList[0].value = d.n();
    }
    void operator()(const Distribution::Discrete& d) {
        id = GenericDistribution::ID_Discrete;
    }
    void operator()(const Distribution::PiecewiseConstant& d) {
        id = GenericDistribution::ID_PiecewiseConstant;
    }
    void operator()(const Distribution::PiecewiseLinear& d) {
        id = GenericDistribution::ID_PiecewiseLinear;
    }
    void operator()(const Distribution::Triangle& d) {
        id = GenericDistribution::ID_Triangle;
        paramList = parameterMap.values(GenericDistribution::ID_Triangle);
        paramList[0].value = d.a();
        paramList[1].value = d.b();
        paramList[2].value = d.c();
    }
    void operator()(const Distribution::Constant& d) {}

    GenericDistribution::DistributionID id;
    QList<ParameterInfo> paramList;
};

// Sets reasonable default limits for the PDF plot, based on current parameters.

struct PdfLimitsVisitor : public boost::static_visitor<>
{
    void operator()(const Distribution::UniformReal& d) {
        x0 = d.a() - (d.b() - d.a());
        x1 = d.b() + (d.b() - d.a());
    }
    void operator()(const Distribution::UniformInt& d) {
        x0 = d.a() - (d.b() - d.a());
        x1 = d.b() + (d.b() - d.a());
    }
    void operator()(const Distribution::Binomial& d) {
        x0 = 0;
        x1 = d.t();
    }
    void operator()(const Distribution::Geometric& d) {
        x0 = 1;
        x1 = 10;
    }
    void operator()(const Distribution::NegativeBinomial& d) {
        x0 = 0;
        x1 = 25;
    }
    void operator()(const Distribution::Poisson& d) {
        x0 = 0;
        x1 = 2 * d.mean();
    }
    void operator()(const Distribution::Exponential& d) {
        x0 = 0;
        x1 = 5;
    }
    void operator()(const Distribution::Gamma& d) {
        x0 = 0;
        x1 = d.alpha() * 2;
    }
    void operator()(const Distribution::Weibull& d) {
        x0 = 0;
        x1 = 5;
    }
    void operator()(const Distribution::ExtremeValue& d) {
        x0 = d.a() - (4 * d.b());
        x1 = d.a() + (6 * d.b());
    }
    void operator()(const Distribution::Beta& d) {
        x0 = 0.01;
        x1 = 0.99;
    }
    void operator()(const Distribution::Laplace& d) {
        x0 = d.mean()- (4 * d.beta());
        x1 = d.mean() + (4 * d.beta());
    }
    void operator()(const Distribution::Normal& d) {
        x0 = d.mean() - 5;
        x1 = d.mean() + 5;
    }
    void operator()(const Distribution::Lognormal& d) {
        x0 = 0;
        x1 = 5;
    }
    void operator()(const Distribution::ChiSquared& d) {
        x0 = 0;
        x1 = 8;
    }
    void operator()(const Distribution::NCChiSquared& d) {
        x0 = 0;
        x1 = 8;
    }
    void operator()(const Distribution::Cauchy& d) {
        x0 = d.a() - 5;
        x1 = d.a() + 5;
    }
    void operator()(const Distribution::FisherF& d) {
        x0 = 0;
        x1 = 5;
    }
    void operator()(const Distribution::StudentT& d) {
        x0 = -5;
        x1 = 5;
    }
    void operator()(const Distribution::Discrete& d) {}
    void operator()(const Distribution::PiecewiseConstant& d) {}
    void operator()(const Distribution::PiecewiseLinear& d) {}
    void operator()(const Distribution::Triangle& d) {
        x0 = d.a();
        x1 = d.c();
    }
    void operator()(const Distribution::Constant& d) {}

    double x0;
    double x1;
};

GenericDistributionDialog::GenericDistributionDialog(const GenericDistribution& d, QWidget *parent)
    : QDialog(parent), currentDist(d)
{
    setWindowTitle(tr("Define Distribution"));
    setWindowFlag(Qt::Tool);
    setWindowIcon(this->style()->standardIcon(static_cast<QStyle::StandardPixmap>(AppStyle::CP_LineEditFunctionIcon)));

    distributionTree = new QTreeView;
    distributionTree->setHeaderHidden(true);
    distributionTree->setEditTriggers(QAbstractItemView::NoEditTriggers);

    label1 = new QLabel;
    label2 = new QLabel;
    label3 = new QLabel;

    sb1 = new QDoubleSpinBox;
    sb2 = new QDoubleSpinBox;
    sb3 = new QDoubleSpinBox;

    sb1->setMinimum(std::numeric_limits<double>::lowest());
    sb2->setMinimum(std::numeric_limits<double>::lowest());
    sb3->setMinimum(std::numeric_limits<double>::lowest());
    sb1->setMaximum(std::numeric_limits<double>::max());
    sb2->setMaximum(std::numeric_limits<double>::max());
    sb3->setMaximum(std::numeric_limits<double>::max());

    pdfPlot = new StandardPlot;
    pdfPlot->setXAxisTitle("Value");
    pdfPlot->setYAxisTitle("Density");

    pdfPlotCurveData = new QwtPointSeriesData;
    pdfPlotCurve = new QwtPlotCurve;
    pdfPlotCurve->setRenderHint(QwtPlotItem::RenderAntialiased, true);
    pdfPlotCurve->attach(pdfPlot);

    QColor baseColor = QWidget::palette().base().color();

    QFormLayout *controlsLayout = new QFormLayout;
    controlsLayout->addRow(label1, sb1);
    controlsLayout->addRow(label2, sb2);
    controlsLayout->addRow(label3, sb3);
    QFrame *controlsFrame = new QFrame;
    controlsFrame->setLayout(controlsLayout);
    controlsFrame->setFrameStyle(QFrame::StyledPanel | QFrame::Plain);
    QPalette controlsFramePalette = controlsFrame->palette();
    controlsFramePalette.setColor(QPalette::Window, baseColor);
    controlsFrame->setPalette(controlsFramePalette);
    controlsFrame->setAutoFillBackground(true);

    // Button Box
    buttonBox = new QDialogButtonBox(QDialogButtonBox::Ok | QDialogButtonBox::Cancel);
    connect(buttonBox, &QDialogButtonBox::accepted, this, &QDialog::accept);
    connect(buttonBox, &QDialogButtonBox::rejected, this, &QDialog::reject);

    // Distribution Controls
    QSplitter *controlsSplitter = new QSplitter;
    controlsSplitter->setOrientation(Qt::Vertical);
    controlsSplitter->setMinimumWidth(300);
    controlsSplitter->addWidget(distributionTree);
    controlsSplitter->addWidget(controlsFrame);

    // Plot Frame
    QFrame *plotFrame = new QFrame;
    QVBoxLayout *plotFrameLayout = new QVBoxLayout;
    plotFrameLayout->addWidget(pdfPlot);
    plotFrame->setLayout(plotFrameLayout);
    plotFrame->setFrameStyle(QFrame::StyledPanel | QFrame::Plain);
    QPalette plotFramePalette = plotFrame->palette();
    plotFramePalette.setColor(QPalette::Window, baseColor);
    plotFrame->setPalette(plotFramePalette);
    plotFrame->setAutoFillBackground(true);

    // Splitter
    QSplitter *mainSplitter = new QSplitter;
    mainSplitter->addWidget(controlsSplitter);
    mainSplitter->addWidget(plotFrame);
    mainSplitter->setOrientation(Qt::Horizontal);

    // Main Layout
    QVBoxLayout *mainLayout = new QVBoxLayout;
    mainLayout->addWidget(mainSplitter);
    mainLayout->addWidget(buttonBox);

    setLayout(mainLayout);

    initializeModel();
    initializeConnections();
    resetControls();
    updatePlot();
}

void GenericDistributionDialog::initializeModel()
{
    distributionModel = new QStandardItemModel;
    auto keys = categoryToIDs.uniqueKeys();
    for (const auto &i : keys) {
        QString categoryStr = categoryMap.value(i);
        QStandardItem *rootItem = new QStandardItem(categoryStr);
        rootItem->setData(-1, Qt::UserRole);
        for (const auto &id : categoryToIDs.values(i)) {
            DistributionInfo info = distributionMap.value(id);
            QStandardItem *item = new QStandardItem(info.fullName);
            item->setData(static_cast<int>(id), Qt::UserRole); // DistributionID
            rootItem->appendRow(item);
        }
        distributionModel->appendRow(rootItem);

        // FontRole
        rootItem->setSelectable(false);
        QFont font = qvariant_cast<QFont>(rootItem->data(Qt::FontRole));
        font.setBold(true);
        rootItem->setData(font, Qt::FontRole);
    }

    distributionTree->setModel(distributionModel);
}

void GenericDistributionDialog::initializeConnections()
{
    connect(distributionTree->selectionModel(), &QItemSelectionModel::currentChanged,
        [=](const QModelIndex& currentIndex, const QModelIndex&)
    {
        QStandardItem *currentItem = distributionModel->itemFromIndex(currentIndex);
        int id = qvariant_cast<int>(currentItem->data(Qt::UserRole)); // DistributionID

        if (id >= 0) {
            currentID = static_cast<GenericDistribution::DistributionID>(id);
            setDistribution(currentID);
            resetControls();
            updatePlot();
        }
    });

    connect(sb1, QOverload<double>::of(&QDoubleSpinBox::valueChanged),
            this, &GenericDistributionDialog::onValueChanged);
    connect(sb2, QOverload<double>::of(&QDoubleSpinBox::valueChanged),
            this, &GenericDistributionDialog::onValueChanged);
    connect(sb3, QOverload<double>::of(&QDoubleSpinBox::valueChanged),
            this, &GenericDistributionDialog::onValueChanged);
}

void GenericDistributionDialog::resetControls()
{
    ParamVisitor pv;
    boost::apply_visitor(pv, currentDist);

    currentID = pv.id;

    // Iterate over the distribution model, select current
    for (int i = 0; i < distributionModel->rowCount(); ++i) {
        QModelIndex parentIndex = distributionModel->index(i, 0);
        if (distributionModel->hasChildren(parentIndex)) {
            for (int j = 0; j < distributionModel->rowCount(parentIndex); ++j) {
                QModelIndex index = distributionModel->index(j, 0, parentIndex);
                QStandardItem *item = distributionModel->itemFromIndex(index);
                int id = qvariant_cast<int>(item->data(Qt::UserRole));
                if (id == static_cast<int>(pv.id)) {
                    distributionTree->expand(parentIndex);
                    distributionTree->selectionModel()->select(index, QItemSelectionModel::Select);
                }
            }
        }
    }

    int n = pv.paramList.length();

    double min = std::numeric_limits<double>::lowest();
    double max = std::numeric_limits<double>::max();

    label1->setVisible(false);
    label2->setVisible(false);
    label3->setVisible(false);

    sb1->blockSignals(true);
    sb2->blockSignals(true);
    sb3->blockSignals(true);

    sb1->setVisible(false);
    sb2->setVisible(false);
    sb3->setVisible(false);

    if (n > 0) {
        ParameterInfo pi1 = pv.paramList[0];
        label1->setText(pi1.name);
        sb1->setMinimum(pi1.min.value_or(min));
        sb1->setMaximum(pi1.max.value_or(max));
        sb1->setDecimals(pi1.precision);
        sb1->setSingleStep(pi1.step);
        sb1->setValue(pi1.value.value());
        label1->setVisible(true);
        sb1->setVisible(true);
        sb1->blockSignals(false);
    }
    if (n > 1) {
        ParameterInfo pi2 = pv.paramList[1];
        label2->setText(pi2.name);
        sb2->setMinimum(pi2.min.value_or(min));
        sb2->setMaximum(pi2.max.value_or(max));
        sb2->setDecimals(pi2.precision);
        sb2->setSingleStep(pi2.step);
        sb2->setValue(pi2.value.value());
        label2->setVisible(true);
        sb2->setVisible(true);
        sb2->blockSignals(false);
    }
    if (n > 2) {
        ParameterInfo pi3 = pv.paramList[2];
        label3->setText(pi3.name);
        sb3->setMinimum(pi3.min.value_or(min));
        sb3->setMaximum(pi3.max.value_or(max));
        sb3->setDecimals(pi3.precision);
        sb3->setSingleStep(pi3.step);
        sb3->setValue(pi3.value.value());
        label3->setVisible(true);
        sb3->setVisible(true);
        sb3->blockSignals(false);
    }
}

void GenericDistributionDialog::setDistribution(GenericDistribution::DistributionID id)
{
    // For selected distributions, set more appropriate defaults.

    switch (id) {
        case GenericDistribution::ID_UniformInt:
            currentDist = Distribution::UniformInt(0, 10);
            break;
        case GenericDistribution::ID_UniformReal:
            currentDist = Distribution::UniformReal(0, 10);
            break;
        case GenericDistribution::ID_Binomial:
            currentDist = Distribution::Binomial(10, 0.5);
            break;
        case GenericDistribution::ID_Geometric:
            currentDist = Distribution::Geometric();
            break;
        case GenericDistribution::ID_NegativeBinomial:
            currentDist = Distribution::NegativeBinomial(10, 0.5);
            break;
        case GenericDistribution::ID_Poisson:
            currentDist = Distribution::Poisson(10);
            break;
        case GenericDistribution::ID_Exponential:
            currentDist = Distribution::Exponential();
            break;
        case GenericDistribution::ID_Gamma:
            currentDist = Distribution::Gamma(5, 1);
            break;
        case GenericDistribution::ID_Weibull:
            currentDist = Distribution::Weibull(1, 1);
            break;
        case GenericDistribution::ID_ExtremeValue:
            currentDist = Distribution::ExtremeValue();
            break;
        case GenericDistribution::ID_Beta:
            currentDist = Distribution::Beta(0.5, 2);
            break;
        case GenericDistribution::ID_Laplace:
            currentDist = Distribution::Laplace();
            break;
        case GenericDistribution::ID_Normal:
            currentDist = Distribution::Normal(0, 1);
            break;
        case GenericDistribution::ID_Lognormal:
            currentDist = Distribution::Lognormal();
            break;
        case GenericDistribution::ID_ChiSquared:
            currentDist = Distribution::ChiSquared(2);
            break;
        case GenericDistribution::ID_NCChiSquared:
            currentDist = Distribution::NCChiSquared();
            break;
        case GenericDistribution::ID_Cauchy:
            currentDist = Distribution::Cauchy();
            break;
        case GenericDistribution::ID_FisherF:
            currentDist = Distribution::FisherF(2, 1);
            break;
        case GenericDistribution::ID_StudentT:
            currentDist = Distribution::StudentT();
            break;
        case GenericDistribution::ID_Discrete:
            currentDist = Distribution::Discrete();
            break;
        case GenericDistribution::ID_PiecewiseConstant:
            currentDist = Distribution::PiecewiseConstant();
            break;
        case GenericDistribution::ID_PiecewiseLinear:
            currentDist = Distribution::PiecewiseLinear();
            break;
        case GenericDistribution::ID_Triangle:
            currentDist = Distribution::Triangle(-1, 0, 1);
            break;
    }
}

GenericDistribution GenericDistributionDialog::getDistribution() const
{
    return currentDist;
}

void GenericDistributionDialog::updatePlot()
{
    BOOST_LOG_SCOPED_THREAD_TAG("Source", "Distribution");

    double step;
    QVector<QPointF> samples;
    QVector<double> xvec;
    QVector<double> yvec;

    // Set curve styles
    DistributionInfo info = distributionMap.value(currentID);
    switch (info.type) {
        case DT::Continuous:
            step = 0.01;
            pdfPlotCurve->setPen(Qt::blue, 1.5);
            pdfPlotCurve->setStyle(QwtPlotCurve::Lines);
            break;
        case DT::Discrete:
            step = 1;
            pdfPlotCurve->setPen(Qt::blue, 2);
            pdfPlotCurve->setStyle(QwtPlotCurve::Sticks);
            break;
        case DT::Step:
            step = 0.01;
            pdfPlotCurve->setPen(Qt::blue, 1.5);
            pdfPlotCurve->setStyle(QwtPlotCurve::Steps);
            break;
    }

    // Reset the plot
    pdfPlotCurveData->setSamples(samples);
    pdfPlotCurve->setData(pdfPlotCurveData);
    pdfPlot->replot();

    // Estimate X-axis limits
    PdfLimitsVisitor pv;
    boost::apply_visitor(pv, currentDist);

    // Populate independent variable vector
    for (double x = pv.x0; x < pv.x1; x += step)
        xvec.push_back(x);
    xvec.push_back(pv.x1); // last value

    // Evaluate the PDF at each value
    try {
        currentDist.pdf(xvec, yvec);
    } catch (const std::exception &e) {
        BOOST_LOG_TRIVIAL(warning) << e.what();
        return;
    }

    if (xvec.size() != yvec.size())
        return;

    for (int i=0; i < xvec.size(); ++i) {
        QPointF point(xvec[i], yvec[i]);
        samples.push_back(point);
    }

    pdfPlotCurveData->setSamples(samples);
    pdfPlotCurve->setData(pdfPlotCurveData);
    pdfPlot->replot();
}

void GenericDistributionDialog::onValueChanged(double)
{
    if (static_cast<int>(currentID) < 0)
        return;

    auto id = static_cast<GenericDistribution::DistributionID>(currentID);

    double a = sb1->value();
    double b = sb2->value();
    double c = sb3->value();

    switch (id) {
        case GenericDistribution::ID_UniformInt:
            currentDist = Distribution::UniformInt(a, b);
            break;
        case GenericDistribution::ID_UniformReal:
            currentDist = Distribution::UniformReal(a, b);
            break;
        case GenericDistribution::ID_Binomial:
            currentDist = Distribution::Binomial(a, b);
            break;
        case GenericDistribution::ID_Geometric:
            currentDist = Distribution::Geometric(a);
            break;
        case GenericDistribution::ID_NegativeBinomial:
            currentDist = Distribution::NegativeBinomial(a, b);
            break;
        case GenericDistribution::ID_Poisson:
            currentDist = Distribution::Poisson(a);
            break;
        case GenericDistribution::ID_Exponential:
            currentDist = Distribution::Exponential(a);
            break;
        case GenericDistribution::ID_Gamma:
            currentDist = Distribution::Gamma(a, b);
            break;
        case GenericDistribution::ID_Weibull:
            currentDist = Distribution::Weibull(a, b);
            break;
        case GenericDistribution::ID_ExtremeValue:
            currentDist = Distribution::ExtremeValue(a, b);
            break;
        case GenericDistribution::ID_Beta:
            currentDist = Distribution::Beta(a, b);
            break;
        case GenericDistribution::ID_Laplace:
            currentDist = Distribution::Laplace(a, b);
            break;
        case GenericDistribution::ID_Normal:
            currentDist = Distribution::Normal(a, b);
            break;
        case GenericDistribution::ID_Lognormal:
            currentDist = Distribution::Lognormal(a, b);
            break;
        case GenericDistribution::ID_ChiSquared:
            currentDist = Distribution::ChiSquared(a);
            break;
        case GenericDistribution::ID_NCChiSquared:
            currentDist = Distribution::NCChiSquared(a, b);
            break;
        case GenericDistribution::ID_Cauchy:
            currentDist = Distribution::Cauchy(a, b);
            break;
        case GenericDistribution::ID_FisherF:
            currentDist = Distribution::FisherF(a, b);
            break;
        case GenericDistribution::ID_StudentT:
            currentDist = Distribution::StudentT(a);
            break;
        case GenericDistribution::ID_Discrete:
            currentDist = Distribution::Discrete();
            break;
        case GenericDistribution::ID_PiecewiseConstant:
            currentDist = Distribution::PiecewiseConstant();
            break;
        case GenericDistribution::ID_PiecewiseLinear:
            currentDist = Distribution::PiecewiseLinear();
            break;
        case GenericDistribution::ID_Triangle:
            currentDist = Distribution::Triangle(a, b, c);
            break;
    }

    updatePlot();
}


// Copyright (c) 2014 The Bitcoin Core developers
// Copyright (c) 2014-2015 The Dash developers
// Copyright (c) 2015-2017 The PIVX developers
// Copyright (c) 2018 LightPayCoin developers
// Copyright (c) 2018 The InVertex developers
// Distributed under the MIT/X11 software license, see the accompanying
// file COPYING or http://www.opensource.org/licenses/mit-license.php.

#include "primitives/transaction.h"
#include "main.h"

#include <boost/test/unit_test.hpp>

BOOST_AUTO_TEST_SUITE(main_tests)

CAmount nMoneySupplyPoWEnd = 13000000 * COIN;

BOOST_AUTO_TEST_CASE(subsidy_limit_test)
{
    CAmount nSum = 0;
    //PoW + PoS Premine
    for (int nHeight = 0; nHeight < 1000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 500 * COIN);
        nSum += nSubsidy;
    }
    //Main PoS phase
    for (int nHeight = 1001; nHeight < 100000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 1 * COIN);
        nSum += nSubsidy;
    }

    for (int nHeight = 100001; nHeight < 110000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 25 * COIN);
        nSum += nSubsidy;
    }

    for (int nHeight = 110001; nHeight < 120000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 100 * COIN);
        nSum += nSubsidy;
    }

    for (int nHeight = 120001; nHeight < 130000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 50 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 130001; nHeight < 140000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 150 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 140001; nHeight < 150000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 75 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 150001; nHeight < 160000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 200 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 160001; nHeight < 170000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 100 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 170001; nHeight < 180000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 300 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 180001; nHeight < 190000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 150 * COIN);
        nSum += nSubsidy;
    }
    for (int nHeight = 190001; nHeight < 200000; nHeight += 1) {
        CAmount nSubsidy = GetBlockValue(nHeight);
        BOOST_CHECK(nSubsidy <= 50 * COIN);
        nSum += nSubsidy;
        BOOST_CHECK(nSum > 0 && nSum <= nMoneySupplyPoWEnd);
    }

}

BOOST_AUTO_TEST_SUITE_END()