"""
Faithful Python reimplementation of rnMRA — Ronhovde & Nussinov (2009).

Multiresolution Absolute Potts Model greedy solver.

This module mirrors the original ``rnMRA`` C++ binary logic:
  - Greedy node-by-node energy minimization via ``communityDetectNZ``
    (cluster.hpp:51–390): cluster-by-cluster iteration, clean flags,
    same-node bounce detection, singleton preference when cost_cur > 0.
  - Community merge step via ``detectAndMergeAllNZ`` (clustermerges.cpp):
    merge+refine loop until convergence.
  - Multiple trials per gamma (``communityBestE``, RNMRA.hpp:72–175),
    keeping the lowest-energy partition.
  - 8 independent replicas per gamma step (``rnMRA``, RNMRA.hpp:375–408).
  - Resolution selection via inter-replica pairwise NMI (``nmviMRA``,
    RNMain.hpp:272–310): max average NMI with spike filter and plateau
    detection (RNMRA.hpp:456–504).
  - Geometric gamma scan from high to low.

Parameters match rnMRA defaults (cluster_clparams.cpp / runRN.pl):
  gamma range   : 19.0 → 0.001
  steps/decade  : 20
  NTrialsMax    : 4
  NReplicasMax  : 8
  NIterMax      : 1000

Reference
---------
Ronhovde P & Nussinov Z (2009) "Multiresolution community detection
for megascale networks by information-based replica correlations."
Physical Review E 80(1): 016109.
"""

from __future__ import annotations

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

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


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

def solve(edges: Sequence[Edge]) -> Partition:
    """Run the full rnMRA pipeline and return {node_label: community_id}."""

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

    # Build adjacency list (integer-indexed)
    adj: list[list[int]] = [[] for _ in range(n_nodes)]
    for u, v in edges:
        ui, vi = node2idx[u], node2idx[v]
        adj[ui].append(vi)
        adj[vi].append(ui)

    # ── MRA parameters (cluster_clparams.cpp / runRN.pl) ─────────────
    gamma_start = 19.0           # gammaStart (cluster_clparams.cpp:57)
    gamma_stop = 0.001           # gammaStop  (cluster_clparams.cpp:58)
    n_decade_steps = 20          # gammaNSteps (cluster_clparams.cpp:59)
    n_trials = 4                 # NTrialsMax
    n_replicas = 8               # NReplicasMax
    max_iter = 1000              # NIterMax
    n_steps_max = n_nodes        # NStepsMax via -nsm:$nNodes (runRN.pl:37)

    # Fixed gamma multiplier (RNMRA.hpp:227): step down by 10^(1/20)
    gamma_log_step = 10.0 ** (1.0 / n_decade_steps)
    epsilon = 1e-10

    # ── Full MRA scan with replicas (RNMRA.hpp:340-617) ──────────────
    # do-while: process gamma → step → check termination
    scan: list[tuple[float, list[list[int]], list[float], float]] = []
    gamma = gamma_start

    for _step_i in range(n_steps_max):
        r_mems: list[list[int]] = []
        r_energies: list[float] = []
        for _r in range(n_replicas):
            mem, e = _best_trial(adj, n_nodes, gamma, n_trials, max_iter)
            r_mems.append(mem)
            r_energies.append(e)
        avg_q = sum(len(set(m)) for m in r_mems) / n_replicas
        scan.append((gamma, r_mems, r_energies, avg_q))

        # Step gamma (RNMRA.hpp:601) — high→low: divide
        gamma /= gamma_log_step

        # Termination (RNMRA.hpp:614-617)
        if gamma <= gamma_stop:
            break
        if avg_q <= 1.0 + epsilon:
            break

    if not scan:
        return {idx2node[i]: 0 for i in range(n_nodes)}

    # ── Resolution selection via inter-replica NMI ────────────────────
    best_gamma_idx, best_replica_idx = _select_resolution(scan, n_nodes)
    chosen = scan[best_gamma_idx][1][best_replica_idx]
    return {idx2node[i]: int(c) for i, c in enumerate(chosen)}


