#include "optimizer_v2.h"
#include "surprise_fast.h"
#include "sp4.h"          // fagso_surprise() for warm-start

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

// ── Asymptotic Surprise: S ≈ m · D(q ‖ ⟨q⟩)  — O(1) ──────────────────────
// Traag, Aldecoa & Delvenne (2015), Eq. 2-3
static double asym_S(double F, double M, double m, double p) {
    if (m <= 0.0 || F <= 0.0) return 0.0;
    double q  = p / m;          // fraction of intra-community edges
    double q0 = M / F;          // expected fraction under ER null
    // Clamp to avoid log(0)
    if (q  <= 0.0) q  = 1e-15;
    if (q  >= 1.0) q  = 1.0 - 1e-15;
    if (q0 <= 0.0) q0 = 1e-15;
    if (q0 >= 1.0) q0 = 1.0 - 1e-15;
    // KL divergence D(q ‖ q0)
    double kl = q * std::log(q / q0) + (1.0 - q) * std::log((1.0 - q) / (1.0 - q0));
    return m * kl;
}

// ═══════════════════════════════════════════════════════════════════════════════
//  M2Solver
// ═══════════════════════════════════════════════════════════════════════════════

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

    // Current partition state
    std::vector<int> comm;
    std::vector<int> comm_size;
    double M = 0.0, p = 0.0;
    double S_exact = 0.0;   // exact Surprise (for comparison / final)
    double S_asym  = 0.0;   // asymptotic Surprise (for fast moves)

    M2Solver(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;
    }

    // Compact community IDs, recompute M, p, S_exact, S_asym
    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 += 1.0;

        S_exact = surprise_fast(F, M, m, p);
        S_asym  = asym_S(F, M, m, p);
    }

    // ── Initialization strategies ────────────────────────────────────────────

    // FAGSO warm-start: run FAGSO and adopt its partition
    void init_fagso(const std::vector<std::pair<int,int>>& edges) {
        comm = fagso_surprise(edges, N);
        recompute();
    }

    // Pre_CS: common-neighbour seeding (Tang et al. 2019)
    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();
    }

    // ── Phase 1: Asymptotic node-movement pass ──────────────────────────────
    // Uses O(1) asymptotic S for ALL candidate evaluations.
    // Returns true if any node was moved.
    bool asym_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];
            int n_c   = comm_size[cur_c];
            int wvc   = w_to(v, cur_c);

            // Candidate communities: neighbours + new singleton
            std::unordered_set<int> cands;
            for (int u : adj[v]) cands.insert(comm[u]);
            cands.insert(next_id);
            cands.erase(cur_c);

            double best_delta = 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 dM = (double)n_d - (double)(n_c - 1);
                double dp = (double)(wvd - wvc);
                double M2 = M + dM;
                double p2 = p + dp;
                if (M2 < 0.0 || p2 < 0.0) continue;

                // O(1) asymptotic ΔS
                double S2 = asym_S(F, M2, m, p2);
                double delta = S2 - S_asym;
                if (delta > best_delta) {
                    best_delta = delta;
                    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_asym += best_delta;
                improved = true;
            }
        }
        recompute();
        return improved;
    }

    // ── Phase 1b: Super-node aggregation pass ───────────────────────────────
    // Compresses current communities into super-nodes, then runs asymptotic
    // node movements on the compressed graph.
    // Adapted from SP4Solver::aggregation_pass (sp4.cpp).
    bool aggregation_pass(std::mt19937& rng) {
        recompute();
        int K = (int)comm_size.size();
        if (K <= 1) return false;

        // Super-node sizes
        std::vector<int> sz(K);
        for (int s = 0; s < K; s++) sz[s] = comm_size[s];

        // Build super-graph adjacency
        std::vector<std::unordered_map<int,int>> sadj(K);
        for (int v = 0; v < N; v++) {
            int cv = comm[v];
            for (int u : adj[v]) {
                int cu = comm[u];
                if (cu != cv) sadj[cv][cu]++;
            }
        }

        // Meta-community partition (each super-node starts 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];

        double Mc = M, pc = p, Sc_asym = S_asym;

        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 = (int)mc_size.size();

            for (int s : order) {
                if (sz[s] == 0) continue;
                int cur_mc = mc[s];
                double N_C = mc_size[cur_mc];
                double N_s = (double)sz[s];
                double w_sC = w_super_to_mc(s, cur_mc);

                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_id = 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;

                    // O(1) asymptotic ΔS
                    double d = asym_S(F, M2, m, p2) - Sc_asym;
                    if (d > best_d) {
                        best_d = d;
                        best_mc_id = cmc;
                        best_Mc = M2;
                        best_pc = p2;
                    }
                }

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

        if (!improved) return false;

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

    // ── Phase 2: Exact node-movement refinement ─────────────────────────────
    // One pass using exact surprise_fast. Corrects asymptotic approximation
    // errors for borderline moves.
    bool exact_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];
            int n_c   = comm_size[cur_c];
            int 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_delta = 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 dM = (double)n_d - (double)(n_c - 1);
                double dp = (double)(wvd - wvc);
                double M2 = M + dM;
                double p2 = p + dp;
                if (M2 < 0.0 || p2 < 0.0) continue;

                // Exact ΔS — only in final refinement
                double S2 = surprise_fast(F, M2, m, p2);
                double delta = S2 - S_exact;
                if (delta > best_delta) {
                    best_delta = delta;
                    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_exact += best_delta;
                improved = true;
            }
        }
        recompute();
        return improved;
    }

    // ── Single M² run ──────────────────────────────────────────────────────
    double run(int init_mode, std::mt19937& rng,
               const std::vector<std::pair<int,int>>& edges) {
        // Initialization
        switch (init_mode) {
            case 0: init_fagso(edges);   break;   // FAGSO warm-start
            case 1: init_pre_cs(rng);    break;   // Pre_CS
            default: init_singleton();   break;   // cold start
        }

        // Phase 1: asymptotic Louvain + aggregation until stable
        bool go = true;
        while (go) {
            go = false;
            while (asym_node_pass(rng)) go = true;
            if (aggregation_pass(rng)) go = true;
        }

        // Phase 2: exact refinement on original graph
        exact_node_pass(rng);

        return S_exact;
    }

    // ── Multi-restart entry point ───────────────────────────────────────────
    std::vector<int> optimize(int n_restarts, int seed,
                              const std::vector<std::pair<int,int>>& edges) {
        std::mt19937 rng(seed >= 0 ? (unsigned)seed : std::random_device{}());

        std::vector<int> best_comm;
        double best_S = -std::numeric_limits<double>::infinity();

        for (int r = 0; r < n_restarts; r++) {
            // r=0: FAGSO,  r=1: Pre_CS,  r=2: singleton,  r>=3: alternate
            int init_mode;
            if (r == 0) init_mode = 0;       // FAGSO
            else if (r % 2 == 1) init_mode = 1; // Pre_CS
            else init_mode = 2;              // singleton

            double S = run(init_mode, rng, edges);
            if (S > best_S) {
                best_S = S;
                best_comm = comm;
            }
        }
        comm = best_comm;
        recompute();
        return comm;
    }
};

// ── Public API ───────────────────────────────────────────────────────────────
std::vector<int> m2_surprise(
    const std::vector<std::pair<int,int>>& edges,
    int n_nodes,
    int n_restarts,
    int seed)
{
    if (n_nodes < 2) return {0};
    M2Solver solver(edges, n_nodes);
    return solver.optimize(n_restarts, seed, edges);
}
