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

#include "modularity_clustering_abstract.h"
#include "../sparse_vector.h"
#include "../graph_utils/edge_list_graphs.h"
#include "../matrix.h"
#include "../rand.h"

namespace dlib
{

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

    namespace impl
    {
        inline double newman_cluster_split (
            dlib::rand& rnd,
            const std::vector<ordered_sample_pair>& edges,
            const matrix<double,0,1>& node_degrees, // k from the Newman paper
            const matrix<double,0,1>& Bdiag,        // diag(B) from the Newman paper
            const double& edge_sum,                 // m from the Newman paper
            matrix<double,0,1>& labels,
            const double eps,
            const unsigned long max_iterations
        )
        /*!
            requires
                - node_degrees.size() == max_index_plus_one(edges)
                - Bdiag.size() == max_index_plus_one(edges)
                - edges must be sorted according to order_by_index()
            ensures
                - This routine splits a graph into two subgraphs using the Newman 
                  clustering method.  
                - returns the modularity obtained when the graph is split according
                  to the contents of #labels. 
                - #labels.size() == node_degrees.size()
                - for all valid i: #labels(i) == -1 or +1
                - if (this function returns 0) then
                    - all the labels are equal, i.e. the graph is not split.
        !*/
        {
            // Scale epsilon so that it is relative to the expected value of an element of a
            // unit vector of length node_degrees.size().
            const double power_iter_eps = eps * std::sqrt(1.0/node_degrees.size());

            // Make a random unit vector and put in labels.
            labels.set_size(node_degrees.size());
            for (long i = 0; i < labels.size(); ++i)
                labels(i) = rnd.get_random_gaussian();
            labels /= length(labels);

            matrix<double,0,1> Bv, Bv_unit;

            // Do the power iteration for a while.
            double eig = -1;
            double offset = 0;
            while (eig < 0)
            {

                // any number larger than power_iter_eps
                double iteration_change = power_iter_eps*2+1; 
                for (unsigned long i = 0; i < max_iterations && iteration_change > power_iter_eps; ++i) 
                {
                    sparse_matrix_vector_multiply(edges, labels, Bv);
                    Bv -= dot(node_degrees, labels)/(2*edge_sum) * node_degrees;

                    if (offset != 0)
                    {
                        Bv -= offset*labels;
                    }


                    const double len = length(Bv);
                    if (len != 0)
                    {
                        Bv_unit = Bv/len;
                        iteration_change = max(abs(labels-Bv_unit));
                        labels.swap(Bv_unit);
                    }
                    else
                    {
                        // Had a bad time, pick another random vector and try it with the
                        // power iteration.
                        for (long i = 0; i < labels.size(); ++i)
                            labels(i) = rnd.get_random_gaussian();
                    }
                }

                eig = dot(Bv,labels);
                // we will repeat this loop if the largest eigenvalue is negative
                offset = eig;
            }


            for (long i = 0; i < labels.size(); ++i)
            {
                if (labels(i) > 0)
                    labels(i) = 1;
                else
                    labels(i) = -1;
            }


            // compute B*labels, store result in Bv.
            sparse_matrix_vector_multiply(edges, labels, Bv);
            Bv -= dot(node_degrees, labels)/(2*edge_sum) * node_degrees;

            // Do some label refinement.  In this step we swap labels if it
            // improves the modularity score.
            bool flipped_label = true;
            while(flipped_label)
            {
                flipped_label = false;
                unsigned long idx = 0;
                for (long i = 0; i < labels.size(); ++i)
                {
                    const double val = -2*labels(i);
                    const double increase = 4*Bdiag(i) + 2*val*Bv(i);

                    // if there is an increase in modularity for swapping this label
                    if (increase > 0)
                    {
                        labels(i) *= -1;
                        while (idx < edges.size() && edges[idx].index1() == (unsigned long)i)
                        {
                            const long j = edges[idx].index2();
                            Bv(j) += val*edges[idx].distance();
                            ++idx;
                        }

                        Bv -= (val*node_degrees(i)/(2*edge_sum))*node_degrees;

                        flipped_label = true;
                    }
                    else
                    {
                        while (idx < edges.size() && edges[idx].index1() == (unsigned long)i)
                        {
                            ++idx;
                        }
                    }
                }
            }


            const double modularity = dot(Bv, labels)/(4*edge_sum);

            return modularity;
        }

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