# ---------------------------------------------------------------------------
# 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 _greedy_solve(
    adj: list[list[int]],
    n_nodes: int,
    gamma: float,
    max_iter: int,
    init_mem: Optional[list[int]] = None,
) -> list[int]:
    """One greedy solve matching ``communityDetectNZ`` (cluster.hpp:51-390).

    - Cluster-by-cluster iteration (deterministic, sorted by cluster id)
    - Clean flag: each node processed at most once per step
    - Same-node bounce detection (cluster.hpp:312)
    - Absolute Potts Model cost: -(1+γ)·edges + γ·size
    - Singleton when cost_cur > 0 (cluster.hpp:218,241)
    """
    if init_mem is not None:
        mem = list(init_mem)
        csz: dict[int, int] = {}
        for c in mem:
            csz[c] = csz.get(c, 0) + 1
    else:
        # Match C++ initSymmetric(bRandomOrder=1) (clusterclasses.cpp:1507-1510):
        # creates singletons then swapClusters(i, randomInt(0,nClusters-1))
        # for each i, randomizing which node sits at which cluster position.
        # This gives each trial/replica a different node processing order.
        perm = list(range(n_nodes))
        random.shuffle(perm)
        mem = perm
        csz = {perm[i]: 1 for i in range(n_nodes)}
    next_id = max(mem) + 1 if mem else n_nodes

    apg = 1.0 + gamma  # (a + |b|) / a  =  1 + γ
    last_node_moved = -2
    this_node_moved = -1
    equal_energy_count = 0

    for _step in range(max_iter):
        moved = False
        clean = [True] * n_nodes

        # Build cluster → node list from current membership
        cmap: dict[int, list[int]] = {}
        for i, c in enumerate(mem):
            cmap.setdefault(c, []).append(i)

        # Iterate cluster-by-cluster (sorted id = deterministic)
        for cid in sorted(cmap.keys()):
            for node in list(cmap.get(cid, [])):
                if not clean[node]:
                    continue

                cur = mem[node]
                nbrs = adj[node]
                if not nbrs:
                    continue

                # Count edges to each neighbouring community
                ec: dict[int, int] = {}
                for nb in nbrs:
                    c = mem[nb]
                    ec[c] = ec.get(c, 0) + 1

                e_cur = ec.get(cur, 0)
                s_cur = csz[cur] - 1  # exclude self
                cost_cur = -apg * e_cur + gamma * s_cur

                best_cost = cost_cur
                best_tgt = cur

                for tgt, e_tgt in ec.items():
                    if tgt == cur:
                        continue
                    cost_tgt = -apg * e_tgt + gamma * csz[tgt]
                    if cost_tgt < best_cost:
                        best_cost = cost_tgt
                        best_tgt = tgt

                # Singleton (cluster.hpp:218,241): cost_cur > 0 → new cluster
                if cost_cur > 0:
                    best_tgt = next_id
                    next_id += 1
                    csz[best_tgt] = 0
                    best_cost = 0  # cost in empty cluster = 0

                if best_tgt != cur:
                    csz[cur] -= 1
                    if csz[cur] == 0:
                        del csz[cur]
                    csz[best_tgt] = csz.get(best_tgt, 0) + 1
                    mem[node] = best_tgt
                    clean[node] = False

                    moved = True

                    # Same-node bounce (cluster.hpp:312)
                    last_node_moved = this_node_moved
                    this_node_moved = node
                    if last_node_moved == this_node_moved:
                        moved = False

                    # Zero-energy move counter (cluster.hpp:324-332)
                    e_change = best_cost - cost_cur
                    if e_change < -1e-15:
                        equal_energy_count = 0
                    elif abs(e_change) < 1e-15:
                        equal_energy_count += 1
                    if equal_energy_count > 1000:
                        moved = False

        if not moved:
            break
    return mem


