"""
Knowledge Graph Visualizer Module

This module provides comprehensive visualization capabilities for knowledge graphs in the
Semantica framework, including interactive network graphs, community visualizations,
centrality-based node sizing, entity/relationship distributions, and relationship matrices.

Key Features:
    - Interactive network graph visualizations
    - Community detection visualization with color coding
    - Centrality-based node sizing and coloring
    - Entity type distribution charts
    - Relationship frequency matrices
    - Multiple layout algorithms (force-directed, hierarchical, circular)
    - Customizable color schemes and node/edge styling

Main Classes:
    - KGVisualizer: Main knowledge graph visualizer coordinator

Example Usage:
    >>> from semantica.visualization import KGVisualizer
    >>> viz = KGVisualizer(layout="force", color_scheme="vibrant")
    >>> fig = viz.visualize_network(graph, output="interactive")
    >>> viz.visualize_communities(graph, communities, file_path="communities.html")
    >>> viz.visualize_centrality(graph, centrality, centrality_type="degree")
    >>> viz.visualize_entity_types(graph, file_path="entity_types.png")
    >>> viz.visualize_relationship_matrix(graph, output="interactive")

Author: Semantica Contributors
License: MIT
"""

from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union

import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

from ..utils.exceptions import ProcessingError
from ..utils.logging import get_logger
from ..utils.progress_tracker import get_progress_tracker
from .utils.color_schemes import ColorPalette, ColorScheme
from .utils.export_formats import (
    export_matplotlib_figure,
    export_plotly_figure,
    save_html,
)
from .utils.layout_algorithms import (
    CircularLayout,
    ForceDirectedLayout,
    HierarchicalLayout,
)