        inline unsigned long newman_cluster_helper (
            dlib::rand& rnd,
            const std::vector<ordered_sample_pair>& edges,
            const matrix<double,0,1>& node_degrees, // k from the Newman paper
            const matrix<double,0,1>& Bdiag,        // diag(B) from the Newman paper
            const double& edge_sum,                 // m from the Newman paper
            std::vector<unsigned long>& labels,
            double modularity_threshold,
            const double eps,
            const unsigned long max_iterations
        )
        /*!
            ensures
                - returns the number of clusters the data was split into
        !*/
        {
            matrix<double,0,1> l;
            const double modularity = newman_cluster_split(rnd,edges,node_degrees,Bdiag,edge_sum,l,eps,max_iterations);


            // We need to collapse the node index values down to contiguous values.  So
            // we use the following two vectors to contain the mappings from input index
            // values to their corresponding index values in each split.
            std::vector<unsigned long> left_idx_map(node_degrees.size());
            std::vector<unsigned long> right_idx_map(node_degrees.size());

            // figure out how many nodes went into each side of the split.
            unsigned long num_left_split = 0;
            unsigned long num_right_split = 0;
            for (long i = 0; i < l.size(); ++i)
            {
                if (l(i) > 0)
                {
                    left_idx_map[i] = num_left_split;
                    ++num_left_split;
                }
                else
                {
                    right_idx_map[i] = num_right_split;
                    ++num_right_split;
                }
            }

            // do a recursive split if it will improve the modularity.
            if (modularity > modularity_threshold && num_left_split > 0 && num_right_split > 0)
            {

                // split the node_degrees and Bdiag matrices into left and right split parts
                matrix<double,0,1> left_node_degrees(num_left_split);
                matrix<double,0,1> right_node_degrees(num_right_split);
                matrix<double,0,1> left_Bdiag(num_left_split);
                matrix<double,0,1> right_Bdiag(num_right_split);
                for (long i = 0; i < l.size(); ++i)
                {
                    if (l(i) > 0)
                    {
                        left_node_degrees(left_idx_map[i]) = node_degrees(i);
                        left_Bdiag(left_idx_map[i]) = Bdiag(i);
                    }
                    else
                    {
                        right_node_degrees(right_idx_map[i]) = node_degrees(i);
                        right_Bdiag(right_idx_map[i]) = Bdiag(i);
                    }
                }


                // put the edges from one side of the split into split_edges
                std::vector<ordered_sample_pair> split_edges;
                modularity_threshold = 0;
                for (unsigned long k = 0; k < edges.size(); ++k)
                {
                    const unsigned long i = edges[k].index1();
                    const unsigned long j = edges[k].index2();
                    const double d = edges[k].distance();
                    if (l(i) > 0 && l(j) > 0)
                    {
                        split_edges.push_back(ordered_sample_pair(left_idx_map[i], left_idx_map[j], d));
                        modularity_threshold += d;
                    }
                }
                modularity_threshold -= sum(left_node_degrees*sum(left_node_degrees))/(2*edge_sum);
                modularity_threshold /= 4*edge_sum;

                unsigned long num_left_clusters;
                std::vector<unsigned long> left_labels;
                num_left_clusters = newman_cluster_helper(rnd,split_edges,left_node_degrees,left_Bdiag,
                                                          edge_sum,left_labels,modularity_threshold,
                                                          eps, max_iterations);

                // now load the other side into split_edges and cluster it as well
                split_edges.clear();
                modularity_threshold = 0;
                for (unsigned long k = 0; k < edges.size(); ++k)
                {
                    const unsigned long i = edges[k].index1();
                    const unsigned long j = edges[k].index2();
                    const double d = edges[k].distance();
                    if (l(i) < 0 && l(j) < 0)
                    {
                        split_edges.push_back(ordered_sample_pair(right_idx_map[i], right_idx_map[j], d));
                        modularity_threshold += d;
                    }
                }
                modularity_threshold -= sum(right_node_degrees*sum(right_node_degrees))/(2*edge_sum);
                modularity_threshold /= 4*edge_sum;

                unsigned long num_right_clusters;
                std::vector<unsigned long> right_labels;
                num_right_clusters = newman_cluster_helper(rnd,split_edges,right_node_degrees,right_Bdiag,
                                                           edge_sum,right_labels,modularity_threshold,
                                                           eps, max_iterations);

                // Now merge the labels from the two splits.
                labels.resize(node_degrees.size());
                for (unsigned long i = 0; i < labels.size(); ++i)
                {
                    // if this node was in the left split
                    if (l(i) > 0)
                    {
                        labels[i] = left_labels[left_idx_map[i]];
                    }
                    else // if this node was in the right split
                    {
                        labels[i] = right_labels[right_idx_map[i]] + num_left_clusters;
                    }
                }


                return num_left_clusters + num_right_clusters;
            }
            else
            {
                labels.assign(node_degrees.size(),0);
                return 1;
            }

        }
    }

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

