visualize.py

import json
import math
import logging
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image, ImageDraw, ImageOps
from matplotlib.offsetbox import OffsetImage, AnnotationBbox
from scipy.stats import gaussian_kde
from scipy.spatial import cKDTree
from pathlib import Path

from config import load_environment, ROOT

# Configure logging
logger = logging.getLogger(__name__)

config = load_environment()
json_folder = config['json_folder']
cache_dir = config['cache_dir']
MY_ID = config.get('my_id')

CANVAS_SIZE = tuple(
    float(p.strip())
    for p in config['canvas_size'].split(',')
)
DPI = config['dpi']
AVATAR_PIXELS = config['avatar_pixels']
AVATAR_ZOOM = config['avatar_zoom']
PEER_WEIGHT = config['peer_weight']
HERO_WEIGHT = config['hero_weight']
HEATMAP_GRID = config['heatmap_grid']

BLURPLE_HERO = "#8C9CD5"
PEER_COLOR = '#efefef'
DARK_GREY = '#1a1a1a' 

def get_processed_img(node_id: str, is_srv: bool = False, border_color: str = '#5865F2'):
    """
    Load, resize, and process avatar/icon image with circular or rounded-square mask.

    Args:
        node_id: Discord user or server ID
        is_srv: True if this is a server icon (rounded square), False for user avatar (circle)
        border_color: Hex color for border

    Returns:
        PIL Image with transparency, or None if file not found
    """
    prefix = "srv" if is_srv else "usr"
    path = cache_dir / f"{prefix}_{node_id}.png"

    if not path.exists():
        logger.debug(f"Image not found: {path}")
        return None

    try:
        img = Image.open(path).convert("RGBA")
        size = (AVATAR_PIXELS, AVATAR_PIXELS)
        img = img.resize(size, Image.LANCZOS)

        # Create mask for circular (user) or rounded-square (server) shape
        mask = Image.new("L", size, 0)
        draw = ImageDraw.Draw(mask)

        if is_srv:
            # Rounded square for servers
            radius = max(16, int(AVATAR_PIXELS / 16))
            draw.rounded_rectangle((0, 0) + size, radius=radius, fill=255)
        else:
            # Perfect circle for users
            draw.ellipse((0, 0) + size, fill=255)

        # Apply mask to image
        output = ImageOps.fit(img, mask.size, centering=(0.5, 0.5))
        output.putalpha(mask)

        # Add border
        draw_border = ImageDraw.Draw(output)
        border_width = max(4, int(AVATAR_PIXELS / 32))

        if is_srv:
            radius = max(8, int(AVATAR_PIXELS / 6))
            draw_border.rounded_rectangle(
                (0, 0) + size, radius=radius, outline=border_color, width=border_width
            )
        else:
            draw_border.ellipse((0, 0) + size, outline=border_color, width=border_width)

        return output

    except Exception as e:
        logger.error(f"Failed to process image {path}: {e}")
        return None

def add_heatmap(ax, pos):
    x = np.array([p[0] for p in pos.values()])
    y = np.array([p[1] for p in pos.values()])
    xy = np.vstack([x, y])
    xi, yi = np.mgrid[x.min()-1:x.max()+1:complex(HEATMAP_GRID), y.min()-1:y.max()+1:complex(HEATMAP_GRID)]
    zi = gaussian_kde(xy)(np.vstack([xi.flatten(), yi.flatten()]))
    ax.pcolormesh(xi, yi, zi.reshape(xi.shape), shading='auto', cmap='magma', alpha=0.12, zorder=-10)

