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

#include "assignment_function_abstract.h"
#include "../matrix.h"
#include <vector>
#include "../optimization/max_cost_assignment.h"

namespace dlib
{

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

    template <
        typename feature_extractor 
        >
    class assignment_function
    {
    public:

        typedef typename feature_extractor::lhs_element lhs_element;
        typedef typename feature_extractor::rhs_element rhs_element;


        typedef std::pair<std::vector<lhs_element>, std::vector<rhs_element> > sample_type;

        typedef std::vector<long> label_type;
        typedef label_type result_type;

        assignment_function()
        {
            weights.set_size(fe.num_features());
            weights = 0;
            bias = 0;
            force_assignment = false;
        }

        explicit assignment_function(
            const matrix<double,0,1>& weights_,
            double bias_
        ) : 
            weights(weights_),
            bias(bias_),
            force_assignment(false)
        {
            // make sure requires clause is not broken
            DLIB_ASSERT(fe.num_features() == static_cast<unsigned long>(weights_.size()),
                "\t assignment_function::assignment_function(weights_)"
                << "\n\t These sizes should match"
                << "\n\t fe.num_features(): " << fe.num_features() 
                << "\n\t weights_.size():   " << weights_.size() 
                << "\n\t this: " << this
                );

        }

        assignment_function(
            const matrix<double,0,1>& weights_,
            double bias_,
            const feature_extractor& fe_
        ) :
            fe(fe_),
            weights(weights_),
            bias(bias_),
            force_assignment(false)
        {
            // make sure requires clause is not broken
            DLIB_ASSERT(fe_.num_features() == static_cast<unsigned long>(weights_.size()),
                "\t assignment_function::assignment_function(weights_,fe_)"
                << "\n\t These sizes should match"
                << "\n\t fe_.num_features(): " << fe_.num_features() 
                << "\n\t weights_.size():    " << weights_.size() 
                << "\n\t this: " << this
                );
        }

        assignment_function(
            const matrix<double,0,1>& weights_,
            double bias_,
            const feature_extractor& fe_,
            bool force_assignment_
        ) :
            fe(fe_),
            weights(weights_),
            bias(bias_),
            force_assignment(force_assignment_)
        {
            // make sure requires clause is not broken
            DLIB_ASSERT(fe_.num_features() == static_cast<unsigned long>(weights_.size()),
                "\t assignment_function::assignment_function(weights_,fe_,force_assignment_)"
                << "\n\t These sizes should match"
                << "\n\t fe_.num_features(): " << fe_.num_features() 
                << "\n\t weights_.size():    " << weights_.size() 
                << "\n\t this: " << this
                );
        }

        const feature_extractor& get_feature_extractor (
        ) const { return fe; }

        const matrix<double,0,1>& get_weights (
        ) const { return weights; }

        double get_bias (
        ) const { return bias; }

        bool forces_assignment (
        ) const { return force_assignment; }

        void predict_assignments (
            const std::vector<lhs_element>& lhs,
            const std::vector<rhs_element>& rhs,
            result_type& assignment
        ) const
        {
            assignment.clear();

            matrix<double> cost;
            unsigned long size;
            if (force_assignment)
            {
                size = std::max(lhs.size(), rhs.size());
            }
            else
            {
                size = rhs.size() + lhs.size();
            }
            cost.set_size(size, size);

            typedef typename feature_extractor::feature_vector_type feature_vector_type;
            feature_vector_type feats;

            // now fill out the cost assignment matrix
            for (long r = 0; r < cost.nr(); ++r)
            {
                for (long c = 0; c < cost.nc(); ++c)
                {
                    if (r < (long)lhs.size() && c < (long)rhs.size())
                    {
                        fe.get_features(lhs[r], rhs[c], feats);
                        cost(r,c) = dot(weights, feats) + bias;
                    }
                    else
                    {
                        cost(r,c) = 0;
                    }
                }
            }


            if (cost.size() != 0)
            {
                // max_cost_assignment() only works with integer matrices, so convert from
                // double to integer.
                const double scale = (std::numeric_limits<dlib::int64>::max()/1000)/max(abs(cost));
                matrix<dlib::int64> int_cost = matrix_cast<dlib::int64>(round(cost*scale));
                assignment = max_cost_assignment(int_cost);
                assignment.resize(lhs.size());
            }

            // adjust assignment so that non-assignments have a value of -1
            for (unsigned long i = 0; i < assignment.size(); ++i)
            {
                if (assignment[i] >= (long)rhs.size())
                    assignment[i] = -1;
            }
        }

        void predict_assignments (
            const sample_type& item,
            result_type& assignment
        ) const
        {
            predict_assignments(item.first, item.second, assignment);
        }

        result_type operator()(
            const std::vector<lhs_element>& lhs,
            const std::vector<rhs_element>& rhs 
        ) const
        {
            result_type temp;
            predict_assignments(lhs,rhs,temp);
            return temp;
        }

        result_type operator() (
            const sample_type& item
        ) const
        {
            return (*this)(item.first, item.second);
        }

    private:


        feature_extractor fe;
        matrix<double,0,1> weights;
        double bias;
        bool force_assignment;
    };

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

    template <
        typename feature_extractor
        >
    void serialize (
        const assignment_function<feature_extractor>& item,
        std::ostream& out
    )
    {
        int version = 2;
        serialize(version, out);
        serialize(item.get_feature_extractor(), out);
        serialize(item.get_weights(), out);
        serialize(item.get_bias(), out);
        serialize(item.forces_assignment(), out);
    }

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

    template <
        typename feature_extractor
        >
    void deserialize (
        assignment_function<feature_extractor>& item,
        std::istream& in 
    )
    {
        feature_extractor fe;
        matrix<double,0,1> weights;
        double bias;
        bool force_assignment;
        int version = 0;
        deserialize(version, in);
        if (version != 2)
            throw serialization_error("Unexpected version found while deserializing dlib::assignment_function.");

        deserialize(fe, in);
        deserialize(weights, in);
        deserialize(bias, in);
        deserialize(force_assignment, in);

        item = assignment_function<feature_extractor>(weights, bias, fe, force_assignment);
    }

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

}

#endif // DLIB_ASSIGNMENT_FuNCTION_Hh_