#include "sp4.h"
#include "surprise_fast.h"

#include <algorithm>
#include <cmath>
#include <limits>
#include <numeric>
#include <random>
#include <unordered_map>
#include <unordered_set>
#include <vector>

// ═══════════════════════════════════════════════════════════════════════════
//  FAGSO  (Jiang et al. 2014 / Nicolini et al. 2016)
// ═══════════════════════════════════════════════════════════════════════════
// Ranks every edge by Jaccard similarity; then sweeps them in decreasing
// order merging two communities whenever ΔSurprise > 0.
// Uses an inter-community edge map so ΔM and Δp are O(1) per candidate.

std::vector<int> fagso_surprise(
    const std::vector<std::pair<int,int>>& edges,
    int n_nodes)
{
    if (n_nodes < 2) return {0};

    double F = (double)n_nodes * (n_nodes - 1) / 2.0;
    double m = (double)edges.size();
    ensure_log_fact((int)std::min((double)LOG_FACT_MAX, F));

    // Build adjacency sets (for Jaccard)
    std::vector<std::unordered_set<int>> adj(n_nodes);
    for (auto& [u, v] : edges) { adj[u].insert(v); adj[v].insert(u); }

    // Jaccard for each edge
    std::vector<std::pair<double, std::pair<int,int>>> ranked;
    ranked.reserve(edges.size());
    for (auto& [u, v] : edges) {
        int common = 0;
        for (int w : adj[u]) if (adj[v].count(w)) common++;
        double jac = (double)common / (adj[u].size() + adj[v].size() - common);
        ranked.push_back({jac, {u, v}});
    }
    std::sort(ranked.begin(), ranked.end(),
              [](auto& a, auto& b){ return a.first > b.first; });

    // All-singleton init
    std::vector<int> comm(n_nodes);
    std::iota(comm.begin(), comm.end(), 0);
    std::vector<int> comm_size(n_nodes, 1);

    // Symmetric inter-community edge map: inter[a][b] = inter[b][a] =
    // number of undirected edges between community a and community b.
    std::vector<std::unordered_map<int,int>> inter(n_nodes);
    for (auto& [u, v] : edges) {
        int a = comm[u], b = comm[v];
        if (a == b) continue;
        inter[a][b]++;
        inter[b][a]++;
    }

    double M = 0.0, p = 0.0;
    double S_cur = surprise_fast(F, M, m, p);

    // Helper: merge community cb into ca
    auto merge_into = [&](int ca, int cb) {
        int eab = inter[ca].count(cb) ? inter[ca][cb] : 0;
        M += (double)comm_size[ca] * comm_size[cb];
        p += (double)eab;
        comm_size[ca] += comm_size[cb];
        comm_size[cb]  = 0;

        // Move all nodes of cb to ca
        for (int v = 0; v < n_nodes; v++) if (comm[v] == cb) comm[v] = ca;

        // Copy cb neighbours before modifying to avoid iterator invalidation
        auto cb_nbrs = inter[cb];
        inter[ca].erase(cb);
        inter[cb].erase(ca);
        for (auto& [nc, cnt] : cb_nbrs) {
            if (nc == ca) continue;
            inter[ca][nc] += cnt;
            inter[nc][ca] += cnt;
            inter[nc].erase(cb);
        }
        inter[cb].clear();
    };

    // Main FAGSO sweep
    for (auto& [jac, uv] : ranked) {
        auto [u, v] = uv;
        int cu = comm[u], cv = comm[v];
        if (cu == cv) continue;

        int key_a = std::min(cu, cv), key_b = std::max(cu, cv);
        int eab = inter[key_a].count(key_b) ? inter[key_a][key_b] : 0;
        double dM    = (double)comm_size[cu] * comm_size[cv];
        double dp    = (double)eab;
        double S_new = surprise_fast(F, M + dM, m, p + dp);
        if (S_new > S_cur) {
            S_cur = S_new;
            merge_into(cu, cv);
        }
    }

    // Compact IDs
    std::unordered_map<int,int> remap;
    for (int v = 0; v < n_nodes; v++) {
        auto [it, ok] = remap.emplace(comm[v], (int)remap.size());
        comm[v] = it->second;
    }
    return comm;
}

// ═══════════════════════════════════════════════════════════════════════════
//  SP4  (Tang et al. 2019)
//  Pre_CS  →  Louvain_process (with super-node aggregation)  →  Post_CS
// ═══════════════════════════════════════════════════════════════════════════

struct SP4Solver {
    int N;
    double F, m;
    std::vector<std::vector<int>> adj;

    // ── Current state ──────────────────────────────────────────────────────
    std::vector<int> comm;
    std::vector<int> comm_size;
    double M = 0.0, p = 0.0, S_cur = 0.0;

