#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/matrix.hpp>

#include <stdio.h>
#include <math.h>

#include <cmath>
#include <memory>
#include <vector>
#include <limits>
#include <random>
#include <cmath>

#include "statespace.hpp"

using namespace boost::numeric::ublas;

int udu(matrix<double> &M_mat, matrix<double> &U_mat, vector<double> &d_vec, unsigned int state_size) {

  unsigned int n = M_mat.size1();
  double *M = &(M_mat(0,0));
  double *U = &(U_mat(0,0));
  double *d = &(d_vec(0));

  // this method stolen from pykalman.sqrt.bierman by Daniel Duckworth (BSD license)
  /*Construct the UDU' decomposition of a positive, semidefinite matrix M

    Parameters
    ----------
    M : [n, n] array
        Matrix to factorize

    Returns
    -------
    UDU : UDU_decomposition of size n
        UDU' representation of M
  */

  // make M upper triangular
  for (unsigned i=0; i < state_size; ++i) {
    for (unsigned idx=i*n; idx < i*n+i; ++idx) {
      M[idx] = 0;
    }
  }

  // initialize U as the identity
  for (unsigned i=0; i < state_size; ++i) {
    for (unsigned idx=i*n; idx < i*n+state_size; ++idx) {
      U[idx] = 0;
    }
  }
  for (unsigned i=0; i < state_size; ++i) {
    U[i*n+i] = 1;
  }

  // note j cannot be unsigned, to handle state_size=0
  for (int j=state_size-1; j >= 1; --j) {
    d[j] = M[j*n+j];
    double alpha = 0.0;
    double beta = 0.0;
    if (d[j] > 0) {
      alpha = 1.0 / d[j];
    } else {
      if (fabs(d[j] > 1e-5) ) {
	printf("WARNING: nonpositive d[%d] %f in udu decomp\n", j, d[j]);
	return -1;
      }
      d[j] = 0.0;
      alpha = 0.0;
    }
    for (unsigned k=0; k < j; ++k) {
      beta = M[k*n+ j ];
      U[k*n+ j] = alpha * beta;
      for (unsigned kk=0; kk <= k; ++kk) {
	M[kk*n+ k] = M[kk*n+k] - beta * U[kk*n+j];
      }
    }
  }
  d[0] = M[0];

  if (d[0] < 0) {
    printf("ERROR: udu decomposition on non-posdef matrix with M(0,0)=%f\n", M[0]);
    return -1;
  }
  return 0;
}

FilterState::FilterState(int max_dimension, double eps_stationary) {
  this->eps_stationary = eps_stationary;
  this->at_fixed_point=false;
  this->alpha = 0;
  this->wasnan = false;

  this->obs_U = matrix<double>(max_dimension, max_dimension);
  this->pred_U = matrix<double>(max_dimension, max_dimension);
  this->tmp_U1 = matrix<double>(max_dimension, max_dimension);
  this->tmp_U2 = matrix<double>(max_dimension, max_dimension);

  // we only ever write to the upper triangle of this matrix, so
  // zeroing the full matrix now guarantees that it will always be zero.
  this->obs_U.clear();
  this->pred_U.clear();
  this->tmp_U1.clear();
  this->tmp_U2.clear();

  this->P = matrix<double>(max_dimension, max_dimension);

  this->obs_d = vector<double>(max_dimension);
  this->pred_d = vector<double>(max_dimension);
  this->obs_d.clear();
  this->pred_d.clear();

  this->gain = vector<double>(max_dimension);
  this->f = vector<double>(max_dimension);
  this->v = vector<double>(max_dimension);

  this->xk = vector<double>(max_dimension);
}

void print_vec(const vector<double> & v) {
  for(unsigned i=0; i < v.size(); ++i) {
    printf("%.3f ", v(i));
  }
  printf("\n");
}

void print_mat(const matrix<double> & m) {
  for(unsigned i=0; i < m.size1(); ++i) {
    for(unsigned j=0; j < m.size2(); ++j) {
      printf("%.3f ", m(i,j));
    }
    printf("\n");
  }
  printf("\n");
}

void write_vec(const char *fname, const vector<double> & v) {
  FILE *f = fopen(fname, "w");
  if (f == NULL)
    {
      printf("Error opening file!\n");
      exit(1);
    }
  for(unsigned i=0; i < v.size(); ++i) {
    fprintf(f, "%.12f ", v(i));
  }
  fprintf(f, "\n");
  fclose(f);

}

void write_mat_col(const char *fname, const matrix<double> & m) {

  FILE *f = fopen(fname, "w");
  if (f == NULL)
    {
      printf("Error opening file!\n");
      exit(1);
    }

  for(unsigned i=0; i < m.size1(); ++i) {
    for(unsigned j=0; j < m.size2(); ++j) {
      fprintf(f, "%.12f ", m(i,j));
    }
    fprintf(f, "\n");
  }
  fprintf(f, "\n");

  fclose(f);
}

void write_stuff(const char *item, unsigned int k, const vector<double> & v) {
  char fname[100];
  snprintf(fname, 100, "matrices/%s_c_%d.txt", item, k);
  write_vec(fname, v);
}

void write_stuff(const char *item, unsigned int k, const matrix<double> & m) {
  char fname[100];
  snprintf(fname, 100, "matrices/%s_c_%d.txt", item, k);
  write_mat_col(fname, m);
}

void print_mat_col(const matrix<double> & m) {
  for(unsigned i=0; i < m.size1(); ++i) {
    for(unsigned j=0; j < m.size2(); ++j) {
      printf("%.3f ", m(i,j));
    }
    printf("\n");
  }
  printf("\n");
}

#include <cblas.h>
void compute_explicit_cov_atlas(FilterState &cache,
				matrix<double> &U,
				vector<double> &d,
				int prev_state_size) {

  matrix<double> &mtmp = cache.tmp_U2;
  matrix<double> &P = cache.P;
  unsigned state_size = cache.state_size;

  if (prev_state_size == -1) {
    prev_state_size = state_size;
  }

  // U has dimension (ss x pss)
  for (unsigned i=0; i < state_size; ++i) {
    for (unsigned j=0; j < prev_state_size; ++j) {
      mtmp(i,j) = U(i,j) * d(j);
    }
  }

  cblas_dgemm(CblasRowMajor,
	      CblasNoTrans,
	      CblasTrans,
	      state_size,  state_size, prev_state_size, 1.0,
	      &(mtmp(0,0)), mtmp.size2(),
	      &(U(0,0)), U.size2(),
	      0.0, &(P(0,0)), P.size2());

}

/* void compute_explicit_cov(FilterState &cache,
			  matrix<double> &U_tmp,
			  vector<double> &d_tmp,
			  int prev_state_size) {

  // this is the original version of this function in ublas.
  // it's slow and superceded by the atlas version above.

  matrix<double> &mtmp = cache.tmp_U2;
  matrix<double> &P = cache.P;
  unsigned state_size = cache.state_size;

  if (prev_state_size == -1) {
    prev_state_size = state_size;
  }

  // construct the cov matrix
  for (unsigned i=0; i < prev_state_size; ++i) {
    subrange(mtmp, 0, state_size, i, i+1) = subrange(U_tmp, 0, state_size, i, i+1);
    subrange(mtmp, 0, state_size, i, i+1) *= d_tmp(i);
  }
  noalias(subrange(P, 0, state_size, 0, state_size))		\
    = prod(subrange(mtmp, 0, state_size, 0, prev_state_size),
	   trans(subrange(U_tmp, 0, state_size, 0, prev_state_size)));
} */

void FilterState::init_priors(StateSpaceModel &ssm) {
  this->pred_d.clear();
  this->xk.clear();

  this->state_size = ssm.prior_mean(&(this->xk(0)));
  this->state_size = ssm.prior_vars(&(this->pred_d(0)));

  this->pred_U.clear();
  for (unsigned i=0; i < ssm.max_dimension; ++i) {
    this->pred_U(i,i) = 1;
  }

  return;
}

void FilterState::init_incremental_state(StateSpaceModel &ssm, int k) {
  /*
    Initialize the state to a high-variance prior on the state variables
    active at some (nonzero) step k. This allows us to start filtering
    somewhere other than the beginning of the sequence.
   */

  this->pred_d.clear();
  this->xk.clear();
  this->state_size = ssm.state_size_at_timestep(k);

  for (unsigned i=0; i < ssm.max_dimension; ++i) {
    this->pred_d(i) = 1e8;
  }

  this->pred_U.clear();
  for (unsigned i=0; i < ssm.max_dimension; ++i) {
    this->pred_U(i,i) = 1;
  }

  return;
}

double kalman_observe_sqrt(StateSpaceModel &ssm, FilterState &cache, int k, double zk) {

  // if needed, use this for isnan: #include <boost/math/special_functions/fpclassify.hpp>
  if (std::isnan(zk)) {
    // printf("declaring nan at timestep %d, val %f\n", k, zk);
    if (!cache.wasnan && cache.at_fixed_point) {
       cache.at_fixed_point = false;
    }
    cache.wasnan = true;

    // no observation, so the predicted state becomes the "observed" state
    cache.obs_d = cache.pred_d;
    cache.obs_U = cache.pred_U;

    return 0;
  } else {
    if (cache.wasnan && cache.at_fixed_point) {
       cache.at_fixed_point = false;
    }
    cache.wasnan = false;
  }

  double alpha = cache.alpha;
  unsigned int state_size = cache.state_size;
  if (!cache.at_fixed_point || !ssm.stationary(k)) {
    cache.at_fixed_point = false;

    matrix<double> &U_old = cache.pred_U;
    vector<double> &d_old = cache.pred_d;

    matrix<double> &U = cache.obs_U;
    vector<double> &d = cache.obs_d;
    vector<double> &K = cache.gain;
    vector<double> &f = cache.f;
    vector<double> &v = cache.v;
    double r = ssm.observation_noise(k);

    K.clear();

    ssm.apply_observation_matrix(U_old, 0,
				 k, &(f(0)), &(v(0)), state_size);

    for (unsigned i=0; i < state_size; ++i) {
      v(i) = d_old(i)*f(i);
    }

    D(write_stuff("U_obs_old", k, U_old);)
      D(write_stuff("d_obs_old", k, d_old);)
      D(write_stuff("f", k, f);)
      D(write_stuff("v", k, v);)

    alpha = r + v(0)*f(0);
    if (alpha > 1e-20) {
       d(0) = d_old(0) * r/alpha;
    } else {
      //printf("step %d correcting initial alpha from %f to 1e-20\n", k, alpha);
      d(0) = d_old(0);
      alpha = 1e-20;
    }
    // printf("   alpha C: %f\n", alpha);
    K(0)=v(0);

    U(0, 0) = U_old(0,0);
    for (unsigned j=1; j < state_size; ++j) {
      double old_alpha = alpha;
      alpha += v(j)*f(j);
      //printf("   alpha C %d: %f\n", j, alpha);
      if (alpha > 1e-20) {
	d(j) = d_old(j) * (old_alpha/alpha);
	//printf("d = %f * %f = %f\n" , d_old(j), old_alpha/alpha, d(j));
      } else {
	//printf("step %d correcting alpha from %f to 1e-20\n", k, alpha);
	d(j) = d_old(j);
	alpha = 1e-20;
      }

      for (unsigned i=0; i < state_size; ++i) {
	U(i, j) = U_old(i,j) - (K(i)/old_alpha)*f(j);
        K(i) += v(j) * U_old(i,j);
      }

    }
    cache.alpha = alpha;

      D(write_stuff("U_obs", k, U);)
      D(write_stuff("d_obs", k, d);)

  }

  // given the Kalman gain from the covariance update, compute
  // the updated mean vector.
  vector<double> &xk = cache.xk;
  double pred_z = ssm.apply_observation_matrix(&(xk(0)), k) + ssm.observation_bias(k);
  cache.pred_z = pred_z;
  double yk = zk - pred_z;
  for (unsigned i=0; i < state_size; ++i) {
    xk(i) += cache.gain(i) * yk/alpha;
  }

  // also compute log marginal likelihood for this observation
  double step_ell = -.5 * log(2*PI*alpha) - .5 * yk*yk / alpha;

  //printf("step %d (C) pred %.4f alpha %.4f z %.4f y %.4f ell %.4f\n", k, pred_z, alpha, zk, yk, step_ell);

  if (std::isnan(step_ell)) {
    printf("step %d (C) pred %.4f alpha %.4f z %.4f y %.4f ell %.4f\n", k, pred_z, alpha, zk, yk, step_ell);
    //print_vec(cache.obs_d);
    //printf("\n");
    //print_vec(cache.xk);
    //exit(-1);
    step_ell = -INFINITY;
  }

  return step_ell;
}

