import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from typing import Tuple, Optional, Union
from ..coins import hadamard_coin, grover_coin, fourier_coin, su_n_coin

class DTQW1D:
    """Discrete Time Quantum Walk in 1D."""
    
    def __init__(self, N: int, coin_type: str = 'hadamard', coin_params: Optional[Tuple] = None):
        """
        Initialize 1D DTQW.
        
        Parameters:
        -----------
        N : int
            Half-size of line (positions -N to N)
        coin_type : str
            Type of coin operator: 'hadamard', 'grover', 'fourier', 'su2'
        coin_params : tuple, optional
            Parameters for SU(2) coin: (theta, phi, lambda)
        """
        self.N = N
        self.positions = np.arange(-N, N + 1)
        self.num_pos = len(self.positions)
        self.coin_dim = 2
        self.total_dim = self.num_pos * self.coin_dim
        self.state = np.zeros(self.total_dim, dtype=complex)
        self.coin_type = coin_type
        self.coin = self._get_coin(coin_type, coin_params)
        self._build_operators()
    
    def _get_coin(self, coin_type: str, params: Optional[Tuple]) -> np.ndarray:
        coin_operators = {
            'hadamard': lambda: hadamard_coin(2),
            'grover': lambda: grover_coin(2),
            'fourier': lambda: fourier_coin(2),
            'su2': lambda: su_n_coin(2, params)
        }
        
        if coin_type not in coin_operators:
            raise ValueError(f"Unknown coin type: {coin_type}")
            
        return coin_operators[coin_type]()
    
    def _build_operators(self):
        # Coin operator
        self.coin_op = np.kron(np.eye(self.num_pos), self.coin)
        
        # Shift operator
        self.shift_op = np.zeros((self.total_dim, self.total_dim), dtype=complex)
        for i in range(self.num_pos):
            # |x, 0> -> |x+1, 0> (right)
            if i + 1 < self.num_pos:
                self.shift_op[self.coin_dim * (i + 1), self.coin_dim * i] = 1
            # |x, 1> -> |x-1, 1> (left)
            if i - 1 >= 0:
                self.shift_op[self.coin_dim * (i - 1) + 1, self.coin_dim * i + 1] = 1
    
    def initialize_state(self, init_pos: int = 0, init_coin: np.ndarray = None):
        """
        Initialize quantum walk state.
        
        Parameters:
        -----------
        init_pos : int
            Initial position
        init_coin : np.ndarray
            Initial coin state (normalized)
        """
        if init_coin is None:
            init_coin = np.array([1, 0]) / np.sqrt(2)
            
        idx = init_pos + self.N
        if 0 <= idx < self.num_pos:
            self.state[self.coin_dim * idx: self.coin_dim * (idx + 1)] = (
                init_coin / np.linalg.norm(init_coin)
            )
        else:
            raise ValueError("Initial position out of bounds")
    
    def step(self):
        """Perform one evolution step."""
        self.state = self.shift_op @ (self.coin_op @ self.state)
    
    def evolve(self, num_steps: int):
        """Evolve for multiple steps."""
        for _ in range(num_steps):
            self.step()
    
    def get_probabilities(self) -> np.ndarray:
        """Get probability distribution over positions."""
        probs = np.zeros(self.num_pos)
        for i in range(self.num_pos):
            probs[i] = np.sum(np.abs(self.state[self.coin_dim * i: self.coin_dim * (i + 1)]) ** 2)
        return probs
    
    def plot(self, **kwargs):
        """Plot probability distribution."""
        probs = self.get_probabilities()
        plt.figure(figsize=(10, 6))
        plt.plot(self.positions, probs, marker='o', linestyle='-', **kwargs)
        plt.xlabel('Position')
        plt.ylabel('Probability')
        plt.title(f'1D Quantum Walk ({self.coin_type} coin)')
        plt.grid(True, alpha=0.3)
        plt.tight_layout()
        return plt.gcf()

