src/visualiser/pyvis_graph.py

"""
PyVis Graph Visualiser
Builds interactive network visualisations for account investigation.
"""
import json
import tempfile
import os
import streamlit.components.v1 as components
import pandas as pd
import networkx as nx
from pyvis.network import Network

# Constants
NODE_SIZE_CENTER = 35
NODE_SIZE_HIGH_RISK = 22
NODE_SIZE_NORMAL = 14
NODE_SIZE_HIGH_PAGERANK = 26
EDGE_WIDTH_MAX = 8
EDGE_WIDTH_SCALE = 1_000_000
PAGERANK_TOP_PCT = 0.01
PHYSICS_GRAVITY = -50
PHYSICS_SPRING = 100
PHYSICS_ITERATIONS = 150

# Color palette
COLOR_CENTER = '#FFD700'       # Gold for center node
COLOR_FRAUD = '#FF4136'        # Red for confirmed fraud
COLOR_HIGH_PR = '#FF851B'      # Orange for high PageRank
COLOR_NORMAL = '#0074D9'       # Blue for normal
COLOR_FRAUD_EDGE = '#FF4136'   # Red for fraudulent edges
COLOR_NORMAL_EDGE = '#AAAAAA'  # Grey for normal edges

PHYSICS_OPTIONS = json.dumps({
    "nodes": {"borderWidth": 2, "shadow": True},
    "edges": {
        "smooth": {"type": "curvedCW", "roundness": 0.2},
        "shadow": True,
        "arrows": {"to": {"enabled": True, "scaleFactor": 0.8}},
    },
    "physics": {
        "forceAtlas2Based": {
            "gravitationalConstant": PHYSICS_GRAVITY,
            "springLength": PHYSICS_SPRING,
        },
        "solver": "forceAtlas2Based",
        "stabilization": {"iterations": PHYSICS_ITERATIONS},
    },
    "interaction": {"hover": True, "tooltipDelay": 100},
})


def build_pyvis_graph(
    subgraph: nx.DiGraph,
    df: pd.DataFrame,
    center_node: str,
    fraud_accounts: set,
    pagerank_scores: dict,
    louvain_partition: dict,
) -> Network:
    """
    Build an interactive PyVis network from a NetworkX subgraph.

    Nodes are sized and colored based on their role (center, fraud, high-PR, normal).
    Edges are scaled by transaction amount and colored by fraud status.

    Returns:
        A pyvis.network.Network instance ready for rendering.
    """
    net = Network(
        height='570px',
        width='100%',
        directed=True,
        notebook=False,
    )

    net.set_options(PHYSICS_OPTIONS)

    # Compute pagerank threshold for highlighting top nodes
    all_pr = sorted(pagerank_scores.values(), reverse=True)
    top_n = max(1, int(len(all_pr) * PAGERANK_TOP_PCT))
    pagerank_threshold = all_pr[min(top_n, len(all_pr) - 1)]

    # --- Add nodes ---
    for node in subgraph.nodes():
        node_data = subgraph.nodes[node]
        is_center = (node == center_node)
        is_fraud = node in fraud_accounts
        is_high_pr = pagerank_scores.get(node, 0) >= pagerank_threshold

        # Size and shape
        if is_center:
            size = NODE_SIZE_CENTER
            shape = 'star'
            color = COLOR_CENTER
        elif is_high_pr:
            size = NODE_SIZE_HIGH_PAGERANK
            shape = 'diamond'
            color = COLOR_HIGH_PR
        elif is_fraud:
            size = NODE_SIZE_HIGH_RISK
            shape = 'dot'
            color = COLOR_FRAUD
        else:
            size = NODE_SIZE_NORMAL
            shape = 'dot'
            color = COLOR_NORMAL

        # Tooltip
        total_sent = node_data.get('total_sent', 0)
        total_received = node_data.get('total_received', 0)
        count_sent = node_data.get('count_sent', 0)
        community = louvain_partition.get(node, 'N/A')
        status = 'FLAGGED' if is_fraud else 'Normal'

        tooltip = (
            f"Account: {node}\n"
            f"Total Sent: {total_sent:,.0f}\n"
            f"Total Received: {total_received:,.0f}\n"
            f"Transactions: {count_sent}\n"
            f"Community: {community}\n"
            f"Status: {status}"
        )

        label = str(node)[:12]

        net.add_node(
            str(node),
            label=label,
            size=size,
            shape=shape,
            color=color,
            title=tooltip,
            borderWidth=2,
            borderWidthSelected=4,
        )

    # --- Add edges ---
    for src, tgt, data in subgraph.edges(data=True):
        amount = data.get('amount', 0)
        is_fraud_edge = data.get('is_laundering', 0) == 1

        edge_width = min(amount / EDGE_WIDTH_SCALE, EDGE_WIDTH_MAX)
        edge_width = max(edge_width, 0.5)

        edge_color = COLOR_FRAUD_EDGE if is_fraud_edge else COLOR_NORMAL_EDGE

        tooltip = (
            f"Amount: {amount:,.0f}\n"
            f"Channel: {data.get('payment_type', '')}\n"
            f"Suspicious: {'Yes' if is_fraud_edge else 'No'}"
        )

        net.add_edge(
            str(src),
            str(tgt),
            value=edge_width,
            title=tooltip,
            color=edge_color,
            arrows='to',
        )

    return net


def render_pyvis(net: Network) -> None:
    """
    Render a PyVis network inside a Streamlit app using an HTML component.
    """
    with tempfile.NamedTemporaryFile(delete=False, suffix='.html', mode='w') as f:
        tmp_path = f.name
    net.save_graph(tmp_path)
    with open(tmp_path, 'r') as f:
        html_content = f.read()
    os.unlink(tmp_path)
    components.html(html_content, height=580, scrolling=False)


def save_pyvis_html(net: Network, path: str) -> None:
    """
    Save a PyVis network to an HTML file on disk.
    """
    net.save_graph(path)