kg_services/visualizer.py

"""Knowledge Graph Visualization Module for Track 2 Implementation.

This module provides interactive network visualization capabilities using Plotly and NetworkX
for the KGraph-MCP platform, implementing simplified visualization as part of Track 2
hackathon submission.
"""

import logging
from typing import Any

import networkx as nx
import numpy as np
import plotly.graph_objects as go
from plotly.graph_objects import Figure

from .ontology import MCPPrompt, MCPTool, PlannedStep

logger = logging.getLogger(__name__)


class KGVisualizer:
    """Professional knowledge graph visualizer for Track 2 submission.
    
    Provides interactive network visualization using Plotly and NetworkX,
    designed to showcase knowledge graph relationships and planning workflows.
    """

    def __init__(self):
        """Initialize the visualizer with professional color schemes."""
        # Professional color palette - more sophisticated and accessible
        self.colors = {
            # Primary brand colors - more sophisticated blues and teals
            "primary": "#1e40af",      # Deep professional blue
            "primary_light": "#3b82f6", # Accessible light blue
            "accent": "#0ea5e9",       # Professional cyan
            "accent_light": "#06b6d4",  # Lighter teal

            # Node type colors - distinct but harmonious
            "tool": "#059669",         # Professional emerald green
            "prompt": "#7c3aed",       # Rich purple
            "step": "#dc2626",         # Professional red
            "query": "#f59e0b",        # Warm amber

            # UI colors - neutral and professional
            "background": "#f8fafc",   # Light neutral background
            "surface": "#ffffff",      # Pure white surface
            "border": "#e2e8f0",       # Subtle border gray
            "text_primary": "#1e293b", # Dark professional text
            "text_secondary": "#64748b", # Medium gray text

            # Status colors - accessible and clear
            "success": "#10b981",      # Success green
            "warning": "#f59e0b",      # Warning amber
            "error": "#ef4444",        # Error red
            "info": "#3b82f6",         # Info blue

            # Interactive states
            "hover": "#fbbf24",        # Gold hover state
            "selected": "#ec4899",     # Pink selection state
            "disabled": "#9ca3af",     # Gray disabled state
        }

        # Layout configuration for professional appearance
        self.layout_config = {
            "showlegend": True,
            "hovermode": "closest",
            "margin": dict(b=20, l=5, r=5, t=40),
            "annotations": [
                dict(
                    text="KGraph-MCP Knowledge Network - Track 2 Visualization",
                    showarrow=False,
                    xref="paper", yref="paper",
                    x=0.005, y=-0.002,
                    xanchor="left", yanchor="bottom",
                    font=dict(color=self.colors["text_secondary"], size=12)
                )
            ],
            "xaxis": dict(showgrid=False, zeroline=False, showticklabels=False),
            "yaxis": dict(showgrid=False, zeroline=False, showticklabels=False),
            "plot_bgcolor": self.colors["background"],
            "paper_bgcolor": self.colors["surface"],
        }

    def create_plan_visualization(self, planned_steps: list[PlannedStep],
                                query: str = "") -> Figure:
        """Create interactive visualization of planned steps.
        
        Args:
            planned_steps: List of PlannedStep objects to visualize
            query: Original user query for context
            
        Returns:
            Plotly Figure with interactive network visualization
        """
        if not planned_steps:
            return self._create_empty_visualization("No planned steps to visualize")

        try:
            # Create NetworkX graph
            G = nx.Graph()

            # Add query node as central hub
            query_text = query[:50] + "..." if len(query) > 50 else query
            G.add_node("query",
                      type="query",
                      label=f"Query: {query_text}",
                      size=20,
                      color=self.colors["query"])

            # Add nodes and edges for each planned step
            for i, step in enumerate(planned_steps):
                step_id = f"step_{i}"
                tool_id = f"tool_{step.tool.tool_id}"
                prompt_id = f"prompt_{step.prompt.prompt_id}"

                # Add step node
                relevance_text = f" (Score: {step.relevance_score:.2f})" if step.relevance_score else ""
                G.add_node(step_id,
                          type="step",
                          label=f"Step {i+1}{relevance_text}",
                          size=15,
                          color=self.colors["step"],
                          relevance=step.relevance_score or 0.0)

                # Add tool node
                tool_label = f"🔧 {step.tool.name}"
                G.add_node(tool_id,
                          type="tool",
                          label=tool_label,
                          description=step.tool.description,
                          tags=", ".join(step.tool.tags) if step.tool.tags else "No tags",
                          size=12,
                          color=self.colors["tool"])

                # Add prompt node
                prompt_label = f"📝 {step.prompt.name}"
                G.add_node(prompt_id,
                          type="prompt",
                          label=prompt_label,
                          description=step.prompt.description,
                          difficulty=step.prompt.difficulty_level,
                          size=10,
                          color=self.colors["prompt"])

                # Add edges
                G.add_edge("query", step_id, weight=2.0)
                G.add_edge(step_id, tool_id, weight=1.5)
                G.add_edge(step_id, prompt_id, weight=1.5)
                G.add_edge(tool_id, prompt_id, weight=1.0)

            # Generate layout using spring layout for better aesthetics
            pos = nx.spring_layout(G, k=3, iterations=50, seed=42)

            # Create Plotly traces
            traces = self._create_network_traces(G, pos)

            # Create figure
            fig = go.Figure(data=traces)
            fig.update_layout(
                **self.layout_config,
                title=dict(
                    text="🧠 KGraph-MCP Planning Network",
                    x=0.5,
                    font=dict(size=20, color=self.colors["text_primary"])
                )
            )

            return fig

        except Exception as e:
            logger.error(f"Error creating plan visualization: {e}")
            return self._create_error_visualization(f"Visualization error: {e!s}")

    def create_tool_ecosystem_visualization(self, tools: list[MCPTool],
                                          prompts: list[MCPPrompt]) -> Figure:
        """Create ecosystem view of all tools and prompts.
        
        Args:
            tools: List of available tools
            prompts: List of available prompts
            
        Returns:
            Plotly Figure with ecosystem visualization
        """
        try:
            G = nx.Graph()

            # Add tool nodes
            for tool in tools:
                tool_id = f"tool_{tool.tool_id}"
                G.add_node(tool_id,
                          type="tool",
                          label=f"🔧 {tool.name}",
                          description=tool.description,
                          tags=", ".join(tool.tags) if tool.tags else "No tags",
                          size=15,
                          color=self.colors["tool"])

            # Add prompt nodes and connect to tools
            for prompt in prompts:
                prompt_id = f"prompt_{prompt.prompt_id}"
                tool_id = f"tool_{prompt.target_tool_id}"

                G.add_node(prompt_id,
                          type="prompt",
                          label=f"📝 {prompt.name}",
                          description=prompt.description,
                          difficulty=prompt.difficulty_level,
                          size=10,
                          color=self.colors["prompt"])

                # Connect prompt to its target tool
                if tool_id in G.nodes():
                    G.add_edge(tool_id, prompt_id, weight=1.0)

            # Group nodes by tags for better layout
            pos = self._create_clustered_layout(G, tools)

            # Create traces
            traces = self._create_network_traces(G, pos)

            # Create figure
            fig = go.Figure(data=traces)
            fig.update_layout(
                **self.layout_config,
                title=dict(
                    text="🌐 KGraph-MCP Tool Ecosystem",
                    x=0.5,
                    font=dict(size=20, color=self.colors["text_primary"])
                )
            )

            return fig

        except Exception as e:
            logger.error(f"Error creating ecosystem visualization: {e}")
            return self._create_error_visualization(f"Ecosystem visualization error: {e!s}")

    def _create_network_traces(self, G: nx.Graph, pos: dict) -> list[go.Scatter]:
        """Create Plotly traces for network visualization."""
        traces = []

        # Create edge traces
        edge_x, edge_y = [], []
        for edge in G.edges():
            x0, y0 = pos[edge[0]]
            x1, y1 = pos[edge[1]]
            edge_x.extend([x0, x1, None])
            edge_y.extend([y0, y1, None])

        edge_trace = go.Scatter(
            x=edge_x, y=edge_y,
            line=dict(width=2, color=self.colors["border"]),
            hoverinfo="none",
            mode="lines",
            name="Connections",
            showlegend=False
        )
        traces.append(edge_trace)

        # Create node traces by type
        node_types = set(G.nodes[node].get("type", "unknown") for node in G.nodes())

        for node_type in node_types:
            nodes_of_type = [node for node in G.nodes()
                           if G.nodes[node].get("type") == node_type]

            if not nodes_of_type:
                continue

            node_x = [pos[node][0] for node in nodes_of_type]
            node_y = [pos[node][1] for node in nodes_of_type]

            # Get node attributes
            node_colors = [G.nodes[node].get("color", self.colors["disabled"])
                          for node in nodes_of_type]
            node_sizes = [G.nodes[node].get("size", 10) for node in nodes_of_type]
            node_labels = [G.nodes[node].get("label", node) for node in nodes_of_type]

            # Create hover text
            hover_texts = []
            for node in nodes_of_type:
                node_data = G.nodes[node]
                hover_text = f"<b>{node_data.get('label', node)}</b><br>"

                if "description" in node_data:
                    hover_text += f"Description: {node_data['description']}<br>"
                if "tags" in node_data:
                    hover_text += f"Tags: {node_data['tags']}<br>"
                if "difficulty" in node_data:
                    hover_text += f"Difficulty: {node_data['difficulty']}<br>"
                if "relevance" in node_data:
                    hover_text += f"Relevance: {node_data['relevance']:.2f}<br>"

                hover_texts.append(hover_text)

            # Create node trace
            node_trace = go.Scatter(
                x=node_x, y=node_y,
                mode="markers+text",
                text=node_labels,
                textposition="middle center",
                textfont=dict(size=10, color=self.colors["surface"]),
                hovertemplate="%{hovertext}<extra></extra>",
                hovertext=hover_texts,
                marker=dict(
                    size=node_sizes,
                    color=node_colors,
                    line=dict(width=2, color=self.colors["surface"]),
                    sizemode="diameter"
                ),
                name=node_type.title(),
                showlegend=True
            )
            traces.append(node_trace)

        return traces

    def _create_clustered_layout(self, G: nx.Graph, tools: list[MCPTool]) -> dict:
        """Create clustered layout based on tool tags."""
        # Group tools by their primary tag
        tag_groups = {}
        for tool in tools:
            primary_tag = tool.tags[0] if tool.tags else "general"
            if primary_tag not in tag_groups:
                tag_groups[primary_tag] = []
            tag_groups[primary_tag].append(f"tool_{tool.tool_id}")

        # Create positions with clustering
        pos = {}
        angle_step = 2 * 3.14159 / len(tag_groups) if tag_groups else 1

        for i, (tag, tool_ids) in enumerate(tag_groups.items()):
            center_x = 3 * np.cos(i * angle_step)
            center_y = 3 * np.sin(i * angle_step)

            # Layout tools in this group
            for j, tool_id in enumerate(tool_ids):
                offset_angle = j * 0.5
                offset_radius = 0.8
                pos[tool_id] = (
                    center_x + offset_radius * np.cos(offset_angle),
                    center_y + offset_radius * np.sin(offset_angle)
                )

        # Position prompts near their tools
        for node in G.nodes():
            if node.startswith("prompt_"):
                # Find connected tools
                connected_tools = [neighbor for neighbor in G.neighbors(node)
                                 if neighbor.startswith("tool_")]
                if connected_tools:
                    tool_pos = pos[connected_tools[0]]
                    # Offset prompt position slightly
                    pos[node] = (tool_pos[0] + 0.3, tool_pos[1] + 0.3)
                else:
                    pos[node] = (0, 0)  # Default position

        return pos

    def _create_empty_visualization(self, message: str) -> Figure:
        """Create empty state visualization."""
        fig = go.Figure()
        fig.add_annotation(
            text=f"📊 {message}",
            xref="paper", yref="paper",
            x=0.5, y=0.5,
            showarrow=False,
            font=dict(size=16, color=self.colors["text_secondary"])
        )
        fig.update_layout(**self.layout_config)
        return fig

    def _create_error_visualization(self, error_msg: str) -> Figure:
        """Create error state visualization."""
        fig = go.Figure()
        fig.add_annotation(
            text=f"⚠️ {error_msg}",
            xref="paper", yref="paper",
            x=0.5, y=0.5,
            showarrow=False,
            font=dict(size=16, color=self.colors["error"])
        )
        fig.update_layout(**self.layout_config)
        return fig

    def create_performance_metrics_chart(self, metrics_data: dict[str, Any]) -> Figure:
        """Create performance metrics visualization for demo purposes."""
        try:
            # Sample metrics for demonstration
            categories = ["Response Time", "Accuracy", "Coverage", "Relevance", "User Satisfaction"]
            values = [95, 88, 92, 90, 94]  # Sample scores

            # Create radar chart
            fig = go.Figure(data=go.Scatterpolar(
                r=values,
                theta=categories,
                fill="toself",
                fillcolor=f"rgba({self._hex_to_rgb(self.colors['primary'])}, 0.3)",
                line=dict(color=self.colors["primary"], width=3),
                marker=dict(size=8, color=self.colors["accent"]),
                name="KGraph-MCP Performance"
            ))

            fig.update_layout(
                polar=dict(
                    radialaxis=dict(
                        visible=True,
                        range=[0, 100],
                        tickmode="linear",
                        tick0=0,
                        dtick=20,
                        gridcolor=self.colors["border"]
                    ),
                    angularaxis=dict(
                        gridcolor=self.colors["border"]
                    )
                ),
                showlegend=True,
                title=dict(
                    text="📈 Platform Performance Metrics",
                    x=0.5,
                    font=dict(size=18, color=self.colors["text_primary"])
                ),
                plot_bgcolor=self.colors["background"],
                paper_bgcolor=self.colors["surface"]
            )

            return fig

        except Exception as e:
            logger.error(f"Error creating performance chart: {e}")
            return self._create_error_visualization(f"Performance chart error: {e!s}")

    def _hex_to_rgb(self, hex_color: str) -> str:
        """Convert hex color to RGB string."""
        hex_color = hex_color.lstrip("#")
        return f"{int(hex_color[0:2], 16)}, {int(hex_color[2:4], 16)}, {int(hex_color[4:6], 16)}"


# Convenience function for easy integration
def create_plan_visualization(planned_steps: list[PlannedStep], query: str = "") -> Figure:
    """Convenience function to create plan visualization."""
    visualizer = KGVisualizer()
    return visualizer.create_plan_visualization(planned_steps, query)


def create_ecosystem_visualization(tools: list[MCPTool], prompts: list[MCPPrompt]) -> Figure:
    """Convenience function to create ecosystem visualization."""
    visualizer = KGVisualizer()
    return visualizer.create_tool_ecosystem_visualization(tools, prompts)