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


#include "track_association_function_abstract.h"
#include <vector>
#include <iostream>
#include "../algs.h"
#include "../serialize.h"
#include "assignment_function.h"

namespace dlib
{

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

    template <
        typename detection_type
        > 
    class feature_extractor_track_association
    {
    public:
        typedef typename detection_type::track_type track_type;
        typedef typename track_type::feature_vector_type feature_vector_type;

        typedef detection_type lhs_element;
        typedef track_type rhs_element;

        feature_extractor_track_association() : num_dims(0), num_nonnegative(0) {}

        explicit feature_extractor_track_association (
            unsigned long num_dims_,
            unsigned long num_nonnegative_
        ) : num_dims(num_dims_), num_nonnegative(num_nonnegative_) {}

        unsigned long num_features(
        ) const { return num_dims; }

        unsigned long num_nonnegative_weights (
        ) const { return num_nonnegative; }

        void get_features (
            const detection_type& det,
            const track_type& track,
            feature_vector_type& feats
        ) const
        {
            track.get_similarity_features(det, feats);
        }

        friend void serialize (const feature_extractor_track_association& item, std::ostream& out) 
        { 
            serialize(item.num_dims, out);
            serialize(item.num_nonnegative, out);
        }

        friend void deserialize (feature_extractor_track_association& item, std::istream& in) 
        {
            deserialize(item.num_dims, in);
            deserialize(item.num_nonnegative, in);
        }

    private:
        unsigned long num_dims;
        unsigned long num_nonnegative;
    };

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

    template <
        typename detection_type_
        >
    class track_association_function
    {
    public:

        typedef detection_type_ detection_type;
        typedef typename detection_type::track_type track_type;
        typedef assignment_function<feature_extractor_track_association<detection_type> > association_function_type;

        track_association_function() {}

        track_association_function (
            const association_function_type& assoc_
        ) : assoc(assoc_)
        {
        }

        const association_function_type& get_assignment_function (
        ) const
        {
            return assoc;
        }

        void operator() (
            std::vector<track_type>& tracks,
            const std::vector<detection_type>& dets
        ) const
        {
            std::vector<long> assignments = assoc(dets, tracks);
            std::vector<bool> updated_track(tracks.size(), false);
            // now update all the tracks with the detections that associated to them.
            for (unsigned long i = 0; i < assignments.size(); ++i)
            {
                if (assignments[i] != -1)
                {
                    tracks[assignments[i]].update_track(dets[i]);
                    updated_track[assignments[i]] = true;
                }
                else
                {
                    track_type new_track;
                    new_track.update_track(dets[i]);
                    tracks.push_back(new_track);
                }
            }

            // Now propagate all the tracks that didn't get any detections.
            for (unsigned long i = 0; i < updated_track.size(); ++i)
            {
                if (!updated_track[i])
                    tracks[i].propagate_track();
            }
        }

        friend void serialize (const track_association_function& item, std::ostream& out)
        {
            int version = 1;
            serialize(version, out);
            serialize(item.assoc, out);
        }
        friend void deserialize (track_association_function& item, std::istream& in)
        {
            int version = 0;
            deserialize(version, in);
            if (version != 1)
                throw serialization_error("Unexpected version found while deserializing dlib::track_association_function.");

            deserialize(item.assoc, in);
        }

    private:

        assignment_function<feature_extractor_track_association<detection_type> > assoc;
    };

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

}

#endif // DLIB_TRACK_ASSOCiATION_FUNCTION_H__

