model/habit.py

import tensorflow as tf
from tensorflow import keras
import numpy as np
from habit.model.layers import (
    MonotonicFunction, 
    ResidualBlock, 
    EnhancedPropertyEncoder, 
    SwishActivation,
    CrossAttentionLayer,
    GatingMechanism
)


# Define custom swish activation function
def swish(x):
    return x * tf.keras.backend.sigmoid(x)

# Register the custom activation
tf.keras.utils.get_custom_objects().update({'swish': tf.keras.layers.Activation(swish)})

class HABIT(keras.Model):
    def __init__(self, embedding_dim=128, num_heads=4, num_monotonic_basis=16, dropout_rate=0.1):
        super(HABIT, self).__init__()
    
        # Model properties
        self.embedding_dim = embedding_dim
        self.num_heads = num_heads  # Store num_heads as an instance variable
        self.num_monotonic_basis = num_monotonic_basis
        self.dropout_rate = dropout_rate
    
        # Enhanced property encoders with residual connections and batch normalization
        self.texture_encoder = self._build_property_encoder(3)
        self.bd_encoder = self._build_property_encoder(1)
        self.oc_encoder = self._build_property_encoder(1)
        self.ksat_encoder = self._build_property_encoder(1)
        
        # Cross-Attention layers for property interactions
        self.texture_bd_cross_attention = CrossAttentionLayer(
            embedding_dim=embedding_dim,
            num_heads=self.num_heads,  # Use the instance variable
            dropout_rate=dropout_rate
        )
        
        self.texture_oc_cross_attention = CrossAttentionLayer(
            embedding_dim=embedding_dim,
            num_heads=max(2, self.num_heads // 2),  # Scale down but ensure at least 2 heads
            dropout_rate=dropout_rate
        )
        
        # Gating mechanisms for controlled information flow
        self.texture_bd_gate = GatingMechanism()
        self.texture_oc_gate = GatingMechanism()
                
        # Fusion layers for enhanced property combinations
        self.texture_bd_fusion = keras.layers.Dense(embedding_dim, activation='swish')
        self.texture_oc_fusion = keras.layers.Dense(embedding_dim, activation='swish')
            
        # Property attention layer
        self.property_attention = keras.layers.MultiHeadAttention(
            num_heads=self.num_heads,  # Use the instance variable
            key_dim=self.embedding_dim // self.num_heads,  # Scale key_dim by num_heads
            value_dim=self.embedding_dim // self.num_heads,  # Scale value_dim by num_heads
            dropout=dropout_rate,
            use_bias=True
        )
    
        # Water potential attention layer
        self.wp_attention = keras.layers.MultiHeadAttention(
            num_heads=self.num_heads,  # Use the instance variable
            key_dim=self.embedding_dim // self.num_heads,  # Scale key_dim by num_heads
            value_dim=self.embedding_dim // self.num_heads,  # Scale value_dim by num_heads
            dropout=dropout_rate,
            use_bias=True
        )
            
        # Layer normalization
        self.property_layer_norm = keras.layers.LayerNormalization(epsilon=1e-6)
        self.wp_layer_norm = keras.layers.LayerNormalization(epsilon=1e-6)
        self.fusion_layer_norm = keras.layers.LayerNormalization(epsilon=1e-6)
    
        # Water potential embedding layer
        self.wp_embedding_layer = keras.layers.Dense(self.embedding_dim, activation='swish')
    
        # Additional layers
        self.soil_dense = keras.layers.Dense(self.embedding_dim, activation='swish')
        self.flatten = keras.layers.Flatten()
        
        # Enhanced feature fusion for property integration
        self.property_fusion = keras.Sequential([
            keras.layers.Dense(self.embedding_dim * 2, activation='swish'),
            keras.layers.BatchNormalization(epsilon=1e-5),
            keras.layers.Dropout(self.dropout_rate),
            keras.layers.Dense(self.embedding_dim, activation='swish')
        ])
    
        # Parameter network and monotonic function
        self.parameter_network = keras.Sequential([
            keras.layers.Dense(self.embedding_dim, activation='swish'),
            keras.layers.BatchNormalization(epsilon=1e-5),
            keras.layers.Dense(self.embedding_dim, activation='swish'),
            keras.layers.BatchNormalization(epsilon=1e-5),
            keras.layers.Dense(self.num_monotonic_basis * 2)
        ])
    
        self.monotonic_function = MonotonicFunction(self.num_monotonic_basis)
    
        # Storage for attention weights
        self._property_attention_weights = None
        self._wp_attention_weights = None
        self._cross_attention_weights = None  # Will store cross-attention weights

    def build(self, input_shape):
        # Ensure all layers are built
        if isinstance(input_shape, list):
            texture_shape, bd_shape, oc_shape, ksat_shape, prop_mask_shape, wp_shape = input_shape
        else:
            # Default input shapes if not provided
            texture_shape = (None, 3)
            bd_shape = (None, 1)
            oc_shape = (None, 1)
            ksat_shape = (None, 1)
            prop_mask_shape = (None, 4)
            wp_shape = (None, None)
    
        # Build property encoders
        self.texture_encoder.build(texture_shape)
        self.bd_encoder.build(bd_shape)
        self.oc_encoder.build(oc_shape)
        self.ksat_encoder.build(ksat_shape)
    
        # Build cross-attention and gating layers
        self.texture_bd_cross_attention.build([texture_shape, bd_shape])
        self.texture_oc_cross_attention.build([texture_shape, oc_shape])
    
        # Ensure get_attention_weights method exists
        if not hasattr(self, 'get_attention_weights'):
            def get_attention_weights(self):
                """Return attention weights for interpretation"""
                return {
                    'property_attention': getattr(self, '_property_attention_weights', None),
                    'wp_attention': getattr(self, '_wp_attention_weights', None),
                    'cross_attention': getattr(self, '_cross_attention_weights', None)
                }
        
            # Bind the method to the instance
            import types
            self.get_attention_weights = types.MethodType(get_attention_weights, self)
    
        super().build(input_shape)
    
    def _build_property_encoder(self, input_dim):
        """
        Build an enhanced property encoder with residual connections and batch normalization
        """
        return EnhancedPropertyEncoder(
            embedding_dim=self.embedding_dim,
            input_dim=input_dim,
            dropout_rate=self.dropout_rate
        )
    
    def call(self, inputs, training=False):
        
        texture, bd, oc, ksat, properties_mask, water_potential = inputs
        # Store inputs for access in loss functions
        self.texture_input = texture
        self.bd_input = bd
        self.oc_input = oc
        self.ksat_input = ksat
        self.properties_mask_input = properties_mask
        self.wp_input = water_potential
        
        # Create mask for valid water potential values
        # Note: In log10 scale, 1.0 kPa = 0.0, which is valid!
        # Only mask out actual padding values (use large negative sentinel)
        wp_mask = tf.cast(tf.greater(water_potential, -99), tf.float32)

        # Feature extraction with enhanced property encoders
        texture_features = self.texture_encoder(texture, training=training)
        bd_features = self.bd_encoder(bd, training=training)
    
        # Ensure texture and bulk density features have correct shape
        texture_features = tf.ensure_shape(texture_features, [None, self.embedding_dim])
        bd_features = tf.ensure_shape(bd_features, [None, self.embedding_dim])
    
        # Process optional properties
        oc_features = self.oc_encoder(oc, training=training) * tf.expand_dims(properties_mask[:, 2], -1)
        ksat_features = self.ksat_encoder(ksat, training=training) * tf.expand_dims(properties_mask[:, 3], -1)
    
        # Safe handling of optional features
        oc_features_safe = oc_features if oc_features.shape[-1] == self.embedding_dim else \
            tf.zeros_like(texture_features, dtype=texture_features.dtype)

        # =================================================================
        # CROSS-ATTENTION: ALWAYS COMPUTE, USE MASKING FOR CONTROL
        # =================================================================
        
        # Initialize cross-attention weights storage
        cross_attention_weights = {}
        
        # 1. TEXTURE ↔ BD CROSS-ATTENTION
        bd_availability = tf.expand_dims(properties_mask[:, 1], -1)
        
        # Compute true cross-property attention (2x2 matrix)
        texture_bd_combined, texture_bd_weights = self.texture_bd_cross_attention(
            inputs=[texture_features, bd_features],
            training=training,
            return_attention_weights=True
        )

        # Extract enhanced features from the 2-property output
        texture_enhanced_raw = texture_bd_combined[:, 0, :]  # First property (texture)
        bd_enhanced_raw = texture_bd_combined[:, 1, :]       # Second property (BD)

        # Apply gating and masking
        texture_bd_gated = self.texture_bd_gate([texture_features, texture_enhanced_raw]) * bd_availability
        bd_texture_gated = self.texture_bd_gate([bd_features, bd_enhanced_raw]) * bd_availability

        # Fuse features
        texture_enhanced_by_bd_raw = self.texture_bd_fusion(texture_bd_gated)
        bd_enhanced_by_texture_raw = self.texture_bd_fusion(bd_texture_gated)
        
        # When BD not available, fall back to original features
        texture_enhanced_by_bd = texture_features + (texture_enhanced_by_bd_raw - texture_features) * bd_availability
        bd_enhanced_by_texture = bd_features + (bd_enhanced_by_texture_raw - bd_features) * bd_availability
        
        # Store attention weights
        cross_attention_weights['texture_bd'] = texture_bd_weights
        
        # 2. TEXTURE ↔ OC CROSS-ATTENTION  
        oc_availability = tf.expand_dims(properties_mask[:, 2], -1)

        # Compute true cross-property attention (2x2 matrix)
        texture_oc_combined, texture_oc_weights = self.texture_oc_cross_attention(
            inputs=[texture_features, oc_features_safe],
            training=training,
            return_attention_weights=True
        )

        # Extract enhanced features from the 2-property output
        texture_enhanced_raw_oc = texture_oc_combined[:, 0, :]  # First property (texture)
        oc_enhanced_raw = texture_oc_combined[:, 1, :]          # Second property (OC)
                
        # Apply gating and masking
        texture_oc_gated = self.texture_oc_gate([texture_features, texture_enhanced_raw_oc]) * oc_availability
        oc_texture_gated = self.texture_oc_gate([oc_features_safe, oc_enhanced_raw]) * oc_availability

        # Fuse features
        texture_enhanced_by_oc_raw = self.texture_oc_fusion(texture_oc_gated)
        oc_enhanced_by_texture_raw = self.texture_oc_fusion(oc_texture_gated)
        
        # When OC not available, fall back to original features
        texture_enhanced_by_oc = texture_features + (texture_enhanced_by_oc_raw - texture_features) * oc_availability
        oc_enhanced_by_texture = oc_features_safe + (oc_enhanced_by_texture_raw - oc_features_safe) * oc_availability
        
        # Store attention weights
        cross_attention_weights['texture_oc'] = texture_oc_weights

        # Store all cross-attention weights
        self._cross_attention_weights = cross_attention_weights

        # =================================================================
        # COMBINE ENHANCED FEATURES
        # =================================================================
        
        # Combine texture enhancements (instead of simple addition, use learned combination)
        combined_texture = texture_enhanced_by_bd + texture_enhanced_by_oc
        combined_texture = self.fusion_layer_norm(combined_texture)
        
        # Use enhanced versions of all properties for global attention
        enhanced_bd = bd_enhanced_by_texture
        enhanced_oc = oc_enhanced_by_texture
        
        # Get batch size
        batch_size = tf.shape(texture)[0]
    
        # Stack ENHANCED property features for global attention mechanism
        # Shape: [batch_size, 4, embedding_dim]
        property_features = tf.stack([
            combined_texture,    # Enhanced texture features
            enhanced_bd,         # Enhanced BD features  
            enhanced_oc,         # Enhanced OC features
            ksat_features        # Ksat features (no cross-attention yet)
        ], axis=1)
    
        # Create attention mask for properties
        # 1 for available properties, 0 for unavailable
        property_attention_mask = tf.expand_dims(
            tf.concat([
                tf.ones_like(properties_mask[:, 0:2]),  # Texture & BD always available
                tf.reshape(properties_mask[:, 2:4], [batch_size, 2])  # OC & Ksat availability
            ], axis=1),
            axis=1
        )
    
        # Apply property-level attention
        property_attention_output, property_weights = self.property_attention(
            query=property_features,
            key=property_features,
            value=property_features,
            attention_mask=property_attention_mask,
            return_attention_scores=True,
            training=training
        )
    
        # Store attention weights
        self._property_attention_weights = property_weights
    
        # Add & normalize
        property_features = self.property_layer_norm(property_features + property_attention_output)
    
        # Flatten and combine property features using enhanced fusion
        flattened_features = self.flatten(property_features)
        soil_embedding = self.property_fusion(flattened_features, training=training)
    
        # Process water potential points
        batch_size = tf.shape(water_potential)[0]
        num_points = tf.shape(water_potential)[1]
    
        # Reshape soil embedding for attention
        soil_embedding_expanded = tf.expand_dims(soil_embedding, axis=1)
    
        # Prepare water potential features
        wp_expanded = tf.expand_dims(water_potential, axis=-1)
    
        # Create embedding for water potential values
        wp_embedding = self.wp_embedding_layer(wp_expanded)
    
        # Create attention mask for water potential
        wp_attention_mask = tf.expand_dims(wp_mask, axis=1)
    
        # Apply water potential attention
        wp_attention_output, wp_weights = self.wp_attention(
            query=soil_embedding_expanded,
            key=wp_embedding,
            value=wp_embedding,
            attention_mask=wp_attention_mask,
            return_attention_scores=True,
            training=training
        )
    
        # Store attention weights
        self._wp_attention_weights = wp_weights
    
        # Add & normalize
        soil_wp_embedding = self.wp_layer_norm(soil_embedding_expanded + wp_attention_output)
        soil_wp_embedding = tf.squeeze(soil_wp_embedding, axis=1)
    
        # Expand soil_wp_embedding for each water potential point
        soil_expanded = tf.repeat(soil_wp_embedding, num_points, axis=0)
        
        # Flatten water potential
        wp_flat = tf.reshape(water_potential, [batch_size * num_points, 1])
        
        # Combine soil and water potential features
        combined = tf.concat([soil_expanded, wp_flat], axis=1)
        
        # Generate parameters using enhanced parameter network
        parameters = self.parameter_network(combined, training=training)
        
        # Apply monotonic function
        water_content_flat = self.monotonic_function([wp_flat, parameters])
        
        # Reshape back to original dimensions
        water_content = tf.reshape(water_content_flat, [batch_size, num_points])
        
        # Apply mask to zero out padded values
        return water_content * wp_mask
            
def get_attention_weights(self):
    """Return attention weights for interpretation"""
    return {
        'property_attention': self._property_attention_weights,
        'wp_attention': self._wp_attention_weights,
        'cross_attention': self._cross_attention_weights
    }