/**
 * \file dcs/math/stats/distribution/map.hpp
 *
 * \brief Markov Arrival Process (MAP).
 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 *
 *

 *
 * Copyright 2009 Marco Guazzone (marco.guazzone@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
 *
 * 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 DCS_MATH_STATS_DISTRIBUTION_MAP_HPP
#define DCS_MATH_STATS_DISTRIBUTION_MAP_HPP

#include <boost/numeric/ublas/expression_types.hpp>
#include <boost/numeric/ublas/io.hpp>
#include <boost/numeric/ublas/matrix.hpp>
#include <boost/numeric/ublas/matrix_expression.hpp>
#include <boost/numeric/ublas/matrix_proxy.hpp>
#include <boost/numeric/ublas/vector.hpp>
#include <boost/numeric/ublas/traits.hpp>
#include <boost/numeric/ublasx/operation/abs.hpp>
#include <boost/numeric/ublasx/operation/cumsum.hpp>
#include <boost/numeric/ublasx/operation/diag.hpp>
#include <boost/numeric/ublasx/operation/inv.hpp>
#include <boost/numeric/ublasx/operation/lu.hpp>
#include <boost/numeric/ublasx/operation/max.hpp>
#include <boost/numeric/ublasx/operation/num_columns.hpp>
#include <boost/numeric/ublasx/operation/num_rows.hpp>
#include <boost/numeric/ublasx/operation/sum.hpp>
#include <boost/numeric/ublasx/operation/which.hpp>
#include <cmath>
#include <cstddef>
#include <dcs/assert.hpp>
#include <dcs/debug.hpp>
#include <dcs/math/policies/policy.hpp>
#include <dcs/math/stats/distribution/base_distribution.hpp>
#include <dcs/math/stats/distribution/exponential.hpp>
#include <dcs/math/stats/function/rand.hpp>
#include <dcs/math/random/any_generator.hpp>
#include <dcs/math/random/base_generator.hpp>
#include <dcs/math/random/uniform_01_adaptor.hpp>
//#include <dcs/math/random/uniform_int_adaptor.hpp>
#include <functional>
#include <iostream>
#include <vector>
#include <stdexcept>

namespace dcs { namespace math { namespace stats {

namespace detail { namespace /*<unnamed>*/ {

/**
 * \brief Equilibrium distribution of a continuous-time Markov chain (CTMC).
 * \param Q Infinitesimal generator of the CTMC.
 * \return The equilibrium distribution of the CTMC.
 *
 * Solve the following system:
 * \f{align}
 * \mathbf{\pi}\mathbf{Q}=\mathbf{0} \\
 * \mathbf{\pi}\mathbf{1}^T=1
 * \f}
 * where:
 * - the first equation is the global balance equation,
 * - the * second one is the normalization condition,
 * - and \f$\mathbf{\pi}\f$ is the steady state probabilities vector
 *   (expressed as a row vector).
 * .
 *
 * \note
 *  Actually, we solve the CTMC by means of the direct method (see
 *  [1,2,3] for possible issues and other existing methods).
 *  From the point of view of implementation, there are two main problems:
 *  -# Maths routing usually assumes that vectors are expressed as column
 *     vector.
 *     So, the global balance equation comes in the "wrong" form: the
 *     unknowns are given as a row vector not a column vector.
 *     This problem is resolved by transposing the equation, i.e.
 *     \f[
 *      \mathbf{Q}^T\mathbf{\pi}^T = \mathbf{0}^T
 *     \f]
 *     where the right hand side is now a column vector of zeros, rather than a
 *     row vector.
 *  -# The redundancy amongst the global balance equations.
 *     To solve this problem we need the additional information of the
 *     normalization condition.
 *     This problem is resolved by replacing one of the global balance equations
 *     by the normalization condition.
 *     In the transposed matrix this corresponds to replacing one row by a row
 *     of 1's.
 *     We usually choose the last row and denote the modified matrix Q_n^T .
 *     Similarly we change the constant terms vector, which was all zeros to be
 *     a column vector with 1 in the last row, and zeros everywhere else.
 *     We denote such a vector, \f$e_n^T\f$.
 *  .
 *
 * References:
 * -# W.J. Stewart.
 *    "Introduction to the Numerical Solution of Markov Chains",
 *    Princeton University Press, 1994.
 * -# W.J. Stewart.
 *    "Performance Modeling and Markov Chains",
 *    Formal Methods of Performance Evaluation, Springer, 2007.
 * -# W.J. Stewart.
 *    "Probability, Markov Chains, Queues, and Simulation",
 *    Princeton University Press, 2009.
 * .
 */
template <typename MatrixExprT>
::boost::numeric::ublas::vector<
	typename ::boost::numeric::ublas::matrix_traits<MatrixExprT>::value_type
> ctmc_solve(::boost::numeric::ublas::matrix_expression<MatrixExprT> const& Q)
{
	namespace ublas = ::boost::numeric::ublas;
	namespace ublasx = ::boost::numeric::ublasx;

	typedef typename ublas::matrix_traits<MatrixExprT>::value_type value_type;
	typedef typename ublas::matrix_traits<MatrixExprT>::size_type size_type;
	typedef ublas::vector<size_type> size_vector_type;
	typedef ublas::vector<value_type> vector_type;

	size_type nr = ublasx::num_rows(Q);
	//size_type nc = ublasx::num_columns(Q);

	size_vector_type z;

	// Is the CTMC irreducible?

	z = ublasx::which(
		ublasx::sum<1>(
			ublasx::abs(Q)
		),
		::std::bind2nd(
			::std::equal_to<value_type>(),
			value_type(0)
		)
	);

	size_type nz;
	nz = ublasx::size(z);

	if (nz > 1)
	{
		DCS_DEBUG_TRACE("Warning: Q is a reducible infinitesimal generator.");
//		b = ublas::zero_vector<value_type>(nr);
	}
	else if (nz == 0)
	{
		z.resize(1, false);
		z(0) = nr-1;
		nz = 1;
	}

	// Create the constant terms vecotr b=e_n^T such that Q_n^T\pi^T=e_n^T
	vector_type b;
	b = ublas::zero_vector<value_type>(nr);

	ublas::matrix<value_type> tmp_Q(Q);

	// normalization condition
	for (size_type j = 0; j < nz; ++j)
	{
		size_type k(z(j));

		b(k) = 1;

		for (size_type i = 0; i < nr; ++i)
		{
			tmp_Q(i,k) = value_type(1);
		}
	}

	ublasx::lu_solve_inplace(ublas::trans(tmp_Q), b);

	return ublas::trans(b);
}

/**
 * \brief Equilibrium distribution of a discrete-time Markov chain (DTMC).
 * \param P Stochastic transition matrix of the DTMC.
 * \return The equilibrium distribution of the DTMC.
 */
template <typename MatrixExprT>
::boost::numeric::ublas::vector<
	typename ::boost::numeric::ublas::matrix_traits<MatrixExprT>::value_type
> dtmc_solve(::boost::numeric::ublas::matrix_expression<MatrixExprT> const& P)
{
	namespace ublas = ::boost::numeric::ublas;
	namespace ublasx = ::boost::numeric::ublasx;

	typedef MatrixExprT matrix_type;
	typedef typename ublas::matrix_traits<matrix_type>::size_type size_type;
	typedef typename ublas::matrix_traits<matrix_type>::value_type value_type;

	size_type nr = ublasx::num_rows(P);
	size_type nc = ublasx::num_columns(P);

	ublas::identity_matrix<value_type> I(nr,nc);

	return ctmc_solve(P-I);
}

/**
 * \brief Embedded discrete-time process of the given MAP.
 *
 * \param D0 The \f$D_0\f$ matrix of the MAP.
 * \param D1 The \f$D_1\f$ matrix of the MAP.
 * \return The probability transition matrix of the embedded process.
 *
 * Compute the stochastic transition matrix \f$P\f$ of the embedded phase
 * process at arrival instants.
 * Specifically, \f$P_{ij} is the probability that the MAP restarts in phase
 * \f$j\f$ if the last absorption occurred in phase \f$i\f$.
 *
 * \note If the MAP is feasible, then the probability transition matrix must be
 *  an irreducible stochastic matrix.
 *  In this case, the stationary distribution at arrival instants denoted by
 *  \f$\pi\f$ is given by \f$\piP=\pi$ and \f$\pie^T=1\f$.
 *  After an event has been generated, the distribution of the MAP is given by
 *  \f$\pi\f$.
 */
template <typename D0MatrixExprT, typename D1MatrixExprT>
::boost::numeric::ublas::matrix<
	typename ::boost::numeric::ublas::promote_traits<
		typename ::boost::numeric::ublas::matrix_traits<D0MatrixExprT>::value_type,
		typename ::boost::numeric::ublas::matrix_traits<D1MatrixExprT>::value_type
	>::promote_type
> embedded_dtmc(::boost::numeric::ublas::matrix_expression<D0MatrixExprT> const& D0,
				::boost::numeric::ublas::matrix_expression<D1MatrixExprT> const& D1)
{

//	typedef typename ::boost::numeric::ublas::promote_traits<
//				typename ::boost::numeric::ublas::matrix_traits<D0MatrixExprT>::value_type,
//				typename ::boost::numeric::ublas::matrix_traits<D1MatrixExprT>::value_type
//			>::promote_type value_type;
//	typedef ::boost::numeric::ublas::matrix<value_type> matrix_type;

	// Compute P=(-D0)^{-1})*D1
	return ::boost::numeric::ublas::prod(
				::boost::numeric::ublasx::inv(-D0),
				D1
			);
}

/// Compute the equilibrium distribution of the embedded discrete-time process.
template <typename D0MatrixExprT, typename D1MatrixExprT>
::boost::numeric::ublas::vector<
	typename ::boost::numeric::ublas::promote_traits<
		typename ::boost::numeric::ublas::matrix_traits<D0MatrixExprT>::value_type,
		typename ::boost::numeric::ublas::matrix_traits<D1MatrixExprT>::value_type
	>::promote_type
> equilibrium_distribution(::boost::numeric::ublas::matrix_expression<D0MatrixExprT> const& D0,
						   ::boost::numeric::ublas::matrix_expression<D1MatrixExprT> const& D1)
{
	return dtmc_solve(embedded_dtmc(D0, D1));
}

template <typename ValueT, typename URNG>
inline
ValueT generate_probability(URNG& rng)
{
	::dcs::math::random::uniform_01_adaptor<URNG&,ValueT> u01_rng(rng);

	return u01_rng();
}

template <typename VectorExprT, typename URNG>
typename ::boost::numeric::ublas::vector_traits<VectorExprT>::size_type generate_initial_state(::boost::numeric::ublas::vector_expression<VectorExprT> const& pi, URNG& rng)
{
	namespace ublas = ::boost::numeric::ublas;
	namespace ublasx = ::boost::numeric::ublasx;

	typedef typename ublas::vector_traits<VectorExprT>::size_type size_type;
	typedef typename ublas::vector_traits<VectorExprT>::value_type value_type;

	size_type n = ublasx::size(pi);

	value_type p = generate_probability<value_type>(rng);

	size_type ret(n); //default to last state

	//NOTE: states start from 1!

	ublas::vector<value_type> cdf = ublasx::cumsum(pi);
	for (size_type i = 0; i < n; ++i)
	{
		if (p <= cdf(i))
		{
			ret = i+1;
			break;
		}
	}
// Alternative (but possibly more time-consuming)
//	ret = ublasx::min(
//			ublasx::which(
//				ublasx::cumsum(pi),
//				::std::bind1st(
//					::std::less_equal<value_type>(),
//					p
//				)
//			)
//		);

	return ret;
}

template <typename VectorExprT, typename URNG>
typename ::boost::numeric::ublas::vector_traits<VectorExprT>::size_type generate_next_state(::boost::numeric::ublas::vector_expression<VectorExprT> const& cdf, URNG& rng)
{
	namespace ublas = ::boost::numeric::ublas;
	namespace ublasx = ::boost::numeric::ublasx;

	typedef typename ublas::vector_traits<VectorExprT>::size_type size_type;
	typedef typename ublas::vector_traits<VectorExprT>::value_type value_type;

	value_type p = generate_probability<value_type>(rng);

	ublas::vector<size_type> ix;

	ix = ublasx::which(
			cdf,
			::std::bind2nd(::std::greater_equal<value_type>(), p)
		);

	if (ix.size() == 0)
	{
		//TODO
		throw ::std::runtime_error("Unexpected failure");
	}

	return ix(0);
}

template <typename ValueT, typename URNG>
ValueT generate_sojourn_time(ValueT rate, URNG& rng)
{
	::dcs::math::stats::exponential_distribution<ValueT> exp(rate);

	return ::dcs::math::stats::rand(exp, rng);
}

}} // Namespace detail::<unnamed>

