import random
import math
from typing import List, Union, Optional, Any
from .grnexus_activations import *


class Layer:
    """Base layer class"""
    def __init__(self):
        self.weights = None
        self.biases = None
        self.trainable = True
        self.cache = {}
    
    def forward(self, input_data):
        raise NotImplementedError("Debes implementar el método forward")
    
    def backward(self, gradient):
        raise NotImplementedError("Debes implementar el método backward")
    
    def trainable(self):
        return self.trainable
    
    def parameters(self):
        return [p for p in [self.weights, self.biases] if p is not None]
    
    def zero_gradients(self):
        if hasattr(self, 'weight_gradient'):
            self.weight_gradient = None
        if hasattr(self, 'bias_gradient'):
            self.bias_gradient = None
    
    def update_parameters(self, learning_rate):
        if hasattr(self, 'weight_gradient') and self.weight_gradient is not None:
            self._update_matrix(self.weights, self.weight_gradient, learning_rate)
        if hasattr(self, 'bias_gradient') and self.bias_gradient is not None:
            self._update_vector(self.biases, self.bias_gradient, learning_rate)
    
    def _update_matrix(self, matrix, gradient, lr):
        for i in range(len(matrix)):
            for j in range(len(matrix[i])):
                matrix[i][j] -= lr * gradient[i][j]
    
    def _update_vector(self, vector, gradient, lr):
        for i in range(len(vector)):
            vector[i] -= lr * gradient[i]


