from tensorflow.keras.layers import Conv2D, MaxPooling2D, GlobalAveragePooling2D, Dense, Dropout, LSTM, Reshape, Resizing, BatchNormalization, Flatten, TimeDistributed, Input
from tensorflow.keras import layers
from tensorflow.keras.applications import ResNet50V2,EfficientNetV2S, MobileNetV2, ConvNeXtTiny, ConvNeXtBase


# Strategy for TPU training
regularization_parameter=0.01
dropout_value=0.2
input_shape=(18,96,64,1)


with strategy.scope():

    # Load ResNet50V2 with pre-trained ImageNet weights and modify for your task
#     base_model = ResNet50V2(weights='imagenet', include_top=False, input_shape=(128,128, 3))
#     base_model.trainable=True
#     for layer in base_model.layers[:-10]:
#         layer.trainable=False
    base_model = hub.KerasLayer('https://www.kaggle.com/models/google/vggish/TensorFlow2/vggish/1')
#     base_model = EfficientNetV2M(weights='imagenet', include_top=False, input_shape=(256, 256, 3))
    
    # Add custom layers
    inputs = Input(shape=input_shape)
    
    # Apply TimeDistributed to wrap the base CNN
    x = TimeDistributed(base_model)(inputs)

    # Add pooling layer to reduce dimensionality for each time step
    x = TimeDistributed(GlobalAveragePooling2D())(x)
    # Now x has shape (batch_size, time_steps, features)
    
#     x = Flatten()(x)
    # Add LSTM Layer
#     x = LSTM(256, return_sequences=True)(x)  # First LSTM with 256 units
    x = LSTM(128, return_sequences=False)(x)  # LSTM with 64 units
    
    # Increase regularization with more dropout and L2
#     x = Dense(1024, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
# #     x = BatchNormalization()(x)
#     x = Dropout(dropout_value)(x)
    x = Dense(512, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
    x = Dropout(dropout_value)(x)
    x = Dense(256, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
    x = Dropout(dropout_value)(x)
    x = Dense(128, activation='relu', kernel_regularizer=l2(regularization_parameter))(x)
    x = Dropout(dropout_value)(x)

    # Output layers for age and gender classification"
    age_output = Dense(3, activation='softmax', name='age')(x)
    gender_output = Dense(1, activation='sigmoid', name='gender')(x)

    # Learning Rate Decay
    reduce_lr = tf.keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss',    # Metric to monitor
        factor=0.1,            # Factor by which to reduce the learning rate
        patience=3,            # Number of epochs with no improvement after which the learning rate is reduced
        min_lr=1e-7            # Lower bound on the learning rate
    )
    # Early Stopping 
    early_stopping = EarlyStopping(monitor='val_loss', patience=20, restore_best_weights=True)
    
    # Create the final model
    model = Model(inputs=inputs, outputs=[age_output, gender_output])

    # Compile the model with appropriate loss functions
    model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
                  loss={'age': 'categorical_crossentropy', 'gender': 'binary_crossentropy'},
                  metrics={'age': 'accuracy', 'gender': 'accuracy'}) ---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[15], line 26
     23 inputs = Input(shape=input_shape)
     25 # Apply TimeDistributed to wrap the base CNN
---> 26 x = TimeDistributed(base_model)(inputs)
     28 # Add pooling layer to reduce dimensionality for each time step
     29 x = TimeDistributed(GlobalAveragePooling2D())(x)

File /usr/local/lib/python3.10/site-packages/keras/src/layers/rnn/time_distributed.py:49, in TimeDistributed.__init__(self, layer, **kwargs)
     47 def __init__(self, layer, **kwargs):
     48     if not isinstance(layer, Layer):
---> 49         raise ValueError(
     50             "Please initialize `TimeDistributed` layer with a "
     51             f"`keras.layers.Layer` instance. Received: {layer}"
     52         )
     53     super().__init__(layer, **kwargs)
     54     self.supports_masking = True

ValueError: Please initialize `TimeDistributed` layer with a `keras.layers.Layer` instance. Received: <tensorflow_hub.keras_layer.KerasLayer object at 0x7ca64194ed70>