template <
	typename RealT = double,
	typename PolicyT = ::dcs::math::policies::policy<>
>
class map_distribution//: public base_distribution<RealT>
{
//	private: typedef base_distribution<RealT> base_type;
//	public: typedef typename base_type::value_type value_type;
//	public: typedef typename base_type::support_type support_type;
	public: typedef RealT value_type;
	public: typedef value_type support_type;
	public: typedef ::boost::numeric::ublas::matrix<support_type> matrix_type;
	public: typedef ::std::size_t size_type;
	private: typedef ::boost::numeric::ublas::vector<support_type> value_vector_type;
	private: typedef ::boost::numeric::ublas::vector<size_type> index_vector_type;

	public: template <typename D0MatrixExprT, typename D1MatrixExprT>
		map_distribution(::boost::numeric::ublas::matrix_expression<D0MatrixExprT> const& D0,
						 ::boost::numeric::ublas::matrix_expression<D1MatrixExprT> const& D1)
	: //base_type(),
	  D0_(D0),
	  D1_(D1)
	{
		// pre: num_rows(D0) == num_rows(D1)
		DCS_ASSERT(
			::boost::numeric::ublasx::num_rows(D0) == ::boost::numeric::ublasx::num_rows(D1),
			throw ::std::invalid_argument("[dcs::math::stats::map_distribution::ctor] Matrices D0 and D1 have different number of rows.")
		);
		// pre: num_cols(D0) == num_cols(D1)
		DCS_ASSERT(
			::boost::numeric::ublasx::num_columns(D0) == ::boost::numeric::ublasx::num_columns(D1),
			throw ::std::invalid_argument("[dcs::math::stats::map_distribution::ctor] Matrices D0 and D1 have different number of columns.")
		);
	}

	// Compiler-generated copy-constructor, copy-assignment, and destructor
	// are fine.

	public: matrix_type const& D0() const
	{
		return D0_;
	}

	public: matrix_type const& D1() const
	{
		return D1_;
	}

/*
	private: value_type do_rand(::dcs::math::random::any_generator<value_type>& rng) const
	{
		return rand_sample(rng);
	}

	private: value_type do_rand(::dcs::math::random::base_generator<value_type>& rng) const
	{
		return rand_sample(rng);
	}
*/

	public: template <typename URNG>
		value_type rand(URNG& rng) const
	{
		::std::vector<value_type> samples;
		samples = rand(rng, 1);

		if (samples.size() == 0)
		{
			throw ::std::runtime_error("[dcs::math::stats::map_distribution::rand] Unexpected failure.");
		}

		return samples[0];
	}

	/**
	 * The following algorithm is an adaptation of the one found in the MAP-QN
	 * toolbox (by G. Casale et al, http://www.cs.wm.edu/MAPQN/).
	 */
	public: template <typename URNG>
		::std::vector<value_type> rand(URNG& rng, size_type n) const
	{
		namespace ublas = ::boost::numeric::ublas;
		namespace ublasx = ::boost::numeric::ublasx;

		const size_type max_path_len0(20); // initial max path length
		const size_type path_len_incr(5); // increment in the max path length

		size_type ns(ublasx::num_rows(D0_)); // # stages
		size_type ns2(ns << 1); // 2*ns

		value_vector_type pi;

		pi = detail::equilibrium_distribution(D0_,  D1_);

		size_type s0; // initial stage

		// Randomly generate the initial stage.
		s0 = detail::generate_initial_state(pi, rng);

		matrix_type P(ns,ns2,0);

		for (size_type b=0; b < 2; ++b)
		{
			for (size_type i = 0; i < ns; ++i)
			{
				for (size_type j = 0; j < ns; ++j)
				{
					P(i,b*ns+j) = (b==0 ? D0_(i,j) : D1_(i,j)) / ::std::abs(D0_(i,i));
				}
			}
		}
		for (size_type i=0; i < ns; ++i)
		{
			P(i,i) = 0;
		}

		// Create CDF
		matrix_type C(ns,ns2);
		for (size_type i = 0; i < ns; ++i)
		{
			ublas::row(C, i) = ublasx::cumsum(ublas::row(P, i));
		}
		C = ublasx::abs(C);

		size_type src_s(s0);
		matrix_type V(ublas::zero_matrix<value_type>(n, max_path_len0));
		size_type max_path_len(max_path_len0);
		for (size_type i = 0; i < n; ++i)
		{
			bool arrival(false);
			size_type last(1);
			V(i,0) = src_s;

			while (!arrival)
			{
				size_type dst_s;
				dst_s = detail::generate_next_state(
							ublas::row(C, src_s-1),
							rng
					);
				++dst_s; // make sure that stage id starts from 1
				if (dst_s > ns)
				{
					arrival = true;
					dst_s -= ns;
					src_s = dst_s;
				}
				else
				{
					V(i,last) = dst_s;
					src_s = dst_s;
					++last;
				}
				if (last >= max_path_len)
				{
					max_path_len += path_len_incr;
					V.resize(n, max_path_len, true);
					for (size_type c = last; c < max_path_len; ++c)
					{
						ublas::column(V, c) = ublas::zero_vector<size_type>(n);
					}
				}
			}
		}

//XXX: not used
//		// Create the vector of absorbing stages (for each sample)
//		value_vector_type L(n);
//
//		for (size_type i = 0; i < n; ++i)
//		{
//			size_type ix;
//
//			ix = ublasx::max(
//					ublasx::which(
//						ublas::row(V, i)
//					)
//				);
//			L(i) = V(i,ix);
//		}

//XXX: not used
//		// Create the vector of initial stages (for each sample)
//		value_vector_type F;
//
//		F = ublas::column(V, 1);

		value_vector_type hold_rates(-ublasx::diag(D0_));

		size_type nv = ublasx::num_columns(V);
		matrix_type H(V); // matrix of times to absorptions
		for (size_type i = 0; i < ns; ++i)
		{
			size_type s = i+1;
			for (size_type j = 0; j < n; ++j)
			{
				for (size_type k = 0; k < nv; ++k)
				{
					if (V(j,k) == s)
					{
						H(j,k) = hold_rates(i);
					}
				}
			}
		}

		::std::vector<value_type> samples(n, 0);

		for (size_type i = 0; i < nv; ++i)
		{
			for (size_type k = 0; k < n; ++k)
			{
				if (H(k,i) > 0)
				{
					samples[k] += detail::generate_sojourn_time<value_type>(H(k,i), rng);
				}
			}
		}

		return samples;
	}

	private: matrix_type D0_;
	private: matrix_type D1_;
};

template <
	typename CharT,
	typename CharTraitsT,
	typename RealT,
	typename PolicyT
>
::std::basic_ostream<CharT,CharTraitsT>& operator<<(::std::basic_ostream<CharT,CharTraitsT>& os, map_distribution<RealT,PolicyT> const& dist)
{
	return os << "MAP("
			  << "D0=" <<  dist.D0()
			  << ",D1=" <<  dist.D1()
			  << ")";
}

}}} // Namespace dcs:math::stats

#endif // DCS_MATH_STATS_DISTRIBUTION_MAP_HPP


/* Copyright (C) 2012-2017 IBM Corp.
 * This program is 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. See accompanying LICENSE file.
 */
/*
 * @file eqtesting.cpp
 * @brief Useful fucntions for equality testing...
 */
#include <NTL/lzz_pXFactoring.h>
NTL_CLIENT
#include "FHE.h"
#include "timing.h"
#include "EncryptedArray.h"

#include <cassert>
#include <cstdio>

// Map all non-zero slots to 1, leaving zero slots as zero.
// Assumes that r=1, and that all the slot contain elements from GF(p^d).
//
// We compute x^{p^d-1} = x^{(1+p+...+p^{d-1})*(p-1)} by setting y=x^{p-1}
// and then outputting y * y^p * ... * y^{p^{d-1}}, with exponentiation to
// powers of p done via Frobenius.
void mapTo01(const EncryptedArray& ea, Ctxt& ctxt)
{
  long p = ctxt.getPtxtSpace();
  if (p != ctxt.getContext().zMStar.getP()) // ptxt space is p^r for r>1
    std::logic_error("mapTo01 not implemented for r>1");

  if (p>2)
    ctxt.power(p-1); // set y = x^{p-1}

  long d = ea.getDegree();
  if (d>1) { // compute the product of the d automorphisms
    std::vector<Ctxt> v(d, ctxt);
    for (long i=1; i<d; i++)
      v[i].frobeniusAutomorph(i);
    totalProduct(ctxt, v);
  }
}

// computes ctxt^{2^d-1} using a method that takes
// O(log d) automorphisms and multiplications
void fastPower(Ctxt& ctxt, long d)
{
  assert(ctxt.getPtxtSpace()==2);
  if (d <= 1) return;

  Ctxt orig = ctxt;

  long k = NumBits(d);
  long e = 1;

  for (long i = k-2; i >= 0; i--) {
    Ctxt tmp1 = ctxt;
    tmp1.smartAutomorph(1L << e);
    ctxt.multiplyBy(tmp1);
    e = 2*e;

    if (bit(d, i)) {
      ctxt.smartAutomorph(2);
      ctxt.multiplyBy(orig);
      e += 1;
    }
  }
}

