/***************************************************************************
* Copyright (c) Johan Mabille, Sylvain Corlay and Wolf Vollprecht          *
* Copyright (c) QuantStack                                                 *
*                                                                          *
* Distributed under the terms of the BSD 3-Clause License.                 *
*                                                                          *
* The full license is in the file LICENSE, distributed with this software. *
****************************************************************************/

/**
 * @brief construct histogram
 */

#ifndef XTENSOR_HISTOGRAM_HPP
#define XTENSOR_HISTOGRAM_HPP

#include "xtensor.hpp"
#include "xsort.hpp"
#include "xset_operation.hpp"
#include "xview.hpp"

using namespace xt::placeholders;

namespace xt
{
    /**
     * @ingroup digitize
     * @brief Return the indices of the bins to which each value in input array belongs.
     *
     * @param data The data.
     * @param bin_edges The bin-edges. It has to be 1-dimensional and monotonic.
     * @param right Indicating whether the intervals include the right or the left bin edge.
     * @return Output array of indices, of same shape as x.
     */
    template <class E1, class E2>
    inline auto digitize(E1&& data, E2&& bin_edges, bool right = false)
    {
        XTENSOR_ASSERT(bin_edges.dimension() == 1);
        XTENSOR_ASSERT(bin_edges.size() >= 2);
        XTENSOR_ASSERT(std::is_sorted(bin_edges.cbegin(), bin_edges.cend()));
        XTENSOR_ASSERT(xt::amin(data)[0] >= bin_edges[0]);
        XTENSOR_ASSERT(xt::amax(data)[0] <= bin_edges[bin_edges.size() - 1]);

        return xt::searchsorted(std::forward<E2>(bin_edges), std::forward<E1>(data), right);
    }

    /**
     * @ingroup histogram
     * @brief Compute the histogram of a set of data.
     *
     * @param data The data.
     * @param bin_edges The bin-edges. It has to be 1-dimensional and monotonic.
     * @param weights Weight factors corresponding to each data-point.
     * @param density If true the resulting integral is normalized to 1. [default: false]
     * @return An one-dimensional xarray<double>, length: bin_edges.size()-1.
     */
    template <class R = double, class E1, class E2, class E3>
    inline auto histogram(E1&& data, E2&& bin_edges, E3&& weights, bool density = false)
    {
        using size_type = common_size_type_t<std::decay_t<E1>, std::decay_t<E2>, std::decay_t<E3>>;
        using value_type = typename std::decay_t<E3>::value_type;

        XTENSOR_ASSERT(data.dimension() == 1);
        XTENSOR_ASSERT(weights.dimension() == 1);
        XTENSOR_ASSERT(bin_edges.dimension() == 1);
        XTENSOR_ASSERT(weights.size() == data.size());
        XTENSOR_ASSERT(bin_edges.size() >= 2);
        XTENSOR_ASSERT(std::is_sorted(bin_edges.cbegin(), bin_edges.cend()));
        XTENSOR_ASSERT(xt::amin(data)[0] >= bin_edges[0]);
        XTENSOR_ASSERT(xt::amax(data)[0] <= bin_edges[bin_edges.size() - 1]);

        xt::xtensor<value_type, 1> count = xt::zeros<value_type>({ bin_edges.size() - 1 });

        auto sorter = xt::argsort(data);

        size_type ibin = 0;

        for (auto& idx : sorter)
        {
            while (data[idx] >= bin_edges[ibin + 1] && ibin < bin_edges.size() - 2)
            {
                ++ibin;
            }
            count[ibin] += weights[idx];
        }

        xt::xtensor<R, 1> prob = xt::cast<R>(count);

        if (density)
        {
            R n = static_cast<R>(data.size());
            for (size_type i = 0; i < bin_edges.size() - 1; ++i)
            {
                prob[i] /= (static_cast<R>(bin_edges[i + 1] - bin_edges[i]) * n);
            }
        }

        return prob;
    }

    /**
     * @ingroup histogram
     * @brief Compute the histogram of a set of data.
     *
     * @param data The data.
     * @param bin_edges The bin-edges.
     * @param density If true the resulting integral is normalized to 1. [default: false]
     * @return An one-dimensional xarray<double>, length: bin_edges.size()-1.
     */
    template <class E1, class E2>
    inline auto histogram(E1&& data, E2&& bin_edges, bool density = false)
    {
        using value_type = typename std::decay_t<E1>::value_type;

        auto n = data.size();

        return histogram(std::forward<E1>(data),
                         std::forward<E2>(bin_edges),
                         xt::ones<value_type>({ n }),
                         density);
    }

