src/q_route/visualizer.py

"""
Q-Route Custom Visualization Suite.
Renders three custom visual assets:
1. Visual 1: Physical QPU Topology Graph with Routing Path Overlay (networkx + matplotlib)
2. Visual 2: Circuit Diagram Before vs After (Qiskit circuit drawer)
3. Visual 3: Generation Quality Bar Chart (Plotly graph objects)
"""

import io
from typing import List, Tuple, Dict, Any, Optional
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import networkx as nx
from PIL import Image

try:
    import plotly.graph_objects as go
    HAS_PLOTLY = True
except ImportError:
    go = None
    HAS_PLOTLY = False

try:
    import qiskit
    from qiskit.visualization import circuit_drawer
    from q_route.evaluator import parse_qasm_string, QuantumCircuit
    HAS_QISKIT_DRAW = True
except Exception:
    HAS_QISKIT_DRAW = False


# Theme Color Definitions
COLOR_BG_DARK = "#0A0A0F"
COLOR_BG_CARD = "#111118"
COLOR_TOPO_EDGE = "#3D4A5C"
COLOR_ROUTE_PATH = "#00D4FF"
COLOR_IMPOSSIBLE = "#FF3D57"
COLOR_NODE_FILL = "#7B2FBE"
COLOR_NODE_BORDER = "#9B4FDE"
COLOR_NODE_TEXT = "#FFFFFF"

COLOR_BAR_SWAPS = "#00D4FF"
COLOR_BAR_DEPTH = "#7B2FBE"
COLOR_BAR_GATES = "#00C853"
COLOR_INVALID_RED = "#FF3D57"
COLOR_GOLD = "#FFD700"


def render_topology_routing_graph(
    coupling_map: List[Tuple[int, int]],
    requested_gate: Optional[Tuple[int, int]] = None,
    routed_path: Optional[List[Tuple[int, int]]] = None,
) -> plt.Figure:
    """
    Render Visual 1: Physical QPU Topology Graph with Routing Path Overlay.
    """
    G = nx.Graph()
    for u, v in coupling_map:
        G.add_edge(u, v)

    fig, ax = plt.subplots(figsize=(8, 5), facecolor=COLOR_BG_DARK)
    ax.set_facecolor(COLOR_BG_DARK)

    # Compute node positions
    pos = nx.spring_layout(G, seed=42) if len(G.nodes) > 6 else nx.kamada_kawai_layout(G)

    # 1. Base Topology Edges
    nx.draw_networkx_edges(
        G, pos, ax=ax, edge_color=COLOR_TOPO_EDGE, width=2.5, alpha=0.8
    )

    # 2. Impossible Requested Gate (Red dashed arc)
    if requested_gate:
        u, v = requested_gate
        if u in pos and v in pos:
            ax.annotate(
                "",
                xy=pos[v],
                xytext=pos[u],
                arrowprops=dict(
                    arrowstyle="<->",
                    color=COLOR_IMPOSSIBLE,
                    linestyle="dashed",
                    linewidth=2.5,
                    connectionstyle="arc3,rad=0.35",
                ),
            )
            mid_x = (pos[u][0] + pos[v][0]) / 2.0
            mid_y = (pos[u][1] + pos[v][1]) / 2.0 + 0.15
            ax.text(
                mid_x,
                mid_y,
                f"REQUESTED: Impossible {u}╌{v}",
                color=COLOR_IMPOSSIBLE,
                fontsize=9,
                fontweight="bold",
                ha="center",
                backgroundcolor="#0A0A0F80",
            )

    # 3. Routed Path (Solid Cyan Overlay)
    if routed_path:
        path_edges = []
        for i in range(len(routed_path) - 1):
            path_edges.append((routed_path[i], routed_path[i + 1]))
        nx.draw_networkx_edges(
            G,
            pos,
            edgelist=path_edges,
            ax=ax,
            edge_color=COLOR_ROUTE_PATH,
            width=5.0,
            alpha=0.9,
        )

    # 4. Nodes
    nx.draw_networkx_nodes(
        G,
        pos,
        ax=ax,
        node_color=COLOR_NODE_FILL,
        edgecolors=COLOR_NODE_BORDER,
        linewidths=2.0,
        node_size=700,
    )
    nx.draw_networkx_labels(
        G,
        pos,
        ax=ax,
        font_color=COLOR_NODE_TEXT,
        font_weight="bold",
        font_size=10,
    )

    ax.set_title(
        "Visual 1: Physical QPU Topology & Q-Route Solution Overlay",
        color="#F0F0FF",
        fontsize=12,
        fontweight="bold",
        pad=12,
    )
    ax.axis("off")
    plt.tight_layout()
    return fig