int kalman_predict_sqrt(StateSpaceModel &ssm, FilterState &cache, int k, bool force_P) {

  unsigned int prev_state_size = cache.state_size;

  vector<double> &tmp = cache.f;

  unsigned int state_size = ssm.apply_transition_matrix( &(cache.xk(0)), k,  &(tmp(0)));
  cache.state_size = state_size;

  vector<double> &xk = cache.xk;
  subrange(xk, 0, state_size) = subrange(tmp, 0, state_size);
  ssm.transition_bias(k, &(xk(0)));

  D(write_stuff("xk_posttransit", k, xk);)

  if (cache.at_fixed_point and ssm.stationary(k)) {
    return 0;
  }

  cache.at_fixed_point = false;

  matrix<double> &U_old = cache.obs_U;
  vector<double> &d_old = cache.obs_d;

  matrix<double> &U_tmp = cache.tmp_U1;
  vector<double> &d_tmp = cache.v;

  // get transition noise into temporary storage
  ssm.transition_noise_diag(k, &(tmp(0)));

  /* pushing the covariance P through the transition model F yields
     FPF'. In a factored representation, this is FUDU'F', so we just need
     to compute FU. */
  unsigned int min_size = std::min(prev_state_size, state_size);

  D(write_stuff("U_pretransit", k, U_old);)

  // COMMENTED OUT: this loop is equivalent to the matrix-valued transition call
  // directly below. I'm leaving it in for debugging and to run speed comparisons.

    /*
  for (int i=0; i < min_size; ++i) {
    // THIS ONLY WORKS IF U_old is in column-major order
    ssm.apply_transition_matrix(&(column(U_old, i)(0) ), k, &(d_tmp(0)) );
    noalias(column(U_tmp, i)) = d_tmp;
    }*/
  subrange(d_tmp, 0, prev_state_size) = subrange(d_old, 0, prev_state_size);
  ssm.apply_transition_matrix(U_old, 0, k, U_tmp, 0, prev_state_size);
  for (unsigned i=prev_state_size; i < state_size; ++i) {
    d_tmp(i) = 0;
    for (unsigned j=0; j < state_size; ++j) {
      U_tmp(j, i) = 0;
    }
  }

    D(write_stuff("d_posttransit", k, d_tmp);)
    D(write_stuff("U_posttransit", k, U_tmp);)

  // if there is transition noise, do the expensive reconstruction/factoring step
  if (force_P || state_size != prev_state_size || norm_2(subrange(tmp, 0, state_size)) > 0) {

    compute_explicit_cov_atlas(cache, U_tmp, d_tmp, prev_state_size);
    // add transition noise
    matrix<double> &P = cache.P;

    D(write_stuff("P_prenoise", k, P);)

    for (unsigned i=0; i < state_size; ++i) {
      P(i,i) += tmp(i);
    }
    //printf("noise at time %d ", k);
    //print_vec(tmp);

      D(write_stuff("P", k, P);)

    // printf("step %d state size %d\n", k, state_size);

    // udu overwrites the cov matrix, so we need to
    // save it if we're going to explicitly use it
    // later on.
    matrix<double> & mtmp = cache.tmp_U2;
    if (force_P) {
      mtmp = P;
    }

    // get the new factored representation
    int err = udu(P, U_tmp, d_tmp, state_size);
    if (err != 0) {
      return err;
    }

      D(write_stuff("d_decomp", k, d_tmp);)
      D(write_stuff("U_decomp", k, U_tmp);)

    if (force_P) {
      P = mtmp;
    }

  }

  // if our factored representation is (almost) the same as the previous invocation,
  // we've reached a stationary state
  matrix<double> &U_cached = cache.pred_U;
  vector<double> &d_cached = cache.pred_d;
  if (ssm.stationary(k)) {
    if (k > 0 && ssm.stationary(k-1)) {
      bool potential_fixed_point = true;
      for (unsigned i=0; i < state_size; ++i) {
	if (std::abs(d_tmp(i) - d_cached(i)) > cache.eps_stationary) {
	  potential_fixed_point=false;
	  break;
	}
	for (unsigned j=0; j < state_size; ++j) {
	  if (std::abs(U_tmp(i,j) - U_cached(i,j)) > cache.eps_stationary) {
	    potential_fixed_point=false;
	    break;
	  }
	}
	if (!potential_fixed_point) {
	  break;
	}
      }
      if (potential_fixed_point) {
	cache.at_fixed_point = true;
      }
    }
  }
  if (!cache.at_fixed_point) {
    subrange(U_cached, 0, state_size, 0, state_size) = subrange(U_tmp, 0, state_size, 0, state_size);
    subrange(d_cached, 0, state_size) = subrange(d_tmp, 0, state_size);
  }
  return 0;
}

double filter_likelihood(StateSpaceModel &ssm, const vector<double> &z) {
  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);
  unsigned int N = z.size();
  double ell = 0;

  D(write_stuff("U_prior", 0, cache.pred_U);)
  D(write_stuff("d_prior", 0, cache.pred_d);)
  D(write_stuff("xk_prior", 0, cache.xk);)

    if (N == 0) {
      return ell;
    }

  double step_ell = kalman_observe_sqrt(ssm, cache, 0, z(0));
  if (std::isinf(step_ell)) {
    return step_ell;
  }
  ell += step_ell;

  D(write_stuff("U_post_obs", 0, cache.obs_U);)
    D(write_stuff("d_post_obs", 0, cache.obs_d);)
    D(write_stuff("xk_post_obs", 0, cache.xk);)

  for (unsigned k=1; k < N; ++k) {
    int err = kalman_predict_sqrt(ssm, cache, k, false);

    if (err != 0) {
      return -INFINITY;
    }

    D(write_stuff("U_post_predict", k, cache.pred_U);)
    D(write_stuff("d_post_predict", k, cache.pred_d);)
    D(write_stuff("xk_post_predict", k, cache.xk);)

    double step_ell = kalman_observe_sqrt(ssm, cache, k, z(k));
    if (std::isinf(step_ell)) {
      return step_ell;
    }
    ell += step_ell;

      D(write_stuff("U_post_obs", k, cache.obs_U);)
      D(write_stuff("d_post_obs", k, cache.obs_d);)
      D(write_stuff("xk_post_obs", k, cache.xk);)

  }
  return ell;
}

double filter_incremental(StateSpaceModel &ssm,
			  const vector<double> &z,
			  double * ells,
			  int filter_start_idx,
			  int incr_start_idx,
			  int incr_end_idx,
			  int update_ells,
			  double step_ell_tol,
			  int * steps_processed,
			  int * errcode) {

  unsigned int N = z.size();
  if ((filter_start_idx < 0) || (incr_start_idx < filter_start_idx) || (incr_start_idx >= N)) {
    printf("fatal: filter_incremental received bad indices (%d, %d, %d) for signal of size %d\n",
	   filter_start_idx, incr_start_idx, incr_end_idx, N);
  }

  FilterState cache(ssm.max_dimension, 1e-10);
  if (filter_start_idx == 0) {
    cache.init_priors(ssm);
  } else {
    cache.init_incremental_state(ssm, filter_start_idx);
  }

  double total_discrepancy = 0;
  double step_ell_discrepancy;
  double step_ell;

  //printf("update ells %d\n", update_ells);

  // do some initial filtering steps to "warm up" the hidden state
  for (unsigned k=filter_start_idx; k < incr_start_idx; ++k) {
    step_ell = kalman_observe_sqrt(ssm, cache, k, z(k));
    if (std::isinf(step_ell)) {
      *errcode = ERR_INCR_NUMERIC;
      return -INFINITY;
    }

    int err = kalman_predict_sqrt(ssm, cache, k+1, false);
    if (err != 0) {
      *errcode = ERR_INCR_NUMERIC;
      return -INFINITY;
    }

    step_ell_discrepancy= step_ell - ells[k];

    if (k == incr_start_idx-1) {
      // check that our likelihood calculations are now in sync with the
      // previous calculation.
      if ( fabs(step_ell_discrepancy) > step_ell_tol ) {
	// printf("warmup %d discrepancy %.10f\n", k, step_ell_discrepancy);
         *errcode = ERR_INCR_INIT;
         return -INFINITY;
      }
    }
  }

  // now start the "normal" filtering process with an observation step
  step_ell = kalman_observe_sqrt(ssm, cache, incr_start_idx, z(incr_start_idx));
  if (std::isinf(step_ell)) {
    *errcode = ERR_INCR_NUMERIC;
    return step_ell;
  }
  step_ell_discrepancy = step_ell - ells[incr_start_idx];
  if (update_ells) {
    ells[incr_start_idx] = step_ell;
  }
  total_discrepancy += step_ell_discrepancy;

  //printf("first obs %d discrepancy %.2f\n", incr_start_idx, step_ell_discrepancy);

  *steps_processed = N - filter_start_idx;
  int match_counter = 0;
  for (unsigned k=incr_start_idx+1; k < N; ++k) {
    int err = kalman_predict_sqrt(ssm, cache, k, false);

    if (err != 0) {
      *errcode = ERR_INCR_NUMERIC;
      return -INFINITY;
    }

    step_ell = kalman_observe_sqrt(ssm, cache, k, z(k));
    if (std::isinf(step_ell)) {
      *errcode = ERR_INCR_NUMERIC;
      return -INFINITY;
    }
    step_ell_discrepancy = step_ell - ells[k];
    if (update_ells) {
      ells[k] = step_ell;
    }
    total_discrepancy += step_ell_discrepancy;
    //printf("step %d discrepancy %.2f\n", k, step_ell_discrepancy);

    // quit filtering if we're past the end idx and the likelihoods seem to
    // have reverted to their previous values
    if ((k > incr_end_idx) && (fabs(step_ell_discrepancy) < step_ell_tol)) {
      *steps_processed = k - filter_start_idx;
      match_counter += 1;
      if (match_counter >= 5) {
	break;
      }
    }
  }
  return total_discrepancy;
}

void step_obs_likelihoods(StateSpaceModel &ssm, const vector<double> &z,
			  vector<double> & ells,
			  vector<double> & preds,
			  vector<double> & alphas) {
  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);
  unsigned int N = z.size();

  if (N == 0) {
    return;
  }

  ells(0) = kalman_observe_sqrt(ssm, cache, 0, z(0));
  preds(0) = cache.pred_z;
  alphas(0) = cache.alpha;
  for (unsigned k=1; k < N; ++k) {
    kalman_predict_sqrt(ssm, cache, k, false);
    ells(k) = kalman_observe_sqrt(ssm, cache, k, z(k));
    preds(k) = cache.pred_z;
    alphas(k) = cache.alpha;
    //printf("got ell %f at step %d\n", ells(k), k);
  }
}

void mean_obs(StateSpaceModel &ssm, vector<double> & result) {
  vector<double> x(ssm.max_dimension);
  vector<double> x2(ssm.max_dimension);
  x.clear();
  x2.clear();

  ssm.prior_mean(&(x(0)));

  for (unsigned k = 0; k < result.size(); ++k) {
    result[k] = ssm.apply_observation_matrix(&(x(0)), k);
    result[k] += ssm.observation_bias(k);

    if (k+1 < result.size()) {
      ssm.apply_transition_matrix(&(x(0)), k+1, &(x2(0)));
      ssm.transition_bias(k+1, &(x2(0)));
    }
    x = x2; // this copy is unnecessary, we could swap
            // pointers instead, but it doesn't
            // matter cause this method is never the
            // performance bottleneck.
  }
}

