import numpy as np
import networkx as nx
import matplotlib.pyplot as plt
from numpy.linalg import eig
from typing import Optional, List, Union

class SzegedyQW:
    """Szegedy quantum walk for complex networks."""
    
    def __init__(self, graph: nx.Graph, initial_nodes: Optional[List] = None):
        """
        Initialize Szegedy quantum walk.
        
        Parameters:
        -----------
        graph : nx.Graph
            NetworkX graph
        initial_nodes : list, optional
            Starting nodes for superposition
        """
        self.graph = graph
        self.nodes = sorted(graph.nodes())
        self.num_nodes = len(self.nodes)
        self.node_to_idx = {node: i for i, node in enumerate(self.nodes)}
        self.degrees = np.array([graph.degree(node, weight='weight') 
                               for node in self.nodes])
        self.degrees[self.degrees == 0] = 1
        
        # Build transition matrix
        adj = nx.to_numpy_array(graph, weight='weight')
        self.P = adj / self.degrees[:, np.newaxis]
        
        # Build edge space
        self.edges = [(u, v) for u in self.nodes for v in graph.neighbors(u)]
        self.num_edges = len(self.edges)
        self.edge_to_idx = {edge: i for i, edge in enumerate(self.edges)}
        self.idx_to_edge = {i: edge for i, edge in enumerate(self.edges)}
        
        # Build quantum walk operators
        self.psi = self._build_discriminator_states()
        self.Pi = self._build_projector()
        self.Ref_A = 2 * self.Pi - np.eye(self.num_edges, dtype=complex)
        self.S = self._build_swap_operator()
        self.Ref_B = self.S @ self.Ref_A @ self.S.T
        self.U = self.Ref_B @ self.Ref_A
        self.initial_state = self._build_initial_state(initial_nodes)
        self.state = self.initial_state.copy()
    
    def _build_discriminator_states(self) -> np.ndarray:
        psi = np.zeros((self.num_nodes, self.num_edges), dtype=complex)
        for u_idx, u in enumerate(self.nodes):
            for v in self.graph.neighbors(u):
                v_idx = self.node_to_idx[v]
                edge_idx = self.edge_to_idx[(u, v)]
                psi[u_idx, edge_idx] = np.sqrt(self.P[u_idx, v_idx])
        return psi
    
    def _build_projector(self) -> np.ndarray:
        return np.dot(self.psi.T, self.psi)
    
    def _build_swap_operator(self) -> np.ndarray:
        S = np.zeros((self.num_edges, self.num_edges), dtype=complex)
        for idx1, (u, v) in enumerate(self.edges):
            rev_edge = (v, u)
            if rev_edge in self.edge_to_idx:
                idx2 = self.edge_to_idx[rev_edge]
                S[idx1, idx2] = 1
        return S
    
    def _build_initial_state(self, initial_nodes: Optional[List]) -> np.ndarray:
        if initial_nodes is None:
            initial_nodes = self.nodes
            
        init_nodes_idx = [self.node_to_idx[node] for node in initial_nodes 
                         if node in self.node_to_idx]
        initial_state = np.zeros(self.num_edges, dtype=complex)
        
        for u_idx in init_nodes_idx:
            initial_state += self.psi[u_idx]
            
        return initial_state / np.linalg.norm(initial_state) * np.sqrt(len(init_nodes_idx))
    
    def reset(self):
        """Reset to initial state."""
        self.state = self.initial_state.copy()
    
    def step(self):
        """Perform one evolution step."""
        self.state = self.U @ self.state
    
    def evolve(self, num_steps: int):
        """Evolve for multiple steps."""
        self.reset()
        for _ in range(num_steps):
            self.step()
    
    def get_probabilities(self, register: str = 'first') -> np.ndarray:
        """
        Get probability distribution.
        
        Parameters:
        -----------
        register : str
            'first' or 'second' register
        """
        probs = np.zeros(self.num_nodes)
        abs_sq = np.abs(self.state) ** 2
        
        for edge_idx in range(self.num_edges):
            u, v = self.idx_to_edge[edge_idx]
            u_idx, v_idx = self.node_to_idx[u], self.node_to_idx[v]
            
            if register == 'first':
                probs[u_idx] += abs_sq[edge_idx]
            else:
                probs[v_idx] += abs_sq[edge_idx]
                
        return probs
    
    def average_probabilities_finite(self, num_steps: int, register: str = 'first') -> np.ndarray:
        """Compute finite-time average probabilities."""
        self.reset()
        avg_probs = np.zeros(self.num_nodes)
        for _ in range(num_steps):
            self.step()
            avg_probs += self.get_probabilities(register)
        return avg_probs / num_steps
    
    def average_probabilities_infinite(self, register: str = 'first') -> np.ndarray:
        """Compute infinite-time average probabilities."""
        eigenvalues, eigenvectors = eig(self.U)
        c = np.dot(eigenvectors.T.conj(), self.initial_state)
        avg_probs = np.zeros(self.num_nodes)
        
        for k in range(self.num_edges):
            u_k = eigenvectors[:, k]
            prob_k = np.zeros(self.num_nodes)
            abs_u_k_sq = np.abs(u_k) ** 2
            
            for edge_idx in range(self.num_edges):
                u, v = self.idx_to_edge[edge_idx]
                u_idx, v_idx = self.node_to_idx[u], self.node_to_idx[v]
                
                if register == 'first':
                    prob_k[u_idx] += abs_u_k_sq[edge_idx]
                else:
                    prob_k[v_idx] += abs_u_k_sq[edge_idx]
                    
            avg_probs += np.abs(c[k]) ** 2 * prob_k
            
        return avg_probs
    
    def plot(self, probs: Optional[np.ndarray] = None, 
             register: str = 'first', layout: str = 'spring', **kwargs):
        """Plot graph with probabilities."""
        if probs is None:
            probs = self.get_probabilities(register)
            
        pos = (nx.spring_layout(self.graph) if layout == 'spring' 
               else nx.circular_layout(self.graph))
        
        plt.figure(figsize=(10, 8))
        node_colors = [probs[self.node_to_idx[node]] for node in self.nodes]
        
        nx.draw(self.graph, pos,
                node_color=node_colors,
                cmap='viridis',
                node_size=500,
                with_labels=True,
                **kwargs)
        
        plt.colorbar(plt.cm.ScalarMappable(
            cmap='viridis',
            norm=plt.Normalize(vmin=min(node_colors), vmax=max(node_colors))
        ), label='Probability')
        
        plt.title(f'Szegedy Quantum Walk (Register: {register})')
        plt.tight_layout()
        return plt.gcf()