src/socialdynamics/visualization.py

from __future__ import annotations

import math

import networkx as nx
import numpy as np
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from .metrics import NullTest
from .simulation import OpinionSimulation

# ── Dark minimal palette ────────────────────────────────────────
PAPER = "#0C0D0F"
PAPER_2 = "#16181C"
INK = "#E8E6E1"
INK_2 = "#A09D96"
MUTED = "#7C7972"
RULE = "rgba(255,255,255,0.12)"
GRID = "rgba(255,255,255,0.05)"
BLUE = "#6C8EF5"
VERMILION = "#F0694A"
OCHRE = "#D9A441"
GREEN = "#5FB88A"
PLUM = "#B39AE8"

# Backwards-compatible aliases used by earlier versions of app.py
TEXT = INK
CYAN = BLUE
ROSE = VERMILION
EMERALD = GREEN
VIOLET = PLUM
AMBER = OCHRE

SANS = "Geist, ui-sans-serif, system-ui, sans-serif"
MONO = "Geist Mono, ui-monospace, SFMono-Regular, Menlo, monospace"
SERIF = SANS

OPINION_SCALE = [
    [0.0, "#3D6BE0"],
    [0.25, "#7E9BEA"],
    [0.5, "#5A5D66"],
    [0.75, "#EE8A70"],
    [1.0, "#E5532F"],
]


def _hex_to_rgb(value: str) -> tuple[int, int, int]:
    value = value.lstrip("#")
    return int(value[0:2], 16), int(value[2:4], 16), int(value[4:6], 16)


def opinion_color(x: float) -> str:
    """Map an opinion in [-1, 1] to the diverging opinion scale."""
    t = (float(np.clip(x, -1.0, 1.0)) + 1.0) / 2.0
    for (t0, c0), (t1, c1) in zip(OPINION_SCALE[:-1], OPINION_SCALE[1:]):
        if t <= t1:
            f = 0.0 if t1 == t0 else (t - t0) / (t1 - t0)
            a, b = _hex_to_rgb(c0), _hex_to_rgb(c1)
            r, g, bl = (round(a[k] + (b[k] - a[k]) * f) for k in range(3))
            return f"#{r:02x}{g:02x}{bl:02x}"
    return OPINION_SCALE[-1][1]


def _hub_radial_layout(graph: nx.Graph, seed: int) -> dict[int, np.ndarray]:
    nodes = list(graph.nodes())
    if not nodes:
        return {}
    if len(nodes) == 1:
        return {nodes[0]: np.array([0.0, 0.0], dtype=float)}

    degrees = dict(graph.degree())
    hub = max(nodes, key=lambda n: (degrees[n], -int(n) if isinstance(n, int) else 0))
    distances = nx.single_source_shortest_path_length(graph, hub)
    max_known = max(distances.values(), default=1)

    rings: dict[int, list[int]] = {}
    for node in nodes:
        rings.setdefault(int(distances.get(node, max_known + 1)), []).append(node)

    rng = np.random.default_rng(seed + 1701)
    pos: dict[int, np.ndarray] = {}
    max_ring = max(rings) if rings else 1

    for ring, ring_nodes in sorted(rings.items()):
        ring_nodes = sorted(ring_nodes, key=lambda n: (-degrees[n], n))
        count = len(ring_nodes)
        phase = float(rng.uniform(0.0, 2.0 * math.pi))
        if ring == 0:
            pos[ring_nodes[0]] = np.array([0.0, 0.0], dtype=float)
            continue
        radius = 0.18 + 0.82 * (ring / max(1, max_ring))
        for idx, node in enumerate(ring_nodes):
            angle = phase + (2.0 * math.pi * idx / max(1, count))
            r = max(0.06, radius + float(rng.uniform(-0.025, 0.025)))
            pos[node] = np.array([1.08 * r * math.cos(angle), r * math.sin(angle)], dtype=float)
    return pos


