//
// Created by Vasiliy Ershov on 08/11/2016.
//

#ifndef PROJECT_GAMMA_POISSON_MODEL_HPP
#define PROJECT_GAMMA_POISSON_MODEL_HPP

#include <common/utils/parallel/openmp_wrapper.h>
#include <boost/math/special_functions/binomial.hpp>
#include <boost/math/special_functions/gamma.hpp>
#include <boost/math/special_functions/trigamma.hpp>
#include <vector>
#include "kmer_data.hpp"
#include "thread_utils.h"
#include "valid_hkmer_generator.hpp"
//

namespace n_gamma_poisson_model {

struct QualFunc {
  double alpha_;
  double beta_;

  double operator()(double x) const { return alpha_ * x + beta_; }

  double GenomicLogLikelihood(double x) const {
    const double val = (*this)(x);
    const double exp_point = exp(val);
    return val - (std::isfinite(exp_point) ? log(1 + exp_point) : val);
  }
};

class GammaDistribution {
 private:
  double shape_;
  double rate_;
  double log_gamma_at_shape_;

 public:
  GammaDistribution(const GammaDistribution&) = default;

  GammaDistribution& operator=(const GammaDistribution&) = default;

  GammaDistribution(const double shape = 1, const double rate = 1)
      : shape_(shape), rate_(rate) {
    log_gamma_at_shape_ = boost::math::lgamma(shape_);
  }

  inline double GetShape() const { return shape_; }

  inline double GetRate() const { return rate_; }

  inline double LogGammaAtShape() const { return log_gamma_at_shape_; }
};

class GammaMixture {
 private:
  GammaDistribution first_;
  GammaDistribution second_;
  double first_weight_;

 public:
  GammaMixture() : first_(1, 1), second_(1, 1), first_weight_(-1) {}

  GammaMixture(const GammaDistribution& first,
                const GammaDistribution& second,
                double firstWeight)
      : first_(first), second_(second), first_weight_(firstWeight) {}

  const GammaDistribution& GetFirst() const { return first_; }

  const GammaDistribution& GetSecond() const { return second_; }

  double GetFirstWeight() const { return first_weight_; }
};

class PoissonGammaDistribution {
 private:
  const GammaDistribution& prior_;
  static std::array<double, 100000> log_gamma_integer_cache_;

 private:
  inline double IntLogGamma(size_t count) const {
    if (count < log_gamma_integer_cache_.size()) {
      return log_gamma_integer_cache_[count];
    } else {
      return boost::math::lgamma(((double)count) + 1);
    }
  }

 public:
  PoissonGammaDistribution(const GammaDistribution& prior) : prior_(prior) {}

  inline double PartialLogLikelihood(size_t count) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();

    double ll = 0.0;
    ll += a * log(b) - (a + (double)count) * log(b + 1);
    ll +=
        boost::math::lgamma(prior_.GetShape() + (double)count) - prior_.LogGammaAtShape();
    return ll;
  }

  inline double LogLikelihood(size_t count) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();

    double ll = 0.0;
    ll += a * log(b) - (a + (double)count) * log(b + 1);
    ll += boost::math::lgamma(prior_.GetShape() + ((double)count)) - IntLogGamma(count) -
          prior_.LogGammaAtShape();

    return ll;
  }

  inline double Quantile(double p) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();
    return boost::math::ibeta_inva(a, 1.0 / (1.0 + b), 1.0 - p);
  }

  inline double Cumulative(size_t count) const {
    const double a = prior_.GetShape();
    const double b = prior_.GetRate();

    return 1.0 - boost::math::ibeta((double)count + 1, a, 1.0 / (1.0 + b));
  }
};

constexpr int RunSizeLimit = 8;

class ParametricClusterModel {
 private:
  GammaMixture prior_;
  QualFunc qual_func_;
  double count_threshold_;
  std::array<double, RunSizeLimit> alphas_;

 public:
 public:
  ParametricClusterModel() : count_threshold_(100000) {}

  double ErrorRate(const int runSize) const {
    auto idx = runSize - 1;
    idx = std::max(idx, 0);
    idx = std::min(idx, (int)(alphas_.size() - 1));
    return alphas_[idx] * (runSize == 0 ? 0.5 : 1);
  }

  double ExpectedErrorRate(const hammer::HKMer& from,
                           const hammer::HKMer& to) const {
    double errRate = 0;
    for (unsigned i = 0; i < hammer::K; ++i) {
      errRate +=
          std::abs(from[i].len - to[i].len) * log(ErrorRate(from[i].len));
      //      errRate += std::abs(from[i].len - to[i].len) *
      //      log(ErrorRate(from[i].len)) - log(1.0 - ErrorRate(from[i].len));
    }
    return exp(errRate);
  }

  ParametricClusterModel(const GammaMixture& prior,
                         const QualFunc& qualFunc,
                         const double countThreshold,
                         const std::array<double, RunSizeLimit>& alphas)
      : prior_(prior), qual_func_(qualFunc), count_threshold_(countThreshold) {
    std::copy(alphas.begin(), alphas.end(), alphas_.begin());
    for (unsigned i = 0; i < RunSizeLimit; ++i) {
      INFO("Run length " << i << " estimated error rate " << alphas_[i]);
    }
  }

  ParametricClusterModel(const ParametricClusterModel& other) = default;

  ParametricClusterModel& operator=(const ParametricClusterModel&) = default;

  double QualityLogPrior(double qual) const {
    return std::max(std::min(qual_func_.GenomicLogLikelihood(qual), -1e-10), -1000.0);
  }

  bool NeedSubcluster(const hammer::KMerStat& stat) const {
    return qual_func_.GenomicLogLikelihood(stat.qual) > -0.1 &&
           stat.count >= count_threshold_;
  }

  const GammaDistribution& GenomicPrior() const { return prior_.GetFirst(); }

  const GammaDistribution& NoisePrior() const { return prior_.GetSecond(); }

  double GenerateLogLikelihood(double expectedNoiseCount,
                               size_t noiseCount) const {
    const auto& prior = NoisePrior();

    GammaDistribution posterior(prior.GetShape() + expectedNoiseCount,
                                 prior.GetRate() + 1);
    return PoissonGammaDistribution(posterior).LogLikelihood(noiseCount);
  }

  double GenomicLogLikelihood(size_t count) const {
    const auto& prior = GenomicPrior();
    const double a = prior.GetShape();
    const double b = prior.GetRate();

    double ll = a * log(b) - (a + (double)count) * log(b + 1);
    ll += boost::math::lgamma(prior.GetShape() + ((double)count)) -
          prior.LogGammaAtShape() - boost::math::lgamma(((double)count) + 1);
    return ll;
  }
};

// this class estimate prior distribution.
class TClusterModelEstimator {
 private:
  const KMerData& data_;
  double threshold_;
  unsigned num_threads_;
  size_t max_terations_;
  bool calc_likelihood_;

 private:

  struct TClusterSufficientStat {
    double count_ = 0;
    double qualtiy_ = 0;
    double genomic_class_prob_ = 0;
  };

  struct TQualityStat {
    double quality_ = 0;
    double class_ = 0;

    TQualityStat(double quality, double cls) : quality_(quality), class_(cls) {}
  };

  struct TRunErrorStats {
    const KMerData* data_;
    std::array<double, RunSizeLimit> error_counts_;
    std::array<double, RunSizeLimit> total_count_;

    TRunErrorStats(const KMerData& data)
        : data_(&data){

          };

    std::array<double, RunSizeLimit> EstimateAlphas(
        size_t priorSize = 100) const {
      const double priors[] = {0.002, 0.004, 0.01, 0.02,
                               0.035, 0.05,  0.09, 0.11};

      std::array<double, RunSizeLimit> alphas;
      for (unsigned i = 0; i < RunSizeLimit; ++i) {
        alphas[i] = (error_counts_[i] + priors[i] * (double)priorSize) /
                    (total_count_[i] + (double)priorSize);
      }
      alphas[0] *= 2;
      return alphas;
    };

    TRunErrorStats& operator+=(const TRunErrorStats& other) {
      if (this != &other) {
        for (unsigned i = 0; i < RunSizeLimit; ++i) {
          error_counts_[i] += other.error_counts_[i];
          total_count_[i] += other.total_count_[i];
        }
      }
      return *this;
    }

    void Add(const std::vector<size_t>& indices, size_t centerIdx) {
      const auto& center = (*data_)[centerIdx].kmer;

      for (auto idx : indices) {
        if (idx == centerIdx) {
          continue;
        }
        double errKmerCount = (double)(*data_)[idx].count;
        const auto& errKmer = (*data_)[idx].kmer;
        for (unsigned i = 0; i < hammer::K; ++i) {
          if (center[i].len > RunSizeLimit) {
            continue;
          }
          const int len = center[i].len - 1;
          total_count_[len] += errKmerCount;
          if (center[i].len != errKmer[i].len) {
            error_counts_[len] += errKmerCount;
          }
        }
      }
      for (unsigned i = 0; i < hammer::K; ++i) {
        if (center[i].len > RunSizeLimit) {
          continue;
        }
        total_count_[center[i].len - 1] += (*data_)[centerIdx].count;
      }
    }
  };

  inline void Expectation(const PoissonGammaDistribution& first,
                          const PoissonGammaDistribution& second,
                          const QualFunc& qualFunc,
                          TClusterSufficientStat& center) const {
    const double logPrior = qualFunc.GenomicLogLikelihood(center.qualtiy_) +
                            log(boost::math::gamma_q(center.count_, threshold_));

    const double firstLL = first.PartialLogLikelihood((size_t)center.count_) + logPrior;
    const double secondLL = second.PartialLogLikelihood((size_t)center.count_) +
                            log(std::max(1.0 - exp(logPrior), 1e-20));

    const double posterior = 1.0 / (1.0 + exp(secondLL - firstLL));
    center.genomic_class_prob_ = posterior;
  }

  inline void QualityExpectation(const QualFunc& qualFunc,
                                 TClusterSufficientStat& center) const {
    center.genomic_class_prob_ =
        exp(qualFunc.GenomicLogLikelihood(center.qualtiy_));
  }

  inline TClusterSufficientStat Create(const size_t centerIdx) const {
    TClusterSufficientStat stat;
    stat.genomic_class_prob_ =
        data_[centerIdx].count > 0
            ? boost::math::gamma_q(data_[centerIdx].count, threshold_)
            : 0;
    stat.count_ = data_[centerIdx].count;
    stat.qualtiy_ = data_[centerIdx].qual;
    return stat;
  }

  std::vector<TClusterSufficientStat> CreateSufficientStats(
      const std::vector<size_t>& clusterCenters) const {
    std::vector<TClusterSufficientStat> clusterSufficientStat;
    clusterSufficientStat.reserve(clusterCenters.size());

    for (size_t i = 0; i < clusterCenters.size(); ++i) {
      const size_t centerIdx = clusterCenters[i];
      auto stat = Create(centerIdx);
      if (stat.count_ > 0) {
        clusterSufficientStat.push_back(stat);
      }
    }
    return clusterSufficientStat;
  }