    /**
     * @ingroup histogram
     * @brief Compute the histogram of a set of data.
     *
     * @param data The data.
     * @param bins The number of bins. [default: 10]
     * @param density If true the resulting integral is normalized to 1. [default: false]
     * @return An one-dimensional xarray<double>, length: bin_edges.size()-1.
     */
    template <class E1>
    inline auto histogram(E1&& data, std::size_t bins = 10, bool density = false)
    {
        using value_type = typename std::decay_t<E1>::value_type;

        auto n = data.size();

        auto bin_edges = histogram_bin_edges(data, xt::ones<value_type>({ n }), bins);

        return histogram(std::forward<E1>(data),
                         std::forward<E1>(bin_edges),
                         xt::ones<value_type>({ n }),
                         density);
    }

    /**
     * @ingroup histogram
     * @brief Compute the histogram of a set of data.
     *
     * @param data The data.
     * @param bins The number of bins.
     * @param weights Weight factors corresponding to each data-point.
     * @param density If true the resulting integral is normalized to 1. [default: false]
     * @return An one-dimensional xarray<double>, length: bin_edges.size()-1.
     */
    template <class E1, class E2>
    inline auto histogram(E1&& data, std::size_t bins, E2&& weights, bool density = false)
    {
        auto bin_edges = histogram_bin_edges(data, weights, bins);

        return histogram(std::forward<E1>(data),
                         std::forward<E1>(bin_edges),
                         std::forward<E2>(weights),
                         density);
    }

    /**
     * @ingroup histogram
     * @brief Defines different algorithms to be used in "histogram_bin_edges"
     */
    enum class histogram_algorithm
    {
        automatic,
        linspace,
        logspace,
        uniform
    };

    /**
     * @ingroup histogram
     * @brief Compute the bin-edges of a histogram of a set of data using different algorithms.
     *
     * @param data The data.
     * @param weights Weight factors corresponding to each data-point.
     * @param left The lower-most edge.
     * @param right The upper-most edge.
     * @param bins The number of bins. [default: 10]
     * @param mode The type of algorithm to use. [default: "auto"]
     * @return An one-dimensional xarray<double>, length: bins+1.
     */
    template <class E1, class E2, class E3>
    inline auto histogram_bin_edges(E1&& data,
                                    E2&& weights,
                                    E3 left,
                                    E3 right,
                                    std::size_t bins = 10,
                                    histogram_algorithm mode = histogram_algorithm::automatic)
    {
        // counter and return type
        using size_type = common_size_type_t<std::decay_t<E1>, std::decay_t<E2>>;
        using value_type = typename std::decay_t<E1>::value_type;
        using weights_type = typename std::decay_t<E2>::value_type;

        // basic checks
        // - rank
        XTENSOR_ASSERT(data.dimension() == 1);
        XTENSOR_ASSERT(weights.dimension() == 1);
        // - size
        XTENSOR_ASSERT(weights.size() == data.size());
        // - bounds
        XTENSOR_ASSERT(left <= xt::amin(data)[0]);
        XTENSOR_ASSERT(right >= xt::amax(data)[0]);
        // - non-empty
        XTENSOR_ASSERT(bins > std::size_t(0));

        // act on different modes
        switch (mode)
        {
            // bins of equal width
            case histogram_algorithm::automatic:
            {
                xt::xtensor<value_type, 1> bin_edges
                    = xt::linspace<value_type>(left, right, bins + 1);
                return bin_edges;
            }

            // bins of equal width
            case histogram_algorithm::linspace:
            {
                xt::xtensor<value_type, 1> bin_edges
                    = xt::linspace<value_type>(left, right, bins + 1);
                return bin_edges;
            }

            // bins of logarithmically increasing size
            case histogram_algorithm::logspace:
            {
                using rhs_value_type
                    = std::conditional_t<xtl::is_integral<value_type>::value, double, value_type>;

                xtensor<value_type, 1> bin_edges = xt::cast<value_type>(
                    xt::logspace<rhs_value_type>(std::log10(left), std::log10(right), bins + 1));

                // TODO: replace above with below after 'xsimd' fix
                // xt::xtensor<value_type,1> bin_edges = xt::logspace<value_type>(
                //     std::log10(left), std::log10(right), bins+1);

                return bin_edges;
            }

            // same amount of data in each bin
            case histogram_algorithm::uniform:
            {
                // indices that sort "data"
                auto sorter = xt::argsort(data);

                // histogram: all of equal 'height'
                // - height
                weights_type w
                    = xt::sum<weights_type>(weights)[0] / static_cast<weights_type>(bins);
                // - apply to all bins
                xt::xtensor<weights_type, 1> count = w * xt::ones<weights_type>({ bins });

                // take cumulative sum, to act as easy look-up
                count = xt::cumsum(count);

                // edges
                // - allocate
                std::vector<size_t> shape = { bins + 1 };
                xt::xtensor<value_type, 1> bin_edges = xtensor<value_type, 1>::from_shape(shape);
                // - first/last edge
                bin_edges[0] = left;
                bin_edges[bins] = right;
                // - cumulative weight
                weights_type cum_weight = static_cast<weights_type>(0);
                // - current bin
                size_type ibin = 0;
                // - loop to find interior bin-edges
                for (size_type i = 0; i < weights.size(); ++i)
                {
                    if (cum_weight >= count[ibin])
                    {
                        bin_edges[ibin + 1] = data[sorter[i]];
                        ++ibin;
                    }
                    cum_weight += weights[sorter[i]];
                }
                return bin_edges;
            }

            // bins of equal width
            default:
            {
                xt::xtensor<value_type, 1> bin_edges
                    = xt::linspace<value_type>(left, right, bins + 1);
                return bin_edges;
            }
        }
    }