def _layout(graph: nx.Graph, seed: int, previous: dict | None = None) -> dict[int, np.ndarray]:
    if graph.number_of_nodes() <= 250:
        return nx.spring_layout(graph, pos=previous, seed=seed, iterations=42, threshold=1e-4)
    return _hub_radial_layout(graph, seed)


def _cached_layout(sim: OpinionSimulation) -> dict[int, np.ndarray]:
    """Layout of the *current* graph; recomputed only after rewiring changed it."""
    version = getattr(sim, "graph_version", 0)
    pos = getattr(sim, "_layout_cache", None)
    if pos is None or getattr(sim, "_layout_version", -1) != version:
        pos = _layout(sim.graph, sim.config.seed, previous=pos)
        sim._layout_cache = pos
        sim._layout_version = version
    return pos


def base_layout(title: str, subtitle: str | None = None, height: int = 520) -> dict:
    """Shared Plotly layout: dark ground, quiet titles, mono ticks, hairline rules."""
    title_text = title if not subtitle else (
        f"{title}<br><span style='font-family:{SANS};font-size:12px;color:{MUTED}'>"
        f"{subtitle}</span>"
    )
    return dict(
        title=dict(text=title_text, x=0.012, xanchor="left", y=0.965, yanchor="top", font=dict(family=SANS, size=15, color=INK)),
        template="plotly_dark",
        paper_bgcolor="rgba(0,0,0,0)",
        plot_bgcolor="rgba(0,0,0,0)",
        font=dict(family=SANS, color=INK_2, size=12),
        height=height,
        margin=dict(l=56, r=24, t=76, b=52),
        hoverlabel=dict(bgcolor="#1A1C20", bordercolor="rgba(255,255,255,0.14)", font=dict(family=MONO, size=11, color=INK)),
        legend=dict(
            orientation="h", yanchor="bottom", y=1.015, x=1, xanchor="right",
            bgcolor="rgba(0,0,0,0)", font=dict(family=SANS, size=11.5, color=INK_2),
        ),
    )


# Kept for compatibility with callers that imported the private name.
_base_layout = base_layout


def style_axes(fig: go.Figure, xtitle: str, ytitle: str, yrange=None, xrange=None, ytickformat=None) -> go.Figure:
    common = dict(
        gridcolor=GRID, zeroline=False, showline=True, linecolor=RULE, linewidth=1,
        ticks="outside", tickcolor=RULE, ticklen=4, color=MUTED,
        tickfont=dict(family=MONO, size=10, color=MUTED),
    )
    fig.update_xaxes(title=dict(text=xtitle, font=dict(size=11.5, color=INK_2)), range=xrange, **common)
    fig.update_yaxes(title=dict(text=ytitle, font=dict(size=11.5, color=INK_2)), range=yrange, tickformat=ytickformat, **common)
    return fig