void obs_var(StateSpaceModel &ssm, vector<double> & result) {
  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);
  compute_explicit_cov_atlas(cache, cache.pred_U, cache.pred_d, -1);

  matrix<double> P = cache.P;

  for (unsigned k = 0; k < result.size(); ++k) {
    ssm.apply_observation_matrix(P, 0,
    				 k, &(cache.f(0)), &(cache.v(0)), cache.f.size());
    result(k) = ssm.apply_observation_matrix(&(cache.f(0)), k);
    result(k) += ssm.observation_noise(k);
    if (k+1 < result.size()) {

      ssm.apply_transition_matrix(P, 0, k+1, cache.obs_U, 0, ssm.max_dimension);
      cache.obs_U = trans(cache.obs_U);
      ssm.apply_transition_matrix(cache.obs_U, 0, k+1, P, 0, ssm.max_dimension);

      vector<double> &tmp = cache.f;
      ssm.transition_noise_diag(k+1, &(tmp(0)));
      for (unsigned i=0; i < ssm.max_dimension; ++i) {
	P(i,i) += tmp(i);
      }
    }
  }
}

void prior_sample(StateSpaceModel &ssm, vector<double> & result, unsigned long seed) {
  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);

  std::mt19937 gen(seed);
  std::normal_distribution<double> randn(0,1);
  /* auto randn = std::bind(std::normal_distribution<double>(0,1),
			     ); */

  printf("seed with %lu, first randn %f\n", seed, randn(gen));

  // sample initial state from the prior
  vector<double> &d = cache.pred_d;
  for (unsigned i=0; i < ssm.max_dimension; ++i) {
    cache.xk(i) += randn(gen) * sqrt(d(i));
  }

  unsigned k = 0;
  result(k) = ssm.apply_observation_matrix(&(cache.xk(0)), k);
  result(k) += ssm.observation_bias(k);

  for (k=1; k < result.size(); ++k) {
    cache.obs_d = cache.pred_d;
    cache.obs_U = cache.pred_U;

    vector<double> &tmp = cache.f;
    unsigned int state_size = ssm.apply_transition_matrix( &(cache.xk(0)), k,  &(tmp(0)));
    vector<double> &xk = cache.xk;
    subrange(xk, 0, state_size) = subrange(tmp, 0, state_size);
    ssm.transition_bias(k, &(xk(0)));

    ssm.transition_noise_diag(k, &(tmp(0)));
    for (unsigned i=0; i < state_size; ++i) {
      cache.xk(i) += randn(gen) * sqrt(tmp(i));
    }

    result(k) = ssm.apply_observation_matrix(&(cache.xk(0)), k);
    result(k) += ssm.observation_bias(k);
    result(k) += randn(gen) * sqrt(ssm.observation_noise(k));

  }
}

double all_filtered_cssm_coef_marginals(TransientCombinedSSM &ssm,
				      const vector<double> &z,
				      vector<double> & step_ells,
				      std::vector<vector<double> > & cmeans,
				      std::vector<vector<double> > & cvars) {
  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);
  ssm.init_coef_priors(cmeans, cvars);

  unsigned int N = z.size();
  double ell = 0;

  if (N == 0) {
    return ell;
  }

  step_ells[0] = kalman_observe_sqrt(ssm, cache, 0, z(0));
  ell += step_ells[0];
  compute_explicit_cov_atlas(cache, cache.obs_U, cache.obs_d, -1);
  ssm.extract_all_coefs(cache, 0, cmeans, cvars);
  for (unsigned k=1; k < N; ++k) {
    kalman_predict_sqrt(ssm, cache, k, false);
    step_ells[k] = kalman_observe_sqrt(ssm, cache, k, z(k));
    ell += step_ells[k];

    compute_explicit_cov_atlas(cache, cache.obs_U, cache.obs_d, -1);
    ssm.extract_all_coefs(cache, k, cmeans, cvars);
  }
  return ell;
}

double tssm_component_means(TransientCombinedSSM &ssm,
			    const vector<double> &z,
			    std::vector<vector<double> > & means) {

  if (means.size() != ssm.n_ssms) {
    printf("component_means() needs exactly one vector for each component SSM\n");
    exit(-1);
  }

  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);

  unsigned int N = z.size();
  double ell = 0;

  if (N == 0) {
    return ell;
  }

  ell += kalman_observe_sqrt(ssm, cache, 0, z(0));
  ssm.extract_component_means(&(cache.xk(0)), 0, means);
  for (unsigned k=1; k < N; ++k) {
    kalman_predict_sqrt(ssm, cache, k, false);
    ell += kalman_observe_sqrt(ssm, cache, k, z(k));
    ssm.extract_component_means(&(cache.xk(0)), k, means);
  }
  return ell;
}

double tssm_component_vars(TransientCombinedSSM &ssm,
			  const vector<double> &z,
			  std::vector<vector<double> > & vars) {

  if (vars.size() != ssm.n_ssms) {
    printf("component_vars() needs exactly one vector for each component SSM\n");
    exit(-1);
  }

  FilterState cache(ssm.max_dimension, 1e-10);
  cache.init_priors(ssm);
  compute_explicit_cov_atlas(cache, cache.pred_U, cache.pred_d, -1);

  unsigned int N = z.size();
  double ell = 0;

  if (N == 0) {
    return ell;
  }

  ell += kalman_observe_sqrt(ssm, cache, 0, z(0));
  compute_explicit_cov_atlas(cache, cache.obs_U, cache.obs_d, -1);
  ssm.extract_component_vars(cache.P, cache.tmp_U2, 0, vars);
  for (unsigned k=1; k < N; ++k) {
    kalman_predict_sqrt(ssm, cache, k, false);
    ell += kalman_observe_sqrt(ssm, cache, k, z(k));
    compute_explicit_cov_atlas(cache, cache.obs_U, cache.obs_d, -1);
    ssm.extract_component_vars(cache.P, cache.tmp_U2, k, vars);
  }
  return ell;
}


// Boost.Geometry (aka GGL, Generic Geometry Library)

// Copyright (c) 2009-2012 Mateusz Loskot, London, UK.
// Copyright (c) 2008-2012 Barend Gehrels, Amsterdam, the Netherlands.
// Copyright (c) 2008-2012 Bruno Lalande, Paris, France.

// Use, modification and distribution is 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_GEOMETRY_ARITHMETIC_DETERMINANT_HPP
#define BOOST_GEOMETRY_ARITHMETIC_DETERMINANT_HPP

#include <cstddef>

#include <boost/geometry/core/access.hpp>
#include <boost/geometry/geometries/concepts/point_concept.hpp>
#include <boost/geometry/util/select_coordinate_type.hpp>

namespace boost { namespace geometry
{

#ifndef DOXYGEN_NO_DETAIL
namespace detail
{

template <typename ReturnType, typename U, typename V>
class calculate_determinant
{
    template <typename T>
    static inline ReturnType rt(T const& v)
    {
        return boost::numeric_cast<ReturnType>(v);
    }

public :

    static inline ReturnType apply(U const& ux, U const& uy
                                 , V const& vx, V const& vy)
    {
        return rt(ux) * rt(vy) - rt(uy) * rt(vx);
    }
};

template <typename ReturnType, typename U, typename V>
inline ReturnType determinant(U const& ux, U const& uy
                            , V const& vx, V const& vy)
{
    return calculate_determinant
        <
            ReturnType, U, V
        >::apply(ux, uy, vx, vy);
}

template <typename ReturnType, typename U, typename V>
inline ReturnType determinant(U const& u, V const& v)
{
    BOOST_CONCEPT_ASSERT( (concept::ConstPoint<U>) );
    BOOST_CONCEPT_ASSERT( (concept::ConstPoint<V>) );

    return calculate_determinant
        <
            ReturnType,
            typename geometry::coordinate_type<U>::type,
            typename geometry::coordinate_type<V>::type
        >::apply(get<0>(u), get<1>(u), get<0>(v), get<1>(v));
}

} // namespace detail
#endif // DOXYGEN_NO_DETAIL

}} // namespace boost::geometry

#endif // BOOST_GEOMETRY_ARITHMETIC_DETERMINANT_HPP


//
// Expansion Hunter
// Copyright 2016-2019 Illumina, Inc.
// All rights reserved.
//
// Author: Xiao Chen <xchen2@illumina.com>
//         Egor Dolzhenko <edolzhenko@illumina.com>
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
//
#include "genotyping/CopyNumberGenotyper.hh"
#include <boost/math/distributions/normal.hpp>
#include <math.h>
#include <numeric>

using boost::optional;
using boost::math::normal_distribution;
using std::vector;

namespace ehunter
{

CopyNumberGenotyper::CopyNumberGenotyper(
    int maxCopyNumber, double depthScaleFactor, double standardDeviationOfCN2,
    const std::vector<double>& meanDepthValues, const std::vector<double>& priorCopyNumberFreq)
    : maxCopyNumber_(maxCopyNumber)
    , depthScaleFactor_(depthScaleFactor)
    , standardDeviationOfCN2_(standardDeviationOfCN2)
    , meanDepthValues_(meanDepthValues)
    , priorCopyNumberFreq_(priorCopyNumberFreq)
{
    if (maxCopyNumber_ + 1 != static_cast<int>(meanDepthValues_.size()))
    {
        throw std::runtime_error("Number of mean values is inconsistent with total copy number states.");
    }

    if (maxCopyNumber_ + 1 != static_cast<int>(priorCopyNumberFreq_.size()))
    {
        throw std::runtime_error("Number of prior frequencies is inconsistent with total copy number states.");
    }
}

boost::optional<int> CopyNumberGenotyper::genotype(double normalizedDepth) const
{
    const double adjustedDepth = normalizedDepth / depthScaleFactor_;
    std::vector<double> likelihoodOfAllCN;
    std::vector<double> pvalueOfAllCN;

    for (int currentGenotype = 0; currentGenotype != maxCopyNumber_ + 1; currentGenotype++)
    {
        std::pair<double, double> likelihoodAndPvalue = genotypeLikelihoodAndPvalue(currentGenotype, adjustedDepth);
        double currentLikelihood = likelihoodAndPvalue.first;
        double currentPvalue = likelihoodAndPvalue.second;
        likelihoodOfAllCN.emplace_back(currentLikelihood);
        pvalueOfAllCN.emplace_back(currentPvalue);
    }

    const std::pair<int, double> bestGenotypeAndPosterior = getBestGenotypeAndPosterior(likelihoodOfAllCN);
    const int bestGenotype = bestGenotypeAndPosterior.first;
    const double posteriorProbabilityOfBestGenotype = bestGenotypeAndPosterior.second;
    const bool posteriorCheck = posteriorProbabilityOfBestGenotype > posteriorProbabilityThreshold_;
    const bool pvalueCheck = pvalueOfAllCN[bestGenotype] > pvalueThreshold_;
    const optional<int> genotype = (posteriorCheck && pvalueCheck) ? bestGenotype : optional<int>();
    return genotype;
}

std::pair<int, double>
CopyNumberGenotyper::getBestGenotypeAndPosterior(const std::vector<double>& likelihoodOfAllCN) const
{
    double sumOfLikelihood = 0;
    for (double likelihoodValue : likelihoodOfAllCN)
    {
        sumOfLikelihood += likelihoodValue;
    }

    assert(!likelihoodOfAllCN.empty());
    auto maxElement = max_element(likelihoodOfAllCN.begin(), likelihoodOfAllCN.end());
    const double maxLikelihood = *maxElement;
    std::pair<int, double> bestGenotypeAndPosterior;
    bestGenotypeAndPosterior.first = distance(likelihoodOfAllCN.begin(), maxElement);
    bestGenotypeAndPosterior.second = maxLikelihood / sumOfLikelihood;

    return bestGenotypeAndPosterior;
}

std::pair<double, double>
CopyNumberGenotyper::genotypeLikelihoodAndPvalue(int currentGenotype, double adjustedDepth) const
{
    assert(currentGenotype < (int)meanDepthValues_.size());
    assert(currentGenotype < (int)priorCopyNumberFreq_.size());

    const double meanValue = meanDepthValues_[currentGenotype];
    // standard deviation for each CN state can be computed from CN2 except for CN0 it is pre-defined
    const double standardDeviation
        = currentGenotype == 0 ? standardDeviationOfCN0_ : standardDeviationOfCN2_ * sqrt(double(currentGenotype) / 2);
    const double priorFreq = priorCopyNumberFreq_[currentGenotype];

    const normal_distribution<> cnDistribution(meanValue, standardDeviation);
    const double genotypeLikelihood = priorFreq * pdf(cnDistribution, adjustedDepth);
    const double cumulativeFrequency = cdf(cnDistribution, adjustedDepth);
    const double pvalue = std::min(cumulativeFrequency, 1 - cumulativeFrequency);

    std::pair<double, double> likelihoodAndPvalue;
    likelihoodAndPvalue.first = genotypeLikelihood;
    likelihoodAndPvalue.second = pvalue;

    return likelihoodAndPvalue;
}
}


