src/models/rings.py

"""Mule-ring detection.

The defining trait of mule accounts is that they operate in *networks*, not in
isolation. The dataset has no transaction edges, so we infer latent rings from
BEHAVIOURAL co-similarity: among the accounts the model FLAGS, we build a k-NN
similarity graph and run community detection. Tight, dense clusters of
behaviourally near-identical high-risk accounts are candidate mule rings — the
unit a fraud team should freeze together to stop circulation.

Honest framing: these are behavioural-similarity rings (accounts that act alike
and score high together), not proven money-flow links.
"""
from __future__ import annotations

import json

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import networkx as nx
import numpy as np
from sklearn.neighbors import NearestNeighbors
from sklearn.preprocessing import StandardScaler

from src import config
from src.data.load import load_raw, split_xy
from src.models.scoring import load_artifacts

K = 5            # neighbours per flagged account
MIN_RING = 3     # a ring needs >=3 accounts
TIER_HEX = {"CRITICAL": "#ff4d4f", "HIGH": "#ff9f43", "MEDIUM": "#febc2e", "LOW": "#2bd97c"}


def _score_all():
    art = load_artifacts()
    X, y = split_xy(load_raw())
    feats = art.builder.transform(X)
    prob = art.model.predict_proba(feats)[:, 1]
    risk = np.array([config.prob_to_risk(p, art.threshold) for p in prob])
    return feats, risk, y.values


def build_ring_graph(feats, risk):
    """k-NN similarity graph among flagged accounts only."""
    fidx = np.where(risk >= 50)[0]
    # Native-NaN features: median-impute for the distance computation only
    # (StandardScaler / k-NN need finite input; the classifier still sees raw NaN).
    vals = feats.values.astype(float)
    col_med = np.nan_to_num(np.nanmedian(vals, axis=0), nan=0.0)
    nanpos = np.where(np.isnan(vals))
    vals[nanpos] = np.take(col_med, nanpos[1])
    Z = StandardScaler().fit_transform(vals)[fidx]
    k = min(K, len(fidx) - 1)
    dist, idx = NearestNeighbors(n_neighbors=k + 1).fit(Z).kneighbors(Z)
    thr = np.median(dist[:, 1:]) * 1.6  # only keep genuinely-similar edges
    G = nx.Graph()
    G.add_nodes_from(range(len(fidx)))
    for i in range(len(fidx)):
        for j, d in zip(idx[i, 1:], dist[i, 1:]):
            if d <= thr:
                G.add_edge(i, int(j), weight=1.0 / (1.0 + float(d)))
    return G, fidx


def detect_rings(G, fidx, risk, y):
    comms = nx.community.louvain_communities(G, weight="weight", seed=config.SEED)
    rings = []
    for cid, comm in enumerate(comms):
        local = sorted(comm)
        if len(local) < MIN_RING:
            continue
        gi = [int(fidx[t]) for t in local]
        rings.append({
            "ring_id": len(rings) + 1, "size": len(gi),
            "mean_risk": round(float(risk[gi].mean()), 1),
            "max_risk": round(float(risk[gi].max()), 1),
            "actual_mules": int(y[gi].sum()),
            "members": [f"ACC-{g:06d}" for g in gi],
            "_local": local,
        })
    rings.sort(key=lambda r: (-r["size"], -r["mean_risk"]))
    for rank, r in enumerate(rings):
        r["ring_id"] = rank + 1
    return rings