def network_figure(sim: OpinionSimulation, max_edges: int = 2600) -> go.Figure:
    graph = sim.graph
    if sim.config.n_agents >= 1000:
        max_edges = min(max_edges, 900)
    elif sim.config.n_agents >= 500:
        max_edges = min(max_edges, 1400)

    pos = _cached_layout(sim)

    edges = list(graph.edges())
    if len(edges) > max_edges:
        rng = np.random.default_rng(sim.config.seed + 73)
        chosen = rng.choice(len(edges), size=max_edges, replace=False)
        edges = [edges[int(i)] for i in chosen]

    edge_x: list[float | None] = []
    edge_y: list[float | None] = []
    for u, v in edges:
        x0, y0 = pos[u]
        x1, y1 = pos[v]
        edge_x.extend((float(x0), float(x1), None))
        edge_y.extend((float(y0), float(y1), None))

    edge_alpha = 0.10 if sim.config.n_agents <= 100 else 0.06
    edge_trace = go.Scattergl(
        x=edge_x, y=edge_y, mode="lines",
        line=dict(width=0.6, color=f"rgba(255,255,255,{edge_alpha})"),
        hoverinfo="skip", name="ties",
    )

    nodes = list(graph.nodes())
    opinions = sim.opinions[nodes]
    base_size = 9.0 if sim.config.n_agents <= 100 else 5.6 if sim.config.n_agents <= 500 else 4.6
    stub = np.array([bool(sim.stubborn_mask[i]) for i in nodes])
    exposed = np.array([bool(sim.misinformation_mask[i]) for i in nodes])
    marker_sizes = np.where(stub, base_size + 4.5, base_size)
    marker_symbols = np.where(stub, "diamond", "circle")
    line_colors = np.where(exposed, OCHRE, np.where(stub, INK, "rgba(12,13,15,0.9)"))
    line_widths = np.where(exposed, 1.8, np.where(stub, 1.1, 0.7))

    hover = []
    for i in nodes:
        tags = []
        if sim.stubborn_mask[i]:
            tags.append("stubborn")
        if sim.misinformation_mask[i]:
            tags.append("exposed")
        role = " · ".join(tags) if tags else "ordinary"
        hover.append(f"<b>agent {i}</b><br>opinion  {sim.opinions[i]:+.3f}<br>degree   {graph.degree(i)}<br>role     {role}")

    node_trace = go.Scattergl(
        x=[float(pos[i][0]) for i in nodes],
        y=[float(pos[i][1]) for i in nodes],
        mode="markers", hovertext=hover, hoverinfo="text", name="agents",
        marker=dict(
            size=marker_sizes, symbol=marker_symbols, color=opinions,
            colorscale=OPINION_SCALE, cmin=-1, cmax=1, showscale=True,
            colorbar=dict(
                title=dict(text="opinion", font=dict(family=SANS, color=MUTED, size=11), side="top"),
                thickness=8, len=0.46, y=0.5, x=1.0, xpad=6,
                tickvals=[-1, -0.5, 0, 0.5, 1], ticktext=["−1", "−.5", "0", "+.5", "+1"],
                tickfont=dict(family=MONO, color=MUTED, size=10), outlinewidth=0, ticks="",
            ),
            line=dict(width=line_widths, color=line_colors),
        ),
    )

    total_edges = graph.number_of_edges()
    shown_edges = len(edges)
    caption = f"{sim.config.n_agents:,} agents · {total_edges:,} ties" + (f" · {shown_edges:,} drawn" if shown_edges < total_edges else "")
    if getattr(sim, "rewired_ties", 0):
        caption += f" · {sim.rewired_ties:,} rewired"
    r = sim.metric_history[-1].get("assortativity", 0.0)
    caption += f" · r = {r:+.2f}"
    fig = go.Figure(data=[edge_trace, node_trace])
    subtitle = f"{sim.config.network} topology" + (" · adaptive (homophilic rewiring)" if sim.config.rewiring_rate > 0 else " · fixed ties")
    fig.update_layout(**base_layout("Network state", subtitle, height=580))
    fig.update_layout(
        margin=dict(l=14, r=18, t=76, b=18),
        showlegend=False,
        xaxis=dict(visible=False, fixedrange=True),
        yaxis=dict(visible=False, fixedrange=True, scaleanchor="x", scaleratio=1),
        uirevision="network",
        annotations=[
            dict(
                text=caption, x=1.0, y=1.07, xref="paper", yref="paper", showarrow=False,
                xanchor="right", yanchor="top", font=dict(family=MONO, size=11, color=MUTED),
            )
        ],
    )
    return fig