class DenseLayer(Layer):
    """Fully connected dense layer"""
    def __init__(self, units, input_dim, activation=None, use_bias=True, weight_init='xavier'):
        super().__init__()
        self.units = units
        self.input_dim = input_dim
        self.use_bias = use_bias
        
        # Handle activation - can be string or object
        if activation is None:
            self.activation = None
        elif isinstance(activation, str):
            # Try to get activation from grnexus_activations module
            try:
                self.activation = globals()[activation]()
            except KeyError:
                raise ValueError(f"Unknown activation function: {activation}")
        else:
            self.activation = activation
        
        # Initialize weights
        self.weights = self._initialize_weights(weight_init, input_dim, units)
        self.biases = [0.0] * units if use_bias else None
    
    def _initialize_weights(self, method, input_dim, units):
        if method == 'xavier':
            std = math.sqrt(2.0 / (input_dim + units))
            return [[self._rand_normal(0, std) for _ in range(units)] for _ in range(input_dim)]
        elif method == 'he':
            std = math.sqrt(2.0 / input_dim)
            return [[self._rand_normal(0, std) for _ in range(units)] for _ in range(input_dim)]
        elif method == 'random':
            return [[random.random() * 2 - 1 for _ in range(units)] for _ in range(input_dim)]
        else:
            return [[0.0] * units for _ in range(input_dim)]
    
    def _rand_normal(self, mean, std_dev):
        u1 = random.random()
        u2 = random.random()
        z0 = math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
        return z0 * std_dev + mean
    
    def forward(self, input_data):
        # Check if batch input
        batch_case = (isinstance(input_data[0], list) and 
                     len(input_data[0]) > 0 and 
                     isinstance(input_data[0][0], (int, float)))
        
        if batch_case:
            self.cache['input'] = [row[:] for row in input_data]
            output_batch = []
            
            for x in input_data:
                # weights is (input_dim x units), x is (input_dim)
                # result should be (units)
                z = [0.0] * self.units
                for i in range(self.units):
                    sum_val = 0.0
                    for j in range(self.input_dim):
                        sum_val += self.weights[j][i] * x[j]
                    z[i] = sum_val
                
                if self.biases:
                    z = [z[i] + self.biases[i] for i in range(len(z))]
                
                if self.activation:
                    if 'pre_activation' not in self.cache:
                        self.cache['pre_activation'] = []
                    self.cache['pre_activation'].append(z[:])
                    z = self.activation(z)
                
                output_batch.append(z)
            
            return output_batch
        else:
            self.cache['input'] = input_data[:]
            # weights is (input_dim x units), input is (input_dim)
            # result should be (units)
            z = [0.0] * self.units
            for i in range(self.units):
                sum_val = 0.0
                for j in range(self.input_dim):
                    sum_val += self.weights[j][i] * input_data[j]
                z[i] = sum_val
            
            if self.biases:
                z = [z[i] + self.biases[i] for i in range(len(z))]
            
            if self.activation:
                self.cache['pre_activation'] = z[:]
                z = self.activation(z)
            
            return z
    
    def _multiply_matrix_vector(self, matrix, vector):
        result = [0.0] * len(matrix)
        for i in range(len(matrix)):
            result[i] = self._dot_product(matrix[i], vector)
        return result
    
    def _dot_product(self, a, b):
        sum_val = 0.0
        for i in range(len(a)):
            sum_val += a[i] * b[i]
        return sum_val
    
    def _add_vectors(self, a, b):
        return [a[i] + b[i] for i in range(len(a))]
    
    def _multiply_elementwise(self, a, b):
        return [a[i] * b[i] for i in range(len(a))]
    
    def _multiply_outer_product(self, vec1, vec2):
        result = []
        for i in range(len(vec1)):
            row = []
            for j in range(len(vec2)):
                row.append(vec1[i] * vec2[j])
            result.append(row)
        return result
    
    def _average_matrices(self, matrices):
        rows, cols = len(matrices[0]), len(matrices[0][0])
        avg = [[0.0 for _ in range(cols)] for _ in range(rows)]
        
        for matrix in matrices:
            for i in range(rows):
                for j in range(cols):
                    avg[i][j] += matrix[i][j]
        
        for i in range(rows):
            for j in range(cols):
                avg[i][j] /= len(matrices)
        
        return avg
    
    def _average_arrays(self, arrays):
        length = len(arrays[0])
        avg = [0.0] * length
        
        for arr in arrays:
            for i in range(length):
                avg[i] += arr[i]
        
        for i in range(length):
            avg[i] /= len(arrays)
        
        return avg
    
    def backward(self, gradient):
        # Check if batch input
        batch_case = (isinstance(gradient[0], list) and 
                     isinstance(gradient[0][0], (int, float)))
        
        gradients = gradient if batch_case else [gradient]
        inputs = self.cache['input'] if batch_case else [self.cache['input']]
        
        input_gradients = []
        weight_gradients = []
        bias_gradients = [] if self.biases else None
        
        for idx, grad in enumerate(gradients):
            current_input = inputs[idx]
            
            if self.activation:
                pre_act = (self.cache['pre_activation'][idx] if batch_case 
                          else self.cache['pre_activation'])
                act_deriv = self.activation(pre_act, derivative=True)
                grad = self._multiply_elementwise(grad, act_deriv)
            
            weight_grad = self._multiply_outer_product(current_input, grad)
            weight_gradients.append(weight_grad)
            
            if self.biases:
                bias_gradients.append(grad[:])
            
            input_grad = self._multiply_matrix_vector(self.weights, grad)
            input_gradients.append(input_grad)
        
        if batch_case:
            self.weight_gradient = self._average_matrices(weight_gradients)
            self.bias_gradient = self._average_arrays(bias_gradients) if self.biases else None
            return input_gradients
        else:
            self.weight_gradient = weight_gradients[0]
            self.bias_gradient = bias_gradients[0] if self.biases else None
            return input_gradients[0]


class ActivationLayer(Layer):
    """Standalone activation layer"""
    def __init__(self, activation):
        super().__init__()
        if isinstance(activation, str):
            try:
                self.activation = globals()[activation]()
            except KeyError:
                raise ValueError(f"Unknown activation function: {activation}")
        else:
            self.activation = activation
        self.trainable = False
    
    def forward(self, input_data):
        self.cache['input'] = input_data[:] if isinstance(input_data, list) else input_data
        if isinstance(input_data[0], list):
            return [self.activation(x) for x in input_data]
        else:
            return self.activation(input_data)
    
    def _multiply_elementwise(self, a, b):
        return [a[i] * b[i] for i in range(len(a))]
    
    def backward(self, gradient):
        input_cache = self.cache['input']
        
        if isinstance(gradient[0], list):
            result = []
            for grad, cached_input in zip(gradient, input_cache):
                activation_deriv = self.activation(cached_input, derivative=True)
                result.append(self._multiply_elementwise(grad, activation_deriv))
            return result
        else:
            activation_deriv = self.activation(input_cache, derivative=True)
            return self._multiply_elementwise(gradient, activation_deriv)