def render_circuit_before_after(
    abstract_qasm: str, routed_qasm: str
) -> plt.Figure:
    """
    Render Visual 2: Qiskit Circuit Diagram Before vs After.
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4.5), facecolor=COLOR_BG_DARK)
    ax1.set_facecolor(COLOR_BG_CARD)
    ax2.set_facecolor(COLOR_BG_CARD)

    qc_abs = parse_qasm_string(abstract_qasm) if HAS_QISKIT_DRAW else None
    qc_rt = parse_qasm_string(routed_qasm) if HAS_QISKIT_DRAW else None

    if HAS_QISKIT_DRAW and isinstance(qc_abs, QuantumCircuit) and isinstance(qc_rt, QuantumCircuit):
        try:
            circuit_drawer(qc_abs, output="mpl", ax=ax1)
            circuit_drawer(qc_rt, output="mpl", ax=ax2)
        except Exception:
            _draw_fallback_text_box(ax1, "BEFORE (Abstract Input)", abstract_qasm, COLOR_IMPOSSIBLE)
            _draw_fallback_text_box(ax2, "AFTER (Q-Route Hardware Compliant)", routed_qasm, COLOR_ROUTE_PATH)
    else:
        _draw_fallback_text_box(ax1, "BEFORE (Abstract Input)", abstract_qasm, COLOR_IMPOSSIBLE)
        _draw_fallback_text_box(ax2, "AFTER (Q-Route Hardware Compliant)", routed_qasm, COLOR_ROUTE_PATH)

    ax1.set_title("BEFORE (Abstract Input — Illegal Gates)", color=COLOR_IMPOSSIBLE, fontsize=11, fontweight="bold")
    ax2.set_title("AFTER (Q-Route Routed — 100% Compliant)", color=COLOR_BAR_GATES, fontsize=11, fontweight="bold")
    plt.tight_layout()
    return fig


def _draw_fallback_text_box(ax: plt.Axes, title: str, code: str, border_color: str):
    ax.axis("off")
    code_lines = "\n".join(code.splitlines()[:12])
    ax.text(
        0.05,
        0.90,
        f"{title}\n\n{code_lines}",
        color="#C8D3E8",
        fontsize=8.5,
        fontfamily="monospace",
        va="top",
        bbox=dict(boxstyle="round,pad=0.5", facecolor=COLOR_BG_CARD, edgecolor=border_color, linewidth=1.5),
    )


def render_quality_bar_chart_mpl(
    model_eval_metrics: Dict[str, Dict[str, Any]],
    theoretical_min_swaps: int = 0,
) -> plt.Figure:
    """Matplotlib fallback rendering engine for Visual 3 Bar Chart."""
    import numpy as np
    models = list(model_eval_metrics.keys())
    swaps = [model_eval_metrics[m].get("swap_count", 0) for m in models]
    depths = [model_eval_metrics[m].get("depth", 0) for m in models]
    gates = [model_eval_metrics[m].get("total_2q_gates", 0) for m in models]
    passes = [model_eval_metrics[m].get("pass_topology", False) for m in models]

    display_names = [f"{m} ✅" if p else f"{m} ❌ [INVALID]" for m, p in zip(models, passes)]

    fig, ax = plt.subplots(figsize=(8.5, max(3.5, len(models) * 0.8)), facecolor=COLOR_BG_DARK)
    ax.set_facecolor(COLOR_BG_CARD)

    y = np.arange(len(models))
    height = 0.25

    ax.barh(y - height, swaps, height, label="SWAP Gates", color=COLOR_BAR_SWAPS)
    ax.barh(y, depths, height, label="Circuit Depth", color=COLOR_BAR_DEPTH)
    ax.barh(y + height, gates, height, label="2-Qubit Gates", color=COLOR_BAR_GATES)

    if theoretical_min_swaps > 0:
        ax.axvline(x=theoretical_min_swaps, color=COLOR_GOLD, linestyle="--", label=f"Min SWAPs ({theoretical_min_swaps})")

    ax.set_yticks(y)
    ax.set_yticklabels(display_names, color="#F0F0FF", fontsize=9.5)
    ax.invert_yaxis()
    ax.set_xlabel("Count", color="#A0A0C0", fontsize=10)
    ax.set_title("Visual 3: Generation Quality & Topology Compliance Comparison", color="#F0F0FF", fontsize=11, fontweight="bold")
    ax.legend(facecolor=COLOR_BG_DARK, edgecolor="#2A2A4A", labelcolor="#F0F0FF", fontsize=8.5)
    ax.tick_params(colors="#A0A0C0")
    for spine in ax.spines.values():
        spine.set_color("#2A2A4A")

    plt.tight_layout()
    return fig


def render_quality_bar_chart(
    model_eval_metrics: Dict[str, Dict[str, Any]],
    theoretical_min_swaps: int = 0,
) -> Any:
    """
    Render Visual 3: Plotly Horizontal Grouped Bar Chart with Matplotlib Fallback.
    """
    if not model_eval_metrics:
        return None

    if HAS_PLOTLY and go:
        try:
            models = list(model_eval_metrics.keys())
            swaps = [model_eval_metrics[m].get("swap_count", 0) for m in models]
            depths = [model_eval_metrics[m].get("depth", 0) for m in models]
            gates = [model_eval_metrics[m].get("total_2q_gates", 0) for m in models]
            passes = [model_eval_metrics[m].get("pass_topology", False) for m in models]

            display_names = [f"{m} ✅" if p else f"{m} ❌ [INVALID]" for m, p in zip(models, passes)]

            fig = go.Figure()

            fig.add_trace(
                go.Bar(
                    y=display_names,
                    x=swaps,
                    name="SWAP Gates",
                    orientation="h",
                    marker=dict(color=COLOR_BAR_SWAPS),
                    text=[f"{s} SWAPs" for s in swaps],
                    textposition="auto",
                )
            )

            fig.add_trace(
                go.Bar(
                    y=display_names,
                    x=depths,
                    name="Circuit Depth",
                    orientation="h",
                    marker=dict(color=COLOR_BAR_DEPTH),
                    text=[f"Depth {d}" for d in depths],
                    textposition="auto",
                )
            )

            fig.add_trace(
                go.Bar(
                    y=display_names,
                    x=gates,
                    name="2-Qubit Gates",
                    orientation="h",
                    marker=dict(color=COLOR_BAR_GATES),
                    text=[f"{g} Gates" for g in gates],
                    textposition="auto",
                )
            )

            if theoretical_min_swaps > 0:
                fig.add_vline(
                    x=theoretical_min_swaps,
                    line_dash="dash",
                    line_color=COLOR_GOLD,
                    annotation_text=f"Theoretical Min SWAPs ({theoretical_min_swaps})",
                    annotation_position="top right",
                    annotation_font_color=COLOR_GOLD,
                )

            fig.update_layout(
                title=dict(
                    text="<b>Visual 3: Generation Quality & Topology Compliance Comparison</b>",
                    font=dict(family="Space Grotesk, sans-serif", size=15, color="#F0F0FF"),
                ),
                barmode="group",
                paper_bgcolor=COLOR_BG_DARK,
                plot_bgcolor=COLOR_BG_CARD,
                font=dict(family="Inter, sans-serif", color="#A0A0C0"),
                xaxis=dict(
                    title="Count",
                    gridcolor="#2A2A4A",
                    showgrid=True,
                    zerolinecolor="#3D3D6B",
                ),
                yaxis=dict(
                    autorange="reversed",
                    gridcolor="#2A2A4A",
                ),
                legend=dict(
                    orientation="h",
                    yanchor="bottom",
                    y=1.02,
                    xanchor="right",
                    x=1,
                    font=dict(color="#F0F0FF"),
                ),
                margin=dict(l=20, r=20, t=60, b=40),
                height=380,
            )
            return fig
        except Exception:
            pass

    return render_quality_bar_chart_mpl(model_eval_metrics, theoretical_min_swaps)