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

    #include "svm_rank_trainer_abstract.h"

    #include "ranking_tools.h"
    #include "../algs.h"
    #include "../optimization.h"
    #include "function.h"
    #include "kernel.h"
    #include "sparse_vector.h"
    #include <iostream>

    namespace dlib
    {

    // ----------------------------------------------------------------------------------------

        template <
            typename matrix_type,
            typename sample_type
            >
        class oca_problem_ranking_svm : public oca_problem<matrix_type >
        {
        public:
            /*
                This class is used as part of the implementation of the svm_rank_trainer
                defined towards the end of this file.
            */

            typedef typename matrix_type::type scalar_type;

            oca_problem_ranking_svm(
                const scalar_type C_,
                const std::vector<ranking_pair<sample_type> >& samples_,
                const bool be_verbose_,
                const scalar_type eps_,
                const unsigned long max_iter
            ) :
                samples(samples_),
                C(C_),
                be_verbose(be_verbose_),
                eps(eps_),
                max_iterations(max_iter)
            {
            }

            virtual scalar_type get_c (
            ) const
            {
                return C;
            }

            virtual long get_num_dimensions (
            ) const
            {
                return max_index_plus_one(samples);
            }

            virtual bool optimization_status (
                scalar_type current_objective_value,
                scalar_type current_error_gap,
                scalar_type current_risk_value,
                scalar_type current_risk_gap,
                unsigned long num_cutting_planes,
                unsigned long num_iterations
            ) const
            {
                if (be_verbose)
                {
                    using namespace std;
                    cout << "objective:     " << current_objective_value << endl;
                    cout << "objective gap: " << current_error_gap << endl;
                    cout << "risk:          " << current_risk_value << endl;
                    cout << "risk gap:      " << current_risk_gap << endl;
                    cout << "num planes:    " << num_cutting_planes << endl;
                    cout << "iter:          " << num_iterations << endl;
                    cout << endl;
                }

                if (num_iterations >= max_iterations)
                    return true;

                if (current_risk_gap < eps)
                    return true;

                return false;
            }

            virtual bool risk_has_lower_bound (
                scalar_type& lower_bound
            ) const
            {
                lower_bound = 0;
                return true;
            }

            virtual void get_risk (
                matrix_type& w,
                scalar_type& risk,
                matrix_type& subgradient
            ) const
            {
                subgradient.set_size(w.size(),1);
                subgradient = 0;
                risk = 0;

                // Note that we want the risk value to be in terms of the fraction of overall
                // rank flips.  So a risk of 0.1 would mean that rank flips happen < 10% of the
                // time.

                std::vector<double> rel_scores;
                std::vector<double> nonrel_scores;
                std::vector<unsigned long> rel_counts;
                std::vector<unsigned long> nonrel_counts;

                unsigned long total_pairs = 0;

                // loop over all the samples and compute the risk and its subgradient at the current solution point w
                for (unsigned long i = 0; i < samples.size(); ++i)
                {
                    rel_scores.resize(samples[i].relevant.size());
                    nonrel_scores.resize(samples[i].nonrelevant.size());

                    for (unsigned long k = 0; k < rel_scores.size(); ++k)
                        rel_scores[k] = dot(samples[i].relevant[k], w);

                    for (unsigned long k = 0; k < nonrel_scores.size(); ++k)
                        nonrel_scores[k] = dot(samples[i].nonrelevant[k], w) + 1;

                    count_ranking_inversions(rel_scores, nonrel_scores, rel_counts, nonrel_counts);

                    total_pairs += rel_scores.size()*nonrel_scores.size();

                    for (unsigned long k = 0; k < rel_counts.size(); ++k)
                    {
                        if (rel_counts[k] != 0)
                        {
                            risk -= rel_counts[k]*rel_scores[k];
                            subtract_from(subgradient, samples[i].relevant[k], rel_counts[k]);
                        }
                    }

                    for (unsigned long k = 0; k < nonrel_counts.size(); ++k)
                    {
                        if (nonrel_counts[k] != 0)
                        {
                            risk += nonrel_counts[k]*nonrel_scores[k];
                            add_to(subgradient, samples[i].nonrelevant[k], nonrel_counts[k]);
                        }
                    }

                }

                const scalar_type scale = 1.0/total_pairs;

                risk *= scale;
                subgradient = scale*subgradient;
            }

        private:

        // -----------------------------------------------------
        // -----------------------------------------------------

            const std::vector<ranking_pair<sample_type> >& samples;
            const scalar_type C;

            const bool be_verbose;
            const scalar_type eps;
            const unsigned long max_iterations;
        };

    // ----------------------------------------------------------------------------------------

        template <
            typename matrix_type,
            typename sample_type,
            typename scalar_type
            >
        oca_problem_ranking_svm<matrix_type, sample_type> make_oca_problem_ranking_svm (
            const scalar_type C,
            const std::vector<ranking_pair<sample_type> >& samples,
            const bool be_verbose,
            const scalar_type eps,
            const unsigned long max_iterations
        )
        {
            return oca_problem_ranking_svm<matrix_type, sample_type>(
                C, samples, be_verbose, eps, max_iterations);
        }

    // ----------------------------------------------------------------------------------------

        template <
            typename K
            >
        class svm_rank_trainer
        {

        public:
            typedef K kernel_type;
            typedef typename kernel_type::scalar_type scalar_type;
            typedef typename kernel_type::sample_type sample_type;
            typedef typename kernel_type::mem_manager_type mem_manager_type;
            typedef decision_function<kernel_type> trained_function_type;

            // You are getting a compiler error on this line because you supplied a non-linear kernel
            // to the svm_rank_trainer object.  You have to use one of the linear kernels with this
            // trainer.
            COMPILE_TIME_ASSERT((is_same_type<K, linear_kernel<sample_type> >::value ||
                                 is_same_type<K, sparse_linear_kernel<sample_type> >::value ));

            svm_rank_trainer (
            )
            {
                C = 1;
                verbose = false;
                eps = 0.001;
                max_iterations = 10000;
                learn_nonnegative_weights = false;
                last_weight_1 = false;
            }

            explicit svm_rank_trainer (
                const scalar_type& C_
            )
            {
                // make sure requires clause is not broken
                DLIB_ASSERT(C_ > 0,
                    "\t svm_rank_trainer::svm_rank_trainer()"
                    << "\n\t C_ must be greater than 0"
                    << "\n\t C_:    " << C_
                    << "\n\t this: " << this
                    );

                C = C_;
                verbose = false;
                eps = 0.001;
                max_iterations = 10000;
                learn_nonnegative_weights = false;
                last_weight_1 = false;
            }

            void set_epsilon (
                scalar_type eps_
            )
            {
                // make sure requires clause is not broken
                DLIB_ASSERT(eps_ > 0,
                    "\t void svm_rank_trainer::set_epsilon()"
                    << "\n\t eps_ must be greater than 0"
                    << "\n\t eps_: " << eps_
                    << "\n\t this: " << this
                    );

                eps = eps_;
            }

            const scalar_type get_epsilon (
            ) const { return eps; }

            unsigned long get_max_iterations (
            ) const { return max_iterations; }

            void set_max_iterations (
                unsigned long max_iter
            )
            {
                max_iterations = max_iter;
            }

            void be_verbose (
            )
            {
                verbose = true;
            }

            void be_quiet (
            )
            {
                verbose = false;
            }

            bool forces_last_weight_to_1 (
            ) const
            {
                return last_weight_1;
            }

            void force_last_weight_to_1 (
                bool should_last_weight_be_1
            )
            {
                last_weight_1 = should_last_weight_be_1;
            }

            void set_oca (
                const oca& item
            )
            {
                solver = item;
            }

            const oca get_oca (
            ) const
            {
                return solver;
            }

            const kernel_type get_kernel (
            ) const
            {
                return kernel_type();
            }

            bool learns_nonnegative_weights (
            ) const { return learn_nonnegative_weights; }

            void set_learns_nonnegative_weights (
                bool value
            )
            {
                learn_nonnegative_weights = value;
            }

            void set_c (
                scalar_type C_
            )
            {
                // make sure requires clause is not broken
                DLIB_ASSERT(C_ > 0,
                    "\t void svm_rank_trainer::set_c()"
                    << "\n\t C_ must be greater than 0"
                    << "\n\t C_:    " << C_
                    << "\n\t this: " << this
                    );

                C = C_;
            }

            const scalar_type get_c (
            ) const
            {
                return C;
            }

            const decision_function<kernel_type> train (
                const std::vector<ranking_pair<sample_type> >& samples
            ) const
            {
                // make sure requires clause is not broken
                DLIB_CASSERT(is_ranking_problem(samples) == true,
                    "\t decision_function svm_rank_trainer::train(samples)"
                    << "\n\t invalid inputs were given to this function"
                    << "\n\t samples.size(): " << samples.size()
                    << "\n\t is_ranking_problem(samples): " << is_ranking_problem(samples)
                    );

                typedef matrix<scalar_type,0,1> w_type;
                w_type w;

                const unsigned long num_dims = max_index_plus_one(samples);

                unsigned long num_nonnegative = 0;
                if (learn_nonnegative_weights)
                {
                    num_nonnegative = num_dims;
                }

                unsigned long force_weight_1_idx = std::numeric_limits<unsigned long>::max();
                if (last_weight_1)
                {
                    force_weight_1_idx = num_dims-1;
                }

                solver( make_oca_problem_ranking_svm<w_type>(C, samples, verbose, eps, max_iterations),
                        w,
                        num_nonnegative,
                        force_weight_1_idx);

                // put the solution into a decision function and then return it
                decision_function<kernel_type> df;
                df.b = 0;
                df.basis_vectors.set_size(1);
                // Copy the results into the output basis vector.  The output vector might be a
                // sparse vector container so we need to use this special kind of copy to
                // handle that case.
                assign(df.basis_vectors(0), matrix_cast<scalar_type>(w));
                df.alpha.set_size(1);
                df.alpha(0) = 1;

                return df;
            }

            const decision_function<kernel_type> train (
                const ranking_pair<sample_type>& sample
            ) const
            {
                return train(std::vector<ranking_pair<sample_type> >(1, sample));
            }

        private:

            scalar_type C;
            oca solver;
            scalar_type eps;
            bool verbose;
            unsigned long max_iterations;
            bool learn_nonnegative_weights;
            bool last_weight_1;
        };

    // ----------------------------------------------------------------------------------------

    }

    #endif // DLIB_SVM_RANK_TrAINER_H__