class DropoutLayer(Layer):
    """Dropout regularization layer"""
    def __init__(self, rate=0.5):
        super().__init__()
        self.rate = rate
        self.mask = None
        self.trainable = False
    
    def forward(self, input_data, training=True):
        self.cache['training'] = training
        if training:
            if isinstance(input_data[0], list):
                self.mask = [[1.0 / (1.0 - self.rate) if random.random() > self.rate else 0.0 
                             for _ in sample] for sample in input_data]
                return [[input_data[i][j] * self.mask[i][j] for j in range(len(input_data[i]))] 
                       for i in range(len(input_data))]
            else:
                self.mask = [1.0 / (1.0 - self.rate) if random.random() > self.rate else 0.0 
                            for _ in input_data]
                return [input_data[i] * self.mask[i] for i in range(len(input_data))]
        else:
            return input_data
    
    def _multiply_elementwise(self, a, b):
        return [a[i] * b[i] for i in range(len(a))]
    
    def _multiply_batch_elementwise(self, batch_a, batch_b):
        return [self._multiply_elementwise(a, b) for a, b in zip(batch_a, batch_b)]
    
    def backward(self, gradient):
        training = self.cache['training']
        if training and self.mask:
            if isinstance(gradient[0], list):
                return self._multiply_batch_elementwise(gradient, self.mask)
            else:
                return self._multiply_elementwise(gradient, self.mask)
        return gradient