///////////////////////////////////////////////////////////////
//  Copyright 2012 John Maddock. Distributed under the Boost
//  Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_

#ifndef BOOST_MP_INT_FUNC_HPP
#define BOOST_MP_INT_FUNC_HPP

#include <boost/multiprecision/number.hpp>

namespace boost{ namespace multiprecision{

namespace default_ops
{

template <class Backend>
inline void eval_qr(const Backend& x, const Backend& y, Backend& q, Backend& r)
{
   eval_divide(q, x, y);
   eval_modulus(r, x, y);
}

template <class Backend, class Integer>
inline Integer eval_integer_modulus(const Backend& x, Integer val)
{
   BOOST_MP_USING_ABS
   using default_ops::eval_modulus;
   using default_ops::eval_convert_to;
   typedef typename boost::multiprecision::detail::canonical<Integer, Backend>::type int_type;
   Backend t;
   eval_modulus(t, x, static_cast<int_type>(val));
   Integer result;
   eval_convert_to(&result, t);
   return abs(result);
}

#ifdef BOOST_MSVC
#pragma warning(push)
#pragma warning(disable:4127)
#endif

template <class B>
inline void eval_gcd(B& result, const B& a, const B& b)
{
   using default_ops::eval_lsb;
   using default_ops::eval_is_zero;
   using default_ops::eval_get_sign;

   int shift;

   B u(a), v(b);

   int s = eval_get_sign(u);

   /* GCD(0,x) := x */
   if(s < 0)
   {
      u.negate();
   }
   else if(s == 0)
   {
      result = v;
      return;
   }
   s = eval_get_sign(v);
   if(s < 0)
   {
      v.negate();
   }
   else if(s == 0)
   {
      result = u;
      return;
   }

   /* Let shift := lg K, where K is the greatest power of 2
   dividing both u and v. */

   unsigned us = eval_lsb(u);
   unsigned vs = eval_lsb(v);
   shift = (std::min)(us, vs);
   eval_right_shift(u, us);
   eval_right_shift(v, vs);

   do
   {
      /* Now u and v are both odd, so diff(u, v) is even.
      Let u = min(u, v), v = diff(u, v)/2. */
      s = u.compare(v);
      if(s > 0)
         u.swap(v);
      if(s == 0)
         break;
      eval_subtract(v, u);
      vs = eval_lsb(v);
      eval_right_shift(v, vs);
   }
   while(true);

   result = u;
   eval_left_shift(result, shift);
}

#ifdef BOOST_MSVC
#pragma warning(pop)
#endif

template <class B>
inline void eval_lcm(B& result, const B& a, const B& b)
{
   typedef typename mpl::front<typename B::unsigned_types>::type ui_type;
   B t;
   eval_gcd(t, a, b);

   if(eval_is_zero(t))
   {
      result = static_cast<ui_type>(0);
   }
   else
   {
      eval_divide(result, a, t);
      eval_multiply(result, b);
   }
   if(eval_get_sign(result) < 0)
      result.negate();
}

}

template <class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer>::type
   divide_qr(const number<Backend, ExpressionTemplates>& x, const number<Backend, ExpressionTemplates>& y,
   number<Backend, ExpressionTemplates>& q, number<Backend, ExpressionTemplates>& r)
{
   using default_ops::eval_qr;
   eval_qr(x.backend(), y.backend(), q.backend(), r.backend());
}

template <class Backend, expression_template_option ExpressionTemplates, class tag, class A1, class A2, class A3, class A4>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer>::type
   divide_qr(const number<Backend, ExpressionTemplates>& x, const multiprecision::detail::expression<tag, A1, A2, A3, A4>& y,
   number<Backend, ExpressionTemplates>& q, number<Backend, ExpressionTemplates>& r)
{
   divide_qr(x, number<Backend, ExpressionTemplates>(y), q, r);
}

template <class tag, class A1, class A2, class A3, class A4, class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer>::type
   divide_qr(const multiprecision::detail::expression<tag, A1, A2, A3, A4>& x, const number<Backend, ExpressionTemplates>& y,
   number<Backend, ExpressionTemplates>& q, number<Backend, ExpressionTemplates>& r)
{
   divide_qr(number<Backend, ExpressionTemplates>(x), y, q, r);
}

template <class tag, class A1, class A2, class A3, class A4, class tagb, class A1b, class A2b, class A3b, class A4b, class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer>::type
   divide_qr(const multiprecision::detail::expression<tag, A1, A2, A3, A4>& x, const multiprecision::detail::expression<tagb, A1b, A2b, A3b, A4b>& y,
   number<Backend, ExpressionTemplates>& q, number<Backend, ExpressionTemplates>& r)
{
   divide_qr(number<Backend, ExpressionTemplates>(x), number<Backend, ExpressionTemplates>(y), q, r);
}

template <class Backend, expression_template_option ExpressionTemplates, class Integer>
inline typename enable_if<mpl::and_<is_integral<Integer>, mpl::bool_<number_category<Backend>::value == number_kind_integer> >, Integer>::type
   integer_modulus(const number<Backend, ExpressionTemplates>& x, Integer val)
{
   using default_ops::eval_integer_modulus;
   return eval_integer_modulus(x.backend(), val);
}

template <class tag, class A1, class A2, class A3, class A4, class Integer>
inline typename enable_if<mpl::and_<is_integral<Integer>, mpl::bool_<number_category<typename multiprecision::detail::expression<tag, A1, A2, A3, A4>::result_type>::value == number_kind_integer> >, Integer>::type
   integer_modulus(const multiprecision::detail::expression<tag, A1, A2, A3, A4>& x, Integer val)
{
   typedef typename multiprecision::detail::expression<tag, A1, A2, A3, A4>::result_type result_type;
   return integer_modulus(result_type(x), val);
}

template <class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer, unsigned>::type
   lsb(const number<Backend, ExpressionTemplates>& x)
{
   using default_ops::eval_lsb;
   return eval_lsb(x.backend());
}

template <class tag, class A1, class A2, class A3, class A4>
inline typename enable_if_c<number_category<typename multiprecision::detail::expression<tag, A1, A2, A3, A4>::result_type>::value == number_kind_integer, unsigned>::type
   lsb(const multiprecision::detail::expression<tag, A1, A2, A3, A4>& x)
{
   typedef typename multiprecision::detail::expression<tag, A1, A2, A3, A4>::result_type number_type;
   number_type n(x);
   using default_ops::eval_lsb;
   return eval_lsb(n.backend());
}

template <class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer, bool>::type
   bit_test(const number<Backend, ExpressionTemplates>& x, unsigned index)
{
   using default_ops::eval_bit_test;
   return eval_bit_test(x.backend(), index);
}

template <class tag, class A1, class A2, class A3, class A4>
inline typename enable_if_c<number_category<typename multiprecision::detail::expression<tag, A1, A2, A3, A4>::result_type>::value == number_kind_integer, bool>::type
   bit_test(const multiprecision::detail::expression<tag, A1, A2, A3, A4>& x, unsigned index)
{
   typedef typename multiprecision::detail::expression<tag, A1, A2, A3, A4>::result_type number_type;
   number_type n(x);
   using default_ops::eval_bit_test;
   return eval_bit_test(n.backend(), index);
}

template <class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer, number<Backend, ExpressionTemplates>&>::type
   bit_set(number<Backend, ExpressionTemplates>& x, unsigned index)
{
   using default_ops::eval_bit_set;
   eval_bit_set(x.backend(), index);
   return x;
}

template <class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer, number<Backend, ExpressionTemplates>&>::type
   bit_unset(number<Backend, ExpressionTemplates>& x, unsigned index)
{
   using default_ops::eval_bit_unset;
   eval_bit_unset(x.backend(), index);
   return x;
}

template <class Backend, expression_template_option ExpressionTemplates>
inline typename enable_if_c<number_category<Backend>::value == number_kind_integer, number<Backend, ExpressionTemplates>&>::type
   bit_flip(number<Backend, ExpressionTemplates>& x, unsigned index)
{
   using default_ops::eval_bit_flip;
   eval_bit_flip(x.backend(), index);
   return x;
}

namespace detail{

//
// Within powm, we need a type with twice as many digits as the argument type, define
// a traits class to obtain that type:
//
template <class Backend>
struct double_precision_type
{
   typedef Backend type;
};

//
// Calculate (a^p)%c:
//
template <class Backend>
void eval_powm(Backend& result, const Backend& a, const Backend& p, const Backend& c)
{
   using default_ops::eval_bit_test;
   using default_ops::eval_get_sign;
   using default_ops::eval_multiply;
   using default_ops::eval_modulus;
   using default_ops::eval_right_shift;

   typedef typename double_precision_type<Backend>::type double_type;
   typedef typename canonical<unsigned char, double_type>::type ui_type;

   double_type x, y(a), b(p), t;
   x = ui_type(1u);

   while(eval_get_sign(b) > 0)
   {
      if(eval_bit_test(b, 0))
      {
         eval_multiply(t, x, y);
         eval_modulus(x, t, c);
      }
      eval_multiply(t, y, y);
      eval_modulus(y, t, c);
      eval_right_shift(b, ui_type(1));
   }
   Backend x2(x);
   eval_modulus(result, x2, c);
}

template <class Backend, class Integer>
void eval_powm(Backend& result, const Backend& a, const Backend& p, Integer c)
{
   typedef typename double_precision_type<Backend>::type double_type;
   typedef typename canonical<unsigned char, double_type>::type ui_type;
   typedef typename canonical<Integer, double_type>::type i1_type;
   typedef typename canonical<Integer, Backend>::type i2_type;

   using default_ops::eval_bit_test;
   using default_ops::eval_get_sign;
   using default_ops::eval_multiply;
   using default_ops::eval_modulus;
   using default_ops::eval_right_shift;

   if(eval_get_sign(p) < 0)
   {
      BOOST_THROW_EXCEPTION(std::runtime_error("powm requires a positive exponent."));
   }

   double_type x, y(a), b(p), t;
   x = ui_type(1u);

   while(eval_get_sign(b) > 0)
   {
      if(eval_bit_test(b, 0))
      {
         eval_multiply(t, x, y);
         eval_modulus(x, t, static_cast<i1_type>(c));
      }
      eval_multiply(t, y, y);
      eval_modulus(y, t, static_cast<i1_type>(c));
      eval_right_shift(b, ui_type(1));
   }
   Backend x2(x);
   eval_modulus(result, x2, static_cast<i2_type>(c));
}

template <class Backend, class Integer>
typename enable_if<is_unsigned<Integer> >::type eval_powm(Backend& result, const Backend& a, Integer b, const Backend& c)
{
   typedef typename double_precision_type<Backend>::type double_type;
   typedef typename canonical<unsigned char, double_type>::type ui_type;

   using default_ops::eval_bit_test;
   using default_ops::eval_get_sign;
   using default_ops::eval_multiply;
   using default_ops::eval_modulus;
   using default_ops::eval_right_shift;

   double_type x, y(a), t;
   x = ui_type(1u);

   while(b > 0)
   {
      if(b & 1)
      {
         eval_multiply(t, x, y);
         eval_modulus(x, t, c);
      }
      eval_multiply(t, y, y);
      eval_modulus(y, t, c);
      b >>= 1;
   }
   Backend x2(x);
   eval_modulus(result, x2, c);
}

template <class Backend, class Integer>
typename enable_if<is_signed<Integer> >::type eval_powm(Backend& result, const Backend& a, Integer b, const Backend& c)
{
   if(b < 0)
   {
      BOOST_THROW_EXCEPTION(std::runtime_error("powm requires a positive exponent."));
   }
   eval_powm(result, a, static_cast<typename make_unsigned<Integer>::type>(b), c);
}

template <class Backend, class Integer1, class Integer2>
typename enable_if<is_unsigned<Integer1> >::type eval_powm(Backend& result, const Backend& a, Integer1 b, Integer2 c)
{
   typedef typename double_precision_type<Backend>::type double_type;
   typedef typename canonical<unsigned char, double_type>::type ui_type;
   typedef typename canonical<Integer1, double_type>::type i1_type;
   typedef typename canonical<Integer2, Backend>::type i2_type;

   using default_ops::eval_bit_test;
   using default_ops::eval_get_sign;
   using default_ops::eval_multiply;
   using default_ops::eval_modulus;
   using default_ops::eval_right_shift;

   double_type x, y(a), t;
   x = ui_type(1u);

   while(b > 0)
   {
      if(b & 1)
      {
         eval_multiply(t, x, y);
         eval_modulus(x, t, static_cast<i1_type>(c));
      }
      eval_multiply(t, y, y);
      eval_modulus(y, t, static_cast<i1_type>(c));
      b >>= 1;
   }
   Backend x2(x);
   eval_modulus(result, x2, static_cast<i2_type>(c));
}

template <class Backend, class Integer1, class Integer2>
typename enable_if<is_signed<Integer1> >::type eval_powm(Backend& result, const Backend& a, Integer1 b, Integer2 c)
{
   if(b < 0)
   {
      BOOST_THROW_EXCEPTION(std::runtime_error("powm requires a positive exponent."));
   }
   eval_powm(result, a, static_cast<typename make_unsigned<Integer1>::type>(b), c);
}

struct powm_func
{
   template <class T, class U, class V>
   void operator()(T& result, const T& b, const U& p, const V& m)const
   {
      eval_powm(result, b, p, m);
   }
};

}

template <class T, class U, class V>
inline typename enable_if<
   mpl::and_<
      mpl::bool_<number_category<T>::value == number_kind_integer>,
      mpl::or_<
         is_number<T>,
         is_number_expression<T>
      >,
      mpl::or_<
         is_number<U>,
         is_number_expression<U>,
         is_integral<U>
      >,
      mpl::or_<
         is_number<V>,
         is_number_expression<V>,
         is_integral<V>
      >
   >,
   detail::expression<detail::function, detail::powm_func, T, U, V> >::type
   powm(const T& b, const U& p, const V& mod)
{
   return detail::expression<detail::function, detail::powm_func, T, U, V>(
      detail::powm_func(), b, p, mod);
}

}} //namespaces

