#pragma once

#include <psemek/util/spatial_array.hpp>
#include <boost/math/special_functions/erf.hpp>

#include <iostream>
#include <cmath>
#include <limits>
#include <vector>
#include <algorithm>

namespace psemek::util
{

	namespace detail
	{

		template <typename T>
		constexpr T min()
		{
			if constexpr (std::is_floating_point_v<T>)
			{
				return -std::numeric_limits<T>::infinity();
			}
			else
			{
				return std::numeric_limits<T>::min();
			}
		}

		template <typename T>
		constexpr T max()
		{
			if constexpr (std::is_floating_point_v<T>)
			{
				return std::numeric_limits<T>::infinity();
			}
			else
			{
				return std::numeric_limits<T>::max();
			}
		}

		inline double normal_cdf(double x)
		{
			return 0.5 * (1.0 + boost::math::erf(x / std::sqrt(2.0)));
		}

		inline double normal_cdf_inv(double x)
		{
			return std::sqrt(2.0) * boost::math::erf_inv(2.0 * x - 1.0);
		}

		template <typename T>
		T sqr(T const & value)
		{
			return value * value;
		}

		template <typename T>
		struct base_statistics
		{
			void push(T const & value, std::size_t count = 1)
			{
				*this = merge(*this, singleton(value, count));
			}

			std::size_t count() const
			{
				return count_;
			}

			T mean() const
			{
				return mean_;
			}

			T variance() const
			{
				return variance_;
			}

			T stddev() const
			{
				return std::sqrt(variance_);
			}

			T percentile(double p, T const & min, T const & max) const
			{
				// Assume truncated normal distribution in the range [min, max]
				//   which is the maximum-entropy distribution on this range
				//   with specified mean and variance
				// See
				//   https://en.wikipedia.org/wiki/Maximum_entropy_probability_distribution#Other_examples
				//   https://en.wikipedia.org/wiki/Truncated_normal_distribution

				T const mu = mean();
				T const sigma = stddev();

				if (sigma == T{0})
					return mu;

				T const alpha = (min - mu) / sigma;
				T const beta = (max - mu) / sigma;

				T const t = std::lerp(detail::normal_cdf(alpha), detail::normal_cdf(beta), p);

				if (t <= T{0})
					return min;

				if (t >= T{1})
					return max;

				return mu + sigma * detail::normal_cdf_inv(t);
			}

			friend base_statistics merge(base_statistics const & s1, base_statistics const & s2)
			{
				// See https://stackoverflow.com/questions/1480626/merging-two-statistical-result-sets

				base_statistics result;
				result.count_ = s1.count_ + s2.count_;
				if (result.count_ > 0)
				{
					result.mean_ = (s1.count_ * s1.mean_ + s2.count_ * s2.mean_) / result.count_;
					result.variance_ = (s1.count_ * (s1.variance_ + sqr(s1.mean_ - result.mean_)) + s2.count_ * (s2.variance_ + sqr(s2.mean_ - result.mean_))) / result.count_;
				}
				return result;
			}

		private:
			std::size_t count_ = 0;
			T mean_ = T{0};
			T variance_ = T{0};

			static base_statistics singleton(T const & value, std::size_t count)
			{
				base_statistics result;
				result.count_ = count;
				result.mean_ = value;
				result.variance_ = 0;
				return result;
			}
		};

	}

	template <typename T>
	struct statistics_lite
	{
		void push(T const & value, std::size_t count = 1);

		std::size_t count() const { return base_.count(); }
		T mean() const { return base_.mean(); }
		T variance() const { return base_.variance(); }
		T stddev() const { return base_.stddev(); }
		T min() const { return min_; }
		T max() const { return max_; }

		T percentile(double p) const;

		template <typename H>
		friend statistics_lite<H> merge(statistics_lite<H> const & s1, statistics_lite<H> const & s2);

	private:
		detail::base_statistics<T> base_;
		T min_ = detail::max<T>();
		T max_ = detail::min<T>();
	};

	template <typename T>
	void statistics_lite<T>::push(T const & value, std::size_t count)
	{
		base_.push(value, count);

		min_ = std::min(min_, value);
		max_ = std::max(max_, value);
	}

	template <typename T>
	std::ostream & operator << (std::ostream & os, statistics_lite<T> const & s)
	{
		os << "mean = " << s.mean() << ", dev = " << s.stddev() << ", range = [" << s.min() << " .. " << s.max() << "]";
		return os;
	}

	template <typename T>
	T statistics_lite<T>::percentile(double p) const
	{
		return base_.percentile(p, min_, max_);
	}

	template <typename T>
	statistics_lite<T> merge(statistics_lite<T> const & s1, statistics_lite<T> const & s2)
	{
		statistics_lite<T> result;
		result.base_ = merge(s1.base_, s2.base_);
		result.min_ = std::min(s1.min_, s2.min_);
		result.max_ = std::max(s1.max_, s2.max_);
		return result;
	}

	template <typename T>
	struct statistics_log_bucket
		: statistics_lite<T>
	{
		statistics_log_bucket(T precision);
		statistics_log_bucket(statistics_log_bucket && other);

		statistics_log_bucket & operator = (statistics_log_bucket && other);

		T precision() const { return precision_; }

		T bucket_size() const { return precision() / std::log(T{2}); }

		// Only accepts non-negative values
		void push(T const & value, std::size_t count = 1);

		T percentile(double p) const;

		statistics_lite<T> & lite() { return *this; }
		statistics_lite<T> const & lite() const { return *this; }

		template <typename H>
		friend statistics_log_bucket<H> merge(statistics_log_bucket<H> const & s1, statistics_log_bucket<H> const & s2);

	private:
		T precision_;

		std::size_t zero_count_ = 0;
		util::spatial_array<detail::base_statistics<T>, 1, int> buckets_;
		mutable util::spatial_array<std::size_t, 1, int> prefix_count_;
	};