class BatchNormLayer(Layer):
    """Batch normalization layer"""
    def __init__(self, epsilon=1e-5, momentum=0.1):
        super().__init__()
        self.epsilon = epsilon
        self.momentum = momentum
        self.running_mean = None
        self.running_var = None
        self.gamma = 1.0
        self.beta = 0.0
        self.trainable = True
    
    def _calculate_mean(self, batch):
        features = len(batch[0])
        means = [0.0] * features
        
        for sample in batch:
            for i, val in enumerate(sample):
                means[i] += val
        
        for i in range(features):
            means[i] /= len(batch)
        
        return means
    
    def _calculate_variance(self, batch, means):
        features = len(batch[0])
        vars = [0.0] * features
        
        for sample in batch:
            for i, val in enumerate(sample):
                vars[i] += (val - means[i])**2
        
        for i in range(features):
            vars[i] /= len(batch)
        
        return vars
    
    def _update_running_stats(self, batch_mean, batch_var):
        if self.running_mean is None:
            self.running_mean = batch_mean[:]
            self.running_var = batch_var[:]
        else:
            self.running_mean = self._multiply_scalar_add_vector(
                (1 - self.momentum), self.running_mean, 
                self.momentum, batch_mean)
            self.running_var = self._multiply_scalar_add_vector(
                (1 - self.momentum), self.running_var, 
                self.momentum, batch_var)
    
    def _normalize_batch(self, batch, means, vars, inv_std=None):
        if inv_std is None:
            inv_std = [1.0 / math.sqrt(v + self.epsilon) for v in vars]
        
        result = []
        for sample in batch:
            normalized_sample = []
            for val, mean, std_inv in zip(sample, means, inv_std):
                normalized_sample.append((val - mean) * std_inv)
            result.append(normalized_sample)
        
        return result
    
    def _center_batch(self, batch, means):
        result = []
        for sample in batch:
            centered_sample = []
            for val, mean in zip(sample, means):
                centered_sample.append(val - mean)
            result.append(centered_sample)
        return result
    
    def _apply_affine_transform(self, normalized):
        if isinstance(normalized[0], list):
            result = []
            for sample in normalized:
                transformed_sample = []
                for val in sample:
                    transformed_sample.append(self.gamma * val + self.beta)
                result.append(transformed_sample)
            return result
        else:
            return [self.gamma * val + self.beta for val in normalized]
    
    def _calculate_gamma_gradient(self, gradient_batch, x_norm_batch):
        grad_sum = [0.0] * len(x_norm_batch[0])
        
        for grad_sample, norm_sample in zip(gradient_batch, x_norm_batch):
            for i, (grad_val, norm_val) in enumerate(zip(grad_sample, norm_sample)):
                grad_sum[i] += grad_val * norm_val
        
        return grad_sum
    
    def _calculate_beta_gradient(self, gradient_batch):
        grad_sum = [0.0] * len(gradient_batch[0])
        
        for grad_sample in gradient_batch:
            for i, grad_val in enumerate(grad_sample):
                grad_sum[i] += grad_val
        
        return grad_sum
    
    def _multiply_batch_elementwise(self, batch_a, scalar_or_vector):
        if isinstance(scalar_or_vector, list):
            result = []
            for sample in batch_a:
                multiplied_sample = []
                for val, mult in zip(sample, scalar_or_vector):
                    multiplied_sample.append(val * mult)
                result.append(multiplied_sample)
            return result
        else:
            result = []
            for sample in batch_a:
                multiplied_sample = [val * scalar_or_vector for val in sample]
                result.append(multiplied_sample)
            return result
    
    def _subtract_vectors(self, a, b):
        return [a[i] - b[i] for i in range(len(a))]
    
    def _multiply_scalar_add_vector(self, scalar1, vec1, scalar2, vec2):
        return [scalar1 * vec1[i] + scalar2 * vec2[i] for i in range(len(vec1))]
    
    def _multiply_elementwise(self, a, b):
        return [a[i] * b[i] for i in range(len(a))]
    
    def forward(self, input_data, training=True):
        batch_case = (isinstance(input_data[0], list) and 
                     isinstance(input_data[0][0], (int, float)))
        input_tensor = input_data if batch_case else [input_data]
        
        if training:
            batch_mean = self._calculate_mean(input_tensor)
            batch_var = self._calculate_variance(input_tensor, batch_mean)
            self._update_running_stats(batch_mean, batch_var)
            
            self.cache['mean'] = batch_mean
            self.cache['var'] = batch_var
            self.cache['inv_std'] = [1.0 / math.sqrt(v + self.epsilon) for v in batch_var]
            self.cache['x_norm'] = self._normalize_batch(input_tensor, batch_mean, batch_var)
            self.cache['x_centered'] = self._center_batch(input_tensor, batch_mean)
        else:
            # Si no hay running stats, usar las del batch actual
            if self.running_mean is None or self.running_var is None:
                batch_mean = self._calculate_mean(input_tensor)
                batch_var = self._calculate_variance(input_tensor, batch_mean)
                self._update_running_stats(batch_mean, batch_var)
                self.cache['x_norm'] = self._normalize_batch(input_tensor, batch_mean, batch_var)
            else:
                running_inv_std = [1.0 / math.sqrt(v + self.epsilon) for v in self.running_var]
                self.cache['x_norm'] = self._normalize_batch(input_tensor, self.running_mean, self.running_var, running_inv_std)
        
        output_tensor = self._apply_affine_transform(
            self.cache['x_norm'] if batch_case else self.cache['x_norm'][0])
        return output_tensor if batch_case else output_tensor[0]
    
    def backward(self, gradient):
        batch_case = (isinstance(gradient[0], list) and 
                     isinstance(gradient[0][0], (int, float)))
        grad_tensor = gradient if batch_case else [gradient]
        
        self.gamma_gradient = self._calculate_gamma_gradient(grad_tensor, self.cache['x_norm'])
        self.beta_gradient = self._calculate_beta_gradient(grad_tensor)
        
        # dx_norm = grad_tensor * gamma (element-wise)
        dx_norm = []
        for grad_sample in grad_tensor:
            dx_norm.append([g * self.gamma for g in grad_sample])
        
        n = len(grad_tensor)
        inv_std = self.cache['inv_std']
        
        dx = []
        for i in range(n):
            sum1 = self._multiply_elementwise(dx_norm[i], inv_std)
            mean_dx_norm = sum(dx_norm[i][j] / n for j in range(len(dx_norm[i])))
            sum2 = [c * s**3 for c, s in zip(self.cache['x_centered'][i], inv_std)]
            sum2 = [val / n for val in sum2]
            sum2 = self._multiply_elementwise(sum2, [mean_dx_norm] * len(sum2))
            
            dx.append(self._subtract_vectors(sum1, sum2))
        
        return dx if batch_case else dx[0]


