from tensorflow import keras
from tensorflow.keras import layers
from keras.layers import Flatten
from keras.layers import Input, Dense, Layer, Dropout, Conv2D, MaxPooling2D, Flatten, multiply
from MILAttentionLayer import MILAttentionLayer

def SimpleModel(instance_shape,bag_size):
    """ Create Keras model for Multiply Instance Learning
    Parameters
    -------------------
    instance_shape (tuple) - shape of 1 instance in the bag
    bag_size (int) - size of the bag
    Returns
    -------------------
     keras.Model
    """
    # Extract features from inputs.
    inputs, embeddings = [], []
    conv1_1 = Conv2D(16, kernel_size=(2,2), activation='relu') 
    conv1_2 = Conv2D(16, kernel_size=(2,2), activation='relu')  
    mpool_1 = MaxPooling2D((2,2))

    conv2_1 = Conv2D(32, kernel_size=(2,2),   activation='relu')  
    conv2_2 = Conv2D(32, kernel_size=(2,2),activation='relu') 
    mpool_2 = MaxPooling2D((2,2))

    fc0 = Dense(512, activation='relu', name='fc0') 
    fc1 = Dense(512, activation='relu', name='fc1') 
    fc2 = Dense(256, activation= 'relu',  name='fc2')
    
    for _ in range(bag_size):
        inp = layers.Input(instance_shape)
        inputs.append(inp)
        x = conv1_1(inp)
        x = conv1_2(x)
        x = mpool_1(x)

        x = conv2_1(x)
        x = conv2_2(x)
        x = mpool_2(x)

        x = Flatten()(x)
        x = fc0(x)
        x = Dropout(0.5)(x)
        x = fc1(x)
        x = Dropout(0.5)(x)
        x = fc2(x)
        x = Dropout(0.2)(x)
        
        embeddings.append(x)

    # Аttention layer.
    alpha = MILAttentionLayer(
        weight_params_dim=1024,
        kernel_regularizer=keras.regularizers.l2(0),# previous - 0.01
        use_gated=True, 
        name="alpha",
    )(embeddings)

    # Multiply attention weights with the input layers.
    multiply_layers = [
        layers.multiply([alpha[i], embeddings[i]]) for i in range(len(alpha))
    ]

    # Concatenate layers.
    concat = layers.concatenate(multiply_layers, axis=1)

    # Classification output node.
    output = layers.Dense(2, activation = 'softmax')(concat)

    return keras.Model(inputs, output) 

import tensorflow as tf
from CustomDataGenerator import CustomDataGenerator
from SimpleModel import SimpleModel

def train_model (train_df, validation_df, model_save_path):
    """
    Train SimpleModel
    
    Parameters
    -------------------
    train_df (pandas DataFrame) - DataFrame with the training data. X (bag of instances) - list of images paths. y -label
    validation_df (pandas DataFrame) - DataFrame with the validation data. X (bag of instances) - list of images paths. y -label
    model_save_path (str) - path for model saving
    Returns
    -------------------
    """
       
    # create generator of the training and validation data
    train_generator = CustomDataGenerator(df = train_df, shuffle = True, augmentations = True )
    validation_generator = CustomDataGenerator (df = validation_df, shuffle = False, augmentations = False )
    
    # Callbacks
    model_checkpoint = tf.keras.callbacks.ModelCheckpoint(
        model_save_path,
        monitor="val_loss",
        verbose=1,
        mode="min",
        save_best_only=True,
        save_weights_only= False)
    
    
    es = tf.keras.callbacks.EarlyStopping(
        monitor="val_loss",
        patience=10,
        verbose=1,
        mode="min")
    
    # optimizer
    opt = tf.keras.optimizers.Adam(learning_rate=1e-3, decay=0.0005, beta_1=0.9, beta_2=0.999)
    
    # create and compile model
    model = SimpleModel(bag_size = 40, instance_shape = (256, 256, 3) )
    model.compile(optimizer = opt, 
    loss='categorical_crossentropy', metrics=["accuracy",tf.keras.metrics.AUC(name = 'AUC'),
                                                        tf.keras.metrics.AUC(curve = 'PR',name = 'PR_AUC'), 
                                                        tf.keras.metrics.Precision(name = 'Precision', class_id = 1),
                                                        tf.keras.metrics.Recall(name = 'Recall',class_id = 1)])
    # model fitting
    model.fit(
        train_generator,
        validation_data = validation_generator ,
        epochs=100,
        batch_size= 1,
        callbacks=[model_checkpoint,es], 
        verbose=1)