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

class DTQWPageRank:
    """Discrete-time quantum walk for node ranking."""
    
    def __init__(self, graph: nx.DiGraph, num_steps: int = 500, 
                 num_runs: int = 100, epsilon: float = 1e-10):
        """
        Initialize DTQW PageRank.
        
        Parameters:
        -----------
        graph : nx.DiGraph
            Directed graph
        num_steps : int
            Evolution steps per run
        num_runs : int
            Number of independent runs
        epsilon : float
            Small value for numerical stability
        """
        self.graph = graph
        self.nodes = sorted(graph.nodes())
        self.N = len(self.nodes)
        self.node_to_idx = {node: i for i, node in enumerate(self.nodes)}
        self.num_steps = num_steps
        self.num_runs = num_runs
        self.epsilon = epsilon
        
        # Build adjacency matrix and compute properties
        self.A = nx.to_numpy_array(graph, weight='weight').astype(float)
        self.in_deg = np.sum(self.A, axis=0)
        self.out_deg = np.sum(self.A, axis=1)
        self.alpha = self.in_deg / (self.in_deg + self.out_deg + self.epsilon)
        
        # SVD decomposition
        self.P, s, Vh = svd(self.A, full_matrices=False)
        self.Lambda = np.diag(s)
        self.Q = Vh.T
        self.U_down = self.P @ np.diag(np.exp(1j * self.Lambda)) @ self.Q
        
        # Build quantum walk operators
        self.dim = 2 * self.N
        self.C = self._build_coin_operator()
        self.S = self._build_shift_operator()
        self.U = self.S @ self.C
        self.initial_state = self._build_initial_state()
    
    def _build_coin_operator(self) -> np.ndarray:
        C = np.zeros((self.dim, self.dim), dtype=complex)
        for i in range(self.N):
            alpha_i = self.alpha[i]
            c11 = np.sqrt(1 / (alpha_i + 1))
            c12 = np.sqrt(alpha_i / (alpha_i + 1))
            coin_block = np.array([[c11, c12], [c12, -c11]], dtype=complex)
            start = 2 * i
            C[start:start + 2, start:start + 2] = coin_block
        return C
    
    def _build_shift_operator(self) -> np.ndarray:
        S = np.zeros((self.dim, self.dim), dtype=complex)
        S[:self.N, :self.N] = np.eye(self.N)
        S[self.N:, self.N:] = self.U_down
        return S
    
    def _build_initial_state(self) -> np.ndarray:
        state = np.zeros(self.dim, dtype=complex)
        amp = 1 / np.sqrt(2 * self.N)
        state[:self.N] = amp
        state[self.N:] = amp
        return state
    
    def _evolve_single_run(self) -> np.ndarray:
        state = self.initial_state.copy()
        for _ in range(self.num_steps):
            state = self.U @ state
            
        probs = np.zeros(self.N)
        probs += np.abs(state[:self.N]) ** 2
        probs += np.abs(state[self.N:]) ** 2
        return probs
    
    def get_ranking(self) -> np.ndarray:
        """Compute quantum ranking scores."""
        total_probs = np.zeros(self.N)
        for _ in range(self.num_runs):
            total_probs += self._evolve_single_run()
            
        ranking = total_probs / self.num_runs
        return ranking / np.sum(ranking)
    
    def get_sorted_ranking(self) -> List[Tuple]:
        """Get sorted ranking as (node, score) pairs."""
        probs = self.get_ranking()
        idx = np.argsort(probs)[::-1]
        return [(self.nodes[j], probs[j]) for j in idx]
    
    def plot(self, probs: Optional[np.ndarray] = None, 
             plot_type: str = 'bar', layout: str = 'spring', **kwargs):
        """Plot ranking results."""
        if probs is None:
            probs = self.get_ranking()
            
        if plot_type == 'bar':
            idx = np.argsort(probs)[::-1]
            sorted_probs = probs[idx]
            sorted_nodes = [str(self.nodes[j]) for j in idx]
            
            plt.figure(figsize=(12, 6))
            plt.bar(sorted_nodes, sorted_probs, **kwargs)
            plt.xlabel('Node')
            plt.ylabel('Quantum Ranking Score')
            plt.title('DTQW Node Ranking Distribution')
            plt.xticks(rotation=45)
            plt.tight_layout()
            
        elif plot_type == 'graph':
            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='Quantum Ranking Score')
            
            plt.title('DTQW Node Ranking on Network')
            plt.tight_layout()
            
        else:
            raise ValueError("plot_type must be 'bar' or 'graph'")
            
        return plt.gcf()