class Conv2DLayer(Layer):
    """2D Convolutional layer"""
    def __init__(self, filters, kernel_size, stride=1, padding=0):
        super().__init__()
        self.filters = filters
        self.kernel_size = kernel_size if isinstance(kernel_size, list) else [kernel_size, kernel_size]
        self.stride = stride
        self.padding = padding
        
        kh, kw = self.kernel_size
        self.kernels = self._initialize_conv_kernels(filters, kh, kw)
        self.biases = [0.0] * filters
    
    def _initialize_conv_kernels(self, filters, kh, kw):
        fan_in = kh * kw
        std = math.sqrt(2.0 / fan_in)
        return [[[self._rand_normal(0, std) for _ in range(kw)] for _ in range(kh)] 
                for _ in range(filters)]
    
    def _rand_normal(self, mean, std_dev):
        u1 = random.random()
        u2 = random.random()
        z0 = math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
        return z0 * std_dev + mean
    
    def _pad_image(self, image, padding):
        h, w = len(image), len(image[0])
        padded_h, padded_w = h + 2 * padding, w + 2 * padding
        
        padded = [[0.0 for _ in range(padded_w)] for _ in range(padded_h)]
        
        for ih in range(h):
            for iw in range(w):
                padded[ih + padding][iw + padding] = image[ih][iw]
        
        return padded
    
    def _convolve_2d(self, input_image):
        h, w = len(input_image), len(input_image[0])
        kh, kw = len(self.kernels[0]), len(self.kernels[0][0])
        
        out_h = int((h + 2 * self.padding - kh) / self.stride + 1)
        out_w = int((w + 2 * self.padding - kw) / self.stride + 1)
        
        output = [[[0.0 for _ in range(self.filters)] for _ in range(out_w)] for _ in range(out_h)]
        padded = self._pad_image(input_image, self.padding) if self.padding > 0 else input_image
        
        for oh in range(out_h):
            for ow in range(out_w):
                roi_start_h = oh * self.stride
                roi_start_w = ow * self.stride
                
                for f in range(self.filters):
                    sum_val = 0.0
                    for kh_off in range(kh):
                        for kw_off in range(kw):
                            ih = roi_start_h + kh_off
                            iw = roi_start_w + kw_off
                            sum_val += padded[ih][iw] * self.kernels[f][kh_off][kw_off]
                    output[oh][ow][f] = sum_val + self.biases[f]
        
        return output
    
    def forward(self, input_data):
        batch_case = (isinstance(input_data[0], list) and 
                     isinstance(input_data[0][0], list) and 
                     isinstance(input_data[0][0][0], (int, float)))
        input_tensor = input_data if batch_case else [input_data]
        
        output_batch = [self._convolve_2d(single_input) for single_input in input_tensor]
        
        return output_batch if batch_case else output_batch[0]
    
    def backward(self, gradient):
        return gradient


class MaxPoolingLayer(Layer):
    """Max pooling layer"""
    def __init__(self, pool_size, stride=None):
        super().__init__()
        self.pool_size = pool_size if isinstance(pool_size, list) else [pool_size, pool_size]
        self.stride = stride if stride else self.pool_size
        self.stride = self.stride if isinstance(self.stride, list) else [self.stride, self.stride]
        self.trainable = False
    
    def _pool_2d(self, input_image):
        h, w = len(input_image), len(input_image[0])
        ph, pw = self.pool_size
        sh, sw = self.stride
        
        out_h = int((h - ph) / sh + 1)
        out_w = int((w - pw) / sw + 1)
        
        output = [[0.0 for _ in range(out_w)] for _ in range(out_h)]
        switch_indices_map = []
        
        for oh in range(out_h):
            for ow in range(out_w):
                pool_start_h = oh * sh
                pool_start_w = ow * sw
                
                max_val = float('-inf')
                max_h, max_w = 0, 0
                
                for ph_off in range(ph):
                    for pw_off in range(pw):
                        ih = pool_start_h + ph_off
                        iw = pool_start_w + pw_off
                        if input_image[ih][iw] > max_val:
                            max_val = input_image[ih][iw]
                            max_h, max_w = ih, iw
                
                output[oh][ow] = max_val
                switch_indices_map.append([max_h, max_w])
        
        if 'switch_indices' not in self.cache:
            self.cache['switch_indices'] = []
        self.cache['switch_indices'].append(switch_indices_map)
        return output
    
    def forward(self, input_data):
        batch_case = (isinstance(input_data[0], list) and 
                     isinstance(input_data[0][0], (int, float)))
        input_tensor = input_data if batch_case else [input_data]
        
        output_batch = [self._pool_2d(single_input) for single_input in input_tensor]
        
        return output_batch if batch_case else output_batch[0]
    
    def backward(self, gradient):
        return gradient


