import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
from tensorflow.keras.applications import ResNet50V2, DenseNet121

# Import attention and decoder from scratch model for consistency
from model_scratch import BahdanauAttention, RNN_Decoder


class CNN_Encoder_Pretrained(keras.Model):

    def __init__(self, embedding_dim=256, backbone='resnet50v2', fine_tune=False, input_shape=(224, 224, 3)):
        super(CNN_Encoder_Pretrained, self).__init__()
        
        self.embedding_dim = embedding_dim
        self.backbone_name = backbone.lower()
        self.fine_tune = fine_tune
        self.input_shape = input_shape
        
        # Load pretrained backbone
        if self.backbone_name == 'resnet50v2':
            # ResNet50V2 with ImageNet weights
            # Input: [batch, 224, 224, 3]
            # Output at conv4_block6_out: [batch, 14, 14, 1024] (for 224x224 input)
            self.backbone = ResNet50V2(
                include_top=False,
                weights='imagenet',
                input_shape=input_shape,
                name='resnet50v2_backbone'
            )
            
            # Extract features from intermediate layer (conv4_block6_out)
            # This preserves spatial information better than global pooling
            # Layer name: 'conv4_block6_out'
            self.feature_layer_name = 'conv4_block6_out'
            
            # Expected feature dimensions at conv4_block6_out
            # For 224x224 input: [batch, 14, 14, 1024]
            # For 299x299 input: [batch, 19, 19, 1024]
            self.feature_channels = 1024
            
        elif self.backbone_name == 'densenet121':
            # DenseNet121 with ImageNet weights
            # Input: [batch, 224, 224, 3]
            # Output at conv4_block (last dense block): [batch, 7, 7, 1024] (for 224x224 input)
            self.backbone = DenseNet121(
                include_top=False,
                weights='imagenet',
                input_shape=input_shape,
                name='densenet121_backbone'
            )
            
            # Extract features from last dense block
            # Layer name: 'conv4_block (last dense block)'
            self.feature_layer_name = 'conv4_block'  # Last dense block
            self.feature_channels = 1024
            
        else:
            raise ValueError(f"Unsupported backbone: {backbone}. Choose 'resnet50v2' or 'densenet121'")
        
        # Freeze or unfreeze backbone layers based on fine_tune flag
        self.backbone.trainable = fine_tune
        if not fine_tune:
            # Freeze all backbone layers
            for layer in self.backbone.layers:
                layer.trainable = False
            print(f"Backbone '{self.backbone_name}' frozen (transfer learning mode)")
        else:
            # Fine-tuning: unfreeze top layers
            # For ResNet50V2: unfreeze last 2 blocks
            # For DenseNet121: unfreeze last dense block
            if self.backbone_name == 'resnet50v2':
                # Unfreeze conv5 blocks (last 3 blocks)
                for layer in self.backbone.layers[-30:]:
                    layer.trainable = True
            elif self.backbone_name == 'densenet121':
                # Unfreeze last dense block
                for layer in self.backbone.layers[-20:]:
                    layer.trainable = True
            print(f"Backbone '{self.backbone_name}' fine-tuning enabled (top layers trainable)")
        
        # Create a model that outputs features from intermediate layer
        # This allows us to extract spatial feature maps instead of global pooled features
        if self.backbone_name == 'resnet50v2':
            # ResNet50V2: Extract features from the last convolutional block
            # The backbone output is already the feature map before global pooling
            # For ResNet50V2 with 224x224 input: output shape is [batch, 7, 7, 2048]
            # But we want conv4_block6_out which is [batch, 14, 14, 1024]
            # Try to get a layer from the 4th block
            try:
                # Try to get conv4_block6_out layer
                feature_layer = self.backbone.get_layer('conv4_block6_out')
            except:
                try:
                    # Alternative: get the last layer of conv4 block
                    # Find the layer that outputs features before conv5
                    for i, layer in enumerate(self.backbone.layers):
                        if 'conv4_block6' in layer.name:
                            feature_layer = layer
                            break
                    else:
                        # Fallback: use the backbone output directly
                        # This will be [batch, 7, 7, 2048] for 224x224 input
                        feature_layer = self.backbone.layers[-1]
                except:
                    # Final fallback: use backbone output
                    feature_layer = self.backbone.layers[-1]
            
            # Create intermediate model
            self.feature_extractor = keras.Model(
                inputs=self.backbone.input,
                outputs=feature_layer.output,
                name='feature_extractor'
            )
            
            # Update feature_channels based on actual output
            dummy_input = tf.zeros((1, *input_shape))
            dummy_output = self.feature_extractor(dummy_input)
            self.feature_channels = dummy_output.shape[-1]
            
        elif self.backbone_name == 'densenet121':
            # DenseNet121: Extract features from the last dense block
            # The backbone output is the feature map: [batch, 7, 7, 1024] for 224x224 input
            # Use the backbone output directly (it's already before global pooling)
            self.feature_extractor = keras.Model(
                inputs=self.backbone.input,
                outputs=self.backbone.output,
                name='feature_extractor'
            )
            
            # Update feature_channels
            dummy_input = tf.zeros((1, *input_shape))
            dummy_output = self.feature_extractor(dummy_input)
            self.feature_channels = dummy_output.shape[-1]
        
        # Reshape layer to convert spatial feature maps to sequence format
        # Input: [batch, height, width, channels]
        # Output: [batch, height*width, channels]
        # The exact dimensions depend on input size and backbone
        # We'll compute this dynamically in the call method
        
        # Dense layer to project features to embedding dimension
        # Input: [batch, spatial_features, feature_channels]
        # Output: [batch, spatial_features, embedding_dim]
        self.fc = layers.Dense(embedding_dim, activation='relu', name='fc_projection')
    
    def call(self, x, training=False):

        features = self.feature_extractor(x, training=training)
        # features shape: [batch, height, width, feature_channels]
        # Example: [batch, 14, 14, 1024]
        
        # Get spatial dimensions
        batch_size = tf.shape(features)[0]
        height = tf.shape(features)[1]
        width = tf.shape(features)[2]
        channels = features.shape[3]  # Static shape: feature_channels
        
        # Reshape to sequence format for attention mechanism
        # [batch, height, width, channels] → [batch, height*width, channels]
        # Example: [batch, 14, 14, 1024] → [batch, 196, 1024]
        features = tf.reshape(features, (batch_size, height * width, channels))

        features = self.fc(features)
        # features shape: [batch, spatial_features, embedding_dim]
        # Example: [batch, 196, 256]
        
        return features
    
    def get_feature_shape(self, input_shape):

        # Create dummy input
        dummy_input = tf.zeros((1, *input_shape))
        
        # Forward pass
        features = self.call(dummy_input, training=False)
        
        return features.shape


