// Copyright (C) 2018  Davis E. King (davis@dlib.net)
// License: Boost Software License   See LICENSE.txt for the full license.
#ifndef DLIB_ISOTONIC_ReGRESSION_H_
#define DLIB_ISOTONIC_ReGRESSION_H_

#include "isotonic_regression_abstract.h"
#include <vector>
#include <utility>
#include <cstddef>

namespace dlib
{
    class isotonic_regression
    {
    public:

        template <
            typename const_iterator, 
            typename iterator
            >
        void operator() (
            const_iterator begin,
            const_iterator end,
            iterator obegin
        )
        {
            do_isotonic_regression(begin, end);

            // unpack blocks to output
            for (auto& block : blocks)
            {
                for (size_t k = 0; k < block.num; ++k)
                    set_val(*obegin++, block.avg);
            }

            blocks.clear();
        }

        void operator() (
            std::vector<double>& vect
        ) { (*this)(vect.begin(), vect.end(), vect.begin()); }

        template <typename T, typename U>
        void operator() (
            std::vector<std::pair<T,U>>& vect
        ) { (*this)(vect.begin(), vect.end(), vect.begin()); }


        template <
            typename const_iterator, 
            typename iterator
            >
        void fit_with_linear_output_interpolation (
            const_iterator begin,
            const_iterator end,
            iterator obegin
        )
        {
            do_isotonic_regression(begin, end);

            // Unpack blocks to output, but here instead of producing the step function
            // output we linearly interpolate.  Note that this actually fits the data less
            // than the step-function, but in many applications might be closer to what you
            // really want when using isotonic_regression than the step function.
            for (size_t i = 0; i < blocks.size(); ++i)
            {
                auto& block = blocks[i];

                double prev = (blocks.front().avg + block.avg)/2;
                if (i > 0)
                    prev = (blocks[i-1].avg+block.avg)/2;

                double next = (blocks.back().avg + block.avg)/2;
                if (i+1 < blocks.size())
                    next = (blocks[i+1].avg+block.avg)/2;

                for (size_t k = 0; k < block.num; ++k)
                {
                    const auto mid = block.num/2.0;
                    if (k < mid)
                    {
                        const double alpha = k/mid;
                        set_val(*obegin++, (1-alpha)*prev + alpha*block.avg);
                    }
                    else
                    {
                        const double alpha = k/mid-1;
                        set_val(*obegin++, alpha*next + (1-alpha)*block.avg);
                    }
                }
            }

            blocks.clear();
        }

        void fit_with_linear_output_interpolation (
            std::vector<double>& vect
        ) { fit_with_linear_output_interpolation(vect.begin(), vect.end(), vect.begin()); }

        template <typename T, typename U>
        void fit_with_linear_output_interpolation (
            std::vector<std::pair<T,U>>& vect
        ) { fit_with_linear_output_interpolation(vect.begin(), vect.end(), vect.begin()); }

    private:

        template <
            typename const_iterator
            >
        void do_isotonic_regression (
            const_iterator begin,
            const_iterator end
        )
        {
            blocks.clear();

            // Do the actual isotonic regression.  The output is a step-function and is
            // stored in the vector of blocks.
            for (auto i = begin; i != end; ++i)
            {
                blocks.emplace_back(get_val(*i));
                while (blocks.size() > 1 && prev_block().avg > current_block().avg)
                {
                    // merge the last two blocks.
                    prev_block() = prev_block() + current_block();
                    blocks.pop_back();
                }
            }
        }


        template <typename T>
        static double get_val(const T& v) { return v;}

        template <typename T, typename U>
        static double get_val(const std::pair<T,U>& v) { return v.second;}

        template <typename T>
        static void set_val(T& v, double val) { v = val;}

        template <typename T, typename U>
        static void set_val(std::pair<T,U>& v, double val) { v.second = val;}



        struct block_t
        {
            block_t(double val) : num(1), avg(val) {}
            block_t(size_t n, double val) : num(n), avg(val) {}

            size_t num;
            double avg;

            inline block_t operator+(const block_t& rhs) const
            {
                return block_t(num+rhs.num,
                    (num*avg + rhs.num*rhs.avg)/(num+rhs.num));
            }
        };

        inline block_t& prev_block() { return blocks[blocks.size()-2]; }
        inline block_t& current_block() { return blocks.back(); }

        std::vector<block_t> blocks;
    };
}

#endif // DLIB_ISOTONIC_ReGRESSION_H_