def _merge(
    adj: list[list[int]],
    n_nodes: int,
    mem: list[int],
    gamma: float,
) -> tuple[list[int], int]:
    """Merge neighbouring community pairs when it lowers energy.

    Matches ``detectAndMergeAllNZ`` (clustermerges.cpp:339-488):
      1. Scan all cross-edges, compute merge energy per pair.
      2. Sort candidates by energy (best = most negative first).
      3. Apply non-conflicting merges (both sides blocked per pass).
      4. Repeat until no more beneficial merges.

    APM merge energy: -(1+γ)·cross + γ·s1·s2  (merge when < 0)
    Returns (mem, n_merges).
    """
    apg = 1.0 + gamma
    total_merges = 0

    while True:
        cmap: dict[int, list[int]] = {}
        for i, c in enumerate(mem):
            cmap.setdefault(c, []).append(i)

        # Compute cross-edges between all neighbour cluster pairs
        cross: dict[tuple[int, int], int] = {}
        for node in range(n_nodes):
            cn = mem[node]
            for nb in adj[node]:
                cm = mem[nb]
                if cn < cm:
                    key = (cn, cm)
                    cross[key] = cross.get(key, 0) + 1

        # Build sorted candidate list (ascending merge energy = best first)
        candidates: list[tuple[float, int, int]] = []
        for (c1, c2), x in cross.items():
            s1, s2 = len(cmap[c1]), len(cmap[c2])
            merge_e = -apg * x + gamma * s1 * s2
            if merge_e < 0:
                candidates.append((merge_e, c1, c2))
        if not candidates:
            break
        candidates.sort()

        # Apply non-conflicting merges (clustermerges.cpp:442)
        merged_clusters: set[int] = set()
        n_this_pass = 0
        for _me, c1, c2 in candidates:
            if c1 in merged_clusters or c2 in merged_clusters:
                continue
            for nd in cmap[c2]:
                mem[nd] = c1
            cmap[c1].extend(cmap[c2])
            del cmap[c2]
            merged_clusters.add(c1)
            merged_clusters.add(c2)
            n_this_pass += 1
            total_merges += 1

        if n_this_pass == 0:
            break

    return mem, total_merges


