src/entity_graph_service.py

"""
Entity / Bank network intelligence — non-ML, structural graph layer built from
the accounts reference dataset (Bank Name, Bank ID, Account Number, Entity ID,
Entity Name). This is the single source of truth for entity-type derivation
and for building the Entity/Bank lookup indices cached on AppState.

Heterogeneous graph modeled here (no networkx object needed — the
Entity-OWNS-Account-HELD_AT-Bank relationship is a simple star/tree, so plain
dict indices give O(1) lookups without the overhead of a graph library):

    Entity --OWNS--> Account --HELD_AT--> Bank

Does not touch ML models, detectors, or risk scoring in any way.
"""
from __future__ import annotations

import pandas as pd

# Order matters — first matching prefix wins. Falls back to "Entity" for any
# unrecognized prefix (e.g. a handful of "Direct #N" rows observed in the data).
ENTITY_TYPE_PREFIXES = [
    ("Corporation", "Corporation"),
    ("Sole Proprietorship", "Sole Proprietorship"),
    ("Partnership", "Partnership"),
    ("Country", "Country"),
    ("Individual", "Individual"),
]

# Thresholds — shared by entity_graph_service and entity_exposure_service so
# the numbers never drift apart.
CROSS_BANK_MULE_THRESHOLD = 3   # entities at >= this many distinct banks are flagged "cross-bank"
MANY_ACCOUNTS_THRESHOLD = 5     # entities owning >= this many accounts are flagged "many accounts"
INSTITUTION_SPREAD_HIGH = 4     # bank_count >= this -> HIGH spread
INSTITUTION_SPREAD_MEDIUM = 2   # bank_count >= this (and < HIGH) -> MEDIUM spread

TOP_BANKS_LIMIT = 50            # default "top banks by volume" list size for the bank dashboard


def derive_entity_type(entity_name: str | None) -> str:
    """Derive a human-readable entity type from the entity_name prefix.
    Moved here from api/routes/investigation.py's get_account_intel so both
    the existing Account Intelligence endpoint and the new entity/bank
    endpoints share one definition."""
    ename = entity_name or ""
    for prefix, label in ENTITY_TYPE_PREFIXES:
        if ename.startswith(prefix):
            return label
    return "Entity"


def parse_bank_country(bank_name: str | None) -> str:
    """Best-effort country label from a bank name like 'Germany Bank #4815' -> 'Germany'.

    CAVEAT: roughly a quarter of the bank names in this dataset don't encode a
    country this way (e.g. 'Bank of New York', 'National Bank of Cleveland',
    'Savings Bank of Omaha' parse to junk like 'National' or ''). Callers
    should treat this as a best-effort issuer label, not verified geography —
    it is exposed in API responses as `country_label`, never `country`.
    """
    name = bank_name or ""
    idx = name.find(" Bank")
    if idx <= 0:
        return ""
    return name[:idx].strip()


def load_accounts_reference(csv_path: str) -> dict[str, dict]:
    """Load the accounts reference CSV into an O(1) account_number -> info dict.
    Moved here from server.py's background_init() inline block. Uses vectorized
    pandas instead of iterrows() for the ~500k-row scale."""
    df = pd.read_csv(
        csv_path,
        usecols=["Bank Name", "Bank ID", "Account Number", "Entity ID", "Entity Name"],
    )
    df = df.astype(str)
    df = df.rename(columns={
        "Bank Name": "bank_name",
        "Bank ID": "bank_id",
        "Entity ID": "entity_id",
        "Entity Name": "entity_name",
    })
    # A handful of account numbers repeat in the raw data; keep the last
    # occurrence, matching the dict-overwrite behavior of the original
    # row-by-row loop this function replaced.
    df = df.drop_duplicates(subset="Account Number", keep="last")
    return df.set_index("Account Number")[["bank_name", "bank_id", "entity_id", "entity_name"]].to_dict("index")