    /**
     * @ingroup histogram
     * @brief Compute the bin-edges of a histogram of a set of data using different algorithms.
     *
     * @param data The data.
     * @param weights Weight factors corresponding to each data-point.
     * @param bins The number of bins. [default: 10]
     * @param mode The type of algorithm to use. [default: "auto"]
     * @return An one-dimensional xarray<double>, length: bins+1.
     */
    template <class E1, class E2>
    inline auto histogram_bin_edges(E1&& data,
                                    E2&& weights,
                                    std::size_t bins = 10,
                                    histogram_algorithm mode = histogram_algorithm::automatic)
    {
        using value_type = typename std::decay_t<E1>::value_type;
        std::array<value_type, 2> left_right;
        left_right = xt::minmax(data)();

        return histogram_bin_edges(std::forward<E1>(data),
                                   std::forward<E2>(weights),
                                   left_right[0],
                                   left_right[1],
                                   bins,
                                   mode);
    }

    /**
     * @ingroup histogram
     * @brief Compute the bin-edges of a histogram of a set of data using different algorithms.
     *
     * @param data The data.
     * @param bins The number of bins. [default: 10]
     * @param mode The type of algorithm to use. [default: "auto"]
     * @return An one-dimensional xarray<double>, length: bins+1.
     */
    template <class E1>
    inline auto histogram_bin_edges(E1&& data,
                                    std::size_t bins = 10,
                                    histogram_algorithm mode = histogram_algorithm::automatic)
    {
        using value_type = typename std::decay_t<E1>::value_type;

        auto n = data.size();
        std::array<value_type, 2> left_right;
        left_right = xt::minmax(data)();

        return histogram_bin_edges(std::forward<E1>(data),
                                   xt::ones<value_type>({ n }),
                                   left_right[0],
                                   left_right[1],
                                   bins,
                                   mode);
    }

    /**
     * @ingroup histogram
     * @brief Compute the bin-edges of a histogram of a set of data using different algorithms.
     *
     * @param data The data.
     * @param left The lower-most edge.
     * @param right The upper-most edge.
     * @param bins The number of bins. [default: 10]
     * @param mode The type of algorithm to use. [default: "auto"]
     * @return An one-dimensional xarray<double>, length: bins+1.
     */
    template <class E1, class E2>
    inline auto histogram_bin_edges(E1&& data,
                                    E2 left,
                                    E2 right,
                                    std::size_t bins = 10,
                                    histogram_algorithm mode = histogram_algorithm::automatic)
    {
        using value_type = typename std::decay_t<E1>::value_type;

        auto n = data.size();

        return histogram_bin_edges(
            std::forward<E1>(data), xt::ones<value_type>({ n }), left, right, bins, mode);
    }

