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

#include "sequence_segmenter_abstract.h"
#include "../matrix.h"
#include "sequence_labeler.h"
#include <vector>

namespace dlib
{
    // This namespace contains implementation details for the sequence_segmenter.
    namespace impl_ss
    {

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

        // BIO/BILOU labels
        const unsigned int BEGIN   = 0;
        const unsigned int INSIDE  = 1;
        const unsigned int OUTSIDE = 2;
        const unsigned int LAST    = 3;
        const unsigned int UNIT    = 4;


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

        template <typename ss_feature_extractor>
        class feature_extractor
        {
            /*!
                WHAT THIS OBJECT REPRESENTS
                    This is a feature extractor for a sequence_labeler.  It serves to map
                    the interface defined by a sequence_labeler into the kind of interface
                    defined for a sequence_segmenter.
            !*/

        public:
            typedef typename ss_feature_extractor::sequence_type sequence_type;

            ss_feature_extractor fe;

            feature_extractor() {}
            feature_extractor(const ss_feature_extractor& ss_fe_) : fe(ss_fe_) {}

            unsigned long num_nonnegative_weights (
            ) const
            {
                const unsigned long NL = ss_feature_extractor::use_BIO_model ? 3 : 5;
                if (ss_feature_extractor::allow_negative_weights)
                {
                    return 0;
                }
                else
                {
                    // We make everything non-negative except for the label transition
                    // features.
                    return num_features() - NL*NL;
                }
            }

            friend void serialize(const feature_extractor& item, std::ostream& out) 
            {
                serialize(item.fe, out);
            }

            friend void deserialize(feature_extractor& item, std::istream& in) 
            {
                deserialize(item.fe, in);
            }

            unsigned long num_features() const
            {
                const unsigned long NL = ss_feature_extractor::use_BIO_model ? 3 : 5;
                if (ss_feature_extractor::use_high_order_features)
                    return NL + NL*NL + (NL*NL+NL)*fe.num_features()*fe.window_size();
                else
                    return NL + NL*NL + NL*fe.num_features()*fe.window_size();
            }

            unsigned long order() const 
            { 
                return 1; 
            }

            unsigned long num_labels() const 
            { 
                if (ss_feature_extractor::use_BIO_model)
                    return 3;
                else
                    return 5;
            }

        private:

            template <typename feature_setter>
            struct dot_functor
            {
                /*!
                    WHAT THIS OBJECT REPRESENTS
                        This class wraps the feature_setter used by a sequence_labeler
                        and turns it into the kind needed by a sequence_segmenter.
                !*/

                dot_functor(feature_setter& set_feature_, unsigned long offset_) : 
                    set_feature(set_feature_), offset(offset_) {}

                feature_setter& set_feature;
                unsigned long offset;

                inline void operator() (
                    unsigned long feat_index
                )
                {
                    set_feature(offset+feat_index);
                }

                inline void operator() (
                    unsigned long feat_index,
                    double feat_value
                )
                {
                    set_feature(offset+feat_index, feat_value);
                }
            };

        public:

            template <typename EXP>
            bool reject_labeling (
                const sequence_type& ,
                const matrix_exp<EXP>& y,
                unsigned long 
            ) const
            {
                if (ss_feature_extractor::use_BIO_model)
                {
                    // Don't allow BIO label patterns that don't correspond to a sensical
                    // segmentation. 
                    if (y.size() > 1 && y(0) == INSIDE && y(1) == OUTSIDE)
                        return true;
                    if (y.size() == 1 && y(0) == INSIDE)
                        return true;
                }
                else
                {
                    // Don't allow BILOU label patterns that don't correspond to a sensical
                    // segmentation. 
                    if (y.size() > 1)
                    {
                        if (y(1) == BEGIN && y(0) == OUTSIDE)
                            return true;
                        if (y(1) == BEGIN && y(0) == UNIT)
                            return true;
                        if (y(1) == BEGIN && y(0) == BEGIN)
                            return true;

                        if (y(1) == INSIDE && y(0) == BEGIN)
                            return true;
                        if (y(1) == INSIDE && y(0) == OUTSIDE)
                            return true;
                        if (y(1) == INSIDE && y(0) == UNIT)
                            return true;

                        if (y(1) == OUTSIDE && y(0) == INSIDE)
                            return true;
                        if (y(1) == OUTSIDE && y(0) == LAST)
                            return true;

                        if (y(1) == LAST && y(0) == INSIDE)
                            return true;
                        if (y(1) == LAST && y(0) == LAST)
                            return true;

                        if (y(1) == UNIT && y(0) == INSIDE)
                            return true;
                        if (y(1) == UNIT && y(0) == LAST)
                            return true;
                    }
                    else
                    {
                        if (y(0) == INSIDE)
                            return true;
                        if (y(0) == LAST)
                            return true;
                    }
                }
                return false;
            }

