//| Copyright Inria May 2015
//| This project has received funding from the European Research Council (ERC) under
//| the European Union's Horizon 2020 research and innovation programme (grant
//| agreement No 637972) - see http://www.resibots.eu
//|
//| Contributor(s):
//|   - Jean-Baptiste Mouret (jean-baptiste.mouret@inria.fr)
//|   - Antoine Cully (antoinecully@gmail.com)
//|   - Konstantinos Chatzilygeroudis (konstantinos.chatzilygeroudis@inria.fr)
//|   - Federico Allocati (fede.allocati@gmail.com)
//|   - Vaios Papaspyros (b.papaspyros@gmail.com)
//|   - Roberto Rama (bertoski@gmail.com)
//|
//| This software is a computer library whose purpose is to optimize continuous,
//| black-box functions. It mainly implements Gaussian processes and Bayesian
//| optimization.
//| Main repository: http://github.com/resibots/limbo
//| Documentation: http://www.resibots.eu/limbo
//|
//| This software is governed by the CeCILL-C license under French law and
//| abiding by the rules of distribution of free software.  You can  use,
//| modify and/ or redistribute the software under the terms of the CeCILL-C
//| license as circulated by CEA, CNRS and INRIA at the following URL
//| "http://www.cecill.info".
//|
//| As a counterpart to the access to the source code and  rights to copy,
//| modify and redistribute granted by the license, users are provided only
//| with a limited warranty  and the software's author,  the holder of the
//| economic rights,  and the successive licensors  have only  limited
//| liability.
//|
//| In this respect, the user's attention is drawn to the risks associated
//| with loading,  using,  modifying and/or developing or reproducing the
//| software by the user in light of its specific status of free software,
//| that may mean  that it is complicated to manipulate,  and  that  also
//| therefore means  that it is reserved for developers  and  experienced
//| professionals having in-depth computer knowledge. Users are therefore
//| encouraged to load and test the software's suitability as regards their
//| requirements in conditions enabling the security of their systems and/or
//| data to be ensured and,  more generally, to use and operate it in the
//| same conditions as regards security.
//|
//| The fact that you are presently reading this means that you have had
//| knowledge of the CeCILL-C license and that you accept its terms.
//|
#ifndef LIMBO_STOP_MAX_PREDICTED_VALUE_HPP
#define LIMBO_STOP_MAX_PREDICTED_VALUE_HPP

#include <iostream>

#include <boost/parameter/aux_/void.hpp>

#include <Eigen/Core>

#include <limbo/tools/macros.hpp>
#include <limbo/tools/random_generator.hpp>

namespace limbo {
    namespace defaults {
        struct stop_maxpredictedvalue {
            ///@ingroup stop_defaults
            BO_PARAM(double, ratio, 0.9);
        };
    }
    namespace stop {
        ///@ingroup stop
        ///Stop once the value for the best sample is above : ratio * (best value predicted by the model)
        ///
        ///Parameter: double ratio
        template <typename Params, typename Optimizer = boost::parameter::void_>
        struct MaxPredictedValue {

            MaxPredictedValue() {}

            template <typename BO, typename AggregatorFunction>
            bool operator()(const BO& bo, const AggregatorFunction& afun) const
            {
                // Prevent instantiation of GPMean if there are no observed samples
                if (bo.observations().size() == 0)
                    return false;

                auto optimizer = _get_optimizer(typename BO::acqui_optimizer_t(), Optimizer());
                Eigen::VectorXd starting_point = tools::random_vector(bo.model().dim_in());
                auto model_optimization =
                    [&](const Eigen::VectorXd& x, bool g) { return opt::no_grad(afun(bo.model().mu(x))); };
                auto x = optimizer(model_optimization, starting_point, true);
                double val = afun(bo.model().mu(x));

                if (bo.observations().size() == 0 || afun(bo.best_observation(afun)) <= Params::stop_maxpredictedvalue::ratio() * val)
                    return false;
                else {
                    std::cout << "stop caused by Max predicted value reached. Threshold: "
                              << Params::stop_maxpredictedvalue::ratio() * val
                              << " max observations: " << afun(bo.best_observation(afun)) << std::endl;
                    return true;
                }
            }

        protected:
            template <typename Model, typename AggregatorFunction>
            struct ModelMeanOptimization {
            public:
                ModelMeanOptimization(const Model& model, const AggregatorFunction& afun, const Eigen::VectorXd& init) : _model(model), _afun(afun), _init(init) {}

                double utility(const Eigen::VectorXd& v) const
                {
                    return _afun(_model.mu(v));
                }

                size_t param_size() const
                {
                    return _model.dim_in();
                }

                const Eigen::VectorXd& init() const
                {
                    return _init;
                }

            protected:
                const Model& _model;
                const AggregatorFunction& _afun;
                const Eigen::VectorXd _init;
            };

            template <typename Model, typename AggregatorFunction>
            inline ModelMeanOptimization<Model, AggregatorFunction> _make_model_mean_optimization(const Model& model, const AggregatorFunction& afun, const Eigen::VectorXd& init) const
            {
                return ModelMeanOptimization<Model, AggregatorFunction>(model, afun, init);
            }

            template <typename BoAcquiOpt>
            inline BoAcquiOpt _get_optimizer(const BoAcquiOpt& bo_acqui_opt, boost::parameter::void_) const
            {
                return bo_acqui_opt;
            }

            template <typename BoAcquiOpt, typename CurrentOpt>
            inline CurrentOpt _get_optimizer(const BoAcquiOpt&, const CurrentOpt& current_opt) const
            {
                return current_opt;
            }
        };
    }
}

#endif


/**
 * \file dcs/algorithm/subset.hpp
 *
 * \brief Generate the power set of a set in lexicographic order.
 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 *
 *

 *
 * Copyright 2013 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_COMMONS_ALGORITHM_SUBSET_HPP
#define DCS_COMMONS_ALGORITHM_SUBSET_HPP

#include <algorithm>
#include <boost/dynamic_bitset.hpp>
#include <boost/math/special_functions/binomial.hpp>
#include <cstddef>
#include <dcs/assert.hpp>
#include <dcs/debug.hpp>
#include <dcs/exception.hpp>
#include <iostream>
#include <iterator>
#include <stdexcept>
#include <vector>

namespace dcs { namespace algorithm {

/**
 * \brief Traits class for subset types
 *
 * \tparam ValueT The value type
 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 */
template <typename ValueT>
struct subset_traits
{
	typedef ValueT value_type;
	typedef ::std::vector<value_type> subset_container;///< DEPRECATED
	typedef typename subset_container::iterator subset_container_iterator;///< DEPRECATED
	typedef typename subset_container::const_iterator subset_container_const_iterator;///< DEPRECATED
	typedef ::std::vector<value_type> element_container;
	typedef typename element_container::iterator element_iterator;
	typedef typename element_container::const_iterator element_const_iterator;
}; // subset_traits

/**
 * \brief Class to generate in lexicographic order all subsets
 *
 * Given a set N={0,1,...,n-1} of n elements, this class iteratively generates
 * all subset S of N, possibly included the empty set, in lexicographic order,
 * that is, to generate a subset containing the i-th element we generate all
 * subset containing the preceding 0,1,...(i-1)-th elements.
 * For instance, for a set of 4 elements, the subset generation in lexicographic
 * order produces the following sequence:
 *

 *  \emptyset,
 *  {0},
 *  {1},
 *  {0,1},
 *  {2},
 *  {0,2},
 *  {1,2},
 *  {0,1,2},
 *  {3},
 *  {0,3},
 *  {1,3},
 *  {2,3},
 *  {0,1,3},
 *  {0,2,3}
 *  {1,2,3},
 *  {0,1,2,3}
 *

 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 */
class lexicographic_subset
{
	public: class const_iterator;

	friend class const_iterator;

	private: typedef lexicographic_subset self_type;
	private: typedef ::boost::dynamic_bitset<> impl_type;
	private: typedef typename impl_type::size_type size_type;
	private: typedef unsigned long word_type;

	public: explicit lexicographic_subset(::std::size_t n, bool empty_set=true)
	: n_(n),
	  empty_set_(empty_set),
	  bits_(n, empty_set ? 0 : 1),
	  has_prev_(false),
	  has_next_(n_ > 0 ? true : false)
	{
		DCS_ASSERT(n_ > 0,
				   DCS_EXCEPTION_THROW(::std::invalid_argument,
									   "Number of elements must be positive"));
	}

	public: ::std::size_t max_size() const
	{
		return n_;
	}

	public: ::std::size_t size() const
	{
		return bits_.count();
	}

	public: ::std::size_t count() const
	{
		const ::std::size_t c = 1 << n_; // 2^n

		return empty_set_ ? c : (c-1);
	}

	public: self_type& operator++()
	{
		DCS_ASSERT(has_next_,
				   DCS_EXCEPTION_THROW(::std::overflow_error,
									   "No following subsets"));

		has_next_ = bits_.count() < n_;

		if (has_next_)
		{
			bits_ = impl_type(n_, bits_.to_ulong()+1);
		}

		has_prev_ = bits_.to_ulong() > (empty_set_ ? 0 : 1);

		return *this;
	}

	public: bool has_next() const
	{
		return has_next_;
	}

	public: self_type& operator--()
	{
		DCS_ASSERT(has_prev_,
				   DCS_EXCEPTION_THROW(::std::underflow_error,
									   "No preceding subsets"));

		has_prev_ = bits_.to_ulong() > (empty_set_ ? 0 : 1);

		if (has_prev_)
		{
			bits_ = impl_type(n_, bits_.to_ulong()-1);
		}

		has_next_ = bits_.to_ulong() < n_;

		return *this;
	}

	public: bool has_prev() const
	{
		return has_prev_;
	}

	public: typename ::std::vector<size_type> operator()() const
	{
		::std::vector<size_type> subset;

		for (size_type pos = bits_.find_first();
			 pos != impl_type::npos;
			 pos = bits_.find_next(pos))
		{
			subset.push_back(pos);
		}

		return subset;
	}

	//public: template <typename ElemT, typename IterT>
	public: template <typename ElemT>
			typename subset_traits<ElemT>::element_container operator()(::std::vector<ElemT> const& v) const
	{
		DCS_ASSERT(v.size() == n_,
				   DCS_EXCEPTION_THROW(::std::invalid_argument,
									   "Size does not match"));

		typename subset_traits<ElemT>::element_container subset;

		for (size_type pos = bits_.find_first();
			 pos != impl_type::npos;
			 pos = bits_.find_next(pos))
		{
			subset.push_back(v[pos]);
		}

		return subset;
	}

	public: template <typename IterT>
			typename subset_traits< typename ::std::iterator_traits<IterT>::value_type >::element_container operator()(IterT first, IterT last) const
	{
			return this->operator()(::std::vector<typename ::std::iterator_traits<IterT>::value_type>(first, last));
	}

	public: const_iterator begin() const
	{
		return const_iterator(this, bits_.any() ? bits_.find_first() : impl_type::npos);
	}

	public: const_iterator end() const
	{
		return const_iterator(this, impl_type::npos);
	}