	template <typename T>
	statistics_log_bucket<T>::statistics_log_bucket(T precision)
		: precision_(precision)
	{}

	template <typename T>
	statistics_log_bucket<T>::statistics_log_bucket(statistics_log_bucket && other)
		: statistics_lite<T>(other.lite())
		, precision_(other.precision_)
		, zero_count_(other.zero_count_)
		, buckets_(std::move(other.buckets_))
		, prefix_count_(std::move(other.prefix_count_))
	{
		other.zero_count_ = 0;
	}

	template <typename T>
	statistics_log_bucket<T> & statistics_log_bucket<T>::operator = (statistics_log_bucket && other)
	{
		if (this == &other)
			return *this;

		static_cast<statistics_lite<T> &>(*this) = other.lite();
		precision_ = other.precision_;
		zero_count_ = other.zero_count_;
		buckets_ = std::move(other.buckets_);
		prefix_count_ = std::move(other.prefix_count_);
		other.zero_count_ = 0;

		return *this;
	}

	template <typename T>
	void statistics_log_bucket<T>::push(T const & value, std::size_t count)
	{
		lite().push(value, count);
		if (value <= T{0}) [[unlikely]]
			++zero_count_;
		else
		{
			int const bucket = std::floor(std::log2(value) / bucket_size());
			buckets_.get(bucket).push(value, count);
		}

		prefix_count_.clear();
	}

	template <typename T>
	T statistics_log_bucket<T>::percentile(double p) const
	{
		if (prefix_count_.empty())
		{
			std::size_t count = zero_count_;

			prefix_count_.get(buckets_.min(0)) = count;
			for (int bucket = buckets_.min(0); bucket < buckets_.max(0); ++bucket)
			{
				count += buckets_.at(bucket).count();
				prefix_count_.get(bucket + 1) = count;
			}
		}

		std::size_t index = std::min<std::size_t>(this->count() - 1, this->count() * p);

		auto it = std::lower_bound(prefix_count_.begin(), prefix_count_.end(), index);

		if (it == prefix_count_.begin())
			return T{0};

		--it;

		int bucket = buckets_.min(0) + (it - prefix_count_.begin());

		double min = std::exp2(bucket * bucket_size());
		double max = std::exp2((bucket + 1) * bucket_size());
		double q = double(index - *it) / (*(it + 1) - *it);

		return buckets_.at(bucket).percentile(q, min, max);
	}

	template <typename T>
	std::ostream & operator << (std::ostream & os, statistics_log_bucket<T> const & s)
	{
		os << s.lite();
		return os;
	}

	template <typename T>
	statistics_log_bucket<T> merge(statistics_log_bucket<T> const & s1, statistics_log_bucket<T> const & s2)
	{
		if (s1.precision() != s2.precision())
			throw util::exception("Cannot merge log-bucket statistics of different precision");

		statistics_log_bucket<T> result(s1.precision());
		result.lite() = merge(s1.lite(), s2.lite());

		for (int bucket = s1.buckets_.min(0); bucket < s1.buckets_.max(0); ++bucket)
		{
			auto const & source = s1.buckets_.at(bucket);
			auto & target = result.buckets_.get(bucket);
			target = merge(target, source);
		}

		for (int bucket = s2.buckets_.min(0); bucket < s2.buckets_.max(0); ++bucket)
		{
			auto const & source = s2.buckets_.at(bucket);
			auto & target = result.buckets_.get(bucket);
			target = merge(target, source);
		}

		return result;
	}

	template <typename T>
	struct statistics
		: statistics_lite<T>
	{
		void push(T const & value, std::size_t count = 1);

		T percentile(double p) const;

		template <typename H>
		friend statistics<H> merge(statistics<H> const & s1, statistics<H> const & s2);

		statistics_lite<T> & lite() { return *this; }
		statistics_lite<T> const & lite() const { return *this; }

	private:
		mutable std::vector<T> values_;
		mutable bool sorted_ = true;
	};

	template <typename T>
	void statistics<T>::push(T const & value, std::size_t count)
	{
		lite().push(value, count);
		values_.insert(values_.end(), value, count);
		sorted_ = false;
	}

	template <typename T>
	T statistics<T>::percentile(double p) const
	{
		if (!sorted_)
		{
			std::sort(values_.begin(), values_.end());
			sorted_ = true;
		}

		return values_[std::min<std::size_t>(values_.size() - 1, values_.size() * p)];
	}

	template <typename T>
	std::ostream & operator << (std::ostream & os, statistics<T> const & s)
	{
		os << s.lite();
		return os;
	}

	template <typename T>
	statistics<T> merge(statistics<T> const & s1, statistics<T> const & s2)
	{
		statistics<T> result;
		result.lite() = merge(s1.lite(), s2.lite());

		result.values_.reserve(s1.values_.size() + s2.values_.size());
		result.sorted_ = s1.sorted_ && s2.sorted_;
		if (result.sorted_)
		{
			std::merge(s1.values_.begin(), s1.values_.end(), s2.values_.begin(), s2.values_.end(), std::back_inserter(result.values_));
		}
		else
		{
			auto it = std::back_inserter(result.values_);
			it = std::copy(s1.values_.begin(), s1.values_.end(), it);
			it = std::copy(s2.values_.begin(), s2.values_.end(), it);
		}

		return result;
	}

}