def _energy(
    adj: list[list[int]],
    mem: list[int],
    gamma: float,
) -> float:
    """APM energy: Σ_c [-(1+γ)·internal_edges + γ·(s choose 2)]."""
    apg = 1.0 + gamma
    cmap: dict[int, set[int]] = {}
    for i, c in enumerate(mem):
        cmap.setdefault(c, set()).add(i)
    e = 0.0
    for _c, nodes in cmap.items():
        s = len(nodes)
        internal = sum(1 for v in nodes for nb in adj[v] if nb in nodes) // 2
        e += -apg * internal + gamma * (s * (s - 1) // 2)
    return e


def _best_trial(
    adj: list[list[int]],
    n_nodes: int,
    gamma: float,
    n_trials: int,
    max_iter: int,
) -> tuple[list[int], float]:
    """Run n_trials greedy+merge+refine cycles; return lowest-energy result.

    Merge loop repeats merge→refine until no more merges (matching C++).
    """
    best_e = float('inf')
    best_mem: list[int] = list(range(n_nodes))
    for _ in range(n_trials):
        m = _greedy_solve(adj, n_nodes, gamma, max_iter)
        # Merge+refine loop until convergence (matches communityBestE)
        for _ml in range(max_iter):
            m, n_merges = _merge(adj, n_nodes, m, gamma)
            if n_merges == 0:
                break
            m = _greedy_solve(adj, n_nodes, gamma, max_iter, init_mem=m)
        e = _energy(adj, m, gamma)
        if e < best_e:
            best_e = e
            best_mem = m
    return best_mem, best_e


def _nmi_mi(m1: list[int], m2: list[int]) -> tuple[float, float]:
    """Normalised Mutual Information and MI between two membership vectors."""
    n = len(m1)
    if n == 0:
        return 0.0, 0.0
    c1 = Counter(m1)
    c2 = Counter(m2)
    joint = Counter(zip(m1, m2))

    def H(cnt: Counter) -> float:
        tot = sum(cnt.values())
        return -sum((v / tot) * math.log(v / tot)
                    for v in cnt.values() if v > 0)

    h1, h2 = H(c1), H(c2)
    mi = sum((c / n) * math.log((c / n) / ((c1[a] / n) * (c2[b] / n)))
             for (a, b), c in joint.items() if c > 0)
    denom = (h1 + h2) / 2
    nmi = mi / denom if denom > 1e-12 else 0.0
    return nmi, mi


def _nmi(m1: list[int], m2: list[int]) -> float:
    """Normalised Mutual Information between two membership vectors."""
    nmi, _mi = _nmi_mi(m1, m2)
    return nmi


def _lowest_energy_replica(energies: list[float]) -> int:
    """Return index of the replica with the lowest energy."""
    best_idx = 0
    best_e = energies[0]
    for j in range(1, len(energies)):
        if energies[j] < best_e:
            best_e = energies[j]
            best_idx = j
    return best_idx


def _select_resolution(
    scan: list[tuple[float, list[list[int]], list[float], float]],
    n_nodes: int,
) -> tuple[int, int]:
    """Inter-replica NMI resolution selection (RNMRA.hpp:456-504).

    For each gamma step, compute average pairwise NMI between all replica
    pairs (``nmviMRA``, RNMain.hpp:272-310).  Select the gamma with the
    highest average inter-replica NMI, using a moving-average spike filter
    and plateau detection.  Returns ``(gamma_index, replica_index)`` where
    the replica is the one with lowest energy at that gamma.
    """
    epsilon = 1e-10
    n_moving_avg = 3            # nMovingAvgLength (RNMRA.hpp:248)
    spike_filter_w = 0.8        # SpikeFilterW in C++

    # ── Compute average pairwise NMI and MI per gamma step ────────────
    avg_nmis: list[float] = []
    avg_mis: list[float] = []
    for _gamma, r_mems, _r_energies, _avg_q in scan:
        nmi_vals: list[float] = []
        mi_vals: list[float] = []
        n_rep = len(r_mems)
        for k in range(n_rep):
            for j in range(k + 1, n_rep):
                nmi_v, mi_v = _nmi_mi(r_mems[k], r_mems[j])
                nmi_vals.append(nmi_v)
                mi_vals.append(mi_v)
        avg_nmis.append(sum(nmi_vals) / len(nmi_vals) if nmi_vals else 0.0)
        avg_mis.append(sum(mi_vals) / len(mi_vals) if mi_vals else 0.0)

    # ── Moving average (circular buffer, RNMRA.hpp:443) ───────────────
    moving_buf = [0.0] * n_moving_avg
    moving_avgs: list[float] = []
    for i, nmi_val in enumerate(avg_nmis):
        moving_buf[i % n_moving_avg] = nmi_val
        filled = min(i + 1, n_moving_avg)
        moving_avgs.append(sum(moving_buf[:filled]) / filled
                           if i < n_moving_avg
                           else sum(moving_buf) / n_moving_avg)

    # ── Max NMI + spike filter + plateau (RNMRA.hpp:456-504) ─────────
    max_nmi = 0.0
    best_gamma_idx = 0
    best_replica_idx = 0
    best_h_count = 0
    current_h_count = 0
    prev_nmi = 0.0

    prev_mi = 0.0

    for i, (_gamma, r_mems, r_energies, avg_q) in enumerate(scan):
        nmi_val = avg_nmis[i]
        mi_val = avg_mis[i]

        # New maximum (RNMRA.hpp:458-472)
        # useEnergy=1 (cluster_clparams.cpp:588) → avg_q check bypassed
        if nmi_val > max_nmi + epsilon:
            # Spike filter: accept only if moving avg exceeds threshold
            if (moving_avgs[i] > spike_filter_w * max_nmi
                    or i < n_moving_avg):
                max_nmi = nmi_val
                # Pick lowest-energy replica (RNMRA.hpp:467)
                min_e_idx = _lowest_energy_replica(r_energies)
                best_gamma_idx = i
                best_replica_idx = min_e_idx
            current_h_count = 1
            best_h_count = 1

        # Plateau at current max (RNMRA.hpp:483-504)
        elif (nmi_val > max_nmi - epsilon
              and abs(nmi_val - prev_nmi) < epsilon):
            # Inner MI stability check (RNMRA.hpp:488)
            if abs(mi_val - prev_mi) < epsilon:
                current_h_count += 1
            else:
                current_h_count = 1
            if current_h_count > best_h_count:
                best_h_count = current_h_count
                min_e_idx = _lowest_energy_replica(r_energies)
                best_gamma_idx = i
                best_replica_idx = min_e_idx
        else:
            current_h_count = 1

        prev_nmi = nmi_val
        prev_mi = mi_val

    return best_gamma_idx, best_replica_idx