	public: class const_iterator: public ::std::iterator< ::std::bidirectional_iterator_tag,
														 typename lexicographic_subset::impl_type::size_type const>
	{
		private: typedef ::std::iterator< ::std::bidirectional_iterator_tag,
										  typename lexicographic_subset::impl_type::size_type const> base_type;
		private: typedef lexicographic_subset container_type;
		private: typedef typename container_type::impl_type bitset_type;
		private: typedef typename container_type::size_type size_type;
		public: typedef typename base_type::value_type value_type;
		public: typedef typename base_type::difference_type difference_type;
		public: typedef typename base_type::pointer pointer;
		public: typedef typename base_type::reference reference;
		public: typedef typename base_type::iterator_category iterator_category;

		public: const_iterator(container_type const* p_subset, size_type pos)
		: p_sub_(p_subset),
		  pos_(pos)
		{
		}

		public: reference operator*() const
		{
			return pos_;
		}

		public: pointer operator->() const
		{
			return &(operator*());
		}

		public: const_iterator& operator++()
		{
			pos_ = p_sub_->bits_.find_next(pos_);

			return *this;
		}

		public: const_iterator operator++(int)
		{
			const_iterator tmp = *this;

			operator++();

			return tmp;
		}

		public: const_iterator& operator--()
		{
			size_type pos = bitset_type::npos;
			for (size_type pos2 = p_sub_->bits_.find_first();
				 pos2 != pos_;
				 pos2 = p_sub_->bits_.find_next(pos2))
			{
				pos = pos2;
			}

			pos_ = pos;

			return *this;
		}

		public: const_iterator operator--(int)
		{
			const_iterator tmp = *this;

			operator--();

			return tmp;
		}

		public: const_iterator operator+(difference_type n) const
		{
			const_iterator it = *this;

			if (n > 0)
			{
				while (n--)
				{
					++it;
				}
			}
			else
			{
				while (n++)
				{
					--it;
				}
			}

			return it;
		}

		public: const_iterator operator-(difference_type n) const
		{
			return operator+(-n);
		}

		public: const_iterator& operator+=(difference_type n)
		{
			if (n > 0)
			{
				while (n--)
				{
					operator++();
				}
			}
			else
			{
				while (n++)
				{
					operator--();
				}
			}

			return *this;
		}

		public: const_iterator& operator-=(difference_type n)
		{
			return operator+=(-n);
		}

		friend
		bool operator==(const_iterator const& lhs, const_iterator const& rhs)
		{
			return lhs.p_sub_ == rhs.p_sub_ && lhs.pos_ == rhs.pos_;
		}

		friend
		bool operator!=(const_iterator const& lhs, const_iterator const& rhs)
		{
			return !(lhs == rhs);
		}

		private: container_type const* p_sub_;
		private: size_type pos_;
	}; // const_iterator

	private: ::std::size_t n_; ///< The max number of elements
	private: bool empty_set_; ///< Flag to enable or disable the inclusion of the empty set
	private: impl_type bits_; ///< The subset implementation
	private: bool has_prev_;
	private: bool has_next_;
}; // lexicographic_subset

/**
 * \brief Class to generate in lexicographic order all subsets of a specific size
 *
 * Given a set N={0,1,...,n-1} of n elements, this class iteratively generates
 * all subset S of N of size 0<=k<=n, in lexicographic order.
 * For instance, for a set of 4 elements, the generation of subset of size 2 in
 * lexicographic order produces the following sequence:
 *

 *  \emptyset,
 *  {0,1},
 *  {0,2},
 *  {1,2},
 *  {0,3},
 *  {1,3},
 *  {2,3}
 *

 *
 * \author Marco Guazzone (marco.guazzone@gmail.com)
 */
class lexicographic_k_subset
{
	public: class const_iterator;

	friend class const_iterator;

	private: typedef lexicographic_k_subset self_type;
	private: typedef ::boost::dynamic_bitset<> impl_type;
	private: typedef typename impl_type::size_type size_type;
	private: typedef unsigned long word_type;

	public: explicit lexicographic_k_subset(::std::size_t n, ::std::size_t k)
	: n_(n),
	  k_(k),
	  bits_(n, (1 << k_)-1),
	  has_prev_(false),
	  has_next_(n_ > 0 && k_ > 0 ? true : false)
	{
		DCS_ASSERT(n_ > 0,
				   DCS_EXCEPTION_THROW(::std::invalid_argument,
									   "Number of elements must be positive"));
		DCS_ASSERT(n_ >= k,
				   DCS_EXCEPTION_THROW(::std::invalid_argument,
									   "Size of subset must be non negative"));
	}

	public: ::std::size_t max_size() const
	{
		return k_;
	}

	public: ::std::size_t size() const
	{
		return bits_.count();
	}

	public: ::std::size_t count() const
	{
		return static_cast<std::size_t>(boost::math::binomial_coefficient<double>(n_, k_));
	}

	public: self_type& operator++()
	{
		DCS_ASSERT(has_next_,
				   DCS_EXCEPTION_THROW(::std::overflow_error,
									   "No following subsets"));

		if (k_ > 0)
		{
			unsigned long v = bits_.to_ulong();

			const unsigned long lo = v & ~(v - 1);       // lowest one bit
			const unsigned long lz = (v + lo) & ~v;      // lowest zero bit above lo
			v |= lz;                     // add lz to the set
			v &= ~(lz - 1);              // reset bits below lz
			v |= (lz / lo / 2) - 1;      // put back right number of bits at end

			has_next_ = !(v & 1 << n_);

			if (has_next_)
			{
				bits_ = impl_type(n_, v);
			}

			has_prev_ = (bits_.to_ulong() != static_cast<unsigned long>((1 << k_)-1));
		}

		return *this;
	}

	public: bool has_next() const
	{
		return has_next_;
	}

	public: self_type& operator--()
	{
		DCS_ASSERT(has_prev_,
				   DCS_EXCEPTION_THROW(::std::underflow_error,
									   "No preceding subsets"));

		if (k_ > 0)
		{
			const unsigned long w = bits_.to_ulong();
			const unsigned long min = (1 << k_)-1;
			unsigned long v = min;
			unsigned long p = v;

			while (v != w)
			{
				p = v;

				const unsigned long lo = v & ~(v - 1);       // lowest one bit
				const unsigned long lz = (v + lo) & ~v;      // lowest zero bit above lo
				v |= lz;                     // add lz to the set
				v &= ~(lz - 1);              // reset bits below lz
				v |= (lz / lo / 2) - 1;      // put back right number of bits at end
			}

			has_prev_ = (p != min) || (w != min);

//			if (has_prev_)
//			{
				bits_ = impl_type(n_, p);
//			}

			has_next_ = !(p & 1 << n_);
		}

		return *this;
	}

	public: bool has_prev() const
	{
		return has_prev_;
	}

	//public: template <typename ElemT, typename IterT>
	public: template <typename ElemT>
			typename subset_traits<ElemT>::element_container operator()(::std::vector<ElemT> const& v) const
	{
		DCS_ASSERT(v.size() == n_,
				   DCS_EXCEPTION_THROW(::std::invalid_argument,
									   "Size does not match"));

		typename subset_traits<ElemT>::element_container subset;

		for (size_type pos = bits_.find_first();
			 pos != impl_type::npos;
			 pos = bits_.find_next(pos))
		{
			subset.push_back(v[pos]);
		}

		return subset;
	}

	public: template <typename IterT>
			typename subset_traits< typename ::std::iterator_traits<IterT>::value_type >::element_container operator()(IterT first, IterT last) const
	{
			return this->operator()(::std::vector<typename ::std::iterator_traits<IterT>::value_type>(first, last));
	}

	public: const_iterator begin() const
	{
		return const_iterator(this, bits_.any() ? bits_.find_first() : impl_type::npos);
	}

	public: const_iterator end() const
	{
		return const_iterator(this, impl_type::npos);
	}

	public: class const_iterator: public ::std::iterator< ::std::bidirectional_iterator_tag,
														 typename lexicographic_k_subset::impl_type::size_type const>
	{
		private: typedef ::std::iterator< ::std::bidirectional_iterator_tag,
										  typename lexicographic_k_subset::impl_type::size_type const> base_type;
		private: typedef lexicographic_k_subset container_type;
		private: typedef typename container_type::impl_type bitset_type;
		private: typedef typename container_type::size_type size_type;
		public: typedef typename base_type::value_type value_type;
		public: typedef typename base_type::difference_type difference_type;
		public: typedef typename base_type::pointer pointer;
		public: typedef typename base_type::reference reference;
		public: typedef typename base_type::iterator_category iterator_category;

		public: const_iterator(container_type const* p_subset, size_type pos)
		: p_sub_(p_subset),
		  pos_(pos)
		{
		}

		public: reference operator*() const
		{
			return pos_;
		}

		public: pointer operator->() const
		{
			return &(operator*());
		}

		public: const_iterator& operator++()
		{
			pos_ = p_sub_->bits_.find_next(pos_);

			return *this;
		}

		public: const_iterator operator++(int)
		{
			const_iterator tmp = *this;

			operator++();

			return tmp;
		}

		public: const_iterator& operator--()
		{
			size_type pos = bitset_type::npos;
			for (size_type pos2 = p_sub_->bits_.find_first();
				 pos2 != pos_;
				 pos2 = p_sub_->bits_.find_next(pos2))
			{
				pos = pos2;
			}

			pos_ = pos;

			return *this;
		}

		public: const_iterator operator--(int)
		{
			const_iterator tmp = *this;

			operator--();

			return tmp;
		}

		public: const_iterator operator+(difference_type n) const
		{
			const_iterator it = *this;

			if (n > 0)
			{
				while (n--)
				{
					++it;
				}
			}
			else
			{
				while (n++)
				{
					--it;
				}
			}

			return it;
		}

		public: const_iterator operator-(difference_type n) const
		{
			return operator+(-n);
		}

		public: const_iterator& operator+=(difference_type n)
		{
			if (n > 0)
			{
				while (n--)
				{
					operator++();
				}
			}
			else
			{
				while (n++)
				{
					operator--();
				}
			}

			return *this;
		}

		public: const_iterator& operator-=(difference_type n)
		{
			return operator+=(-n);
		}

		friend
		bool operator==(const_iterator const& lhs, const_iterator const& rhs)
		{
			return lhs.p_sub_ == rhs.p_sub_ && lhs.pos_ == rhs.pos_;
		}

		friend
		bool operator!=(const_iterator const& lhs, const_iterator const& rhs)
		{
			return !(lhs == rhs);
		}

		private: container_type const* p_sub_;
		private: size_type pos_;
	}; // const_iterator

	private: ::std::size_t n_; ///< The number of elements of the set
	private: ::std::size_t k_; ///< The max number of elements of the subset
	private: impl_type bits_; ///< The subset implementation
	private: bool has_prev_;
	private: bool has_next_;
}; // lexicographic_k_subset

template <typename CharT, typename CharTraitsT>
::std::basic_ostream<CharT,CharTraitsT>& operator<<(::std::basic_ostream<CharT,CharTraitsT>& os,
													lexicographic_subset const& subset)
{
	os << '(';

	if (subset.size() > 0)
	{
		if (subset.size() > 1)
		{
			::std::copy(subset.begin(),
						subset.end()-1,
						::std::ostream_iterator< ::std::size_t >(os, " "));
		}

		os << *(subset.end()-1);
	}

	os << ')';

	return os;
}

