import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from typing import Union, List, Dict, Any
from .walks.dtqw import DTQW1D, DTQW2D, DTQW3D
from .walks.graph_walk import GraphDTQW
from .walks.szegedy import SzegedyQW

def plot_probability_distribution(walker: Union[DTQW1D, DTQW2D, DTQW3D, GraphDTQW, SzegedyQW],
                                 **kwargs) -> plt.Figure:
    """
    Universal probability distribution plotting function.
    
    Parameters:
    -----------
    walker : quantum walk instance
        Any quantum walk object with get_probabilities() method
    **kwargs : dict
        Additional plotting parameters
        
    Returns:
    --------
    plt.Figure
        Matplotlib figure object
    """
    probs = walker.get_probabilities()
    
    if isinstance(walker, DTQW1D):
        return walker.plot(**kwargs)
    elif isinstance(walker, DTQW2D):
        return walker.plot(**kwargs)
    elif isinstance(walker, DTQW3D):
        return walker.plot(**kwargs)
    elif isinstance(walker, (GraphDTQW, SzegedyQW)):
        return walker.plot(probs, **kwargs)
    else:
        raise ValueError(f"Unsupported walker type: {type(walker)}")

def analyze_convergence(walker: Union[GraphDTQW, SzegedyQW], 
                       max_steps: int = 100,
                       metric: str = 'variation') -> Dict[str, Any]:
    """
    Analyze convergence behavior of quantum walk.
    
    Parameters:
    -----------
    walker : quantum walk instance
        Graph-based quantum walk
    max_steps : int
        Maximum steps to analyze
    metric : str
        Convergence metric: 'variation' or 'entropy'
        
    Returns:
    --------
    dict
        Convergence analysis results
    """
    if not hasattr(walker, 'evolve') or not hasattr(walker, 'get_probabilities'):
        raise ValueError("Walker must have evolve() and get_probabilities() methods")
    
    walker.reset()
    previous_probs = walker.get_probabilities()
    convergence_data = {
        'steps': [],
        'variation': [],
        'entropy': []
    }
    
    for step in range(1, max_steps + 1):
        walker.step()
        current_probs = walker.get_probabilities()
        
        # Total variation distance
        variation = 0.5 * np.sum(np.abs(current_probs - previous_probs))
        
        # Shannon entropy
        entropy = -np.sum(current_probs * np.log(current_probs + 1e-10))
        
        convergence_data['steps'].append(step)
        convergence_data['variation'].append(variation)
        convergence_data['entropy'].append(entropy)
        
        previous_probs = current_probs
    
    return convergence_data

def validate_graph(graph: nx.Graph, 
                  check_connected: bool = True,
                  check_weights: bool = True) -> Dict[str, Any]:
    """
    Validate graph properties for quantum walks.
    
    Parameters:
    -----------
    graph : nx.Graph
        Input graph
    check_connected : bool
        Check if graph is connected
    check_weights : bool
        Check for negative weights
        
    Returns:
    --------
    dict
        Validation results
    """
    results = {
        'is_valid': True,
        'warnings': [],
        'properties': {}
    }
    
    # Basic properties
    results['properties']['num_nodes'] = graph.number_of_nodes()
    results['properties']['num_edges'] = graph.number_of_edges()
    results['properties']['is_directed'] = graph.is_directed()
    
    # Check connectivity
    if check_connected and not nx.is_connected(graph.to_undirected() if graph.is_directed() else graph):
        results['warnings'].append("Graph is not connected")
        results['is_valid'] = False
    
    # Check weights
    if check_weights:
        for u, v, data in graph.edges(data=True):
            weight = data.get('weight', 1.0)
            if weight < 0:
                results['warnings'].append(f"Negative weight found on edge ({u}, {v})")
                results['is_valid'] = False
                break
    
    # Degree statistics
    degrees = [d for _, d in graph.degree()]
    results['properties']['avg_degree'] = np.mean(degrees)
    results['properties']['max_degree'] = np.max(degrees)
    results['properties']['min_degree'] = np.min(degrees)
    
    return results

def compare_classical_quantum(graph: nx.Graph, 
                            alpha: float = 0.85) -> Dict[str, Any]:
    """
    Compare classical and quantum PageRank results.
    
    Parameters:
    -----------
    graph : nx.Graph
        Input graph
    alpha : float
        Damping factor
        
    Returns:
    --------
    dict
        Comparison results
    """
    # Classical PageRank
    classical_pr = nx.pagerank(graph, alpha=alpha)
    classical_scores = np.array([classical_pr[node] for node in sorted(graph.nodes())])
    
    # Quantum PageRank
    quantum_pr = QuantumPageRank(graph, alpha=alpha)
    quantum_scores = quantum_pr.get_ranking()
    
    # Comparison metrics
    correlation = np.corrcoef(classical_scores, quantum_scores)[0, 1]
    kl_divergence = np.sum(classical_scores * np.log(classical_scores / (quantum_scores + 1e-10)))
    
    return {
        'classical_scores': classical_scores,
        'quantum_scores': quantum_scores,
        'correlation': correlation,
        'kl_divergence': kl_divergence,
        'sorted_classical': sorted(classical_pr.items(), key=lambda x: x[1], reverse=True),
        'sorted_quantum': quantum_pr.get_sorted_ranking()
    }

def plot_convergence_analysis(convergence_data: Dict[str, Any], 
                            figsize: Tuple[int, int] = (12, 5)) -> plt.Figure:
    """
    Plot convergence analysis results.
    
    Parameters:
    -----------
    convergence_data : dict
        Output from analyze_convergence()
    figsize : tuple
        Figure size
        
    Returns:
    --------
    plt.Figure
        Matplotlib figure
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)
    
    # Variation plot
    ax1.plot(convergence_data['steps'], convergence_data['variation'], 
             'b-', linewidth=2)
    ax1.set_xlabel('Steps')
    ax1.set_ylabel('Total Variation Distance')
    ax1.set_title('Convergence: Variation Distance')
    ax1.grid(True, alpha=0.3)
    
    # Entropy plot
    ax2.plot(convergence_data['steps'], convergence_data['entropy'],
             'r-', linewidth=2)
    ax2.set_xlabel('Steps')
    ax2.set_ylabel('Shannon Entropy')
    ax2.set_title('Convergence: Entropy')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    return fig

def export_results(walker: Union[DTQW1D, DTQW2D, DTQW3D, GraphDTQW, SzegedyQW],
                  filename: str, format: str = 'npy'):
    """
    Export quantum walk results to file.
    
    Parameters:
    -----------
    walker : quantum walk instance
        Quantum walk object
    filename : str
        Output filename
    format : str
        Export format: 'npy', 'csv', or 'json'
    """
    probs = walker.get_probabilities()
    
    if format == 'npy':
        np.save(filename, probs)
    elif format == 'csv':
        np.savetxt(filename, probs, delimiter=',')
    elif format == 'json':
        import json
        if hasattr(walker, 'nodes'):
            # Graph-based walker
            results = {str(node): float(prob) for node, prob in zip(walker.nodes, probs)}
        else:
            # Lattice walker
            results = {'probabilities': probs.tolist()}
        with open(filename, 'w') as f:
            json.dump(results, f, indent=2)
    else:
        raise ValueError(f"Unsupported format: {format}")