def resolve_overlaps(
    pos: dict, min_dist: float = 0.18, max_iters: int = 300, damping: float = 0.6
) -> dict:
    """
    Resolve overlapping nodes using spatial indexing (O(n log n) per iteration).

    This replaces the naive O(n²) approach with cKDTree for faster performance
    on graphs with >100 nodes.

    Args:
        pos: Dictionary mapping node IDs to (x, y) coordinates
        min_dist: Minimum distance before nodes are considered overlapping
        max_iters: Maximum number of iterations
        damping: Fraction of movement to apply (0-1, lower = more stable)

    Returns:
        Updated position dictionary with resolved overlaps
    """
    # Convert numpy arrays to lists for easier manipulation
    for k, v in pos.items():
        if isinstance(v, np.ndarray):
            pos[k] = v.tolist()

    nodes = list(pos.keys())
    n_nodes = len(nodes)

    if n_nodes < 2:
        return pos

    logger.debug(f"Resolving overlaps for {n_nodes} nodes (max_iters={max_iters})...")

    for iteration in range(max_iters):
        # Build KD-tree for spatial indexing
        points = np.array([pos[node] for node in nodes])
        tree = cKDTree(points)

        moved = False

        # Query neighbors within min_dist for each node
        for i, node_u in enumerate(nodes):
            neighbors_idx = tree.query_ball_point(points[i], r=min_dist)

            for j in neighbors_idx:
                if i >= j:  # Avoid duplicate work and self-comparison
                    continue

                node_v = nodes[j]
                ux, uy = pos[node_u]
                vx, vy = pos[node_v]
                dx = ux - vx
                dy = uy - vy
                dist = math.hypot(dx, dy)

                # Handle coincident points
                if dist < 1e-6:
                    dx = (np.random.rand() - 0.5) * 1e-2
                    dy = (np.random.rand() - 0.5) * 1e-2
                    dist = math.hypot(dx, dy)

                if dist < min_dist:
                    overlap = (min_dist - dist) / 2.0
                    nx_disp = dx / dist * overlap
                    ny_disp = dy / dist * overlap

                    pos[node_u][0] += nx_disp * damping
                    pos[node_u][1] += ny_disp * damping
                    pos[node_v][0] -= nx_disp * damping
                    pos[node_v][1] -= ny_disp * damping

                    moved = True

        if not moved:
            logger.debug(f"Overlap resolution converged after {iteration + 1} iterations")
            break

    return pos

def render_graph(G: nx.Graph, pos: dict, filename: Path, show_servers: bool = True):
    """
    Render a network graph to a PNG file.

    Args:
        G: NetworkX graph object
        pos: Dictionary mapping nodes to (x, y) coordinates
        filename: Output file path
        show_servers: If False, hides server nodes and edges
    """
    fig = plt.figure(figsize=CANVAS_SIZE, facecolor=DARK_GREY)
    ax = plt.gca()
    
    # Remove all borders and margins
    ax.set_facecolor(DARK_GREY)
    for spine in ax.spines.values():
        spine.set_visible(False)
    ax.margins(0)
    
    add_heatmap(ax, pos)

    # Draw edges
    for u, v, d in G.edges(data=True):
        is_me = u == MY_ID or v == MY_ID
        is_srv_edge = d.get('edge_type') == 'server'

        if not show_servers and is_srv_edge:
            continue

        color = BLURPLE_HERO if is_me else PEER_COLOR
        alpha = 0.6 if is_me else 0.25
        width = HERO_WEIGHT if is_me else PEER_WEIGHT
        style = 'dashed' if (is_srv_edge and is_me) else ('solid' if is_me else 'dotted')

        nx.draw_networkx_edges(
            G, pos, edgelist=[(u, v)], width=width, edge_color=color,
            style=style, alpha=alpha, ax=ax
        )

    # Draw nodes
    for node, (x, y) in pos.items():
        d = G.nodes[node]

        if not show_servers and d['type'] == 'server':
            continue

        if d.get('img'):
            oi = OffsetImage(d['img'], zoom=AVATAR_ZOOM)
            ax.add_artist(AnnotationBbox(oi, (x, y), frameon=False))
        else:
            marker = 's' if d['type'] == 'server' else 'o'
            ax.scatter(
                x, y, s=1200, c=BLURPLE_HERO, marker=marker,
                edgecolors='white', alpha=0.8
            )

    # Add labels
    for node, (x, y) in pos.items():
        if G.nodes[node]['type'] != 'user':
            continue

        label = G.nodes[node].get('label', '')
        if not label:
            continue

        offset_points = -36 if G.nodes[node].get('img') else -18
        bbox = dict(boxstyle="round,pad=0.35", facecolor=(0, 0, 0, 0.6), edgecolor='none')

        ax.annotate(
            label, xy=(x, y), xytext=(0, offset_points), textcoords='offset points',
            ha='center', color='white', fontsize=11, fontweight='bold', alpha=0.95,
            bbox=bbox, clip_on=False, zorder=5
        )

    plt.axis('off')
    plt.savefig(filename, dpi=DPI, bbox_inches='tight', pad_inches=0, facecolor=DARK_GREY)
    plt.close()
    logger.info(f"Saved visualization to {filename}")

