/* ========================================================= */
/* Louvain algorithm GraphHPC-2016 reference implementation  */
/* ========================================================= */
#include "defs.h"

struct sigma_t {
    weight_t in, tot;

    sigma_t()
    {
        in = 0;
        tot = 0;
    }
};

typedef map<vertex_id_t, sigma_t> sigma_map_t;

weight_t total_edge_weight(map_graph_t &graph)
{
    weight_t m = 0;
    for (map_graph_t::iterator i = graph.begin(); i != graph.end(); i++) {
        for (map_edges_t::iterator j = i->second.begin(); j != i->second.end(); j++) {
            m += j->second;
        }
    }

    return m;
}

weight_t modularity(map_graph_t &graph, vector<vertex_id_t> component_id)
{
    // calculate sigma_in and sigma_t for each component
    sigma_map_t sigma;
    for (map_graph_t::iterator i = graph.begin(); i != graph.end(); i++) {
        vertex_id_t cu = component_id[i->first];
        for (map_edges_t::iterator j = i->second.begin(); j != i->second.end(); j++) {
            vertex_id_t cv = component_id[j->first];
            weight_t weight = j->second;
            if (cu == cv) {
                sigma[cu].in += weight;
                sigma[cu].tot += weight;
            } else {
                sigma[cu].tot += weight;
                sigma[cv].tot += weight;
            }
        }
    }

    // calculate modularity
    weight_t m = total_edge_weight(graph);
    weight_t q = 0;
    for (sigma_map_t::iterator i = sigma.begin(); i != sigma.end(); i++) {
        weight_t sigma_in = i->second.in, sigma_tot = i->second.tot;
        q += sigma_in / m - pow(sigma_tot / m, 2);
    }

    return q;
}

void init_component_id(map_graph_t &graph, vector<vertex_id_t> &component_id)
{
    for (map_graph_t::iterator i = graph.begin(); i != graph.end(); i++) {
        component_id[i->first] = i->first;
    }
}

void calculate_incident_edge_weights(map_graph_t &graph, vector<weight_t> &weights)
{
    for (map_graph_t::iterator i = graph.begin(); i != graph.end(); i++) {
        vertex_id_t u = i->first;
        weights[u] = 0;
        for (map_edges_t::iterator j = i->second.begin(); j != i->second.end(); j++) {
            vertex_id_t v = j->first;
            weight_t weight = j->second;
            weights[u] += (u == v ? 1 : 2) * weight;
        }
    }
}

void compress(map_graph_t &graph, vector<vertex_id_t> &component_id)
{
    map_graph_t new_graph;
    for (map_graph_t::iterator i = graph.begin(); i != graph.end(); i++) {
        vertex_id_t cu = component_id[i->first];
        for (map_edges_t::iterator j = i->second.begin(); j != i->second.end(); j++) {
            vertex_id_t cv = component_id[j->first];
            weight_t weight = j->second;
            new_graph[cu][cv] += weight;
        }
    }

    graph = new_graph;
}

void translate(vector<vertex_id_t> &component_id, vector<vertex_id_t> &new_component_id)
{
    for (vector<vertex_id_t>::iterator i = component_id.begin(); i != component_id.end(); i++) {
        *i = new_component_id[*i];
    }
}

weight_t delta_q(weight_t sign, weight_t wu, weight_t wuc, weight_t sigma_tot, weight_t m)
{
    return sign * wuc / m - (pow(sigma_tot + sign * (wu - wuc), 2) - pow(sigma_tot, 2)) / pow(m, 2);
}

/* FIXME */
void build_map_graph_structure(graph_t &graph, map_graph_t &map_graph)
{
    for (vertex_id_t u = 0; u < graph.n; u++) {
        for (edge_id_t e = graph.rowsIndices[u]; e < graph.rowsIndices[u + 1]; e++) {
            vertex_id_t v = graph.endV[e];
            weight_t weight = graph.weights[e];
            map_graph[u][v] += weight;
        }
    }
}

/* FIXME */
void community_detection(size_t n, map_graph_t &graph, vector<vertex_id_t> &component_id)
{
    // initially each vertex is in its own component
    init_component_id(graph, component_id);

    // get the sum of all edge weights (NB: not divided by 2!)
    weight_t m = total_edge_weight(graph);

    // calculate initial sigma_total
    vector<weight_t> sigma_tot(n, 0), w(n, 0);
    calculate_incident_edge_weights(graph, sigma_tot);

    while (true) {
        // initialize component mapping for current graph
        vector<vertex_id_t> new_component_id(component_id);
        init_component_id(graph, new_component_id);

        // calculate weights of vertices in current graph
        calculate_incident_edge_weights(graph, w);

        // stage: try to move each vertex into neighbour's community

        //cout << "stage: modularity = " << modularity(graph, new_component_id) << endl;

        bool more, changed = false;
        do {
            more = false;

            for (map_graph_t::iterator i = graph.begin(); i != graph.end(); i++) {
                // calculate weights w_{u->c}
                typedef map<vertex_id_t, weight_t> wu_t;
                wu_t wu;
                vertex_id_t u = i->first;
                for (map_edges_t::iterator j = i->second.begin(); j != i->second.end(); j++) {
                    vertex_id_t v = j->first;
                    if (u != v) {
                        wu[new_component_id[v]] += 2 * j->second;
					}
                }

                vertex_id_t c1 = new_component_id[u]; // current component
                weight_t wuc1 = wu[c1]; // weight of connection to current component

                // find best candidate component for moving vertex u
                weight_t best_delta = 0;
                vertex_id_t best_component = c1;
                for (wu_t::iterator j = wu.begin(); j != wu.end(); j++) {
                    vertex_id_t c2 = j->first; // candidate component
                    if (c2 != c1) {
                        weight_t wuc2 = j->second; // weight of connection to candidate component
                        weight_t delta = // gain in modularity if u is moved from c1 to c2
                                delta_q(-1, w[u], wuc1, sigma_tot[c1], m) +
                                delta_q(+1, w[u], wuc2, sigma_tot[c2], m);
                        if (delta > best_delta) {
                            best_delta = delta;
                            best_component = c2;
                        }
                    }
                }

                // move u to another component if there is a positive gain
                vertex_id_t c2 = best_component;
                if (c1 != c2) {
                    sigma_tot[c1] -= w[u] - wuc1;
                    sigma_tot[c2] += w[u] - wu[c2];
                    new_component_id[u] = c2;
                    more = true;
                }
            }

            //cout << "step: modularity = " << modularity(graph, new_component_id) << endl;

            changed |= more;
        } while (more);
		
        if (!changed) {
            // no vertices have been moved, the algorithm has converged - exit
            return;
        }
		
        // build compressed graph
        compress(graph, new_component_id);
        // update component ids of vertices not in current graph
        translate(component_id, new_component_id);
    }
}