class DTQW2D:
    """Discrete Time Quantum Walk in 2D."""
    
    def __init__(self, Nx: int, Ny: int, coin_type: str = 'grover', coin_params: Optional[Tuple] = None):
        self.Nx = Nx
        self.Ny = Ny
        self.positions_x = np.arange(-Nx, Nx + 1)
        self.positions_y = np.arange(-Ny, Ny + 1)
        self.num_x = len(self.positions_x)
        self.num_y = len(self.positions_y)
        self.coin_dim = 4  # up, down, left, right
        self.total_dim = self.num_x * self.num_y * self.coin_dim
        self.state = np.zeros(self.total_dim, dtype=complex)
        self.coin_type = coin_type
        self.coin = self._get_coin(coin_type, coin_params)
        self.coin_op = np.kron(np.eye(self.num_x * self.num_y), self.coin)
    
    def _get_coin(self, coin_type: str, params: Optional[Tuple]) -> np.ndarray:
        coin_operators = {
            'hadamard': lambda: hadamard_coin(4),
            'grover': lambda: grover_coin(4),
            'fourier': lambda: fourier_coin(4)
        }
        
        if coin_type not in coin_operators:
            raise ValueError(f"Unknown coin type: {coin_type}")
            
        return coin_operators[coin_type]()
    
    def initialize_state(self, init_pos: Tuple[int, int] = (0, 0), 
                        init_coin: np.ndarray = None):
        if init_coin is None:
            init_coin = np.ones(4) / 2
            
        idx_x = init_pos[0] + self.Nx
        idx_y = init_pos[1] + self.Ny
        
        if 0 <= idx_x < self.num_x and 0 <= idx_y < self.num_y:
            flat_idx = (idx_x * self.num_y + idx_y) * self.coin_dim
            self.state[flat_idx: flat_idx + self.coin_dim] = (
                init_coin / np.linalg.norm(init_coin)
            )
        else:
            raise ValueError("Initial position out of bounds")
    
    def step(self):
        self.state = self.coin_op @ self.state
        
        state_reshaped = self.state.reshape((self.num_x, self.num_y, self.coin_dim))
        new_state = np.zeros_like(state_reshaped)
        
        # Shift operations
        new_state[:, :-1, 0] = state_reshaped[:, 1:, 0]  # Up
        new_state[:, 1:, 1] = state_reshaped[:, :-1, 1]  # Down  
        new_state[1:, :, 2] = state_reshaped[:-1, :, 2]  # Left
        new_state[:-1, :, 3] = state_reshaped[1:, :, 3]  # Right
        
        self.state = new_state.flatten()
    
    def evolve(self, num_steps: int):
        for _ in range(num_steps):
            self.step()
    
    def get_probabilities(self) -> np.ndarray:
        state_reshaped = self.state.reshape((self.num_x, self.num_y, self.coin_dim))
        return np.sum(np.abs(state_reshaped) ** 2, axis=2)
    
    def plot(self, **kwargs):
        probs = self.get_probabilities()
        extent = [self.positions_x[0], self.positions_x[-1], 
                 self.positions_y[0], self.positions_y[-1]]
        
        plt.figure(figsize=(10, 8))
        plt.imshow(probs.T, origin='lower', cmap='YlGnBu', 
                  extent=extent, aspect='equal', **kwargs)
        plt.colorbar(label='Probability')
        plt.xlabel('X Position')
        plt.ylabel('Y Position')
        plt.title(f'2D Quantum Walk ({self.coin_type} coin)')
        plt.tight_layout()
        return plt.gcf()

