kgvis/plot.py

from __future__ import annotations

from typing import Iterable, Tuple

import networkx as nx
import plotly.graph_objects as go

from .schema import GraphData


def _build_layout(graph: GraphData, mode: str) -> dict[str, tuple[float, float]]:
    if not graph.nodes:
        return {}

    graph_nx = nx.DiGraph()
    for node in graph.nodes:
        graph_nx.add_node(node.id)
    for edge in graph.edges:
        graph_nx.add_edge(edge.source, edge.target, predicate=edge.predicate)

    if mode == "bbox":
        explicit_positions: dict[str, tuple[float, float]] = {}
        for node in graph.nodes:
            center = _center_from_location(node.attributes)
            if center is not None:
                explicit_positions[node.id] = center

        if explicit_positions and len(explicit_positions) == len(graph.nodes):
            return explicit_positions

        if explicit_positions:
            return _spring_layout(graph_nx, pos=explicit_positions, fixed=explicit_positions.keys())

    return _layout_force(graph_nx)


def _spring_layout(
    graph_nx: nx.DiGraph,
    pos: dict[str, tuple[float, float]] | None = None,
    fixed: Iterable[str] | None = None,
) -> dict[str, tuple[float, float]]:
    n = max(graph_nx.number_of_nodes(), 1)
    k = 0.9 / (n**0.5)
    return nx.spring_layout(
        graph_nx,
        seed=42,
        k=k,
        iterations=300,
        pos=pos,
        fixed=fixed,
    )


def _center_from_location(attributes: dict) -> Tuple[float, float] | None:
    if not isinstance(attributes, dict):
        return None

    bbox = attributes.get("bbox")
    if isinstance(bbox, dict):
        try:
            x1 = float(bbox["x1"])
            y1 = float(bbox["y1"])
            x2 = float(bbox["x2"])
            y2 = float(bbox["y2"])
            return _to_plot_center(x1, y1, x2, y2)
        except (KeyError, TypeError, ValueError):
            pass

    raw = attributes.get("location_raw") or attributes.get("location")
    parsed = _parse_location_value(raw)
    if parsed:
        x1, y1, x2, y2 = parsed
        return _to_plot_center(x1, y1, x2, y2)
    return None


def _parse_location_value(value) -> Tuple[float, float, float, float] | None:
    parts: list[float] = []
    if isinstance(value, str):
        try:
            parts = [float(p.strip()) for p in value.split(",") if p.strip()]
        except ValueError:
            return None
    elif isinstance(value, (list, tuple)):
        try:
            parts = [float(p) for p in value]
        except (TypeError, ValueError):
            return None
    else:
        return None

    if len(parts) < 4:
        return None
    return parts[0], parts[1], parts[2], parts[3]


def _to_plot_center(x1: float, y1: float, x2: float, y2: float) -> Tuple[float, float]:
    center_x = (x1 + x2) / 2
    center_y = (y1 + y2) / 2
    # Normalize to plot coords: y=0 at bottom. Flip so y=0 is top like images.
    return center_x, 1 - center_y


def _normalize_positions(
    pos: dict[str, tuple[float, float]],
    padding: float = 0.05,
) -> dict[str, tuple[float, float]]:
    if not pos:
        return {}
    xs = [p[0] for p in pos.values()]
    ys = [p[1] for p in pos.values()]
    min_x, max_x = min(xs), max(xs)
    min_y, max_y = min(ys), max(ys)
    span_x = max(max_x - min_x, 1e-6)
    span_y = max(max_y - min_y, 1e-6)
    scale_x = (1 - 2 * padding) / span_x
    scale_y = (1 - 2 * padding) / span_y
    normalized = {}
    for node_id, (x, y) in pos.items():
        nx = padding + (x - min_x) * scale_x
        ny = padding + (y - min_y) * scale_y
        normalized[node_id] = (nx, ny)
    return normalized