class LSTMLayer(Layer):
    """Long Short-Term Memory layer"""
    def __init__(self, units, input_size):
        super().__init__()
        self.units = units
        self.input_size = input_size
        self.hidden_size = units
        
        # Initialize weights for gates
        self.wf = self._initialize_weights('xavier', input_size, units)
        self.uf = self._initialize_weights('xavier', self.hidden_size, units)
        self.bf = [0.0] * units
        
        self.wi = self._initialize_weights('xavier', input_size, units)
        self.ui = self._initialize_weights('xavier', self.hidden_size, units)
        self.bi = [0.0] * units
        
        self.wo = self._initialize_weights('xavier', input_size, units)
        self.uo = self._initialize_weights('xavier', self.hidden_size, units)
        self.bo = [0.0] * units
        
        self.wc = self._initialize_weights('xavier', input_size, units)
        self.uc = self._initialize_weights('xavier', self.hidden_size, units)
        self.bc = [0.0] * units
    
    def _initialize_weights(self, method, input_dim, units):
        std = math.sqrt(2.0 / (input_dim + units))
        return [[self._rand_normal(0, std) for _ in range(input_dim)] for _ in range(units)]
    
    def _rand_normal(self, mean, std_dev):
        u1 = random.random()
        u2 = random.random()
        z0 = math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
        return z0 * std_dev + mean
    
    def _add_vectors(self, a, b):
        return [a[i] + b[i] for i in range(len(a))]
    
    def _multiply_elementwise(self, a, b):
        return [a[i] * b[i] for i in range(len(a))]
    
    def _multiply_matrix_vector(self, matrix, vector):
        result = [0.0] * len(matrix)
        for i in range(len(matrix)):
            result[i] = sum(matrix[i][j] * vector[j] for j in range(len(vector)))
        return result
    
    def _lstm_step(self, input_t, prev_hidden, prev_cell):
        # Forget gate
        f_input = self._add_vectors(
            self._multiply_matrix_vector(self.wf, input_t),
            self._multiply_matrix_vector(self.uf, prev_hidden)
        )
        f_input = self._add_vectors(f_input, self.bf)
        f_gate = Sigmoid()(f_input)
        
        # Input gate
        i_input = self._add_vectors(
            self._multiply_matrix_vector(self.wi, input_t),
            self._multiply_matrix_vector(self.ui, prev_hidden)
        )
        i_input = self._add_vectors(i_input, self.bi)
        i_gate = Sigmoid()(i_input)
        
        # Output gate
        o_input = self._add_vectors(
            self._multiply_matrix_vector(self.wo, input_t),
            self._multiply_matrix_vector(self.uo, prev_hidden)
        )
        o_input = self._add_vectors(o_input, self.bo)
        o_gate = Sigmoid()(o_input)
        
        # Candidate cell state
        c_input = self._add_vectors(
            self._multiply_matrix_vector(self.wc, input_t),
            self._multiply_matrix_vector(self.uc, prev_hidden)
        )
        c_input = self._add_vectors(c_input, self.bc)
        candidate = Tanh()(c_input)
        
        # New cell state
        new_cell = self._add_vectors(
            self._multiply_elementwise(f_gate, prev_cell),
            self._multiply_elementwise(i_gate, candidate)
        )
        
        # New hidden state
        tanh_cell = Tanh()(new_cell)
        new_hidden = self._multiply_elementwise(o_gate, tanh_cell)
        
        return new_hidden, new_cell
    
    def _transpose_batch_sequences(self, sequences):
        seq_len = len(sequences)
        batch_size = len(sequences[0])
        input_size = len(sequences[0][0])
        
        transposed = [[[0.0 for _ in range(input_size)] for _ in range(seq_len)] for _ in range(batch_size)]
        
        for t in range(seq_len):
            for b in range(batch_size):
                for i in range(input_size):
                    transposed[b][t][i] = sequences[t][b][i]
        
        return transposed
    
    def forward(self, input_sequence):
        batch_case = (isinstance(input_sequence[0][0], list) and 
                     isinstance(input_sequence[0][0][0], (int, float)))
        sequences = self._transpose_batch_sequences(input_sequence) if batch_case else [input_sequence]
        
        outputs_batch = []
        for single_sequence in sequences:
            hidden_state = [0.0] * self.units
            cell_state = [0.0] * self.units
            outputs = []
            
            for input_t in single_sequence:
                hidden_state, cell_state = self._lstm_step(input_t, hidden_state, cell_state)
                outputs.append(hidden_state[:])
            
            outputs_batch.append(outputs)
        
        if batch_case:
            return self._transpose_batch_sequences(outputs_batch)
        else:
            return outputs_batch[0]
    
    def backward(self, gradient):
        return gradient


