from typing import Optional, Tuple, Union
import numpy as np
from ..core import Tensor
from ..core.backend import backend

class Layer:
    """Base class for all neural network layers"""
    
    def __init__(self):
        self.trainable: bool = True
        self.training: bool = True
        
    def __call__(self, inputs: Tensor) -> Tensor:
        return self.forward(inputs)
        
    def forward(self, inputs: Tensor) -> Tensor:
        raise NotImplementedError
        
    def backward(self, grad: Tensor) -> Tensor:
        raise NotImplementedError
        
    def train(self) -> None:
        self.training = True
        
    def eval(self) -> None:
        self.training = False
        
    def state_dict(self) -> dict:
        """Get layer state"""
        state = {}
        if hasattr(self, 'weight'):
            state['weight'] = self.weight
        if hasattr(self, 'bias'):
            state['bias'] = self.bias
        return state
        
    def load_state_dict(self, state_dict: dict) -> None:
        """Load layer state"""
        if 'weight' in state_dict:
            self.weight = state_dict['weight']
        if 'bias' in state_dict:
            self.bias = state_dict['bias']

class Linear(Layer):
    """Fully connected layer"""
    
    def __init__(self, in_features: int, out_features: int, bias: bool = True):
        super().__init__()
        self.in_features = in_features
        self.out_features = out_features
        
        # Initialize weights using Kaiming initialization
        scale = np.sqrt(2.0 / in_features)
        self.weight = Tensor(
            np.random.normal(0, scale, (in_features, out_features)),
            requires_grad=True
        )
        
        self.bias = Tensor(
            np.zeros(out_features),
            requires_grad=True
        ) if bias else None
            
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass using current backend"""
        output = backend.matmul(inputs, self.weight)
        if self.bias is not None:
            output = backend.add(output, self.bias)
        return output

class Conv2d(Layer):
    """2D Convolution layer"""
    
    def __init__(self, in_channels: int, out_channels: int, 
                 kernel_size: Union[int, Tuple[int, int]], 
                 stride: Union[int, Tuple[int, int]] = 1,
                 padding: Union[int, Tuple[int, int]] = 0,
                 bias: bool = True):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride = stride if isinstance(stride, tuple) else (stride, stride)
        self.padding = padding if isinstance(padding, tuple) else (padding, padding)
        
        # Initialize weights
        scale = np.sqrt(2.0 / (in_channels * self.kernel_size[0] * self.kernel_size[1]))
        self.weight = Tensor(
            np.random.normal(0, scale, 
                           (out_channels, in_channels, *self.kernel_size)),
            requires_grad=True
        )
        
        self.bias = Tensor(
            np.zeros(out_channels),
            requires_grad=True
        ) if bias else None
            
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass using optimal backend for convolution"""
        return backend.conv2d(
            inputs,
            self.weight,
            self.bias,
            self.stride,
            self.padding
        )

class MaxPool2d(Layer):
    """2D max pooling layer"""
    
    def __init__(self, kernel_size: Union[int, Tuple[int, int]],
                 stride: Optional[Union[int, Tuple[int, int]]] = None,
                 padding: Union[int, Tuple[int, int]] = 0):
        super().__init__()
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride = stride if stride is not None else self.kernel_size
        self.stride = self.stride if isinstance(self.stride, tuple) else (self.stride, self.stride)
        self.padding = padding if isinstance(padding, tuple) else (padding, padding)
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass using optimal backend for pooling"""
        return backend.max_pool2d(
            inputs,
            self.kernel_size,
            self.stride,
            self.padding
        )
class BatchNorm2d(Layer):
    """2D Batch Normalization layer"""
    
    def __init__(self, num_features: int, eps: float = 1e-5, momentum: float = 0.1):
        super().__init__()
        self.num_features = num_features
        self.eps = eps
        self.momentum = momentum
        
        # Learnable parameters
        self.weight = Tensor(np.ones(num_features), requires_grad=True)  # gamma
        self.bias = Tensor(np.zeros(num_features), requires_grad=True)   # beta
        
        # Running statistics
        self.running_mean = Tensor(np.zeros(num_features), requires_grad=False)
        self.running_var = Tensor(np.ones(num_features), requires_grad=False)
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass with batch normalization"""
        if self.training:
            # Calculate batch statistics
            batch_mean = inputs.mean(axis=(0, 2, 3))
            batch_var = inputs.var(axis=(0, 2, 3))
            
            # Update running statistics
            self.running_mean = (1 - self.momentum) * self.running_mean + \
                               self.momentum * batch_mean
            self.running_var = (1 - self.momentum) * self.running_var + \
                              self.momentum * batch_var
        else:
            batch_mean = self.running_mean
            batch_var = self.running_var
        
        # Normalize
        x_norm = (inputs - batch_mean.reshape(1, -1, 1, 1)) / \
                np.sqrt(batch_var.reshape(1, -1, 1, 1) + self.eps)
        
        # Scale and shift
        return self.weight.reshape(1, -1, 1, 1) * x_norm + \
               self.bias.reshape(1, -1, 1, 1)

class AdaptiveAvgPool2d(Layer):
    """2D Adaptive Average Pooling layer"""
    
    def __init__(self, output_size: Union[int, Tuple[int, int]]):
        super().__init__()
        self.output_size = output_size if isinstance(output_size, tuple) else (output_size, output_size)
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass using adaptive average pooling"""
        return backend.adaptive_avg_pool2d(inputs, self.output_size)

class Embedding(Layer):
    """Embedding layer that maps indices to dense vectors"""
    
    def __init__(self, num_embeddings: int, embedding_dim: int):
        super().__init__()
        self.num_embeddings = num_embeddings
        self.embedding_dim = embedding_dim
        
        # Initialize weights with Xavier/Glorot initialization
        scale = np.sqrt(6.0 / (num_embeddings + embedding_dim))
        self.weight = Tensor(
            np.random.uniform(-scale, scale, (num_embeddings, embedding_dim)),
            requires_grad=True
        )
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass using embedding lookup"""
        return backend.embedding(inputs, self.weight)