#endif


///////////////////////////////////////////////////////////////////////////////
// importance_sampling::example::sampler.cpp                                 //
//                                                                           //
//  Copyright 2009 Erwann Rogard. Distributed under 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)         //
///////////////////////////////////////////////////////////////////////////////
#include <vector>
#include <cmath>
#include <algorithm>
#include <iterator>
#include <boost/range.hpp>
#include <boost/accumulators/framework/extractor.hpp>
#include <boost/accumulators/accumulators.hpp>
#include <boost/accumulators/statistics/stats.hpp>
#include <boost/random/mersenne_twister.hpp>
#include <boost/random/variate_generator.hpp>
#include <boost/format.hpp>
#include <boost/function.hpp>
#include <boost/typeof/typeof.hpp>
#include <boost/utility/result_of.hpp>
#include <boost/math/tools/precision.hpp>
#include <boost/iterator/transform_iterator.hpp>
#include <boost/lambda/bind.hpp>
#include <boost/lambda/lambda.hpp>
#include <boost/ref.hpp>
#include <boost/foreach.hpp>

#include <boost/fusion/include/at_key.hpp>
#include <boost/fusion/include/pair.hpp>
#include <boost/fusion/include/map.hpp>
#include <boost/fusion/include/make_map.hpp>

#include <boost/math/tools/precision.hpp>

#include <boost/statistics/detail/accumulator/statistics/proportion_less_than.hpp>
#include <boost/statistics/detail/distribution_common/distributions/reference/wrapper.hpp>
#include <boost/statistics/detail/distribution_common/meta/random/generator.hpp>
#include <boost/statistics/detail/distribution_common/functor/log_unnormalized_pdf.hpp>
#include <boost/statistics/detail/distribution_toolkit/distributions/normal/include.hpp>
#include <boost/statistics/detail/distribution_toolkit/map_pdf/ratio_pdf/include.hpp>

#include <boost/statistics/detail/non_parametric/kolmogorov_smirnov/check_convergence.hpp>

#include <boost/statistics/detail/fusion/at_key/functor.hpp>
#include <boost/statistics/detail/fusion/at_key/range.hpp>
#include <boost/statistics/detail/non_parametric/kolmogorov_smirnov/statistic.hpp>

#include <boost/statistics/detail/importance_sampling/weights/prepare_weights.hpp>
#include <boost/statistics/detail/importance_sampling/random/include.hpp>
#include <boost/statistics/detail/importance_sampling/statistics/percentage_effective_sample_size.hpp>

void example_sampler(std::ostream& os){
    os << "->example_sampler : \n";

    // Sample from N(x|mu,sigma^2)N(x|mu,sigma^2) = N(x|mu,sigma^2/2),
    // using SIR with N(x|mu+sigma,sigma^2) as proposal density.
    // The quality of the sample is assessed by a series of
    // kolmorov-distances along the the sample size of the targets.
    using namespace boost;
    namespace stat = boost::statistics::detail;
    namespace dist = stat::distribution;
    namespace tk = stat::distribution::toolkit;
    namespace is = stat::importance_sampling;

    typedef std::string                                 str_;
    typedef double                                      val_;
    typedef std::vector<val_>                           vals_;
    typedef range_iterator<vals_>::type                 vals_it_;
    typedef math::normal_distribution<val_>             dist_;
    typedef mt19937                                     urng_;
    typedef is::prepare_weights<val_>                   prepare_weights_;

    typedef stat::accumulator::tag::percentage_effective_sample_size tag_ess_;
    typedef stat::accumulator::tag::proportion_less_than tag_plt_;
    typedef boost::accumulators::stats<tag_ess_,tag_plt_> stats_;
    typedef boost::accumulators::accumulator_set<val_,stats_> acc_;
    typedef std::size_t                                 size_;

    typedef boost::mpl::int_<0> k0_;
    typedef boost::mpl::int_<1> k1_;
	typedef boost::fusion::result_of::make_map<k0_,k1_,val_,val_>::type data_;
    typedef std::vector<data_>                          vec_data_;
    typedef range_iterator<vec_data_>::type             vec_data_it_;

    typedef stat::fusion::at_key::meta_range<
        vec_data_it_,
        k0_
    >::type range1_;
    typedef is::sampler<range1_,val_>                   	is_sampler_;
    typedef boost::variate_generator<urng_&,is_sampler_> 	vg_;

    // Constants
    const unsigned n_p          = 5e4;
    const val_ max_log          = 100.0;
    const val_ mu               = 0.0;
    const val_ sigma            = 1.0;
    const val_ t_mu             = mu + sigma;
    const val_ t_sigma          = sigma/sqrt(2.0);
    const val_ eps = boost::math::tools::epsilon<val_>();

    const unsigned n_loops  = 7;
    const unsigned n1       = 1e1;
    const unsigned n2       = 1e1;

    prepare_weights_ prepare_weights( max_log );
    dist_ p_d( mu, sigma );       // proposal
    dist_ t_d( t_mu, t_sigma );   // target

    urng_ urng;

    vec_data_ proposals(n_p);
    {
        // Generate proposal values and their log_pdf
        // {(x,log_pdf(x)):i=1,...,n_p}
		BOOST_AUTO(vg,dist::make_random_generator(urng,p_d));
        BOOST_FOREACH(data_& data,proposals){
        	val_ x = vg();
			boost::fusion::at_key<k0_>(data) = x;
            boost::fusion::at_key<k1_>(data) = log_unnormalized_pdf(p_d,x);
		}
    }
    vals_ is_weights( n_p );
    {
        // is_weights <- log_pdf(t_d) - log_pdf(p_d)

        vals_it_ i = boost::begin( is_weights );
        BOOST_FOREACH(const data_& d, proposals)
        {
            *i = log_unnormalized_pdf( t_d, fusion::at_key<k0_>(d) )
                - fusion::at_key<k1_>(d);
            ++i;
        }

    }
    {
        //is_weights <- exp(is_weights + c )
        prepare_weights(
            boost::begin(is_weights),
            boost::end(is_weights)
        );
        os << "weights : " << std::endl
            << prepare_weights_::header << " = "
            << prepare_weights << std::endl;
    }
    {
        acc_ acc = std::for_each(
            boost::begin(is_weights),
            boost::end(is_weights),
            acc_(( stat::accumulator::keyword::threshold = eps ))
        );
        val_ ess = boost::accumulators::extract_result<tag_ess_>(acc);
        val_ plt_eps = boost::accumulators::extract_result<tag_plt_>(acc);
        const std::string str
        	= (boost::format("(ess,plt_eps) = (%1%,%2%)")% ess % plt_eps).str();
        os << str << std::endl;
    }
    {
        range1_ r1 = stat::fusion::at_key::make_range<k0_>(
            boost::begin(proposals),
            boost::end(proposals)
        );
        vg_ vg(
            urng,
            is_sampler_(
                is_weights,
                r1
            )
        );
        vec_data_ targets;
        {
            os << "proposal : " << description(p_d) << std::endl;
            os << "target : " 	<< description(t_d) << std::endl;
			namespace ks = boost::statistics::detail::kolmogorov_smirnov;
			typedef ks::check_convergence<val_> check_;
	        check_ check;
			check(n_loops,n1,n2,t_d,vg,os);

        }
    }

    os << "<-" << std::endl;

}


<<<<<<< HEAD
/*    Copyright (c) 2010-2018, Delft University of Technology
=======
/*    Copyright (c) 2010-2019, Delft University of Technology
>>>>>>> origin/master
 *    All rigths reserved
 *
 *    This file is part of the Tudat. Redistribution and use in source and
 *    binary forms, with or without modification, are permitted exclusively
 *    under the terms of the Modified BSD license. You should have received
 *    a copy of the license with this file. If not, please or visit:
 *    http://tudat.tudelft.nl/LICENSE.
 *
 *    References
 *
 *
 */

#include <cmath>

#include <Eigen/Geometry>

#include "Tudat/Mathematics/BasicMathematics/rotationAboutArbitraryAxis.h"