    inline unsigned long newman_cluster (
        const std::vector<ordered_sample_pair>& edges,
        std::vector<unsigned long>& labels,
        const double eps = 1e-4,
        const unsigned long max_iterations = 2000
    )
    {
        // make sure requires clause is not broken
        DLIB_ASSERT(is_ordered_by_index(edges),
                    "\t unsigned long newman_cluster()"
                    << "\n\t Invalid inputs were given to this function"
        );

        labels.clear();
        if (edges.size() == 0)
            return 0;

        const unsigned long num_nodes = max_index_plus_one(edges);

        // compute the node_degrees vector, edge_sum value, and diag(B).
        matrix<double,0,1> node_degrees(num_nodes);
        matrix<double,0,1> Bdiag(num_nodes);
        Bdiag = 0;
        double edge_sum = 0;
        node_degrees = 0;
        for (unsigned long i = 0; i < edges.size(); ++i)
        {
            node_degrees(edges[i].index1()) += edges[i].distance();
            edge_sum += edges[i].distance();
            if (edges[i].index1() == edges[i].index2())
                Bdiag(edges[i].index1()) += edges[i].distance();
        }
        edge_sum /= 2;
        Bdiag -= squared(node_degrees)/(2*edge_sum);


        dlib::rand rnd;
        return impl::newman_cluster_helper(rnd,edges,node_degrees,Bdiag,edge_sum,labels,0,eps,max_iterations);
    }

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

    inline unsigned long newman_cluster (
        const std::vector<sample_pair>& edges,
        std::vector<unsigned long>& labels,
        const double eps = 1e-4,
        const unsigned long max_iterations = 2000
    )
    {
        std::vector<ordered_sample_pair> oedges;
        convert_unordered_to_ordered(edges, oedges);
        std::sort(oedges.begin(), oedges.end(), &order_by_index<ordered_sample_pair>);

        return newman_cluster(oedges, labels, eps, max_iterations);
    }

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