  std::vector<TQualityStat> CreateQualityStats(
      const std::vector<std::vector<size_t>>& clusters,
      const std::vector<size_t>& clusterCenters) const {
    std::vector<TQualityStat> qualities;
    qualities.reserve(clusterCenters.size());

    for (size_t i = 0; i < clusterCenters.size(); ++i) {
      const size_t centerIdx = clusterCenters[i];
      if (data_[centerIdx].count >= threshold_) {
        for (auto idx : clusters[i]) {
          if (idx != centerIdx) {
            qualities.push_back(TQualityStat(data_[idx].qual, 0));
          }
        }
        qualities.push_back(TQualityStat(data_[centerIdx].qual, 1));
      }
    }
    return qualities;
  }

  template <bool WEIGHTED = true>
  class TCountsStat {
   private:
    double count_ = 0;
    double count2_ = 0;
    double weight_ = 0;

   public:
    void Add(const TClusterSufficientStat& stat) {
      const double w = (WEIGHTED ? stat.genomic_class_prob_ : 1.0);
      count_ += w * stat.count_;
      count2_ += w * stat.count_ * stat.count_;
      weight_ += w;
    }

    TCountsStat& operator+=(const TCountsStat& other) {
      if (this != &other) {
        count_ += other.count_;
        count2_ += other.count2_;
        weight_ += other.weight_;
      }
      return *this;
    }

    double GetWeightedSum() const { return count_; }

    double GetWeightedSum2() const { return count2_; }

    double GetWeight() const { return weight_; }
  };

  class TLogGammaStat {
   private:
    double genomic_shape_;
    double non_genomic_shape_;
    double genomic_log_gamma_sum_ = 0;
    double non_genomic_log_gamma_sum_ = 0;

   public:
    TLogGammaStat(double genomicShape, double nonGenomicShape)
        : genomic_shape_(genomicShape), non_genomic_log_gamma_sum_(nonGenomicShape) {}

    void Add(const TClusterSufficientStat& stat) {
      genomic_log_gamma_sum_ += stat.genomic_class_prob_ *
                            boost::math::lgamma(stat.count_ + genomic_shape_);
      non_genomic_log_gamma_sum_ +=
          (1.0 - stat.genomic_class_prob_) *
          boost::math::lgamma(stat.count_ + non_genomic_shape_);
    }

    TLogGammaStat& operator+=(const TLogGammaStat& other) {
      if (this != &other) {
        genomic_log_gamma_sum_ += other.genomic_log_gamma_sum_;
        non_genomic_log_gamma_sum_ += other.non_genomic_log_gamma_sum_;
      }
      return *this;
    }

    double GetGenomicLogGammaSum() const { return genomic_log_gamma_sum_; }

    double GetNonGenomicLogGammaSum() const { return non_genomic_log_gamma_sum_; }
  };

  class TQualityLogitLinearRegressionPoint {
   private:
    // p(genomic) = exp(Alpha qual + beta) / (1.0 + exp(Alpha qual + beta))
    QualFunc func_;

    double likelihood_ = 0;

    double der_alpha_ = 0;
    double der_beta_ = 0;

    double der2_alpha_ = 0;
    double der2_beta_ = 0;
    double der2_alpha_beta_ = 0;

   public:
    TQualityLogitLinearRegressionPoint(QualFunc func) : func_(func) {}

    void Add(const TClusterSufficientStat& statistic) {
      Add(statistic.genomic_class_prob_, statistic.qualtiy_);
    }

    void Add(const TQualityStat& statistic) {
      Add(statistic.class_, statistic.quality_);
    }

    void Add(const double firstClassProb, double qual) {
      const double val = func_(qual);
      const double expPoint = exp(val);
      const double p =
          std::isfinite(expPoint) ? expPoint / (1.0 + expPoint) : 1.0;

      der_alpha_ += (firstClassProb - p) * qual;
      der_beta_ += firstClassProb - p;

      der2_alpha_ -= sqr(qual) * p * (1 - p);
      der2_beta_ -= p * (1 - p);
      der2_alpha_beta_ -= qual * p * (1 - p);

      likelihood_ += firstClassProb * val -
                    (std::isfinite(expPoint) ? log(1 + expPoint) : val);
    }

    TQualityLogitLinearRegressionPoint& operator+=(
        const TQualityLogitLinearRegressionPoint& other) {
      if (this != &other) {
        likelihood_ += other.likelihood_;

        der_alpha_ += other.der_alpha_;
        der_beta_ += other.der_beta_;

        der2_alpha_ += other.der2_alpha_;
        der2_beta_ += other.der2_beta_;
        der2_alpha_beta_ += other.der2_alpha_beta_;
      }
      return *this;
    }

    double GetLikelihood() const { return likelihood_; }

    double GetDerAlpha() const { return der_alpha_; }

    double GetDerBeta() const { return der_beta_; }

    double GetDer2Alpha() const { return der2_alpha_; }

    double GetDer2Beta() const { return der2_beta_; }

    double GetDer2AlphaBeta() const { return der2_alpha_beta_; }
  };

  QualFunc Update(const QualFunc& current,
                   const TQualityLogitLinearRegressionPoint& pointStats) const {
    const double dera = pointStats.GetDerAlpha();
    const double derb = pointStats.GetDerBeta();

    const double daa = pointStats.GetDer2Alpha() + 1e-3;
    const double dbb = pointStats.GetDer2Beta() + 1e-3;
    const double dab = pointStats.GetDer2AlphaBeta();
    const double det = daa * dbb - sqr(dab);

    double stepAlpha = (dbb * dera - dab * derb) / det;
    double stepBeta = (daa * derb - dab * dera) / det;

    INFO("Quality estimation iteration gradient: " << dera << " " << derb);
    INFO("Quality estimation likelihood: " << pointStats.GetLikelihood());

    return {current.alpha_ - stepAlpha, current.beta_ - stepBeta};
  }

  class TGammaDerivativesStats {
   private:
    double first_class_shift_;
    double second_class_shift_;

    double digamma_sum_first_ = 0;
    double trigamma_sum_first_ = 0;

    double digamma_sum_second_ = 0;
    double trigamma_sum_second_ = 0;

   public:
    TGammaDerivativesStats(double firstShift, double secondShift)
        : first_class_shift_(firstShift), second_class_shift_(secondShift) {}

    void Add(const TClusterSufficientStat& statistic) {
      const double p = statistic.genomic_class_prob_;
      digamma_sum_first_ +=
          p > 1e-3 ? p * boost::math::digamma(statistic.count_ + first_class_shift_)
                   : 0;
      trigamma_sum_first_ +=
          p > 1e-3
              ? p * boost::math::trigamma(statistic.count_ + first_class_shift_)
              : 0;

      digamma_sum_second_ +=
          p < (1.0 - 1e-3) ? (1.0 - p) * boost::math::digamma(statistic.count_ +
                                                              second_class_shift_)
                           : 0;
      trigamma_sum_second_ +=
          p < (1.0 - 1e-3) ? (1.0 - p) * boost::math::trigamma(statistic.count_ +
                                                               second_class_shift_)
                           : 0;
    }

    TGammaDerivativesStats& operator+=(const TGammaDerivativesStats& other) {
      if (this != &other) {
        digamma_sum_first_ += other.digamma_sum_first_;
        trigamma_sum_first_ += other.trigamma_sum_first_;

        digamma_sum_second_ = other.digamma_sum_second_;
        trigamma_sum_second_ += other.trigamma_sum_second_;
      }
      return *this;
    }

    double GetDigammaSumFirst() const { return digamma_sum_first_; }

    double GetTrigammaSumFirst() const { return trigamma_sum_first_; }

    double GetDigammaSumSecond() const { return digamma_sum_second_; }

    double GetTrigammaSumSecond() const { return trigamma_sum_second_; }
  };

  static inline double sqr(double x) { return x * x; }

  struct TDirection {
    double Direction;
    double GradientNorm;
    double Mu;
  };

  static TDirection MoveDirection(double shape, const double weightedSum,
                                  const double weight, const double digammaSum,
                                  const double trigammaSum,
                                  double regularizer = 1e-4) {
    const double mu = weight / weightedSum;
    const double digammaAtShape = boost::math::digamma(shape);
    const double trigammaAtShape = boost::math::trigamma(shape);

    const double b = mu * shape;

    const double der =
        weight * (log(b) - log(b + 1) - digammaAtShape) + digammaSum;
    const double der2 =
        trigammaSum + weight * (1.0 / shape - mu / (b + 1) - trigammaAtShape);

    return {-der / (der2 + regularizer), std::abs(der), mu};
  }

  double Likelihood(GammaDistribution& prior, double weightedSum,
                    double weight, double lgammaSum) {
    const double a = prior.GetShape();
    const double b = prior.GetRate();
    return weight * a * (log(b) - log(b + 1)) + weightedSum * log(b + 1) +
           lgammaSum - weight * prior.LogGammaAtShape();
  }

 public:
  TClusterModelEstimator(const KMerData& data, double threshold,
                         unsigned num_threads = 16,
                         size_t maxIterations = 40,
                         bool calcLikelihood = false)
      : data_(data),
        threshold_(threshold),
        num_threads_(num_threads),
        max_terations_(maxIterations),
        calc_likelihood_(calcLikelihood) {}

  static inline GammaDistribution Update(const GammaDistribution& point,
                                          const TDirection& direction,
                                          double minShape = 0.01) {
    double shape = std::max(point.GetShape() + direction.Direction, minShape);
    double rate = shape * direction.Mu;
    return GammaDistribution(shape, rate);
  }

  static GammaDistribution MomentMethodEstimator(const double sum,
                                                  const double sum2,
                                                  const double weight) {
    const double m = sum / weight;
    const double var = sum2 / weight - m * m;
    const double rate = 1.0 / std::max(var / m - 1, 1e-3);
    const double shape = m * rate;
    return GammaDistribution(shape, rate);
  }

