src/entity_cluster_service.py

"""
Entity Cluster / Bridge Intelligence — read-only aggregation over the
ALREADY-COMPUTED Louvain community partition and betweenness centrality
scores (both loaded from data/checkpoints/*.pkl at startup, see server.py).

This module introduces no new graph computation: it only re-projects two
existing per-account signals (community_id, betweenness) onto the
Entity/Bank structural layer built by entity_graph_service.py, exactly the
way entity_exposure_service.py re-projects the transaction graph onto an
entity's blast radius. Never touches src/ml/ or src/detectors/.
"""
from __future__ import annotations

from src.entity_graph_service import INSTITUTION_SPREAD_MEDIUM


def attach_cluster_aggregates(
    entities_by_id: dict,
    accounts_by_number: dict,
    louvain_partition: dict,
    betweenness_scores: dict,
) -> None:
    """For each entity, derive:
    - community_id: the majority-vote Louvain community among its owned
      accounts (accounts with no partition entry are skipped).
    - community_confidence: fraction of scored accounts that agreed with
      the winning community (1.0 for single-account entities).
    - bridge_score: the max betweenness centrality among its owned accounts
      — a structural "how much does this entity sit between communities"
      signal, independent of risk scoring.
    Mutates entities_by_id in place. Entities with no matching accounts in
    either signal get community_id=None / bridge_score=0.0.
    """
    for ent in entities_by_id.values():
        votes: dict[str, int] = {}
        best_betweenness = 0.0
        for account_number in ent["account_numbers"]:
            community_id = louvain_partition.get(account_number)
            if community_id is not None:
                votes[community_id] = votes.get(community_id, 0) + 1
            score = betweenness_scores.get(account_number)
            if score is not None and score > best_betweenness:
                best_betweenness = score

        if votes:
            winning_community, winning_count = max(votes.items(), key=lambda kv: kv[1])
            # Stringify — Louvain community ids are plain ints; routes/JSON
            # treat community_id as an opaque string key throughout.
            ent["community_id"] = str(winning_community)
            ent["community_confidence"] = round(winning_count / sum(votes.values()), 3)
        else:
            ent["community_id"] = None
            ent["community_confidence"] = 0.0

        ent["bridge_score"] = round(best_betweenness, 6)


def build_clusters_index(entities_by_id: dict) -> dict[str, dict]:
    """Group entities by their attached community_id into per-cluster
    aggregates — the "Mule Ring" unit surfaced by /network/clusters.

    `risk_concentration` is a heuristic (share of member entities flagged
    HIGH/CRITICAL by the existing worst-account risk rule), not a verified
    fraud rate — there is no reliable ground truth at this granularity.
    """
    clusters: dict[str, dict] = {}

    for entity_id, ent in entities_by_id.items():
        community_id = ent.get("community_id")
        if community_id is None:
            continue
        cluster = clusters.get(community_id)
        if cluster is None:
            cluster = {
                "community_id": community_id,
                "entity_ids": [],
                "entity_count": 0,
                "account_count": 0,
                "bank_ids": set(),
                "high_risk_entity_count": 0,
            }
            clusters[community_id] = cluster
        cluster["entity_ids"].append(entity_id)
        cluster["entity_count"] += 1
        cluster["account_count"] += len(ent["account_numbers"])
        cluster["bank_ids"].update(ent["bank_ids"])
        if ent.get("risk_tier") in ("HIGH", "CRITICAL"):
            cluster["high_risk_entity_count"] += 1

    for cluster in clusters.values():
        cluster["bank_ids"] = sorted(cluster["bank_ids"])
        cluster["bank_count"] = len(cluster["bank_ids"])
        cluster["risk_concentration"] = round(
            cluster["high_risk_entity_count"] / cluster["entity_count"], 3
        ) if cluster["entity_count"] else 0.0

    return clusters


def rank_bridge_entities(entities_by_id: dict, limit: int = 50) -> list[dict]:
    """Entities that are both structural bridges (high betweenness among
    their owned accounts) AND already spread across >= INSTITUTION_SPREAD_MEDIUM
    banks — the "Cross-Bank Bridge" leaderboard. Sorted by bridge_score desc.
    """
    candidates = [
        ent for ent in entities_by_id.values()
        if ent.get("bridge_score", 0.0) > 0.0 and len(ent["bank_ids"]) >= INSTITUTION_SPREAD_MEDIUM
    ]
    candidates.sort(key=lambda ent: ent["bridge_score"], reverse=True)
    return candidates[:limit]