namespace tudat
{
namespace basic_mathematics
{

//! Compute rotation of point about arbitrary axis
Eigen::Vector3d computeRotationOfPointAboutArbitraryAxis(
        const Eigen::Vector3d& originOfRotation,
        const double angleOfRotation,
        const Eigen::Vector3d& axisOfRotation,
        const Eigen::Vector3d& initialPositionOfPoint )
{

    //Declare and initialize rotation matrix
    Eigen::Matrix3d rotationMatrix = Eigen::Matrix3d::Zero( );

    // Compute rotation matrix using AngleAxis object.
    rotationMatrix = Eigen::AngleAxisd( angleOfRotation, axisOfRotation.normalized( ) );

    // Compute initial of position of point with respect to origin of rotation.
    const Eigen::Vector3d initialPositionOfPointWithRespectToOriginOfRotation =
            initialPositionOfPoint - originOfRotation;

    // Compute rotation of point about axis of rotation with respect to origin of rotation.
    const Eigen::Vector3d rotatedPositionWithRespectToOriginOfRotation =
            rotationMatrix * initialPositionOfPointWithRespectToOriginOfRotation;

    //Return position with respect to the chosen arbitrary origin after rotation about
    //arbitrary axis.
    return rotatedPositionWithRespectToOriginOfRotation + originOfRotation;

}

//! Compute rotation of vector about arbitrary axis
Eigen::Vector3d computeRotationOfVectorAboutArbitraryAxis(
        const Eigen::Vector3d& originOfRotation,
        const double angleOfRotation,
        const Eigen::Vector3d& axisOfRotation,
        const Eigen::Vector3d& initialPositionOfVectorTail,
        const Eigen::Vector3d& initialVector )
{

    // Compute rotation of the tail of vector. Resulted position is with respect to the chosen
    // arbitrary origin.
    Eigen::Vector3d rotatedPositionOfVectorTail =
            computeRotationOfPointAboutArbitraryAxis( originOfRotation, angleOfRotation,
                                                      axisOfRotation, initialPositionOfVectorTail );

    // Compute rotation of the head of vector. Resulted position is with respect to the chosen
    // arbitrary origin.
    Eigen::Vector3d rotatedPositionOfVectorHead =
            computeRotationOfPointAboutArbitraryAxis( originOfRotation, angleOfRotation,
                                                      axisOfRotation,
                                                      initialPositionOfVectorTail + initialVector );

    // Return rotated vector with respect to the chosen arbitrary origin
    return rotatedPositionOfVectorHead - rotatedPositionOfVectorTail;

}

} // namespace basic_mathematics
} // namespace tudat


/* Copyright (c) 2017, Julian Straub <jstraub@csail.mit.edu>
 * Licensed under the MIT license. See the license file LICENSE.
 */
// g++ -Wall -std=c++1z -I /usr/include/eigen3/ main.cpp -o test
#include <random>
#include <cmath>
#include <iostream>
#include <Eigen/Dense>
#include "vmf.hpp"
#include "vmfPrior.hpp"
#include "normal.hpp"
#include "sample.hpp"

int main() {

  std::mt19937 rnd(1);

  vMF<float,3> vmfA(Eigen::Vector3f(1,0,0), 500);
  vMF<float,3> vmfB(Eigen::Vector3f(cos(5*M_PI/180.),sin(5*M_PI/180.),0), 500);
  Eigen::Matrix3f SigmaO = 0.0001*Eigen::Matrix3f::Identity();
  Normal<float,3> gaussO(SigmaO);
  float tauO = 500.;

  size_t N=100;
  std::vector<std::vector<Eigen::Vector3f>> n; // normals
  std::vector<std::vector<Eigen::Vector3f>> xn; // normal observations
  std::vector<std::vector<Eigen::Vector3f>> x; // loc observations
  std::vector<std::vector<Eigen::Vector3f>> p; // plane location
  for (size_t i=0; i<N; ++i) {
    n.push_back(std::vector<Eigen::Vector3f>());
    xn.push_back(std::vector<Eigen::Vector3f>());
    x.push_back(std::vector<Eigen::Vector3f>());
    p.push_back(std::vector<Eigen::Vector3f>());
    for (size_t j=0; j<N; ++j) {
      if (i<N/2) {
        n[i].push_back(vmfA.sample(rnd));
        xn[i].push_back(vmfA.sample(rnd));
        x[i].push_back(Eigen::Vector3f(1.,(i-N*0.25)/float(0.25*N),(j-N*0.5)/float(0.5*N)));
      }
      if (i>=N/2) {
        n[i].push_back(vmfB.sample(rnd));
        xn[i].push_back(vmfB.sample(rnd));
        x[i].push_back(Eigen::Vector3f((N*0.75-i)/float(0.25*N),1.,(j-N*0.5)/float(0.5*N)));
      }
      p[i].push_back(x[i][j]);
      x[i][j] += gaussO.sample(rnd);
      p[i][j] += gaussO.sample(rnd);
//      std::cout << i << " " << j << ": "  <<  x[i][j][0] << "\t" << x[i][j][1] << std::endl;
    }
  }
  std::cout << "have " << n.size() << " input data" << std::endl;

  std::vector<Eigen::Vector3f> xSum(1, Eigen::Vector3f::Zero());
  std::vector<std::vector<uint32_t>> z(n.size());
  for (size_t i=0; i<n.size(); ++i)
    for (size_t j=0; j<n[i].size(); ++j)  {
      xSum[0] += n[i][j];
      z[i].push_back(0);
    }
  std::vector<float> counts(1, n.size()*n.size());
  std::vector<vMF<float,3>> vmfs;
  vMFprior<float> base(Eigen::Vector3f(0,0,1), .1, 0.0);
  float logAlpha = log(10.);
  float lambda = .10;

  vmfs.push_back(base.sample(rnd));
  for (size_t it=0; it<10000; ++it) {
    // sample labels | parameters
    size_t K = vmfs.size();
    for (size_t i=0; i<n.size(); ++i) {
      for (size_t j=0; j<n[i].size(); ++j) {
        Eigen::VectorXf logPdfs(K+1);
        Eigen::VectorXf pdfs(K+1);

        Eigen::VectorXf neighNs = Eigen::VectorXf::Zero(K);
//        if (i+1<N) neighNs[z[i+1][j]] += 1.f;
//        if (i>=1)  neighNs[z[i-1][j]] += 1.f;
//        if (j+1<N) neighNs[z[i][j+1]] += 1.f;
//        if (j>=1)  neighNs[z[i][j-1]] += 1.f;

        if (i+1<N) neighNs[z[i+1][j]] += n[i+1][j].dot(n[i][j]);
        if (i>=1)  neighNs[z[i-1][j]] += n[i-1][j].dot(n[i][j]);
        if (j+1<N) neighNs[z[i][j+1]] += n[i][j+1].dot(n[i][j]);
        if (j>=1)  neighNs[z[i][j-1]] += n[i][j-1].dot(n[i][j]);

//        if (i+1<N) neighNs[z[i+1][j]] += vmfs[z[i+1][j]].mu_.dot(n[i][j]);
//        if (i>=1)  neighNs[z[i-1][j]] += vmfs[z[i-1][j]].mu_.dot(n[i][j]);
//        if (j+1<N) neighNs[z[i][j+1]] += vmfs[z[i][j+1]].mu_.dot(n[i][j]);
//        if (j>=1)  neighNs[z[i][j-1]] += vmfs[z[i][j-1]].mu_.dot(n[i][j]);

//        if (i+1<N) neighNs[z[i+1][j]] += vmfs[z[i+1][j]].mu_.dot(vmfs[z[i][j]].mu_);
//        if (i>=1)  neighNs[z[i-1][j]] += vmfs[z[i-1][j]].mu_.dot(vmfs[z[i][j]].mu_);
//        if (j+1<N) neighNs[z[i][j+1]] += vmfs[z[i][j+1]].mu_.dot(vmfs[z[i][j]].mu_);
//        if (j>=1)  neighNs[z[i][j-1]] += vmfs[z[i][j-1]].mu_.dot(vmfs[z[i][j]].mu_);

        for (size_t k=0; k<K; ++k) {
          logPdfs[k] = lambda*(neighNs[k]-4);
          if (z[i][j] == k) {
            // TODO what if last in cluster
            logPdfs[k] += log(counts[k]-1)+vmfs[k].logPdf(n[i][j]);
          } else {
            logPdfs[k] += log(counts[k])+vmfs[k].logPdf(n[i][j]);
          }
        }
        logPdfs[K] = logAlpha + base.logMarginal(n[i][j]);
        logPdfs = logPdfs.array() - logSumExp<float>(logPdfs);
        pdfs = logPdfs.array().exp();
        size_t zPrev = z[i][j];
        z[i][j] = sampleDisc(pdfs, rnd);
        //      std::cout << z[i] << " " << K << ": " << pdfs.transpose() << std::endl;
        if (z[i][j] == K) {
          vmfs.push_back(base.posterior(n[i][j],1).sample(rnd));
          counts.push_back(0);
          xSum.push_back(Eigen::Vector3f::Zero());
          K++;
        }
        if (zPrev != z[i][j]) {
          counts[zPrev] --;
          counts[z[i][j]] ++;
          xSum[zPrev] -= n[i][j];
          xSum[z[i][j]] += n[i][j];
        }
      }
    }
//    std::cout << "sample parameters" << std::endl;
    // sample parameters | labels
//    for (size_t i=0; i<n.size(); ++i) {
//      xSum[z[i]] += n[i]; // TODO: can fold in above as well
//    }
    for (size_t k=0; k<K; ++k) {
      if (counts[k] > 0) {
        vmfs[k] = base.posterior(xSum[k],counts[k]).sample(rnd);
      }
    }
    std::cout << "counts " << K << ": ";
    for (size_t k=0; k<K; ++k) if (counts[k] > 0) std::cout << counts[k] << " ";
    std::cout << "\ttaus: " ;
    for (size_t k=0; k<K; ++k) if (counts[k] > 0) std::cout << vmfs[k].tau_ << " ";
    std::cout << std::endl;
//    for (size_t k=0; k<K; ++k)
//      if (counts[k] > 0) {
//        std::cout << vmfs[k].mu_.transpose() << std::endl;
//      }

    // sample ns
    for (size_t k=0; k<K; ++k)
      xSum[k] = Eigen::Vector3f::Zero();
    for (size_t i=0; i<n.size(); ++i) {
      for (size_t j=0; j<n[i].size(); ++j) {
        Eigen::Vector3f mu = xn[i][j]*tauO + vmfs[z[i][j]].mu_*vmfs[z[i][j]].tau_;
        n[i][j] = vMF<float,3>(mu).sample(rnd);
        xSum[z[i][j]] += n[i][j];
      }
    }
//    std::cout << n[N/4][N/2].transpose() << "\t" << n[(3*N)/4][N/2].transpose() << std::endl;
//    std::cout << xn[N/4][N/2].transpose() << "\t" << xn[(3*N)/4][N/2].transpose() << std::endl;

    // sample locations
    for (size_t i=0; i<x.size(); ++i) {
      for (size_t j=0; j<x[i].size(); ++j) {
        Eigen::Matrix3f SigmaPl;
        Eigen::Matrix3f Info =  SigmaO.inverse();
        Eigen::Vector3f xi = SigmaO.ldlt().solve(x[i][j]);
        if (i+1<N && z[i][j] == z[i+1][j]) {
          SigmaPl = vmfs[z[i+1][j]].mu_*vmfs[z[i+1][j]].mu_.transpose();
          Info += SigmaPl;
          xi += SigmaPl*p[i+1][j];
        }
        if (i>=1 && z[i][j] == z[i-1][j])  {
          SigmaPl = vmfs[z[i-1][j]].mu_*vmfs[z[i-1][j]].mu_.transpose();
          Info += SigmaPl;
          xi += SigmaPl*p[i-1][j];
        }
        if (j+1<N && z[i][j] == z[i][j+1]) {
          SigmaPl = vmfs[z[i][j+1]].mu_*vmfs[z[i][j+1]].mu_.transpose();
          Info += SigmaPl;
          xi += SigmaPl*p[i][j+1];
        }
        if (j>=1 && z[i][j] == z[i][j-1])  {
          SigmaPl = vmfs[z[i][j-1]].mu_*vmfs[z[i][j-1]].mu_.transpose();
          Info += SigmaPl;
          xi += SigmaPl*p[i][j-1];
        }
        Eigen::Matrix3f Sigma = Info.inverse();
        Eigen::Vector3f mu = Sigma*xi;
//        std::cout << xi.transpose() << " " << mu.transpose() << std::endl;
        p[i][j] = Normal<float,3>(mu, Sigma).sample(rnd);
      }
    }
//    std::cout << p[N/4][N/2].transpose() << "\t" << p[(3*N)/4][N/2].transpose() << std::endl;
//    std::cout << x[N/4][N/2].transpose() << "\t" << x[(3*N)/4][N/2].transpose() << std::endl;
  }
  return 0;
}