class KGVisualizer:
    """
    Knowledge graph visualizer.

    Provides various visualization methods for knowledge graphs including:
    - Interactive network graphs
    - Community visualizations
    - Centrality visualizations
    - Entity/relationship distribution charts
    """

    def __init__(self, **config):
        """
        Initialize knowledge graph visualizer.

        Args:
            **config: Configuration options:
                - layout: Layout algorithm ("force", "hierarchical", "circular")
                - color_scheme: Color scheme name
                - node_size: Base node size
                - edge_width: Edge width
        """
        self.logger = get_logger("kg_visualizer")
        self.config = config
        self.progress_tracker = get_progress_tracker()

        self.layout_type = config.get("layout", "force")
        color_scheme_name = config.get("color_scheme", "default")
        try:
            self.color_scheme = ColorScheme[color_scheme_name.upper()]
        except (KeyError, AttributeError):
            self.color_scheme = ColorScheme.DEFAULT
        self.node_size = config.get("node_size", 10)
        self.edge_width = config.get("edge_width", 1)

        # Initialize layout algorithms
        self.force_layout = ForceDirectedLayout(**config)
        self.hierarchical_layout = HierarchicalLayout(**config)
        self.circular_layout = CircularLayout(**config)

    def visualize_network(
        self,
        graph: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options,
    ) -> Optional[Any]:
        """
        Visualize knowledge graph as interactive network.

        Args:
            graph: Knowledge graph dictionary with entities and relationships
            output: Output type ("interactive", "html", "png", "svg")
            file_path: Output file path (required for non-interactive)
            **options: Additional visualization options

        Returns:
            Plotly figure (if interactive) or None
        """
        tracking_id = self.progress_tracker.start_tracking(
            module="visualization",
            submodule="KGVisualizer",
            message="Visualizing knowledge graph network",
        )

        try:
            self.logger.info("Visualizing knowledge graph network")

            # Extract entities and relationships
            self.progress_tracker.update_tracking(
                tracking_id, message="Extracting entities and relationships..."
            )
            entities = graph.get("entities", [])
            relationships = graph.get("relationships", [])

            if not entities:
                self.progress_tracker.stop_tracking(
                    tracking_id, status="failed", message="No entities found in graph"
                )
                raise ProcessingError("No entities found in graph")

            # Build node and edge lists
            self.progress_tracker.update_tracking(
                tracking_id, message="Building node and edge lists..."
            )
            nodes = self._extract_nodes(entities)
            edges = self._extract_edges(relationships, entities)

            self.progress_tracker.update_tracking(
                tracking_id, message="Generating visualization..."
            )
            result = self._visualize_network_plotly(
                nodes, edges, output, file_path, **options
            )

            self.progress_tracker.stop_tracking(
                tracking_id,
                status="completed",
                message=f"Visualization generated: {len(nodes)} nodes, {len(edges)} edges",
            )
            return result
        except Exception as e:
            self.progress_tracker.stop_tracking(
                tracking_id, status="failed", message=str(e)
            )
            raise

    def visualize_communities(
        self,
        graph: Dict[str, Any],
        communities: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options,
    ) -> Optional[Any]:
        """
        Visualize knowledge graph with community coloring.

        Args:
            graph: Knowledge graph dictionary
            communities: Community detection results with node_assignments
            output: Output type
            file_path: Output file path
            **options: Additional options

        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing knowledge graph communities")

        entities = graph.get("entities", [])
        relationships = graph.get("relationships", [])

        # Get community assignments
        node_assignments = communities.get("node_assignments", {})
        num_communities = communities.get(
            "num_communities", len(set(node_assignments.values()))
        )

        # Get colors for communities
        community_colors = ColorPalette.get_community_colors(
            num_communities, self.color_scheme
        )

        # Build nodes with community colors
        nodes = self._extract_nodes(entities)
        for node in nodes:
            entity_id = node.get("id", "")
            community_id = node_assignments.get(entity_id, 0)
            node["community"] = community_id
            node["color"] = community_colors[community_id % len(community_colors)]

        edges = self._extract_edges(relationships, entities)

        return self._visualize_network_plotly(
            nodes,
            edges,
            output,
            file_path,
            community_colors=community_colors,
            **options,
        )

    def visualize_centrality(
        self,
        graph: Dict[str, Any],
        centrality: Dict[str, Any],
        centrality_type: str = "degree",
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options,
    ) -> Optional[Any]:
        """
        Visualize knowledge graph with centrality-based node sizing/coloring.

        Args:
            graph: Knowledge graph dictionary
            centrality: Centrality calculation results
            centrality_type: Type of centrality ("degree", "betweenness", "closeness", "eigenvector")
            output: Output type
            file_path: Output file path
            **options: Additional options

        Returns:
            Visualization figure or None
        """
        self.logger.info(
            f"Visualizing knowledge graph with {centrality_type} centrality"
        )

        entities = graph.get("entities", [])
        relationships = graph.get("relationships", [])

        # Get centrality scores
        centrality_scores = centrality.get("centrality", {})
        if not centrality_scores:
            # Try to get from rankings
            rankings = centrality.get("rankings", [])
            centrality_scores = {r["node"]: r["score"] for r in rankings}

        # Build nodes with centrality information
        nodes = self._extract_nodes(entities)
        max_centrality = max(centrality_scores.values()) if centrality_scores else 1.0

        for node in nodes:
            entity_id = node.get("id", "")
            score = centrality_scores.get(entity_id, 0.0)
            node["centrality"] = score
            node["size"] = (
                self.node_size * (1 + 5 * (score / max_centrality))
                if max_centrality > 0
                else self.node_size
            )

        edges = self._extract_edges(relationships, entities)

        return self._visualize_network_plotly(
            nodes, edges, output, file_path, size_by_centrality=True, **options
        )

    def visualize_entity_types(
        self,
        graph: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options,
    ) -> Optional[Any]:
        """
        Visualize entity type distribution.

        Args:
            graph: Knowledge graph dictionary
            output: Output type
            file_path: Output file path
            **options: Additional options

        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing entity type distribution")

        entities = graph.get("entities", [])

        # Count entity types
        type_counts = {}
        for entity in entities:
            entity_type = entity.get("type") or entity.get("entity_type") or "Unknown"
            type_counts[entity_type] = type_counts.get(entity_type, 0) + 1

        fig = px.bar(
            x=list(type_counts.keys()),
            y=list(type_counts.values()),
            labels={"x": "Entity Type", "y": "Count"},
            title="Entity Type Distribution",
        )

        if output == "interactive":
            return fig
        elif file_path:
            export_plotly_figure(
                fig, file_path, format=output if output != "interactive" else "html"
            )
            return None

    def visualize_relationship_matrix(
        self,
        graph: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options,
    ) -> Optional[Any]:
        """
        Visualize relationship frequency matrix between entity types.

        Args:
            graph: Knowledge graph dictionary
            output: Output type
            file_path: Output file path
            **options: Additional options

        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing relationship matrix")

        entities = graph.get("entities", [])
        relationships = graph.get("relationships", [])

        # Build entity type map
        entity_type_map = {}
        for entity in entities:
            entity_id = entity.get("id") or entity.get("entity_id", "")
            entity_type = entity.get("type") or entity.get("entity_type", "Unknown")
            entity_type_map[entity_id] = entity_type

        # Count relationships between types
        type_matrix = {}
        for rel in relationships:
            source_id = rel.get("source") or rel.get("subject", "")
            target_id = rel.get("target") or rel.get("object", "")

            source_type = entity_type_map.get(source_id, "Unknown")
            target_type = entity_type_map.get(target_id, "Unknown")

            key = (source_type, target_type)
            type_matrix[key] = type_matrix.get(key, 0) + 1

        # Convert to matrix format
        source_types = sorted(set(t[0] for t in type_matrix.keys()))
        target_types = sorted(set(t[1] for t in type_matrix.keys()))

        matrix = np.zeros((len(source_types), len(target_types)))
        for i, source_type in enumerate(source_types):
            for j, target_type in enumerate(target_types):
                matrix[i, j] = type_matrix.get((source_type, target_type), 0)

        fig = go.Figure(
            data=go.Heatmap(
                z=matrix,
                x=target_types,
                y=source_types,
                colorscale="Viridis",
                text=matrix,
                texttemplate="%{text}",
                textfont={"size": 10},
            )
        )
        fig.update_layout(
            title="Relationship Frequency Matrix",
            xaxis_title="Target Entity Type",
            yaxis_title="Source Entity Type",
        )

        if output == "interactive":
            return fig
        elif file_path:
            export_plotly_figure(
                fig, file_path, format=output if output != "interactive" else "html"
            )
            return None

    def _extract_nodes(self, entities: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Extract nodes from entities."""
        nodes = []
        for entity in entities:
            node = {
                "id": entity.get("id") or entity.get("entity_id", ""),
                "label": entity.get("text")
                or entity.get("label")
                or entity.get("name", ""),
                "type": entity.get("type") or entity.get("entity_type", "entity"),
                "metadata": entity.get("metadata", {}),
            }
            nodes.append(node)
        return nodes

    def _extract_edges(
        self, relationships: List[Dict[str, Any]], entities: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Extract edges from relationships."""
        edges = []
        entity_map = {e.get("id") or e.get("entity_id", ""): e for e in entities}

        for rel in relationships:
            source_id = rel.get("source") or rel.get("subject", "")
            target_id = rel.get("target") or rel.get("object", "")

            if source_id in entity_map and target_id in entity_map:
                edge = {
                    "source": source_id,
                    "target": target_id,
                    "label": rel.get("type")
                    or rel.get("relationship_type")
                    or rel.get("predicate", ""),
                    "metadata": rel.get("metadata", {}),
                }
                edges.append(edge)

        return edges

    def _visualize_network_plotly(
        self,
        nodes: List[Dict[str, Any]],
        edges: List[Dict[str, Any]],
        output: str,
        file_path: Optional[Path],
        **options,
    ) -> Optional[Any]:
        """Create Plotly network visualization."""
        # Compute layout
        node_ids = [n["id"] for n in nodes]
        edge_tuples = [(e["source"], e["target"]) for e in edges]

        if self.layout_type == "hierarchical":
            pos = self.hierarchical_layout.compute_layout(
                node_ids, edge_tuples, **options
            )
        elif self.layout_type == "circular":
            pos = self.circular_layout.compute_layout(node_ids, edge_tuples, **options)
        else:
            pos = self.force_layout.compute_layout(node_ids, edge_tuples, **options)

        # Get entity type colors
        entity_types = list(set(n.get("type", "entity") for n in nodes))
        type_colors = ColorPalette.get_entity_type_colors(
            entity_types, self.color_scheme
        )

        # Prepare edge traces
        edge_x = []
        edge_y = []
        for edge in edges:
            source_pos = pos.get(edge["source"])
            target_pos = pos.get(edge["target"])
            if source_pos and target_pos:
                edge_x.extend([source_pos[0], target_pos[0], None])
                edge_y.extend([source_pos[1], target_pos[1], None])

        edge_trace = go.Scatter(
            x=edge_x,
            y=edge_y,
            line=dict(width=self.edge_width, color="#888"),
            hoverinfo="none",
            mode="lines",
        )

        # Prepare node traces
        node_x = [pos[n["id"]][0] for n in nodes if n["id"] in pos]
        node_y = [pos[n["id"]][1] for n in nodes if n["id"] in pos]
        node_text = [n["label"] for n in nodes if n["id"] in pos]
        node_colors = [
            type_colors.get(n.get("type", "entity"), "#888")
            for n in nodes
            if n["id"] in pos
        ]
        node_sizes = [n.get("size", self.node_size) for n in nodes if n["id"] in pos]

        node_trace = go.Scatter(
            x=node_x,
            y=node_y,
            mode="markers+text",
            hoverinfo="text",
            text=node_text,
            textposition="middle center",
            marker=dict(
                size=node_sizes, color=node_colors, line=dict(width=2, color="white")
            ),
        )

        fig = go.Figure(
            data=[edge_trace, node_trace],
            layout=go.Layout(
                title="Knowledge Graph Network",
                showlegend=False,
                hovermode="closest",
                margin=dict(b=20, l=5, r=5, t=40),
                xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                yaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
            ),
        )

        if output == "interactive":
            return fig
        elif file_path:
            export_plotly_figure(
                fig, file_path, format=output if output != "interactive" else "html"
            )
            return None
