src/graph_builder.py

"""
================================================================================
GRAPH BUILDER MODULE - Transaction Network Analysis
================================================================================

PURPOSE:
  Constructs and analyzes transaction flow graphs for AML investigations.
  Transforms raw transaction data into NetworkX MultiDiGraph structures with
  sophisticated analysis capabilities (PageRank, betweenness centrality,
  community detection, temporal analysis).

KEY RESPONSIBILITIES:
  1. build_graph()              - Convert transaction DataFrame to MultiDiGraph
  2. attach_node_features()     - Attach account-level features to nodes
  3. compute_pagerank()         - Calculate importance scores weighted by amounts
  4. compute_betweenness()      - Find high-traffic intermediary accounts
  5. compute_louvain()          - Identify account clusters via community detection
  6. to_simple_graph()          - Collapse multi-edges by summing transaction amounts
  7. get_subgraph()             - Extract N-hop neighborhood trimmed by transaction volume
  8. get_laundering_related_nodes() - Expand nodes to include complete fraud schemes
  9. build_temporal_graph()     - Generate time-sorted edge list for animations
  10. compute_graph_stats()     - Calculate network density, degree distribution

DESIGN DECISIONS:
  - MultiDiGraph retains parallel transactions (same source→target, different amounts)
  - Edge weights = transaction amounts (supports value-based analysis)
  - Betweenness uses proportional k (10-100) for performance on large graphs
  - Subgraph trimming prioritizes transaction volume over node degree
  - LRU cache (maxsize=1) on centrality functions prevents recomputation

DEPENDENCIES:
  - networkx: Graph structures and centrality algorithms
  - python-louvain: Community detection for account clustering
  - pandas: DataFrame operations and temporal analysis

USAGE EXAMPLE:
  G = build_graph(transactions_df)
  G = attach_node_features(G, account_features_df)
  pagerank_scores = compute_pagerank(id(G), G)
  sub = get_subgraph(G, center_account_id, hops=2, max_nodes=60)
================================================================================
"""
import functools
import networkx as nx
import pandas as pd
import community as community_louvain

from src.config_loader import get_config


def build_graph(df: pd.DataFrame) -> nx.MultiDiGraph:
    """Build a directed multigraph — each transaction is its own edge."""
    G = nx.MultiDiGraph()
    for _, row in df.iterrows():
        G.add_edge(
            row['source'], row['target'],
            amount=row['amount'],
            payment_type=row['payment_type'],
            is_laundering=int(row['is_laundering']),
            timestamp=str(row['timestamp']),
        )
    return G


def attach_node_features(G: nx.MultiDiGraph, node_features_df: pd.DataFrame) -> nx.MultiDiGraph:
    if 'account' in node_features_df.columns:
        feature_dict = node_features_df.set_index('account').to_dict('index')
    else:
        feature_dict = node_features_df.to_dict('index')
    nx.set_node_attributes(G, feature_dict)
    return G


@functools.lru_cache(maxsize=1)
def compute_pagerank(_G_id: int, _G_ref) -> dict:
    """Compute PageRank weighted by transaction amounts. _G_id is a dummy used only for cache key."""
    cfg = get_config()['graph']
    # Convert to simple DiGraph for PageRank (weight = sum of amounts)
    Gs = nx.DiGraph()
    for u, v, data in _G_ref.edges(data=True):
        if Gs.has_edge(u, v):
            Gs[u][v]['weight'] += data.get('amount', 0)
        else:
            Gs.add_edge(u, v, weight=data.get('amount', 0))
    return nx.pagerank(Gs, alpha=cfg['pagerank_alpha'], max_iter=cfg['pagerank_max_iter'])


@functools.lru_cache(maxsize=1)
def compute_betweenness(_G_id: int, _G_ref) -> dict:
    n = _G_ref.number_of_nodes()
    k = min(100, max(10, n))  # Reduced from 500 for faster startup
    Gs = nx.DiGraph(_G_ref)  # collapse to simple for betweenness
    return nx.betweenness_centrality(Gs, k=k, normalized=True)


@functools.lru_cache(maxsize=1)
def compute_louvain(_G_id: int, _G_ref) -> dict:
    G_simple = nx.Graph()
    for u, v, data in _G_ref.edges(data=True):
        w = data.get('amount', 1)
        if G_simple.has_edge(u, v):
            G_simple[u][v]['weight'] += w
        else:
            G_simple.add_edge(u, v, weight=w)
    return community_louvain.best_partition(G_simple, weight='weight')


def to_simple_graph(G: nx.MultiDiGraph) -> nx.DiGraph:
    """Collapse multiedges into a DiGraph summing amounts."""
    Gs = nx.DiGraph()
    for u, v, data in G.edges(data=True):
        if Gs.has_edge(u, v):
            Gs[u][v]['amount'] += data.get('amount', 0)
            Gs[u][v]['tx_count'] = Gs[u][v].get('tx_count', 1) + 1
        else:
            Gs.add_edge(u, v, amount=data.get('amount', 0), tx_count=1,
                        is_laundering=data.get('is_laundering', 0),
                        payment_type=data.get('payment_type', ''))
    return Gs