template <typename CharT, typename CharTraitsT>
::std::basic_ostream<CharT,CharTraitsT>& operator<<(::std::basic_ostream<CharT,CharTraitsT>& os,
													lexicographic_k_subset const& subset)
{
	os << '(';

	if (subset.size() > 0)
	{
		if (subset.size() > 1)
		{
			::std::copy(subset.begin(),
						subset.end()-1,
						::std::ostream_iterator< ::std::size_t >(os, " "));
		}

		os << *(subset.end()-1);
	}

	os << ')';

	return os;
}

template <typename BidiIterT, typename SubsetT>
inline
typename subset_traits< typename ::std::iterator_traits<BidiIterT>::value_type >::element_container
next_subset(BidiIterT first,
			BidiIterT last,
			SubsetT& subset)
{
	typedef typename ::std::iterator_traits<BidiIterT>::value_type value_type;
	typedef typename subset_traits<value_type>::element_container element_container;

	element_container subs = subset(first, last);

	++subset;

	return subs;
}

template <typename BidiIterT, typename SubsetT>
inline
typename subset_traits< typename ::std::iterator_traits<BidiIterT>::value_type >::element_container
prev_subset(BidiIterT first,
			BidiIterT last,
			SubsetT& subset)
{
	typedef typename ::std::iterator_traits<BidiIterT>::value_type value_type;
	typedef typename subset_traits<value_type>::element_container element_container;

	element_container subs = subset(first, last);

	--subset;

	return subs;
}

/*
inline
::std::size_t count_subsets(n)
{
	const ::std::size_t m = n >> 1;
	const bool flag = ((m << 1) - n) > 0;

	::std::size_t c = 0;
	for (std::size_t i = 0; i < m; ++i)
	{
		c +=
	}
}
*/

}} // Namespace dcs::algorithm

#endif // DCS_COMMONS_ALGORITHM_SUBSET_HPP


// Copyright 2018 The Simons Foundation, Inc. - 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.

#ifndef NETKET_GRAPH_HAMILTONIAN_CC
#define NETKET_GRAPH_HAMILTONIAN_CC

#include <Eigen/Dense>
#include <array>
#include <unordered_map>
#include <vector>
#include "Graph/graph.hpp"
#include "Hilbert/abstract_hilbert.hpp"
#include "Utils/json_helper.hpp"
#include "abstract_operator.hpp"
#include "local_operator.hpp"

namespace netket {

// Graph Hamiltonian on an arbitrary graph
class GraphOperator : public AbstractOperator {
  const AbstractHilbert &hilbert_;

  // Arbitrary graph
  const AbstractGraph &graph_;

  LocalOperator operator_;

  const int nvertices_;

 public:
  using SiteType = std::vector<int>;
  using OMatType = LocalOperator::MatType;
  using OVecType = std::vector<OMatType>;
  using VectorType = AbstractOperator::VectorType;
  using VectorRefType = AbstractOperator::VectorRefType;
  using VectorConstRefType = AbstractOperator::VectorConstRefType;

  explicit GraphOperator(const AbstractHilbert &hilbert,
                         OVecType siteops = OVecType(),
                         OVecType bondops = OVecType(),
                         std::vector<int> bondops_colors = std::vector<int>())
      : hilbert_(hilbert),
        graph_(hilbert.GetGraph()),
        operator_(hilbert),
        nvertices_(hilbert.Size()) {
    // Create the local operator as the sum of all site and bond operators
    // Ensure that at least one of SiteOps and BondOps was initialized
    if (!siteops.size() && !bondops.size()) {
      throw InvalidInputError("Must input at least SiteOps or BondOps");
    }

    std::vector<int> op_color;
    if (bondops_colors.size() == 0) {
      op_color = std::vector<int>(bondops.size(), 0);
    } else {
      op_color = bondops_colors;
    }

    // Site operators
    if (siteops.size() > 0) {
      for (int i = 0; i < nvertices_; i++) {
        for (std::size_t j = 0; j < siteops.size(); j++) {
          operator_ += LocalOperator(hilbert_, siteops[j], std::vector<int>{i});
        }
      }
    }

    // Bond operators
    if (bondops.size() != op_color.size()) {
      throw InvalidInputError(
          "The bond Hamiltonian definition is inconsistent."
          "The sizes of BondOps and BondOpColors do not match.");
    }

    if (bondops.size() > 0) {
      // Use EdgeColors to populate operators
      for (auto const &kv : graph_.EdgeColors()) {
        for (std::size_t c = 0; c < op_color.size(); c++) {
          if (op_color[c] == kv.second && kv.first[0] < kv.first[1]) {
            std::vector<int> edge = {kv.first[0], kv.first[1]};
            operator_ += LocalOperator(hilbert_, bondops[c], edge);
          }
        }
      }
    }
  }

  // Constructor to be used when overloading operators
  explicit GraphOperator(const AbstractHilbert &hilbert,
                         const LocalOperator &lop)
      : hilbert_(hilbert),
        graph_(hilbert.GetGraph()),
        operator_(lop),
        nvertices_(hilbert.Size()) {}

  friend GraphOperator operator+(const GraphOperator &lhs,
                                 const GraphOperator &rhs) {
    assert(rhs.graph_.Size() == lhs.graph_.Size());

    auto lop = lhs.operator_;
    auto rop = rhs.operator_;

    return GraphOperator(lhs.GetHilbert(), lop + rop);
  }

  void FindConn(VectorConstRefType v, std::vector<Complex> &mel,
                std::vector<std::vector<int>> &connectors,
                std::vector<std::vector<double>> &newconfs) const override {
    operator_.FindConn(v, mel, connectors, newconfs);
  }

  const AbstractHilbert &GetHilbert() const noexcept override {
    return hilbert_;
  }
};  // namespace netket
}  // namespace netket
#endif


//////////////////////////////////////////////////////////////////////////////
//
// (C) Copyright Ion Gaztanaga 2015-2016.
// 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)
//
// See http://www.boost.org/libs/move for documentation.
//
//////////////////////////////////////////////////////////////////////////////

//! \file

#ifndef BOOST_MOVE_DETAIL_MERGE_SORT_HPP
#define BOOST_MOVE_DETAIL_MERGE_SORT_HPP

#ifndef BOOST_CONFIG_HPP
#  include <boost/config.hpp>
#endif
#
#if defined(BOOST_HAS_PRAGMA_ONCE)
#  pragma once
#endif

#include <boost/move/detail/config_begin.hpp>
#include <boost/move/detail/workaround.hpp>

#include <boost/move/utility_core.hpp>
#include <boost/move/algo/move.hpp>
#include <boost/move/algo/detail/merge.hpp>
#include <boost/move/detail/iterator_traits.hpp>
#include <boost/move/adl_move_swap.hpp>
#include <boost/move/detail/destruct_n.hpp>
#include <boost/move/algo/detail/insertion_sort.hpp>
#include <cassert>

#if defined(BOOST_CLANG) || (defined(BOOST_GCC) && (BOOST_GCC >= 40600))
#pragma GCC diagnostic push
//#pragma GCC diagnostic ignored "-Wsign-conversion"
#endif

namespace boost {
namespace movelib {

// @cond

static const unsigned MergeSortInsertionSortThreshold = 16;

template <class RandIt, class Compare>
void inplace_stable_sort(RandIt first, RandIt last, Compare comp)
{
   typedef typename iterator_traits<RandIt>::size_type  size_type;
   if (size_type(last - first) <= size_type(MergeSortInsertionSortThreshold)) {
      insertion_sort(first, last, comp);
      return;
   }
   RandIt middle = first + (last - first) / 2;
   inplace_stable_sort(first, middle, comp);
   inplace_stable_sort(middle, last, comp);
   merge_bufferless_ONlogN_recursive
      (first, middle, last, size_type(middle - first), size_type(last - middle), comp);
}

// @endcond

template<class RandIt, class RandIt2, class Compare>
void merge_sort_copy( RandIt first, RandIt last
                   , RandIt2 dest, Compare comp)
{
   typedef typename iterator_traits<RandIt>::size_type         size_type;

   size_type const count = size_type(last - first);
   if(count <= MergeSortInsertionSortThreshold){
      insertion_sort_copy(first, last, dest, comp);
   }
   else{
      size_type const half = size_type(count/2u);
      merge_sort_copy(first + half, last        , dest+half   , comp);
      merge_sort_copy(first       , first + half, first + half, comp);
      merge_with_right_placed
         ( first + half, first + half + half
         , dest, dest+half, dest + count
         , comp);
   }
}

template<class RandIt, class RandItRaw, class Compare>
void merge_sort_uninitialized_copy( RandIt first, RandIt last
                                 , RandItRaw uninitialized
                                 , Compare comp)
{
   typedef typename iterator_traits<RandIt>::size_type       size_type;
   typedef typename iterator_traits<RandIt>::value_type value_type;

   size_type const count = size_type(last - first);
   if(count <= MergeSortInsertionSortThreshold){
      insertion_sort_uninitialized_copy(first, last, uninitialized, comp);
   }
   else{
      size_type const half = count/2;
      merge_sort_uninitialized_copy(first + half, last, uninitialized + half, comp);
      destruct_n<value_type, RandItRaw> d(uninitialized+half);
      d.incr(size_type(count-half));
      merge_sort_copy(first, first + half, first + half, comp);
      uninitialized_merge_with_right_placed
         ( first + half, first + half + half
         , uninitialized, uninitialized+half, uninitialized+count
         , comp);
      d.release();
   }
}

template<class RandIt, class RandItRaw, class Compare>
void merge_sort( RandIt first, RandIt last, Compare comp
               , RandItRaw uninitialized)
{
   typedef typename iterator_traits<RandIt>::size_type       size_type;
   typedef typename iterator_traits<RandIt>::value_type      value_type;

   size_type const count = size_type(last - first);
   if(count <= MergeSortInsertionSortThreshold){
      insertion_sort(first, last, comp);
   }
   else{
      size_type const half = size_type(count/2u);
      size_type const rest = size_type(count -  half);
      RandIt const half_it = first + half;
      RandIt const rest_it = first + rest;

      merge_sort_uninitialized_copy(half_it, last, uninitialized, comp);
      destruct_n<value_type, RandItRaw> d(uninitialized);
      d.incr(rest);
      merge_sort_copy(first, half_it, rest_it, comp);
      merge_with_right_placed
         ( uninitialized, uninitialized + rest
         , first, rest_it, last, antistable<Compare>(comp));
   }
}

///@cond

template<class RandIt, class RandItRaw, class Compare>
void merge_sort_with_constructed_buffer( RandIt first, RandIt last, Compare comp, RandItRaw buffer)
{
   typedef typename iterator_traits<RandIt>::size_type       size_type;

   size_type const count = size_type(last - first);
   if(count <= MergeSortInsertionSortThreshold){
      insertion_sort(first, last, comp);
   }
   else{
      size_type const half = size_type(count/2);
      size_type const rest = size_type(count -  half);
      RandIt const half_it = first + half;
      RandIt const rest_it = first + rest;

      merge_sort_copy(half_it, last, buffer, comp);
      merge_sort_copy(first, half_it, rest_it, comp);
      merge_with_right_placed
         (buffer, buffer + rest
         , first, rest_it, last, antistable<Compare>(comp));
   }
}

template<typename RandIt, typename Pointer,
         typename Distance, typename Compare>
void stable_sort_ONlogN_recursive(RandIt first, RandIt last, Pointer buffer, Distance buffer_size, Compare comp)
{
   typedef typename iterator_traits<RandIt>::size_type  size_type;
   if (size_type(last - first) <= size_type(MergeSortInsertionSortThreshold)) {
      insertion_sort(first, last, comp);
   }
   else {
      const size_type len = size_type(last - first) / 2u;
      const RandIt middle = first + len;
      if (len > ((buffer_size+1)/2)){
         stable_sort_ONlogN_recursive(first, middle, buffer, buffer_size, comp);
         stable_sort_ONlogN_recursive(middle, last, buffer, buffer_size, comp);
      }
      else{
         merge_sort_with_constructed_buffer(first, middle, comp, buffer);
         merge_sort_with_constructed_buffer(middle, last, comp, buffer);
      }
      merge_adaptive_ONlogN_recursive(first, middle, last,
         size_type(middle - first),
         size_type(last - middle),
         buffer, buffer_size,
         comp);
   }
}

template<typename BidirectionalIterator, typename Compare, typename RandRawIt>
void stable_sort_adaptive_ONlogN2(BidirectionalIterator first,
		                           BidirectionalIterator last,
		                           Compare comp,
                                 RandRawIt uninitialized,
                                 std::size_t uninitialized_len)
{
   typedef typename iterator_traits<BidirectionalIterator>::value_type  value_type;

   ::boost::movelib::adaptive_xbuf<value_type, RandRawIt> xbuf(uninitialized, uninitialized_len);
   xbuf.initialize_until(uninitialized_len, *first);
   stable_sort_ONlogN_recursive(first, last, uninitialized, uninitialized_len, comp);
}

///@endcond

}} //namespace boost {  namespace movelib{