def _layout_force(graph_nx: nx.DiGraph) -> dict[str, tuple[float, float]]:
    if graph_nx.number_of_nodes() == 0:
        return {}

    undirected = graph_nx.to_undirected()
    components = [list(c) for c in nx.connected_components(undirected)]
    components.sort(key=len, reverse=True)

    count = len(components)
    cols = max(1, int((count**0.5) + 0.999))
    rows = (count + cols - 1) // cols
    cell_w = 1 / cols
    cell_h = 1 / rows
    cell_pad = 0.08

    positions: dict[str, tuple[float, float]] = {}
    for idx, nodes in enumerate(components):
        sub = graph_nx.subgraph(nodes)
        if len(nodes) == 1:
            pos_comp = {nodes[0]: (0.5, 0.5)}
        else:
            pos_comp = _spring_layout(sub)
        pos_comp = _normalize_positions(pos_comp, padding=0.1)

        col = idx % cols
        row = idx // cols
        x0 = col * cell_w
        y0 = row * cell_h
        scale_x = cell_w * (1 - 2 * cell_pad)
        scale_y = cell_h * (1 - 2 * cell_pad)
        for node_id, (x, y) in pos_comp.items():
            px = x0 + cell_pad * cell_w + x * scale_x
            py = y0 + cell_pad * cell_h + y * scale_y
            positions[node_id] = (px, py)

    return _normalize_positions(positions, padding=0.04)


def build_figure(
    graph: GraphData,
    *,
    show_nodes: bool,
    show_edges: bool,
    show_node_labels: bool,
    show_edge_labels: bool,
    selected_node_id: str | None,
    layout_mode: str,
    background_image: str | None = None,
    background_aspect: float | None = None,
    background_opacity: float = 0.35,
    show_bboxes: bool = False,
    highlight_node_id: str | None = None,
) -> go.Figure:
    pos = _build_layout(graph, layout_mode)
    image_frame = _image_frame(background_aspect) if background_aspect else None
    image_domain = _image_domain(background_aspect) if background_aspect else None
    if layout_mode == "bbox" and image_domain:
        pos = _scale_positions(pos, image_domain)

    edge_x: list[float] = []
    edge_y: list[float] = []
    edge_text_x: list[float] = []
    edge_text_y: list[float] = []
    edge_text: list[str] = []

    if show_edges:
        for edge in graph.edges:
            if edge.source not in pos or edge.target not in pos:
                continue
            x0, y0 = pos[edge.source]
            x1, y1 = pos[edge.target]
            edge_x.extend([x0, x1, None])
            edge_y.extend([y0, y1, None])
            edge_text_x.append((x0 + x1) / 2)
            edge_text_y.append((y0 + y1) / 2)
            edge_text.append(edge.predicate)

    edge_trace = None
    if show_edges:
        edge_trace = go.Scatter(
            x=edge_x,
            y=edge_y,
            mode="lines",
            line={"width": 1.2, "color": "#9aa4b2"},
            hoverinfo="none",
        )

    node_x: list[float] = []
    node_y: list[float] = []
    node_text: list[str] = []
    node_hover: list[str] = []
    node_color: list[str] = []
    node_ids: list[str] = []

    if show_nodes:
        for node in graph.nodes:
            if node.id not in pos:
                continue
            x, y = pos[node.id]
            node_x.append(x)
            node_y.append(y)
            node_ids.append(node.id)
            node_text.append(node.label if show_node_labels else "")
            node_hover.append(f"{node.label} ({node.id})")
            if selected_node_id and node.id == selected_node_id:
                node_color.append("#ff7b7b")
            else:
                node_color.append("#3a7bd5")

    node_trace = None
    if show_nodes:
        node_trace = go.Scatter(
            x=node_x,
            y=node_y,
            mode="markers+text" if show_node_labels else "markers",
            text=node_text,
            textposition="top center",
            textfont={"size": 12, "color": "#1f2a44"},
            hovertext=node_hover,
            hoverinfo="text",
            marker={
                "size": 16,
                "color": node_color,
                "line": {"width": 1.2, "color": "#0f172a"},
            },
            customdata=node_ids,
        )

    traces: list[go.Scatter] = []
    if edge_trace is not None:
        traces.append(edge_trace)
    if node_trace is not None:
        traces.append(node_trace)

    if show_edges and show_edge_labels and edge_text:
        edge_label_trace = go.Scatter(
            x=edge_text_x,
            y=edge_text_y,
            mode="text",
            text=edge_text,
            textfont={"size": 11, "color": "#475569"},
            hoverinfo="none",
        )
        traces.append(edge_label_trace)

    fig = go.Figure(data=traces)
    layout_update = dict(
        showlegend=False,
        hovermode="closest",
        margin={"l": 10, "r": 10, "t": 10, "b": 10},
        xaxis={"visible": False, "range": [0, 1]},
        yaxis={"visible": False, "range": [0, 1]},
        plot_bgcolor="#ffffff",
        annotations=_edge_annotations(graph, pos) if show_edges else [],
        dragmode="pan",
    )
    if layout_mode == "bbox" and image_domain:
        width, height = image_domain
        layout_update["xaxis"] = {
            "visible": False,
            "range": [0, width],
            "constrain": "domain",
        }
        layout_update["yaxis"] = {
            "visible": False,
            "range": [0, height],
            "scaleanchor": "x",
            "scaleratio": 1,
            "constrain": "domain",
        }
    fig.update_layout(**layout_update)

    if show_bboxes:
        shapes = _bbox_shapes(
            graph,
            image_domain=image_domain if layout_mode == "bbox" else None,
            frame=image_frame if layout_mode != "bbox" else None,
            highlight_node_id=highlight_node_id,
        )
        if shapes:
            fig.update_layout(shapes=shapes)

    if background_image:
        if layout_mode == "bbox" and image_domain:
            x0, y0, width, height = 0, 0, image_domain[0], image_domain[1]
        elif image_frame:
            x0, y0, width, height = image_frame
        else:
            x0, y0, width, height = 0, 0, 1, 1
        fig.update_layout(
            images=[
                dict(
                    source=background_image,
                    xref="x",
                    yref="y",
                    x=x0,
                    y=y0 + height,
                    sizex=width,
                    sizey=height,
                    sizing="stretch",
                    xanchor="left",
                    yanchor="top",
                    opacity=max(0.0, min(background_opacity, 1.0)),
                    layer="below",
                )
            ]
        )
    return fig