def render_and_export(G, fidx, rings, risk, y):
    pos = nx.spring_layout(G, seed=config.SEED, k=0.5, iterations=120)
    ring_of = {}
    for r in rings:
        for t in r["_local"]:
            ring_of[t] = r["ring_id"]

    palette = ["#2dd4bf", "#ff9f43", "#a78bfa", "#ff4d4f", "#38bdf8", "#f472b6"]
    fig, ax = plt.subplots(figsize=(11, 7.2), facecolor="#0b1119")
    ax.set_facecolor("#0b1119")
    for u, v in G.edges():
        ax.plot([pos[u][0], pos[v][0]], [pos[u][1], pos[v][1]],
                color="#33475f", lw=0.7, alpha=0.55, zorder=1)
    nodes = list(G.nodes())
    for r in rings:
        col = palette[(r["ring_id"] - 1) % len(palette)]
        rn = r["_local"]
        xs = [pos[n][0] for n in rn]; ys = [pos[n][1] for n in rn]
        sizes = [70 + risk[int(fidx[n])] * 3 for n in rn]
        ax.scatter(xs, ys, color=col, s=sizes, edgecolors="#0b1119", linewidths=0.8,
                   zorder=3, label=f"Ring {r['ring_id']} · {r['size']} accts")
    # unringed flagged nodes (singletons/pairs) in muted grey
    ringed = {n for r in rings for n in r["_local"]}
    other = [n for n in nodes if n not in ringed]
    if other:
        ax.scatter([pos[n][0] for n in other], [pos[n][1] for n in other],
                   color="#5b6b7f", s=50, zorder=2, label="unclustered")
    mx = [pos[n][0] for n in nodes if y[int(fidx[n])] == 1]
    my = [pos[n][1] for n in nodes if y[int(fidx[n])] == 1]
    ax.scatter(mx, my, s=320, facecolors="none", edgecolors="#e7eef6", linewidths=1.4,
               zorder=4, label="confirmed mule")
    ax.set_title(f"{len(rings)} mule rings surfaced among the {len(nodes)} flagged accounts\n"
                 "behavioural-similarity graph · Louvain community detection",
                 color="#e7eef6", fontsize=13.5, weight="bold")
    ax.legend(loc="upper left", facecolor="#131b27", edgecolor="#26344a", labelcolor="#e7eef6",
              fontsize=9, framealpha=0.9)
    ax.axis("off"); fig.tight_layout()
    fig.savefig(config.ROOT / "docs" / "mule_rings.png", dpi=150, facecolor="#0b1119", bbox_inches="tight")
    plt.close(fig)

    g_nodes = [{"id": int(fidx[n]), "acc": f"ACC-{int(fidx[n]):06d}",
                "x": float(pos[n][0]), "y": float(pos[n][1]),
                "risk": float(risk[int(fidx[n])]), "tier": config.risk_tier(risk[int(fidx[n])]),
                "mule": int(y[int(fidx[n])]), "ring": int(ring_of.get(n, 0))} for n in nodes]
    g_edges = [[int(u), int(v)] for u, v in G.edges()]
    (config.ARTIFACTS_DIR / "ring_graph.json").write_text(json.dumps({"nodes": g_nodes, "edges": g_edges}))


def main():
    config.ensure_dirs()
    feats, risk, y = _score_all()
    G, fidx = build_ring_graph(feats, risk)
    rings = detect_rings(G, fidx, risk, y)
    render_and_export(G, fidx, rings, risk, y)

    mules_in_rings = sum(r["actual_mules"] for r in rings)
    summary = {
        "n_flagged": int((risk >= 50).sum()), "total_mules": int(y.sum()),
        "n_rings": len(rings), "mules_in_rings": int(mules_in_rings),
        "largest_ring": rings[0]["size"] if rings else 0,
        "rings": [{k: v for k, v in r.items() if k != "_local"} for r in rings],
    }
    (config.ARTIFACTS_DIR / "rings.json").write_text(json.dumps(summary, indent=2))
    print(f"Rings: {len(rings)} | flagged: {summary['n_flagged']} | "
          f"mules in rings: {mules_in_rings}/{summary['total_mules']}")
    for r in rings[:6]:
        print(f"  Ring {r['ring_id']}: {r['size']} accts, mean_risk {r['mean_risk']}, "
              f"{r['actual_mules']} confirmed mules")


if __name__ == "__main__":
    main()