def evolution_figure(sim: OpinionSimulation) -> go.Figure:
    history = sim.history_array
    steps = np.arange(history.shape[0])
    mean = history.mean(axis=1)
    p10 = np.quantile(history, 0.10, axis=1)
    p25 = np.quantile(history, 0.25, axis=1)
    p75 = np.quantile(history, 0.75, axis=1)
    p90 = np.quantile(history, 0.90, axis=1)

    # A deterministic sample of individual trajectories, coloured by current sign.
    rng = np.random.default_rng(sim.config.seed + 5)
    sample = rng.choice(history.shape[1], size=min(36, history.shape[1]), replace=False)
    fig = go.Figure()
    for sign, color, name in ((1, "rgba(240,105,74,0.30)", "agents ending +"), (-1, "rgba(108,142,245,0.30)", "agents ending −")):
        xs: list[float | None] = []
        ys: list[float | None] = []
        for idx in sample:
            if (history[-1, idx] >= 0) == (sign > 0):
                xs.extend(steps.tolist() + [None])
                ys.extend(history[:, idx].tolist() + [None])
        fig.add_trace(go.Scatter(x=xs, y=ys, mode="lines", line=dict(width=0.8, color=color), name=name, hoverinfo="skip", showlegend=False))

    fig.add_trace(go.Scatter(x=steps, y=p90, mode="lines", line=dict(width=0), showlegend=False, hoverinfo="skip"))
    fig.add_trace(go.Scatter(x=steps, y=p10, mode="lines", fill="tonexty", fillcolor="rgba(255,255,255,0.04)", line=dict(width=0), name="10–90%", hovertemplate="10–90% band<extra></extra>"))
    fig.add_trace(go.Scatter(x=steps, y=p75, mode="lines", line=dict(width=0), showlegend=False, hoverinfo="skip"))
    fig.add_trace(go.Scatter(x=steps, y=p25, mode="lines", fill="tonexty", fillcolor="rgba(255,255,255,0.08)", line=dict(width=0), name="25–75%", hovertemplate="25–75% band<extra></extra>"))
    fig.add_trace(go.Scatter(x=steps, y=mean, mode="lines", line=dict(width=2.4, color=INK), name="mean", hovertemplate="step %{x}<br>mean %{y:+.3f}<extra></extra>"))
    fig.add_hline(y=0, line_width=1, line_dash="dot", line_color=RULE)
    fig.update_layout(**base_layout("Opinion evolution", "Mean, quantile bands, and 36 sampled agents", height=540))
    style_axes(fig, "Simulation step", "Opinion", yrange=[-1.05, 1.05])
    return fig


def metric_figure(sim: OpinionSimulation, null: NullTest | None = None) -> go.Figure:
    """Two panels: network-blind distribution scores, and structure vs filter."""
    steps = [int(m["step"]) for m in sim.metric_history]
    series = lambda key: [m[key] for m in sim.metric_history]  # noqa: E731
    fig = make_subplots(
        rows=2, cols=1, shared_xaxes=True, vertical_spacing=0.13,
        subplot_titles=(
            "Distribution · network-blind",
            "Structure (who sits next to whom) vs filter (what agents hear)",
        ),
    )
    for key, label, color, dash in (
        ("polarization", "polarization Var(x)", VERMILION, "solid"),
        ("bimodality", "bimodality coefficient", OCHRE, "solid"),
    ):
        fig.add_trace(go.Scatter(
            x=steps, y=series(key), mode="lines", name=label, legendgroup="dist",
            line=dict(width=2.0, color=color, dash=dash),
            hovertemplate=f"step %{{x}}<br>{label} %{{y:.3f}}<extra></extra>",
        ), row=1, col=1)
    fig.add_hline(y=5 / 9, line_width=1, line_dash="dot", line_color=RULE, row=1, col=1,
                  annotation_text="BC = 5/9", annotation_font=dict(family=MONO, size=10, color=MUTED),
                  annotation_position="top left")

    if null is not None and null.n_permutations:
        fig.add_hrect(y0=null.null_low, y1=null.null_high, fillcolor="rgba(255,255,255,0.06)", line_width=0,
                      row=2, col=1, layer="below")
        fig.add_annotation(
            x=steps[-1], y=null.null_high, xref="x2", yref="y2", xanchor="right", yanchor="bottom",
            text="shuffled-opinion 95% band", showarrow=False, font=dict(family=MONO, size=10, color=MUTED),
        )
    for key, label, color, dash, width in (
        ("assortativity", "edge assortativity r", BLUE, "solid", 2.4),
        ("neighbour_corr", "neighbour-mean correlation", GREEN, "dash", 1.6),
        ("selective_exposure", "selective exposure (set by β)", PLUM, "dot", 2.0),
    ):
        fig.add_trace(go.Scatter(
            x=steps, y=series(key), mode="lines", name=label, legendgroup="struct",
            line=dict(width=width, color=color, dash=dash),
            hovertemplate=f"step %{{x}}<br>{label} %{{y:+.3f}}<extra></extra>",
        ), row=2, col=1)
    fig.add_hline(y=0, line_width=1, line_dash="dot", line_color=RULE, row=2, col=1)

    fig.update_layout(**base_layout("Collective diagnostics", None, height=640))
    fig.update_layout(margin=dict(l=56, r=24, t=96, b=52), legend=dict(y=1.08))
    fig.update_annotations(font=dict(family=SANS, size=12, color=INK_2))
    common = dict(gridcolor=GRID, zeroline=False, showline=True, linecolor=RULE, ticks="outside",
                  tickcolor=RULE, ticklen=4, tickfont=dict(family=MONO, size=10, color=MUTED))
    fig.update_xaxes(**common)
    fig.update_yaxes(**common)
    fig.update_xaxes(title=dict(text="Simulation step", font=dict(size=11.5, color=INK_2)), row=2, col=1)
    fig.update_yaxes(range=[0, 1.02], title=dict(text="Score", font=dict(size=11.5, color=INK_2)), row=1, col=1)
    fig.update_yaxes(range=[-0.5, 1.02], title=dict(text="Correlation / share", font=dict(size=11.5, color=INK_2)), row=2, col=1)
    return fig


