"""
Ontology Visualizer Module

This module provides comprehensive visualization capabilities for ontologies in the
Semantica framework, including class hierarchy trees, property graphs, ontology structure
networks, class-property matrices, and ontology metrics dashboards.

Key Features:
    - Class hierarchy tree visualizations
    - Property graphs with domain/range relationships
    - Ontology structure network visualizations
    - Class-property relationship matrices
    - Ontology metrics dashboards
    - Support for multiple input formats (dict, SemanticNetwork, OntologyGenerator result)
    - Graphviz integration for hierarchy visualization
    - Interactive and static output formats

Main Classes:
    - OntologyVisualizer: Main ontology visualizer coordinator

Example Usage:
    >>> from semantica.visualization import OntologyVisualizer
    >>> viz = OntologyVisualizer(color_scheme="default")
    >>> fig = viz.visualize_hierarchy(ontology, output="interactive")
    >>> viz.visualize_properties(ontology, file_path="properties.html")
    >>> viz.visualize_structure(ontology, output="interactive")
    >>> viz.visualize_class_property_matrix(ontology, file_path="matrix.png")
    >>> viz.visualize_metrics(ontology, output="interactive")
    >>> viz.visualize_semantic_model(semantic_model, file_path="model.html")

Author: Semantica Contributors
License: MIT
"""

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

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

try:
    import graphviz
except ImportError:
    graphviz = None

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