// ===> This function only works for p=2, r=1 <===
// Test if prefixes of bits in slots are all zero: Set slot j of res[i] to 0
// if bits 0..i of j'th slot in ctxt are all zero, else it is set to 1
// It is assumed that res and the res[i]'s are initialized by the caller.
// Complexity: O(d + n log d) smart automorphisms
//             O(n d)
void incrementalZeroTest(Ctxt* res[], const EncryptedArray& ea,
			 const Ctxt& ctxt, long n)
{
  FHE_TIMER_START;
  long nslots = ea.size();
  long d = ea.getDegree();

  // compute linearized polynomial coefficients

  vector< vector<ZZX> > Coeff;
  Coeff.resize(n);

  for (long i = 0; i < n; i++) {
    // coeffients for mask on bits 0..i
    // L[j] = X^j for j = 0..i, L[j] = 0 for j = i+1..d-1

    vector<ZZX> L;
    L.resize(d);

    for (long j = 0; j <= i; j++)
      SetCoeff(L[j], j);

    vector<ZZX> C;

    ea.buildLinPolyCoeffs(C, L);

    Coeff[i].resize(d);
    for (long j = 0; j < d; j++) {
      // Coeff[i][j] = to the encoding that has C[j] in all slots
      // FIXME: maybe encrtpted array should have this functionality
      //        built in
      vector<ZZX> T;
      T.resize(nslots);
      for (long s = 0; s < nslots; s++) T[s] = C[j];
      ea.encode(Coeff[i][j], T);
    }
  }

  vector<Ctxt> Conj(d, ctxt);
  // initialize Cong[j] to ctxt^{2^j}
  for (long j = 0; j < d; j++) {
    Conj[j].smartAutomorph(1L << j);
  }

  for (long i = 0; i < n; i++) {
    res[i]->clear();
    for (long j = 0; j < d; j++) {
      Ctxt tmp = Conj[j];
      tmp.multByConstant(Coeff[i][j]);
      *res[i] += tmp;
    }

    // *res[i] now has 0..i in each slot
    // next, we raise to the power 2^d-1

    fastPower(*res[i], d);
  }
  FHE_TIMER_STOP;
}

#ifdef DEBUG_TEST
/************************** debugging code below *********************/
void printBits(const vector<ZZX>& v, long n)
{
  long len = v.size();
  if (n>50 || len>32) return;
  for (long i = 0; i < len; i++) {
    for (long j = n-1; j >= 0; j--)
      cout << coeff(v[i], j);
    cout << " ";
  }
  cout << "\n";
}

void  TestIt(long c, long k, long w, long L, long m, long n)
{
  FHEcontext context(m, 2, 1); // p = 2, r = 1
  long d = context.zMStar.getOrdP();

  buildModChain(context, L, c);

  context.zMStar.printout();
  cerr << endl;
#ifdef DEBUG
  cerr << context << endl;
#endif

  FHESecKey secretKey(context);
  const FHEPubKey& publicKey = secretKey;
  secretKey.GenSecKey(w); // A Hamming-weight-w secret key

  ZZX G;

  G = makeIrredPoly(2, d);
  // G = context.alMod.getFactorsOverZZ()[0];

  cerr << "generating key-switching matrices... ";
  addFrbMatrices(secretKey);
  addSome1DMatrices(secretKey);
  cerr << "done\n";

  cerr << "computing masks and tables for rotation...";
  EncryptedArray ea(context, G);
  cerr << "done\n";

  long nslots = ea.size();

  if (n <= 0 || n > d) n = d;

  vector<ZZX> v;
  v.resize(nslots);
  for (long i = 0; i < nslots; i++) {
    GF2X f;
    random(f, n);
    conv(v[i], f);
  }

  printBits(v, n);

  Ctxt ctxt(publicKey);
  ea.encrypt(ctxt, publicKey, v);
  // ctxt encrypts a vector where each slots is a random
  // polynomial of degree < n

  Ctxt* res[n];
  for (long j = 0; j < n; j++) res[j] = new Ctxt(publicKey); // allocate

  resetAllTimers();

  incrementalZeroTest(res, ea, ctxt, n);

  for (long j = 0; j < n; j++) {
    vector<ZZX> v1;
    ea.decrypt(*res[j], secretKey, v1);
    printBits(v1, n);
  }

  for (long j = 0; j < n; j++) delete res[j]; // cleanup
}

void usage(char *prog)
{
  cerr << "Usage: "<<prog<<" [ optional parameters ]...\n";
  cerr << "  optional parameters have the form 'attr1=val1 attr2=val2 ...'\n";
  cerr << "  e.g, 'R=4 L=9 k=80'\n\n";
  cerr << "  c is number of columns in the key-switching matrices [default=2]\n";
  cerr << "  k is the security parameter [default=80]\n";
  cerr << "  L is the # of primes in the modulus chain [default=20]\n";
  cerr << "  m is a specific modulus\n";
  cerr << "  n is the number of masks (defaults to all)\n";
  exit(0);
}

int main(int argc, char *argv[])
{
  argmap_t argmap;
  argmap["c"] = "2";
  argmap["k"] = "80";
  argmap["L"] = "20";
  argmap["m"] = "0";
  argmap["n"] = "0";

  // get parameters from the command line
  if (!parseArgs(argc, argv, argmap)) usage(argv[0]);

  long c = atoi(argmap["c"]);
  long k = atoi(argmap["k"]);
  long L = atoi(argmap["L"]);
  long chosen_m = atoi(argmap["m"]);
  long n = atoi(argmap["n"]);

  long w = 64; // Hamming weight of secret key
  //  long L = z*R; // number of levels

  long m = FindM(k, L, c, 2, 1, 0, chosen_m, true);

  setTimersOn();
  TestIt(c, k, w, L, m, n);

  cerr << endl;
  printAllTimers();
  cerr << endl;

}
// call to get our running test case:
// eqtesting_x m=20485
// eqtesting_x m=105 for quick testing

#endif // #ifdef DEBUG_TEST


///////////////////////////////////////////////////////////////////////////////
//  Copyright 2014 Anton Bikineev
//  Copyright 2014 Christopher Kormanyos
//  Copyright 2014 John Maddock
//  Copyright 2014 Paul Bristow
//  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)

#ifndef BOOST_MATH_HYPERGEOMETRIC_0F1_HPP
#define BOOST_MATH_HYPERGEOMETRIC_0F1_HPP

#include <boost/math/policies/policy.hpp>
#include <boost/math/policies/error_handling.hpp>
#include <boost/math/special_functions/detail/hypergeometric_series.hpp>
#include <boost/math/special_functions/detail/hypergeometric_0F1_bessel.hpp>

namespace boost { namespace math { namespace detail {

   template <class T>
   struct hypergeometric_0F1_cf
   {
      //
      // We start this continued fraction at b on index -1
      // and treat the -1 and 0 cases as special cases.
      // We do this to avoid adding the continued fraction result
      // to 1 so that we can accurately evaluate for small results
      // as well as large ones.  See http://functions.wolfram.com/07.17.10.0002.01
      //
      T b, z;
      int k;
      hypergeometric_0F1_cf(T b_, T z_) : b(b_), z(z_), k(-2) {}
      typedef std::pair<T, T> result_type;

      result_type operator()()
      {
         ++k;
         if (k <= 0)
            return std::make_pair(z / b, 1);
         return std::make_pair(-z / ((k + 1) * (b + k)), 1 + z / ((k + 1) * (b + k)));
      }
   };

   template <class T, class Policy>
   T hypergeometric_0F1_cf_imp(T b, T z, const Policy& pol, const char* function)
   {
      hypergeometric_0F1_cf<T> evaluator(b, z);
      boost::uintmax_t max_iter = policies::get_max_series_iterations<Policy>();
      T cf = tools::continued_fraction_b(evaluator, policies::get_epsilon<T, Policy>(), max_iter);
      policies::check_series_iterations<T>(function, max_iter, pol);
      return cf;
   }

   template <class T, class Policy>
   inline T hypergeometric_0F1_imp(const T& b, const T& z, const Policy& pol)
   {
      const char* function = "boost::math::hypergeometric_0f1<%1%,%1%>(%1%, %1%)";
      BOOST_MATH_STD_USING

         // some special cases
         if (z == 0)
            return T(1);

      if ((b <= 0) && (b == floor(b)))
         return policies::raise_pole_error<T>(
            function,
            "Evaluation of 0f1 with nonpositive integer b = %1%.", b, pol);

      if (z < -5 && b > -5)
      {
         // Series is alternating and divergent, need to do something else here,
         // Bessel function relation is much more accurate, unless |b| is similarly
         // large to |z|, otherwise the CF formula suffers from cancellation when
         // the result would be very small.
         if (fabs(z / b) > 4)
            return hypergeometric_0F1_bessel(b, z, pol);
         return hypergeometric_0F1_cf_imp(b, z, pol, function);
      }
      // evaluation through Taylor series looks
      // more precisious than Bessel relation:
      // detail::hypergeometric_0f1_bessel(b, z, pol);
      return detail::hypergeometric_0F1_generic_series(b, z, pol);
   }

} // namespace detail

template <class T1, class T2, class Policy>
inline typename tools::promote_args<T1, T2>::type hypergeometric_0F1(T1 b, T2 z, const Policy& /* pol */)
{
   BOOST_FPU_EXCEPTION_GUARD
      typedef typename tools::promote_args<T1, T2>::type result_type;
   typedef typename policies::evaluation<result_type, Policy>::type value_type;
   typedef typename policies::normalise<
      Policy,
      policies::promote_float<false>,
      policies::promote_double<false>,
      policies::discrete_quantile<>,
      policies::assert_undefined<> >::type forwarding_policy;
   return policies::checked_narrowing_cast<result_type, Policy>(
      detail::hypergeometric_0F1_imp<value_type>(
         static_cast<value_type>(b),
         static_cast<value_type>(z),
         forwarding_policy()),
      "boost::math::hypergeometric_0F1<%1%>(%1%,%1%)");
}

template <class T1, class T2>
inline typename tools::promote_args<T1, T2>::type hypergeometric_0F1(T1 b, T2 z)
{
   return hypergeometric_0F1(b, z, policies::policy<>());
}

} } // namespace boost::math

#endif // BOOST_MATH_HYPERGEOMETRIC_HPP


/*
 * Copyright (c) 2009 Carnegie Mellon University.
 *     All rights reserved.
 *
 *  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.
 *
 *
 */

/**
 * This implements the classical "k-means" clustering algorithm.
 *
 * It takes as input file a series of lines where each line is a comma separated
 * or space separated list of values representing a vector. For instance:
 *
 * \verbatim
 * 1.1, 1.5, 0.9
 * 0.3, 0.4, -1.1
 * ...
 * \endverbatim
 *
 * It constructs a graph with a single vertex for each data point and simply
 * uses the "Map-Reduce" scheme to perform a k-means clustering of all
 * the datapoints.
 */

#include <boost/config/warning_disable.hpp>
#include <boost/spirit/include/qi.hpp>
#include <boost/spirit/include/phoenix_core.hpp>
#include <boost/spirit/include/phoenix_operator.hpp>
#include <boost/spirit/include/phoenix_stl.hpp>
#include <boost/tokenizer.hpp>

#include <limits>
#include <vector>
#include <map>
#include <iostream>
#include <stdlib.h>

#include <graphlab.hpp>

size_t NUM_CLUSTERS = 0;
bool IS_SPARSE = false;

struct cluster {
  cluster(): count(0), changed(false) { }
  std::vector<double> center;
  std::map<size_t, double> center_sparse;
  size_t count;
  bool changed;

  void save(graphlab::oarchive& oarc) const {
    oarc << center << count << changed << center_sparse;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> center >> count >> changed >> center_sparse;
  }
};

std::vector<cluster> CLUSTERS;

// the current cluster to initialize
size_t KMEANS_INITIALIZATION;

struct vertex_data{
  std::vector<double> point;
  std::map<size_t, double> point_sparse;
  size_t best_cluster;
  double best_distance;
  bool changed;

  void save(graphlab::oarchive& oarc) const {
    oarc << point << best_cluster << best_distance << changed << point_sparse;
  }
  void load(graphlab::iarchive& iarc) {
    iarc >> point >> best_cluster >> best_distance >> changed >> point_sparse;
  }
};