def distribution_figure(sim: OpinionSimulation) -> go.Figure:
    edges = np.linspace(-1, 1, 31)
    centers = (edges[:-1] + edges[1:]) / 2.0
    width = float(edges[1] - edges[0])
    n = max(1, sim.opinions.size)
    current, _ = np.histogram(sim.opinions, bins=edges)
    initial, _ = np.histogram(sim.initial_opinions, bins=edges)
    current = current / n
    initial = initial / n

    fig = go.Figure()
    fig.add_vrect(x0=-0.2, x1=0.2, fillcolor="rgba(255,255,255,0.03)", line_width=0, layer="below")
    fig.add_trace(go.Bar(
        x=centers, y=current, width=width * 0.84, name="current",
        marker=dict(color=[opinion_color(c) for c in centers], line=dict(width=0)),
        hovertemplate="opinion %{x:+.2f}<br>share %{y:.1%}<extra>current</extra>",
    ))
    fig.add_trace(go.Scatter(
        x=edges, y=np.r_[initial, initial[-1]], mode="lines", line=dict(shape="hv", width=1.4, color=INK),
        name="initial", hoverinfo="skip",
    ))
    fig.update_layout(**base_layout("Opinion distribution", "Current shares (bars) against the initial state (line)", height=540))
    fig.update_layout(bargap=0)
    style_axes(fig, "Opinion", "Population share", xrange=[-1.02, 1.02], ytickformat=".0%")
    return fig