    SP4Solver(const std::vector<std::pair<int,int>>& edges, int n_nodes)
        : N(n_nodes), adj(n_nodes)
    {
        F = (double)N * (N - 1) / 2.0;
        m = (double)edges.size();
        for (auto& [u,v] : edges) { adj[u].push_back(v); adj[v].push_back(u); }
        ensure_log_fact((int)std::min((double)LOG_FACT_MAX, F));
    }

    int w_to(int v, int c) const {
        int cnt = 0;
        for (int u : adj[v]) cnt += (comm[u] == c);
        return cnt;
    }

    void recompute() {
        std::unordered_map<int,int> remap;
        for (int v = 0; v < N; v++) {
            auto [it,ok] = remap.emplace(comm[v], (int)remap.size());
            comm[v] = it->second;
        }
        int K = (int)remap.size();
        comm_size.assign(K, 0);
        for (int v = 0; v < N; v++) comm_size[comm[v]]++;
        M = 0.0;
        for (int s : comm_size) M += (double)s * (s - 1) / 2.0;
        p = 0.0;
        for (int v = 0; v < N; v++)
            for (int u : adj[v]) if (comm[u] == comm[v] && u > v) p++;
        S_cur = surprise_fast(F, M, m, p);
    }

    // ── Pre_CS: seed by most-common-neighbour ──────────────────────────────
    void init_pre_cs(std::mt19937& rng) {
        comm.resize(N); std::iota(comm.begin(), comm.end(), 0);
        std::vector<int> order(N); std::iota(order.begin(), order.end(), 0);
        std::shuffle(order.begin(), order.end(), rng);
        std::vector<int> cn(N, 0);
        for (int v : order) {
            std::fill(cn.begin(), cn.end(), 0);
            for (int u : adj[v]) for (int w : adj[u]) if (w != v) cn[w]++;
            int best_u = -1, best_cn = 0;
            for (int u : adj[v]) if (cn[u] > best_cn) { best_cn = cn[u]; best_u = u; }
            if (best_u >= 0) comm[v] = comm[best_u];
        }
        recompute();
    }

    void init_singleton() {
        comm.resize(N); std::iota(comm.begin(), comm.end(), 0); recompute();
    }

    // ── Node-movement pass on original graph ──────────────────────────────
    bool node_pass(std::mt19937& rng) {
        bool improved = false;
        std::vector<int> order(N); std::iota(order.begin(), order.end(), 0);
        std::shuffle(order.begin(), order.end(), rng);
        int next_id = (int)comm_size.size();

        for (int v : order) {
            int cur_c = comm[v], n_c = comm_size[cur_c], wvc = w_to(v, cur_c);
            std::unordered_set<int> cands;
            for (int u : adj[v]) cands.insert(comm[u]);
            cands.insert(next_id); cands.erase(cur_c);

            double best_d = 0.0; int best_c = cur_c;
            double best_M = M,   best_p = p;

            for (int cand : cands) {
                int n_d = (cand < (int)comm_size.size()) ? comm_size[cand] : 0;
                int wvd = (cand < (int)comm_size.size()) ? w_to(v, cand)   : 0;
                double M2 = M + (double)n_d - (double)(n_c - 1);
                double p2 = p + (double)(wvd - wvc);
                if (M2 < 0.0 || p2 < 0.0) continue;
                double d = surprise_fast(F, M2, m, p2) - S_cur;
                if (d > best_d) { best_d = d; best_c = cand; best_M = M2; best_p = p2; }
            }

            if (best_c != cur_c) {
                if (best_c == next_id) { comm_size.push_back(1); next_id++; }
                else comm_size[best_c]++;
                comm_size[cur_c]--;
                comm[v] = best_c; M = best_M; p = best_p; S_cur += best_d;
                improved = true;
            }
        }
        recompute();
        return improved;
    }