// Copyright (c) 2017 Franka Emika GmbH
// Use of this source code is governed by the Apache-2.0 license, see LICENSE
#include <array>
#include <atomic>
#include <cmath>
#include <functional>
#include <iostream>
#include <iterator>
#include <mutex>
#include <thread>

#include <Eigen/Core>
#include <Eigen/Eigen>

#include <franka/duration.h>
#include <franka/exception.h>
#include <franka/model.h>
#include <franka/rate_limiting.h>
#include <franka/robot.h>

namespace {
template <class T, size_t N>
std::ostream& operator<<(std::ostream& ostream, const std::array<T, N>& array) {
  ostream << "[";
  std::copy(array.cbegin(), array.cend() - 1, std::ostream_iterator<T>(ostream, ","));
  std::copy(array.cend() - 1, array.cend(), std::ostream_iterator<T>(ostream));
  ostream << "]";
  return ostream;
}
}  // anonymous namespace

int main(int argc, char** argv) {
  // Check whether the required arguments were passed.
  if (argc != 2) {
    std::cerr << "Usage: " << argv[0] << " <robot-hostname>" << std::endl;
    return -1;
  }

  // Initialize data fields for the print thread.
  struct {
    std::mutex mutex;
    bool has_data;
    franka::RobotState robot_state;
    std::array<double, 7> tau_d_calculated;
    std::array<double, 7> tau_d_last;
    std::array<double, 7> gravity;
    std::array<double, 7> coriolis;
    std::array<double, 42> jacobian_array;
  } print_data{};
  std::atomic_bool running{true};
  const double print_rate = 10.0;
  // Start print thread.
  std::thread print_thread([print_rate, &print_data, &running]() {
    while (running) {
      // Sleep to achieve the desired print rate.
      std::this_thread::sleep_for(
          std::chrono::milliseconds(static_cast<int>((1.0 / print_rate * 1000.0))));

      // Try to lock data to avoid read write collisions.
      if (print_data.mutex.try_lock()) {
        if (print_data.has_data) {
          // std::array<double, 7> tau_error{};
          // double error_rms(0.0);
          // std::array<double, 7> tau_d_actual{};
          // for (size_t i = 0; i < 7; ++i) {
          //   tau_d_actual[i] = print_data.tau_d_last[i] + print_data.gravity[i];
          //   tau_error[i] = tau_d_actual[i] - print_data.robot_state.tau_J[i];
          //   error_rms += std::pow(tau_error[i], 2.0) / tau_error.size();
          // }
          // error_rms = std::sqrt(error_rms);
          const Eigen::Matrix<double, 7, 1> mytorques(print_data.robot_state.tau_ext_hat_filtered.data());
          const Eigen::Matrix<double, 6, 7> jacobian(print_data.jacobian_array.data());

          const Eigen::MatrixXd jacobianPinv1 = ( jacobian * jacobian.transpose() ).inverse() * jacobian;
          const Eigen::VectorXd mywrench1= jacobianPinv1 * mytorques;

          Eigen::FullPivLU<Eigen::Matrix<double, 6, 6>> lu_decomp_ = Eigen::FullPivLU<Eigen::Matrix<double, 6, 6>>();
          lu_decomp_.compute( jacobian * jacobian.transpose() );
          const Eigen::MatrixXd jacobianPinv2 = lu_decomp_.inverse() * jacobian;
          const Eigen::VectorXd mywrench2 = jacobianPinv2 * mytorques;

          Eigen::JacobiSVD<Eigen::Matrix<double, 7, 6>> svdT = Eigen::JacobiSVD<Eigen::Matrix<double, 7, 6>>();
          svdT.setThreshold(0.08); //critical singular value threshold is 0.08
          svdT.compute(jacobian.transpose(), Eigen::ComputeThinU | Eigen::ComputeThinV);
          const Eigen::VectorXd mywrench3 = svdT.solve(mytorques);

          const Eigen::MatrixXd U = svdT.matrixU();
          const Eigen::MatrixXd V = svdT.matrixV();
          const Eigen::VectorXd Svec  = svdT.singularValues();
          Eigen::Matrix<double, 6, 6> S = Eigen::Matrix<double, 6, 6>::Zero();
          S.diagonal() = Svec;
          Eigen::Matrix<double, 6, 6> Sinv = Eigen::Matrix<double, 6, 6>::Zero();
          // Sinv = S.inverse();
          for(int i=0; i<6; i++){
            if (Svec(i)>=0.08){
              Sinv(i,i) = 1.0 / Svec(i);
            }
          }
          const Eigen::MatrixXd jacobianPinv4 = V * Sinv * U.transpose();
          const Eigen::VectorXd mywrench4 = jacobianPinv4 * mytorques;

          // Print data to console
          std::cout << "control_command_success_rate: " <<  print_data.robot_state.control_command_success_rate << std::endl
                    << "joint configuration: " <<  print_data.robot_state.q << std::endl
                    // << "jacobian: \n" << jacobian << std::endl
                    // << "U: \n" << U << std::endl
                    // << "V: \n" << V << std::endl
                    // << "S: \n" << S << std::endl
                    // << "Sinv: \n" << Sinv << std::endl
                    // << "jacobianPinv: \n" << jacobianPinv1.row(0) << std::endl
                    // << "jacobianPinv: \n" << jacobianPinv2.row(0) << std::endl
                    // << "jacobianPinv: \n" << jacobianPinv4.row(0) << std::endl
                    // << "jacobian-singularValues: " << svdT.singularValues().transpose() << std::endl
                    << "External torque, filtered: " <<  print_data.robot_state.tau_ext_hat_filtered << std::endl
                    // << "end-effector wrench (base frame)     : " <<  print_data.robot_state.K_F_ext_hat_K << std::endl
                    << "end-effector wrench (stiffness frame): " <<  print_data.robot_state.O_F_ext_hat_K << std::endl
                    << "end-effector wrench (stiffness frame): " <<  mywrench1.transpose() << std::endl
                    << "end-effector wrench (stiffness frame): " <<  mywrench2.transpose() << std::endl
                    << "end-effector wrench (stiffness frame): " <<  mywrench3.transpose() << std::endl
                    << "end-effector wrench (stiffness frame): " <<  mywrench4.transpose() << std::endl
                    // << "tau_error [Nm]: " << tau_error << std::endl
                    // << "tau_commanded [Nm]: " << tau_d_actual << std::endl
                    // << "tau_measured [Nm]: " << print_data.robot_state.tau_J << std::endl
                    // << "root mean square of tau_error [Nm]: " << error_rms << std::endl
                    << "-----------------------" << std::endl
                    ;
          print_data.has_data = false;
        }
        print_data.mutex.unlock();
      }
    }
  });

  try {
    franka::Robot robot(argv[1]);
    robot.setJointImpedance({{3000, 3000, 3000, 2500, 2500, 2000, 2000}}); //values taken from https://github.com/frankaemika/libfranka/blob/master/examples/examples_common.cpp#L18
    robot.setCartesianImpedance({{3000, 3000, 3000, 300, 300, 300}}); //values taken from https://github.com/frankaemika/libfranka/blob/master/examples/examples_common.cpp#L19
    // robot.setJointImpedance({{100,100,100,100,100,100,100}});
    // robot.setCartesianImpedance({{100,100,100,10,10,10}});
    robot.setCollisionBehavior(
        {{20.0, 20.0, 18.0, 18.0, 16.0, 14.0, 12.0}}, {{20.0, 20.0, 18.0, 18.0, 16.0, 14.0, 12.0}},
        {{20.0, 20.0, 18.0, 18.0, 16.0, 14.0, 12.0}}, {{20.0, 20.0, 18.0, 18.0, 16.0, 14.0, 12.0}},
        {{20.0, 20.0, 20.0, 25.0, 25.0, 25.0}}, {{20.0, 20.0, 20.0, 25.0, 25.0, 25.0}},
        {{20.0, 20.0, 20.0, 25.0, 25.0, 25.0}}, {{20.0, 20.0, 20.0, 25.0, 25.0, 25.0}});

    //TODO: be careful with this mode!
    robot.setCollisionBehavior(
        {{2000.0, 2000.0, 1800.0, 1800.0, 1600.0, 1400.0, 1200.0}}, {{2000.0, 2000.0, 1800.0, 1800.0, 1600.0, 1400.0, 1200.0}},
        {{2000.0, 2000.0, 1800.0, 1800.0, 1600.0, 1400.0, 1200.0}}, {{2000.0, 2000.0, 1800.0, 1800.0, 1600.0, 1400.0, 1200.0}},
        {{2000.0, 2000.0, 2000.0, 2500.0, 2500.0, 2500.0}}, {{2000.0, 2000.0, 2000.0, 2500.0, 2500.0, 2500.0}},
        {{2000.0, 2000.0, 2000.0, 2500.0, 2500.0, 2500.0}}, {{2000.0, 2000.0, 2000.0, 2500.0, 2500.0, 2500.0}});

    franka::Model model = robot.loadModel();

    // const std::array<double, 7> d_gains = {{50.0, 50.0, 50.0, 50.0, 30.0, 25.0, 15.0}};
    const std::array<double, 7> d_gains = {{15.0, 15.0, 15.0, 15.0, 9.0, 8.0, 4.0}};
    // const std::array<double, 7> d_gains = {{5.0, 5.0, 5.0, 5.0, 3.0, 2.5, 1.5}};
    // const std::array<double, 7> d_gains = {{0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0}}; //TODO pure gravity compensation mode (with coriolis) performs very bad

    // Define callback for the joint torque control loop.
    std::function<franka::Torques(const franka::RobotState&, franka::Duration)>
        impedance_control_callback =
            [&print_data, &model, d_gains](
                const franka::RobotState& state, franka::Duration /*period*/) -> franka::Torques {

      const std::array<double, 7> gravity = model.gravity(state);
      const std::array<double, 7> coriolis = model.coriolis(state);
      const std::array<double, 42> jacobian_array = model.zeroJacobian(franka::Frame::kEndEffector, state);

      // Compute torque command
      std::array<double, 7> tau_d_calculated;
      for (size_t i = 0; i < 7; i++) {
        tau_d_calculated[i] = - d_gains[i] * state.dq[i] + coriolis[i];
      }

      // The following line is only necessary for printing the rate limited torque. As we activated
      // rate limiting for the control loop (activated by default), the torque would anyway be
      // adjusted!
      std::array<double, 7> tau_d_rate_limited =
          franka::limitRate(franka::kMaxTorqueRate, tau_d_calculated, state.tau_J_d);

      // Update data to print.
      if (print_data.mutex.try_lock()) {
        print_data.has_data = true;
        print_data.robot_state = state;
        print_data.tau_d_calculated = tau_d_calculated;
        print_data.tau_d_last = tau_d_rate_limited;
        print_data.gravity = gravity;
        print_data.coriolis = coriolis;
        print_data.jacobian_array = jacobian_array;
        print_data.mutex.unlock();
      }

      // Send torque command.
      return tau_d_rate_limited;
    };

    // Start real-time control loop.
    robot.control(impedance_control_callback);

  } catch (const franka::Exception& ex) {
    running = false;
    std::cerr << ex.what() << std::endl;
  }

  if (print_thread.joinable()) {
    print_thread.join();
  }
  return 0;
}


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