if __name__ == '__main__':
    logger.info("Building relationship graph from friend profiles...")

    G = nx.Graph()

    # Load all friend profiles and build graph
    profile_count = 0
    for json_file in json_folder.glob('*.json'):
        try:
            with open(json_file, 'r') as f:
                data = json.load(f)

            u = data.get('user', {})
            uid = u.get('id')

            if not uid:
                logger.warning(f"No user ID in {json_file.name}")
                continue

            accent = f"#{u.get('accent_color'):06x}" if u.get('accent_color') else BLURPLE_HERO

            G.add_node(
                uid,
                label=u.get('global_name') or u.get('username'),
                img=get_processed_img(uid, False, accent),
                type='user'
            )

            # Add mutual friends
            for f_d in data.get('mutual_friends', []):
                fid = f_d.get('id')
                if not fid:
                    continue

                if fid not in G:
                    G.add_node(
                        fid,
                        type='user',
                        img=None,
                        label=f_d.get('global_name') or f_d.get('username')
                    )

                G.add_edge(uid, fid, edge_type='friend')

            # Add mutual guilds
            for g_d in data.get('mutual_guilds', []):
                gid = g_d.get('id')
                if not gid:
                    continue

                if gid not in G:
                    G.add_node(
                        gid,
                        type='server',
                        img=get_processed_img(gid, True, BLURPLE_HERO)
                    )

                G.add_edge(uid, gid, edge_type='server')

            profile_count += 1

        except json.JSONDecodeError as e:
            logger.error(f"Invalid JSON in {json_file.name}: {e}")
        except Exception as e:
            logger.error(f"Error processing {json_file.name}: {e}")

    logger.info(f"Built graph from {profile_count} profiles: {G.number_of_nodes()} nodes, {G.number_of_edges()} edges")

    # Calculate layout
    logger.info("Calculating social physics (spring layout)...")

    n_nodes = G.number_of_nodes()
    initial_pos = {}

    if MY_ID in G:
        # Pin the user at the center
        initial_pos[MY_ID] = np.array([0.0, 0.0])
        neighbors = list(G.neighbors(MY_ID))
        m = len(neighbors)

        for idx, nb in enumerate(neighbors):
            angle = 2 * np.pi * idx / m if m > 0 else 0
            initial_pos[nb] = np.array([math.cos(angle) * 1.0, math.sin(angle) * 1.0])

        # Random positions for other nodes
        for node in G.nodes():
            if node not in initial_pos:
                initial_pos[node] = np.random.randn(2) * 2.0

        k_val = 1.2 if n_nodes < 50 else 4.8
        min_dist = 0.7 * max(0.25, 50.0 / max(50, n_nodes))

        pos = nx.spring_layout(
            G, pos=initial_pos, fixed=[MY_ID], k=k_val,
            iterations=1250, seed=42, threshold=1e-5
        )
        logger.debug(f"Spring layout with MY_ID pinned, k={k_val}, min_dist={min_dist}")
    else:
        logger.warning("MY_ID not found in graph; using unpinned layout")
        pos = nx.spring_layout(G, k=4.8, iterations=1250, seed=42, threshold=1e-5)
        min_dist = 0.6

    pos = resolve_overlaps(pos, min_dist=min_dist, max_iters=400, damping=0.6)

    # Render visualizations
    output_dir = ROOT / 'output'
    output_dir.mkdir(exist_ok=True)

    render_graph(G, pos, output_dir / "servers-and-friends.png", show_servers=True)
    render_graph(G, pos, output_dir / "friends-only.png", show_servers=False)

    # Print analysis
    if MY_ID and MY_ID in G:
        logger.info("\n--- Top Mutual Connections ---")

        my_neighbors = set(G.neighbors(MY_ID))
        my_friend_neighbors = set(
            n for n in my_neighbors if G.nodes[n]['type'] == 'user'
        )
        my_server_neighbors = set(
            n for n in my_neighbors if G.nodes[n]['type'] == 'server'
        )

        mutual_friends_counts = []
        mutual_servers_counts = []

        for node in G.nodes():
            if G.nodes[node]['type'] != 'user' or node == MY_ID:
                continue

            neigh = set(G.neighbors(node))
            mutual_friends = len(neigh & my_friend_neighbors)
            mutual_servers = len(neigh & my_server_neighbors)

            if mutual_friends > 0:
                mutual_friends_counts.append((node, mutual_friends))
            if mutual_servers > 0:
                mutual_servers_counts.append((node, mutual_servers))

        top_mutual_friends = sorted(
            mutual_friends_counts, key=lambda x: x[1], reverse=True
        )[:3]
        top_mutual_servers = sorted(
            mutual_servers_counts, key=lambda x: x[1], reverse=True
        )[:3]

        if top_mutual_friends:
            logger.info("Top mutual friends:")
            for node, cnt in top_mutual_friends:
                label = G.nodes[node].get('label', node)
                logger.info(f"  {label}: {cnt} mutual friends")
        else:
            logger.info("No mutual friends found.")

        if top_mutual_servers:
            logger.info("Top mutual servers:")
            for node, cnt in top_mutual_servers:
                label = G.nodes[node].get('label', node)
                logger.info(f"  {label}: {cnt} mutual servers")
        else:
            logger.info("No mutual servers found.")
    else:
        logger.warning("MY_ID not set or not found in graph; skipping mutual analysis")

    logger.info("Visualization complete!")