"""
Network analysis examples using quantum walks.
"""

import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from quantumwalks import GraphDTQW, SzegedyQW, DTQWPageRank, QuantumPageRank

def example_graph_walk():
    """Graph quantum walk example."""
    print("Running graph quantum walk...")
    
    # Create a complex network
    graph = nx.connected_watts_strogatz_graph(20, 4, 0.3)
    
    # Initialize quantum walk
    gqw = GraphDTQW(graph, coin_type='grover', initial_node=0)
    gqw.evolve(10)
    
    # Plot results
    fig = gqw.plot(layout='spring')
    plt.show()
    
    # Compare finite vs infinite averages
    finite_avg = gqw.average_probabilities_finite(20)
    infinite_avg = gqw.average_probabilities_infinite()
    
    print("Finite average:", finite_avg)
    print("Infinite average:", infinite_avg)

def example_szegedy_walk():
    """Szegedy quantum walk example."""
    print("Running Szegedy quantum walk...")
    
    # Use Karate club graph
    graph = nx.karate_club_graph()
    
    # Initialize Szegedy walk
    sqw = SzegedyQW(graph, initial_nodes=[0, 1])
    sqw.evolve(34)
    
    # Plot both registers
    fig1 = sqw.plot(register='first')
    fig2 = sqw.plot(register='second')
    plt.show()

def example_pagerank_comparison():
    """Compare classical and quantum PageRank."""
    print("Running PageRank comparison...")
    
    # Create test graph
    graph = nx.erdos_renyi_graph(15, 0.3)
    
    # Classical PageRank
    classical_pr = nx.pagerank(graph)
    
    # Quantum PageRank
    qpr = QuantumPageRank(graph)
    quantum_pr = qpr.get_ranking()
    
    # DTQW PageRank
    dtqw_pr = DTQWPageRank(graph.to_directed(), num_steps=100, num_runs=10)
    dtqw_scores = dtqw_pr.get_ranking()
    
    # Plot comparison
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(15, 5))
    
    nodes = sorted(graph.nodes())
    
    ax1.bar(nodes, [classical_pr[n] for n in nodes])
    ax1.set_title('Classical PageRank')
    
    ax2.bar(nodes, quantum_pr)
    ax2.set_title('Quantum PageRank')
    
    ax3.bar(nodes, dtqw_scores)
    ax3.set_title('DTQW PageRank')
    
    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    example_graph_walk()
    example_szegedy_walk()
    example_pagerank_comparison()