    // ── Super-node aggregation pass ────────────────────────────────────────
    // Builds a compressed graph where each community = one super-node,
    // then runs a node-movement pass on that compressed graph.
    // Returns true if any improvement was found.
    bool aggregation_pass(std::mt19937& rng) {
        // Map community IDs to super-node indices 0..K-1
        recompute();
        int K = (int)comm_size.size();
        if (K <= 1) return false;

        // super_size[s] = original nodes in super-node s  (= comm_size[s])
        std::vector<int> sz(K);
        for (int s = 0; s < K; s++) sz[s] = comm_size[s];

        // Build super-graph adjacency (weighted)
        // super_adj[s][t] = edges between original communities s and t
        std::vector<std::unordered_map<int,int>> sadj(K);
        std::vector<int> sintra(K, 0);  // intra-edges within super-node
        for (int v = 0; v < N; v++) {
            int cv = comm[v];
            for (int u : adj[v]) {
                int cu = comm[u];
                if (cu == cv) { if (u > v) sintra[cv]++; }
                else sadj[cv][cu]++;
            }
        }
        // sadj[s][t] = directed count → divide by 2 for undirected
        // (we'll use directed and halve when computing Δp)

        // Partition of super-nodes into meta-communities (start: each alone)
        std::vector<int> mc(K); std::iota(mc.begin(), mc.end(), 0);
        std::vector<double> mc_size(K);
        for (int s = 0; s < K; s++) mc_size[s] = sz[s];

        // Current M and p (same as after recompute)
        double Mc = M, pc = p, Sc = S_cur;

        // sadj[s][t] = directed edges from community s to community t.
        // For undirected original graph, sadj[s][t] = undirected edges between s and t
        // (counted once per side), so summing sadj[s][*] for neighbours in a
        // meta-community gives the correct undirected edge count without /2.
        auto w_super_to_mc = [&](int s, int meta_c) -> double {
            double cnt = 0.0;
            for (auto& [t, w] : sadj[s])
                if (mc[t] == meta_c) cnt += w;
            return cnt;
        };

        bool improved = false;
        bool go = true;
        while (go) {
            go = false;
            std::vector<int> order(K); std::iota(order.begin(), order.end(), 0);
            std::shuffle(order.begin(), order.end(), rng);
            int next_mc = K;

            for (int s : order) {
                int cur_mc = mc[s];
                double N_C = mc_size[cur_mc];   // total original nodes in meta-comm C
                double N_s = (double)sz[s];      // original nodes in super-node s
                double w_sC = w_super_to_mc(s, cur_mc);  // edges from s to rest of C

                // Candidate meta-communities = neighbour meta-comms + new singleton
                std::unordered_set<int> cands;
                for (auto& [t, w] : sadj[s]) cands.insert(mc[t]);
                cands.insert(next_mc); cands.erase(cur_mc);

                double best_d = 0.0; int best_mc = cur_mc;
                double best_Mc = Mc, best_pc = pc;

                for (int cmc : cands) {
                    double N_D = (cmc < (int)mc_size.size()) ? mc_size[cmc] : 0.0;
                    double w_sD = (cmc < (int)mc_size.size()) ? w_super_to_mc(s, cmc) : 0.0;

                    double dM = ((N_D + N_s) * (N_D + N_s - 1) / 2.0 - N_D * (N_D - 1) / 2.0) +
                                ((N_C - N_s) * (N_C - N_s - 1) / 2.0 - N_C * (N_C - 1) / 2.0);
                    double dp = w_sD - w_sC;
                    double M2 = Mc + dM, p2 = pc + dp;
                    if (M2 < 0.0 || p2 < 0.0) continue;

                    double d = surprise_fast(F, M2, m, p2) - Sc;
                    if (d > best_d) { best_d = d; best_mc = cmc; best_Mc = M2; best_pc = p2; }
                }

                if (best_mc != cur_mc) {
                    if (best_mc == next_mc) {
                        mc_size.push_back(N_s); next_mc++;
                    } else {
                        mc_size[best_mc] += N_s;
                    }
                    mc_size[cur_mc] -= N_s;
                    mc[s] = best_mc; Mc = best_Mc; pc = best_pc; Sc += best_d;
                    go = true; improved = true;
                }
            }
        }

        if (!improved) return false;

        // Map back: comm[v] = mc[comm[v]]
        for (int v = 0; v < N; v++) comm[v] = mc[comm[v]];
        recompute();
        return true;
    }

    // ── Single SP4 run ─────────────────────────────────────────────────────
    double run(bool pre_cs, std::mt19937& rng) {
        if (pre_cs) init_pre_cs(rng); else init_singleton();

        // Phase 1: node movement until convergence, then aggregation
        bool go = true;
        while (go) {
            go = false;
            while (node_pass(rng)) go = true;
            if (aggregation_pass(rng)) go = true;
        }

        // Phase 2 (Post_CS): one final node-movement pass on original graph
        node_pass(rng);
        return S_cur;
    }

    // ── Multi-restart ──────────────────────────────────────────────────────
    std::vector<int> optimize(int n_restarts, int seed) {
        std::mt19937 rng(seed >= 0 ? (unsigned)seed : std::random_device{}());
        std::vector<int> best; double best_S = -std::numeric_limits<double>::infinity();
        for (int r = 0; r < n_restarts; r++) {
            double S = run(r % 2 == 0, rng);
            if (S > best_S) { best_S = S; best = comm; }
        }
        comm = best; recompute(); return comm;
    }
};

std::vector<int> sp4_surprise(
    const std::vector<std::pair<int,int>>& edges,
    int n_nodes, int n_restarts, int seed)
{
    if (n_nodes < 2) return {0};
    SP4Solver solver(edges, n_nodes);
    return solver.optimize(n_restarts, seed);
}
