#include "optimizer.h"
#include "surprise_fast.h"

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

// ── Asymptotic Surprise (KL-divergence, for merge-pass speed) ────────────────
static double asymptotic_surprise(double F, double M, double m, double p) {
    if (m <= 0.0 || F <= 0.0) return 0.0;
    double q  = p / m;
    double q0 = M / F;
    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;
    return m * (q * std::log(q / q0) + (1.0 - q) * std::log((1.0 - q) / (1.0 - q0)));
}

// ── Optimizer ─────────────────────────────────────────────────────────────────

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

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

    Optimizer(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& e : edges) {
            adj[e.first ].push_back(e.second);
            adj[e.second].push_back(e.first );
        }
        ensure_log_fact((int)std::min((double)LOG_FACT_MAX, F));
    }

    // edges from v to community c
    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 and recompute M, p, S_cur from scratch
    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_cur = surprise_fast(F, M, m, p);
    }

    // ── Pre_CS initialization (tang2019) ──────────────────────────────────
    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 ────────────────────────────────────────────────
    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];
            int n_c    = comm_size[cur_c];
            int wvc    = w_to(v, cur_c);

            // Candidate communities: neighbors of v + 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_new = M, best_p_new = 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 M_new = M + dM;
                double p_new = p + dp;
                if (M_new < 0.0 || p_new < 0.0) continue;

                double S_new = surprise_fast(F, M_new, m, p_new);
                double delta = S_new - S_cur;
                if (delta > best_delta) {
                    best_delta = delta;
                    best_c     = cand;
                    best_M_new = M_new;
                    best_p_new = p_new;
                }
            }

            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_new;
                p       = best_p_new;
                S_cur  += best_delta;
                improved = true;
            }
        }
        recompute();   // compact + guard against drift
        return improved;
    }

    // ── Community merge pass (uses asymptotic S for O(1) ΔS) ─────────────
    bool merge_pass() {
        bool any = false;

        std::vector<int> cs;
        {
            std::unordered_set<int> seen(comm.begin(), comm.end());
            cs.assign(seen.begin(), seen.end());
        }

        bool local = true;
        while (local) {
            local = false;
            // Refresh
            {
                std::unordered_set<int> seen(comm.begin(), comm.end());
                cs.assign(seen.begin(), seen.end());
            }
            std::sort(cs.begin(), cs.end());

            for (int i = 0; i < (int)cs.size() && !local; i++) {
                for (int j = i + 1; j < (int)cs.size() && !local; j++) {
                    int ca = cs[i], cb = cs[j];
                    int na = comm_size[ca], nb = comm_size[cb];

                    // edges from ca to cb (undirected count)
                    int eab = 0;
                    for (int v = 0; v < N; v++)
                        if (comm[v] == ca)
                            for (int u : adj[v])
                                if (comm[u] == cb) eab++;

                    double dM    = (double)na * nb;
                    double M_new = M + dM;
                    double p_new = p + (double)eab;

                    // Use exact Surprise for final merge decision
                    double S_new = surprise_fast(F, M_new, m, p_new);
                    if (S_new > S_cur) {
                        for (int v = 0; v < N; v++)
                            if (comm[v] == cb) { comm[v] = ca; }
                        comm_size[ca] += nb;
                        comm_size[cb]  = 0;
                        M     = M_new;
                        p     = p_new;
                        S_cur = S_new;
                        local = true;
                        any   = true;
                    }
                }
            }
        }
        if (any) recompute();
        return any;
    }

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

        bool go = true;
        while (go) {
            go = node_pass(rng);
            if (!go) go = merge_pass();
        }
        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_comm;
        double best_S = -std::numeric_limits<double>::infinity();

        for (int r = 0; r < n_restarts; r++) {
            // Alternate Pre_CS (even) and singleton (odd) starts
            double S = run(r % 2 == 0, rng);
            if (S > best_S) {
                best_S    = S;
                best_comm = comm;
            }
        }
        comm = best_comm;
        recompute();
        return comm;
    }
};

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