class DTQW3D:
    """Discrete Time Quantum Walk in 3D."""
    
    def __init__(self, Nx: int, Ny: int, Nz: int, 
                 coin_type: str = 'grover', coin_params: Optional[Tuple] = None):
        self.Nx = Nx
        self.Ny = Ny
        self.Nz = Nz
        self.positions_x = np.arange(-Nx, Nx + 1)
        self.positions_y = np.arange(-Ny, Ny + 1)
        self.positions_z = np.arange(-Nz, Nz + 1)
        self.num_x = len(self.positions_x)
        self.num_y = len(self.positions_y)
        self.num_z = len(self.positions_z)
        self.coin_dim = 6  # ±x, ±y, ±z
        self.total_dim = self.num_x * self.num_y * self.num_z * self.coin_dim
        self.state = np.zeros(self.total_dim, dtype=complex)
        self.coin_type = coin_type
        self.coin = self._get_coin(coin_type, coin_params)
        self.coin_op = np.kron(np.eye(self.num_x * self.num_y * self.num_z), self.coin)
    
    def _get_coin(self, coin_type: str, params: Optional[Tuple]) -> np.ndarray:
        coin_operators = {
            'grover': lambda: grover_coin(6),
            'fourier': lambda: fourier_coin(6),
            'su_n': lambda: su_n_coin(6, params)
        }
        
        if coin_type not in coin_operators:
            raise ValueError(f"Unknown coin type: {coin_type}")
            
        return coin_operators[coin_type]()
    
    def initialize_state(self, init_pos: Tuple[int, int, int] = (0, 0, 0),
                        init_coin: np.ndarray = None):
        if init_coin is None:
            init_coin = np.ones(6) / np.sqrt(6)
            
        idx_x = init_pos[0] + self.Nx
        idx_y = init_pos[1] + self.Ny
        idx_z = init_pos[2] + self.Nz
        
        if (0 <= idx_x < self.num_x and 0 <= idx_y < self.num_y and 
            0 <= idx_z < self.num_z):
            flat_idx = ((idx_x * self.num_y * self.num_z + 
                        idx_y * self.num_z + idx_z) * self.coin_dim)
            self.state[flat_idx: flat_idx + self.coin_dim] = (
                init_coin / np.linalg.norm(init_coin)
            )
        else:
            raise ValueError("Initial position out of bounds")
    
    def step(self):
        self.state = self.coin_op @ self.state
        
        state_reshaped = self.state.reshape((self.num_x, self.num_y, 
                                           self.num_z, self.coin_dim))
        new_state = np.zeros_like(state_reshaped)
        
        # Shift operations
        new_state[:-1, :, :, 0] = state_reshaped[1:, :, :, 0]  # +x
        new_state[1:, :, :, 1] = state_reshaped[:-1, :, :, 1]  # -x
        new_state[:, :-1, :, 2] = state_reshaped[:, 1:, :, 2]  # +y
        new_state[:, 1:, :, 3] = state_reshaped[:, :-1, :, 3]  # -y
        new_state[:, :, :-1, 4] = state_reshaped[:, :, 1:, 4]  # +z
        new_state[:, :, 1:, 5] = state_reshaped[:, :, :-1, 5]  # -z
        
        self.state = new_state.flatten()
    
    def evolve(self, num_steps: int):
        for _ in range(num_steps):
            self.step()
    
    def get_probabilities(self) -> np.ndarray:
        state_reshaped = self.state.reshape((self.num_x, self.num_y, 
                                           self.num_z, self.coin_dim))
        return np.sum(np.abs(state_reshaped) ** 2, axis=3)
    
    def plot(self, threshold: float = 1e-5, **kwargs):
        probs = self.get_probabilities()
        x, y, z = np.meshgrid(self.positions_x, self.positions_y, 
                             self.positions_z, indexing='ij')
        mask = probs > threshold
        
        fig = plt.figure(figsize=(12, 9))
        ax = fig.add_subplot(111, projection='3d')
        scatter = ax.scatter(x[mask], y[mask], z[mask], 
                           c=probs[mask], s=probs[mask] * 1000, 
                           cmap='viridis', **kwargs)
        plt.colorbar(scatter, label='Probability')
        ax.set_xlabel('X Position')
        ax.set_ylabel('Y Position')
        ax.set_zlabel('Z Position')
        ax.set_title(f'3D Quantum Walk ({self.coin_type} coin)')
        plt.tight_layout()
        return plt.gcf()