"""
Visualization utilities for PM4PY integration.

This module provides functions to visualize process models and other
outputs generated by PM4PY algorithms, wrapped in a Heraclitus-friendly interface.
"""
from typing import Optional, Dict, Any, List, Union, Tuple, BinaryIO
import warnings
import os
import tempfile

from heraclitus.data import EventLog
from heraclitus.utils.pm4py_converter import check_pm4py_available, to_pm4py, apply_pm4py_algorithm

try:
    import pm4py
    from pm4py.visualization.petri_net import visualizer as pn_visualizer
    from pm4py.visualization.bpmn import visualizer as bpmn_visualizer
    from pm4py.visualization.heuristics_net import visualizer as hn_visualizer
    from pm4py.visualization.dfg import visualizer as dfg_visualizer
    from pm4py.visualization.process_tree import visualizer as pt_visualizer
    from pm4py.algo.discovery.alpha import algorithm as alpha_miner
    from pm4py.algo.discovery.inductive import algorithm as inductive_miner
    from pm4py.algo.discovery.heuristics import algorithm as heuristics_miner
    from pm4py.algo.discovery.dfg import algorithm as dfg_discovery
    PM4PY_AVAILABLE = True
except ImportError:
    PM4PY_AVAILABLE = False


def visualize_process_model(
    event_log: EventLog,
    algorithm: str = "inductive",
    output_path: Optional[str] = None,
    show_labels: bool = True,
    show_performance: bool = False,
    layout: str = "horizontal",
    **kwargs
) -> Optional[str]:
    """
    Generate and visualize a process model from an event log.
    
    Args:
        event_log: The Heraclitus EventLog to analyze
        algorithm: The algorithm to use for process discovery
            Options: "alpha", "inductive", "heuristics", "dfg"
        output_path: Path to save the visualization. If None, a temporary file is used.
        show_labels: Whether to show activity labels in the visualization
        show_performance: Whether to show performance metrics in the visualization
        layout: Layout direction for the graph
        **kwargs: Additional parameters for the discovery algorithm
    
    Returns:
        Path to the saved visualization file or None if visualization failed
    
    Raises:
        ImportError: If PM4PY is not installed
        ValueError: If an unknown algorithm is specified
    """
    check_pm4py_available()
    
    # Convert to PM4PY format
    pm4py_log = to_pm4py(event_log)
    
    # Select algorithm and visualizer
    discovery_algo = None
    visualizer = None
    
    if algorithm == "alpha":
        net, initial_marking, final_marking = alpha_miner.apply(pm4py_log, **kwargs)
        parameters = pn_visualizer.Variants.PERFORMANCE.value["parameters"] if show_performance else None
        gviz = pn_visualizer.apply(net, initial_marking, final_marking, parameters=parameters)
        visualizer = pn_visualizer
    
    elif algorithm == "inductive":
        tree = inductive_miner.apply_tree(pm4py_log, **kwargs)
        
        # For process tree visualization
        if "tree" in kwargs and kwargs["tree"]:
            gviz = pt_visualizer.apply(tree)
            visualizer = pt_visualizer
        else:
            # Convert to Petri net
            net, initial_marking, final_marking = pm4py.convert_to_petri_net(tree)
            parameters = pn_visualizer.Variants.PERFORMANCE.value["parameters"] if show_performance else None
            gviz = pn_visualizer.apply(net, initial_marking, final_marking, parameters=parameters)
            visualizer = pn_visualizer
    
    elif algorithm == "heuristics":
        heu_net = heuristics_miner.apply_heu(pm4py_log, **kwargs)
        parameters = {
            hn_visualizer.Variants.PYDOTPLUS.value.Parameters.SHOW_LABELS: show_labels,
            hn_visualizer.Variants.PYDOTPLUS.value.Parameters.LAYOUT_DIRECTION: layout
        }
        gviz = hn_visualizer.apply(heu_net, parameters=parameters)
        visualizer = hn_visualizer
    
    elif algorithm == "dfg":
        dfg = dfg_discovery.apply(pm4py_log, **kwargs)
        parameters = {
            dfg_visualizer.Variants.FREQUENCY.value.Parameters.SHOW_LABELS: show_labels,
            dfg_visualizer.Variants.FREQUENCY.value.Parameters.LAYOUT_DIRECTION: layout
        }
        if show_performance:
            gviz = dfg_visualizer.apply(dfg, log=pm4py_log, variant=dfg_visualizer.Variants.PERFORMANCE, parameters=parameters)
        else:
            gviz = dfg_visualizer.apply(dfg, log=pm4py_log, variant=dfg_visualizer.Variants.FREQUENCY, parameters=parameters)
        visualizer = dfg_visualizer
    
    else:
        raise ValueError(f"Unknown algorithm: {algorithm}. "
                         "Available options: alpha, inductive, heuristics, dfg")
    
    # Handle output path
    if output_path is None:
        # Create a temporary file
        temp_dir = tempfile.gettempdir()
        output_path = os.path.join(temp_dir, f"process_model_{algorithm}.png")
    
    # Save visualization
    visualizer.save(gviz, output_path)
    
    return output_path


def convert_to_bpmn(
    event_log: EventLog,
    output_path: Optional[str] = None,
    algorithm: str = "inductive",
    **kwargs
) -> Optional[str]:
    """
    Convert an event log to a BPMN diagram.
    
    Args:
        event_log: The Heraclitus EventLog to convert
        output_path: Path to save the BPMN visualization. If None, a temporary file is used.
        algorithm: The algorithm to use for process discovery before BPMN conversion
            Options: "inductive" (recommended), "alpha"
        **kwargs: Additional parameters for the discovery algorithm
    
    Returns:
        Path to the saved BPMN visualization file or None if conversion failed
    
    Raises:
        ImportError: If PM4PY is not installed
        ValueError: If an unsupported algorithm is specified
    """
    check_pm4py_available()
    
    # Convert to PM4PY format
    pm4py_log = to_pm4py(event_log)
    
    # Discover process model
    if algorithm == "alpha":
        net, initial_marking, final_marking = alpha_miner.apply(pm4py_log, **kwargs)
    elif algorithm == "inductive":
        tree = inductive_miner.apply_tree(pm4py_log, **kwargs)
        net, initial_marking, final_marking = pm4py.convert_to_petri_net(tree)
    else:
        raise ValueError(f"Unsupported algorithm for BPMN conversion: {algorithm}. "
                         "Available options: inductive, alpha")
    
    # Convert to BPMN
    bpmn_graph = pm4py.convert_to_bpmn(net, initial_marking, final_marking)
    
    # Handle output path
    if output_path is None:
        temp_dir = tempfile.gettempdir()
        output_path = os.path.join(temp_dir, "process_model_bpmn.png")
    
    # Visualize BPMN
    gviz = bpmn_visualizer.apply(bpmn_graph)
    bpmn_visualizer.save(gviz, output_path)
    
    return output_path