//use edges when edge weight file is given
struct edge_data {
  double weight;

  edge_data() :
      weight(0.0) {
  }
  explicit edge_data(double w) :
      weight(w) {
  }

  void save(graphlab::oarchive& oarc) const {
    oarc << weight;
  }
  void load(graphlab::iarchive& iarc) {
    iarc >> weight;
  }
};

// helper function to compute distance between points
double sqr_distance(const std::vector<double>& a,
                    const std::vector<double>& b) {
  ASSERT_EQ(a.size(), b.size());
  double total = 0;
  for (size_t i = 0;i < a.size(); ++i) {
    double d = a[i] - b[i];
    total += d * d;
  }
  return total;
}

double sqr_distance(const std::map<size_t, double>& a,
                    const std::map<size_t, double>& b) {
  double total = 0.0;
  for(std::map<size_t, double>::const_iterator iter = a.begin();
      iter != a.end(); ++iter){
    size_t id = (*iter).first;
    double val = (*iter).second;
    if(b.find(id) != b.end()){
      double d = val - b.at(id);
      total += d*d;
    }else{
      total += val * val;
    }
  }
  for(std::map<size_t, double>::const_iterator iter = b.begin();
      iter != b.end(); ++iter){
    double val = (*iter).second;
    if(a.find((*iter).first) == a.end()){
      total += val * val;
    }
  }

  return total;

////   cosine distance is better for sparse datapoints?
//    double ip = 0.0;
//    double lenA = 0.0;
//    double lenB = 0.0;
//    for(std::map<size_t, double>::const_iterator iter = a.begin();
//        iter != a.end(); ++iter){
//      size_t id = (*iter).first;
//      double val = (*iter).second;
//      if(b.find(id) != b.end()){
//        ip += val * b.at(id);
//      }
//      lenA += val*val;
//    }
//
//    if(ip == 0.0 || lenA == 0.0)
//      return 1.0;
//
//    for(std::map<size_t, double>::const_iterator iter = b.begin();
//        iter != b.end(); ++iter){
//      double val = (*iter).second;
//      lenB += val * val;
//    }
//
//    if(lenB == 1.0)
//      return 1.0;
//
//    return 1.0 - ip/(sqrt(lenA)*sqrt(lenB));

}

// helper function to add two vectors
std::vector<double>& plus_equal_vector(std::vector<double>& a,
                                       const std::vector<double>& b) {
  ASSERT_EQ(a.size(), b.size());
  for (size_t i = 0;i < a.size(); ++i) {
    a[i] += b[i];
  }
  return a;
}

// helper function to add two vectors
std::map<size_t, double>& plus_equal_vector(std::map<size_t, double>& a,
                                       const std::map<size_t, double>& b) {
  for(std::map<size_t, double>::const_iterator iter = b.begin();
    iter != b.end(); ++iter){
    size_t id = (*iter).first;
    double val = (*iter).second;
    if(a.find(id) != a.end()){
      a[id] += b.at(id);
    }else{
      a.insert(std::make_pair(id, val));
    }
  }
  return a;
}

// helper function to scale a vector vectors
std::vector<double>& scale_vector(std::vector<double>& a, double d) {
  for (size_t i = 0;i < a.size(); ++i) {
    a[i] *= d;
  }
  return a;
}

// helper function to scale a vector vectors
std::map<size_t, double>& scale_vector(std::map<size_t, double>& a, double d) {
  for(std::map<size_t, double>::iterator iter = a.begin();
    iter != a.end(); ++iter){
  size_t id = (*iter).first;
  double val = (*iter).second;
  a[id] = val*d;
//    (*iter).second *= d;
  }
  return a;
}

typedef graphlab::distributed_graph<vertex_data, edge_data> graph_type;

graphlab::atomic<graphlab::vertex_id_type> NEXT_VID;

// Read a line from a file and creates a vertex
bool vertex_loader(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;
  namespace qi = boost::spirit::qi;
  namespace ascii = boost::spirit::ascii;
  namespace phoenix = boost::phoenix;
  vertex_data vtx;
  const bool success = qi::phrase_parse
    (line.begin(), line.end(),
     //  Begin grammar
     (
      (qi::double_[phoenix::push_back(phoenix::ref(vtx.point), qi::_1)] % -qi::char_(",") )
      )
     ,
     //  End grammar
     ascii::space);

  if (!success) return false;
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(NEXT_VID.inc_ret_last(1), vtx);
  return true;
}

// Read a line from a file and creates a vertex
bool vertex_loader_sparse(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;

  vertex_data vtx;
  boost::char_separator<char> sep(" ");
  boost::tokenizer< boost::char_separator<char> > tokens(line, sep);
  BOOST_FOREACH (const std::string& t, tokens) {
    std::string::size_type pos = t.find(":");
    if(pos > 0){
      size_t id = (size_t)std::atoi(t.substr(0, pos).c_str());
      double val = std::atof(t.substr(pos+1, t.length() - pos -1).c_str());
      vtx.point_sparse.insert(std::make_pair(id, val));
    }
  }
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(NEXT_VID.inc_ret_last(1), vtx);
  return true;
}

// Read a line from a file and creates a vertex
bool vertex_loader_with_id(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;
  size_t id = 0;
  namespace qi = boost::spirit::qi;
  namespace ascii = boost::spirit::ascii;
  namespace phoenix = boost::phoenix;
  vertex_data vtx;
  const bool success = qi::phrase_parse
    (line.begin(), line.end(),
     //  Begin grammar
     (
      qi::ulong_[phoenix::ref(id) = qi::_1] >> -qi::char_(",") >>
      (qi::double_[phoenix::push_back(phoenix::ref(vtx.point), qi::_1)] % -qi::char_(",") )
      )
     ,
     //  End grammar
     ascii::space);

  if (!success) return false;
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(id, vtx);
  return true;
}

// Read a line from a file and creates a vertex
bool vertex_loader_with_id_sparse(graph_type& graph, const std::string& fname,
                   const std::string& line) {
  if (line.empty()) return true;

  vertex_data vtx;
  size_t id = 0;
  boost::char_separator<char> sep(" ");
  boost::tokenizer<boost::char_separator<char> > tokens(line, sep);
  bool first = true;
  BOOST_FOREACH (const std::string& t, tokens) {
    if(first){
      id = (size_t)std::atoi(t.c_str());
      first = false;
    }else{
      std::string::size_type pos = t.find(":");
      if(pos > 0){
        size_t id = (size_t)std::atoi(t.substr(0, pos).c_str());
        double val = std::atof(t.substr(pos+1, t.length() - pos -1).c_str());
        vtx.point_sparse.insert(std::make_pair(id, val));
      }
    }
  }
  vtx.best_cluster = (size_t)(-1);
  vtx.best_distance = std::numeric_limits<double>::infinity();
  vtx.changed = false;
  graph.add_vertex(id, vtx);
  return true;
}

//call this when edge weight file is given.
//each line should be [source id] [target id] [weight].
//directions of edges are ignored.
bool edge_loader(graph_type& graph, const std::string& filename,
    const std::string& textline) {
  if (textline.empty())
    return true;
  std::stringstream strm(textline);
  size_t source_vid = 0;
  size_t target_vid = 0;
  double weight = 0.0;
  strm >> source_vid;
  strm.ignore(1);
  strm >> target_vid;
  strm.ignore(1);
  strm >> weight;
  if(source_vid != target_vid)
    graph.add_edge(source_vid, target_vid, edge_data(weight));
  return true;
}

// A set of Map Reduces to compute the maximum and minimum vector sizes
// to ensure that all vectors have the same length
struct max_point_size_reducer: public graphlab::IS_POD_TYPE {
  size_t max_point_size;

  static max_point_size_reducer get_max_point_size(const graph_type::vertex_type& v) {
    max_point_size_reducer r;
    r.max_point_size = v.data().point.size();
    return r;
  }

  max_point_size_reducer& operator+=(const max_point_size_reducer& other) {
    max_point_size = std::max(max_point_size, other.max_point_size);
    return *this;
  }
};

struct min_point_size_reducer: public graphlab::IS_POD_TYPE {
  size_t min_point_size;

  static min_point_size_reducer get_min_point_size(const graph_type::vertex_type& v) {
    min_point_size_reducer r;
    r.min_point_size = v.data().point.size();
    return r;
  }

  min_point_size_reducer& operator+=(const min_point_size_reducer& other) {
    min_point_size = std::min(min_point_size, other.min_point_size);
    return *this;
  }
};

/*
 * This transform vertices call is only used during
 * the initialization phase. It computes distance to
 * cluster[KMEANS_INITIALIZATION] and assigns itself
 * to the new cluster KMEANS_INITIALIZATION if the new distance
 * is smaller that its previous cluster assignment
 */
void kmeans_pp_initialization(graph_type::vertex_type& v) {
  double d = sqr_distance(v.data().point,
                          CLUSTERS[KMEANS_INITIALIZATION].center);
  if (v.data().best_distance > d) {
    v.data().best_distance = d;
    v.data().best_cluster = KMEANS_INITIALIZATION;
  }
}

void kmeans_pp_initialization_sparse(graph_type::vertex_type& v) {
  double d = sqr_distance(v.data().point_sparse,
                          CLUSTERS[KMEANS_INITIALIZATION].center_sparse);
  if (v.data().best_distance > d) {
    v.data().best_distance = d;
    v.data().best_cluster = KMEANS_INITIALIZATION;
  }
}

/*
 * Draws a random sample from the data points that is
 * proportionate to the "best distance" stored in the vertex.
 */
struct random_sample_reducer {
  std::vector<double> vtx;
  double weight;

  random_sample_reducer():weight(0) { }
  random_sample_reducer(const std::vector<double>& vtx,
                        double weight):vtx(vtx),weight(weight) { }

  static random_sample_reducer get_weight(const graph_type::vertex_type& v) {
    if (v.data().best_cluster == (size_t)(-1)) {
      return random_sample_reducer(v.data().point, 1);
    }
    else {
      return random_sample_reducer(v.data().point,
                                   v.data().best_distance);
    }
  }

  random_sample_reducer& operator+=(const random_sample_reducer& other) {
    double totalweight = weight + other.weight;
    // if any weight is too small, just quit
    if (totalweight <= 0) return *this;

    double myp = weight / (weight + other.weight);
    if (graphlab::random::bernoulli(myp)) {
      weight += other.weight;
      return *this;
    }
    else {
      vtx = other.vtx;
      weight += other.weight;
      return *this;
    }
  }

  void save(graphlab::oarchive &oarc) const {
    oarc << vtx << weight;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> vtx >> weight;
  }
};

struct random_sample_reducer_sparse{
  std::map<size_t, double> vtx;
  double weight;

  random_sample_reducer_sparse():weight(0) { }
  random_sample_reducer_sparse(const std::map<size_t, double>& vtx,
                        double weight):vtx(vtx),weight(weight) { }

  static random_sample_reducer_sparse get_weight(const graph_type::vertex_type& v) {
    if (v.data().best_cluster == (size_t)(-1)) {
      return random_sample_reducer_sparse(v.data().point_sparse, 1);
    }
    else {
      return random_sample_reducer_sparse(v.data().point_sparse,
                                   v.data().best_distance);
    }
  }