def get_laundering_related_nodes(G: nx.MultiDiGraph, nodes: set[str]) -> set[str]:
    """
    Finds other nodes connected to the given set of nodes that are part of the
    same laundering patterns (cycles, mule networks) or are connected via
    transactions explicitly marked as laundering (is_laundering == 1).
    """
    # Import AppState inline to prevent circular dependencies
    from src.state import AppState
    extra_nodes = set()
    alerts = getattr(AppState, 'alerts', None) or []

    # Pre-extract all laundering groups from AppState.alerts to ensure members
    # of the same scheme (e.g., all participants in a round-trip or mule network)
    # are included together when visualizing related accounts.
    groups = []
    for alert in alerts:
        sub_alerts = alert.get('sub_alerts', [alert])
        for sa in sub_alerts:
            group = set()
            # Round Tripping (cycle) members
            if 'cycle_members' in sa:
                group.update(sa['cycle_members'])
            # Mule Network members & bridges
            if 'member_accounts' in sa:
                group.update(sa['member_accounts'])
            if 'bridge_accounts' in sa:
                group.update(sa['bridge_accounts'])

            if group:
                group.add(sa.get('account'))
                group.add(alert.get('account'))
                groups.append(group)

    # Iteratively expand current_nodes: if any laundering group partially overlaps
    # with current_nodes, include all members. This ensures network visualization
    # shows complete suspicion patterns, not isolated pieces.
    current_nodes = set(nodes)
    changed = True
    while changed:
        changed = False
        for group in groups:
            if not group.issubset(current_nodes) and (group & current_nodes):
                current_nodes.update(group)
                extra_nodes.update(group)
                changed = True

    # Check edges marked is_laundering==1 to include counterparties involved in
    # explicitly-detected suspicious transactions, even if they're not part of
    # the same alert group.
    for n in list(current_nodes):
        if n in G:
            # Check outgoing edges
            for nbr in G.successors(n):
                if nbr not in current_nodes:
                    for edge_data in G[n][nbr].values():
                        if int(edge_data.get('is_laundering', 0)) == 1:
                            extra_nodes.add(nbr)
                            current_nodes.add(nbr)
                            break
            # Check incoming edges
            for nbr in G.predecessors(n):
                if nbr not in current_nodes:
                    for edge_data in G[nbr][n].values():
                        if int(edge_data.get('is_laundering', 0)) == 1:
                            extra_nodes.add(nbr)
                            current_nodes.add(nbr)
                            break

    return extra_nodes


def get_subgraph(G: nx.MultiDiGraph, center_node: str, hops: int = 2, max_nodes: int = 60) -> nx.DiGraph:
    """BFS subgraph trimmed by edge weight (not raw degree)."""
    if center_node not in G:
        return None

    nodes_to_include = {center_node}
    frontier = {center_node}

    # BFS to collect nodes within N hops
    for _ in range(hops):
        next_frontier = set()
        for n in frontier:
            next_frontier |= set(G.successors(n))
            next_frontier |= set(G.predecessors(n))
        nodes_to_include |= next_frontier
        frontier = next_frontier
        if len(nodes_to_include) >= max_nodes:
            break

    nodes_list = list(nodes_to_include)
    # If subgraph exceeds max_nodes, prune by total transaction volume (sum of edge amounts).
    # This preserves high-activity accounts over peripheral ones, more meaningful than
    # pruning by node degree which could retain suspicious coordinators.
    if len(nodes_list) > max_nodes:
        def node_weight(n):
            return sum(d.get('amount', 0) for _, _, d in G.edges(n, data=True))
        nodes_list.sort(key=node_weight, reverse=True)
        nodes_list = nodes_list[:max_nodes]
        # Always include the center node regardless of weight
        if center_node not in nodes_list:
            nodes_list[-1] = center_node

    sub = G.subgraph(nodes_list)
    return to_simple_graph(sub)


def build_temporal_graph(df: pd.DataFrame) -> list:
    """Return time-sorted edge list for fund-trail animation."""
    df_sorted = df.sort_values('timestamp')
    return [
        {
            'source': row['source'],
            'target': row['target'],
            'amount': row['amount'],
            'timestamp': str(row['timestamp']),
            'payment_type': row.get('payment_type', ''),
            'is_laundering': int(row.get('is_laundering', 0)),
        }
        for _, row in df_sorted.iterrows()
    ]


def compute_graph_stats(G: nx.MultiDiGraph) -> dict:
    Gs = to_simple_graph(G)
    node_count = Gs.number_of_nodes()
    edge_count = Gs.number_of_edges()
    density = nx.density(Gs)
    avg_degree = sum(d for _, d in Gs.degree()) / max(node_count, 1)
    return {
        'node_count': node_count,
        'edge_count': edge_count,
        'density': round(density, 6),
        'avg_degree': round(avg_degree, 2),
    }