            template <typename feature_setter, typename EXP>
            void get_features (
                feature_setter& set_feature,
                const sequence_type& x,
                const matrix_exp<EXP>& y,
                unsigned long position
            ) const
            {
                unsigned long offset = 0;

                const int window_size = fe.window_size();

                const int base_dims = fe.num_features();
                for (int i = 0; i < window_size; ++i)
                {
                    const long pos = i-window_size/2 + static_cast<long>(position);
                    if (0 <= pos && pos < (long)x.size())
                    {
                        const unsigned long off1 = y(0)*base_dims;
                        dot_functor<feature_setter> fs1(set_feature, offset+off1);
                        fe.get_features(fs1, x, pos);

                        if (ss_feature_extractor::use_high_order_features && y.size() > 1)
                        {
                            const unsigned long off2 = num_labels()*base_dims + (y(0)*num_labels()+y(1))*base_dims;
                            dot_functor<feature_setter> fs2(set_feature, offset+off2);
                            fe.get_features(fs2, x, pos);
                        }
                    }

                    if (ss_feature_extractor::use_high_order_features)
                        offset += num_labels()*base_dims + num_labels()*num_labels()*base_dims;
                    else
                        offset += num_labels()*base_dims;
                }

                // Pull out an indicator feature for the type of transition between the
                // previous label and the current label.
                if (y.size() > 1)
                    set_feature(offset + y(1)*num_labels() + y(0));

                offset += num_labels()*num_labels();
                // pull out an indicator feature for the current label.  This is the per
                // label bias.
                set_feature(offset + y(0));
            }
        };

    } // end namespace impl_ss

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

    template <
        typename feature_extractor
        >
    unsigned long total_feature_vector_size (
        const feature_extractor& fe
    )
    {
        const unsigned long NL = feature_extractor::use_BIO_model ? 3 : 5;
        if (feature_extractor::use_high_order_features)
            return NL + NL*NL + (NL*NL+NL)*fe.num_features()*fe.window_size();
        else
            return NL + NL*NL + NL*fe.num_features()*fe.window_size();
    }

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

    template <
        typename feature_extractor
        >
    class sequence_segmenter
    {
    public:
        typedef typename feature_extractor::sequence_type sample_sequence_type;
        typedef std::vector<std::pair<unsigned long, unsigned long> > segmented_sequence_type;


        sequence_segmenter()
        {
#ifdef ENABLE_ASSERTS
            const feature_extractor& fe = labeler.get_feature_extractor().fe;
            DLIB_ASSERT(fe.window_size() >= 1 && fe.num_features() >= 1,
                "\t sequence_segmenter::sequence_segmenter()"
                << "\n\t An invalid feature extractor was supplied."
                << "\n\t fe.window_size():  " << fe.window_size() 
                << "\n\t fe.num_features(): " << fe.num_features() 
                << "\n\t this: " << this
            );
#endif
        }

        explicit sequence_segmenter(
            const matrix<double,0,1>& weights
        ) : 
            labeler(weights)
        {
#ifdef ENABLE_ASSERTS
            const feature_extractor& fe = labeler.get_feature_extractor().fe;
            // make sure requires clause is not broken
            DLIB_ASSERT(total_feature_vector_size(fe) == (unsigned long)weights.size(),
                "\t sequence_segmenter::sequence_segmenter(weights)"
                << "\n\t These sizes should match"
                << "\n\t total_feature_vector_size(fe):  " << total_feature_vector_size(fe) 
                << "\n\t weights.size(): " << weights.size() 
                << "\n\t this: " << this
                );
            DLIB_ASSERT(fe.window_size() >= 1 && fe.num_features() >= 1,
                "\t sequence_segmenter::sequence_segmenter()"
                << "\n\t An invalid feature extractor was supplied."
                << "\n\t fe.window_size():  " << fe.window_size() 
                << "\n\t fe.num_features(): " << fe.num_features() 
                << "\n\t this: " << this
            );
#endif
        }

