"""
3D Feature Extractor Component for Neutrophil Maturation Classification

This module implements a configurable 3D CNN feature extractor that extends the
existing 2D ResNet-based architecture to handle 3D volumetric data. It serves as
the foundation for contrastive learning and 3D classification by providing shared
feature representations from 3D microscopy volumes.

The 3D FeatureExtractor converts Conv2D layers to Conv3D with appropriate kernel
sizes while maintaining residual connections and batch normalization patterns.
"""

import tensorflow as tf
from tensorflow.keras.layers import (
    Input, Conv3D, BatchNormalization, Add, ReLU, LeakyReLU, PReLU,
    Lambda, GlobalAveragePooling3D
)
from tensorflow.keras.models import Model
from tensorflow.keras.activations import get as get_activation

# Import DynamicResidualScaling from the same models package
from .dynamic_residual_scaling import DynamicResidualScaling, create_adaptive_lambda_scaling


@tf.keras.utils.register_keras_serializable()
class FeatureExtractor3D(tf.keras.Model):
    """
    Configurable 3D CNN Feature Extractor for Neutrophil Classification
    
    This class extends the 2D FeatureExtractor to handle 3D volumetric data,
    providing shared feature representations for contrastive learning and
    3D classification tasks.
    
    The 3D feature extractor supports:
    - 3D convolutional blocks with configurable depth and filter sizes
    - Dynamic residual scaling strategies for bias correction
    - Multiple activation functions and initialization strategies
    - Backward compatibility with existing TOML configurations
    - Memory-efficient 3D processing for typical volume sizes (69³, 103³)
    
    Architecture Overview:
    ----------------------
    Input -> 3D Conv Blocks (with optional residual connections) -> 3D Feature Tensor
    
    Each 3D Conv Block contains:
    - Multiple Conv3D layers (configurable count)
    - Optional BatchNormalization
    - Optional residual connections with 3D projection
    - Configurable activation functions
    - Dynamic residual scaling (if enabled)
    
    The output is the 3D feature tensor BEFORE GlobalAveragePooling3D, allowing
    contrastive heads and other architectures to apply their own pooling strategies.
    """
    
    def __init__(self, config=None, input_shape=None, input_shape_config=None, name='feature_extractor_3d', **kwargs):
        """
        Initialize the 3D FeatureExtractor
        
        Args:
            config (dict): Configuration dictionary containing:
                - conv_blocks: Configuration for 3D convolutional blocks
                - input_shape: Shape of input tensor (optional, can be inferred)
            input_shape (tuple): Input shape (D, H, W, C), overrides config if provided
            input_shape_config (tuple): Alternative parameter name for input_shape (for serialization compatibility)
            name (str): Name of the model
            **kwargs: Additional keyword arguments for tf.keras.Model
        """
        super(FeatureExtractor3D, self).__init__(name=name, **kwargs)
        
        # Store configuration
        self.config = config or {}
        # Handle both input_shape and input_shape_config parameters for backward compatibility
        self.input_shape_config = input_shape_config or input_shape
        
        # Validate and set default configuration
        self._validate_and_set_defaults()
        
        # Build the 3D feature extraction layers
        self._build_feature_layers()
    
    def _validate_and_set_defaults(self):
        """Validate configuration and set defaults for 3D architecture"""
        
        # Default configuration for 3D ResNet architecture
        default_config = {
            'conv_blocks': {
                'filters': [32, 64, 128],
                'kernel_size': [3, 3, 3],  # 3D kernel size
                'activation': {'name': 'leaky_relu', 'alpha': 0.1},
                'padding': 'same',
                'use_batch_norm': True,
                'use_residual': True,
                'conv_layers': 2,
                'residual_kernel_size': [[1, 1, 1], [1, 1, 1]],  # 3D residual kernels
                'kernel_initializer': 'he_normal',
                'bias_initializer': 'zeros',
                'residual_scaling': {
                    'strategy': 'fixed_scaling',
                    'params': {
                        'alpha': 0.5,
                        'scale': 0.2,
                        'momentum': 0.99,
                        'epsilon': 1e-6,
                        'min_scale': 0.1,
                        'max_scale': 2.0
                    }
                }
            }
        }
        
        # Merge user config with defaults
        if 'conv_blocks' not in self.config:
            self.config['conv_blocks'] = default_config['conv_blocks']
        else:
            # Merge conv_blocks configuration
            for key, value in default_config['conv_blocks'].items():
                if key not in self.config['conv_blocks']:
                    self.config['conv_blocks'][key] = value
        
        # Validate required parameters
        conv_config = self.config['conv_blocks']
        required_params = ['filters', 'kernel_size', 'activation']
        for param in required_params:
            if param not in conv_config:
                raise ValueError(f"Missing required parameter '{param}' in conv_blocks configuration")
        
        # Ensure filters is a list
        if not isinstance(conv_config['filters'], list):
            raise ValueError("'filters' must be a list of integers")
        
        # Validate 3D kernel size
        kernel_size = conv_config['kernel_size']
        if isinstance(kernel_size, int):
            conv_config['kernel_size'] = [kernel_size, kernel_size, kernel_size]
        elif len(kernel_size) == 2:
            # Convert 2D to 3D by replicating the kernel size
            conv_config['kernel_size'] = [kernel_size[0], kernel_size[1], kernel_size[1]]
        elif len(kernel_size) != 3:
            raise ValueError("'kernel_size' must be int, 2-tuple, or 3-tuple for 3D convolution")
        
        # Validate residual_kernel_size for 3D
        if 'residual_kernel_size' in conv_config:
            residual_kernels = conv_config['residual_kernel_size']
            expected_length = len(conv_config['filters'])
            
            # Convert 2D residual kernels to 3D if necessary
            for i, kernel in enumerate(residual_kernels):
                if isinstance(kernel, int):
                    residual_kernels[i] = [kernel, kernel, kernel]
                elif len(kernel) == 2:
                    residual_kernels[i] = [kernel[0], kernel[1], kernel[1]]
                elif len(kernel) != 3:
                    print(f"Warning: Invalid residual kernel size at index {i}. Using [1,1,1].")
                    residual_kernels[i] = [1, 1, 1]
            
            if len(residual_kernels) != expected_length:
                print(f"Warning: residual_kernel_size length ({len(residual_kernels)}) "
                      f"doesn't match expected length ({expected_length}). Using default [1,1,1] for missing entries.")
                # Pad with default if too short
                if len(residual_kernels) < expected_length:
                    residual_kernels.extend([[1, 1, 1]] * (expected_length - len(residual_kernels)))
    
    def _get_weight_initializer(self, initializer_name):
        """Get weight initializer based on name with stability improvements for 3D"""
        if initializer_name is None:
            return 'he_normal'  # Default stable choice for ReLU networks
        
        initializer_name = initializer_name.lower()
        
        # Map common initializer names to TensorFlow initializers
        initializer_map = {
            'he_normal': tf.keras.initializers.HeNormal(seed=42),
            'he_uniform': tf.keras.initializers.HeUniform(seed=42),
            'xavier_normal': tf.keras.initializers.GlorotNormal(seed=42),
            'glorot_normal': tf.keras.initializers.GlorotNormal(seed=42),
            'xavier_uniform': tf.keras.initializers.GlorotUniform(seed=42),
            'glorot_uniform': tf.keras.initializers.GlorotUniform(seed=42),
            'lecun_normal': tf.keras.initializers.LecunNormal(seed=42),
            'lecun_uniform': tf.keras.initializers.LecunUniform(seed=42),
            'random_normal': tf.keras.initializers.RandomNormal(stddev=0.02, seed=42),
            'random_uniform': tf.keras.initializers.RandomUniform(minval=-0.05, maxval=0.05, seed=42),
            'truncated_normal': tf.keras.initializers.TruncatedNormal(stddev=0.02, seed=42),
            'zeros': 'zeros',
            'ones': 'ones'
        }
        
        if initializer_name in initializer_map:
            return initializer_map[initializer_name]
        else:
            print(f"Warning: Unknown initializer '{initializer_name}', using 'he_normal'")
            return tf.keras.initializers.HeNormal(seed=42)
    
    def _get_activation_function(self, activation):
        """Convert activation function specification to a callable function or string"""
        if isinstance(activation, str):
            return activation
        elif isinstance(activation, dict):
            if 'name' not in activation:
                print(f"Warning: Activation dictionary missing 'name' key. Using 'relu' as default.")
                return 'relu'
            
            name = activation['name'].lower()
            
            if name == 'leaky_relu':
                alpha = activation.get('alpha', 0.3)
                return LeakyReLU(alpha=alpha)
            elif name == 'prelu':
                return PReLU()
            elif name == 'elu':
                alpha = activation.get('alpha', 1.0)
                return tf.keras.layers.ELU(alpha=alpha)
            elif name == 'selu':
                return 'selu'
            elif name == 'swish' or name == 'silu':
                return tf.keras.activations.swish
            elif name == 'gelu':
                return tf.keras.activations.gelu
            elif name == 'relu6':
                return tf.keras.layers.ReLU(max_value=6)
            else:
                try:
                    return get_activation(name)
                except ValueError:
                    print(f"Warning: Unknown activation function '{name}'. Using 'relu' as default.")
                    return 'relu'
        else:
            if callable(activation):
                return activation
            print(f"Warning: Unsupported activation type {type(activation)}. Using 'relu' as default.")
            return 'relu'
    
    def _build_feature_layers(self):
        """Build the 3D convolutional feature extraction layers"""
        
        conv_config = self.config['conv_blocks']
        
        # Get configuration parameters
        filters_list = conv_config['filters']
        kernel_size_base = tuple(conv_config.get('kernel_size', [3, 3, 3]))
        activation_config = conv_config.get('activation', 'relu')
        padding_base = conv_config.get('padding', 'same')
        use_batch_norm_base = conv_config.get('use_batch_norm', False)
        use_residual_base = conv_config.get('use_residual', False)
        conv_layers_count = conv_config.get('conv_layers', 2)
        residual_kernel_sizes_list = conv_config.get('residual_kernel_size', [])
        
        # Get weight initializers
        conv_kernel_init_name = conv_config.get('kernel_initializer', 'he_normal')
        conv_bias_init_name = conv_config.get('bias_initializer', 'zeros')
        conv_initializer = self._get_weight_initializer(conv_kernel_init_name)
        conv_bias_initializer = conv_bias_init_name
        
        # Get residual scaling configuration
        residual_scaling_config = conv_config.get('residual_scaling', {})
        self.residual_scaling_strategy = residual_scaling_config.get('strategy', 'standard')
        self.residual_scaling_params = residual_scaling_config.get('params', {})
        
        # Store layers for building the model
        self.conv_layers = []
        
        # Build layers for each block
        for i, current_filters in enumerate(filters_list):
            block_layers = {}
            
            # 3D ResNet downsampling: first block keeps spatial size, subsequent blocks downsample by 2
            downsample_stride = (2, 2, 2) if i > 0 else (1, 1, 1)
            
            # Build conv layers for this block
            block_conv_layers = []
            block_bn_layers = []
            block_activation_layers = []
            
            for i_conv in range(conv_layers_count):
                # First conv in each block uses downsampling stride, others use (1,1,1)
                conv_strides = downsample_stride if i_conv == 0 else (1, 1, 1)
                
                conv_layer = Conv3D(
                    current_filters,
                    kernel_size_base,
                    strides=conv_strides,
                    activation=None,  # Activation applied separately
                    padding=padding_base,
                    kernel_initializer=conv_initializer,
                    bias_initializer=conv_bias_initializer,
                    use_bias=not use_batch_norm_base,  # No bias when using BatchNorm
                    name=f'main_block_{i}_conv3d_{i_conv}'
                )
                block_conv_layers.append(conv_layer)
                
                if use_batch_norm_base:
                    bn_layer = BatchNormalization(
                        momentum=0.99,
                        epsilon=1e-3,
                        center=True,
                        scale=True,
                        beta_initializer='zeros',
                        gamma_initializer='ones',
                        name=f'bn_main_block_{i}_conv3d_{i_conv}'
                    )
                    block_bn_layers.append(bn_layer)
                else:
                    block_bn_layers.append(None)
                
                # Apply activation for all but the last conv layer in the main path
                if i_conv < conv_layers_count - 1:
                    activation_fn = self._get_activation_function(activation_config)
                    # Set name for layer-based activations
                    if hasattr(activation_fn, '__call__') and hasattr(activation_fn, 'name'):
                        if hasattr(activation_fn, '__class__') and hasattr(activation_fn.__class__, '__name__'):
                            layer_type = activation_fn.__class__.__name__.lower()
                            activation_fn.name = f'{layer_type}_main_block_{i}_conv3d_{i_conv}'
                    block_activation_layers.append(activation_fn)
                else:
                    block_activation_layers.append(None)
            
            block_layers['conv_layers'] = block_conv_layers
            block_layers['bn_layers'] = block_bn_layers
            block_layers['activation_layers'] = block_activation_layers
            
            # Build residual projection if needed
            if use_residual_base:
                res_kernel_idx = i - 1
                res_kernel_size = tuple(residual_kernel_sizes_list[res_kernel_idx] if res_kernel_idx < len(residual_kernel_sizes_list) and res_kernel_idx >= 0 else [1, 1, 1])
                
                residual_projection = Conv3D(
                    current_filters,
                    res_kernel_size,
                    strides=downsample_stride,
                    padding='same',
                    activation=None,
                    kernel_initializer=conv_initializer,
                    bias_initializer=conv_bias_initializer,
                    use_bias=not use_batch_norm_base,
                    name=f'conv3d_skip_projection_block_{i}'
                )
                block_layers['residual_projection'] = residual_projection
                
                if use_batch_norm_base:
                    residual_bn = BatchNormalization(
                        momentum=0.99,
                        epsilon=1e-3,
                        center=True,
                        scale=True,
                        beta_initializer='zeros',
                        gamma_initializer='ones',
                        name=f'bn_skip_projection_block_{i}'
                    )
                    block_layers['residual_bn'] = residual_bn
                else:
                    block_layers['residual_bn'] = None
                
                # Create residual scaling layer
                if self.residual_scaling_strategy == 'standard':
                    block_layers['residual_scaling'] = None
                elif self.residual_scaling_strategy == 'fixed_scaling':
                    scale = self.residual_scaling_params.get('scale', 0.5)
                    block_layers['residual_scaling'] = Lambda(
                        lambda inputs: inputs[0] + inputs[1] * scale,
                        name=f'lambda_fixed_scale_3d_block_{i}'
                    )
                elif self.residual_scaling_strategy == 'dynamic_adaptive':
                    alpha = self.residual_scaling_params.get('alpha', 0.5)
                    block_layers['residual_scaling'] = DynamicResidualScaling(
                        strategy='adaptive_norm',
                        alpha=alpha,
                        name=f'dyn_scale_adaptive_3d_block_{i}'
                    )
                elif self.residual_scaling_strategy == 'dynamic_learned':
                    alpha = self.residual_scaling_params.get('alpha', 0.5)
                    block_layers['residual_scaling'] = DynamicResidualScaling(
                        strategy='learned_balance',
                        alpha=alpha,
                        name=f'dyn_scale_learned_3d_block_{i}'
                    )
                elif self.residual_scaling_strategy == 'lambda_adaptive':
                    alpha = self.residual_scaling_params.get('alpha', 0.5)
                    scaling_func = create_adaptive_lambda_scaling(alpha=alpha)
                    block_layers['residual_scaling'] = Lambda(
                        scaling_func,
                        name=f'lambda_adaptive_scale_3d_block_{i}'
                    )
                else:
                    raise ValueError(f"Unknown residual scaling strategy: {self.residual_scaling_strategy}")
            else:
                block_layers['residual_projection'] = None
                block_layers['residual_bn'] = None
                block_layers['residual_scaling'] = None
            
            # Final block activation
            final_activation_fn = self._get_activation_function(activation_config)
            if hasattr(final_activation_fn, '__call__') and hasattr(final_activation_fn, 'name'):
                if hasattr(final_activation_fn, '__class__') and hasattr(final_activation_fn.__class__, '__name__'):
                    layer_type = final_activation_fn.__class__.__name__.lower()
                    final_activation_fn.name = f'{layer_type}_final_3d_block_{i}'
            block_layers['final_activation'] = final_activation_fn
            
            # Store block configuration
            self.conv_layers.append(block_layers)
        
        print(f"✓ FeatureExtractor3D built with {len(filters_list)} convolutional blocks")
        print(f"  - Filters: {filters_list}")
        print(f"  - Kernel size: {kernel_size_base}")
        print(f"  - Residual connections: {use_residual_base}")
        print(f"  - Batch normalization: {use_batch_norm_base}")
        print(f"  - Residual scaling strategy: {self.residual_scaling_strategy}")
        if self.residual_scaling_params and self.residual_scaling_strategy != 'standard':
            print(f"  - Residual scaling params: {self.residual_scaling_params}")
    
    def build(self, inputs_shape):
        """
        Build all layers with the given input shape
        
        This method builds all the layers in the 3D feature extractor with the specified
        input shape, ensuring that layer output shapes are properly computed and
        available for model summary functions.
        
        Args:
            inputs_shape: Input shape tuple (batch_size, depth, height, width, channels)
                        or (depth, height, width, channels)
        """
        # Handle input shape format - add batch dimension if not present
        if len(inputs_shape) == 4:
            # Input shape without batch dimension, add None for batch
            shape_with_batch = (None,) + inputs_shape
        else:
            # Input shape already includes batch dimension
            shape_with_batch = inputs_shape
        
        conv_config = self.config['conv_blocks']
        filters_list = conv_config['filters']
        use_residual_base = conv_config.get('use_residual', False)
        conv_layers_count = conv_config.get('conv_layers', 2)
        
        # Track current shape through the network (with batch dimension)
        current_shape = shape_with_batch
        
        # Build each convolutional block
        for i, current_filters in enumerate(filters_list):
            block_layers = self.conv_layers[i]
            
            # Standard ResNet downsampling: first block keeps spatial size, subsequent blocks downsample by 2
            downsample_stride = (2, 2, 2) if i > 0 else (1, 1, 1)
            
            # Store the input shape for this block (for residual connections)
            block_input_shape = current_shape
            
            # Build conv layers for this block
            for i_conv in range(conv_layers_count):
                conv_strides = downsample_stride if i_conv == 0 else (1, 1, 1)
                
                # Build conv layer
                conv_layer = block_layers['conv_layers'][i_conv]
                conv_layer.build(current_shape)
                current_shape = conv_layer.compute_output_shape(current_shape)
                
                # Build batch normalization layer if present
                if block_layers['bn_layers'][i_conv] is not None:
                    bn_layer = block_layers['bn_layers'][i_conv]
                    bn_layer.build(current_shape)
                    current_shape = bn_layer.compute_output_shape(current_shape)
            
            # Build residual projection layers if needed
            if use_residual_base:
                # Build residual projection layer
                if block_layers['residual_projection'] is not None:
                    proj_layer = block_layers['residual_projection']
                    proj_layer.build(block_input_shape)
                    proj_output_shape = proj_layer.compute_output_shape(block_input_shape)
                    
                    # Build residual batch normalization if present
                    if block_layers['residual_bn'] is not None:
                        residual_bn = block_layers['residual_bn']
                        residual_bn.build(proj_output_shape)
        
        # Mark as built
        self.built = True
    
    def call(self, inputs, training=None, **kwargs):
        """
        Forward pass through the 3D feature extractor
        
        Args:
            inputs: Input tensor of shape (batch_size, depth, height, width, channels)
            training: Boolean indicating training mode
        
        Returns:
            Feature tensor before global average pooling of shape 
            (batch_size, feature_depth, feature_height, feature_width, feature_channels)
        """
        if training is None:
            training = tf.keras.backend.learning_phase()

        x = inputs
        
        conv_config = self.config['conv_blocks']
        filters_list = conv_config['filters']
        use_residual_base = conv_config.get('use_residual', False)
        conv_layers_count = conv_config.get('conv_layers', 2)
        
        # Process each convolutional block
        for i, current_filters in enumerate(filters_list):
            with tf.name_scope(f'feature_3d_block_{i}'):
                block_layers = self.conv_layers[i]
                block_input_tensor = x  # Input to the current block for potential skip connection
                current_path_tensor = x  # Tensor being transformed within the main path
                
                # Apply conv layers in the main path
                for i_conv in range(conv_layers_count):
                    with tf.name_scope(f'conv3d_{i_conv}'):
                        # Apply 3D convolution
                        current_path_tensor = block_layers['conv_layers'][i_conv](current_path_tensor)
                        
                        # Apply batch normalization if enabled
                        if block_layers['bn_layers'][i_conv] is not None:
                            current_path_tensor = block_layers['bn_layers'][i_conv](current_path_tensor, training=training)
                        
                        # Apply activation for all but the last conv layer
                        if block_layers['activation_layers'][i_conv] is not None:
                            current_path_tensor = block_layers['activation_layers'][i_conv](current_path_tensor)
                
                # Handle residual connections
                if use_residual_base:
                    projected_skip_tensor = block_input_tensor
                    needs_projection = False
                    
                    # Determine stride used in this block
                    downsample_stride = (2, 2, 2) if i > 0 else (1, 1, 1)
                    
                    # Need projection if stride changes or channels change
                    if downsample_stride != (1, 1, 1) or block_input_tensor.shape[-1] != current_filters:
                        needs_projection = True
                    
                    if needs_projection:
                        projected_skip_tensor = block_layers['residual_projection'](block_input_tensor)
                        if block_layers['residual_bn'] is not None:
                            projected_skip_tensor = block_layers['residual_bn'](projected_skip_tensor, training=training)
                    else:
                        projected_skip_tensor = block_input_tensor
                    
                    # Apply residual scaling strategy
                    if self.residual_scaling_strategy == 'standard':
                        x = Add(name=f'add_standard_3d_block_{i}')([current_path_tensor, projected_skip_tensor])
                    elif block_layers['residual_scaling'] is not None:
                        x = block_layers['residual_scaling']([current_path_tensor, projected_skip_tensor])
                    else:
                        # Fallback to standard add
                        x = Add(name=f'add_fallback_3d_block_{i}')([current_path_tensor, projected_skip_tensor])
                    
                    # Apply final block activation AFTER Add or scaling
                    x = block_layers['final_activation'](x)
                else:
                    # Not a residual block, apply final activation to main path
                    x = block_layers['final_activation'](current_path_tensor)
        
        return x
    
    def get_feature_shape(self, input_shape):
        """
        Calculate the output feature shape given an input shape
        
        Args:
            input_shape: Input shape tuple (depth, height, width, channels)
        
        Returns:
            Output feature shape tuple (feature_depth, feature_height, feature_width, feature_channels)
        """
        # Build a temporary model to calculate output shape
        temp_input = Input(shape=input_shape)
        temp_output = self(temp_input)
        return temp_output.shape[1:]  # Remove batch dimension
    
    def build_model(self, input_shape):
        """
        Build a functional Keras model from the 3D feature extractor
        
        Args:
            input_shape: Input shape tuple (depth, height, width, channels)
        
        Returns:
            tf.keras.Model: Functional model
        """
        inputs = Input(shape=input_shape, name='feature_extractor_3d_input')
        features = self(inputs)
        
        model = Model(inputs=inputs, outputs=features, name=self.name + '_model')
        return model
    
    def get_config(self):
        """Get configuration for serialization"""
        config = super(FeatureExtractor3D, self).get_config()
        config.update({
            'config': self.config,
            'input_shape_config': self.input_shape_config
        })
        return config
    
    @classmethod
    def from_config(cls, config):
        """Create FeatureExtractor3D from configuration"""
        return cls(**config)


