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

class GraphDTQW:
    """Discrete Time Quantum Walk on arbitrary graphs."""
    
    def __init__(self, graph: nx.Graph, coin_type: str = 'grover', 
                 initial_node: Optional[Union[int, str]] = None):
        """
        Initialize quantum walk on graph.
        
        Parameters:
        -----------
        graph : nx.Graph
            NetworkX graph
        coin_type : str
            Type of coin: 'grover' or 'fourier'
        initial_node : int or str, optional
            Starting node
        """
        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)}
        
        # Build adjacency matrix
        adj = nx.to_numpy_array(graph, dtype=int)
        self.degrees = [int(np.sum(adj[i])) for i in range(self.num_nodes)]
        self.node_index = np.cumsum([0] + self.degrees[:-1])
        self.total_dim = sum(self.degrees)
        
        # Build coin maps
        self.coin_maps = []
        for i in range(self.num_nodes):
            mapping = {}
            count = 0
            for j in range(self.num_nodes):
                w = adj[i, j]
                if w > 0:
                    mapping[j] = list(range(count, count + w))
                    count += w
            self.coin_maps.append(mapping)
        
        # Build operators
        self.S = self._build_shift_operator()
        self.C = self._build_coin_operator(coin_type)
        self.U = self.S @ self.C
        self.initial_state = self._build_initial_state(initial_node)
        self.state = self.initial_state.copy()
    
    def _build_shift_operator(self) -> np.ndarray:
        S = np.zeros((self.total_dim, self.total_dim), dtype=complex)
        for i in range(self.num_nodes):
            for j in self.coin_maps[i]:
                for local_i, local_j in zip(self.coin_maps[i][j], 
                                          self.coin_maps[j].get(i, [])):
                    g_i = self.node_index[i] + local_i
                    g_j = self.node_index[j] + local_j
                    S[g_j, g_i] = 1
        return S
    
    def _build_coin_operator(self, coin_type: str) -> np.ndarray:
        total_coin = np.zeros((0, 0), dtype=complex)
        for d in self.degrees:
            if d == 0:
                continue
            if coin_type == 'grover':
                C_i = np.full((d, d), 2 / d, dtype=complex)
                np.fill_diagonal(C_i, 2 / d - 1)
            elif coin_type == 'fourier':
                omega = np.exp(2j * np.pi / d)
                C_i = np.zeros((d, d), dtype=complex)
                for row in range(d):
                    for col in range(d):
                        C_i[row, col] = omega ** (row * col)
                C_i /= np.sqrt(d)
            else:
                raise ValueError(f"Unknown coin_type: {coin_type}")
            
            # Direct sum
            if total_coin.shape[0] == 0:
                total_coin = C_i
            else:
                rows, cols = total_coin.shape
                new_rows, new_cols = C_i.shape
                top = np.hstack((total_coin, np.zeros((rows, new_cols), dtype=complex)))
                bottom = np.hstack((np.zeros((new_rows, cols), dtype=complex), C_i))
                total_coin = np.vstack((top, bottom))
        return total_coin
    
    def _build_initial_state(self, initial_node: Optional[Union[int, str]]) -> np.ndarray:
        state = np.zeros(self.total_dim, dtype=complex)
        
        if initial_node is not None:
            i = (self.node_to_idx[initial_node] if initial_node in self.node_to_idx 
                 else initial_node)
            d = self.degrees[i]
            if d > 0:
                amp = 1 / np.sqrt(d)
                start = self.node_index[i]
                state[start:start + d] = amp
        else:
            idx = 0
            norm = np.sqrt(sum(d for d in self.degrees if d > 0))
            for d in self.degrees:
                if d > 0:
                    amp = 1 / np.sqrt(d) / norm
                    state[idx:idx + d] = amp
                idx += d
        return state
    
    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) -> np.ndarray:
        """Get probability distribution over nodes."""
        probs = np.zeros(self.num_nodes)
        abs_state_sq = np.abs(self.state) ** 2
        for i in range(self.num_nodes):
            start = self.node_index[i]
            end = start + self.degrees[i]
            probs[i] = np.sum(abs_state_sq[start:end])
        return probs
    
    def average_probabilities_finite(self, num_steps: int) -> 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()
        return avg_probs / num_steps
    
    def average_probabilities_infinite(self) -> 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.total_dim):
            u_k = eigenvectors[:, k]
            prob_k = np.zeros(self.num_nodes)
            abs_u_k_sq = np.abs(u_k) ** 2
            
            for i in range(self.num_nodes):
                start = self.node_index[i]
                end = start + self.degrees[i]
                prob_k[i] = np.sum(abs_u_k_sq[start:end])
                
            avg_probs += np.abs(c[k]) ** 2 * prob_k
            
        return avg_probs
    
    def plot(self, probs: Optional[np.ndarray] = None, layout: str = 'spring', **kwargs):
        """Plot graph with node probabilities."""
        if probs is None:
            probs = self.get_probabilities()
            
        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('Graph Quantum Walk Probability Distribution')
        plt.tight_layout()
        return plt.gcf()