  ParametricClusterModel Estimate(
      const std::vector<std::vector<size_t>>& clusters,
      const std::vector<size_t>& clusterCenter, const bool useEM = false,
      const size_t sample = 0) {
    if (sample && clusters.size() > sample) {
      std::vector<std::vector<size_t>> sampledClusters;
      std::vector<size_t> sampledCenters;
    }

    const auto qualityFunc = [&]() -> QualFunc {
      auto qualStats = CreateQualityStats(clusters, clusterCenter);

      QualFunc cursor = {-1e-5, 0.0};

      for (unsigned i = 0; i < 15; ++i) {
        const auto qualDerStats =
            n_computation_utils::TAdditiveStatisticsCalcer<
                TQualityStat, TQualityLogitLinearRegressionPoint>(qualStats,
                                                                  num_threads_)
                .Calculate([&]() -> TQualityLogitLinearRegressionPoint {
                  return TQualityLogitLinearRegressionPoint(cursor);
                });

        cursor = Update(cursor, qualDerStats);

        if ((std::abs(qualDerStats.GetDerAlpha()) +
             std::abs(qualDerStats.GetDerBeta())) < 1e-2) {
          break;
        }
      }

      INFO("Quality function: " << cursor.alpha_ << "q + " << cursor.beta_);
      return cursor;
    }();

    auto alphas = [&]() -> std::array<double, RunSizeLimit> {
      TRunErrorStats errorStats =
          n_computation_utils::ParallelStatisticsCalcer<TRunErrorStats>(
              num_threads_)
              .Calculate(
                  clusters.size(),
                  [&]() -> TRunErrorStats { return TRunErrorStats(data_); },
                  [&](TRunErrorStats& stat, size_t k) {
                    if (data_[clusterCenter[k]].count >= threshold_) {
                      stat.Add(clusters[k], clusterCenter[k]);
                    }
                  });
      return errorStats.EstimateAlphas();
    }();

    std::vector<TClusterSufficientStat> clusterSufficientStat =
        CreateSufficientStats(clusterCenter);

    const auto totalStats =
        n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                     TCountsStat<false>>(
            clusterSufficientStat, num_threads_)
            .Calculate([]() -> TCountsStat<false> {
              return TCountsStat<false>();
            });

#pragma omp parallel for num_threads(num_threads_)
    for (size_t k = 0; k < clusterSufficientStat.size(); ++k) {
      QualityExpectation(qualityFunc, clusterSufficientStat[k]);
    }

    auto countsStats =
        n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                     TCountsStat<true>>(
            clusterSufficientStat, num_threads_)
            .Calculate([]() -> TCountsStat<true> {
              return TCountsStat<true>();
            });

    GammaDistribution genomicPrior = [&]() -> GammaDistribution {
      const double m = countsStats.GetWeightedSum() / countsStats.GetWeight();
      const double var =
          countsStats.GetWeightedSum2() / countsStats.GetWeight() - m * m;
      const double rate = 1.0 / std::max(var / m - 1, 1e-3);
      const double shape = m * rate;
      return GammaDistribution(shape, rate);
    }();

    GammaDistribution nonGenomicPrior = [&]() -> GammaDistribution {
      const double m =
          (totalStats.GetWeightedSum() - countsStats.GetWeightedSum()) /
          (totalStats.GetWeight() - countsStats.GetWeight());
      const double var =
          (totalStats.GetWeightedSum2() - countsStats.GetWeightedSum2()) /
              (totalStats.GetWeight() - countsStats.GetWeight()) -
          m * m;
      const double rate = 1.0 / std::max(var / m - 1, 1e-3);
      const double shape = m * rate;
      return GammaDistribution(shape, rate);
    }();

    for (unsigned i = 0, steps = 0; i < max_terations_; ++i, ++steps) {
      auto gammaDerStats =
          n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                       TGammaDerivativesStats>(
              clusterSufficientStat, num_threads_)
              .Calculate([&]() -> TGammaDerivativesStats {
                return TGammaDerivativesStats(genomicPrior.GetShape(),
                                              nonGenomicPrior.GetShape());
              });

      auto genomicDirection = MoveDirection(
          genomicPrior.GetShape(), countsStats.GetWeightedSum(),
          countsStats.GetWeight(), gammaDerStats.GetDigammaSumFirst(),
          gammaDerStats.GetTrigammaSumFirst());

      auto nonGenomicDirection = MoveDirection(
          nonGenomicPrior.GetShape(),
          totalStats.GetWeightedSum() - countsStats.GetWeightedSum(),
          totalStats.GetWeight() - countsStats.GetWeight(),
          gammaDerStats.GetDigammaSumSecond(),
          gammaDerStats.GetTrigammaSumSecond());

      auto gradientNorm =
          genomicDirection.GradientNorm + nonGenomicDirection.GradientNorm;

      INFO("Iteration #" << i << " gradient norm " << gradientNorm);

      genomicPrior = Update(genomicPrior, genomicDirection);
      nonGenomicPrior = Update(nonGenomicPrior, nonGenomicDirection);

      if (calc_likelihood_) {
        auto logGammaStats =
            n_computation_utils::TAdditiveStatisticsCalcer<TClusterSufficientStat,
                                                         TLogGammaStat>(
                clusterSufficientStat, num_threads_)
                .Calculate([&]() -> TLogGammaStat {
                  return TLogGammaStat(genomicPrior.GetShape(),
                                       nonGenomicPrior.GetShape());
                });

        INFO("Genomic likelihood: " << Likelihood(
                 genomicPrior, countsStats.GetWeightedSum(),
                 countsStats.GetWeight(),
                 logGammaStats.GetGenomicLogGammaSum()));

        INFO("NonGenomic likelihood: " << Likelihood(
                 nonGenomicPrior,
                 totalStats.GetWeightedSum() - countsStats.GetWeightedSum(),
                 totalStats.GetWeight() - countsStats.GetWeight(),
                 logGammaStats.GetNonGenomicLogGammaSum()));
      }

      {
        INFO("Genomic gamma prior estimation step: shape "
             << genomicPrior.GetShape() << " and rate "
             << genomicPrior.GetRate());
        INFO("Nongenomic gamma prior estimation step: shape "
             << nonGenomicPrior.GetShape() << " and rate "
             << nonGenomicPrior.GetRate());
      }

      double shapeDiff = std::abs(genomicDirection.Direction) +
                         std::abs(nonGenomicDirection.Direction);

      if (useEM) {
        if ((shapeDiff < 1e-2) || gradientNorm < 1e-1 ||
            (steps == 5 && (i < max_terations_ - 10))) {
          PoissonGammaDistribution genomic(genomicPrior);
          PoissonGammaDistribution nonGenomic(nonGenomicPrior);
#pragma omp parallel for num_threads(num_threads_)
          for (size_t k = 0; k < clusterSufficientStat.size(); ++k) {
            Expectation(genomic, nonGenomic, qualityFunc,
                        clusterSufficientStat[k]);
          }

          countsStats = n_computation_utils::TAdditiveStatisticsCalcer<
                            TClusterSufficientStat, TCountsStat<true>>(
                            clusterSufficientStat, num_threads_)
                            .Calculate([]() -> TCountsStat<true> {
                              return TCountsStat<true>();
                            });
          steps = 0;
        }
      } else {
        if ((shapeDiff < 1e-4) || gradientNorm < 1e-2) {
          break;
        }
      }
    }

    INFO("Genomic gamma prior genomic estimated with shape "
         << genomicPrior.GetShape() << " and rate " << genomicPrior.GetRate());
    INFO("Nongenomic Gamma prior estimated with shape "
         << nonGenomicPrior.GetShape() << " and rate "
         << nonGenomicPrior.GetRate());

    return ParametricClusterModel(
        GammaMixture(genomicPrior, nonGenomicPrior,
                      countsStats.GetWeight() / totalStats.GetWeight()),
        qualityFunc, threshold_, alphas);
  }

  static GammaDistribution EstimatePrior(const std::vector<size_t>& counts) {
    const size_t observations = counts.size();
    double sum = 0;
    double sum2 = 0;
    for (auto count : counts) {
      sum += (double)count;
      sum2 += (double)count * (double)count;
    }

    GammaDistribution prior =
        TClusterModelEstimator::MomentMethodEstimator(sum, sum2, (double)observations);

    for (unsigned i = 0, steps = 0; i < 10; ++i, ++steps) {
      double digammaSum = 0;
      double trigammaSum = 0;
      for (auto count : counts) {
        digammaSum += boost::math::digamma((double)count + prior.GetShape());
        trigammaSum += boost::math::trigamma((double)count + prior.GetShape());
      }

      auto direction = MoveDirection(prior.GetShape(), sum, (double)observations,
                                     digammaSum, trigammaSum);

      const double shapeDiff = std::abs(direction.Direction);
      if (shapeDiff < 1e-3 || (direction.GradientNorm < 1e-4)) {
        break;
      }
      prior = Update(prior, direction, 1e-2);
    }
    return prior;
  }
};

}  // namespace NGammaPoissonModel

#endif  // PROJECT_GAMMA_POISSON_MODEL_HPP


#define EIGEN_USE_MKL_ALL

#include <iostream> // for standard output
#include <fstream> // for file input
#include <iomanip> // some nice printing functions
#include <random> // for randomness
#include <boost/algorithm/string.hpp> // string manipulation library
#include <boost/filesystem.hpp> // filesystem path manipulation library
#include <boost/program_options.hpp> // options parsing library
#include <eigen3/Eigen/Dense> // linear algebra library

#include "constants.h"
#include "qp-math.h"
#include "gates.h"
#include "nv-math.h"
#include "nv-control.h"

using namespace std;
using namespace Eigen;
namespace fs = boost::filesystem;
namespace po = boost::program_options;