def _image_frame(aspect: float) -> tuple[float, float, float, float]:
    if aspect <= 0:
        return 0, 0, 1, 1
    if aspect >= 1:
        width = 1.0
        height = 1.0 / aspect
        pad_x = 0.0
        pad_y = (1.0 - height) / 2
    else:
        width = aspect
        height = 1.0
        pad_x = (1.0 - width) / 2
        pad_y = 0.0
    return pad_x, pad_y, width, height


def _image_domain(aspect: float) -> tuple[float, float]:
    if aspect <= 0:
        return 1.0, 1.0
    if aspect >= 1:
        return aspect, 1.0
    return 1.0, 1.0 / aspect


def _apply_image_frame(
    pos: dict[str, tuple[float, float]],
    frame: tuple[float, float, float, float],
) -> dict[str, tuple[float, float]]:
    if not pos:
        return {}
    x0, y0, width, height = frame
    adjusted = {}
    for node_id, (x, y) in pos.items():
        adjusted[node_id] = (x0 + x * width, y0 + y * height)
    return adjusted


def _scale_positions(
    pos: dict[str, tuple[float, float]],
    domain: tuple[float, float],
) -> dict[str, tuple[float, float]]:
    if not pos:
        return {}
    width, height = domain
    scaled = {}
    for node_id, (x, y) in pos.items():
        scaled[node_id] = (x * width, y * height)
    return scaled