  random_sample_reducer_sparse& operator+=(const random_sample_reducer_sparse& other) {
    double totalweight = weight + other.weight;
    // if any weight is too small, just quit
    if (totalweight <= 0) return *this;

    double myp = weight / (weight + other.weight);
    if (graphlab::random::bernoulli(myp)) {
      weight += other.weight;
      return *this;
    }
    else {
      vtx = other.vtx;
      weight += other.weight;
      return *this;
    }
  }

  void save(graphlab::oarchive &oarc) const {
    oarc << vtx << weight;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> vtx >> weight;
  }
};

/*
 * This transform vertices call is used during the
 * actual k-means iteration. It computes distance to
 * all "changed" clusters and reassigns itself if necessary
 */
void kmeans_iteration(graph_type::vertex_type& v) {
  // if current vertex's cluster was modified, we invalidate the distance.
  // and we need to recompute to all existing clusters
  // otherwise, we just need to recompute to changed cluster centers.
  size_t prev_asg = v.data().best_cluster;
  if (CLUSTERS[v.data().best_cluster].changed) {
    // invalidate. recompute to all
    v.data().best_cluster = (size_t)(-1);
    v.data().best_distance = std::numeric_limits<double>::infinity();
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      if (CLUSTERS[i].center.size() > 0 || CLUSTERS[i].center_sparse.size() > 0) {
        double d = 0.0;
        if(IS_SPARSE == true)
          d = sqr_distance(v.data().point_sparse, CLUSTERS[i].center_sparse);
        else
          d = sqr_distance(v.data().point, CLUSTERS[i].center);
        if (d < v.data().best_distance) {
          v.data().best_distance = d;
          v.data().best_cluster = i;
        }
      }
    }
  }
  else {
    // just compute distance to what has changed
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      if (CLUSTERS[i].changed &&
          (CLUSTERS[i].center.size() > 0 || CLUSTERS[i].center_sparse.size() > 0)) {
        double d = 0.0;
        if(IS_SPARSE == true)
          d = sqr_distance(v.data().point_sparse, CLUSTERS[i].center_sparse);
        else
          d= sqr_distance(v.data().point, CLUSTERS[i].center);
        if (d < v.data().best_distance) {
          v.data().best_distance = d;
          v.data().best_cluster = i;
        }
      }
    }
  }
  v.data().changed = (prev_asg != v.data().best_cluster);
}

//gathered information
//used when edge weight file is given
struct neighbor_info {
  std::map<size_t, double> cw_map;

  neighbor_info() :
      cw_map() {
  }
  neighbor_info(size_t clst, double weight) :
      cw_map() {
    cw_map.insert(std::make_pair(clst, weight));
  }

  neighbor_info& operator+=(const neighbor_info& other) {
    for (std::map<size_t, double>::const_iterator iter = other.cw_map.begin();
        iter != other.cw_map.end(); iter++) {
      size_t clst = iter->first;
      if (cw_map.find(clst) == cw_map.end()) {
        cw_map.insert(std::make_pair(clst, iter->second));
      } else {
        cw_map[clst] += iter->second;
      }
    }
    return *this;
  }

  void save(graphlab::oarchive& oarc) const {
    oarc << cw_map;
  }
  void load(graphlab::iarchive& iarc) {
    iarc >> cw_map;
  }
};

//used when edge weight file is given
class cluster_assignment: public graphlab::ivertex_program<graph_type,
    neighbor_info>, public graphlab::IS_POD_TYPE {
public:
  //gather on all the edges
  edge_dir_type gather_edges(icontext_type& context,
      const vertex_type& vertex) const {
    return graphlab::ALL_EDGES;
  }

  //for each edge gather the weights and the assigned clusters of the neighbors
  neighbor_info gather(icontext_type& context, const vertex_type& vertex,
      edge_type& edge) const {
    if (edge.source().id() == vertex.id()) { //out edge
      return neighbor_info(edge.target().data().best_cluster,
          edge.data().weight);
    } else { //in edge
      return neighbor_info(edge.source().data().best_cluster,
          edge.data().weight);
    }
  }

  //assign a cluster, considering the clusters of neighbors
  void apply(icontext_type& context, vertex_type& vertex,
      const gather_type& total) {
    size_t past_clst = vertex.data().best_cluster;
    vertex.data().best_cluster = (size_t) (-1);
    vertex.data().best_distance = std::numeric_limits<double>::infinity();
    for (size_t i = 0; i < NUM_CLUSTERS; ++i) {
      if (CLUSTERS[i].center.size() > 0 || CLUSTERS[i].center_sparse.size() > 0) {
        double d = 0.0;
        if(IS_SPARSE == true)
          d = sqr_distance(vertex.data().point_sparse, CLUSTERS[i].center_sparse);
        else
          d = sqr_distance(vertex.data().point, CLUSTERS[i].center);
        //consider neighbors
        const std::map<size_t, double>& cw_map = total.cw_map;
        for (std::map<size_t, double>::const_iterator iter = cw_map.begin();
            iter != cw_map.end(); iter++) {
          size_t neighbor_cluster = iter->first;
          double total_wieght = iter->second;
          if (i == neighbor_cluster)
            d -= total_wieght;
        }
        if (d < vertex.data().best_distance) {
          vertex.data().best_distance = d;
          vertex.data().best_cluster = i;
        }
      }
    }
    vertex.data().changed = (past_clst != vertex.data().best_cluster);
  }

  //send signals to the neighbors when the cluster assignment has changed
  edge_dir_type scatter_edges(icontext_type& context,
      const vertex_type& vertex) const {
    if (vertex.data().changed)
      return graphlab::ALL_EDGES;
    else
      return graphlab::NO_EDGES;
  }

  void scatter(icontext_type& context, const vertex_type& vertex,
      edge_type& edge) const {
  }
};

/*
 * computes new cluster centers
 * Also accumulates a counter counting the number of vertices which
 * assignments changed.
 */
struct cluster_center_reducer {
  std::vector<cluster> new_clusters;
  size_t num_changed;
  double cost;

  cluster_center_reducer():new_clusters(NUM_CLUSTERS), num_changed(0), cost(0) { }

  static cluster_center_reducer get_center(const graph_type::vertex_type& v) {
    cluster_center_reducer cc;
    ASSERT_NE(v.data().best_cluster, (size_t)(-1));

    if(IS_SPARSE == true)
      cc.new_clusters[v.data().best_cluster].center_sparse = v.data().point_sparse;
    else
      cc.new_clusters[v.data().best_cluster].center = v.data().point;
    cc.new_clusters[v.data().best_cluster].count = 1;
    cc.num_changed = v.data().changed;
    cc.cost = v.data().best_distance;
    return cc;
  }

  cluster_center_reducer& operator+=(const cluster_center_reducer& other) {
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      if (new_clusters[i].count == 0) new_clusters[i] = other.new_clusters[i];
      else if (other.new_clusters[i].count > 0) {
        if(IS_SPARSE == true)
          plus_equal_vector(new_clusters[i].center_sparse, other.new_clusters[i].center_sparse);
        else
          plus_equal_vector(new_clusters[i].center, other.new_clusters[i].center);
        new_clusters[i].count += other.new_clusters[i].count;
      }
    }
    num_changed += other.num_changed;
    cost += other.cost;
    return *this;
  }

  void save(graphlab::oarchive& oarc) const {
    oarc << new_clusters << num_changed <<cost;
  }

  void load(graphlab::iarchive& iarc) {
    iarc >> new_clusters >> num_changed >> cost;
  }
};

struct vertex_writer {
  std::string save_vertex(graph_type::vertex_type v) {
    std::stringstream strm;
    for (size_t i = 0;i < v.data().point.size(); ++i) {
      strm << v.data().point[i] << "\t";
    }
    strm << v.data().best_distance << "\t";
    strm << v.data().best_cluster << "\n";
    strm.flush();
    return strm.str();
  }

  std::string save_edge(graph_type::edge_type e) { return ""; }
};

struct vertex_writer_sparse {
  std::string save_vertex(graph_type::vertex_type v) {
    std::stringstream strm;
    for(std::map<size_t, double>::iterator iter = v.data().point_sparse.begin();
        iter != v.data().point_sparse.end();++iter){
      strm << (*iter).first << ":" << (*iter).second << " ";
    }
    strm << v.data().best_cluster << "\n";
    strm.flush();
    return strm.str();
  }

  std::string save_edge(graph_type::edge_type e) { return ""; }
};

struct vertex_writer_with_id {
  std::string save_vertex(graph_type::vertex_type v) {
    std::stringstream strm;
    strm << v.id() << "\t";
    strm << v.data().best_cluster+1 << "\n";
    strm.flush();
    return strm.str();
  }

  std::string save_edge(graph_type::edge_type e) { return ""; }
};