#ifndef CorrelationAcceptanceVisitor_HPP
#define CorrelationAcceptanceVisitor_HPP

#include <boost/graph/graph_traits.hpp>

// Parent class
#include "Visitors/AcceptanceVisitors/AcceptanceVisitorParent.h"
#include "Visitors/AcceptanceVisitors/VarianceFunctor.hpp"
#include "Visitors/AcceptanceVisitors/AverageFunctor.hpp"

// Custom
#include "ITKHelpers/ITKHelpers.h"

// ITK
#include "itkAddImageFilter.h"
#include "itkMultiplyImageFilter.h"
#include "itkImageRegion.h"

/**

 */
template <typename TGraph, typename TImage>
struct CorrelationAcceptanceVisitor : public AcceptanceVisitorParent<TGraph>
{
  TImage* Image;
  Mask* MaskImage;

  const unsigned int HalfWidth;

  float Threshold;

  typedef typename boost::graph_traits<TGraph>::vertex_descriptor VertexDescriptorType;

  CorrelationAcceptanceVisitor(TImage* const image, Mask* const mask,
                               const unsigned int halfWidth, const float threshold, const std::string& visitorName = "CorrelationAcceptanceVisitor") :
  AcceptanceVisitorParent<TGraph>(visitorName), Image(image), MaskImage(mask),
    HalfWidth(halfWidth), Threshold(threshold)
  {

  }

  bool AcceptMatch(VertexDescriptorType target, VertexDescriptorType source, float& computedEnergy) const override
  {
    itk::Index<2> targetPixel = ITKHelpers::CreateIndex(target);
    itk::ImageRegion<2> targetRegion = ITKHelpers::GetRegionInRadiusAroundPixel(targetPixel, HalfWidth);

    itk::Index<2> sourcePixel = ITKHelpers::CreateIndex(source);
    itk::ImageRegion<2> sourceRegion = ITKHelpers::GetRegionInRadiusAroundPixel(sourcePixel, HalfWidth);

    typedef itk::Image<float, 2> FloatImageType;

    typedef itk::VectorMagnitudeImageFilter<TImage, FloatImageType> VectorMagnitudeFilterType;
    typename VectorMagnitudeFilterType::Pointer magnitudeFilter = VectorMagnitudeFilterType::New();
    magnitudeFilter->SetInput(Image);
    magnitudeFilter->Update();

    std::vector<itk::Offset<2> > validOffsets = MaskImage->GetValidOffsetsInRegion(targetRegion);

    FloatImageType::Pointer sourceImage = FloatImageType::New();
//     sourceImage->SetRegions(ITKHelpers::CornerRegion(sourceRegion.GetSize()));
//     sourceImage->Allocate();
    ITKHelpers::ExtractRegion(magnitudeFilter->GetOutput(), sourceRegion, sourceImage.GetPointer());

    FloatImageType::Pointer targetImage = FloatImageType::New();
//     sourceImage->SetRegions(ITKHelpers::CornerRegion(targetRegion.GetSize()));
//     sourceImage->Allocate();
    ITKHelpers::ExtractRegion(magnitudeFilter->GetOutput(), targetRegion, targetImage.GetPointer());

    std::vector<itk::Index<2> > validIndices = ITKHelpers::OffsetsToIndices(validOffsets);

    VarianceFunctor varianceFunctor;
    AverageFunctor averageFunctor;
    /////////// Target region //////////
    std::vector<FloatImageType::PixelType> validPixelsTargetRegion = ITKHelpers::GetPixelValues(targetImage.GetPointer(), validIndices);
    typename TypeTraits<FloatImageType::PixelType>::LargerType targetMean = averageFunctor(validPixelsTargetRegion);
    typename TypeTraits<FloatImageType::PixelType>::LargerType targetStandardDeviation = sqrt(varianceFunctor(validPixelsTargetRegion));

    typedef itk::AddImageFilter <FloatImageType, FloatImageType, FloatImageType> AddImageFilterType;
    AddImageFilterType::Pointer targetAddImageFilter = AddImageFilterType::New();
    targetAddImageFilter->SetInput(targetImage);
    targetAddImageFilter->SetConstant2(-1.0f * targetMean);
    targetAddImageFilter->Update();

    typedef itk::MultiplyImageFilter<FloatImageType, FloatImageType, FloatImageType> MultiplyImageFilterType;
    MultiplyImageFilterType::Pointer targetMultiplyImageFilter = MultiplyImageFilterType::New();
    targetMultiplyImageFilter->SetInput(targetImage);
    targetMultiplyImageFilter->SetConstant(1.0f/targetStandardDeviation);
    targetMultiplyImageFilter->Update();

    /////////// Source region //////////
    std::vector<FloatImageType::PixelType> validPixelsSourceRegion = ITKHelpers::GetPixelValues(sourceImage.GetPointer(), validIndices);
    typename TypeTraits<FloatImageType::PixelType>::LargerType sourceMean = averageFunctor(validPixelsSourceRegion);
    typename TypeTraits<FloatImageType::PixelType>::LargerType sourceStandardDeviation = sqrt(varianceFunctor(validPixelsSourceRegion));

    AddImageFilterType::Pointer sourceAddImageFilter = AddImageFilterType::New();
    sourceAddImageFilter->SetInput(sourceImage);
    sourceAddImageFilter->SetConstant2(-1.0f * sourceMean);
    sourceAddImageFilter->Update();

    MultiplyImageFilterType::Pointer sourceMultiplyImageFilter = MultiplyImageFilterType::New();
    sourceMultiplyImageFilter->SetInput(sourceImage);
    sourceMultiplyImageFilter->SetConstant(1.0f/sourceStandardDeviation);
    sourceMultiplyImageFilter->Update();

    // Initialize
    computedEnergy = 0.0f;

    for(std::vector<itk::Index<2> >::const_iterator iter = validIndices.begin(); iter != validIndices.end(); ++iter)
    {
      computedEnergy += (sourceMultiplyImageFilter->GetOutput()->GetPixel(*iter) * targetMultiplyImageFilter->GetOutput()->GetPixel(*iter));
    }

    computedEnergy /= static_cast<float>(validIndices.size());

    if(computedEnergy < Threshold)
      {
      std::cout << this->VisitorName << ": Match accepted (" << computedEnergy << " is less than " << Threshold << ")" << std::endl << std::endl;
      return true;
      }
    else
      {
      std::cout << this->VisitorName << ": Match rejected (" << computedEnergy << " is greater than " << Threshold << ")" << std::endl << std::endl;
      return false;
      }
  };

};

#endif


//---------------------------------------------------------------------------//
//!
//! \file   Utility_PQLAQuadrature.hpp
//! \author Philip Britt
//! \brief  PQLA Direction Quadrature handler declaration
//!
//---------------------------------------------------------------------------//

#ifndef UTILITY_PQLA_QUADRATURE
#define UTILITY_PQLA_QUADRATURE

// Boost Includes
#include <boost/serialization/version.hpp>
#include <boost/serialization/export.hpp>
#include <boost/serialization/shared_ptr.hpp>

// FRENSIE includes
#include "Utility_Vector.hpp"
#include "Utility_Tuple.hpp"
#include "Utility_Array.hpp"
#include "Utility_ExplicitSerializationTemplateInstantiationMacros.hpp"
#include "Utility_SerializationHelpers.hpp"

namespace Utility{

struct SphericalTriangle
{
  /*! Vector that contains a tuple representing spherical triangle parameters.
   * \details First element of the tuple is an array that contains the
   * 2-norm direction representing a vertex of the triangle.
   * Second is the length of the spherical triangle side opposite from that vertex
   * (or angle that the 2 other vertices of the triangle make with each other).
   * Third is the angle made from the sides of the spherical triangle from that vertex.
   * Note their order does NOT matter so long as they are consistent with the above definition
   */
  std::vector<std::tuple<std::array<double, 3>, double, double>> triangle_parameter_vector;

  //! Area of the triangle
  double area;

    //! Processes spherical triangle information
  void computeAndStoreTriangleParameters(std::vector<std::array<double, 3>>& vertex_vector);

  // Serialize the data
  template<typename Archive>
  void serialize(Archive& ar, const unsigned version)
  {
    ar & BOOST_SERIALIZATION_NVP(triangle_parameter_vector);
    ar & BOOST_SERIALIZATION_NVP(area);
  }
};

class PQLAQuadrature
{

  public:

  //! Constructor
  PQLAQuadrature(unsigned quadrature_order);

  //! Destructor
  ~PQLAQuadrature()
  { /* ... */ }

  //! Find which triangle bin a direction vector is in
  size_t findTriangleBin(const std::array<double, 3>& direction) const;

  //! Find which triangle bin a direction vector is in
  size_t findTriangleBin(const double x_direction, const double y_direction, const double z_direction) const;

  //! Return the order of the quadrature
  unsigned getQuadratureOrder() const;

  //! Get the total number of triangles
  size_t getNumberOfTriangles() const;

  //! Get the area of a specific spherical triangle
  double getTriangleArea(const size_t triangle_index) const;

  /*! Get a random direction from within a spherical triangle (evenly distributed probability) - reference here
   * \details reference: Stratified Sampling of Spherical Triangles, James Arvo, SIGGRAPH '95
   */
  void sampleIsotropicallyFromTriangle(std::array<double, 3>& direction_vector,
                                       const size_t triangle_index) const;

  const std::vector<SphericalTriangle>& getSphericalTriangleVector() const;

  private:

  //! Default constructor (for archiving)
  PQLAQuadrature()
  { /* ... */ }

  //! Vector operation for the purpose of sampleIsotropicallyFromTriangle
  void isotropicSamplingVectorOperation(const std::array<double, 3>& vertex_1,
                                        const std::array<double, 3>& vertex_2,
                                        std::array<double, 3>& result_vector) const;

  //! Converts direction vector to 1-norm normalized vector
  void normalizeVectorToOneNorm(const std::array<double, 3>& direction_2_norm,
                                 std::array<double, 3>& direction_1_norm) const;

  //! Converts direction vector to 1-norm normalized vector
  void normalizeVectorToOneNorm(const double x_direction,
                                  const double y_direction,
                                  const double z_direction,
                                  std::array<double, 3>& direction_1_norm) const;

  //! Take lower bounding plane indices of direction vector to form triangle index
  size_t calculatePositiveTriangleBinIndex(const unsigned i_x, const unsigned i_y, const unsigned i_z) const;

  //! Take direction signs to calculate secondary index
  size_t findSecondaryIndex(const bool x_sign, const bool y_sign, const bool z_sign) const;

  //! Quadrature order
  unsigned d_quadrature_order;

  //! Vector that stores POSITIVE DOMAIN spherical triangles
  std::vector<SphericalTriangle> d_spherical_triangle_vector;

  //! Serialize the data
  template<typename Archive>
  void serialize(Archive& ar, const unsigned version);

  //! Declare the boost serialization access object as a friend
  friend class boost::serialization::access;

};

// Serialize the data
template<typename Archive>
void PQLAQuadrature::serialize(Archive& ar, const unsigned version)
{
  // Serialize the member data
  ar & BOOST_SERIALIZATION_NVP(d_quadrature_order);
  ar & BOOST_SERIALIZATION_NVP(d_spherical_triangle_vector);
}

} // end Utility namespace

BOOST_SERIALIZATION_CLASS_VERSION(PQLAQuadrature, Utility, 0);
BOOST_SERIALIZATION_CLASS_EXPORT_STANDARD_KEY(PQLAQuadrature, Utility);
EXTERN_EXPLICIT_CLASS_SERIALIZE_INST(Utility, PQLAQuadrature);

#endif // end UTILITY_PQLA_QUADRATURE

//---------------------------------------------------------------------------//
// end Utility_PQLADiscetization.hpp
//---------------------------------------------------------------------------//