def create_feature_extractor_3d_from_config(config, input_shape=None):
    """
    Factory function to create a FeatureExtractor3D from configuration
    
    This function provides a convenient way to create a 3D FeatureExtractor
    from TOML configuration files, converting 2D configurations to 3D where needed.
    
    Args:
        config (dict): Configuration dictionary with 'model' section
        input_shape (tuple): Input shape (depth, height, width, channels)
    
    Returns:
        FeatureExtractor3D: Configured 3D feature extractor instance
    
    Example:
        >>> import toml
        >>> config = toml.load('config/resnet_3d.toml')
        >>> feature_extractor = create_feature_extractor_3d_from_config(
        ...     config, input_shape=(69, 69, 69, 1)
        ... )
        >>> features = feature_extractor(inputs)
    """
    model_config = config.get('model', {}) if 'model' in config else config
    
    # Convert 2D configuration to 3D if necessary
    if 'conv_blocks' in model_config:
        conv_blocks = model_config['conv_blocks']
        
        # Convert 2D kernel sizes to 3D
        if 'kernel_size' in conv_blocks:
            kernel_size = conv_blocks['kernel_size']
            if isinstance(kernel_size, list) and len(kernel_size) == 2:
                conv_blocks['kernel_size'] = [kernel_size[0], kernel_size[1], kernel_size[1]]
        
        # Convert 2D residual kernel sizes to 3D
        if 'residual_kernel_size' in conv_blocks:
            residual_kernels = conv_blocks['residual_kernel_size']
            for i, kernel in enumerate(residual_kernels):
                if len(kernel) == 2:
                    residual_kernels[i] = [kernel[0], kernel[1], kernel[1]]
    
    feature_extractor = FeatureExtractor3D(
        config=model_config,
        input_shape=input_shape
    )
    
    return feature_extractor