class QuantumPageRank:
    """Quantum PageRank using Szegedy quantum walk."""
    
    def __init__(self, graph: nx.Graph, alpha: float = 0.85):
        """
        Initialize Quantum PageRank.
        
        Parameters:
        -----------
        graph : nx.Graph
            Network graph
        alpha : float
            Damping factor
        """
        self.graph = graph
        self.nodes = sorted(graph.nodes())
        self.N = len(self.nodes)
        self.node_to_idx = {node: i for i, node in enumerate(self.nodes)}
        self.alpha = alpha
        
        # Build Google matrix
        adj = nx.to_numpy_array(graph).astype(float)
        degrees = np.sum(adj, axis=1)
        degrees[degrees == 0] = 1.0
        
        P = adj / degrees[:, np.newaxis]
        ones = np.ones((self.N, self.N))
        self.G = alpha * P + (1 - alpha) / self.N * ones
        self.sqrtG = np.sqrt(self.G)
        
        # Build Gram matrices
        self.Gram = np.zeros((self.N, self.N))
        for m in range(self.N):
            for j in range(self.N):
                self.Gram[m, j] = np.sum(self.sqrtG[:, m] * self.sqrtG[:, j])
                
        self.H = np.sqrt(self.G * self.G.T)
        
        # Build reduced unitary
        I = np.eye(self.N)
        RefA = 2 * self.Gram - I
        self.U_red = np.block([
            [np.zeros((self.N, self.N)), -I],
            [RefA, 2 * self.H]
        ])
        
        self.initial_vec = np.hstack([np.full(self.N, 1 / np.sqrt(self.N)), 
                                    np.zeros(self.N)])
    
    def get_ranking(self) -> np.ndarray:
        """Compute Quantum PageRank scores."""
        evals, evecs = eig(self.U_red)
        c = np.dot(evecs.T.conj(), self.initial_vec)
        ranking = np.zeros(self.N)
        
        for i in range(self.N):
            for k in range(2 * self.N):
                uA = np.abs(evecs[:self.N, k]) ** 2
                uB = np.abs(evecs[self.N:, k]) ** 2
                contrib = np.dot(uA, self.G[i, :]) + uB[i]
                ranking[i] += np.abs(c[k]) ** 2 * contrib
                
        return ranking / np.sum(ranking)
    
    def get_sorted_ranking(self) -> List[Tuple]:
        """Get sorted ranking as (node, score) pairs."""
        probs = self.get_ranking()
        idx = np.argsort(probs)[::-1]
        return [(self.nodes[j], probs[j]) for j in idx]
    
    def plot(self, probs: Optional[np.ndarray] = None,
             plot_type: str = 'bar', layout: str = 'spring', **kwargs):
        """Plot ranking results."""
        if probs is None:
            probs = self.get_ranking()
            
        if plot_type == 'bar':
            idx = np.argsort(probs)[::-1]
            sorted_probs = probs[idx]
            sorted_nodes = [str(self.nodes[j]) for j in idx]
            
            plt.figure(figsize=(12, 6))
            plt.bar(sorted_nodes, sorted_probs, **kwargs)
            plt.xlabel('Node')
            plt.ylabel('Quantum PageRank Score')
            plt.title('Quantum PageRank Distribution')
            plt.xticks(rotation=45)
            plt.tight_layout()
            
        elif plot_type == 'graph':
            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='Quantum PageRank Score')
            
            plt.title('Quantum PageRank on Network')
            plt.tight_layout()
            
        else:
            raise ValueError("plot_type must be 'bar' or 'graph'")
            
        return plt.gcf()