int main(const int arg_num, const char *arg_vec[]) {

  // -------------------------------------------------------------------------------------
  // Set input options
  // -------------------------------------------------------------------------------------

  const uint help_text_length = 90;

  unsigned long long int seed;

  po::options_description general("General options", help_text_length);
  general.add_options()
    ("help,h", "produce help message")
    ("seed", po::value<unsigned long long int>(&seed)->default_value(0),
     "seed for random number generator")
    ;

  bool pair_search;
  bool coherence_scan;
  bool coherence_signal;
  bool angular_coherence_signal;
  bool rotation;
  bool coupling;
  bool iswap;
  bool swap;
  bool identity;
  bool swap_nvst;
  bool larmor_identity;
  bool initialize;
  bool initialize_x;
  bool initialize_larmor;
  bool testing_mode;

  po::options_description simulations("Available simulations",help_text_length);
  simulations.add_options()
    ("pair_search",
     po::value<bool>(&pair_search)->default_value(false)->implicit_value(true),
     "identify and characterize addressable larmor pairs")
    ("scan", po::value<bool>(&coherence_scan)->default_value(false)->implicit_value(true),
     "perform NV coherence scan for effective larmor frequencies")
    ("signal",
     po::value<bool>(&coherence_signal)->default_value(false)->implicit_value(true),
     "measure NV coherence signal as a function of the coupling constant"
     " at the frequency of a target nucleus")
    ("angular_signal", po::value<bool>(&angular_coherence_signal)
     ->default_value(false)->implicit_value(true),
     "measure NV coherence signal as a function of the coupling constant"
     " and angle of a control field at the frequency of a target nucleus")
    ("rotate", po::value<bool>(&rotation)->default_value(false)->implicit_value(true),
     "rotate an individual nucleus")
    ("couple", po::value<bool>(&coupling)->default_value(false)->implicit_value(true),
     "couple an individual nucleus to NV center")
    ("iswap", po::value<bool>(&iswap)->default_value(false)->implicit_value(true),
     "compute iSWAP fidelities")
    ("swap", po::value<bool>(&swap)->default_value(false)->implicit_value(true),
     "compute SWAP fidelities")
    ("identity", po::value<bool>(&identity)->default_value(false)->implicit_value(true),
     "compute fidelity of an identity operation on a single spin")
    ("swap_nvst", po::value<bool>(&swap_nvst)->default_value(false)->implicit_value(true),
     "compute SWAP_NVST fidelity")
    ("larmor_identity",
     po::value<bool>(&larmor_identity)->default_value(false)->implicit_value(true),
     "compute fidelity of an identity operation on a larmor qubit")
    ("initialize",
     po::value<bool>(&initialize)->default_value(false)->implicit_value(true),
     "compute fidelity of a deterministic initialization of a thermalized nucleus"
     " into |u> or |d>")
    ("initialize_x",
     po::value<bool>(&initialize_x)->default_value(false)->implicit_value(true),
     "compute fidelity of a probabalistic initialization of a thermalized nucleus"
     " into |u> +/- |d>")
    ("initialize_larmor",
     po::value<bool>(&initialize_larmor)->default_value(false)->implicit_value(true),
     "compute fidelity of a probabalistic initialization of a larmor pair"
     " from |dd> into |ud> +/- |du>")
    ("test" ,po::value<bool>(&testing_mode)->default_value(false)->implicit_value(true),
     "enable testing mode")
    ;

  double c13_abundance;
  double c13_factor;
  double nuclear_isolation;
  double larmor_isolation;
  double larmor_isolation_in_kHz;
  double min_hyperfine_xy;
  double min_hyperfine_xy_in_kHz;
  uint max_cluster_size;
  double hyperfine_cutoff;
  double hyperfine_cutoff_in_kHz;
  int ms;
  uint k_DD_int;
  axy_harmonic k_DD;
  double static_Bz;
  double static_Bz_in_gauss;
  double identity_time;
  double scale_factor;
  double integration_factor;
  bool no_nn;

  po::options_description simulation_options("Simulation options",help_text_length);
  simulation_options.add_options()
    ("c13_factor", po::value<double>(&c13_factor)->default_value(1),
     "abundance of C-13 relative to its natural abundance")
    ("nuclear_isolation", po::value<double>(&nuclear_isolation)->default_value(100),
     "maximum internuclear coupling factor for an 'isolated' nucleus (Hz)")
    ("larmor_isolation", po::value<double>(&larmor_isolation_in_kHz)->default_value(0),
     "isolation factor of parallel component of hyperfine field"
     " for an 'isolated' nucleus (kHz)")
    ("min_hyperfine_xy", po::value<double>(&min_hyperfine_xy_in_kHz)->default_value(0),
     "minimum magnitude of hyperfine field perpendicular to the NV axis"
     " for a larmor pair to be counted in a pair search (kHz)")
    ("max_cluster_size", po::value<uint>(&max_cluster_size)->default_value(6),
     "maximum allowable size of C-13 clusters")
    ("hyperfine_cutoff", po::value<double>(&hyperfine_cutoff_in_kHz)->default_value(10),
     "set cutoff scale for hyperfine field (kHz)")
    ("ms", po::value<int>(&ms)->default_value(1),
     "NV center spin state used with |0> for an effective two-level system (+/-1)")
    ("k_DD", po::value<uint>(&k_DD_int)->default_value(1),
     "resonance harmonic used in spin addressing (1 or 3)")
    ("static_Bz", po::value<double>(&static_Bz_in_gauss)->default_value(140.1,"140.1"),
     "strength of static magnetic field along the NV axis (gauss)")
    ("identity_time", po::value<double>(&identity_time)->default_value(1),
     "time of identity operation (s)")
    ("scale_factor", po::value<double>(&scale_factor)->default_value(10),
     "factor used to define different scales (i.e. if a << b, then a = b/scale_factor)")
    ("integration_factor", po::value<double>(&integration_factor)->default_value(10),
     "factor used to determine size of integration step size")
    ("no_nn" ,po::value<bool>(&no_nn)->default_value(false)->implicit_value(true),
     "turn off internuclear couplings")
    ;

  vector<uint> target_nuclei;
  bool target_pairs;
  double angle;
  double angle_over_pi;
  double target_polar;
  double target_pitch_over_pi;
  double target_azimuth;
  double target_azimuth_over_pi;
  double nv_polar;
  double nv_pitch_over_pi;
  double nv_azimuth;
  double nv_azimuth_over_pi;

  po::options_description addressing_options("Nuclear spin addressing options"
                                             " (all angles are in units of pi radians)",
                                             help_text_length);
  addressing_options.add_options()
    ("target", po::value<vector<uint>>(&target_nuclei)->multitoken(),
     "indices of nuclei to target")
    ("target_pairs",
     po::value<bool>(&target_pairs)->default_value(false)->implicit_value(true),
     "target only one nucleus in each larmor pair")
    ("angle", po::value<double>(&angle_over_pi)->default_value(1),
     "angle of operation to perform")
    ("target_pitch", po::value<double>(&target_pitch_over_pi)->default_value(0),
     "pitch (angle above x-y plane) of target rotation axis")
    ("target_azimuth", po::value<double>(&target_azimuth_over_pi)->default_value(0),
     "azimuthal angle of target rotation axis")
    ("nv_pitch", po::value<double>(&nv_pitch_over_pi)->default_value(0.5),
     "pitch (angle above x-y plane) of NV rotation axis")
    ("nv_azimuth", po::value<double>(&nv_azimuth_over_pi)->default_value(0),
     "azimuthal angle of NV rotation axis")
    ;

  uint coherence_bins;
  double scan_time;
  double scan_time_in_ms;
  double f_DD;
  uint angular_resolution;
  double signal_gB_factor;
  double max_f_factor;

  po::options_description scan_options("Coherence measurement options",help_text_length);
  scan_options.add_options()
    ("bins", po::value<uint>(&coherence_bins)->default_value(500),
     "number of bins in coherence measurements")
    ("scan_time", po::value<double>(&scan_time_in_ms)->default_value(1),
     "time for each measurement in the coherence scan (microseconds)")
    ("f_DD", po::value<double>(&f_DD)->default_value(0.06,"0.06"),
     "magnitude of fourier component used in coherence scanning")
    ("angular_resolution", po::value<uint>(&angular_resolution)->default_value(20),
     "angular resolution for magnetic field direction during signal measurement"
     " (1/[pi radians])")
    ("signal_gB_factor", po::value<double>(&signal_gB_factor)->default_value(100),
     "sets signal field strength to static_gBz/signal_gB_factor")
    ("max_f_factor", po::value<double>(&max_f_factor)->default_value(0.5),
     "factor scaling maximum value of f_DD in coherence signal measurement")
    ;

  string lattice_file;

  po::options_description file_io("File IO",help_text_length);
  file_io.add_options()
    ("lattice_file", po::value<string>(&lattice_file),
     "input file defining system configuration")
    ;

  bool print_lattice;
  bool target_info;

  po::options_description print_options("Available printing options",help_text_length);
  print_options.add_options()
    ("print_lattice",
     po::value<bool>(&print_lattice)->default_value(false)->implicit_value(true),
     "print C-13 lattice to a file")
    ("target_info",
     po::value<bool>(&target_info)->default_value(false)->implicit_value(true),
     "print information about target nuclei")
    ;

  po::options_description all("Allowed options");
  all.add(general);
  all.add(simulations);
  all.add(simulation_options);
  all.add(addressing_options);
  all.add(scan_options);
  all.add(file_io);
  all.add(print_options);

  // collect inputs
  po::variables_map inputs;
  po::store(parse_command_line(arg_num, arg_vec, all), inputs);
  po::notify(inputs);

  // if requested, print help text
  if (inputs.count("help")) {
    cout << all;
    return 0;
  }

  // -------------------------------------------------------------------------------------
  // Run a sanity check on inputs
  // -------------------------------------------------------------------------------------

  // determine whether certain options were used
  bool using_input_lattice = inputs.count("lattice_file");
  bool set_hyperfine_cutoff = !inputs["hyperfine_cutoff"].defaulted();
  bool set_c13_factor = !inputs["c13_factor"].defaulted();
  bool set_target_nuclei = inputs.count("target");

  // make sure we are either printing something, or performing a simulation
  if (!testing_mode) {
    const bool printing = print_lattice || target_info;
    if (!printing) {
      if (int(pair_search)
          + int(coherence_scan)
          + int(coherence_signal)
          + int(angular_coherence_signal)
          + int(rotation)
          + int(coupling)
          + int(iswap)
          + int(swap)
          + int(identity)
          + int(swap_nvst)
          + int(larmor_identity)
          + int(initialize)
          + int(initialize_x)
          + int(initialize_larmor)
          != 1) {
        cout << "Please choose one simulation to perform\n";
        return -1;
      }
    }
  }

  // check lattice options
  assert(!(print_lattice && using_input_lattice));
  assert(!(using_input_lattice && !fs::exists(lattice_file)));
  assert(!(using_input_lattice && set_hyperfine_cutoff));
  assert(!(using_input_lattice && set_c13_factor));

  // check targeting options
  assert(!(set_target_nuclei && target_pairs));

  // verify validity of other values
  assert(hyperfine_cutoff_in_kHz > 0);
  assert(c13_factor >= 0);
  assert(nuclear_isolation >= 0);
  assert(larmor_isolation_in_kHz >= 0);
  assert(min_hyperfine_xy_in_kHz >= 0);
  assert(max_cluster_size > 0);
  assert(ms == 1 || ms == -1);
  assert((k_DD_int == 1) || (k_DD_int == 3));
  assert(scale_factor > 1);
  assert(integration_factor > 1);

  if (coherence_scan) {
    assert(coherence_bins > 0);
    assert(scan_time_in_ms > 0);
  }
  assert(angular_resolution > 2);
  assert(signal_gB_factor > 1);
  assert(max_f_factor <= 1);

  // set some variables based on iputs
  c13_abundance = c13_factor*c13_natural_abundance;
  larmor_isolation = larmor_isolation_in_kHz * 1e3;
  min_hyperfine_xy = min_hyperfine_xy_in_kHz * 1e3;
  hyperfine_cutoff = hyperfine_cutoff_in_kHz * 1e3;
  k_DD = (k_DD_int == 1 ? first : third);
  static_Bz = static_Bz_in_gauss * gauss;
  scan_time = scan_time_in_ms * 1e-3;

  angle = angle_over_pi * pi;
  target_polar = pi/2 - target_pitch_over_pi * pi;
  target_azimuth = target_azimuth_over_pi * pi;
  nv_polar = pi/2 - nv_pitch_over_pi * pi;
  nv_azimuth = nv_azimuth_over_pi * pi;

  uniform_real_distribution<double> rnd(0.0,1.0); // uniform distribution on [0,1)
  mt19937_64 generator(seed); // use and seed the 64-bit Mersenne Twister 19937 generator

  // -------------------------------------------------------------------------------------
  // Construct lattice of nuclei
  // -------------------------------------------------------------------------------------

  vector<Vector3d> nuclei;

  if (!using_input_lattice) { // place nuclei at lattice sites

    // number of fcc cells to simulate out from the origin
    const int cell_radius
      = round(pow(abs(g_e*g_C13)/(4*pi*a0*a0*a0*hyperfine_cutoff),1.0/3));

    // set positions of nuclei at lattice sites
    for (uint b: {0,1}) {
      for (int l = -2*cell_radius; l <= 2*cell_radius; l++) {
        for (int m = -2*cell_radius; m <= 2*cell_radius; m++) {
          for (int n = -2*cell_radius; n <= 2*cell_radius; n++) {
            if (rnd(generator) < c13_abundance) { // check for C-13 isotopic abundance
              if (l != 0 || m != 0 || n != 0) { // don't place C-13 nucleus on NV sites
                const Vector3d pos = b*ao+l*a1+m*a2+n*a3;
                // only place C-13 nuclei with a hyperfine field strength above the cutoff
                if (hyperfine(pos).norm() > hyperfine_cutoff) {
                  nuclei.push_back(pos);
                }
              }
            }
          }
        }
      }
    }
    cout << "Placed " << nuclei.size() << " C-13 nuclei\n\n";
    if (nuclei.size() == 0) return 0;

    if (print_lattice) {
      for (uint i = 0; i < nuclei.size(); i++) {
        cout << nuclei.at(i)(0) << " "
             << nuclei.at(i)(1) << " "
             << nuclei.at(i)(2) << endl;
      }
    }

  } else { // if using_input_lattice, read in the lattice

    string line;
    ifstream lattice(lattice_file);

    // get C-13 positions
    double x,y,z;
    while (getline(lattice,line,' ')) {
      x = stod(line);
      getline(lattice,line,' ');
      y = stod(line);
      getline(lattice,line);
      z = stod(line);
      nuclei.push_back((Vector3d() << x,y,z).finished());
    }
    lattice.close();

    // assert that no C-13 nuclei lie at the NV lattice sites
    for (uint i = 0; i < nuclei.size(); i++) {
      if ((nuclei.at(i) == n_pos) || (nuclei.at(i) == e_pos)) {
        cout << "The input lattice places a C-13 nucleus at one of the NV lattice sites!"
             << endl;
        return -1;
      }
    }
  }

  // -------------------------------------------------------------------------------------
  // Characterize nuclei and targets
  // -------------------------------------------------------------------------------------

  // identify nuclei which cannot be addressed
  vector<uint> addressable_nuclei;
  vector<uint> unaddressable_nuclei;
  for (uint n = 0; n < nuclei.size(); n++) {
    if (can_address(nuclei,n)) {
      addressable_nuclei.push_back(n);
    } else {
      unaddressable_nuclei.push_back(n);
    }
  }
  if (unaddressable_nuclei.size() > 0) {
    cout << "The following nuclei cannot be addressed:";
    for (uint n: unaddressable_nuclei) {
      cout << " " << n;
    }
    cout << endl << endl;
  }

  // identify larmor pairs
  vector<uint> larmor_nuclei;
  vector<vector<uint>> larmor_pairs;
  for (uint i = 0; i < addressable_nuclei.size(); i++) {
    const uint n_i = addressable_nuclei.at(i);
    for (uint j = i+1; j < addressable_nuclei.size(); j++) {
      const uint n_j = addressable_nuclei.at(j);
      if (is_larmor_pair(nuclei,n_i,n_j)) {
        larmor_nuclei.push_back(n_i);
        larmor_nuclei.push_back(n_j);
        larmor_pairs.push_back({n_i,n_j});
        continue;
      }
    }
  }

  // if we are only doing a search for isolated larmor pairs, do it now
  if (pair_search) {
    cout << "Starting search for isolated larmor pairs..." << endl;
    cout << "idx1 idx2 hyperfine_xy" << endl;
    for (vector<uint> larmor_pair: larmor_pairs) {
      const uint ln_1 = larmor_pair.at(0);
      const uint ln_2 = larmor_pair.at(1);
      const double A_z = hyperfine_z(nuclei.at(ln_1));
      const double A_xy = hyperfine_xy(nuclei.at(ln_1));

      if (A_xy <= min_hyperfine_xy) continue;

      bool isolated_pair = (strong_field_coupling(nuclei, ln_1, ln_2)
                            <= nuclear_isolation);
      if (!isolated_pair) continue;

      for (uint n = 0; n < nuclei.size(); n++) {
        if (n == ln_1 || n == ln_2) continue;
        if (strong_field_coupling(nuclei, n, ln_1) > nuclear_isolation ||
            strong_field_coupling(nuclei, n, ln_2) > nuclear_isolation ||
            (abs(hyperfine_z(nuclei.at(n)) - A_z) < larmor_isolation)) {
          isolated_pair = false;
          break;
        }
      }
      if (!isolated_pair) continue;

      cout << ln_1 << " " << ln_2 << " " << A_xy << endl;
    }
    return 0;
  }

  // determine which nuclei to target
  if (!set_target_nuclei) {
    if (target_pairs) target_nuclei = larmor_nuclei;
    else target_nuclei = addressable_nuclei;

  } else { // if (set_target_nuclei)
    vector<uint> unaddressable_targets;
    for (uint t_i = 0; t_i < target_nuclei.size(); t_i++) {
      if (in_vector(target_nuclei.at(t_i),unaddressable_nuclei)) {
        unaddressable_targets.push_back(target_nuclei.at(t_i));
        target_nuclei.erase(target_nuclei.begin()+t_i);
        t_i--;
      }
    }
    if (unaddressable_targets.size() > 0) {
      cout << "(WARNING) Ignoring following target nuclei:";
      for (uint n: unaddressable_targets) {
        cout << " " << n;
      }
      cout << endl << endl;
    }
  }
  if (target_nuclei.size() == 0) {
    cout << "There are no target nuclei." << endl;
    return 0;
  }

  // identify targeted larmor pairs
  vector<vector<uint>> targeted_larmor_pairs;
  for (uint i = 0; i < target_nuclei.size(); i++) {
    const uint n_i = target_nuclei.at(i);
    for (uint j = i+1; j < target_nuclei.size(); j++) {
      const uint n_j = target_nuclei.at(j);
      if (is_larmor_pair(nuclei,n_i,n_j)) {
        targeted_larmor_pairs.push_back({n_i,n_j});
        continue;
      }
    }
  }

  if (swap_nvst || larmor_identity) {
    if (targeted_larmor_pairs.size() == 0) {
      if (larmor_pairs.size() == 0) {
        cout << "There are no larmor pairs in this system\n";
      } else {
        cout << "No larmor pairs are targeted\n";
      }
      return -1;
    }
  }

  // -------------------------------------------------------------------------------------
  // Cluster C-13 nuclei
  // -------------------------------------------------------------------------------------

  // unless we are performing a coherence scan, we will be grouping together clusters by
  //  the larmor frequencies of the nuclei, so first we check whether doing so is possible
  //  for the given min_cluster_size_cap
  const uint min_cluster_size_cap = smallest_possible_cluster_size(nuclei);
  const bool cluster_larmor_pairs =
    !(coherence_scan || (max_cluster_size < min_cluster_size_cap));
  if (!coherence_scan) {
    cout << "The minimum cluster size cap is " << min_cluster_size_cap << endl;
    if (!testing_mode && (max_cluster_size < min_cluster_size_cap)) return -1;
  }

  const vector<vector<uint>> clusters = cluster_nuclei(nuclei, max_cluster_size,
                                                       cluster_larmor_pairs);
  const double cluster_coupling = get_cluster_coupling(nuclei,clusters);

  cout << "Nuclei grouped into " << clusters.size() << " clusters"
       << " with a coupling factor of " << cluster_coupling << " Hz\n";

  // collect and print histogram of cluster sizes
  vector<uint> size_hist(largest_cluster_size(clusters));
  for (uint i = 0; i < clusters.size(); i++) {
    size_hist.at(clusters.at(i).size()-1) += 1;
  }
  cout << "Cluster size histogram:\n";
  for (uint i = 0; i < size_hist.size(); i++) {
    cout << "  " << i+1 << ": " << size_hist.at(i) << endl;
  }
  cout << endl;

  // initialize nv_system object
  const nv_system nv(nuclei, clusters, ms, g_C13*static_Bz, k_DD,
                     scale_factor, integration_factor, no_nn);

  // -------------------------------------------------------------------------------------
  // Print info about target nuclei
  // -------------------------------------------------------------------------------------

  if (target_info) {
    cout << "Target info:\n\n";
    for (uint n: target_nuclei) {
      cout << "index: " << n << endl
           << "position (nm): "
           << in_crystal_basis(nv.nuclei.at(n)).transpose() * a0/2 / nm << endl
           << "effective_larmor (kHz): " << effective_larmor(nv,n).norm() * 1e-3 << endl
           << "hyperfine (kHz): " << hyperfine(nv,n).norm() * 1e-3 << endl
           << "hyperfine_perp (kHz): " << hyperfine_perp(nv,n).norm() * 1e-3 << endl
           << endl;
    }
    for (uint i = 0; i < target_nuclei.size(); i++) {
      const Vector3d pos_i = nv.nuclei.at(target_nuclei.at(i));
      for (uint j = i + 1; j < target_nuclei.size(); j++) {
        const Vector3d pos_j = nv.nuclei.at(target_nuclei.at(j));
        cout << "indices: " << i << " " << j << endl
             << "displacement (nm): "
             << in_crystal_basis(pos_j-pos_i).transpose() * a0/2 / nm << endl
             << "coupling (Hz): " << coupling_strength(pos_i,pos_j) << endl
             << "strong field coupling (Hz): "
             <<  strong_field_coupling(pos_i,pos_j) << endl
             << endl;
      }
    }
  }

  // -------------------------------------------------------------------------------------
  // Coherence scan
  // -------------------------------------------------------------------------------------

  if (coherence_scan) {
    cout << endl
         << "Larmor and hyperfine frequency data:" << endl
         << "# format: w_larmor A_perp" << endl;

    double w_max = 0, w_min = DBL_MAX;
    for (uint i = 0; i < nv.nuclei.size(); i++) {
      const double A_perp = hyperfine_perp(nv,i).norm();
      const double w_larmor = effective_larmor(nv,i).norm();
      cout << i << " " << w_larmor << " " << A_perp << endl;

      if (w_larmor < w_min) w_min = w_larmor;
      if (w_larmor > w_max) w_max = w_larmor;
    }

    cout << endl
         << "Coherence scan results:" << endl
         << "# format: w_scan coherence" << endl;
    const double w_range = max(w_max - w_min, 1e5);
    const double w_start = max(w_min - w_range/10, 0.);
    const double w_end = w_max + w_range/10;
    for (uint i = 0; i < coherence_bins; i++) {
      const double w_scan = w_start + (i+0.5)*(w_end-w_start)/coherence_bins;
      const double coherence = coherence_measurement(nv, w_scan, f_DD, scan_time);
      cout << w_scan << " " << coherence << endl;
    }

  }

  // -------------------------------------------------------------------------------------
  // Coherence signal
  // -------------------------------------------------------------------------------------

  if (coherence_signal || angular_coherence_signal) {
    cout << "Coherence signal results:" << endl
         << "# k_DD: " << nv.k_DD << endl;
    cout << "# format: f_DD coherence(s)" << endl;

    for (uint target: target_nuclei) {
      // only perform this measurement once for larmor pairs
      vector<uint> targets = {target};
      for (uint nn: target_nuclei) {
        if (nn == target) continue;
        if (is_larmor_pair(nv,nn,target)) targets.push_back(nn);
      }
      if (targets.size() == 2 && targets.at(0) > targets.at(1)) continue;

      cout << "# targets:";
      for (uint nn: targets) cout << " " << nn;
      cout << endl;
      if (angular_coherence_signal) {
        cout << "# azimuths/pi:";
        for (uint nn: targets) {
          const double angular_pos = atan2(dot(nv.nuclei.at(nn),yhat),
                                           dot(nv.nuclei.at(nn),xhat));
          cout << " " << angular_pos/pi;
        }
        cout << endl;
      }

      // determine control field to use during coherence signal measurement
      const double w_signal = effective_larmor(nv,target).norm();
      const double phi_dec = [&]() -> double {
        if (!angular_coherence_signal) return angle;
        else return 0;
      }();
      const control_fields controls(nv.static_gBz/signal_gB_factor*xhat,
                                    w_signal, phi_dec);
      cout << "# phi_dec/pi: " << phi_dec/pi << endl;

      // fix coherence signal measurement time
      const double measurement_time = 4*pi / (axy_f_max(nv.k_DD) * max_f_factor
                                              * hyperfine_perp(nv,target).norm() / 4);
      cout << "# measurement_time (ms): " << measurement_time * 1e3 << endl;

      for (uint ii = 0; ii < coherence_bins; ii++) {
        const double f_DD = (ii+0.5)/coherence_bins * axy_f_max(nv.k_DD) * max_f_factor;

        if (!angular_coherence_signal) {
          const double coherence =
            coherence_measurement(nv, w_signal, f_DD, measurement_time, controls);
          cout << f_DD << " " << coherence << endl;

        } else {
          cout << f_DD;
          for (uint jj = 0; jj < angular_resolution; jj++) {
            const double phi_DD = (jj+0.5)/angular_resolution * pi;
            const double coherence =
              coherence_measurement(nv, w_signal, f_DD, measurement_time,
                                    controls, phi_DD);
            cout << " " << coherence;
          }
          cout << endl;
        }
      }
    }
  }

  // -------------------------------------------------------------------------------------
  // Individual addressing -- control
  // -------------------------------------------------------------------------------------

  if (rotation) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      const Vector3d target_axis = axis(target_polar,target_azimuth);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = rotate_target(nv, target, angle, target_axis, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Individual addressing -- NV coupling
  // -------------------------------------------------------------------------------------

  if (coupling) {
    cout << "target fidelity time pulses\n";
    const Vector3d nv_axis = axis(nv_polar, nv_azimuth);
    for (uint target: target_nuclei) {
      const Vector3d target_axis = axis(target_polar, target_azimuth);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = couple_target(nv, target, angle, nv_axis, target_axis, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // NV/nucleus iSWAP fidelity
  // -------------------------------------------------------------------------------------

  if (iswap) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = iSWAP(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // SWAP fidelity: NV electron spin and single nuclear spin
  // -------------------------------------------------------------------------------------

  if (swap) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = SWAP(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of identity operation on a single spin
  // -------------------------------------------------------------------------------------

  if (identity) {
    cout << "target fidelity\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = target_identity(nv, target, identity_time, exact);
      }
      const uint cluster_target = get_index_in_cluster(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {cluster_target}) << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // SWAP fidelity: NV electron spin and singlet-triplet subspace of two nuclear spins
  // -------------------------------------------------------------------------------------

  if (swap_nvst) {
    cout << "idx1 idx2 fidelity time pulses\n";
    for (vector<uint> idxs: targeted_larmor_pairs) {
      const uint idx1 = idxs.at(0);
      const uint idx2 = idxs.at(1);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = SWAP_NVST(nv, idx1, idx2, exact);
      }
      const uint ss_idx1 = get_index_in_subsystem(nv,idx1);
      const uint ss_idx2 = get_index_in_subsystem(nv,idx2);
      cout << idx1 << " " << idx2 << " "
           << protocol_fidelity(P, {0, ss_idx1, ss_idx2}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of identity operation on a larmor qubit
  // -------------------------------------------------------------------------------------

  if (larmor_identity) {
    cout << "idx1 idx2 fidelity\n";
    const bool targeting_pair = true;
    for (vector<uint> idxs: targeted_larmor_pairs) {
      const uint idx1 = idxs.at(0);
      const uint idx2 = idxs.at(1);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = target_identity(nv, idx1, identity_time, exact, targeting_pair);
      }
      const uint cluster_idx1 = get_index_in_cluster(nv,idx1);
      const uint cluster_idx2 = get_index_in_cluster(nv,idx2);
      cout << idx1 << " " << idx2 << " "
           << protocol_fidelity(P, {cluster_idx1, cluster_idx2}) << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of deterministically initializing a thermalized nucleus into |u> or |d>
  // -------------------------------------------------------------------------------------

  if (initialize) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = initialize_spin(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of probabalistically initializing a thermalized nucleus into |u> +/- |d>
  // -------------------------------------------------------------------------------------

  if (initialize_x) {
    cout << "target fidelity time pulses\n";
    for (uint target: target_nuclei) {
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = initialize_spin_X(nv, target, exact);
      }
      const uint subsystem_target = get_index_in_subsystem(nv, target);
      cout << target << " "
           << protocol_fidelity(P, {0, subsystem_target}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

  // -------------------------------------------------------------------------------------
  // Fidelity of probabalistically initializing a larmor pair from |dd> into |ud> +/- |du>
  // -------------------------------------------------------------------------------------

  if (initialize_larmor) {
    cout << "idx1 idx2 fidelity time pulses\n";
    for (vector<uint> idxs: targeted_larmor_pairs) {
      const uint idx1 = idxs.at(0);
      const uint idx2 = idxs.at(1);
      vector<protocol> P(2);
      for (bool exact : {true,false}) {
        P.at(exact) = initialize_larmor_qubit(nv, idx1, idx2, exact);
      }
      const uint ss_idx1 = get_index_in_subsystem(nv,idx1);
      const uint ss_idx2 = get_index_in_subsystem(nv,idx2);
      cout << idx1 << " " << idx2 << " "
           << protocol_fidelity(P, {0, ss_idx1, ss_idx2}) << " "
           << P.at(false).time << " "
           << P.at(false).pulses << endl;
    }
  }

}


//  Copyright John Maddock 2010.
//  Copyright Paul A. Bristow 2010.

//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#ifndef BOOST_STATS_INVERSE_GAUSSIAN_HPP
#define BOOST_STATS_INVERSE_GAUSSIAN_HPP

#ifdef _MSC_VER
#pragma warning(disable: 4512) // assignment operator could not be generated
#endif

// http://en.wikipedia.org/wiki/Normal-inverse_Gaussian_distribution
// http://mathworld.wolfram.com/InverseGaussianDistribution.html

// The normal-inverse Gaussian distribution
// also called the Wald distribution (some sources limit this to when mean = 1).

// It is the continuous probability distribution
// that is defined as the normal variance-mean mixture where the mixing density is the
// inverse Gaussian distribution. The tails of the distribution decrease more slowly
// than the normal distribution. It is therefore suitable to model phenomena
// where numerically large values are more probable than is the case for the normal distribution.

// The Inverse Gaussian distribution was first studied in relationship to Brownian motion.
// In 1956 M.C.K. Tweedie used the name 'Inverse Gaussian' because there is an inverse
// relationship between the time to cover a unit distance and distance covered in unit time.

// Examples are returns from financial assets and turbulent wind speeds.
// The normal-inverse Gaussian distributions form
// a subclass of the generalised hyperbolic distributions.

// See also

// http://en.wikipedia.org/wiki/Normal_distribution
// http://www.itl.nist.gov/div898/handbook/eda/section3/eda3661.htm
// Also:
// Weisstein, Eric W. "Normal Distribution."
// From MathWorld--A Wolfram Web Resource.
// http://mathworld.wolfram.com/NormalDistribution.html

// http://www.jstatsoft.org/v26/i04/paper General class of inverse Gaussian distributions.
// ig package - withdrawn but at http://cran.r-project.org/src/contrib/Archive/ig/

// http://www.stat.ucl.ac.be/ISdidactique/Rhelp/library/SuppDists/html/inverse_gaussian.html
// R package for dinverse_gaussian, ...

// http://www.statsci.org/s/inverse_gaussian.s  and http://www.statsci.org/s/inverse_gaussian.html

//#include <boost/math/distributions/fwd.hpp>
#include <boost/math/special_functions/erf.hpp> // for erf/erfc.
#include <boost/math/distributions/complement.hpp>
#include <boost/math/distributions/detail/common_error_handling.hpp>
#include <boost/math/distributions/normal.hpp>
#include <boost/math/distributions/gamma.hpp> // for gamma function
// using boost::math::gamma_p;

#include <boost/math/tools/tuple.hpp>
//using std::tr1::tuple;
//using std::tr1::make_tuple;
#include <boost/math/tools/roots.hpp>
//using boost::math::tools::newton_raphson_iterate;

#include <utility>

namespace boost{ namespace math{

template <class RealType = double, class Policy = policies::policy<> >
class inverse_gaussian_distribution
{
public:
   typedef RealType value_type;
   typedef Policy policy_type;

   inverse_gaussian_distribution(RealType mean = 1, RealType scale = 1)
      : m_mean(mean), m_scale(scale)
   { // Default is a 1,1 inverse_gaussian distribution.
     static const char* function = "boost::math::inverse_gaussian_distribution<%1%>::inverse_gaussian_distribution";

     RealType result;
     detail::check_scale(function, scale, &result, Policy());
     detail::check_location(function, mean, &result, Policy());
   }

   RealType mean()const
   { // alias for location.
      return m_mean; // aka mu
   }

   // Synonyms, provided to allow generic use of find_location and find_scale.
   RealType location()const
   { // location, aka mu.
      return m_mean;
   }
   RealType scale()const
   { // scale, aka lambda.
      return m_scale;
   }

   RealType shape()const
   { // shape, aka phi = lambda/mu.
      return m_scale / m_mean;
   }

private:
   //
   // Data members:
   //
   RealType m_mean;  // distribution mean or location, aka mu.
   RealType m_scale;    // distribution standard deviation or scale, aka lambda.
}; // class normal_distribution

typedef inverse_gaussian_distribution<double> inverse_gaussian;

template <class RealType, class Policy>
inline const std::pair<RealType, RealType> range(const inverse_gaussian_distribution<RealType, Policy>& /*dist*/)
{ // Range of permissible values for random variable x, zero to max.
   using boost::math::tools::max_value;
   return std::pair<RealType, RealType>(static_cast<RealType>(0.), max_value<RealType>()); // - to + max value.
}

template <class RealType, class Policy>
inline const std::pair<RealType, RealType> support(const inverse_gaussian_distribution<RealType, Policy>& /*dist*/)
{ // Range of supported values for random variable x, zero to max.
  // This is range where cdf rises from 0 to 1, and outside it, the pdf is zero.
   using boost::math::tools::max_value;
   return std::pair<RealType, RealType>(static_cast<RealType>(0.),  max_value<RealType>()); // - to + max value.
}

template <class RealType, class Policy>
inline RealType pdf(const inverse_gaussian_distribution<RealType, Policy>& dist, const RealType& x)
{ // Probability Density Function
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType scale = dist.scale();
   RealType mean = dist.mean();
   RealType result = 0;
   static const char* function = "boost::math::pdf(const inverse_gaussian_distribution<%1%>&, %1%)";
   if(false == detail::check_scale(function, scale, &result, Policy()))
   {
      return result;
   }
   if(false == detail::check_location(function, mean, &result, Policy()))
   {
      return result;
   }
   if(false == detail::check_positive_x(function, x, &result, Policy()))
   {
      return result;
   }

   if (x == 0)
   {
     return 0; // Convenient, even if not defined mathematically.
   }

   result =
     sqrt(scale / (constants::two_pi<RealType>() * x * x * x))
    * exp(-scale * (x - mean) * (x - mean) / (2 * x * mean * mean));
   return result;
} // pdf

template <class RealType, class Policy>
inline RealType cdf(const inverse_gaussian_distribution<RealType, Policy>& dist, const RealType& x)
{ // Cumulative Density Function.
   BOOST_MATH_STD_USING  // for ADL of std functions.

   RealType scale = dist.scale();
   RealType mean = dist.mean();
   static const char* function = "boost::math::cdf(const inverse_gaussian_distribution<%1%>&, %1%)";
   RealType result = 0;
   if(false == detail::check_scale(function, scale, &result, Policy()))
   {
      return result;
   }
   if(false == detail::check_location(function, mean, &result, Policy()))
   {
      return result;
   }
   if(false == detail::check_positive_x(function, x, &result, Policy()))
   {
     return result;
   }
   if (x == 0)
   {
     return 0; // Convenient, even if not defined mathematically.
   }
   // Problem with this formula for large scale > 1000 or small x,
   //result = 0.5 * (erf(sqrt(scale / x) * ((x / mean) - 1) / constants::root_two<RealType>(), Policy()) + 1)
   //  + exp(2 * scale / mean) / 2
   //  * (1 - erf(sqrt(scale / x) * (x / mean + 1) / constants::root_two<RealType>(), Policy()));
   // so use normal distribution version:
   // Wikipedia CDF equation http://en.wikipedia.org/wiki/Inverse_Gaussian_distribution.

   normal_distribution<RealType> n01;

   RealType n0 = sqrt(scale / x);
   n0 *= ((x / mean) -1);
   RealType n1 = cdf(n01, n0);
   RealType expfactor = exp(2 * scale / mean);
   RealType n3 = - sqrt(scale / x);
   n3 *= (x / mean) + 1;
   RealType n4 = cdf(n01, n3);
   result = n1 + expfactor * n4;
   return result;
} // cdf

template <class RealType, class Policy>
struct inverse_gaussian_quantile_functor
{

  inverse_gaussian_quantile_functor(const boost::math::inverse_gaussian_distribution<RealType, Policy> dist, RealType const& p)
    : distribution(dist), prob(p)
  {
  }
  boost::math::tuple<RealType, RealType> operator()(RealType const& x)
  {
    RealType c = cdf(distribution, x);
    RealType fx = c - prob;  // Difference cdf - value - to minimize.
    RealType dx = pdf(distribution, x); // pdf is 1st derivative.
    // return both function evaluation difference f(x) and 1st derivative f'(x).
    return boost::math::make_tuple(fx, dx);
  }
  private:
  const boost::math::inverse_gaussian_distribution<RealType, Policy> distribution;
  RealType prob;
};

template <class RealType, class Policy>
struct inverse_gaussian_quantile_complement_functor
{
    inverse_gaussian_quantile_complement_functor(const boost::math::inverse_gaussian_distribution<RealType, Policy> dist, RealType const& p)
    : distribution(dist), prob(p)
  {
  }
  boost::math::tuple<RealType, RealType> operator()(RealType const& x)
  {
    RealType c = cdf(complement(distribution, x));
    RealType fx = c - prob;  // Difference cdf - value - to minimize.
    RealType dx = -pdf(distribution, x); // pdf is 1st derivative.
    // return both function evaluation difference f(x) and 1st derivative f'(x).
    //return std::tr1::make_tuple(fx, dx); if available.
    return boost::math::make_tuple(fx, dx);
  }
  private:
  const boost::math::inverse_gaussian_distribution<RealType, Policy> distribution;
  RealType prob;
};

namespace detail
{
  template <class RealType>
  inline RealType guess_ig(RealType p, RealType mu = 1, RealType lambda = 1)
  { // guess at random variate value x for inverse gaussian quantile.
      BOOST_MATH_STD_USING
      using boost::math::policies::policy;
      // Error type.
      using boost::math::policies::overflow_error;
      // Action.
      using boost::math::policies::ignore_error;

      typedef policy<
        overflow_error<ignore_error> // Ignore overflow (return infinity)
      > no_overthrow_policy;

    RealType x; // result is guess at random variate value x.
    RealType phi = lambda / mu;
    if (phi > 2.)
    { // Big phi, so starting to look like normal Gaussian distribution.
      //    x=(qnorm(p,0,1,true,false) - 0.5 * sqrt(mu/lambda)) / sqrt(lambda/mu);
      // Whitmore, G.A. and Yalovsky, M.
      // A normalising logarithmic transformation for inverse Gaussian random variables,
      // Technometrics 20-2, 207-208 (1978), but using expression from
      // V Seshadri, Inverse Gaussian distribution (1998) ISBN 0387 98618 9, page 6.

      normal_distribution<RealType, no_overthrow_policy> n01;
      x = mu * exp(quantile(n01, p) / sqrt(phi) - 1/(2 * phi));
     }
    else
    { // phi < 2 so much less symmetrical with long tail,
      // so use gamma distribution as an approximation.
      using boost::math::gamma_distribution;

      // Define the distribution, using gamma_nooverflow:
      typedef gamma_distribution<RealType, no_overthrow_policy> gamma_nooverflow;

      gamma_distribution<RealType, no_overthrow_policy> g(static_cast<RealType>(0.5), static_cast<RealType>(1.));

      // gamma_nooverflow g(static_cast<RealType>(0.5), static_cast<RealType>(1.));
      // R qgamma(0.2, 0.5, 1)  0.0320923
      RealType qg = quantile(complement(g, p));
      //RealType qg1 = qgamma(1.- p, 0.5, 1.0, true, false);
      x = lambda / (qg * 2);
      //
      if (x > mu/2) // x > mu /2?
      { // x too large for the gamma approximation to work well.
        //x = qgamma(p, 0.5, 1.0); // qgamma(0.270614, 0.5, 1) = 0.05983807
        RealType q = quantile(g, p);
       // x = mu * exp(q * static_cast<RealType>(0.1));  // Said to improve at high p
       // x = mu * x;  // Improves at high p?
        x = mu * exp(q / sqrt(phi) - 1/(2 * phi));
      }
    }
    return x;
  }  // guess_ig
} // namespace detail

template <class RealType, class Policy>
inline RealType quantile(const inverse_gaussian_distribution<RealType, Policy>& dist, const RealType& p)
{
   BOOST_MATH_STD_USING  // for ADL of std functions.
   // No closed form exists so guess and use Newton Raphson iteration.

   RealType mean = dist.mean();
   RealType scale = dist.scale();
   static const char* function = "boost::math::quantile(const inverse_gaussian_distribution<%1%>&, %1%)";

   RealType result = 0;
   if(false == detail::check_scale(function, scale, &result, Policy()))
      return result;
   if(false == detail::check_location(function, mean, &result, Policy()))
      return result;
   if(false == detail::check_probability(function, p, &result, Policy()))
      return result;
   if (p == 0)
   {
     return 0; // Convenient, even if not defined mathematically?
   }
   if (p == 1)
   { // overflow
      result = policies::raise_overflow_error<RealType>(function,
        "probability parameter is 1, but must be < 1!", Policy());
      return result; // std::numeric_limits<RealType>::infinity();
   }

  RealType guess = detail::guess_ig(p, dist.mean(), dist.scale());
  using boost::math::tools::max_value;

  RealType min = 0.; // Minimum possible value is bottom of range of distribution.
  RealType max = max_value<RealType>();// Maximum possible value is top of range.
  // int digits = std::numeric_limits<RealType>::digits; // Maximum possible binary digits accuracy for type T.
  // digits used to control how accurate to try to make the result.
  // To allow user to control accuracy versus speed,
  int get_digits = policies::digits<RealType, Policy>();// get digits from policy,
  boost::uintmax_t m = policies::get_max_root_iterations<Policy>(); // and max iterations.
  using boost::math::tools::newton_raphson_iterate;
  result =
    newton_raphson_iterate(inverse_gaussian_quantile_functor<RealType, Policy>(dist, p), guess, min, max, get_digits, m);
   return result;
} // quantile

template <class RealType, class Policy>
inline RealType cdf(const complemented2_type<inverse_gaussian_distribution<RealType, Policy>, RealType>& c)
{
   BOOST_MATH_STD_USING  // for ADL of std functions.

   RealType scale = c.dist.scale();
   RealType mean = c.dist.mean();
   RealType x = c.param;
   static const char* function = "boost::math::cdf(const complement(inverse_gaussian_distribution<%1%>&), %1%)";
   // infinite arguments not supported.
   //if((boost::math::isinf)(x))
   //{
   //  if(x < 0) return 1; // cdf complement -infinity is unity.
   //  return 0; // cdf complement +infinity is zero
   //}
   // These produce MSVC 4127 warnings, so the above used instead.
   //if(std::numeric_limits<RealType>::has_infinity && x == std::numeric_limits<RealType>::infinity())
   //{ // cdf complement +infinity is zero.
   //  return 0;
   //}
   //if(std::numeric_limits<RealType>::has_infinity && x == -std::numeric_limits<RealType>::infinity())
   //{ // cdf complement -infinity is unity.
   //  return 1;
   //}
   RealType result = 0;
   if(false == detail::check_scale(function, scale, &result, Policy()))
      return result;
   if(false == detail::check_location(function, mean, &result, Policy()))
      return result;
   if(false == detail::check_positive_x(function, x, &result, Policy()))
      return result;

   normal_distribution<RealType> n01;
   RealType n0 = sqrt(scale / x);
   n0 *= ((x / mean) -1);
   RealType cdf_1 = cdf(complement(n01, n0));

   RealType expfactor = exp(2 * scale / mean);
   RealType n3 = - sqrt(scale / x);
   n3 *= (x / mean) + 1;

   //RealType n5 = +sqrt(scale/x) * ((x /mean) + 1); // note now positive sign.
   RealType n6 = cdf(complement(n01, +sqrt(scale/x) * ((x /mean) + 1)));
   // RealType n4 = cdf(n01, n3); // =
   result = cdf_1 - expfactor * n6;
   return result;
} // cdf complement

template <class RealType, class Policy>
inline RealType quantile(const complemented2_type<inverse_gaussian_distribution<RealType, Policy>, RealType>& c)
{
   BOOST_MATH_STD_USING  // for ADL of std functions

   RealType scale = c.dist.scale();
   RealType mean = c.dist.mean();
   static const char* function = "boost::math::quantile(const complement(inverse_gaussian_distribution<%1%>&), %1%)";
   RealType result = 0;
   if(false == detail::check_scale(function, scale, &result, Policy()))
      return result;
   if(false == detail::check_location(function, mean, &result, Policy()))
      return result;
   RealType q = c.param;
   if(false == detail::check_probability(function, q, &result, Policy()))
      return result;

   RealType guess = detail::guess_ig(q, mean, scale);
   // Complement.
   using boost::math::tools::max_value;

  RealType min = 0.; // Minimum possible value is bottom of range of distribution.
  RealType max = max_value<RealType>();// Maximum possible value is top of range.
  // int digits = std::numeric_limits<RealType>::digits; // Maximum possible binary digits accuracy for type T.
  // digits used to control how accurate to try to make the result.
  int get_digits = policies::digits<RealType, Policy>();
  boost::uintmax_t m = policies::get_max_root_iterations<Policy>();
  using boost::math::tools::newton_raphson_iterate;
  result =
    newton_raphson_iterate(inverse_gaussian_quantile_complement_functor<RealType, Policy>(c.dist, q), guess, min, max, get_digits, m);
   return result;
} // quantile

template <class RealType, class Policy>
inline RealType mean(const inverse_gaussian_distribution<RealType, Policy>& dist)
{ // aka mu
   return dist.mean();
}

template <class RealType, class Policy>
inline RealType scale(const inverse_gaussian_distribution<RealType, Policy>& dist)
{ // aka lambda
   return dist.scale();
}

template <class RealType, class Policy>
inline RealType shape(const inverse_gaussian_distribution<RealType, Policy>& dist)
{ // aka phi
   return dist.shape();
}

template <class RealType, class Policy>
inline RealType standard_deviation(const inverse_gaussian_distribution<RealType, Policy>& dist)
{
  BOOST_MATH_STD_USING
  RealType scale = dist.scale();
  RealType mean = dist.mean();
  RealType result = sqrt(mean * mean * mean / scale);
  return result;
}

template <class RealType, class Policy>
inline RealType mode(const inverse_gaussian_distribution<RealType, Policy>& dist)
{
  BOOST_MATH_STD_USING
  RealType scale = dist.scale();
  RealType  mean = dist.mean();
  RealType result = mean * (sqrt(1 + (9 * mean * mean)/(4 * scale * scale))
      - 3 * mean / (2 * scale));
  return result;
}

template <class RealType, class Policy>
inline RealType skewness(const inverse_gaussian_distribution<RealType, Policy>& dist)
{
  BOOST_MATH_STD_USING
  RealType scale = dist.scale();
  RealType  mean = dist.mean();
  RealType result = 3 * sqrt(mean/scale);
  return result;
}

template <class RealType, class Policy>
inline RealType kurtosis(const inverse_gaussian_distribution<RealType, Policy>& dist)
{
  RealType scale = dist.scale();
  RealType  mean = dist.mean();
  RealType result = 15 * mean / scale -3;
  return result;
}

template <class RealType, class Policy>
inline RealType kurtosis_excess(const inverse_gaussian_distribution<RealType, Policy>& dist)
{
  RealType scale = dist.scale();
  RealType  mean = dist.mean();
  RealType result = 15 * mean / scale;
  return result;
}

} // namespace math
} // namespace boost

// This include must be at the end, *after* the accessors
// for this distribution have been defined, in order to
// keep compilers that support two-phase lookup happy.
#include <boost/math/distributions/detail/derived_accessors.hpp>

#endif // BOOST_STATS_INVERSE_GAUSSIAN_HPP


/*-------------determinant.cpp------------------------------------------------//
*
* Purpose: Simple multiplication to help visualize eigenvectors
*
*   Notes: compile with:
*            g++ -I /usr/include/eigen3/ eigentest.cpp -Wno-ignored-attributes -Wno-deprecated-declarations
*
*-----------------------------------------------------------------------------*/

#include <iostream>
#include <Eigen/Core>
#include <vector>
#include <fstream>

using namespace Eigen;

int main(){

    // Opening file to writing
    std::ofstream output;
    output.open("out.dat");

    int size = 2, count=0;

    // setting up positions of all the desired points to transform
    double x[size*size*size], y[size*size*size], z[size*size*size],
           x2[size*size*size], y2[size*size*size], z2[size*size*size],
           xval = -1, yval = -1, zval = -1;
    MatrixXd Arr(3, 3);
    MatrixXd pos(3, 1);

    // Putting in "random" values for matrix
    Arr << 1, 2, 0,
           2, 1, 0,
           0, 0, -3;
/*
    Arr << 1, 0, 0,
           0, 1, 0,
           0, 0, 1;
*/
    std::cout << Arr << '\n';

    // Creating initial x, y, and z locations
    for (int i = 0; i < size; ++i){
        for (int j = 0; j < size; ++j){
            for (int k = 0; k < size; ++k){
                xval = -1 + 2 * ((double)i / ((double)size - 1));
                yval = -1 + 2 * ((double)j / ((double)size - 1));
                zval = -1 + 2 * ((double)k / ((double)size - 1));
                x[count] = xval;
                y[count] = yval;
                z[count] = zval;

                // Performing multiplication / setting up vector
                pos(0) = xval;
                pos(1) = yval;
                pos(2) = zval;
                pos = Arr * pos;

                // Storing values in x2, y2, and z2
                x2[count] = pos(0);
                y2[count] = pos(1);
                z2[count] = pos(2);
                count += 1;
            }
        }
    }
    count = 8;
    for (int i = 0; i < count; ++i){
        std::cout << x[i] << '\t' << y[i] << '\t' << z[i] << '\n';
        pos(0) = x[i];
        pos(1) = y[i];
        pos(2) = z[i];

        pos = Arr * pos;

        // Storing values in x2, y2, and z2
        x2[i] = pos(0);
        y2[i] = pos(1);
        z2[i] = pos(2);

    }
/*

    // Creating initial x, y, and z locations
    for (int i = 0; i < size; ++i){
        for (int j = 0; j < size; ++j){
            for (int k = 0; k < size; ++k){
                xval = -1 + 2 * ((double)i / (double)size);
                yval = -1 + 2 * ((double)j / (double)size);
                zval = -1 + 2 * ((double)k / (double)size);
                x[count] = xval;
                y[count] = yval;
                z[count] = zval;

                // Performing multiplication / setting up vector
                pos(0) = xval;
                pos(1) = yval;
                pos(2) = zval;
                pos = Arr * pos;

                // Storing values in x2, y2, and z2
                x2[count] = pos(0);
                y2[count] = pos(1);
                z2[count] = pos(2);
                count += 1;
            }
        }
    }

    // Writing to file in correct format
    count = 0;
    for (int i = 0; i < size; ++i){
        for (int j = 0; j < size; ++j){
            for (int k = 0; k < size; ++k){
                output << x[count] << '\t' << y[count]
                       << '\t' << z[count] << '\t' << count << '\n';
                count++;
            }
        }
    }

    output << '\n' << '\n';

    count = 0;
    for (int i = 0; i < size; ++i){
        for (int j = 0; j < size; ++j){
            for (int k = 0; k < size; ++k){
                output << x2[count] << '\t' << y2[count]
                       << '\t' << z2[count] << '\t' << count << '\n';
                count++;
            }
        }
    }
*/

    int frames = 60;
    count = 0;
    double xvel[size*size*size], yvel[size*size*size], zvel[size*size*size];
    double vid_time = 1.0;

    count = 8;

    for (int i = 0; i < count; ++i){
        xvel[i] = (x[i] - x2[i]) / vid_time;
        yvel[i] = (y[i] - y2[i]) / vid_time;
        zvel[i] = (z[i] - z2[i]) / vid_time;
    }

    for (int f = 0; f < frames; ++f){
        for (int i = 0; i < count; ++i){
            xval = x[i] + ((x2[i] - x[i])
                   * ((double)f / (double)frames));
            yval = y[i] + ((y2[i] - y[i])
                   * ((double)f / (double)frames));
            zval = z[i] + ((z2[i] - z[i])
                   * ((double)f / (double)frames));

            output << xval << '\t' << yval
                   << '\t' << zval << '\t'
                   << xvel[i] << '\t' << yvel[i] << '\t'
                   << zvel[i] << '\t' << 1 << '\t' << i << '\n';

        }
        output << '\n' << '\n';
    }

    /*
    for (int i = 0; i < size; ++i){
        for (int j = 0; j < size; ++j){
            for (int k = 0; k < size; ++k){
                xvel[count] = (x[count] - x2[count]) / vid_time;
                yvel[count] = (y[count] - y2[count]) / vid_time;
                zvel[count] = (z[count] - z2[count]) / vid_time;
                count = count + 1;
            }
        }
    }

    for (int f = 0; f < frames; ++f){
        count = 0;
        for (int i = 0; i < size; ++i){
            for (int j = 0; j < size; ++j){
                for (int k = 0; k < size; ++k){
                    xval = x[count] + ((x2[count] - x[count])
                           * ((double)f / (double)frames));
                    yval = y[count] + ((y2[count] - y[count])
                           * ((double)f / (double)frames));
                    zval = z[count] + ((z2[count] - z[count])
                           * ((double)f / (double)frames));

                    output << xval << '\t' << yval
                           << '\t' << zval << '\t'
                           << xvel[count] << '\t' << yvel[count] << '\t'
                           << zvel[count] << '\t' << 1 << '\t' << count << '\n';
                    count++;
                }
            }
        }

        output << '\n' << '\n';
    }
    */

    output.close();

}


/***********************************************************************
This file is part of the librjmcmc project source files.

Copyright : Institut Geographique National (2008-2012)
Contributors : Mathieu Brédif, Olivier Tournaire, Didier Boldo
email : librjmcmc@ign.fr

This software is a generic C++ library for stochastic optimization.

This software is governed by the CeCILL license under French law and
abiding by the rules of distribution of free software. You can use,
modify and/or redistribute the software under the terms of the CeCILL
license as circulated by CEA, CNRS and INRIA at the following URL
"http://www.cecill.info".

As a counterpart to the access to the source code and rights to copy,
modify and redistribute granted by the license, users are provided only
with a limited warranty and the software's author, the holder of the
economic rights, and the successive licensors have only limited liability.

In this respect, the user's attention is drawn to the risks associated
with loading, using, modifying and/or developing or reproducing the
software by the user in light of its specific status of free software,
that may mean that it is complicated to manipulate, and that also
therefore means that it is reserved for developers and experienced
professionals having in-depth computer knowledge. Users are therefore
encouraged to load and test the software's suitability as regards their
requirements in conditions enabling the security of their systems and/or
data to be ensured and, more generally, to use and operate it in the
same conditions as regards security.

The fact that you are presently reading this means that you have had
knowledge of the CeCILL license and that you accept its terms.

***********************************************************************/

#ifndef GEOMETRY_RECTANGLE_2_INTEGRATED_FLUX_HPP
#define GEOMETRY_RECTANGLE_2_INTEGRATED_FLUX_HPP

#include "rjmcmc/geometry/Rectangle_2.hpp"
#include "rjmcmc/geometry/Segment_2_iterator.hpp"
#include "rjmcmc/geometry/Iso_rectangle_2_Segment_2_clip.hpp"
#include <boost/gil/image.hpp>
#include <boost/gil/extension/matis/float_images.hpp>

template<typename K, typename OrientedImage, typename Segment, typename Functor>
void integrated_flux(const OrientedImage& v, const Segment& s0, Functor& f)
{
    typedef typename OrientedImage::view_t view_t;
    typedef typename view_t::xy_locator xy_locator;

    view_t view = v.view();
    int x0 = v.x0();
    int y0 = v.y0();

    int x1 = x0+view.width();
    int y1 = y0+view.height();
    Segment s(s0);
    typename K::Iso_rectangle_2 bbox(x0,y0,x1,y1);
    if(!clip(bbox,s)) return;

    geometry::Segment_2_iterator<K> it(s);

    xy_locator loc = view.xy_at(
            (typename xy_locator::x_coord_t) (it.x()-x0),
            (typename xy_locator::y_coord_t) (it.y()-y0)
            );

    boost::gil::point2<std::ptrdiff_t> movement[2] = {
        boost::gil::point2<std::ptrdiff_t> (it.step(0), 0),
        boost::gil::point2<std::ptrdiff_t> (0, it.step(1))
    };
    typename K::Vector_2 edge(s.target()-s.source());
    typename K::Vector_2 normal(edge.y(),-edge.x());
    for (; !it.end() ; ++it)
    {
        if(it.x()>=x0 && it.x()<x1 && it.y()>=y0 && it.y()<y1) {
            typename K::RT length = it.length();
            f(length,normal,loc);
        }
        loc += movement[it.axis()];
    }
}

class Flux_functor
{
public:
    template<typename RT, typename Vector, typename Locator>
    void operator()(RT length, const Vector& normal, const Locator& loc)
    {
        Vector g(boost::gil::at_c<0> (*loc), boost::gil::at_c<1> (*loc));
        double dot = length*geometry::to_double(normal * g);
        // res += std::abs(dot)-2.;
        m_value += std::max(0.,dot);
    }
    Flux_functor() : m_value(0) {}
    double value() const { return m_value; }
private:
    double m_value;
};

template<typename OrientedImage, typename K>
double integrated_flux(const OrientedImage& view, const geometry::Rectangle_2<K>& r)
{
    Flux_functor f;
    integrated_flux<K>(view,r.segment(0),f);
    integrated_flux<K>(view,r.segment(1),f);
    integrated_flux<K>(view,r.segment(2),f);
    integrated_flux<K>(view,r.segment(3),f);
    return f.value();
}

#endif // GEOMETRY_RECTANGLE_2_INTEGRATED_FLUX_HPP