int main(int argc, char** argv) {
  std::cout << "Computes a K-means clustering of data.\n\n";

  graphlab::command_line_options clopts
    ("K-means clustering. The input data file is provided by the "
     "--data argument which is non-optional. The format of the data file is a "
     "collection of lines, where each line contains a comma or white-space "
     "separated lost of numeric values representing a vector. Every line "
     "must have the same number of values. The required --clusters=N "
     "argument denotes the number of clusters to generate. To store the output "
     "see the --output-cluster and --output-data arguments");

  std::string datafile;
  std::string outcluster_file;
  std::string outdata_file;
  std::string edgedata_file;
  size_t MAX_ITERATION = 0;
  bool use_id = false;
  clopts.attach_option("data", datafile,
                       "Input file. Each line holds a white-space or comma separated numeric vector");
  clopts.attach_option("clusters", NUM_CLUSTERS,
                       "The number of clusters to create.");
  clopts.attach_option("output-clusters", outcluster_file,
                       "If set, will write a file containing cluster centers "
                       "to this filename. This must be on the local filesystem "
                       "and must be accessible to the root node.");
  clopts.attach_option("output-data", outdata_file,
                       "If set, will output a copy of the input data with an additional "
                       "two columns. The first added column is the distance to assigned "
		       "center and the last is the assigned cluster centers. The output "
                       "will be written to a sequence of filenames where each file is "
                       "prefixed by this value. This may be on HDFS.");
  clopts.attach_option("sparse", IS_SPARSE,
                       "If set to true, will use a sparse vector representation."
                       "The file format is [feature id]:[value] [feature id]:[value] ..."
                       ", where [feature id] must be positive integer or zero.");
  clopts.attach_option("id", use_id,
                       "If set to true, will use ids for data points. The id of a data point "
                       "must be written at the head of each line of the input data. "
                       "The output data will consist of two columns: the first one "
                       "denotes the ids; the second one denotes the assigned clusters.");
  clopts.attach_option("pairwise-reward", edgedata_file,
                       "If set, will consider pairwise rewards when clustering. "
                       "Each line of the file beginning with the argument holds [id1] [id2] "
                       "[reward]. This mode must be used with --id option.");
  clopts.attach_option("max-iteration", MAX_ITERATION,
                       "The max number of iterations");

  if(!clopts.parse(argc, argv)) return EXIT_FAILURE;
  if (datafile == "") {
    std::cout << "--data is not optional\n";
    return EXIT_FAILURE;
  }
  if (NUM_CLUSTERS == 0) {
    std::cout << "--clusters is not optional\n";
    return EXIT_FAILURE;
  }
  if(edgedata_file.size() > 0){
    if(use_id == false){
      std::cout << "--id is not optional when you use edge data\n";
      return EXIT_FAILURE;
    }
  }

  graphlab::mpi_tools::init(argc, argv);
  graphlab::distributed_control dc;
  // load graph
  graph_type graph(dc, clopts);
  NEXT_VID = (((graphlab::vertex_id_type)1 << 31) / dc.numprocs()) * dc.procid();
  if(IS_SPARSE == true){
    if(use_id){
      graph.load(datafile, vertex_loader_with_id_sparse);
    }else{
      graph.load(datafile, vertex_loader_sparse);
    }
  }else{
    if(use_id){
      graph.load(datafile, vertex_loader_with_id);
    }else{
      graph.load(datafile, vertex_loader);
    }
  }
  if(edgedata_file.size() > 0){
    graph.load(edgedata_file, edge_loader);
  }
  graph.finalize();
  dc.cout() << "Number of datapoints: " << graph.num_vertices() << std::endl;

  if (graph.num_vertices() < NUM_CLUSTERS) {
    dc.cout() << "More clusters than datapoints! Cannot proceed" << std::endl;
    return EXIT_FAILURE;
  }

  dc.cout() << "Validating data...";

  CLUSTERS.resize(NUM_CLUSTERS);
  // make sure all have the same array length
  if(IS_SPARSE == false){
    size_t max_p_size = graph.map_reduce_vertices<max_point_size_reducer>
                                  (max_point_size_reducer::get_max_point_size).max_point_size;
    size_t min_p_size = graph.map_reduce_vertices<min_point_size_reducer>
                                  (min_point_size_reducer::get_min_point_size).min_point_size;
    if (max_p_size != min_p_size) {
      dc.cout() << "Data has dimensionality ranging from " << min_p_size << " to " << max_p_size
                << "! K-means cannot proceed!" << std::endl;
      return EXIT_FAILURE;
    }
    // allocate clusters
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      CLUSTERS[i].center.resize(max_p_size);
    }
  }

  dc.cout() << "Initializing using Kmeans++\n";
  // ok. perform kmeans++ initialization
  for (KMEANS_INITIALIZATION = 0;
       KMEANS_INITIALIZATION < NUM_CLUSTERS;
       ++KMEANS_INITIALIZATION) {

    if(IS_SPARSE == true){
      random_sample_reducer_sparse rs = graph.map_reduce_vertices<random_sample_reducer_sparse>
                                        (random_sample_reducer_sparse::get_weight);
      CLUSTERS[KMEANS_INITIALIZATION].center_sparse = rs.vtx;
      graph.transform_vertices(kmeans_pp_initialization_sparse);
    }else{
      random_sample_reducer rs = graph.map_reduce_vertices<random_sample_reducer>
                                        (random_sample_reducer::get_weight);
      CLUSTERS[KMEANS_INITIALIZATION].center = rs.vtx;
      graph.transform_vertices(kmeans_pp_initialization);
    }
  }

  // "reset" all clusters
  for (size_t i = 0; i < NUM_CLUSTERS; ++i) CLUSTERS[i].changed = true;
  // perform Kmeans iteration

  dc.cout() << "Running Kmeans...\n";
  bool clusters_changed = true;
  size_t iteration_count = 0;
  while(clusters_changed) {
		if(MAX_ITERATION > 0 && iteration_count >= MAX_ITERATION)
			break;

    cluster_center_reducer cc = graph.map_reduce_vertices<cluster_center_reducer>
                                    (cluster_center_reducer::get_center);
    // the first round (iteration_count == 0) is not so meaningful
    // since I am just recomputing the centers from the output of the KMeans++
    // initialization
    if (iteration_count > 0) {
      dc.cout() << "Kmeans iteration " << iteration_count << ": " <<
                 "# points with changed assignments = " << cc.num_changed <<
		 " total cost: " << cc.cost << std::endl;
    }
    for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
      double d = cc.new_clusters[i].count;
      if(IS_SPARSE){
        if (d > 0) scale_vector(cc.new_clusters[i].center_sparse, 1.0 / d);
        if (cc.new_clusters[i].count == 0 && CLUSTERS[i].count > 0) {
          dc.cout() << "Cluster " << i << " lost" << std::endl;
          CLUSTERS[i].center_sparse.clear();
          CLUSTERS[i].count = 0;
          CLUSTERS[i].changed = false;
        }
        else {
          CLUSTERS[i] = cc.new_clusters[i];
          CLUSTERS[i].changed = true;
        }
      }else{
        if (d > 0) scale_vector(cc.new_clusters[i].center, 1.0 / d);
        if (cc.new_clusters[i].count == 0 && CLUSTERS[i].count > 0) {
          dc.cout() << "Cluster " << i << " lost" << std::endl;
          CLUSTERS[i].center.clear();
          CLUSTERS[i].count = 0;
          CLUSTERS[i].changed = false;
        }
        else {
          CLUSTERS[i] = cc.new_clusters[i];
          CLUSTERS[i].changed = true;
        }
      }
    }
    clusters_changed = iteration_count == 0 || cc.num_changed > 0;

    if(edgedata_file.size() > 0){
      clopts.engine_args.set_option("factorized", true);
      graphlab::omni_engine<cluster_assignment> engine(dc, graph, "async", clopts);
      engine.signal_all();
      engine.start();
    }else{
      graph.transform_vertices(kmeans_iteration);
    }

    ++iteration_count;
  }

  if (!outcluster_file.empty() && dc.procid() == 0) {
    dc.cout() << "Writing Cluster Centers..." << std::endl;
    std::ofstream fout(outcluster_file.c_str());
    if(IS_SPARSE){
      for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
        if(use_id)
          fout << i+1 << "\t";
        for (std::map<size_t, double>::iterator iter = CLUSTERS[i].center_sparse.begin();
             iter != CLUSTERS[i].center_sparse.end();++iter) {
          fout << (*iter).first << ":" << (*iter).second << " ";
        }
        fout << "\n";
      }
    }else{
      for (size_t i = 0;i < NUM_CLUSTERS; ++i) {
        if(use_id)
          fout << i+1 << "\t";
        for (size_t j = 0; j < CLUSTERS[i].center.size(); ++j) {
          fout << CLUSTERS[i].center[j] << " ";
        }
        fout << "\n";
      }
    }
  }

  if (!outdata_file.empty()) {
    dc.cout() << "Writing Data with cluster assignments...\n" << std::endl;
    if(use_id){
      graph.save(outdata_file, vertex_writer_with_id(), false, true, false, 1);
    }else{
      if(IS_SPARSE == true)
        graph.save(outdata_file, vertex_writer_sparse(), false, true, false, 1);
      else
        graph.save(outdata_file, vertex_writer(), false, true, false, 1);
    }
  }

  graphlab::mpi_tools::finalize();
}


#include <cfloat>
#include <cmath>

#include <boost/math/distributions/normal.hpp>

#include <DataTypes/DataTypeTuple.h>
#include <DataTypes/DataTypesDecimal.h>
#include <DataTypes/DataTypesNumber.h>
#include <Columns/ColumnTuple.h>
#include <Columns/ColumnsNumber.h>
#include <Functions/FunctionFactory.h>
#include <Functions/FunctionHelpers.h>
#include <Functions/IFunction.h>
#include <Functions/castTypeToEither.h>
#include <Interpreters/castColumn.h>

namespace DB
{

namespace ErrorCodes
{
    extern const int ILLEGAL_TYPE_OF_ARGUMENT;
}

template <typename Impl>
class FunctionMinSampleSize : public IFunction
{
public:
    static constexpr auto name = Impl::name;

    static FunctionPtr create(ContextPtr) { return std::make_shared<FunctionMinSampleSize<Impl>>(); }

    String getName() const override { return name; }

    size_t getNumberOfArguments() const override { return Impl::num_args; }
    ColumnNumbers getArgumentsThatAreAlwaysConstant() const override
    {
        return ColumnNumbers(std::begin(Impl::const_args), std::end(Impl::const_args));
    }

    bool useDefaultImplementationForNulls() const override { return false; }
    bool useDefaultImplementationForConstants() const override { return true; }
    bool isSuitableForShortCircuitArgumentsExecution(const DataTypesWithConstInfo & /*arguments*/) const override { return false; }

    static DataTypePtr getReturnType()
    {
        auto float_64_type = std::make_shared<DataTypeNumber<Float64>>();

        DataTypes types{
            float_64_type,
            float_64_type,
            float_64_type,
        };

        Strings names{
            "minimum_sample_size",
            "detect_range_lower",
            "detect_range_upper",
        };

        return std::make_shared<DataTypeTuple>(std::move(types), std::move(names));
    }

    DataTypePtr getReturnTypeImpl(const DataTypes & arguments) const override
    {
        Impl::validateArguments(arguments);
        return getReturnType();
    }

    ColumnPtr executeImpl(const ColumnsWithTypeAndName & arguments, const DataTypePtr &, size_t input_rows_count) const override
    {
        return Impl::execute(arguments, input_rows_count);
    }
};

static bool isBetweenZeroAndOne(Float64 v)
{
    return v >= 0.0 && v <= 1.0 && fabs(v - 0.0) >= DBL_EPSILON && fabs(v - 1.0) >= DBL_EPSILON;
}

struct ContinousImpl
{
    static constexpr auto name = "minSampleSizeContinous";
    static constexpr size_t num_args = 5;
    static constexpr size_t const_args[] = {2, 3, 4};

    static void validateArguments(const DataTypes & arguments)
    {
        for (size_t i = 0; i < arguments.size(); ++i)
        {
            if (!isNativeNumber(arguments[i]))
            {
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The {}th Argument of function {} must be a number.", i + 1, name);
            }
        }
    }

    static ColumnPtr execute(const ColumnsWithTypeAndName & arguments, size_t input_rows_count)
    {
        auto float_64_type = std::make_shared<DataTypeFloat64>();
        auto baseline_argument = arguments[0];
        baseline_argument.column = baseline_argument.column->convertToFullColumnIfConst();
        auto baseline_column_untyped = castColumnAccurate(baseline_argument, float_64_type);
        const auto * baseline_column = checkAndGetColumn<ColumnVector<Float64>>(*baseline_column_untyped);
        const auto & baseline_column_data = baseline_column->getData();

        auto sigma_argument = arguments[1];
        sigma_argument.column = sigma_argument.column->convertToFullColumnIfConst();
        auto sigma_column_untyped = castColumnAccurate(sigma_argument, float_64_type);
        const auto * sigma_column = checkAndGetColumn<ColumnVector<Float64>>(*sigma_column_untyped);
        const auto & sigma_column_data = sigma_column->getData();

        const IColumn & col_mde = *arguments[2].column;
        const IColumn & col_power = *arguments[3].column;
        const IColumn & col_alpha = *arguments[4].column;

        auto res_min_sample_size = ColumnFloat64::create();
        auto & data_min_sample_size = res_min_sample_size->getData();
        data_min_sample_size.reserve(input_rows_count);

        auto res_detect_lower = ColumnFloat64::create();
        auto & data_detect_lower = res_detect_lower->getData();
        data_detect_lower.reserve(input_rows_count);

        auto res_detect_upper = ColumnFloat64::create();
        auto & data_detect_upper = res_detect_upper->getData();
        data_detect_upper.reserve(input_rows_count);

        /// Minimal Detectable Effect
        const Float64 mde = col_mde.getFloat64(0);
        /// Sufficient statistical power to detect a treatment effect
        const Float64 power = col_power.getFloat64(0);
        /// Significance level
        const Float64 alpha = col_alpha.getFloat64(0);

        boost::math::normal_distribution<> nd(0.0, 1.0);

        for (size_t row_num = 0; row_num < input_rows_count; ++row_num)
        {
            /// Mean of control-metric
            Float64 baseline = baseline_column_data[row_num];
            /// Standard deviation of conrol-metric
            Float64 sigma = sigma_column_data[row_num];

            if (!std::isfinite(baseline) || !std::isfinite(sigma) || !isBetweenZeroAndOne(mde) || !isBetweenZeroAndOne(power)
                || !isBetweenZeroAndOne(alpha))
            {
                data_min_sample_size.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_lower.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_upper.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                continue;
            }

            Float64 delta = baseline * mde;

            using namespace boost::math;
            /// https://towardsdatascience.com/required-sample-size-for-a-b-testing-6f6608dd330a
            /// \frac{2\sigma^{2} * (Z_{1 - alpha /2} + Z_{power})^{2}}{\Delta^{2}}
            Float64 min_sample_size
                = 2 * std::pow(sigma, 2) * std::pow(quantile(nd, 1.0 - alpha / 2) + quantile(nd, power), 2) / std::pow(delta, 2);

            data_min_sample_size.emplace_back(min_sample_size);
            data_detect_lower.emplace_back(baseline - delta);
            data_detect_upper.emplace_back(baseline + delta);
        }

        return ColumnTuple::create(Columns{std::move(res_min_sample_size), std::move(res_detect_lower), std::move(res_detect_upper)});
    }
};

struct ConversionImpl
{
    static constexpr auto name = "minSampleSizeConversion";
    static constexpr size_t num_args = 4;
    static constexpr size_t const_args[] = {1, 2, 3};