def convert_2d_to_3d_weights(model_2d, model_3d, expand_mode='replicate'):
    """
    Convert 2D model weights to 3D model weights for transfer learning.
    
    Args:
        model_2d: Trained 2D FeatureExtractor model
        model_3d: 3D FeatureExtractor model to initialize
        expand_mode: Method to expand 2D kernels to 3D
                    - 'replicate': Replicate 2D kernel along depth dimension
                    - 'center': Place 2D kernel in center slice, zero elsewhere
                    - 'average': Average 2D kernel across depth dimension
    
    Returns:
        model_3d: 3D model with converted weights
    """
    print(f"Converting 2D weights to 3D using '{expand_mode}' mode...")
    
    # Get layer mappings between 2D and 3D models
    layer_2d_dict = {layer.name: layer for layer in model_2d.layers}
    layer_3d_dict = {layer.name: layer for layer in model_3d.layers}
    
    converted_count = 0
    
    for layer_name, layer_3d in layer_3d_dict.items():
        # Find corresponding 2D layer
        layer_name_2d = layer_name.replace('conv3d', 'conv').replace('_3d_', '_')
        
        if layer_name_2d in layer_2d_dict:
            layer_2d = layer_2d_dict[layer_name_2d]
            
            # Convert Conv3D weights
            if isinstance(layer_3d, Conv3D) and hasattr(layer_2d, 'get_weights'):
                weights_2d = layer_2d.get_weights()
                if len(weights_2d) > 0:
                    kernel_2d = weights_2d[0]  # Shape: (H, W, in_channels, out_channels)
                    
                    if expand_mode == 'replicate':
                        # Replicate 2D kernel along depth dimension
                        kernel_3d = np.tile(kernel_2d[np.newaxis, :, :, :, :], 
                                          (kernel_2d.shape[0], 1, 1, 1, 1))
                    elif expand_mode == 'center':
                        # Place 2D kernel in center slice
                        depth = layer_3d.kernel_size[0]
                        kernel_3d = np.zeros((depth, kernel_2d.shape[0], kernel_2d.shape[1], 
                                           kernel_2d.shape[2], kernel_2d.shape[3]))
                        center_slice = depth // 2
                        kernel_3d[center_slice] = kernel_2d
                    elif expand_mode == 'average':
                        # Average 2D kernel across depth dimension
                        depth = layer_3d.kernel_size[0]
                        kernel_3d = np.tile(kernel_2d[np.newaxis, :, :, :, :] / depth, 
                                          (depth, 1, 1, 1, 1))
                    
                    # Update weights
                    new_weights = [kernel_3d]
                    if len(weights_2d) > 1:  # Include bias if present
                        new_weights.append(weights_2d[1])
                    
                    layer_3d.set_weights(new_weights)
                    converted_count += 1
            
            # Convert BatchNormalization weights (these are compatible)
            elif isinstance(layer_3d, BatchNormalization) and hasattr(layer_2d, 'get_weights'):
                weights_2d = layer_2d.get_weights()
                if len(weights_2d) > 0:
                    layer_3d.set_weights(weights_2d)
                    converted_count += 1
    
    print(f"✓ Converted {converted_count} layers from 2D to 3D")
    return model_3d