class GRULayer(Layer):
    """Gated Recurrent Unit layer"""
    def __init__(self, units, input_size):
        super().__init__()
        self.units = units
        self.input_size = input_size
        self.hidden_size = units
        
        # Initialize weights for gates
        self.wr = self._initialize_weights('xavier', input_size, units)
        self.ur = self._initialize_weights('xavier', self.hidden_size, units)
        self.br = [0.0] * units
        
        self.wz = self._initialize_weights('xavier', input_size, units)
        self.uz = self._initialize_weights('xavier', self.hidden_size, units)
        self.bz = [0.0] * units
        
        self.wh = self._initialize_weights('xavier', input_size, units)
        self.uh = self._initialize_weights('xavier', self.hidden_size, units)
        self.bh = [0.0] * units
    
    def _initialize_weights(self, method, input_dim, units):
        std = math.sqrt(2.0 / (input_dim + units))
        return [[self._rand_normal(0, std) for _ in range(input_dim)] for _ in range(units)]
    
    def _rand_normal(self, mean, std_dev):
        u1 = random.random()
        u2 = random.random()
        z0 = math.sqrt(-2 * math.log(u1)) * math.cos(2 * math.pi * u2)
        return z0 * std_dev + mean
    
    def _add_vectors(self, a, b):
        return [a[i] + b[i] for i in range(len(a))]
    
    def _multiply_elementwise(self, a, b):
        return [a[i] * b[i] for i in range(len(a))]
    
    def _multiply_matrix_vector(self, matrix, vector):
        result = [0.0] * len(matrix)
        for i in range(len(matrix)):
            result[i] = sum(matrix[i][j] * vector[j] for j in range(len(vector)))
        return result
    
    def _gru_step(self, input_t, prev_hidden):
        # Reset gate
        r_input = self._add_vectors(
            self._multiply_matrix_vector(self.wr, input_t),
            self._multiply_matrix_vector(self.ur, prev_hidden)
        )
        r_input = self._add_vectors(r_input, self.br)
        r_gate = Sigmoid()(r_input)
        
        # Update gate
        z_input = self._add_vectors(
            self._multiply_matrix_vector(self.wz, input_t),
            self._multiply_matrix_vector(self.uz, prev_hidden)
        )
        z_input = self._add_vectors(z_input, self.bz)
        z_gate = Sigmoid()(z_input)
        
        # Hidden candidate
        rh_hidden = self._multiply_elementwise(r_gate, prev_hidden)
        h_input = self._add_vectors(
            self._multiply_matrix_vector(self.wh, input_t),
            self._multiply_matrix_vector(self.uh, rh_hidden)
        )
        h_input = self._add_vectors(h_input, self.bh)
        h_tilde = Tanh()(h_input)
        
        # New hidden state
        one_minus_z = [1.0 - z_gate[i] for i in range(len(z_gate))]
        term1 = self._multiply_elementwise(one_minus_z, h_tilde)
        term2 = self._multiply_elementwise(z_gate, prev_hidden)
        new_hidden = self._add_vectors(term1, term2)
        
        return new_hidden
    
    def _transpose_batch_sequences(self, sequences):
        seq_len = len(sequences)
        batch_size = len(sequences[0])
        input_size = len(sequences[0][0])
        
        transposed = [[[0.0 for _ in range(input_size)] for _ in range(seq_len)] for _ in range(batch_size)]
        
        for t in range(seq_len):
            for b in range(batch_size):
                for i in range(input_size):
                    transposed[b][t][i] = sequences[t][b][i]
        
        return transposed
    
    def forward(self, input_sequence):
        batch_case = (isinstance(input_sequence[0][0], list) and 
                     isinstance(input_sequence[0][0][0], (int, float)))
        sequences = self._transpose_batch_sequences(input_sequence) if batch_case else [input_sequence]
        
        outputs_batch = []
        for single_sequence in sequences:
            hidden_state = [0.0] * self.units
            outputs = []
            
            for input_t in single_sequence:
                hidden_state = self._gru_step(input_t, hidden_state)
                outputs.append(hidden_state[:])
            
            outputs_batch.append(outputs)
        
        if batch_case:
            return self._transpose_batch_sequences(outputs_batch)
        else:
            return outputs_batch[0]
    
    def backward(self, gradient):
        return gradient


