from tensorflow import keras

class FCNutr(keras.Model):
    def __init__(self, nutr_names, crop_size, hidden_dim, num_shared_hidden, use_dropout):
        super().__init__()
        
        self.cnn = keras.applications.InceptionV3(
            include_top=False, 
            weights="imagenet", 
            input_shape=crop_size + (3,), 
            pooling='avg'
        )

        shared_layers = []
        for i in range(num_shared_hidden):
            if use_dropout:
                shared_layers.append(keras.layers.Dropout(0.2 if i == 0 else 0.5))
            shared_layers.append(keras.layers.Dense(hidden_dim, activation="relu"))
            
        self.shared = keras.Sequential(shared_layers, name='shared')

        self.multitask_heads = []
        for name in nutr_names:
            head_layers = [
                keras.layers.Dense(hidden_dim, activation="relu"),
                keras.layers.Dense(1, activation="relu")
            ]

            if use_dropout:
                head_layers.insert(0, keras.layers.Dropout(0.5))
                head_layers.insert(2, keras.layers.Dropout(0.5))

            self.multitask_heads.append(keras.Sequential(head_layers, name=name))

    def call(self, inputs, training=False):
        x = self.cnn(inputs, training=training)
        x = self.shared(x, training=training)
        return {head.name: head(x, training=training) for head in self.multitask_heads}