    namespace impl
    {
        inline std::vector<unsigned long> remap_labels (
            const std::vector<unsigned long>& labels,
            unsigned long& num_labels
        )
        /*!
            ensures
                - This function takes labels and produces a mapping which maps elements of
                  labels into the most compact range in [0, max] as possible.  In particular,
                  there won't be any unused integers in the mapped range.
                - #num_labels == the number of distinct values in labels.
                - returns a vector V such that:
                    - V.size() == labels.size()
                    - max(mat(V))+1 == num_labels.
                    - for all valid i,j:
                        - if (labels[i] == labels[j]) then
                            - V[i] == V[j]
                        - else
                            - V[i] != V[j]
        !*/
        {
            std::map<unsigned long, unsigned long> temp;
            for (unsigned long i = 0; i < labels.size(); ++i)
            {
                if (temp.count(labels[i]) == 0)
                {
                    const unsigned long next = temp.size();
                    temp[labels[i]] = next;
                }
            }

            num_labels = temp.size();

            std::vector<unsigned long> result(labels.size());
            for (unsigned long i = 0; i < labels.size(); ++i)
            {
                result[i] = temp[labels[i]];
            }
            return result;
        }
    }

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

    inline double modularity (
        const std::vector<sample_pair>& edges,
        const std::vector<unsigned long>& labels
    )
    {
        const unsigned long num_nodes = max_index_plus_one(edges);
        // make sure requires clause is not broken
        DLIB_ASSERT(labels.size() == num_nodes,
                    "\t double modularity()"
                    << "\n\t Invalid inputs were given to this function"
        );

        unsigned long num_labels;
        const std::vector<unsigned long>& labels_ = dlib::impl::remap_labels(labels,num_labels);

        std::vector<double> cluster_sums(num_labels,0);
        std::vector<double> k(num_nodes,0);

        double Q = 0;
        double m = 0;
        for (unsigned long i = 0; i < edges.size(); ++i)
        {
            const unsigned long n1 = edges[i].index1();
            const unsigned long n2 = edges[i].index2();
            k[n1] += edges[i].distance();
            if (n1 != n2)
                k[n2] += edges[i].distance();

            if (n1 != n2)
                m += edges[i].distance();
            else
                m += edges[i].distance()/2;

            if (labels_[n1] == labels_[n2])
            {
                if (n1 != n2)
                    Q += 2*edges[i].distance();
                else
                    Q += edges[i].distance();
            }
        }

        if (m == 0)
            return 0;

        for (unsigned long i = 0; i < labels_.size(); ++i)
        {
            cluster_sums[labels_[i]] += k[i];
        }

        for (unsigned long i = 0; i < labels_.size(); ++i)
        {
            Q -= k[i]*cluster_sums[labels_[i]]/(2*m);
        }

        return 1.0/(2*m)*Q;
    }

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

    inline double modularity (
        const std::vector<ordered_sample_pair>& edges,
        const std::vector<unsigned long>& labels
    )
    {
        const unsigned long num_nodes = max_index_plus_one(edges);
        // make sure requires clause is not broken
        DLIB_ASSERT(labels.size() == num_nodes,
                    "\t double modularity()"
                    << "\n\t Invalid inputs were given to this function"
        );


        unsigned long num_labels;
        const std::vector<unsigned long>& labels_ = dlib::impl::remap_labels(labels,num_labels);

        std::vector<double> cluster_sums(num_labels,0);
        std::vector<double> k(num_nodes,0);

        double Q = 0;
        double m = 0;
        for (unsigned long i = 0; i < edges.size(); ++i)
        {
            const unsigned long n1 = edges[i].index1();
            const unsigned long n2 = edges[i].index2();
            k[n1] += edges[i].distance();
            m += edges[i].distance();
            if (labels_[n1] == labels_[n2])
            {
                Q += edges[i].distance();
            }
        }

        if (m == 0)
            return 0;

        for (unsigned long i = 0; i < labels_.size(); ++i)
        {
            cluster_sums[labels_[i]] += k[i];
        }

        for (unsigned long i = 0; i < labels_.size(); ++i)
        {
            Q -= k[i]*cluster_sums[labels_[i]]/m;
        }

        return 1.0/m*Q;
    }

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

}

#endif // DLIB_MODULARITY_ClUSTERING__H__