    /**
     * Count number of occurrences of each value in array of non-negative ints.
     *
     * The number of bins (of size 1) is one larger than the largest value in x.
     * If minlength is specified, there will be at least this number of bins in
     * the output array (though it will be longer if necessary, depending on the
     * contents of x). Each bin gives the number of occurrences of its index
     * value in x. If weights is specified the input array is weighted by it,
     * i.e. if a value ``n`` is found at position ``i``, ``out[n] += weight[i]``
     * instead of ``out[n] += 1``.
     *
     * @param data the 1D container with integers to count into bins
     * @param weights a 1D container with the same number of elements as ``data``
     * @param minlength The minlength
     *
     * @return 1D container with the bincount
     */
    template <class E1, class E2, XTL_REQUIRES(is_xexpression<std::decay_t<E2>>)>
    inline auto bincount(E1&& data, E2&& weights, std::size_t minlength = 0)
    {
        using result_value_type = typename std::decay_t<E2>::value_type;
        using input_value_type = typename std::decay_t<E1>::value_type;
        using size_type = typename std::decay_t<E1>::size_type;

        static_assert(xtl::is_integral<typename std::decay_t<E1>::value_type>::value,
                      "Bincount data has to be integral type.");
        XTENSOR_ASSERT(data.dimension() == 1);
        XTENSOR_ASSERT(weights.dimension() == 1);

        std::array<input_value_type, 2> left_right;
        left_right = xt::minmax(data)();

        if (left_right[0] < input_value_type(0))
        {
            XTENSOR_THROW(std::runtime_error,
                "Data argument for bincount can only contain positive integers!");
        }

        xt::xtensor<result_value_type, 1> res = xt::zeros<result_value_type>(
            { (std::max)(minlength, std::size_t(left_right[1] + 1)) });

        for (size_type i = 0; i < data.size(); ++i)
        {
            res(data(i)) += weights(i);
        }

        return res;
    }

    template <class E1>
    inline auto bincount(E1&& data, std::size_t minlength = 0)
    {
        return bincount(std::forward<E1>(data),
                        xt::ones<typename std::decay_t<E1>::value_type>(data.shape()),
                        minlength);
    }

    /**
     * Get the number of items in each bin, given the fraction of items per bin.
     * The output is such that the total number of items of all bins is exactly "N".
     *
     * @param N the number of items to distribute
     * @param weights fraction of items per bin: a 1D container whose size is the number of bins
     *
     * @return 1D container with the number of items per bin
     */
    template <class E>
    inline xt::xtensor<size_t, 1> bin_items(size_t N, E&& weights)
    {
        if (weights.size() <= std::size_t(1))
        {
            xt::xtensor<size_t, 1> n = N * xt::ones<size_t>({1});
            return n;
        }

        using value_type = typename std::decay_t<E>::value_type;

        XTENSOR_ASSERT(xt::all(weights >= static_cast<value_type>(0)));
        XTENSOR_ASSERT(xt::sum(weights)() > static_cast<value_type>(0));

        xt::xtensor<double, 1> P = xt::cast<double>(weights) / static_cast<double>(xt::sum(weights)());
        xt::xtensor<size_t, 1> n = xt::ceil(static_cast<double>(N) * P);

        if (xt::sum(n)() == N)
        {
            return n;
        }

        xt::xtensor<size_t, 1> d = xt::zeros<size_t>(P.shape());
        xt::xtensor<size_t, 1> sorter = xt::argsort(P);
        sorter = xt::view(sorter, xt::range(P.size(), _, -1));
        sorter = xt::view(sorter, xt::range(0, xt::sum(n)(0) - N));
        xt::view(d, xt::keep(sorter)) = 1;
        n -= d;

        return n;
    }

    /**
     * Get the number of items in each bin, with each bin having approximately the same number of
     * items in it,under the constraint that the total number of items of all bins is exactly "N".
     *
     * @param N the number of items to distribute
     * @param bins the number of bins
     *
     * @return 1D container with the number of items per bin
     */
    inline xt::xtensor<size_t,1> bin_items(size_t N, size_t bins)
    {
        return bin_items(N, xt::ones<double>({bins}));
    }
}

#endif
