"""
Faithful Python reimplementation of Duch & Arenas (2005) extremal optimization.

Recursive bipartition of a graph by modularity maximisation using τ-EO
(extremal optimisation with power-law rank selection).

Algorithm
---------
1. Start with all nodes in a single community.
2. Attempt to split every community into two groups using τ-EO.
3. Accept the split only when it increases global modularity Q.
4. Repeat until no split improves Q.

τ-EO bipartition
~~~~~~~~~~~~~~~~~
For a subset of k nodes to be split:
  a. Random bipartition into two groups.
  b. Compute the fitness (modularity contribution) of every node.
  c. Rank nodes by fitness (ascending → worst first).
  d. Select a node from rank distribution P(r) ∝ r^{-τ},  τ = 1 + 1/ln(N).
  e. Move the selected node to the other group.
  f. Track the best ΔQ seen; revert to it when patience runs out.

Reference
---------
Duch J & Arenas A (2005) "Community detection in complex networks
using extremal optimization." Physical Review E 72(2): 027104.
"""

from __future__ import annotations

import math
import random
from typing import Dict, Optional, Sequence, Tuple

Edge = Tuple[str, str]
Partition = Dict[str, int]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def solve(edges: Sequence[Edge]) -> Partition:
    """Run DA extremal optimisation and return {node_label: community_id}."""

    node2idx, idx2node = _node_index_maps(edges)
    n = len(node2idx)
    if n == 0:
        return {}

    # Build adjacency list & degree vector
    adj: list[list[int]] = [[] for _ in range(n)]
    for u, v in edges:
        ui, vi = node2idx[u], node2idx[v]
        adj[ui].append(vi)
        adj[vi].append(ui)

    deg = [len(adj[i]) for i in range(n)]
    m = sum(deg) // 2
    if m == 0:
        return {idx2node[i]: 0 for i in range(n)}
    two_m = 2 * m

    # ── Recursive splitting ───────────────────────────────────────────
    comm = [0] * n
    next_cid = 1
    communities: dict[int, list[int]] = {0: list(range(n))}

    changed = True
    while changed:
        changed = False
        for cid in list(communities.keys()):
            nodes = communities[cid]
            if len(nodes) <= 1:
                continue

            result = _eo_split(nodes, adj, deg, m, two_m, n)
            if result is None:
                continue

            g0, g1 = result
            new_cid = next_cid
            next_cid += 1
            for nd in g1:
                comm[nd] = new_cid
            communities[cid] = g0
            communities[new_cid] = g1
            changed = True
            break  # restart — community map changed

    return {idx2node[i]: comm[i] for i in range(n)}


# ---------------------------------------------------------------------------
# Internals
# ---------------------------------------------------------------------------

def _node_index_maps(edges: Sequence[Edge]):
    node2idx: Dict[str, int] = {}
    idx2node: Dict[int, str] = {}
    idx = 0
    for u, v in edges:
        if u not in node2idx:
            node2idx[u] = idx
            idx2node[idx] = u
            idx += 1
        if v not in node2idx:
            node2idx[v] = idx
            idx2node[idx] = v
            idx += 1
    return node2idx, idx2node