    static void validateArguments(const DataTypes & arguments)
    {
        size_t arguments_size = arguments.size();
        for (size_t i = 0; i < arguments_size; ++i)
        {
            if (!isFloat(arguments[i]))
            {
                throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The {}th argument of function {} must be a float.", i + 1, name);
            }
        }
    }

    static ColumnPtr execute(const ColumnsWithTypeAndName & arguments, size_t input_rows_count)
    {
        auto first_argument_column = castColumnAccurate(arguments[0], std::make_shared<DataTypeFloat64>());

        if (const ColumnConst * const col_p1_const = checkAndGetColumnConst<ColumnVector<Float64>>(first_argument_column.get()))
        {
            const Float64 left_value = col_p1_const->template getValue<Float64>();
            return process<true>(arguments, &left_value, input_rows_count);
        }
        else if (const ColumnVector<Float64> * const col_p1 = checkAndGetColumn<ColumnVector<Float64>>(first_argument_column.get()))
        {
            return process<false>(arguments, col_p1->getData().data(), input_rows_count);
        }
        else
        {
            throw Exception(ErrorCodes::ILLEGAL_TYPE_OF_ARGUMENT, "The first argument of function {} must be a float.", name);
        }
    }

    template <bool const_p1>
    static ColumnPtr process(const ColumnsWithTypeAndName & arguments, const Float64 * col_p1, const size_t input_rows_count)
    {
        const IColumn & col_mde = *arguments[1].column;
        const IColumn & col_power = *arguments[2].column;
        const IColumn & col_alpha = *arguments[3].column;

        auto res_min_sample_size = ColumnFloat64::create();
        auto & data_min_sample_size = res_min_sample_size->getData();
        data_min_sample_size.reserve(input_rows_count);

        auto res_detect_lower = ColumnFloat64::create();
        auto & data_detect_lower = res_detect_lower->getData();
        data_detect_lower.reserve(input_rows_count);

        auto res_detect_upper = ColumnFloat64::create();
        auto & data_detect_upper = res_detect_upper->getData();
        data_detect_upper.reserve(input_rows_count);

        /// Minimal Detectable Effect
        const Float64 mde = col_mde.getFloat64(0);
        /// Sufficient statistical power to detect a treatment effect
        const Float64 power = col_power.getFloat64(0);
        /// Significance level
        const Float64 alpha = col_alpha.getFloat64(0);

        boost::math::normal_distribution<> nd(0.0, 1.0);

        for (size_t row_num = 0; row_num < input_rows_count; ++row_num)
        {
            /// Proportion of control-metric
            Float64 p1;

            if constexpr (const_p1)
            {
                p1 = col_p1[0];
            }
            else if constexpr (!const_p1)
            {
                p1 = col_p1[row_num];
            }

            if (!std::isfinite(p1) || !isBetweenZeroAndOne(mde) || !isBetweenZeroAndOne(power) || !isBetweenZeroAndOne(alpha))
            {
                data_min_sample_size.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_lower.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                data_detect_upper.emplace_back(std::numeric_limits<Float64>::quiet_NaN());
                continue;
            }

            Float64 q1 = 1.0 - p1;
            Float64 p2 = p1 + mde;
            Float64 q2 = 1.0 - p2;
            Float64 p_bar = (p1 + p2) / 2.0;
            Float64 q_bar = 1.0 - p_bar;

            using namespace boost::math;
            /// https://towardsdatascience.com/required-sample-size-for-a-b-testing-6f6608dd330a
            /// \frac{(Z_{1-alpha/2} * \sqrt{2*\bar{p}*\bar{q}} + Z_{power} * \sqrt{p1*q1+p2*q2})^{2}}{\Delta^{2}}
            Float64 min_sample_size
                = std::pow(
                      quantile(nd, 1.0 - alpha / 2.0) * std::sqrt(2.0 * p_bar * q_bar) + quantile(nd, power) * std::sqrt(p1 * q1 + p2 * q2),
                      2)
                / std::pow(mde, 2);

            data_min_sample_size.emplace_back(min_sample_size);
            data_detect_lower.emplace_back(p1 - mde);
            data_detect_upper.emplace_back(p1 + mde);
        }

        return ColumnTuple::create(Columns{std::move(res_min_sample_size), std::move(res_detect_lower), std::move(res_detect_upper)});
    }
};

void registerFunctionMinSampleSize(FunctionFactory & factory)
{
    factory.registerFunction<FunctionMinSampleSize<ContinousImpl>>();
    factory.registerFunction<FunctionMinSampleSize<ConversionImpl>>();
}

}


#include <chrono>
#include <string>
#include <iostream>
#include <algorithm>
#include <pngwriter.h>
#include <boost/program_options.hpp>
#include "mandelbrot.hpp"

namespace po = boost::program_options;

int main(int argc, char* argv[]) {
    std::string out_file;
    double scale, cx, cy;
    size_t max_n, width, height, n_colors, n_threads, nx, ny;
    po::options_description desc("Allowed options");
    desc.add_options()
            ("help,h", "Print this message")
            ("out,o", po::value(&out_file)->default_value("result.png"), "Output file")
            ("scale,s", po::value(&scale)->default_value(1), "Scale")
            ("cx", po::value(&cx)->default_value(-0.28676842048), "X center coordinate")
            ("cy", po::value(&cy)->default_value(0), "Y center coordinate")
            ("iter,i", po::value(&max_n)->default_value(5000), "Maximum number of iterations")
            ("width", po::value(&width)->default_value(1920), "Image width")
            ("height", po::value(&height)->default_value(1080), "Image height")
            ("colors", po::value(&n_colors)->default_value(10), "Number of colors")
            ("threads,t", po::value(&n_threads)->default_value(1), "Number of threads")
            ("nx", po::value(&nx)->default_value(16), "Split into nx columns")
            ("ny", po::value(&ny)->default_value(16), "Split into ny rows");
    po::variables_map vm;
    po::store(po::parse_command_line(argc, argv, desc), vm);
    po::notify(vm);
    if (vm.count("help")) {
        std::cout << desc << std::endl;
        return 0;
    }
    auto m = mtms::mandelbrot_set<double>(width, height, n_colors);
    const auto start = std::chrono::high_resolution_clock::now();
    m.run(n_threads, nx, ny, scale * std::max(3. / width, 2. / height), {cx, cy}, max_n);
    const auto run_end = std::chrono::high_resolution_clock::now();
    pngwriter png(static_cast<int>(width), static_cast<int>(height), 0, out_file.c_str());
    for (size_t i = 0; i < height; ++i)
        for (size_t j = 0; j < width; ++j) {
            const auto color = m.color(i, j);
            png.plot(static_cast<int>(j), static_cast<int>(i), color.r, color.g, color.b);
        }
    png.close();
    const auto write_end = std::chrono::high_resolution_clock::now();
    std::cout << "Run time: "
              << std::chrono::duration_cast<std::chrono::nanoseconds>(run_end - start).count() << " ns\n"
              << "Write time : "
              << std::chrono::duration_cast<std::chrono::nanoseconds>(write_end - run_end).count() << " ns\n"
              << "All time: "
              << std::chrono::duration_cast<std::chrono::nanoseconds>(write_end - start).count() << " ns" << std::endl;
    return 0;
}


///////////////////////////////////////////////////////////////////////////////
// weighted_extended_p_square.hpp
//
//  Copyright 2005 Daniel Egloff. 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)

#ifndef BOOST_ACCUMULATORS_STATISTICS_WEIGHTED_EXTENDED_P_SQUARE_HPP_DE_01_01_2006
#define BOOST_ACCUMULATORS_STATISTICS_WEIGHTED_EXTENDED_P_SQUARE_HPP_DE_01_01_2006

#include <vector>
#include <functional>
#include <boost/range/begin.hpp>
#include <boost/range/end.hpp>
#include <boost/range/iterator_range.hpp>
#include <boost/iterator/transform_iterator.hpp>
#include <boost/iterator/counting_iterator.hpp>
#include <boost/iterator/permutation_iterator.hpp>
#include <boost/parameter/keyword.hpp>
#include <boost/mpl/placeholders.hpp>
#include <boost/accumulators/framework/accumulator_base.hpp>
#include <boost/accumulators/framework/extractor.hpp>
#include <boost/accumulators/numeric/functional.hpp>
#include <boost/accumulators/framework/parameters/sample.hpp>
#include <boost/accumulators/framework/depends_on.hpp>
#include <boost/accumulators/statistics_fwd.hpp>
#include <boost/accumulators/statistics/count.hpp>
#include <boost/accumulators/statistics/sum.hpp>
#include <boost/accumulators/statistics/times2_iterator.hpp>
#include <boost/accumulators/statistics/extended_p_square.hpp>