#if defined(BOOST_CLANG) || (defined(BOOST_GCC) && (BOOST_GCC >= 40600))
#pragma GCC diagnostic pop
#endif

#include <boost/move/detail/config_end.hpp>

#endif //#ifndef BOOST_MOVE_DETAIL_MERGE_SORT_HPP


#include "ai_base.hpp"

#include <boost/lexical_cast.hpp>

#include <cmath>
#include <fstream>
#include <iostream>
#include <random>
#include <vector>
#include <algorithm>

class my_ai
  : public aiwc::ai_base
{
  static constexpr double PI = 3.1415926535;

public:
  my_ai(std::string server_ip, std::size_t port, std::string realm, std::string key, std::string datapath)
    : aiwc::ai_base(std::move(server_ip), port, std::move(realm), std::move(key), std::move(datapath))
    , robot_wheels{}
  {
    std::cout << "I am ready." << std::endl;
  }

private:
  // this function is called at the beginning of a game
  void init()
  {
    // from here, you have access to game-specific constant information such as field dimensions
    // check example 'general_check-variables_cpp' to see what information are available

    // you can initialize some customvariables here
  }

  // copy coordinates from frames to different variables just for convenience
  auto get_coord(const aiwc::frame& f)
  {
    //Get data in the frame
    decltype(cur_posture) cur;
    decltype(cur_posture_op) cur_op;
      std::array<double, 2> pos_ball = { (*f.opt_coordinates).ball.x, (*f.opt_coordinates).ball.y };

    for(std::size_t id = 0; id < 5; ++id) {
      cur[id][X]  = (*f.opt_coordinates).robots[MYTEAM][id].x;
      cur[id][Y]  = (*f.opt_coordinates).robots[MYTEAM][id].y;
      cur[id][TH] = (*f.opt_coordinates).robots[MYTEAM][id].th;
      cur[id][ACTIVE] = (*f.opt_coordinates).robots[MYTEAM][id].active;
      cur[id][TOUCH] = (*f.opt_coordinates).robots[MYTEAM][id].touch;

      cur_op[id][X]  = (*f.opt_coordinates).robots[OPPONENT][id].x;
      cur_op[id][Y]  = (*f.opt_coordinates).robots[OPPONENT][id].y;
      cur_op[id][TH] = (*f.opt_coordinates).robots[OPPONENT][id].th;
      cur_op[id][ACTIVE] = (*f.opt_coordinates).robots[MYTEAM][id].active;
      cur_op[id][TOUCH] = (*f.opt_coordinates).robots[MYTEAM][id].touch;
    }

    return std::make_tuple(cur, cur_op, pos_ball);
  }

  // this function is called at each timestep. in 'f', current step's information are stored
  // check example 'general_check-variables_cpp' to see what information are available here
  // you should implement an AI soccer algorithm that sets robot wheel velocities for each timestep here
  void update(const aiwc::frame& f)
  {
	  frames.push_back(f);
    if(f.reset_reason == aiwc::GAME_START) {
      previous_frame = f;
      std::tie(cur_posture, cur_posture_op, cur_ball) = get_coord(f);
      return;
    }
    else if(f.reset_reason == aiwc::EPISODE_END) {
      // EPISODE_END is sent instead of GAME_END when 'repeat' option is set to 'true'
      // to mark the end of episode
      // you can reinitialize the parameters, count the number of episodes done, etc. here

      // this example does not do anything at episode end
    }
    else if(f.reset_reason == aiwc::HALFTIME) {
      // halftime is met - from next frame, received_frame.half_passed will be set to True
      // although the simulation switches sides,
      // coordinates and images given to your AI soccer algorithm will stay the same
      // that your team is red and located on left side whether it is 1st half or 2nd half

      // this example does not do anything at halftime
    }
    else if(f.reset_reason == aiwc::GAME_END) {
      return;
    }

    // update current robot and ball postures
    std::tie(cur_posture, cur_posture_op, cur_ball) = get_coord(f);
    find_closest_robot();

    // array for holding wheel speeds
    std::array<double, 10> ws;

    // act differently based on the current game state
    switch(f.game_state) {
		//////////////////////////////////////////////////////////////////////
      case aiwc::STATE_DEFAULT:
        {
          // robots functions in STATE_DEFAULT
          goalkeeper(0);
          defender(1);
          defender(2);
          forward(3);
          forward(4);
        }
        break;
		//////////////////////////////////////////////////////////////////////
      case aiwc::STATE_KICKOFF:
        {
          // if the ball belongs to my team, initiate kickoff
          if (f.ball_ownership)
            set_target_position(4, 0, 0, 1.4, 3.0, 0.4, false);
        }
        break;
		//////////////////////////////////////////////////////////////////////
      case aiwc::STATE_GOALKICK:
        {
          // if the ball belongs to my team,
          // drive the goalkeeper to kick the ball
          if (f.ball_ownership)
            set_wheel_velocity(0, info.max_linear_velocity[0], info.max_linear_velocity[0], true);
        }
        break;
		//////////////////////////////////////////////////////////////////////
      case aiwc::STATE_CORNERKICK:
        {
          // just play as simple as possible
          goalkeeper(0);
          defender(1);
          defender(2);
          forward(3);
          forward(4);
        }
        break;
		//////////////////////////////////////////////////////////////////////
      case aiwc::STATE_PENALTYKICK:
        {
          // if the ball belongs to my team,
          // drive the forward to kick the ball
          if (f.ball_ownership)
            set_wheel_velocity(4, info.max_linear_velocity[4], info.max_linear_velocity[4], true);
        }
        break;
      default:
        break;
    }

    // hand over current frame as previous frame to next step
    prev_ball = cur_ball;
    prev_posture = cur_posture;
    prev_posture_op = cur_posture_op;
    previous_frame = f;

    // collect the robot wheel speeds
    for(std::size_t id = 0; id < 5; ++id) {
      ws[2*id    ] = robot_wheels[id][0]; // left
      ws[2*id + 1] = robot_wheels[id][1]; // right
    }

    // send wheel speed data to the simulator
    set_wheel(ws); // this function is defined at ai_base.cpp(examples/common)
  }

  // this function is called when the game is over
  // you can save some data here if necessary
  void finish()
  {
    // you have less than 30 seconds before it's killed.
    std::ofstream ofs(datapath + "/result.txt");

    // write something
    ofs << "time, a.score, b.score, ";
    ofs << "ball.x, ball.y" << std::endl;

    // close the file
    ofs.close();
  }

  // All functions below are used for rulebased robot controls used in update()
  // convert degree to radian
  double d2r(double deg) {
    return deg * PI / 180;
  }

  // convert radian to degree
  double r2d(double rad) {
    return rad * 180 / PI;
  }

  // measure the distance between two coordinates (x1, y1) and (x2, y2)
  double dist(double x1, double x2, double y1, double y2)
  {
    const auto dx = x1 - x2;
    const auto dy = y1 - y2;
    return std::sqrt(dx * dx + dy * dy);
  }

  // convert radian in (-inf, inf) range to (-PI, PI) range
  double trim_radian(double rad)
  {
    double adj_rad = rad;
    while (adj_rad > PI)
      adj_rad -= 2 * PI;
    while (adj_rad < -PI)
      adj_rad += 2 * PI;
    return adj_rad;
  }

  // set the left and right wheel velocities of robot with id 'id'
  // 'max_velocity' scales the velocities up to the point where at least one of wheel is operating at max velocity
  void set_wheel_velocity(std::size_t id, double l, double r, bool max_velocity)
  {
    double multiplier = 1;

    // wheel velocities need to be scaled so that none of wheels exceed the maximum velocity available
    // otherwise, the velocity above the limit will be set to the max velocity by the simulation program
    // if that happens, the velocity ratio between left and right wheels will be changed that the robot may not execute
    // turning actions correctly.
    if (std::abs(l) > info.max_linear_velocity[id] || std::abs(r) > info.max_linear_velocity[id] || max_velocity) {
      if (std::abs(l) > std::abs(r)) {
        multiplier = info.max_linear_velocity[id] / std::abs(l);
      }
      else {
        multiplier = info.max_linear_velocity[id] / std::abs(r);
      }

      l *= multiplier;
      r *= multiplier;
    }

    robot_wheels[id] = {l, r};
  }

  // let the robot with id 'id' move to a target position (x, y)
  // the trajectory to reach the target position is determined by several different parameters
  void set_target_position(std::size_t id, double x, double y, double scale, double mult_lin, double mult_ang, bool max_velocity)
  {
    const double damping = 0.35;
    double ka = 0;
    int sign = 1;

    // calculate how far the target position is from the robot
    const double dx = x - cur_posture[id][X];
    const double dy = y - cur_posture[id][Y];
    const double d_e = std::sqrt(dx * dx + dy * dy);

    // calculate how much the direction is off
    const double desired_th = (dx == 0 && dy == 0) ? (PI / 2) : std::atan2(dy, dx);
    double d_th = desired_th - cur_posture[id][TH];
    while(d_th > PI) d_th -= 2 * PI;
    while(d_th < -PI) d_th += 2 * PI;

    // based on how far the target position is, set a parameter that
    // decides how much importance should be put into changing directions
    // farther the target is, less need to change directions fastly
    if(d_e > 1) {        ka = 17/90; }
    else if(d_e > 0.5) { ka = 19/90; }
    else if(d_e > 0.3) { ka = 21/90; }
    else if(d_e > 0.2) { ka = 23/90; }
    else               { ka = 25/90; }

    // if the target position is at rear of the robot, drive backward instead
    if(d_th > d2r(95)) {
      d_th -= PI;
      sign = -1;
    }
    else if(d_th < d2r(-95)) {
      d_th += PI;
      sign = -1;
    }

    // if the direction is off by more than 85 degrees,
    // make a turn first instead of start moving toward the target
    if(std::abs(d_th) > d2r(85)) {
      set_wheel_velocity(id, -mult_ang * d_th, mult_ang * d_th, false);
    }
    // otherwise
    else {
      // scale the angular velocity further down if the direction is off by less than 40 degrees
      if(d_e < 5.0 && std::abs(d_th) < d2r(40)) {
	      ka = 0.1;
      }
      ka *= 4;

      // set the wheel velocity
      // 'sign' determines the direction [forward, backward]
      // 'scale' scales the overall velocity at which the robot is driving
      // 'mult_lin' scales the linear velocity at which the robot is driving
      // larger distance 'd_e' scales the base linear velocity higher
      // 'damping' slows the linear velocity down
      // 'mult_ang' and 'ka' scales the angular velocity at which the robot is driving
      // larger angular difference 'd_th' scales the base angular velocity higher
      // if 'max_velocity' is true, the overall velocity is scaled to the point
      // where at least one wheel is operating at maximum velocity
      set_wheel_velocity(id,
	       sign * scale * (mult_lin * (1 / (1 + std::exp(-3 * d_e)) - damping) - mult_ang * ka * d_th),
	       sign * scale * (mult_lin * (1 / (1 + std::exp(-3 * d_e)) - damping) + mult_ang * ka * d_th), max_velocity);
    }
  }

  // predict where the ball will be located after 'steps' steps
  std::array<double, 2> predict_ball_location(int steps)
  {
    double dx = cur_ball[X] - prev_ball[X];
    double dy = cur_ball[Y] - prev_ball[Y];
    return {cur_ball[X] + steps * dx, cur_ball[Y] + steps * dy};
  }

  // check if a certain position is inside the penalty area of 'team'
  bool in_penalty_area(std::array<double, 2> obj, std::size_t team)
  {
    if (std::abs(obj[Y]) > info.penalty_area[Y] / 2)
      return false;

    if (team == MYTEAM)
      return (obj[X] < -info.field[X] / 2 + info.penalty_area[X]);
    else
      return (obj[X] > info.field[X] / 2 - info.penalty_area[X]);
  }

  // check if the robot with id 'id' has a chance to shoot
  bool shoot_chance(size_t id)
  {
    double dx = cur_ball[X] - cur_posture[id][X];
    double dy = cur_ball[Y] - cur_posture[id][Y];

    // if the ball is located further on left than the robot, it will be hard to shoot
  	if (dx < 0) {
  		return false;
  	}

    // if the robot->ball direction aligns with opponent's goal, the robot can shoot
    double y = (info.field[X] / 2 - cur_ball[X]) * dy / dx + cur_posture[id][Y];
  	if (abs(y) < info.goal[Y] / 2) {
  		return true;
  	}
  	else {
  		return false;
  	}
  }

  // find a defender and a forward closest to the ball
  void find_closest_robot(void)
  {
    // find the closest defender
    int min_idx = 0;
    double min_dist = 9999.99;
    double def_min_dist = 9999.99;

    for (int i = 1; i < 3; i++) {
      auto measured_dist = dist(cur_ball[X], cur_posture[i][X], cur_ball[Y], cur_posture[i][Y]);
      if (measured_dist < min_dist) {
        min_dist = measured_dist;
        min_idx = i;
        def_idx = min_idx;
        def_min_dist = min_dist;
      }
    }

    // find the closest forward
    min_idx = 0;
    min_dist = 9999.99;
    double fwd_min_dist = 9999.99;

    for (int i = 3; i < 5; i++) {
      auto measured_dist = dist(cur_ball[X], cur_posture[i][X], cur_ball[Y], cur_posture[i][Y]);
      if (measured_dist < min_dist) {
        min_dist = measured_dist;
        min_idx = i;
        fwd_idx = min_idx;
        fwd_min_dist = min_dist;
      }
    }
  }

  // turn to face 'desired_th' direction
  void angle(std::size_t id, double desired_th)
  {
    double mult_ang = 0.4;
    double sign = 1;

    double d_th = desired_th - cur_posture[id][TH];
    d_th = trim_radian(d_th);

    // the robot instead puts the direction rear if the angle difference is large
    if (d_th > d2r(95)) {
        d_th -= PI;
        sign = -1;
    }
    else if (d_th < d2r(-95)) {
        d_th += PI;
        sign = -1;
    }
    set_wheel_velocity(id, -mult_ang * d_th, mult_ang * d_th, false);
  }

  // returns the angle toward a specific position from current robot posture
  double direction_angle(std::size_t id, double x, double y)
  {
    double dx = x - cur_posture[id][X];
    double dy = y - cur_posture[id][Y];

    if ((dx == 0) && (dy == 0))
      return (PI / 2);
    else
      return std::atan2(dy, dx);
  }

  // let the robot face toward specific direction
  void face_specific_position(std::size_t id, double x, double y)
  {
    double dx = x - cur_posture[id][X];
    double dy = y - cur_posture[id][Y];

    double desired_th;
    if ((dx == 0) && (dy == 0))
      desired_th = PI / 2;
    else
      desired_th = std::atan2(dy, dx);

    angle(id, desired_th);
  }

  // check if the ball is coming toward the robot
  bool ball_coming_toward_robot(std::size_t id)
  {
    double x_dir = std::abs(cur_posture[id][X] - prev_ball[X]) > std::abs(cur_posture[id][X] - cur_ball[X]);
    double y_dir = std::abs(cur_posture[id][Y] - prev_ball[Y]) > std::abs(cur_posture[id][Y] - cur_ball[Y]);

    // ball is coming closer
    if ((x_dir) && (y_dir))
      return true;
    else
      return false;
  }

  // a basic goalkeeper rulebased algorithm
  void goalkeeper(std::size_t id)
  {
    // default desired position
    double x = (-info.field[X] / 2) + (info.robot_size[id] / 2) + 0.05;
    double y = std::max(std::min(cur_ball[Y], (info.goal[Y] / 2 - info.robot_size[id] / 2)), -info.goal[Y] / 2 + info.robot_size[id] / 2);
    std::array<double, 2> a = {cur_posture[id][X], cur_posture[id][Y]};

    // if the robot is inside the goal, try to get out
    if (cur_posture[id][X] < -info.field[X] / 2){
      if (cur_posture[id][Y] < 0)
        set_target_position(id, x, cur_posture[id][Y] + 0.2, 1.4, 5.0, 0.4, false);
      else
        set_target_position(id, x, cur_posture[id][Y] - 0.2, 1.4, 5.0, 0.4, false);
    }
    // if the goalkeeper is outside the penalty area
    else if (!in_penalty_area(a, MYTEAM)){
      // return to the desired position
      set_target_position(id, x, y, 1.4, 5.0, 0.4, true);
    }
    // if the goalkeeper is inside the penalty area
    else{
      // if the ball is inside the penalty area
      if (in_penalty_area(cur_ball, MYTEAM)){
        // if the ball is behind the goalkeeper
        if (cur_ball[X] < cur_posture[id][X]){
          // if the ball is not blocking the goalkeeper's path
          if (std::abs(cur_ball[Y] - cur_posture[id][Y]) > 2 * info.robot_size[id]){
            // try to get ahead of the ball
            set_target_position(id, cur_ball[X] - info.robot_size[id], cur_posture[id][Y], 1.4, 5.0, 0.4, false);
          }
          else{
            // just give up and try not to make a suicidal goal
            angle(id, PI / 2);
          }
        }
        // if the ball is ahead of the goalkeeper
        else{
          double desired_th = direction_angle(id, cur_ball[X], cur_ball[Y]);
          double rad_diff = trim_radian(desired_th - cur_posture[id][TH]);
          // if the robot direction is too away from the ball direction
          if (rad_diff > PI / 3){
            // give up kicking the ball and block the goalpost
            set_target_position(id, x, y, 1.4, 5.0, 0.4, false);
          }
          else{
            // try to kick the ball away from the goal
            set_target_position(id, cur_ball[X], cur_ball[Y], 1.4, 3.0, 0.8, true);
          }
        }
      }
      // if the ball is not in the penalty area
      else{
        // if the ball is within alert range and y position is not too different
        if ((cur_ball[X] < -info.field[X] / 2 + 1.5 * info.penalty_area[X])
        && (std::abs(cur_ball[Y]) < 1.5 * info.penalty_area[Y] / 2)
        && (std::abs(cur_ball[Y] - cur_posture[id][Y]) < 0.2))
          face_specific_position(id, cur_ball[X], cur_ball[Y]);
        // otherwise
        else
          set_target_position(id, x, y, 1.4, 5.0, 0.4, true);
      }
    }
  }

  // a basic defender rulebased algorithm
  void defender(std::size_t id)
  {
    // if the robot is inside the goal, try to get out
    if (cur_posture[id][X] < -info.field[X] / 2){
      if (cur_posture[id][Y] < 0)
        set_target_position(id, -0.7 * info.field[X] / 2, cur_posture[id][Y] + 0.2, 1.4, 3.5, 0.6, false);
      else
        set_target_position(id, -0.7 * info.field[X] / 2, cur_posture[id][Y] - 0.2, 1.4, 3.5, 0.6, false);
      return;
    }
    // the defender may try to shoot if condition meets
    if ((id == def_idx) && (shoot_chance(id)) && (cur_ball[X] < 0.3 * info.field[X] / 2)){
      set_target_position(id, cur_ball[X], cur_ball[Y], 1.4, 5.0, 0.4, true);
      return;
    }
    // if this defender is closer to the ball than the other defender
    if (id == def_idx){
      // ball is on our side
      if (cur_ball[X] < 0){
        // if the robot can push the ball toward opponent's side, do it
        if (cur_posture[id][X] < cur_ball[X] - info.ball_radius)
          set_target_position(id, cur_ball[X], cur_ball[Y], 1.4, 5.0, 0.4, true);
        else{
          // otherwise go behind the ball
          if (std::abs(cur_ball[Y] - cur_posture[id][Y]) > 0.3)
            set_target_position(id, std::max(cur_ball[X] - 0.5, -info.field[X] / 2 + info.robot_size[id] / 2), cur_ball[Y], 1.4, 3.5, 0.6, false);
          else
            set_target_position(id, std::max(cur_ball[X] - 0.5, -info.field[X] / 2 + info.robot_size[id] / 2), cur_posture[id][Y], 1.4, 3.5, 0.6, false);
        }
      }
      else
        set_target_position(id, -0.7 * info.field[X] / 2, cur_ball[Y], 1.4, 3.5, 0.4, false);
    }
    // if this defender is not closer to the ball than the other defender
    else{
      // ball is on our side
      if (cur_ball[X] < 0){
        // ball is on our left
        if (cur_ball[Y] > info.goal[Y] / 2 + 0.15)
          set_target_position(id, std::max(cur_ball[X] - 0.5, -info.field[X] / 2 + info.robot_size[id] / 2 + 0.1), info.goal[Y] / 2 + 0.15, 1.4, 3.5, 0.4, false);
        // ball is on our right
        else if (cur_ball[Y] < -info.goal[Y] / 2 - 0.15)
          set_target_position(id, std::max(cur_ball[X] - 0.5, -info.field[X] / 2 + info.robot_size[id] / 2 + 0.1), -info.goal[Y] / 2 - 0.15, 1.4, 3.5, 0.4, false);
        // ball is in center
        else
          set_target_position(id, std::max(cur_ball[X] - 0.5, -info.field[X] / 2 + info.robot_size[id] / 2 + 0.1), cur_ball[Y], 1.4, 3.5, 0.4, false);
      }
      else{
        // ball is on right side
        if (cur_ball[Y] < 0)
          set_target_position(id, -0.7 * info.field[X] / 2, std::min(cur_ball[Y] + 0.5, info.field[Y] / 2 - info.robot_size[id] / 2), 1.4, 3.5, 0.4, false);
        // ball is on left side
        else
          set_target_position(id, -0.7 * info.field[X] / 2, std::max(cur_ball[Y] - 0.5, -info.field[Y] / 2 + info.robot_size[id] / 2), 1.4, 3.5, 0.4, false);
      }
    }
  }

  // a basic forward rulebased algorithm
  void forward(std::size_t id)
  {
    // if the robot is blocking the ball's path toward opponent side
    if ((cur_ball[X] > -0.3 * info.field[X] / 2) && (cur_ball[X] < 0.3 * info.field[X] / 2) && (cur_posture[id][X] > cur_ball[X] + 0.1) && (std::abs(cur_posture[id][Y] - cur_ball[Y]) < 0.3)){
      if (cur_ball[Y] < 0)
        set_target_position(id, cur_posture[id][X] - 0.25, cur_ball[Y] + 0.75, 1.4, 3.0, 0.8, false);
      else
        set_target_position(id, cur_posture[id][X] - 0.25, cur_ball[Y] - 0.75, 1.4, 3.0, 0.8, false);
      return;
    }

    // if the robot can shoot from current position
    if ((id == fwd_idx) && (shoot_chance(id))){
      std::array<double, 2> pred_ball = predict_ball_location(2);
      set_target_position(id, pred_ball[X], pred_ball[Y], 1.4, 5.0, 0.4, true);
      return;
    }

    // if the ball is coming toward the robot, seek for shoot chance
    if ((id == fwd_idx) && (ball_coming_toward_robot(id))){
      double dx = cur_ball[X] - prev_ball[X];
      double dy = cur_ball[Y] - prev_ball[Y];
      double pred_x = (cur_posture[id][Y] - cur_ball[Y]) * dx / dy + cur_ball[X];
      double steps = (cur_posture[id][Y] - cur_ball[Y]) / dy;

      // if the ball will be located in front of the robot
      if (pred_x > cur_posture[id][X]){
        double pred_dist = pred_x - cur_posture[id][X];
        // if the predicted ball location is close enough
        if ((pred_dist > 0.1) && (pred_dist < 0.3) && (steps < 10)){
          // find the direction towards the opponent goal and look toward it
          double goal_angle = direction_angle(id, info.field[X] / 2, 0);
          angle(id, goal_angle);
          return;
        }
      }
    }

    // if this forward is closer to the ball than the other forward
    if (id == fwd_idx){
      if (cur_ball[X] > -0.3 * info.field[X] / 2){
        // if the robot can push the ball toward opponent's side, do it
        if (cur_posture[id][X] < cur_ball[X] - info.ball_radius)
          set_target_position(id, cur_ball[X], cur_ball[Y], 1.4, 5.0, 0.4, true);
        else{
          // otherwise go behind th ball
          if (std::abs(cur_ball[Y] - cur_posture[id][Y]) > 0.3)
            set_target_position(id, cur_ball[X] - 0.2, cur_ball[Y], 1.4, 3.5, 0.6, false);
          else
            set_target_position(id, cur_ball[X] - 0.2, cur_posture[id][Y], 1.4, 3.5, 0.6, false);
        }
      }
      else
        set_target_position(id, -0.1 * info.field[X] / 2, cur_ball[Y], 1.4, 3.5, 0.4, false);
    }
    // if this forward is closer to the ball than the other forward
    else{
      if (cur_ball[X] > -0.3 * info.field[X] / 2){
        // ball is on our right
        if (cur_ball[Y] < 0)
          set_target_position(id, cur_ball[X] - 0.25, info.goal[Y] / 2, 1.4, 3.5, 0.4, false);
        // ball is on our left
        else
          set_target_position(id, cur_ball[X] - 0.25, -info.goal[Y] / 2, 1.4, 3.5, 0.4, false);
      }
      else{
        // ball is on right side
        if (cur_ball[Y] < 0)
          set_target_position(id, -0.1 * info.field[X] / 2, std::min(-cur_ball[Y] - 0.5, info.field[Y] / 2 - info.robot_size[id] / 2), 1.4, 3.5, 0.4, false);
        // ball is on left side
        else
          set_target_position(id, -0.1 * info.field[X] / 2, std::max(-cur_ball[Y] + 0.5, -info.field[Y] / 2 + info.robot_size[id] / 2), 1.4, 3.5, 0.4, false);
      }
    }
  }

private: // member variable
  aiwc::frame previous_frame;

  std::array<std::array<double, 5>, 5> cur_posture;	//X,Y,THETA,ACTIVE,TOUCH
  std::array<std::array<double, 5>, 5> cur_posture_op;
  std::array<std::array<double, 5>, 5> prev_posture;
  std::array<std::array<double, 5>, 5> prev_posture_op;
  std::array<double, 2> prev_ball;
  std::array<double, 2> cur_ball;
  double time;

  std::array<std::array<double, 2>, 5> robot_wheels;

  std::array<bool, 5> touch;
  size_t def_idx;
  size_t fwd_idx;

  std::vector<aiwc::frame> frames;
};