def build_entity_bank_indices(accounts_by_number: dict[str, dict]) -> tuple[dict, dict, dict]:
    """Single O(n) pass over the already-loaded accounts_by_number dict (no
    second CSV read). Returns (entities_by_id, banks_by_id, network_summary).

    Also tracks, per bank, how many of its entities own more than one account
    AT THAT BANK specifically (needed for the bank profile's
    multi_account_entity_count) — computed here, in the same pass, since this
    is the only place account-level (entity_id, bank_id) pairs are visible
    together.
    """
    entities: dict[str, dict] = {}
    banks: dict[str, dict] = {}
    # bank_id -> entity_id -> count of accounts that entity holds at that bank
    bank_entity_account_counts: dict[str, dict[str, int]] = {}

    for account_number, info in accounts_by_number.items():
        entity_id = info.get("entity_id")
        entity_name = info.get("entity_name", "")
        bank_id = info.get("bank_id")
        bank_name = info.get("bank_name", "")

        ent = entities.get(entity_id)
        if ent is None:
            ent = {
                "entity_id": entity_id,
                "entity_name": entity_name,
                "entity_type": derive_entity_type(entity_name),
                "account_numbers": [],
                "bank_ids": set(),
                "bank_names": set(),
            }
            entities[entity_id] = ent
        ent["account_numbers"].append(account_number)
        ent["bank_ids"].add(bank_id)
        ent["bank_names"].add(bank_name)

        bank = banks.get(bank_id)
        if bank is None:
            bank = {
                "bank_id": bank_id,
                "bank_name": bank_name,
                "country_label": parse_bank_country(bank_name),
                "account_numbers": [],
                "entity_ids": set(),
            }
            banks[bank_id] = bank
        bank["account_numbers"].append(account_number)
        bank["entity_ids"].add(entity_id)

        per_entity_counts = bank_entity_account_counts.setdefault(bank_id, {})
        per_entity_counts[entity_id] = per_entity_counts.get(entity_id, 0) + 1

    # Freeze sets into sorted lists for stable, JSON-friendly output
    for ent in entities.values():
        ent["bank_ids"] = sorted(ent["bank_ids"])
        ent["bank_names"] = sorted(ent["bank_names"])
    for bank_id, bank in banks.items():
        bank["entity_ids"] = sorted(bank["entity_ids"])
        bank["multi_account_entity_count"] = sum(
            1 for count in bank_entity_account_counts.get(bank_id, {}).values() if count > 1
        )

    entity_type_counts: dict[str, int] = {}
    multi_account_entities = 0
    cross_bank_entities = 0
    for ent in entities.values():
        entity_type_counts[ent["entity_type"]] = entity_type_counts.get(ent["entity_type"], 0) + 1
        if len(ent["account_numbers"]) > 1:
            multi_account_entities += 1
        if len(ent["bank_ids"]) >= CROSS_BANK_MULE_THRESHOLD:
            cross_bank_entities += 1

    network_summary = {
        "total_entities": len(entities),
        "total_accounts": len(accounts_by_number),
        "total_banks": len(banks),
        "multi_account_entities": multi_account_entities,
        "cross_bank_entities": cross_bank_entities,
        "entity_type_counts": entity_type_counts,
    }

    return entities, banks, network_summary


def build_bank_profiles(entities_by_id: dict, banks_by_id: dict) -> tuple[dict[str, dict], list[str]]:
    """Per-bank aggregate stats: account/entity counts, entity-type breakdown,
    and multi-account-entity count (read off the precomputed field set by
    build_entity_bank_indices). Also returns `top_banks_by_volume` — the top
    TOP_BANKS_LIMIT bank_ids by account_count — for the frontend bank
    dashboard's default listing, since there are tens of thousands of banks,
    far too many to list unranked."""
    entity_type_by_id = {eid: ent["entity_type"] for eid, ent in entities_by_id.items()}

    profiles: dict[str, dict] = {}
    for bank_id, bank in banks_by_id.items():
        corp = partner = sole = 0
        for eid in bank["entity_ids"]:
            etype = entity_type_by_id.get(eid, "Entity")
            if etype == "Corporation":
                corp += 1
            elif etype == "Partnership":
                partner += 1
            elif etype == "Sole Proprietorship":
                sole += 1

        profiles[bank_id] = {
            "bank_id": bank_id,
            "bank_name": bank["bank_name"],
            "country_label": bank["country_label"],
            "account_count": len(bank["account_numbers"]),
            "entity_count": len(bank["entity_ids"]),
            "corporation_count": corp,
            "partnership_count": partner,
            "sole_proprietorship_count": sole,
            "multi_account_entity_count": bank.get("multi_account_entity_count", 0),
        }

    top_banks_by_volume = sorted(
        profiles.keys(), key=lambda bid: profiles[bid]["account_count"], reverse=True
    )[:TOP_BANKS_LIMIT]

    return profiles, top_banks_by_volume