class OntologyVisualizer:
    """
    Ontology visualizer.
    
    Provides visualization methods for ontologies including:
    - Class hierarchy trees
    - Property graphs
    - Ontology structure networks
    - Class-property matrices
    """
    
    def __init__(self, **config):
        """
        Initialize ontology visualizer.
        
        Args:
            **config: Configuration options:
                - color_scheme: Color scheme name
                - node_size: Base node size
        """
        self.logger = get_logger("ontology_visualizer")
        self.config = config
        self.progress_tracker = get_progress_tracker()
        
        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", 15)
    
    def visualize_hierarchy(
        self,
        ontology: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options
    ) -> Optional[Any]:
        """
        Visualize class hierarchy as tree.
        
        Args:
            ontology: Ontology dictionary with classes, or SemanticNetwork object, 
                     or ontology generator result
            output: Output type ("interactive", "html", "png", "svg", "dot")
            file_path: Output file path
            **options: Additional options
            
        Returns:
            Visualization figure or None
        """
        tracking_id = self.progress_tracker.start_tracking(
            module="visualization",
            submodule="OntologyVisualizer",
            message="Visualizing ontology class hierarchy"
        )
        
        try:
            self.logger.info("Visualizing ontology class hierarchy")
            
            # Handle different input formats
            self.progress_tracker.update_tracking(tracking_id, message="Extracting classes from ontology...")
            if hasattr(ontology, "classes"):
                # OntologyGenerator result object
                classes = ontology.classes if hasattr(ontology, "classes") else []
            elif isinstance(ontology, dict):
                classes = ontology.get("classes", ontology.get("class_definitions", []))
            else:
                classes = []
            
            # If no classes, try to extract from semantic network
            if not classes and isinstance(ontology, dict):
                # Check if it's a semantic model or semantic network
                semantic_network = ontology.get("semantic_network", ontology.get("network"))
                if semantic_network:
                    classes = self._extract_classes_from_semantic_network(semantic_network)
            
            if not classes:
                self.progress_tracker.stop_tracking(tracking_id, status="failed",
                                                  message="No classes found in ontology")
                raise ProcessingError("No classes found in ontology. Please provide classes or a semantic network.")
            
            # If output is dot and graphviz is available, use it
            if output == "dot" and graphviz is not None and file_path:
                self.progress_tracker.update_tracking(tracking_id, message="Generating Graphviz visualization...")
                result = self._visualize_hierarchy_graphviz(classes, file_path, **options)
                self.progress_tracker.stop_tracking(tracking_id, status="completed",
                                                  message=f"Hierarchy visualized: {len(classes)} classes")
                return result
            
            # Build hierarchy tree
            self.progress_tracker.update_tracking(tracking_id, message="Building hierarchy tree...")
            hierarchy = self._build_hierarchy_tree(classes)
            
            self.progress_tracker.update_tracking(tracking_id, message="Generating visualization...")
            result = self._visualize_hierarchy_plotly(hierarchy, classes, output, file_path, **options)
            
            self.progress_tracker.stop_tracking(tracking_id, status="completed",
                                              message=f"Hierarchy visualized: {len(classes)} classes")
            return result
        except Exception as e:
            self.progress_tracker.stop_tracking(tracking_id, status="failed", message=str(e))
            raise
    
    def visualize_properties(
        self,
        ontology: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options
    ) -> Optional[Any]:
        """
        Visualize property graph showing properties and their domains/ranges.
        
        Args:
            ontology: Ontology dictionary, SemanticNetwork, or ontology generator result
            output: Output type
            file_path: Output file path
            **options: Additional options
            
        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing ontology properties")
        
        # Handle different input formats
        if hasattr(ontology, "properties"):
            properties = ontology.properties if hasattr(ontology, "properties") else []
            classes = ontology.classes if hasattr(ontology, "classes") else []
        elif isinstance(ontology, dict):
            properties = ontology.get("properties", ontology.get("property_definitions", []))
            classes = ontology.get("classes", ontology.get("class_definitions", []))
        else:
            properties = []
            classes = []
        
        # If no properties, try to extract from semantic network
        if not properties and isinstance(ontology, dict):
            semantic_network = ontology.get("semantic_network", ontology.get("network"))
            if semantic_network:
                properties = self._extract_properties_from_semantic_network(semantic_network)
        
        if not properties:
            raise ProcessingError("No properties found in ontology")
        
        return self._visualize_properties_plotly(properties, classes, output, file_path, **options)
    
    def visualize_structure(
        self,
        ontology: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options
    ) -> Optional[Any]:
        """
        Visualize ontology structure as network.
        
        Args:
            ontology: Ontology dictionary
            output: Output type
            file_path: Output file path
            **options: Additional options
            
        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing ontology structure")
        
        classes = ontology.get("classes", [])
        properties = ontology.get("properties", [])
        
        # Build nodes and edges
        nodes = []
        edges = []
        
        # Add class nodes
        for cls in classes:
            nodes.append({
                "id": cls.get("name") or cls.get("uri", ""),
                "label": cls.get("label") or cls.get("name", ""),
                "type": "class"
            })
            
            # Add hierarchy edges
            parent = cls.get("parent") or cls.get("subClassOf")
            if parent:
                edges.append({
                    "source": cls.get("name") or cls.get("uri", ""),
                    "target": parent,
                    "type": "subClassOf"
                })
        
        # Add property nodes
        for prop in properties:
            prop_name = prop.get("name") or prop.get("uri", "")
            nodes.append({
                "id": prop_name,
                "label": prop.get("label") or prop.get("name", ""),
                "type": "property"
            })
            
            # Add domain edges
            domain = prop.get("domain")
            if domain:
                edges.append({
                    "source": prop_name,
                    "target": domain,
                    "type": "domain"
                })
            
            # Add range edges
            range_val = prop.get("range")
            if range_val:
                edges.append({
                    "source": prop_name,
                    "target": range_val,
                    "type": "range"
                })
        
        return self._visualize_structure_plotly(nodes, edges, output, file_path, **options)
    
    def visualize_class_property_matrix(
        self,
        ontology: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options
    ) -> Optional[Any]:
        """
        Visualize class-property matrix showing which properties belong to which classes.
        
        Args:
            ontology: Ontology dictionary
            output: Output type
            file_path: Output file path
            **options: Additional options
            
        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing class-property matrix")
        
        classes = ontology.get("classes", [])
        properties = ontology.get("properties", [])
        
        # Build matrix
        class_names = [cls.get("name") or cls.get("label", "") for cls in classes]
        property_names = [prop.get("name") or prop.get("label", "") for prop in properties]
        
        matrix = []
        for cls in classes:
            row = []
            cls_props = cls.get("properties", [])
            for prop in properties:
                prop_name = prop.get("name") or prop.get("uri", "")
                # Check if property belongs to class (via domain or direct property list)
                domain = prop.get("domain")
                has_prop = prop_name in cls_props or domain == cls.get("name")
                row.append(1 if has_prop else 0)
            matrix.append(row)
        
        fig = go.Figure(data=go.Heatmap(
            z=matrix,
            x=property_names,
            y=class_names,
            colorscale='Blues',
            text=matrix,
            texttemplate='%{text}',
            textfont={"size": 8}
        ))
        fig.update_layout(
            title="Class-Property Matrix",
            xaxis_title="Properties",
            yaxis_title="Classes",
            width=800,
            height=600
        )
        
        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_metrics(
        self,
        ontology: Dict[str, Any],
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options
    ) -> Optional[Any]:
        """
        Visualize ontology metrics dashboard.
        
        Args:
            ontology: Ontology dictionary
            output: Output type
            file_path: Output file path
            **options: Additional options
            
        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing ontology metrics")
        
        classes = ontology.get("classes", [])
        properties = ontology.get("properties", [])
        
        # Calculate metrics
        num_classes = len(classes)
        num_properties = len(properties)
        
        # Calculate hierarchy depth
        max_depth = 0
        for cls in classes:
            depth = self._calculate_class_depth(cls, classes)
            max_depth = max(max_depth, depth)
        
        fig = make_subplots(
            rows=1, cols=3,
            subplot_titles=("Class Count", "Property Count", "Max Depth"),
            specs=[[{"type": "indicator"}, {"type": "indicator"}, {"type": "indicator"}]]
        )
        
        fig.add_trace(go.Indicator(
            mode="number",
            value=num_classes,
            title={"text": "Classes"},
            domain={'x': [0, 0.33], 'y': [0, 1]}
        ), row=1, col=1)
        
        fig.add_trace(go.Indicator(
            mode="number",
            value=num_properties,
            title={"text": "Properties"},
            domain={'x': [0.33, 0.66], 'y': [0, 1]}
        ), row=1, col=2)
        
        fig.add_trace(go.Indicator(
            mode="number",
            value=max_depth,
            title={"text": "Max Depth"},
            domain={'x': [0.66, 1], 'y': [0, 1]}
        ), row=1, col=3)
        
        fig.update_layout(title="Ontology Metrics Dashboard")
        
        if output == "interactive":
            return fig
        elif file_path:
            export_plotly_figure(fig, file_path, format=output if output != "interactive" else "html")
            return None
    
    def _build_hierarchy_tree(self, classes: List[Dict[str, Any]]) -> Dict[str, List[str]]:
        """Build hierarchy tree structure."""
        tree = {}
        root_classes = []
        
        for cls in classes:
            cls_name = cls.get("name") or cls.get("uri", "")
            parent = cls.get("parent") or cls.get("subClassOf")
            
            if parent:
                if parent not in tree:
                    tree[parent] = []
                tree[parent].append(cls_name)
            else:
                root_classes.append(cls_name)
            
            if cls_name not in tree:
                tree[cls_name] = []
        
        return tree
    
    def _calculate_class_depth(self, cls: Dict[str, Any], all_classes: List[Dict[str, Any]]) -> int:
        """Calculate depth of class in hierarchy."""
        parent = cls.get("parent") or cls.get("subClassOf") or cls.get("superClassOf")
        if not parent:
            return 1
        
        # Find parent class
        for p_cls in all_classes:
            cls_name = p_cls.get("name") or p_cls.get("uri") or p_cls.get("label", "")
            if cls_name == parent:
                return 1 + self._calculate_class_depth(p_cls, all_classes)
        
        return 1
    
    def _extract_classes_from_semantic_network(self, semantic_network: Any) -> List[Dict[str, Any]]:
        """Extract class definitions from semantic network."""
        classes = []
        
        # Handle SemanticNetwork dataclass
        if hasattr(semantic_network, "nodes"):
            # Group nodes by type to form classes
            type_groups = {}
            for node in semantic_network.nodes:
                node_type = node.type if hasattr(node, "type") else "Unknown"
                if node_type not in type_groups:
                    type_groups[node_type] = []
                type_groups[node_type].append(node)
            
            # Create class definitions
            for node_type, nodes in type_groups.items():
                classes.append({
                    "name": node_type,
                    "label": node_type,
                    "uri": f"#{node_type}",
                    "instances": len(nodes),
                    "properties": list(set(
                        prop for node in nodes 
                        for prop in (node.properties.keys() if hasattr(node, "properties") else [])
                    ))
                })
        
        # Handle dictionary format
        elif isinstance(semantic_network, dict):
            nodes = semantic_network.get("nodes", [])
            type_groups = {}
            for node in nodes:
                node_type = node.get("type") if isinstance(node, dict) else (
                    node.type if hasattr(node, "type") else "Unknown"
                )
                if node_type not in type_groups:
                    type_groups[node_type] = []
                type_groups[node_type].append(node)
            
            for node_type, nodes in type_groups.items():
                classes.append({
                    "name": node_type,
                    "label": node_type,
                    "uri": f"#{node_type}",
                    "instances": len(nodes)
                })
        
        return classes
    
    def visualize_semantic_model(
        self,
        semantic_model: Any,
        output: str = "interactive",
        file_path: Optional[Union[str, Path]] = None,
        **options
    ) -> Optional[Any]:
        """
        Visualize semantic model from ontology generator.
        
        This method extracts and visualizes both the ontology structure
        and the underlying semantic network that generated it.
        
        Args:
            semantic_model: Semantic model from OntologyGenerator or semantic network
            output: Output type
            file_path: Output file path
            **options: Additional options
            
        Returns:
            Visualization figure or None
        """
        self.logger.info("Visualizing semantic model")
        
        # Handle OntologyGenerator result
        if hasattr(semantic_model, "semantic_network"):
            # Visualize the semantic network
            from .semantic_network_visualizer import SemanticNetworkVisualizer
            sem_net_viz = SemanticNetworkVisualizer(**self.config)
            return sem_net_viz.visualize_network(
                semantic_model.semantic_network,
                output=output,
                file_path=file_path,
                **options
            )
        
        # Handle dictionary format with semantic_network
        elif isinstance(semantic_model, dict):
            if "semantic_network" in semantic_model:
                from .semantic_network_visualizer import SemanticNetworkVisualizer
                sem_net_viz = SemanticNetworkVisualizer(**self.config)
                return sem_net_viz.visualize_network(
                    semantic_model["semantic_network"],
                    output=output,
                    file_path=file_path,
                    **options
                )
            # Otherwise treat as ontology
            else:
                return self.visualize_structure(semantic_model, output, file_path, **options)
        
        # Handle direct semantic network
        else:
            from .semantic_network_visualizer import SemanticNetworkVisualizer
            sem_net_viz = SemanticNetworkVisualizer(**self.config)
            return sem_net_viz.visualize_network(
                semantic_model,
                output=output,
                file_path=file_path,
                **options
            )
    
    def _visualize_hierarchy_plotly(
        self,
        hierarchy: Dict[str, List[str]],
        classes: List[Dict[str, Any]],
        output: str,
        file_path: Optional[Path],
        **options
    ) -> Optional[Any]:
        """Create Plotly hierarchy visualization."""
        # Build tree structure for plotly
        # This is a simplified tree visualization
        # For full tree, would need more complex layout
        
        # Get root classes
        all_class_names = {cls.get("name") or cls.get("uri", ""): cls for cls in classes}
        root_classes = [name for name in all_class_names.keys() if name not in [c for children in hierarchy.values() for c in children]]
        
        if not root_classes:
            root_classes = list(all_class_names.keys())[:1] if all_class_names else []
        
        # Build node and edge lists for tree
        nodes = []
        edges = []
        
        def add_node_and_children(cls_name, level=0, x_offset=0):
            nodes.append({
                "name": cls_name,
                "level": level,
                "x": x_offset,
                "y": -level
            })
            
            children = hierarchy.get(cls_name, [])
            child_width = 1.0 / max(len(children), 1)
            for i, child in enumerate(children):
                child_x = x_offset - 0.5 + (i + 0.5) * child_width
                edges.append({
                    "source": cls_name,
                    "target": child
                })
                add_node_and_children(child, level + 1, child_x)
        
        # Start from root
        root_x = 0.5
        root_width = 1.0 / max(len(root_classes), 1)
        for i, root in enumerate(root_classes):
            root_x = (i + 0.5) * root_width
            add_node_and_children(root, 0, root_x)
        
        # Create visualization
        edge_x = []
        edge_y = []
        node_positions = {n["name"]: (n["x"], n["y"]) for n in nodes}
        
        for edge in edges:
            source_pos = node_positions.get(edge["source"])
            target_pos = node_positions.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=2, color='#888'),
            hoverinfo='none',
            mode='lines'
        )
        
        node_x = [n["x"] for n in nodes]
        node_y = [n["y"] for n in nodes]
        node_text = [n["name"] for n in nodes]
        
        node_trace = go.Scatter(
            x=node_x, y=node_y,
            mode='markers+text',
            text=node_text,
            textposition="middle center",
            marker=dict(size=self.node_size*10, color='lightblue', line=dict(width=2, color='darkblue'))
        )
        
        fig = go.Figure(
            data=[edge_trace, node_trace],
            layout=go.Layout(
                title='Ontology Class Hierarchy',
                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
    
    def _visualize_hierarchy_graphviz(
        self,
        classes: List[Dict[str, Any]],
        file_path: Path,
        **options
    ) -> None:
        """Create Graphviz hierarchy visualization."""
        if graphviz is None:
            raise ProcessingError("Graphviz not available. Install with: pip install graphviz")
        
        dot = graphviz.Digraph(comment='Ontology Hierarchy')
        
        # Add nodes
        for cls in classes:
            cls_name = cls.get("name") or cls.get("label", "")
            dot.node(cls_name, cls_name)
        
        # Add edges
        for cls in classes:
            cls_name = cls.get("name") or cls.get("label", "")
            parent = cls.get("parent") or cls.get("subClassOf")
            if parent:
                dot.edge(parent, cls_name)
        
        dot.render(str(file_path), format='svg', cleanup=True)
        self.logger.info(f"Saved Graphviz hierarchy to {file_path}")
    
    def _visualize_properties_plotly(
        self,
        properties: List[Dict[str, Any]],
        classes: List[Dict[str, Any]],
        output: str,
        file_path: Optional[Path],
        **options
    ) -> Optional[Any]:
        """Create Plotly property visualization."""
        # Build property graph
        nodes = []
        edges = []
        
        # Add class nodes
        for cls in classes:
            cls_name = cls.get("name") or cls.get("uri", "")
            nodes.append({"id": cls_name, "label": cls.get("label") or cls_name, "type": "class"})
        
        # Add property nodes and edges
        for prop in properties:
            prop_name = prop.get("name") or prop.get("uri", "")
            nodes.append({"id": prop_name, "label": prop.get("label") or prop_name, "type": "property"})
            
            domain = prop.get("domain")
            if domain:
                edges.append({"source": prop_name, "target": domain, "type": "domain"})
            
            range_val = prop.get("range")
            if range_val:
                edges.append({"source": prop_name, "target": range_val, "type": "range"})
        
        # Use similar approach as KG visualizer
        from .kg_visualizer import KGVisualizer
        kg_viz = KGVisualizer(**self.config)
        return kg_viz._visualize_network_plotly(nodes, edges, output, file_path, **options)
    
    def _extract_properties_from_semantic_network(self, semantic_network: Any) -> List[Dict[str, Any]]:
        """Extract property definitions from semantic network."""
        properties = []
        
        # Handle SemanticNetwork dataclass
        if hasattr(semantic_network, "edges"):
            # Extract unique edge labels as properties
            property_types = set()
            for edge in semantic_network.edges:
                edge_label = edge.label if hasattr(edge, "label") else ""
                if edge_label and edge_label not in property_types:
                    property_types.add(edge_label)
                    properties.append({
                        "name": edge_label,
                        "label": edge_label,
                        "uri": f"#{edge_label}",
                        "domain": "Thing",  # Default domain
                        "range": "Thing"    # Default range
                    })
        
        # Handle dictionary format
        elif isinstance(semantic_network, dict):
            edges = semantic_network.get("edges", semantic_network.get("relationships", []))
            property_types = set()
            for edge in edges:
                edge_label = edge.get("label") if isinstance(edge, dict) else (
                    edge.label if hasattr(edge, "label") else ""
                )
                if edge_label and edge_label not in property_types:
                    property_types.add(edge_label)
                    properties.append({
                        "name": edge_label,
                        "label": edge_label,
                        "uri": f"#{edge_label}",
                        "domain": edge.get("domain", "Thing") if isinstance(edge, dict) else "Thing",
                        "range": edge.get("range", "Thing") if isinstance(edge, dict) else "Thing"
                    })
        
        return properties
    
    def _visualize_structure_plotly(
        self,
        nodes: List[Dict[str, Any]],
        edges: List[Dict[str, Any]],
        output: str,
        file_path: Optional[Path],
        **options
    ) -> Optional[Any]:
        """Create Plotly structure visualization."""
        from .kg_visualizer import KGVisualizer
        kg_viz = KGVisualizer(**self.config)
        return kg_viz._visualize_network_plotly(nodes, edges, output, file_path, **options)