int main(int argc, char *argv[])
{
  if(argc < 6) {
    std::cerr << "Usage " << argv[0] << " server_ip port realm key datapath" << std::endl;
    return -1;
  }

  const auto& server_ip = std::string(argv[1]);
  const auto& port      = boost::lexical_cast<std::size_t>(argv[2]);
  const auto& realm     = std::string(argv[3]);
  const auto& key       = std::string(argv[4]);
  const auto& datapath  = std::string(argv[5]);

  my_ai ai(server_ip, port, realm, key, datapath);

  ai.run();

  return 0;
}


#pragma once

#include <Eigen/Core>

namespace ipc {

bool inexact_point_edge_ccd_2D(
    const Eigen::Vector2d& p_t0,
    const Eigen::Vector2d& e0_t0,
    const Eigen::Vector2d& e1_t0,
    const Eigen::Vector2d& p_t1,
    const Eigen::Vector2d& e0_t1,
    const Eigen::Vector2d& e1_t1,
    double& toi,
    const double conservative_rescaling);

} // namespace ipc


#define BOOST_TEST_MODULE "test_axis_aligned_plane_shape"

#ifdef BOOST_TEST_DYN_LINK
#include <boost/test/unit_test.hpp>
#else
#include <boost/test/included/unit_test.hpp>
#endif