def _eo_split(
    node_list: list[int],
    adj: list[list[int]],
    deg: list[int],
    m: int,
    two_m: int,
    n_total: int,
) -> Optional[tuple[list[int], list[int]]]:
    """Try to bipartition *node_list* using τ-EO.

    Returns ``(group0, group1)`` or ``None`` if the split does not
    improve modularity.
    """
    k = len(node_list)
    if k <= 1:
        return None

    # τ uses ln(N_total) as in the paper (Duch & Arenas 2005, Sec. II)
    tau = 1.0 + 1.0 / math.log(max(n_total, 3))
    max_steps = k * k          # no cap — faithfulness over speed
    patience = max(k * 10, 500)

    # Random bipartition
    shuffled = list(node_list)
    random.shuffle(shuffled)
    grp: dict[int, int] = {}
    for i, nd in enumerate(shuffled):
        grp[nd] = 0 if i < k // 2 else 1

    # Degree sums per group
    d_grp = [0, 0]
    for nd in node_list:
        d_grp[grp[nd]] += deg[nd]

    # Edges from each node to each group (only within the subset)
    e_to: dict[int, list[int]] = {}
    for nd in node_list:
        e = [0, 0]
        for nb in adj[nd]:
            if nb in grp:
                e[grp[nb]] += 1
        e_to[nd] = e

    # Internal edge counts per group (tracked incrementally)
    l_grp = [0, 0]
    for nd in node_list:
        l_grp[grp[nd]] += e_to[nd][grp[nd]]
    l_grp[0] //= 2
    l_grp[1] //= 2

    # Total edges within the subset (constant across moves)
    total_sub = sum(e_to[nd][0] + e_to[nd][1] for nd in node_list) // 2
    d_all = d_grp[0] + d_grp[1]

    def _delta_q() -> float:
        q_before = total_sub / m - (d_all / two_m) ** 2
        q_after = (l_grp[0] / m - (d_grp[0] / two_m) ** 2 +
                   l_grp[1] / m - (d_grp[1] / two_m) ** 2)
        return q_after - q_before

    # Fitness per node: λ_i = q_i / k_i  (normalized by degree)
    # q_i = modularity contribution of node i
    # Paper: "normalized by its number of connections"
    fit_val: dict[int, float] = {}
    for nd in node_list:
        g = grp[nd]
        q_i = e_to[nd][g] / two_m - deg[nd] * d_grp[g] / (two_m * two_m)
        fit_val[nd] = q_i / deg[nd] if deg[nd] > 0 else 0.0

    best_dq = _delta_q()
    best_grp = dict(grp)
    no_improve = 0

    # Precompute τ-CDF for rank selection
    cum_w: list[float] = []
    s = 0.0
    for r in range(1, k + 1):
        s += r ** (-tau)
        cum_w.append(s)
    total_w = s

    for _ in range(max_steps):
        # Rank nodes by fitness (ascending = worst first)
        ranked = sorted(node_list, key=lambda nd: fit_val[nd])

        # τ-selection: draw rank from P(r) ∝ r^{-τ}
        target = random.random() * total_w
        idx = k - 1  # fallback: best node (if FP rounding overshoots)
        for i, cw in enumerate(cum_w):
            if target <= cw:
                idx = i
                break

        nd = ranked[idx]
        old_g = grp[nd]
        new_g = 1 - old_g

        # Pre-move edge counts
        e_nd_old = e_to[nd][old_g]
        e_nd_new = e_to[nd][new_g]

        # Update internal edge counts
        l_grp[old_g] -= e_nd_old
        l_grp[new_g] += e_nd_new

        # Update degree sums
        d_grp[old_g] -= deg[nd]
        d_grp[new_g] += deg[nd]

        # Move node and update neighbours
        grp[nd] = new_g
        for nb in adj[nd]:
            if nb in grp:
                e_to[nb][old_g] -= 1
                e_to[nb][new_g] += 1
                g_nb = grp[nb]
                q_nb = (e_to[nb][g_nb] / two_m
                        - deg[nb] * d_grp[g_nb] / (two_m * two_m))
                fit_val[nb] = q_nb / deg[nb] if deg[nb] > 0 else 0.0

        q_nd = (e_to[nd][new_g] / two_m
                - deg[nd] * d_grp[new_g] / (two_m * two_m))
        fit_val[nd] = q_nd / deg[nd] if deg[nd] > 0 else 0.0

        dq = _delta_q()
        if dq > best_dq + 1e-12:
            best_dq = dq
            best_grp = dict(grp)
            no_improve = 0
        else:
            no_improve += 1
            if no_improve >= patience:
                break

    if best_dq <= 1e-12:
        return None

    g0 = [nd for nd in node_list if best_grp[nd] == 0]
    g1 = [nd for nd in node_list if best_grp[nd] == 1]
    if not g0 or not g1:
        return None
    return g0, g1
