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

#include "one_vs_all_trainer_abstract.h"

#include "one_vs_all_decision_function.h"
#include <vector>

#include "multiclass_tools.h"

#include <sstream>
#include <iostream>

#include "../any.h"
#include <map>
#include <set>
#include "../threads.h"

namespace dlib
{

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

    template <
        typename any_trainer,
        typename label_type_ = double
        >
    class one_vs_all_trainer
    {
    public:
        typedef label_type_ label_type;

        typedef typename any_trainer::sample_type sample_type;
        typedef typename any_trainer::scalar_type scalar_type;
        typedef typename any_trainer::mem_manager_type mem_manager_type;

        typedef one_vs_all_decision_function<one_vs_all_trainer> trained_function_type;

        one_vs_all_trainer (
        ) : 
            verbose(false),
            num_threads(4)
        {}

        void set_trainer (
            const any_trainer& trainer
        )
        {
            default_trainer = trainer;
            trainers.clear();
        }

        void set_trainer (
            const any_trainer& trainer,
            const label_type& l
        )
        {
            trainers[l] = trainer;
        }

        void be_verbose (
        )
        {
            verbose = true;
        }

        void be_quiet (
        )
        {
            verbose = false;
        }

        void set_num_threads (
            unsigned long num
        )
        {
            num_threads = num;
        }

        unsigned long get_num_threads (
        ) const
        {
            return num_threads;
        }

        struct invalid_label : public dlib::error 
        { 
            invalid_label(const std::string& msg, const label_type& l_
                ) : dlib::error(msg), l(l_) {};

            virtual ~invalid_label(
            ) throw() {}

            label_type l;
        };

        trained_function_type train (
            const std::vector<sample_type>& all_samples,
            const std::vector<label_type>& all_labels
        ) const
        {
            // make sure requires clause is not broken
            DLIB_ASSERT(is_learning_problem(all_samples,all_labels),
                "\t trained_function_type one_vs_all_trainer::train(all_samples,all_labels)"
                << "\n\t invalid inputs were given to this function"
                << "\n\t all_samples.size():     " << all_samples.size() 
                << "\n\t all_labels.size():      " << all_labels.size() 
                );

            const std::vector<label_type> distinct_labels = select_all_distinct_labels(all_labels);

            // make sure we have a trainer object for each of the label types.
            for (unsigned long i = 0; i < distinct_labels.size(); ++i)
            {
                const label_type l = distinct_labels[i];
                const typename binary_function_table::const_iterator itr = trainers.find(l);

                if (itr == trainers.end() && default_trainer.is_empty())
                {
                    std::ostringstream sout;
                    sout << "In one_vs_all_trainer, no trainer registered for the " << l << " label.";
                    throw invalid_label(sout.str(), l);
                }
            }


            // now do the training
            parallel_for_helper helper(all_samples,all_labels,default_trainer,trainers,verbose,distinct_labels);
            parallel_for(num_threads, 0, distinct_labels.size(), helper, 500);

            if (helper.error_message.size() != 0)
            {
                throw dlib::error("binary trainer threw while training one vs. all classifier.  Error was: " + helper.error_message);
            }
            return trained_function_type(helper.dfs);
        }

    private:

        typedef std::map<label_type, any_trainer> binary_function_table;
        struct parallel_for_helper
        {
            parallel_for_helper(
                const std::vector<sample_type>& all_samples_,
                const std::vector<label_type>& all_labels_,
                const any_trainer& default_trainer_,
                const binary_function_table& trainers_,
                const bool verbose_,
                const std::vector<label_type>& distinct_labels_ 
            ) : 
                all_samples(all_samples_),
                all_labels(all_labels_),
                default_trainer(default_trainer_),
                trainers(trainers_), 
                verbose(verbose_),
                distinct_labels(distinct_labels_)
            {}

            void operator()(long i) const 
            {
                try
                {
                    std::vector<scalar_type> labels;

                    const label_type l = distinct_labels[i];

                    // setup one of the one vs all training sets
                    for (unsigned long k = 0; k < all_samples.size(); ++k)
                    {
                        if (all_labels[k] == l)
                            labels.push_back(+1);
                        else 
                            labels.push_back(-1);
                    }


                    if (verbose)
                    {
                        auto_mutex lock(class_mutex);
                        std::cout << "Training classifier for " << l << " vs. all" << std::endl;
                    }

                    any_trainer trainer;
                    // now train a binary classifier using the samples we selected
                    { auto_mutex lock(class_mutex); 
                    const typename binary_function_table::const_iterator itr = trainers.find(l);
                    if (itr != trainers.end())
                        trainer = itr->second;
                    else 
                        trainer = default_trainer;
                    }

                    any_decision_function<sample_type,scalar_type> binary_df = trainer.train(all_samples, labels);

                    auto_mutex lock(class_mutex);
                    dfs[l] = binary_df;
                }
                catch (std::exception& e)
                {
                    auto_mutex lock(class_mutex);
                    error_message = e.what();
                }
            }

            mutable typename trained_function_type::binary_function_table dfs;
            mutex class_mutex;
            mutable std::string error_message;

            const std::vector<sample_type>& all_samples;
            const std::vector<label_type>& all_labels;
            const any_trainer& default_trainer;
            const binary_function_table& trainers;
            const bool verbose;
            const std::vector<label_type>& distinct_labels;
        };

        any_trainer default_trainer;

        binary_function_table trainers;

        bool verbose;
        unsigned long num_threads;

    };

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

}

#endif // DLIB_ONE_VS_ALL_TRAiNER_Hh_