import numpy as np
from scipy.stats import unitary_group
from typing import Optional, Tuple, Union

def hadamard_coin(dim: int = 2) -> np.ndarray:
    """
    Hadamard coin operator.
    
    Parameters:
    -----------
    dim : int
        Dimension of coin space (2 or 4)
        
    Returns:
    --------
    np.ndarray
        Hadamard coin operator matrix
    """
    if dim == 2:
        return (1 / np.sqrt(2)) * np.array([[1, 1], [1, -1]])
    elif dim == 4:
        H2 = hadamard_coin(2)
        return np.kron(H2, H2)
    else:
        raise ValueError("Hadamard coin defined for dim=2 or 4")

def grover_coin(dim: int) -> np.ndarray:
    """
    Grover coin (diffusion operator).
    
    Parameters:
    -----------
    dim : int
        Dimension of coin space
        
    Returns:
    --------
    np.ndarray
        Grover coin operator matrix
    """
    return (2 / dim) * np.ones((dim, dim)) - np.eye(dim)

def fourier_coin(dim: int) -> np.ndarray:
    """
    Fourier (DFT) coin operator.
    
    Parameters:
    -----------
    dim : int
        Dimension of coin space
        
    Returns:
    --------
    np.ndarray
        Fourier coin operator matrix
    """
    k = np.arange(dim)
    return (1 / np.sqrt(dim)) * np.exp(2j * np.pi * np.outer(k, k) / dim)

def su_n_coin(dim: int, params: Optional[Tuple] = None) -> np.ndarray:
    """
    General SU(n) coin operator.
    
    Parameters:
    -----------
    dim : int
        Dimension of coin space
    params : tuple, optional
        Parameters for SU(2): (theta, phi, lambda)
        
    Returns:
    --------
    np.ndarray
        SU(n) coin operator matrix
    """
    if dim == 2:
        if params is None:
            params = (np.pi / 4, 0, 0)
        theta, phi, lambda_ = params
        return np.array([
            [np.cos(theta), -np.exp(1j * lambda_) * np.sin(theta)],
            [np.exp(1j * phi) * np.sin(theta), np.exp(1j * (phi + lambda_)) * np.cos(theta)]
        ])
    else:
        # For higher dimensions, generate random unitary if no parameters provided
        if params is None:
            return unitary_group.rvs(dim)
        else:
            raise ValueError("Custom SU(n) for dim>2 not implemented; provide matrix or use default")