        sequence_segmenter(
            const matrix<double,0,1>& weights,
            const feature_extractor& fe
        ) :
            labeler(weights, impl_ss::feature_extractor<feature_extractor>(fe))
        {
            // make sure requires clause is not broken
            DLIB_ASSERT(total_feature_vector_size(fe) == (unsigned long)weights.size(),
                "\t sequence_segmenter::sequence_segmenter(weights,fe)"
                << "\n\t These sizes should match"
                << "\n\t total_feature_vector_size(fe):  " << total_feature_vector_size(fe) 
                << "\n\t weights.size(): " << weights.size() 
                << "\n\t this: " << this
                );
            DLIB_ASSERT(fe.window_size() >= 1 && fe.num_features() >= 1,
                "\t sequence_segmenter::sequence_segmenter()"
                << "\n\t An invalid feature extractor was supplied."
                << "\n\t fe.window_size():  " << fe.window_size() 
                << "\n\t fe.num_features(): " << fe.num_features() 
                << "\n\t this: " << this
            );
        }

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

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

        segmented_sequence_type operator() (
            const sample_sequence_type& x
        ) const
        {
            segmented_sequence_type y;
            segment_sequence(x,y);
            return y;
        }

        void segment_sequence (
            const sample_sequence_type& x,
            segmented_sequence_type& y
        ) const
        {
            y.clear();
            std::vector<unsigned long> labels;
            labeler.label_sequence(x, labels);

            if (feature_extractor::use_BIO_model)
            {
                // Convert from BIO tagging to the explicit segments representation.
                for (unsigned long i = 0; i < labels.size(); ++i)
                {
                    if (labels[i] == impl_ss::BEGIN)
                    {
                        const unsigned long begin = i;
                        ++i;
                        while (i < labels.size() && labels[i] == impl_ss::INSIDE)
                            ++i;

                        y.push_back(std::make_pair(begin, i));
                        --i;
                    }
                }
            }
            else
            {
                // Convert from BILOU tagging to the explicit segments representation.
                for (unsigned long i = 0; i < labels.size(); ++i)
                {
                    if (labels[i] == impl_ss::BEGIN)
                    {
                        const unsigned long begin = i;
                        ++i;
                        while (i < labels.size() && labels[i] == impl_ss::INSIDE)
                            ++i;

                        y.push_back(std::make_pair(begin, i+1));
                    }
                    else if (labels[i] == impl_ss::UNIT)
                    {
                        y.push_back(std::make_pair(i, i+1));
                    }
                }
            }
        }

        friend void serialize(const sequence_segmenter& item, std::ostream& out)
        {
            int version = 1;
            serialize(version, out);

            // Save these just so we can compare them when we deserialize and make
            // sure the feature_extractor being used is compatible with the model being
            // loaded.
            serialize(feature_extractor::use_BIO_model, out);
            serialize(feature_extractor::use_high_order_features, out);
            serialize(total_feature_vector_size(item.get_feature_extractor()), out);

            serialize(item.labeler, out);
        }

        friend void deserialize(sequence_segmenter& item, std::istream& in)
        {
            int version = 0;
            deserialize(version, in);
            if (version != 1)
                throw serialization_error("Unexpected version found while deserializing dlib::sequence_segmenter.");

            // Try to check if the saved model is compatible with the current feature
            // extractor.
            bool use_BIO_model, use_high_order_features;
            unsigned long dims;
            deserialize(use_BIO_model, in);
            deserialize(use_high_order_features, in);
            deserialize(dims, in);
            if (use_BIO_model != feature_extractor::use_BIO_model)
            {
                throw serialization_error("Incompatible feature extractor found while deserializing "
                    "dlib::sequence_segmenter. Wrong value of use_BIO_model.");
            }
            if (use_high_order_features != feature_extractor::use_high_order_features)
            {
                throw serialization_error("Incompatible feature extractor found while deserializing "
                    "dlib::sequence_segmenter. Wrong value of use_high_order_features.");
            }
            if (dims != total_feature_vector_size(item.get_feature_extractor()))
            {
                throw serialization_error("Incompatible feature extractor found while deserializing "
                    "dlib::sequence_segmenter. Wrong value of total_feature_vector_size().");
            }

            deserialize(item.labeler, in);
        }

    private:
        sequence_labeler<impl_ss::feature_extractor<feature_extractor> > labeler;
    };

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

}

#endif // DLIB_SEQUENCE_SeGMENTER_H___