namespace boost { namespace accumulators
{

namespace impl
{
    ///////////////////////////////////////////////////////////////////////////////
    // weighted_extended_p_square_impl
    //  multiple quantile estimation with weighted samples
    /**
        @brief Multiple quantile estimation with the extended \f$P^2\f$ algorithm for weighted samples

        This version of the extended \f$P^2\f$ algorithm extends the extended \f$P^2\f$ algorithm to
        support weighted samples. The extended \f$P^2\f$ algorithm dynamically estimates several
        quantiles without storing samples. Assume that \f$m\f$ quantiles
        \f$\xi_{p_1}, \ldots, \xi_{p_m}\f$ are to be estimated. Instead of storing the whole sample
        cumulative distribution, the algorithm maintains only \f$m+2\f$ principal markers and
        \f$m+1\f$ middle markers, whose positions are updated with each sample and whose heights
        are adjusted (if necessary) using a piecewise-parablic formula. The heights of the principal
        markers are the current estimates of the quantiles and are returned as an iterator range.

        For further details, see

        K. E. E. Raatikainen, Simultaneous estimation of several quantiles, Simulation, Volume 49,
        Number 4 (October), 1986, p. 159-164.

        The extended \f$ P^2 \f$ algorithm generalizes the \f$ P^2 \f$ algorithm of

        R. Jain and I. Chlamtac, The P^2 algorithm for dynamic calculation of quantiles and
        histograms without storing observations, Communications of the ACM,
        Volume 28 (October), Number 10, 1985, p. 1076-1085.

        @param extended_p_square_probabilities A vector of quantile probabilities.
    */
    template<typename Sample, typename Weight>
    struct weighted_extended_p_square_impl
      : accumulator_base
    {
        typedef typename numeric::functional::multiplies<Sample, Weight>::result_type weighted_sample;
        typedef typename numeric::functional::average<weighted_sample, std::size_t>::result_type float_type;
        typedef std::vector<float_type> array_type;
        // for boost::result_of
        typedef iterator_range<
            detail::lvalue_index_iterator<
                permutation_iterator<
                    typename array_type::const_iterator
                  , detail::times2_iterator
                >
            >
        > result_type;

        template<typename Args>
        weighted_extended_p_square_impl(Args const &args)
          : probabilities(
                boost::begin(args[extended_p_square_probabilities])
              , boost::end(args[extended_p_square_probabilities])
            )
          , heights(2 * probabilities.size() + 3)
          , actual_positions(heights.size())
          , desired_positions(heights.size())
        {
        }

        template<typename Args>
        void operator ()(Args const &args)
        {
            std::size_t cnt = count(args);
            std::size_t sample_cell = 1; // k
            std::size_t num_quantiles = this->probabilities.size();

            // m+2 principal markers and m+1 middle markers
            std::size_t num_markers = 2 * num_quantiles + 3;

            // first accumulate num_markers samples
            if(cnt <= num_markers)
            {
                this->heights[cnt - 1] = args[sample];
                this->actual_positions[cnt - 1] = args[weight];

                // complete the initialization of heights (and actual_positions) by sorting
                if(cnt == num_markers)
                {
                    // TODO: we need to sort the initial samples (in heights) in ascending order and
                    // sort their weights (in actual_positions) the same way. The following lines do
                    // it, but there must be a better and more efficient way of doing this.
                    typename array_type::iterator it_begin, it_end, it_min;

                    it_begin = this->heights.begin();
                    it_end   = this->heights.end();

                    std::size_t pos = 0;

                    while (it_begin != it_end)
                    {
                        it_min = std::min_element(it_begin, it_end);
                        std::size_t d = std::distance(it_begin, it_min);
                        std::swap(*it_begin, *it_min);
                        std::swap(this->actual_positions[pos], this->actual_positions[pos + d]);
                        ++it_begin;
                        ++pos;
                    }

                    // calculate correct initial actual positions
                    for (std::size_t i = 1; i < num_markers; ++i)
                    {
                        actual_positions[i] += actual_positions[i - 1];
                    }
                }
            }
            else
            {
                if(args[sample] < this->heights[0])
                {
                    this->heights[0] = args[sample];
                    this->actual_positions[0] = args[weight];
                    sample_cell = 1;
                }
                else if(args[sample] >= this->heights[num_markers - 1])
                {
                    this->heights[num_markers - 1] = args[sample];
                    sample_cell = num_markers - 1;
                }
                else
                {
                    // find cell k = sample_cell such that heights[k-1] <= sample < heights[k]

                    typedef typename array_type::iterator iterator;
                    iterator it = std::upper_bound(
                        this->heights.begin()
                      , this->heights.end()
                      , args[sample]
                    );

                    sample_cell = std::distance(this->heights.begin(), it);
                }

                // update actual position of all markers above sample_cell
                for(std::size_t i = sample_cell; i < num_markers; ++i)
                {
                    this->actual_positions[i] += args[weight];
                }

                // compute desired positions
                {
                    this->desired_positions[0] = this->actual_positions[0];
                    this->desired_positions[num_markers - 1] = sum_of_weights(args);
                    this->desired_positions[1] = (sum_of_weights(args) - this->actual_positions[0]) * probabilities[0]
                                              / 2. + this->actual_positions[0];
                    this->desired_positions[num_markers - 2] = (sum_of_weights(args) - this->actual_positions[0])
                                                            * (probabilities[num_quantiles - 1] + 1.)
                                                            / 2. + this->actual_positions[0];

                    for (std::size_t i = 0; i < num_quantiles; ++i)
                    {
                        this->desired_positions[2 * i + 2] = (sum_of_weights(args) - this->actual_positions[0])
                                                          * probabilities[i] + this->actual_positions[0];
                    }

                    for (std::size_t i = 1; i < num_quantiles; ++i)
                    {
                        this->desired_positions[2 * i + 1] = (sum_of_weights(args) - this->actual_positions[0])
                                                      * (probabilities[i - 1] + probabilities[i])
                                                      / 2. + this->actual_positions[0];
                    }
                }

                // adjust heights and actual_positions of markers 1 to num_markers - 2 if necessary
                for (std::size_t i = 1; i <= num_markers - 2; ++i)
                {
                    // offset to desired position
                    float_type d = this->desired_positions[i] - this->actual_positions[i];

                    // offset to next position
                    float_type dp = this->actual_positions[i + 1] - this->actual_positions[i];

                    // offset to previous position
                    float_type dm = this->actual_positions[i - 1] - this->actual_positions[i];

                    // height ds
                    float_type hp = (this->heights[i + 1] - this->heights[i]) / dp;
                    float_type hm = (this->heights[i - 1] - this->heights[i]) / dm;

                    if((d >= 1 && dp > 1) || (d <= -1 && dm < -1))
                    {
                        short sign_d = static_cast<short>(d / std::abs(d));

                        float_type h = this->heights[i] + sign_d / (dp - dm) * ((sign_d - dm)*hp + (dp - sign_d) * hm);

                        // try adjusting heights[i] using p-squared formula
                        if(this->heights[i - 1] < h && h < this->heights[i + 1])
                        {
                            this->heights[i] = h;
                        }
                        else
                        {
                            // use linear formula
                            if(d > 0)
                            {
                                this->heights[i] += hp;
                            }
                            if(d < 0)
                            {
                                this->heights[i] -= hm;
                            }
                        }
                        this->actual_positions[i] += sign_d;
                    }
                }
            }
        }

        result_type result(dont_care) const
        {
            // for i in [1,probabilities.size()], return heights[i * 2]
            detail::times2_iterator idx_begin = detail::make_times2_iterator(1);
            detail::times2_iterator idx_end = detail::make_times2_iterator(this->probabilities.size() + 1);

            return result_type(
                make_permutation_iterator(this->heights.begin(), idx_begin)
              , make_permutation_iterator(this->heights.begin(), idx_end)
            );
        }

    private:
        array_type probabilities;         // the quantile probabilities
        array_type heights;               // q_i
        array_type actual_positions;      // n_i
        array_type desired_positions;     // d_i
    };

} // namespace impl

///////////////////////////////////////////////////////////////////////////////
// tag::weighted_extended_p_square
//
namespace tag
{
    struct weighted_extended_p_square
      : depends_on<count, sum_of_weights>
      , extended_p_square_probabilities
    {
        typedef accumulators::impl::weighted_extended_p_square_impl<mpl::_1, mpl::_2> impl;
    };
}

///////////////////////////////////////////////////////////////////////////////
// extract::weighted_extended_p_square
//
namespace extract
{
    extractor<tag::weighted_extended_p_square> const weighted_extended_p_square = {};

    BOOST_ACCUMULATORS_IGNORE_GLOBAL(weighted_extended_p_square)
}

using extract::weighted_extended_p_square;

}} // namespace boost::accumulators

#endif


//
// Created by beck on 24/9/18.
//
#include <stdio.h>
#include <limits>
#include <string.h>
#include <Eigen/Dense>
#include <ros/ros.h>
#include "kalman.h"
#include "geometry_msgs/Point32.h"
#include "sensor_msgs/PointCloud.h"

ros::Subscriber radar_cloud_sub;
ros::Publisher filtered_cloud_pub;

bool filter_initialized = false;

KalmanFilter kf;

void
initialize_radar_filter(double distance)
{
    int n = 2;
    int m = 1;

    double dt = 1.0 / 14; // 14Hz from the radar

    Eigen::MatrixXd A(n, n);
    Eigen::MatrixXd H(m, n);
    Eigen::MatrixXd Q(n, n);
    Eigen::MatrixXd R(m, m);
    Eigen::MatrixXd P(n, n);

    // constant velocity model
    A << 1, dt,
         0, 1;
    H << 1, 0;
    Q << 1, 0,
         0, 1;
    R << 0.01;
    P << 1, 0,
         0, 1;

    // initialize the filter
    KalmanFilter kf_init(dt, A, H, Q, R, P);
    Eigen::VectorXd x0(n);
    x0 << distance, 0;

    // create a prediction with Kalman filter
    kf_init.init(ros::Time::now().toSec(), x0);
    kf = kf_init;
	filter_initialized = true;
	ROS_INFO("initialized the radar kalman filter at %f", distance);
}

// change the ID to the sorted order
void
publish_message(const int id, const std_msgs::Header& header)
{
    geometry_msgs::Point32 point;

	point.x = kf.state()[0]; // estimated state
	point.y = kf.state()[1]; // estimated velocity
	point.z = id;

	sensor_msgs::PointCloud ptCloud;
	ptCloud.header = header;
	ptCloud.points.push_back(point);
	filtered_cloud_pub.publish(ptCloud);
}

void
radar_callback(const sensor_msgs::PointCloud::ConstPtr pc_ptr)
{
    int n = 2;
    int m = 1;
    int s = pc_ptr->points.size();
    int i;

    if (s > 0) {
        if (!filter_initialized) {
            double distance = 9.0;

            for ( i = 0; i < s; ++i) {
                double x_ = pc_ptr->points[i].x;
                if (x_ > 9 && x_ < 10)
                    distance = x_;
            }
            initialize_radar_filter(distance);
        }
        else
        {
            double minChiSquare = std::numeric_limits<double>::max();
            int chi_index = s;
            bool found_nearest_point = false;
			int id = 0;

            // compare the chi-square to select into one
            for ( i = 0; i < s; ++i) {
                Eigen::VectorXd z(m);
                z << pc_ptr->points[i].x;

                double chiSquare = kf.chiSquare(z);
                if (chiSquare < minChiSquare) {
                    minChiSquare = chiSquare;
                    chi_index = i;
                }
            }
            double residual = pc_ptr->points[chi_index].x - kf.state()[0];
            if (residual < 1 && residual > -1)
                found_nearest_point = true;

            if (found_nearest_point) {
                ROS_INFO("found the number at %d with X2 %f", chi_index, minChiSquare);

                kf.propagate();

                Eigen::VectorXd z_sel(m);
                z_sel << pc_ptr->points[chi_index].x;
                kf.update(z_sel);

                publish_message(id, pc_ptr->header);
            }
			else {
				ROS_INFO("radar does not track the obstacle");
				// kf.propagate();

                publish_message(id, pc_ptr->header);
			}
        }
    }
	else {
		ROS_INFO("no point available");
	}
}

int main(int argc, char **argv)
{
    ros::init(argc, argv, "radar_preprocess");
    ros::NodeHandle n = ros::NodeHandle("~");

    radar_cloud_sub
    = n.subscribe<sensor_msgs::PointCloud>("/inf24g/inf24radar",
                                            10,
                                            radar_callback);
    filtered_cloud_pub
    = n.advertise<sensor_msgs::PointCloud>("/filtered_radar", 10);

    ros::spin();

    return 0;
}