# Re-export decoder and attention for convenience
# These are imported from model_scratch to ensure identical architecture
__all__ = ['CNN_Encoder_Pretrained', 'BahdanauAttention', 'RNN_Decoder']


# Example usage and testing
if __name__ == "__main__":
    # Test model instantiation and forward pass
    print("Testing Transfer Learning Encoder-Decoder Model...")
    
    # Model hyperparameters
    EMBEDDING_DIM = 256
    UNITS = 512
    VOCAB_SIZE = 10000
    BATCH_SIZE = 4
    
    # Test ResNet50V2 encoder
    print("\n=== Testing ResNet50V2 Encoder ===")
    encoder_resnet = CNN_Encoder_Pretrained(
        embedding_dim=EMBEDDING_DIM,
        backbone='resnet50v2',
        fine_tune=False,
        input_shape=(224, 224, 3)
    )
    print(f"ResNet50V2 encoder created with embedding_dim={EMBEDDING_DIM}")
    
    # Test forward pass
    dummy_img_resnet = tf.random.normal((BATCH_SIZE, 224, 224, 3))
    print(f"Input image shape: {dummy_img_resnet.shape}")
    
    features_resnet = encoder_resnet(dummy_img_resnet, training=False)
    print(f"Encoder output (features) shape: {features_resnet.shape}")
    
    # Test DenseNet121 encoder
    print("\n=== Testing DenseNet121 Encoder ===")
    encoder_densenet = CNN_Encoder_Pretrained(
        embedding_dim=EMBEDDING_DIM,
        backbone='densenet121',
        fine_tune=False,
        input_shape=(224, 224, 3)
    )
    print(f"DenseNet121 encoder created with embedding_dim={EMBEDDING_DIM}")
    
    # Test forward pass
    dummy_img_densenet = tf.random.normal((BATCH_SIZE, 224, 224, 3))
    print(f"Input image shape: {dummy_img_densenet.shape}")
    
    features_densenet = encoder_densenet(dummy_img_densenet, training=False)
    print(f"Encoder output (features) shape: {features_densenet.shape}")
    
    # Test decoder compatibility
    print("\n=== Testing Decoder Compatibility ===")
    decoder = RNN_Decoder(
        embedding_dim=EMBEDDING_DIM,
        units=UNITS,
        vocab_size=VOCAB_SIZE,
        rnn_type='lstm'
    )
    print(f"Decoder created with units={UNITS}, vocab_size={VOCAB_SIZE}")
    
    # Test with ResNet features
    decoder_states = decoder.reset_state(batch_size=BATCH_SIZE)
    hidden = decoder_states[0]
    carry = decoder_states[1] if len(decoder_states) > 1 else None
    
    dummy_token = tf.constant([[1], [2], [3], [4]])
    predictions, new_hidden, new_carry, attn_weights = decoder(
        dummy_token, features_resnet, hidden, carry, training=False
    )
    print(f"Decoder predictions shape: {predictions.shape}")
    print(f"Decoder attention weights shape: {attn_weights.shape}")
    
    # Test fine-tuning mode
    print("\n=== Testing Fine-tuning Mode ===")
    encoder_finetune = CNN_Encoder_Pretrained(
        embedding_dim=EMBEDDING_DIM,
        backbone='resnet50v2',
        fine_tune=True,
        input_shape=(224, 224, 3)
    )
    
    trainable_count = sum([tf.size(w).numpy() for w in encoder_finetune.trainable_weights])
    total_count = sum([tf.size(w).numpy() for w in encoder_finetune.weights])
    print(f"Fine-tuning mode: {trainable_count}/{total_count} parameters trainable")
    
    print("\n✅ Model test completed successfully!")