def influence_figure(sim: OpinionSimulation) -> go.Figure:
    nodes = np.array(list(sim.graph.nodes()), dtype=int)
    degrees = np.array([sim.graph.degree(int(i)) for i in nodes], dtype=float)
    opinions = sim.opinions[nodes]
    stub = sim.stubborn_mask[nodes]
    exposure = sim.misinformation_mask[nodes]
    sizes = np.where(stub, 12.0, 7.5)
    symbols = np.where(stub, "diamond", "circle")
    jitter = np.random.default_rng(sim.config.seed + 11).uniform(-0.22, 0.22, size=degrees.size)

    hover = [
        f"<b>agent {int(i)}</b><br>degree   {int(d)}<br>opinion  {o:+.3f}<br>stubborn {'yes' if s else 'no'}<br>exposed  {'yes' if e else 'no'}"
        for i, d, o, s, e in zip(nodes, degrees, opinions, stub, exposure)
    ]

    fig = go.Figure(
        go.Scattergl(
            x=degrees + jitter, y=opinions, mode="markers", hovertext=hover, hoverinfo="text", name="agents",
            marker=dict(
                size=sizes, symbol=symbols, color=opinions, colorscale=OPINION_SCALE, cmin=-1, cmax=1,
                opacity=np.where(exposure, 1.0, 0.78),
                line=dict(width=np.where(exposure, 1.6, 0.5), color=np.where(exposure, OCHRE, np.where(stub, INK, "rgba(12,13,15,0.9)"))),
                showscale=False,
            ),
        )
    )
    fig.add_hline(y=0, line_width=1, line_dash="dot", line_color=RULE)
    fig.update_layout(**base_layout("Influence map", "Connectivity against current opinion · amber ring = exposed", height=540))
    style_axes(fig, "Node degree", "Opinion", yrange=[-1.05, 1.05])
    return fig


def sweep_figure(summary: list[dict[str, float]]) -> go.Figure:
    """Filter vs chamber: selective exposure and assortativity against beta, with 95% CIs."""
    fig = make_subplots(rows=1, cols=2, shared_yaxes=True, horizontal_spacing=0.06,
                        subplot_titles=("Fixed ties", "Homophilic rewiring"))
    rates = sorted({row["rewiring_rate"] for row in summary})
    for col, rate in enumerate(rates[:2], start=1):
        rows = [r for r in summary if r["rewiring_rate"] == rate]
        x = [r["confirmation_bias"] for r in rows]
        for key, label, color, dash in (
            ("selective_exposure", "selective exposure (filter)", PLUM, "dot"),
            ("assortativity", "edge assortativity r (structure)", BLUE, "solid"),
        ):
            mean = np.array([r[f"{key}_mean"] for r in rows])
            ci = np.array([r[f"{key}_ci95"] for r in rows])
            rgb = ",".join(str(c) for c in _hex_to_rgb(color))
            fig.add_trace(go.Scatter(x=x + x[::-1], y=list(mean + ci) + list((mean - ci)[::-1]), fill="toself", mode="lines",
                                     fillcolor=f"rgba({rgb},0.14)", line=dict(width=0), hoverinfo="skip",
                                     showlegend=False), row=1, col=col)
            fig.add_trace(go.Scatter(x=x, y=mean, mode="lines+markers", name=label, legendgroup=key,
                                     showlegend=col == 1, line=dict(width=2.2, color=color, dash=dash),
                                     marker=dict(size=6, color=color),
                                     customdata=ci, hovertemplate=f"β %{{x:.2f}}<br>{label} %{{y:+.3f}} ± %{{customdata:.3f}}<extra></extra>"),
                          row=1, col=col)
        fig.add_hline(y=0, line_width=1, line_dash="dot", line_color=RULE, row=1, col=col)
    if len(rates) > 1:
        fig.layout.annotations[1].text = f"Homophilic rewiring · rate {rates[1]:.2f}"
    fig.update_layout(**base_layout("Filter vs chamber", "Mean ± 95% CI across seeds; all other parameters held fixed", height=500))
    fig.update_layout(legend=dict(orientation="h", yanchor="top", y=-0.17, x=0, xanchor="left"), margin=dict(l=56, r=24, t=96, b=96))
    fig.update_annotations(font=dict(family=SANS, size=12, color=INK_2))
    common = dict(gridcolor=GRID, zeroline=False, showline=True, linecolor=RULE, ticks="outside",
                  tickcolor=RULE, ticklen=4, tickfont=dict(family=MONO, size=10, color=MUTED))
    fig.update_xaxes(title=dict(text="Confirmation bias (slider value)", font=dict(size=11.5, color=INK_2)), range=[-0.03, 1.03], **common)
    fig.update_yaxes(range=[-0.35, 1.05], **common)
    fig.update_yaxes(title=dict(text="Final-state score", font=dict(size=11.5, color=INK_2)), row=1, col=1)
    return fig