#include <mjolnir/core/SimulatorTraits.hpp>
#include <mjolnir/core/BoundaryCondition.hpp>
#include <mjolnir/core/System.hpp>
#include <mjolnir/core/AxisAlignedPlane.hpp>
#include <mjolnir/util/make_unique.hpp>
#include <random>
#include <cmath>

BOOST_AUTO_TEST_CASE(AxisAlignedPlane_PositiveZ_geometry_unlimited)
{
    mjolnir::LoggerManager::set_default_logger("test_AxisAlignedPlane");
    using traits_type   = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type     = typename traits_type::real_type;
    using coord_type    = typename traits_type::coordinate_type;
    using boundary_type = typename traits_type::boundary_type;

    constexpr real_type tol = 1e-8;

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<double> uni(-100, 100);

    boundary_type bdry;
    mjolnir::AxisAlignedPlane<traits_type, mjolnir::PositiveZDirection<traits_type>> xyplane(0.0);

    for(std::size_t i=0; i<100000; ++i)
    {
        const coord_type pos(uni(mt), uni(mt), uni(mt));
        const real_type  dist = xyplane.calc_distance(pos, bdry);
        const coord_type fdir = xyplane.calc_force_direction(pos, bdry);

        BOOST_TEST(fdir[0] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(fdir[1] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(fdir[2] == 1.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(dist == pos[2], boost::test_tools::tolerance(tol));
    }
}

BOOST_AUTO_TEST_CASE(AxisAlignedPlane_NegativeZ_geometry_unlimited)
{
    mjolnir::LoggerManager::set_default_logger("test_AxisAlignedPlane");
    using traits_type   = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using real_type     = typename traits_type::real_type;
    using coord_type    = typename traits_type::coordinate_type;
    using boundary_type = typename traits_type::boundary_type;

    constexpr real_type tol = 1e-8;

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<double> uni(-100, 100);

    boundary_type bdry;
    mjolnir::AxisAlignedPlane<traits_type, mjolnir::NegativeZDirection<traits_type>> xyplane(0.0);

    for(std::size_t i=0; i<100000; ++i)
    {
        const coord_type pos(uni(mt), uni(mt), uni(mt));
        const real_type  dist = xyplane.calc_distance(pos, bdry);
        const coord_type fdir = xyplane.calc_force_direction(pos, bdry);

        BOOST_TEST(fdir[0] ==  0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(fdir[1] ==  0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(fdir[2] == -1.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(dist == -pos[2], boost::test_tools::tolerance(tol));
    }
}

BOOST_AUTO_TEST_CASE(AxisAlignedPlane_geometry_periodic)
{
    mjolnir::LoggerManager::set_default_logger("test_AxisAlignedPlane");
    using traits_type   = mjolnir::SimulatorTraits<double, mjolnir::CuboidalPeriodicBoundary>;
    using real_type     = typename traits_type::real_type;
    using coord_type    = typename traits_type::coordinate_type;
    using boundary_type = typename traits_type::boundary_type;

    constexpr real_type tol = 1e-8;

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<double> uni(0, 100);

    const coord_type lower(0, 0, 0), upper(100, 100, 100);
    boundary_type bdry(lower, upper);
    mjolnir::AxisAlignedPlane<traits_type, mjolnir::PositiveZDirection<traits_type>> xyplane(0.0);

    for(std::size_t i=0; i<100000; ++i)
    {
        const coord_type pos(uni(mt), uni(mt), uni(mt));
        const real_type  dist = xyplane.calc_distance(pos, bdry);
        const coord_type fdir = xyplane.calc_force_direction(pos, bdry);

        const real_type d    = (pos[2] <= 50.0) ? pos[2] : pos[2] - 100.0;

        BOOST_TEST(fdir[0] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(fdir[1] == 0.0, boost::test_tools::tolerance(tol));
        BOOST_TEST(fdir[2] == 1.0, boost::test_tools::tolerance(tol));

        if(std::abs(pos[2] - 50.0) < tol)
        {
            BOOST_TEST(std::abs(dist) == 50.0, boost::test_tools::tolerance(tol));
        }
        else
        {
            BOOST_TEST(dist == d, boost::test_tools::tolerance(tol));
        }
    }
}

template<typename realT>
struct dummy_potential
{
    realT cutoff;
    std::size_t N;

    std::vector<std::size_t> participants() const
    {
        std::vector<std::size_t> retval(N);
        std::iota(retval.begin(), retval.end(), 0);
        return retval;
    }

    realT max_cutoff_length() const noexcept {return cutoff;}

};

BOOST_AUTO_TEST_CASE(AxisAlignedPlane_neighbors_unlimited)
{
    mjolnir::LoggerManager::set_default_logger("test_AxisAlignedPlane");
    using traits_type   = mjolnir::SimulatorTraits<double, mjolnir::UnlimitedBoundary>;
    using coord_type    = typename traits_type::coordinate_type;
    using boundary_type = typename traits_type::boundary_type;
    using system_type   = mjolnir::System<traits_type>;

    std::mt19937 mt(123456789);
    std::uniform_real_distribution<double> uni(-10, 10);

    system_type sys(1000, boundary_type{});
    for(std::size_t i=0; i<sys.size(); ++i)
    {
        sys.mass(i)     = 0.0;
        sys.position(i) = coord_type(uni(mt), uni(mt), uni(mt));
        sys.velocity(i) = coord_type(0, 0, 0);
        sys.force(i)    = coord_type(0, 0, 0);
    }

    dummy_potential<double> dummy; dummy.cutoff = 1.0; dummy.N = 1000;
    mjolnir::AxisAlignedPlane<traits_type, mjolnir::PositiveZDirection<traits_type>>
        xyplane(0.0, 0.0); // no mergin here
    xyplane.initialize(sys, dummy);

    const auto neighbors = xyplane.neighbors();

    for(std::size_t i : neighbors)
    {
        if(std::abs(sys.position(i)[2]) > 1.0)
        {
            std::cerr << i << ", " << sys.position(i)[2] << std::endl;
        }
        BOOST_TEST(std::abs(sys.position(i)[2]) <= 1.0);
    }

    for(std::size_t i=0; i<sys.size(); ++i)
    {
        if(std::abs(sys.position(i)[2]) > 1.0){continue;}
        const bool found =
            std::find(neighbors.begin(), neighbors.end(), i) != neighbors.end();
        BOOST_TEST(found);
    }
}


//
//  Copyright (C) 2016 Greg Landrum
//
//   @@ All Rights Reserved @@
//  This file is part of the RDKit.
//  The contents are covered by the terms of the BSD license
//  which is included in the file license.txt, found at the root
//  of the RDKit source tree.
//

#ifdef _MSC_VER
// disable warnings about getenv in visual C++
#define _CRT_SECURE_NO_WARNINGS
#endif

#include <iostream>
#include <fstream>

#include <RDGeneral/BoostStartInclude.h>
#include <boost/algorithm/string.hpp>
#include <boost/algorithm/string/trim.hpp>
#include <boost/lexical_cast.hpp>
#include <RDGeneral/BoostEndInclude.h>

#include <RDGeneral/Invariant.h>
#include <RDGeneral/RDLog.h>
#include <RDGeneral/utils.h>
#include <RDGeneral/StreamOps.h>

#include <GraphMol/RDKitBase.h>
#include <GraphMol/SmilesParse/SmilesParse.h>
#include <GraphMol/FileParsers/FileParsers.h>
#include <GraphMol/FileParsers/MolSupplier.h>

#include <GraphMol/Descriptors/MolDescriptors3D.h>

using namespace RDKit;
using namespace RDKit::Descriptors;

bool compare(const std::string &inm, double ref, double val,
             double tol = 1e-3) {
  if (fabs(ref - val) > tol) {
    std::cerr << "value mismatch: " << inm << " " << ref << " " << val
              << std::endl;
  }
  return fabs(ref - val) < tol;
}

void testPMI1() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    Basic PMI tests." << std::endl;

  std::string pathName = getenv("RDBASE");
  std::string sdfName =
      pathName + "/Code/GraphMol/Descriptors/test_data/PBF_egfr.sdf";

  RDKit::SDMolSupplier reader(sdfName, true, false);
  std::string fName =
      pathName + "/Code/GraphMol/Descriptors/test_data/PMI_egfr.out";
  std::ifstream instrm(fName.c_str());
  int nDone = 0;
  while (!reader.atEnd()) {
    RDKit::ROMol *m = reader.next();
    TEST_ASSERT(m);
    RDKit::ROMol mcpy(*m);
    std::string nm;
    m->getProp("_Name", nm);
    std::string inm;
    instrm >> inm;
    TEST_ASSERT(inm == nm);
    double val;
    double pmi1_m, pmi2_m, pmi3_m, pmi1_nom, pmi2_nom, pmi3_nom;
    instrm >> pmi1_m;
    instrm >> pmi2_m;
    instrm >> pmi3_m;
    instrm >> pmi1_nom;
    instrm >> pmi2_nom;
    instrm >> pmi3_nom;

    val = RDKit::Descriptors::PMI1(*m);
    TEST_ASSERT(compare(inm, pmi1_m, val));
    val = RDKit::Descriptors::PMI2(*m);
    TEST_ASSERT(compare(inm, pmi2_m, val));
    val = RDKit::Descriptors::PMI3(*m);
    TEST_ASSERT(compare(inm, pmi3_m, val));

    val = RDKit::Descriptors::PMI1(*m, -1, false);
    TEST_ASSERT(compare(inm, pmi1_nom, val));
    val = RDKit::Descriptors::PMI2(*m, -1, false);
    TEST_ASSERT(compare(inm, pmi2_nom, val));
    val = RDKit::Descriptors::PMI3(*m, -1, false);
    TEST_ASSERT(compare(inm, pmi3_nom, val));

    // now try doing it in the reverse order to make sure caching doesn't
    // screw up.
    val = RDKit::Descriptors::PMI1(mcpy, -1, false);
    TEST_ASSERT(compare(inm, pmi1_nom, val));
    val = RDKit::Descriptors::PMI2(mcpy, -1, false);
    TEST_ASSERT(compare(inm, pmi2_nom, val));
    val = RDKit::Descriptors::PMI3(mcpy, -1, false);
    TEST_ASSERT(compare(inm, pmi3_nom, val));
    val = RDKit::Descriptors::PMI1(mcpy);
    TEST_ASSERT(compare(inm, pmi1_m, val));
    val = RDKit::Descriptors::PMI2(mcpy);
    TEST_ASSERT(compare(inm, pmi2_m, val));
    val = RDKit::Descriptors::PMI3(mcpy);
    TEST_ASSERT(compare(inm, pmi3_m, val));

    delete m;
    ++nDone;
  }
  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

void testPMIEdges() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    PMI edge cases." << std::endl;

  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/linear.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::PMI1(*m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::PMI2(*m);
    TEST_ASSERT(fabs(val) >= 10);
    val = RDKit::Descriptors::PMI3(*m);
    TEST_ASSERT(val >= 10);
    TEST_ASSERT(RDKit::Descriptors::PMI3(*m) - RDKit::Descriptors::PMI2(*m) <
                1e-2);

    delete m;
  }
  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/linear_2atom.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::PMI1(*m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::PMI2(*m);
    TEST_ASSERT(fabs(val) >= 1);
    val = RDKit::Descriptors::PMI3(*m);
    TEST_ASSERT(val >= 1);
    TEST_ASSERT(RDKit::Descriptors::PMI3(*m) - RDKit::Descriptors::PMI2(*m) <
                1e-2);

    delete m;
  }
  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/planar.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);

    TEST_ASSERT(RDKit::Descriptors::PMI2(*m) - RDKit::Descriptors::PMI1(*m) <
                1e-2);
    TEST_ASSERT(RDKit::Descriptors::PMI3(*m) - RDKit::Descriptors::PMI1(*m) >
                10);
    delete m;
  }
  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/planar_3atom.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);

    TEST_ASSERT(RDKit::Descriptors::PMI2(*m) - RDKit::Descriptors::PMI1(*m) <
                1e-2);
    TEST_ASSERT(RDKit::Descriptors::PMI3(*m) - RDKit::Descriptors::PMI1(*m) >
                1);
    delete m;
  }

  {
    RDKit::RWMol m;
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addConformer(new RDKit::Conformer(m.getNumAtoms()));
    double val = RDKit::Descriptors::PMI1(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::PMI2(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::PMI3(m);
    TEST_ASSERT(fabs(val) < 1e-4);
  }

  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

void testPMI2() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    More PMI/NPR tests." << std::endl;

  std::string pathName = getenv("RDBASE");
  std::string sdfName =
      pathName + "/Code/GraphMol/Descriptors/test_data/pmi.sdf";

  RDKit::SDMolSupplier reader(sdfName, true, false);
  while (!reader.atEnd()) {
    RDKit::ROMol *mnoh = reader.next();
    TEST_ASSERT(mnoh);
    bool explicitOnly = false, addCoords = true;
    RDKit::ROMol *m = MolOps::addHs(*mnoh, explicitOnly, addCoords);
    delete mnoh;
    double pmi1 = RDKit::Descriptors::PMI1(*m);
    double pmi2 = RDKit::Descriptors::PMI2(*m);
    double pmi3 = RDKit::Descriptors::PMI3(*m);

    double npr1 = RDKit::Descriptors::NPR1(*m);
    double npr2 = RDKit::Descriptors::NPR2(*m);

    // tolerances are coarse because the reference values come from MOE
    // and the placement of Hs is not identical
    TEST_ASSERT(compare("pmi1", m->getProp<double>("pmi1"), pmi1, pmi1 / 100));
    TEST_ASSERT(compare("pmi2", m->getProp<double>("pmi2"), pmi2, pmi2 / 100));
    TEST_ASSERT(compare("pmi3", m->getProp<double>("pmi3"), pmi3, pmi3 / 100));

    TEST_ASSERT(compare("npr1", m->getProp<double>("npr1"), npr1, npr1 / 100));
    TEST_ASSERT(compare("npr2", m->getProp<double>("npr2"), npr2, npr2 / 100));
    delete m;
  }
  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

void testNPR1() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    Basic NPR tests." << std::endl;

  std::string pathName = getenv("RDBASE");
  std::string sdfName =
      pathName + "/Code/GraphMol/Descriptors/test_data/PBF_egfr.sdf";
  RDKit::SDMolSupplier reader(sdfName, true, false);

  int nDone = 0;
  while (!reader.atEnd()) {
    RDKit::ROMol *m = reader.next();
    TEST_ASSERT(m);
    RDKit::ROMol mcpy(*m);
    std::string nm;
    m->getProp("_Name", nm);

    double val;
    double pmi1_m, pmi2_m, pmi3_m, pmi1_nom, pmi2_nom, pmi3_nom;
    pmi1_m = RDKit::Descriptors::PMI1(*m);
    pmi2_m = RDKit::Descriptors::PMI2(*m);
    pmi3_m = RDKit::Descriptors::PMI3(*m);
    pmi1_nom = RDKit::Descriptors::PMI1(*m, -1, false);
    pmi2_nom = RDKit::Descriptors::PMI2(*m, -1, false);
    pmi3_nom = RDKit::Descriptors::PMI3(*m, -1, false);

    val = RDKit::Descriptors::NPR1(*m);
    compare(nm, pmi1_m / pmi3_m, val);
    val = RDKit::Descriptors::NPR2(*m);
    compare(nm, pmi2_m / pmi3_m, val);

    val = RDKit::Descriptors::NPR1(*m, -1, false);
    compare(nm, pmi1_nom / pmi3_nom, val);
    val = RDKit::Descriptors::NPR2(*m, -1, false);
    compare(nm, pmi2_nom / pmi3_nom, val);

    delete m;
    ++nDone;
  }
  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

void testNPREdges() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    NPR edge cases." << std::endl;

  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/linear.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::NPR1(*m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::NPR2(*m);
    TEST_ASSERT(fabs(val - 1) < 1e-4);

    delete m;
  }
  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/linear_2atom.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::NPR1(*m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::NPR2(*m);
    TEST_ASSERT(fabs(val - 1) < 1e-4);

    delete m;
  }
  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/planar.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::NPR1(*m);
    TEST_ASSERT(fabs(val - 0.5) < 1e-4);
    val = RDKit::Descriptors::NPR2(*m);
    TEST_ASSERT(fabs(val - 0.5) < 1e-4);

    delete m;
  }
  {
    std::string pathName = getenv("RDBASE");
    std::string sdfName =
        pathName + "/Code/GraphMol/Descriptors/test_data/planar_3atom.mol";

    RDKit::ROMol *m = MolFileToMol(sdfName);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::NPR1(*m);
    TEST_ASSERT(fabs(val - 0.5) < 1e-4);
    val = RDKit::Descriptors::NPR2(*m);
    TEST_ASSERT(fabs(val - 0.5) < 1e-4);

    delete m;
  }
  {
    RDKit::RWMol m;
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addConformer(new RDKit::Conformer(m.getNumAtoms()));
    double val = RDKit::Descriptors::NPR1(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::NPR2(m);
    TEST_ASSERT(fabs(val) < 1e-4);
  }

  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

void test3DVals() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    3D descriptors." << std::endl;
  std::string rdbase = getenv("RDBASE");
  {  // a disc (benzene)
    std::string fName =
        rdbase + "/Code/GraphMol/MolTransforms/test_data/github1262_1.mol";
    RWMol *m = MolFileToMol(fName, true, false);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::radiusOfGyration(*m);
    TEST_ASSERT(fabs(val - 1.511) < 1e-2);
    val = RDKit::Descriptors::eccentricity(*m);
    TEST_ASSERT(fabs(val - 0.866) < 1e-2);
    val = RDKit::Descriptors::asphericity(*m);
    TEST_ASSERT(fabs(val - 0.25) < 1e-2);
    val = RDKit::Descriptors::spherocityIndex(*m);
    TEST_ASSERT(fabs(val) < 1e-2);

    delete m;
  }
  {  // a rod (dimethyl acetylene)
    std::string fName =
        rdbase + "/Code/GraphMol/MolTransforms/test_data/github1262_2.mol";
    RWMol *m = MolFileToMol(fName, true, false);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::radiusOfGyration(*m);
    TEST_ASSERT(fabs(val - 1.686) < 1e-2);
    val = RDKit::Descriptors::eccentricity(*m);
    TEST_ASSERT(fabs(val - 1.0) < 1e-2);
    val = RDKit::Descriptors::asphericity(*m);
    TEST_ASSERT(fabs(val - 0.875) < 1e-2);
    val = RDKit::Descriptors::spherocityIndex(*m);
    // nothing really precise to say here
    TEST_ASSERT((0 < val) && (val < 0.25));

    delete m;
  }
  {  // adamantane
    std::string fName =
        rdbase + "/Code/GraphMol/MolTransforms/test_data/github1262_3.mol";
    RWMol *m = MolFileToMol(fName, true, false);
    TEST_ASSERT(m);
    double val;

    val = RDKit::Descriptors::radiusOfGyration(*m);
    TEST_ASSERT(fabs(val - 1.827) < 1e-2);
    val = RDKit::Descriptors::eccentricity(*m);
    TEST_ASSERT(fabs(val) < 1e-2);
    val = RDKit::Descriptors::asphericity(*m);
    TEST_ASSERT(fabs(val) < 1e-2);
    val = RDKit::Descriptors::spherocityIndex(*m);
    TEST_ASSERT(fabs(val - 1) < 1e-2);

    delete m;
  }

  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

void test3DEdges() {
  BOOST_LOG(rdErrorLog) << "-------------------------------------" << std::endl;
  BOOST_LOG(rdErrorLog) << "    3D descriptor edge cases." << std::endl;
  {  // octahedron
    RDKit::RWMol m;
    m.addAtom(new RDKit::Atom(1));
    m.addAtom(new RDKit::Atom(1));
    m.addAtom(new RDKit::Atom(1));
    m.addAtom(new RDKit::Atom(1));
    m.addAtom(new RDKit::Atom(1));
    m.addAtom(new RDKit::Atom(1));
    m.addConformer(new RDKit::Conformer(m.getNumAtoms()));
    m.getConformer().setAtomPos(0, RDGeom::Point3D(1, 0, 0));
    m.getConformer().setAtomPos(1, RDGeom::Point3D(-1, 0, 0));
    m.getConformer().setAtomPos(2, RDGeom::Point3D(0, 1, 0));
    m.getConformer().setAtomPos(3, RDGeom::Point3D(0, -1, 0));
    m.getConformer().setAtomPos(4, RDGeom::Point3D(0, 0, 1));
    m.getConformer().setAtomPos(5, RDGeom::Point3D(0, 0, -1));
    double val;
    val = RDKit::Descriptors::radiusOfGyration(m);
    TEST_ASSERT(fabs(val) > 0.1);
    val = RDKit::Descriptors::eccentricity(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::asphericity(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::spherocityIndex(m);
    TEST_ASSERT(fabs(1. - val) < 1e-4);
  }

  {
    RDKit::RWMol m;
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addAtom(new RDKit::Atom(6));
    m.addConformer(new RDKit::Conformer(m.getNumAtoms()));
    double val;
    val = RDKit::Descriptors::radiusOfGyration(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::inertialShapeFactor(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::eccentricity(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::asphericity(m);
    TEST_ASSERT(fabs(val) < 1e-4);
    val = RDKit::Descriptors::spherocityIndex(m);
    TEST_ASSERT(fabs(val) < 1e-4);
  }

  BOOST_LOG(rdErrorLog) << "  done" << std::endl;
}

//-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*
//
//-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*-*
int main() {
  RDLog::InitLogs();
  testPMI1();
  testPMI2();

#if 1
  testNPR1();
  testPMIEdges();
  testNPREdges();
  test3DVals();
  test3DEdges();
#endif
}


#include <EigenUtils.h>
#include <CudaUtils.h>
#include <iostream>

Eigen::MatrixXf EigenUtils::toEigen(const CudaUtils::MatrixF &x)
{
	int rows = x.rows;
	int cols = x.cols;
	Eigen::MatrixXf r(rows, cols);
	float *p = x.data;
	for (int i = 0; i < rows; i++)
	{
		for (int j = 0; j < cols; j++)
		{
			r(i, j) = p[cols * i + j];
		}
	}
	return r;
}

CudaUtils::MatrixF EigenUtils::toMatrix(const Eigen::MatrixXf &x)
{
	int rows = x.rows();
	int cols = x.cols();
	CudaUtils::MatrixF r = CudaUtils::allocateMatrixOnHostF(rows, cols);
	float *p = r.data;
	for (int i = 0; i < rows; i++)
	{
		for (int j = 0; j < cols; j++)
		{
			p[cols * i + j] = x(i, j);
		}
	}
	return r;
}

bool EigenUtils::almostEqual(const Eigen::MatrixXf &lhs, const Eigen::MatrixXf &rhs, float tol, int verbose)
{
	if (lhs.rows() != rhs.rows())
	{
		if (verbose > 0)
			std::cout << "rows does not match.\n";
		return false;
	}
	if (lhs.cols() != rhs.cols())
	{
		if (verbose > 0)
			std::cout << "rows does not match.\n";
		return false;
	}
	float max_err = 0.0f;
	for (int j = 0; j < lhs.rows(); j++)
		for (int i = 0; i < lhs.cols(); i++)
		{
			float t = fabs(lhs(j, i) - rhs(j, i));
			if (t > max_err)
				max_err = t;
		}
	if (max_err > tol)
	{
		if (verbose > 0)
			std::cout << "max err greater then tol. max_err = " << max_err << " tol = " << tol << "\n";
		return false;
	}
	if (verbose > 1)
		std::cout << "max_err = " << max_err << " tol = " << tol << "\n";
	return true;
}


#ifndef PRESSIODEMOAPPS_SET_ZERO_HPP_
#define PRESSIODEMOAPPS_SET_ZERO_HPP_

#include <vector>
#include <Eigen/Dense>
#include <Eigen/Sparse>

namespace pressiodemoapps{

template<class value_type, int I0, int I1>
void set_zero(Eigen::Matrix<value_type, I0, I1, Eigen::ColMajor> & a){
  a.setZero();
}

template<class value_type, int I0, int I1>
void set_zero(Eigen::Matrix<value_type, I0, I1, Eigen::RowMajor> & a){
  a.setZero();
}

template<class value_type, class IntType>
void set_zero(Eigen::SparseMatrix<value_type, Eigen::RowMajor, IntType> & M)
{
  auto values = M.valuePtr();
  for (int i=0; i<M.nonZeros(); ++i){
    values[i] = static_cast<value_type>(0);
  }
}

template<class value_type, class IntType>
void set_zero(Eigen::SparseMatrix<value_type, Eigen::ColMajor, IntType> & M)
{
  auto values = M.valuePtr();
  for (int i=0; i<M.nonZeros(); ++i){
    values[i] = static_cast<value_type>(0);
  }
}

template<class value_type>
void set_zero(Eigen::Ref<Eigen::Matrix<value_type, -1, 1>> & a)
{
  for (int i=0; i<a.rows(); ++i){
    a(i) = static_cast<value_type>(0);
  }
}

template<class value_type>
void set_zero(Eigen::Ref<Eigen::Matrix<value_type, -1, -1, Eigen::ColMajor>> & a)
{
  for (int i=0; i<a.rows(); ++i){
    for (int j=0; j<a.cols(); ++j){
      a(i,j) = static_cast<value_type>(0);
    }
  }
}

template<class value_type>
void set_zero(Eigen::Ref<Eigen::Matrix<value_type, -1, -1, Eigen::RowMajor>> & a)
{
  for (int i=0; i<a.rows(); ++i){
    for (int j=0; j<a.cols(); ++j){
      a(i,j) = static_cast<value_type>(0);
    }
  }
}

#ifdef PRESSIODEMOAPPS_ENABLE_OPENMP
template<class value_type, class IntType>
void set_zero_omp(Eigen::SparseMatrix<value_type, Eigen::RowMajor, IntType> & M)
{
  auto values = M.valuePtr();
#pragma omp for schedule(auto)
  for (int i=0; i<M.nonZeros(); ++i){
    values[i] = static_cast<value_type>(0);
  }
}

template<class value_type, class IntType>
void set_zero_omp(Eigen::SparseMatrix<value_type, Eigen::ColMajor, IntType> & M)
{
  auto values = M.valuePtr();
#pragma omp for schedule(auto)
  for (int i=0; i<M.nonZeros(); ++i){
    values[i] = static_cast<value_type>(0);
  }
}

template<class value_type>
void set_zero_omp(Eigen::Matrix<value_type, -1, 1> & V)
{
#pragma omp for schedule(auto)
   for (decltype(V.size()) i=0; i<V.size(); ++i){
     V(i) = static_cast<value_type>(0);
   }
}

template<class value_type>
void set_zero_omp(Eigen::Ref<Eigen::Matrix<value_type, -1, 1>> & V)
{
#pragma omp for schedule(auto)
   for (decltype(V.size()) i=0; i<V.size(); ++i){
     V(i) = static_cast<value_type>(0);
   }
}
#endif

}
#endif