class LayerNorm(Layer):
    """Layer Normalization"""
    
    def __init__(self, normalized_shape: Union[int, Tuple[int, ...]], eps: float = 1e-5):
        super().__init__()
        self.normalized_shape = normalized_shape if isinstance(normalized_shape, tuple) else (normalized_shape,)
        self.eps = eps
        
        # Learnable parameters
        self.weight = Tensor(np.ones(normalized_shape), requires_grad=True)  # gamma
        self.bias = Tensor(np.zeros(normalized_shape), requires_grad=True)   # beta
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass with layer normalization"""
        # Calculate mean and variance along the normalization dimensions
        axes = tuple(range(-len(self.normalized_shape), 0))
        mean = inputs.mean(axis=axes, keepdims=True)
        var = inputs.var(axis=axes, keepdims=True)
        
        # Normalize
        x_norm = (inputs - mean) / np.sqrt(var + self.eps)
        
        # Scale and shift
        shape = [1] * (inputs.dim() - len(self.normalized_shape)) + list(self.normalized_shape)
        return self.weight.reshape(shape) * x_norm + self.bias.reshape(shape)

class Dropout(Layer):
    """Dropout layer"""
    
    def __init__(self, p: float = 0.5):
        super().__init__()
        if not 0 <= p < 1:
            raise ValueError("Dropout probability must be in range [0, 1)")
        self.p = p
        self.mask: Optional[Tensor] = None
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass with dropout during training"""
        if not self.training or self.p == 0:
            return inputs
            
        # Generate dropout mask
        self.mask = Tensor(
            np.random.binomial(1, 1-self.p, inputs.shape).astype(np.float32)
        )
        
        # Apply mask and scale using backend
        scaled_inputs = backend.div(inputs * self.mask, (1 - self.p))
        return scaled_inputs
        
    def backward(self, grad: Tensor) -> Tensor:
        """Backward pass applies the same mask"""
        if self.mask is not None:
            return grad * self.mask / (1 - self.p)
        return grad

class Sequential(Layer):
    """Sequential container for layers"""
    
    def __init__(self, *layers: Layer):
        super().__init__()
        self.layers = list(layers)
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass through all layers in sequence"""
        for layer in self.layers:
            inputs = layer(inputs)
        return inputs
        
    def train(self) -> None:
        """Set all layers to training mode"""
        super().train()
        for layer in self.layers:
            layer.train()
            
    def eval(self) -> None:
        """Set all layers to evaluation mode"""
        super().eval()
        for layer in self.layers:
            layer.eval()
            
    def state_dict(self) -> dict:
        """Get state of all layers"""
        return {f'layer{i}': layer.state_dict()
                for i, layer in enumerate(self.layers)}
                
    def load_state_dict(self, state_dict: dict) -> None:
        """Load state for all layers"""
        for i, layer in enumerate(self.layers):
            key = f'layer{i}'
            if key in state_dict:
                layer.load_state_dict(state_dict[key])

    def __getitem__(self, idx: int) -> Layer:
        """Get layer by index"""
        return self.layers[idx]
        
    def __len__(self) -> int:
        """Get number of layers"""
        return len(self.layers)

# Factory functions
def get_activation(name: str) -> Layer:
    """Get activation layer by name"""
    from .activations import ReLU, Sigmoid, Tanh
    
    activations = {
        'relu': ReLU,
        'sigmoid': Sigmoid,
        'tanh': Tanh
    }
    
    name = name.lower()
    if name not in activations:
        raise ValueError(f"Unknown activation function: {name}")
        
    return activations[name]()

class ConvTranspose2d(Layer):
    """2D Transposed Convolution layer"""
    
    def __init__(self, in_channels: int, out_channels: int,
                 kernel_size: Union[int, Tuple[int, int]],
                 stride: Union[int, Tuple[int, int]] = 1,
                 padding: Union[int, Tuple[int, int]] = 0,
                 bias: bool = True):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size)
        self.stride = stride if isinstance(stride, tuple) else (stride, stride)
        self.padding = padding if isinstance(padding, tuple) else (padding, padding)
        
        # Initialize weights
        scale = np.sqrt(2.0 / (out_channels * self.kernel_size[0] * self.kernel_size[1]))
        self.weight = Tensor(
            np.random.normal(0, scale, (in_channels, out_channels, *self.kernel_size)),
            requires_grad=True
        )
        self.bias = Tensor(
            np.zeros(out_channels),
            requires_grad=True
        ) if bias else None
        
    def forward(self, inputs: Tensor) -> Tensor:
        """Forward pass using backend's transposed convolution"""
        return backend.conv_transpose2d(inputs, self.weight, self.bias, self.stride, self.padding)

class Reshape(Layer):
    """Reshape layer"""
    
    def __init__(self, *shape):
        super().__init__()
        self.shape = shape
        
    def forward(self, inputs: Tensor) -> Tensor:
        if self.shape[0] == -1:
            batch_size = inputs.shape[0]
            shape = (batch_size,) + self.shape[1:]
        else:
            shape = self.shape
        return inputs.reshape(shape)

class Flatten(Layer):
    """Flatten layer"""
    
    def forward(self, inputs: Tensor) -> Tensor:
        batch_size = inputs.shape[0]
        return inputs.reshape(batch_size, -1)