def compute_entity_risk(account_numbers: list[str], features_by_account: dict) -> dict:
    """Worst-account rule: an entity's risk is the highest risk among the
    accounts it owns. Prefers the most-refined score available per account
    (hybrid_score > gnn_fraud_score > the base XGBoost risk_score) so this
    automatically improves once Phase 2 (GNN/hybrid) finishes in the
    background — it reads whatever AppState.features_by_account has at call
    time, no separate model invocation.

    risk_tier is NOT recomputed from a static score cutoff here — it's read
    straight off the risk-driver account's own precomputed risk_tier (see
    src/risk_tiers.py, set in server.py alongside risk_score), which is a
    percentile-rank classification, not a fixed score>=N threshold. This
    keeps exactly one place in the codebase deciding what counts as
    CRITICAL/HIGH/MEDIUM/LOW."""
    best_score = 0
    best_account = None
    scored = 0
    for account_number in account_numbers:
        feat = features_by_account.get(account_number)
        if not feat:
            continue
        scored += 1
        score = feat.get("risk_score", 0) or 0
        hybrid = feat.get("hybrid_score")
        gnn = feat.get("gnn_fraud_score")
        if hybrid:
            score = max(score, round(hybrid * 100))
        elif gnn:
            score = max(score, round(gnn * 100))
        if score > best_score or best_account is None:
            best_score, best_account = score, account_number
    driver_feat = features_by_account.get(best_account) or {}
    return {
        "risk_score": best_score,
        "risk_tier": driver_feat.get("risk_tier", "LOW"),
        "risk_driver_account": best_account,
        "accounts_scored": scored,
    }


def attach_risk_aggregates(
    entities_by_id: dict,
    banks_by_id: dict,
    bank_profiles_cache: dict,
    features_by_account: dict,
) -> dict:
    """Precomputes entity- and bank-level risk once at startup (O(1) reads
    thereafter) by reusing the existing model's account-level scores —
    introduces no new scoring logic, only aggregation over what's already in
    features_by_account. Mutates entities_by_id / banks_by_id /
    bank_profiles_cache in place and returns summary counters to merge into
    network_summary_cache."""
    high_risk_entity_count = 0
    bank_high_risk_counts: dict[str, int] = {}

    for entity_id, ent in entities_by_id.items():
        risk = compute_entity_risk(ent["account_numbers"], features_by_account)
        ent["risk_score"] = risk["risk_score"]
        ent["risk_tier"] = risk["risk_tier"]
        ent["risk_driver_account"] = risk["risk_driver_account"]
        if risk["risk_tier"] in ("HIGH", "CRITICAL"):
            high_risk_entity_count += 1
            for bank_id in ent["bank_ids"]:
                bank_high_risk_counts[bank_id] = bank_high_risk_counts.get(bank_id, 0) + 1

    for bank_id, bank in banks_by_id.items():
        bank_risk_score = 0
        bank_risk_tier = "LOW"
        for entity_id in bank["entity_ids"]:
            ent = entities_by_id.get(entity_id)
            if ent and ent.get("risk_score", 0) > bank_risk_score:
                bank_risk_score = ent["risk_score"]
                bank_risk_tier = ent.get("risk_tier", "LOW")  # inherit, don't reclassify
        bank["risk_score"] = bank_risk_score
        bank["risk_tier"] = bank_risk_tier
        bank["high_risk_entity_count"] = bank_high_risk_counts.get(bank_id, 0)
        profile = bank_profiles_cache.get(bank_id)
        if profile is not None:
            profile["risk_score"] = bank_risk_score
            profile["risk_tier"] = bank["risk_tier"]
            profile["high_risk_entity_count"] = bank["high_risk_entity_count"]

    return {"high_risk_entity_count": high_risk_entity_count}