class SoftmaxLayer(Layer):
    """Softmax activation layer"""
    def __init__(self):
        super().__init__()
        self.trainable = False
    
    def _compute_softmax(self, x):
        max_val = max(x)
        exps = [math.exp(val - max_val) for val in x]
        sum_exps = sum(exps)
        return [exp / sum_exps for exp in exps]
    
    def forward(self, input_data):
        batch_case = (isinstance(input_data[0], list) and 
                     isinstance(input_data[0][0], (int, float)))
        input_tensor = input_data if batch_case else [input_data]
        
        output_batch = [self._compute_softmax(x) for x in input_tensor]
        
        return output_batch if batch_case else output_batch[0]
    
    def backward(self, gradient):
        return gradient


class EmbeddingLayer(Layer):
    """Embedding layer - converts word indices to dense vectors"""
    def __init__(self, vocab_size, embedding_dim, padding_idx=None):
        super().__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx
        self.trainable = False  # Set to False for now (no gradient update)
        
        # Xavier initialization
        import math
        limit = math.sqrt(6.0 / (vocab_size + embedding_dim))
        self.weights = [[random.uniform(-limit, limit) for _ in range(embedding_dim)] 
                       for _ in range(vocab_size)]
    
    def forward(self, input_data):
        # Check if input is a batch of sequences
        batch_case = isinstance(input_data[0], list)
        input_tensor = input_data if batch_case else [input_data]
        
        output_batch = []
        for sequence in input_tensor:
            embedded_sequence = []
            for idx in sequence:
                # Convert to int and clamp to valid range
                idx_int = int(idx)
                idx_int = max(0, min(idx_int, self.vocab_size - 1))
                embedded_sequence.append(self.weights[idx_int][:])
            output_batch.append(embedded_sequence)
        
        return output_batch if batch_case else output_batch[0]
    
    def backward(self, gradient):
        # For now, pass gradient through
        # Full backprop through embeddings would require tracking indices
        return gradient


class FlattenLayer(Layer):
    """Flatten layer - handles 2D and 3D tensors"""
    def __init__(self):
        super().__init__()
        self.trainable = False
        self.input_shape = None
    
    def forward(self, input_data):
        # Handle batch of sequences (3D: batch x sequence x features)
        if (isinstance(input_data[0], list) and 
            isinstance(input_data[0][0], list) and 
            isinstance(input_data[0][0][0], (int, float))):
            self.input_shape = (len(input_data), len(input_data[0]), len(input_data[0][0]))
            # Flatten each sample in the batch
            return [[item for sublist in sample for item in sublist] for sample in input_data]
        # Handle batch of vectors (2D: batch x features)
        elif isinstance(input_data[0], list) and isinstance(input_data[0][0], (int, float)):
            self.input_shape = (len(input_data), len(input_data[0]))
            return input_data
        # Handle single sequence (2D: sequence x features)
        elif isinstance(input_data[0], list):
            self.input_shape = (len(input_data), len(input_data[0]))
            return [[item for sublist in input_data for item in sublist]]
        # Handle single vector (1D: features)
        else:
            self.input_shape = (len(input_data),)
            return [input_data]
    
    def backward(self, gradient):
        return gradient


class ReshapeLayer(Layer):
    """Reshape layer"""
    def __init__(self, shape):
        super().__init__()
        self.target_shape = shape
        self.trainable = False
    
    def _reshape_tensor(self, tensor, new_shape):
        flattened = self._flatten(tensor)
        return self._build_tensor(flattened, new_shape)
    
    def _flatten(self, tensor):
        if not isinstance(tensor[0], list):
            return tensor
        else:
            result = []
            for item in tensor:
                if isinstance(item, list):
                    result.extend(self._flatten(item))
                else:
                    result.append(item)
            return result
    
    def _build_tensor(self, flat_array, shape):
        if len(shape) == 1:
            return flat_array
        elif len(shape) == 2:
            rows, cols = shape
            result = []
            for i in range(rows):
                row = []
                for j in range(cols):
                    row.append(flat_array[i * cols + j])
                result.append(row)
            return result
        else:
            size = shape[0]
            remaining_shape = shape[1:]
            remaining_size = 1
            for dim in remaining_shape:
                remaining_size *= dim
            
            result = []
            for i in range(size):
                sub_array = flat_array[i * remaining_size:(i + 1) * remaining_size]
                result.append(self._build_tensor(sub_array, remaining_shape))
            return result
    
    def forward(self, input_data):
        return self._reshape_tensor(input_data, self.target_shape)
    
    def backward(self, gradient):
        # This would need to reshape back to original shape
        # For now, just return the gradient as is
        return gradient