def _bbox_from_attrs(attributes: dict) -> tuple[float, float, float, float] | None:
    if not isinstance(attributes, dict):
        return None
    bbox = attributes.get("bbox")
    if isinstance(bbox, dict):
        try:
            x1 = float(bbox["x1"])
            y1 = float(bbox["y1"])
            x2 = float(bbox["x2"])
            y2 = float(bbox["y2"])
            return x1, y1, x2, y2
        except (KeyError, TypeError, ValueError):
            pass

    raw = attributes.get("location_raw") or attributes.get("location")
    parsed = _parse_location_value(raw)
    if parsed:
        return parsed
    return None


def _bbox_shapes(
    graph: GraphData,
    image_domain: tuple[float, float] | None,
    frame: tuple[float, float, float, float] | None,
    highlight_node_id: str | None,
) -> list[dict]:
    shapes: list[dict] = []
    for node in graph.nodes:
        bbox = _bbox_from_attrs(node.attributes)
        if not bbox:
            continue
        x1, y1, x2, y2 = bbox
        # Convert from image coords (origin top-left) to plot coords (origin bottom-left).
        px1, px2 = x1, x2
        py1, py2 = 1 - y2, 1 - y1

        if image_domain:
            width, height = image_domain
            px1 *= width
            px2 *= width
            py1 *= height
            py2 *= height
        elif frame:
            fx, fy, fw, fh = frame
            px1 = fx + px1 * fw
            px2 = fx + px2 * fw
            py1 = fy + py1 * fh
            py2 = fy + py2 * fh

        is_highlight = highlight_node_id and node.id == highlight_node_id
        line_color = "#ef4444" if is_highlight else "rgba(242,95,76,0.55)"
        fill_color = "rgba(239,68,68,0.18)" if is_highlight else "rgba(242,95,76,0.08)"
        line_width = 2.4 if is_highlight else 1.2

        shapes.append(
            dict(
                type="rect",
                xref="x",
                yref="y",
                x0=min(px1, px2),
                x1=max(px1, px2),
                y0=min(py1, py2),
                y1=max(py1, py2),
                line={"color": line_color, "width": line_width},
                fillcolor=fill_color,
                layer="below",
            )
        )
    return shapes


def empty_figure(message: str) -> go.Figure:
    fig = go.Figure()
    fig.add_annotation(
        text=message,
        x=0.5,
        y=0.5,
        xref="paper",
        yref="paper",
        xanchor="center",
        yanchor="middle",
        showarrow=False,
        font={"size": 16, "color": "#64748b"},
    )
    fig.update_layout(
        xaxis={"visible": False},
        yaxis={"visible": False},
        plot_bgcolor="#ffffff",
        margin={"l": 20, "r": 20, "t": 20, "b": 20},
    )
    return fig


def _edge_annotations(graph: GraphData, pos: dict[str, tuple[float, float]]):
    annotations = []
    for edge in graph.edges:
        if edge.source not in pos or edge.target not in pos:
            continue
        x0, y0 = pos[edge.source]
        x1, y1 = pos[edge.target]
        dx = x1 - x0
        dy = y1 - y0
        length = (dx**2 + dy**2) ** 0.5
        if length == 0:
            continue

        # Shorten the arrow so it doesn't overlap the node markers.
        shrink = min(0.05, length * 0.25)
        ux = dx / length
        uy = dy / length
        tail_x = x0 + ux * shrink
        tail_y = y0 + uy * shrink
        head_x = x1 - ux * shrink
        head_y = y1 - uy * shrink

        annotations.append(
            dict(
                x=head_x,
                y=head_y,
                ax=tail_x,
                ay=tail_y,
                xref="x",
                yref="y",
                axref="x",
                ayref="y",
                showarrow=True,
                arrowhead=3,
                arrowsize=1.15,
                arrowwidth=1.2,
                arrowcolor="#94a3b8",
                opacity=0.9,
            )
        )
    return annotations