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

    #include "segment_image_abstract.h"
    #include "../algs.h"
    #include <vector>
    #include "../geometry.h"
    #include "../disjoint_subsets.h"
    #include "../set.h"

    namespace dlib
    {

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

        namespace impl
        {
            template <typename T>
            inline T edge_diff_uint(
                const T& a,
                const T& b
            )
            {
                if (a > b)
                    return a - b;
                else
                    return b - a;
            }

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

            template <typename T, typename enabled = void>
            struct edge_diff_funct
            {
                typedef double diff_type;

                template <typename pixel_type>
                double operator()(
                    const pixel_type& a,
                    const pixel_type& b
                ) const
                {
                    return length(pixel_to_vector<double>(a) - pixel_to_vector<double>(b));
                }
            };

            template <>
            struct edge_diff_funct<uint8,void>
            {
                typedef uint8 diff_type;
                uint8 operator()( const uint8& a, const uint8& b) const { return edge_diff_uint(a,b); }
            };

            template <>
            struct edge_diff_funct<uint16,void>
            {
                typedef uint16 diff_type;
                uint16 operator()( const uint16& a, const uint16& b) const { return edge_diff_uint(a,b); }
            };

            template <>
            struct edge_diff_funct<uint32,void>
            {
                typedef uint32 diff_type;
                uint32 operator()( const uint32& a, const uint32& b) const { return edge_diff_uint(a,b); }
            };

            template <>
            struct edge_diff_funct<double,void>
            {
                typedef double diff_type;
                double operator()( const double& a, const double& b) const { return std::abs(a-b); }
            };

            template <typename T>
            struct edge_diff_funct<T, typename enable_if<is_matrix<T> >::type>
            {
                typedef double diff_type;
                double operator()(
                    const T& a,
                    const T& b
                ) const
                {
                    return length(a-b);
                }
            };

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

            template <typename T>
            struct graph_image_segmentation_data_T
            {
                graph_image_segmentation_data_T() : component_size(1), internal_diff(0) {}
                unsigned long component_size;
                T internal_diff;
            };

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

            template <typename T>
            struct segment_image_edge_data_T
            {
                segment_image_edge_data_T (){}

                segment_image_edge_data_T (
                    const rectangle& rect,
                    const point& p1,
                    const point& p2,
                    const T& diff_
                ) :
                    idx1(p1.y()*rect.width() + p1.x()),
                    idx2(p2.y()*rect.width() + p2.x()),
                    diff(diff_)
                {}

                bool operator<(const segment_image_edge_data_T& item) const
                { return diff < item.diff; }

                unsigned long idx1;
                unsigned long idx2;
                T diff;
            };

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

            template <typename image_view_type>
            struct uint8_or_uint16_pixels
            {
                typedef typename image_view_type::pixel_type pixel_type;
                const static bool value = is_same_type<pixel_type,uint8>::value ||
                    is_same_type<pixel_type,uint16>::value;
            };

            // This is an overload of get_pixel_edges() that is optimized to segment images
            // with 8bit or 16bit  pixels very quickly.  We do this by using a radix sort
            // instead of quicksort.
            template <typename in_image_type, typename T>
            typename enable_if<uint8_or_uint16_pixels<in_image_type> >::type
            get_pixel_edges (
                const in_image_type& in_img,
                std::vector<segment_image_edge_data_T<T> >& sorted_edges
            )
            {
                typedef typename in_image_type::pixel_type ptype;
                typedef T diff_type;
                std::vector<unsigned long> counts(std::numeric_limits<ptype>::max()+1, 0);

                edge_diff_funct<ptype> edge_diff;

                border_enumerator be(get_rect(in_img), 1);
                // we are going to do a radix sort on the edge weights.  So the first step
                // is to accumulate them into count.
                const rectangle area = get_rect(in_img);
                while (be.move_next())
                {
                    const long r = be.element().y();
                    const long c = be.element().x();
                    const ptype pix = in_img[r][c];
                    if (area.contains(c-1,r))   counts[edge_diff(pix, in_img[r  ][c-1])] += 1;
                    if (area.contains(c+1,r))   counts[edge_diff(pix, in_img[r  ][c+1])] += 1;
                    if (area.contains(c  ,r-1)) counts[edge_diff(pix, in_img[r-1][c  ])] += 1;
                    if (area.contains(c  ,r+1)) counts[edge_diff(pix, in_img[r+1][c  ])] += 1;
                }
                for (long r = 1; r+1 < in_img.nr(); ++r)
                {
                    for (long c = 1; c+1 < in_img.nc(); ++c)
                    {
                        const ptype pix = in_img[r][c];
                        counts[edge_diff(pix, in_img[r-1][c+1])] += 1;
                        counts[edge_diff(pix, in_img[r  ][c+1])] += 1;
                        counts[edge_diff(pix, in_img[r+1][c  ])] += 1;
                        counts[edge_diff(pix, in_img[r+1][c+1])] += 1;
                    }
                }

                const unsigned long num_edges = shrink_rect(area,1).area()*4 + in_img.nr()*2*3 - 4 + (in_img.nc()-2)*2*3;
                typedef segment_image_edge_data_T<T> segment_image_edge_data;
                sorted_edges.resize(num_edges);

                // integrate counts.  The idea is to have sorted_edges[counts[i]] be the location that edges
                // with an edge_diff of i go.  So counts[0] == 0, counts[1] == number of 0 edge diff edges, etc.
                unsigned long prev = counts[0];
                for (unsigned long i = 1; i < counts.size(); ++i)
                {
                    const unsigned long temp = counts[i];
                    counts[i] += counts[i-1];
                    counts[i-1] -= prev;
                    prev = temp;
                }
                counts[counts.size()-1] -= prev;

                // now build a sorted list of all the edges
                be.reset();
                while(be.move_next())
                {
                    const point p = be.element();
                    const long r = p.y();
                    const long c = p.x();
                    const ptype pix = in_img[r][c];
                    if (area.contains(c-1,r))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r  ][c-1]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c-1,r),diff);
                    }

                    if (area.contains(c+1,r))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r  ][c+1]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r),diff);
                    }

                    if (area.contains(c  ,r-1))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r-1][c  ]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c  ,r-1),diff);
                    }

                    if (area.contains(c  ,r+1))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r+1][c  ]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c  ,r+1),diff);
                    }
                }
                // same thing as the above loop but now we do it on the interior of the image and therefore
                // don't have to include the boundary checking if statements used above.
                for (long r = 1; r+1 < in_img.nr(); ++r)
                {
                    for (long c = 1; c+1 < in_img.nc(); ++c)
                    {
                        const point p(c,r);
                        const ptype pix = in_img[r][c];
                        diff_type diff;

                        diff = edge_diff(pix, in_img[r  ][c+1]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r),diff);
                        diff = edge_diff(pix, in_img[r-1][c+1]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r-1),diff);
                        diff = edge_diff(pix, in_img[r+1][c+1]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c+1,r+1),diff);
                        diff = edge_diff(pix, in_img[r+1][c  ]);
                        sorted_edges[counts[diff]++] = segment_image_edge_data(area,p,point(c  ,r+1),diff);
                    }
                }
            }

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

            // This is the general purpose version of get_pixel_edges().  It handles all pixel types.
            template <typename in_image_type, typename T>
            typename disable_if<uint8_or_uint16_pixels<in_image_type> >::type
            get_pixel_edges (
                const in_image_type& in_img,
                std::vector<segment_image_edge_data_T<T> >& sorted_edges
            )
            {
                const rectangle area = get_rect(in_img);
                sorted_edges.reserve(area.area()*4);

                typedef typename in_image_type::pixel_type ptype;
                edge_diff_funct<ptype> edge_diff;
                typedef T diff_type;
                typedef segment_image_edge_data_T<T> segment_image_edge_data;

                border_enumerator be(get_rect(in_img), 1);

                // now build a sorted list of all the edges
                be.reset();
                while(be.move_next())
                {
                    const point p = be.element();
                    const long r = p.y();
                    const long c = p.x();
                    const ptype& pix = in_img[r][c];
                    if (area.contains(c-1,r))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r  ][c-1]);
                        sorted_edges.push_back(segment_image_edge_data(area,p,point(c-1,r),diff));
                    }

                    if (area.contains(c+1,r))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r  ][c+1]);
                        sorted_edges.push_back(segment_image_edge_data(area,p,point(c+1,r),diff));
                    }

                    if (area.contains(c  ,r-1))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r-1][c  ]);
                        sorted_edges.push_back( segment_image_edge_data(area,p,point(c  ,r-1),diff));
                    }
                    if (area.contains(c  ,r+1))
                    {
                        const diff_type diff = edge_diff(pix, in_img[r+1][c  ]);
                        sorted_edges.push_back( segment_image_edge_data(area,p,point(c  ,r+1),diff));
                    }
                }
                // same thing as the above loop but now we do it on the interior of the image and therefore
                // don't have to include the boundary checking if statements used above.
                for (long r = 1; r+1 < in_img.nr(); ++r)
                {
                    for (long c = 1; c+1 < in_img.nc(); ++c)
                    {
                        const point p(c,r);
                        const ptype& pix = in_img[r][c];
                        diff_type diff;

                        diff = edge_diff(pix, in_img[r  ][c+1]);
                        sorted_edges.push_back( segment_image_edge_data(area,p,point(c+1,r),diff));
                        diff = edge_diff(pix, in_img[r+1][c+1]);
                        sorted_edges.push_back( segment_image_edge_data(area,p,point(c+1,r+1),diff));
                        diff = edge_diff(pix, in_img[r+1][c  ]);
                        sorted_edges.push_back( segment_image_edge_data(area,p,point(c  ,r+1),diff));
                        diff = edge_diff(pix, in_img[r-1][c+1]);
                        sorted_edges.push_back( segment_image_edge_data(area,p,point(c+1,r-1),diff));
                    }
                }

                std::sort(sorted_edges.begin(), sorted_edges.end());

            }

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

        } // end of namespace impl

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

        template <
            typename in_image_type,
            typename out_image_type
            >
        void segment_image (
            const in_image_type& in_img_,
            out_image_type& out_img_,
            const double k = 200,
            const unsigned long min_size = 10
        )
        {
            using namespace dlib::impl;
            typedef typename image_traits<in_image_type>::pixel_type ptype;
            typedef typename edge_diff_funct<ptype>::diff_type diff_type;

            // make sure requires clause is not broken
            DLIB_ASSERT(is_same_object(in_img_, out_img_) == false,
                "\t void segment_image()"
                << "\n\t The input images can't be the same object."
                );

            COMPILE_TIME_ASSERT(is_unsigned_type<typename image_traits<out_image_type>::pixel_type>::value);

            const_image_view<in_image_type> in_img(in_img_);
            image_view<out_image_type> out_img(out_img_);

            out_img.set_size(in_img.nr(), in_img.nc());
            // don't bother doing anything if the image is too small
            if (in_img.nr() < 2 || in_img.nc() < 2)
            {
                assign_all_pixels(out_img,0);
                return;
            }

            disjoint_subsets sets;
            sets.set_size(in_img.size());

            std::vector<segment_image_edge_data_T<diff_type> > sorted_edges;
            get_pixel_edges(in_img, sorted_edges);

            std::vector<graph_image_segmentation_data_T<diff_type> > data(in_img.size());

            // now start connecting blobs together to make a minimum spanning tree.
            for (unsigned long i = 0; i < sorted_edges.size(); ++i)
            {
                const unsigned long idx1 = sorted_edges[i].idx1;
                const unsigned long idx2 = sorted_edges[i].idx2;

                unsigned long set1 = sets.find_set(idx1);
                unsigned long set2 = sets.find_set(idx2);
                if (set1 != set2)
                {
                    const diff_type diff = sorted_edges[i].diff;
                    const diff_type tau1 = static_cast<diff_type>(k/data[set1].component_size);
                    const diff_type tau2 = static_cast<diff_type>(k/data[set2].component_size);

                    const diff_type mint = std::min(data[set1].internal_diff + tau1,
                                                    data[set2].internal_diff + tau2);
                    if (diff <= mint)
                    {
                        const unsigned long new_set = sets.merge_sets(set1, set2);
                        data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                        data[new_set].internal_diff = diff;
                    }
                }
            }

            // now merge any really small blobs
            if (min_size != 0)
            {
                for (unsigned long i = 0; i < sorted_edges.size(); ++i)
                {
                    const unsigned long idx1 = sorted_edges[i].idx1;
                    const unsigned long idx2 = sorted_edges[i].idx2;

                    unsigned long set1 = sets.find_set(idx1);
                    unsigned long set2 = sets.find_set(idx2);
                    if (set1 != set2 && (data[set1].component_size < min_size || data[set2].component_size < min_size))
                    {
                        const unsigned long new_set = sets.merge_sets(set1, set2);
                        data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                        //data[new_set].internal_diff = sorted_edges[i].diff;
                    }
                }
            }

            unsigned long idx = 0;
            for (long r = 0; r < out_img.nr(); ++r)
            {
                for (long c = 0; c < out_img.nc(); ++c)
                {
                    out_img[r][c] = sets.find_set(idx++);
                }
            }
        }

    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------
    //                     Candidate object location generation code.
    // ----------------------------------------------------------------------------------------
    // ----------------------------------------------------------------------------------------

        namespace impl
        {
            struct edge_data
            {
                double edge_diff;
                unsigned long set1;
                unsigned long set2;
                bool operator<(const edge_data& item) const
                {
                    return edge_diff < item.edge_diff;
                }
            };

            template <
                typename in_image_type,
                typename diff_type
                >
            void find_basic_candidate_object_locations (
                const in_image_type& in_img,
                const std::vector<dlib::impl::segment_image_edge_data_T<diff_type> >& sorted_edges,
                std::vector<rectangle>& out_rects,
                std::vector<edge_data>& edges,
                const double k,
                const unsigned long min_size
            )
            {
                using namespace dlib::impl;

                std::vector<dlib::impl::segment_image_edge_data_T<diff_type> > rejected_edges;
                rejected_edges.reserve(sorted_edges.size());

                out_rects.clear();
                edges.clear();

                // don't bother doing anything if the image is too small
                if (in_img.nr() < 2 || in_img.nc() < 2)
                {
                    return;
                }

                disjoint_subsets sets;
                sets.set_size(in_img.size());

                std::vector<graph_image_segmentation_data_T<diff_type> > data(in_img.size());

                std::pair<unsigned long,unsigned long> last_blob_edge(std::numeric_limits<unsigned long>::max(),
                                                                      std::numeric_limits<unsigned long>::max());;
                // now start connecting blobs together to make a minimum spanning tree.
                for (unsigned long i = 0; i < sorted_edges.size(); ++i)
                {
                    const unsigned long idx1 = sorted_edges[i].idx1;
                    const unsigned long idx2 = sorted_edges[i].idx2;

                    unsigned long set1 = sets.find_set(idx1);
                    unsigned long set2 = sets.find_set(idx2);
                    if (set1 != set2)
                    {
                        const diff_type diff = sorted_edges[i].diff;
                        const diff_type tau1 = static_cast<diff_type>(k/data[set1].component_size);
                        const diff_type tau2 = static_cast<diff_type>(k/data[set2].component_size);

                        const diff_type mint = std::min(data[set1].internal_diff + tau1,
                            data[set2].internal_diff + tau2);
                        if (diff <= mint)
                        {
                            const unsigned long new_set = sets.merge_sets(set1, set2);
                            data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                            data[new_set].internal_diff = diff;
                        }
                        else
                        {
                            // Don't bother keeping multiple edges from the same pair of blobs, we
                            // only need one for what we will do later.
                            if (std::make_pair(set1,set2) != last_blob_edge)
                            {
                                segment_image_edge_data_T<diff_type> temp = sorted_edges[i];
                                temp.idx1 = set1;
                                temp.idx2 = set2;
                                rejected_edges.push_back(temp);
                                last_blob_edge = std::make_pair(set1,set2);
                            }
                        }
                    }
                }

                // merge small blobs
                for (unsigned long i = 0; i < rejected_edges.size(); ++i)
                {
                    const unsigned long idx1 = rejected_edges[i].idx1;
                    const unsigned long idx2 = rejected_edges[i].idx2;

                    unsigned long set1 = sets.find_set(idx1);
                    unsigned long set2 = sets.find_set(idx2);
                    rejected_edges[i].idx1 = set1;
                    rejected_edges[i].idx2 = set2;
                    if (set1 != set2 && (data[set1].component_size < min_size || data[set2].component_size < min_size))
                    {
                        const unsigned long new_set = sets.merge_sets(set1, set2);
                        data[new_set].component_size = data[set1].component_size + data[set2].component_size;
                        data[new_set].internal_diff = rejected_edges[i].diff;
                    }
                }

                // find bounding boxes of each blob
                std::map<unsigned long, rectangle> boxes;
                std::map<unsigned long, unsigned long> box_id_map;
                unsigned long idx = 0;
                for (long r = 0; r < in_img.nr(); ++r)
                {
                    for (long c = 0; c < in_img.nc(); ++c)
                    {
                        const unsigned long id = sets.find_set(idx++);
                        // Accumulate the current point into its box and if it is the first point
                        // in the box then also record the id number for this box.
                        if ((boxes[id] += point(c,r)).area() == 1)
                            box_id_map[id] = boxes.size()-1;
                    }
                }

                // copy boxes into out_rects
                out_rects.resize(boxes.size());
                for (std::map<unsigned long,rectangle>::iterator i = boxes.begin(); i != boxes.end(); ++i)
                {
                    out_rects[box_id_map[i->first]] = i->second;
                }

                // Now find the edges between the boxes
                typedef dlib::memory_manager<char>::kernel_2c mm_type;
                dlib::set<std::pair<unsigned long, unsigned long>, mm_type>::kernel_1a neighbors_final;
                for (unsigned long i = 0; i < rejected_edges.size(); ++i)
                {
                    const unsigned long idx1 = rejected_edges[i].idx1;
                    const unsigned long idx2 = rejected_edges[i].idx2;

                    unsigned long set1 = sets.find_set(idx1);
                    unsigned long set2 = sets.find_set(idx2);
                    if (set1 != set2)
                    {
                        std::pair<unsigned long, unsigned long> p = std::make_pair(set1,set2);
                        if (!neighbors_final.is_member(p))
                        {
                            neighbors_final.add(p);

                            edge_data temp;
                            const diff_type mint = std::min(data[set1].internal_diff ,
                                                            data[set2].internal_diff );
                            temp.edge_diff = rejected_edges[i].diff - mint;
                            temp.set1 = box_id_map[set1];
                            temp.set2 = box_id_map[set2];
                            edges.push_back(temp);
                        }
                    }
                }

                std::sort(edges.begin(), edges.end());
            }
        } // end namespace impl

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

        template <typename alloc>
        void remove_duplicates (
            std::vector<rectangle,alloc>& rects
        )
        {
            std::sort(rects.begin(), rects.end(), std::less<rectangle>());
            unsigned long num_unique = 1;
            for (unsigned long i = 1; i < rects.size(); ++i)
            {
                if (rects[i] != rects[i-1])
                {
                    rects[num_unique++] = rects[i];
                }
            }
            if (rects.size() != 0)
                rects.resize(num_unique);
        }

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

        template <
            typename in_image_type,
            typename EXP
            >
        void find_candidate_object_locations (
            const in_image_type& in_img_,
            std::vector<rectangle>& rects,
            const matrix_exp<EXP>& kvals,
            const unsigned long min_size = 20,
            const unsigned long max_merging_iterations = 50
        )
        {
            // make sure requires clause is not broken
            DLIB_ASSERT(is_vector(kvals) && kvals.size() > 0,
                "\t void find_candidate_object_locations()"
                << "\n\t Invalid inputs were given to this function."
                << "\n\t is_vector(kvals): " << is_vector(kvals)
                << "\n\t kvals.size():     " << kvals.size()
                );

            typedef dlib::memory_manager<char>::kernel_2c mm_type;
            typedef dlib::set<rectangle, mm_type>::kernel_1a set_of_rects;

            using namespace dlib::impl;
            typedef typename image_traits<in_image_type>::pixel_type ptype;
            typedef typename edge_diff_funct<ptype>::diff_type diff_type;

            const_image_view<in_image_type> in_img(in_img_);

            // don't bother doing anything if the image is too small
            if (in_img.nr() < 2 || in_img.nc() < 2)
            {
                return;
            }

            std::vector<edge_data> edges;
            std::vector<rectangle> working_rects;
            std::vector<segment_image_edge_data_T<diff_type> > sorted_edges;
            get_pixel_edges(in_img, sorted_edges);

            disjoint_subsets sets;

            for (long j = 0; j < kvals.size(); ++j)
            {
                const double k = kvals(j);

                find_basic_candidate_object_locations(in_img, sorted_edges, working_rects, edges, k, min_size);
                rects.insert(rects.end(), working_rects.begin(), working_rects.end());

                // Now iteratively merge all the rectangles we have and record the results.
                // Note that, unlike what is described in the paper
                //    Segmentation as Selective Search for Object Recognition" by Koen E. A. van de Sande, et al.
                // we don't use any kind of histogram/SIFT like thing to order the edges
                // between the blobs.  Here we simply order by the pixel difference value.
                // Additionally, note that we keep progressively merging boxes in the outer
                // loop rather than performing just a single iteration as indicated in the
                // paper.
                set_of_rects detected_rects;
                bool did_merge = true;
                for (unsigned long iter = 0; did_merge && iter < max_merging_iterations; ++iter)
                {
                    did_merge = false;
                    sets.clear();
                    sets.set_size(working_rects.size());

                    // recursively merge neighboring blobs until we have merged everything
                    for (unsigned long i = 0; i < edges.size(); ++i)
                    {
                        edge_data temp = edges[i];

                        temp.set1 = sets.find_set(temp.set1);
                        temp.set2 = sets.find_set(temp.set2);
                        if (temp.set1 != temp.set2)
                        {
                            rectangle merged_rect = working_rects[temp.set1] + working_rects[temp.set2];
                            // Skip merging this pair of blobs if it was merged in a previous
                            // iteration.  Doing this lets us consider other possible blob
                            // merges.
                            if (!detected_rects.is_member(merged_rect))
                            {
                                const unsigned long new_set = sets.merge_sets(temp.set1, temp.set2);
                                rects.push_back(merged_rect);
                                working_rects[new_set] = merged_rect;
                                did_merge = true;
                                detected_rects.add(merged_rect);
                            }
                        }
                    }
                }
            }

            remove_duplicates(rects);
        }

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

        template <
            typename in_image_type
            >
        void find_candidate_object_locations (
            const in_image_type& in_img,
            std::vector<rectangle>& rects
        )
        {
            find_candidate_object_locations(in_img, rects, linspace(50, 200, 3));
        }

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

    }

    #endif // DLIB_SEGMENT_ImAGE_Hh_