import torch


def train_model(model, weight_sharing, aux_loss_lambda, optimizer, scheduler, early_stop_patience, X_train, y_train,
                X_val, y_val, mini_batch_size, epochs, criterion, verbose=0):
    """
    Train a model

    Parameters:
    model (Module): The model to train
    weight_sharing (boolean): Whether to use weight sharing
    aux_loss_lambda (float): The weight given to the auxiliary loss
    optimizer (Optimizer): The optimizer used to train the model
    scheduler (Scheduler): The scheduler used to adapt the learning rate
    early_stop_patience (int): The number of epochs we wait without improvement before stopping
    X_train (Tensor): The data used to train the model
    y_train (Tensor): The labels used to train the model
    X_val (Tensor): The data used to validate the model
    y_val (Tensor): The labels used to validate the model
    mini_batch_size (int): The size of each mini-batch
    epochs (int): The number of epochs
    criterion (Criterion): The criterion used to compute the losses
    versbose (int): Used to define the quantity of logs to print while training

    Returns:
    list(float): All the train losses
    list(float): All the validation losses
    """

    all_train_loss = []
    all_val_loss = []
    for e in range(epochs):
        model.train()
        train_loss = 0
        for b in range(0, X_train.size(0), mini_batch_size):
            output_target, output_class_l, output_class_r = model(X_train.narrow(0, b, mini_batch_size),
                                                                  weight_sharing)  # Forward pass
            loss_target = criterion(output_target,
                                    y_train.narrow(0, b, mini_batch_size)[:, 0])  # The loss wrt. the binary target
            loss_class_l = criterion(output_class_l,
                                     y_train.narrow(0, b, mini_batch_size)[:, 1])  # The loss wrt. the left image number
            loss_class_r = criterion(output_class_r, y_train.narrow(0, b, mini_batch_size)[:, 2])  # The loss wrt. the right image number

            # We compute the total loss (no auxiliary loss if aux_loss_lambda is 0)
            # If we set aux_loss_lambda to a very big value, it will basicaly shadow the target_loss (tried during validation)
            loss = loss_target + aux_loss_lambda * 0.5 * (loss_class_l + loss_class_r)

            train_loss += loss.item()
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        # For the validation, we change the mode of the model to eval()
        model.eval()
        val_loss = 0
        with torch.no_grad():
            for b in range(0, X_val.size(0), mini_batch_size):
                output_target, output_class_l, output_class_r = model(X_val.narrow(0, b, mini_batch_size),
                                                                      weight_sharing)
                loss_target = criterion(output_target, y_val.narrow(0, b, mini_batch_size)[:, 0])
                loss_class_l = criterion(output_class_l, y_val.narrow(0, b, mini_batch_size)[:, 1])
                loss_class_r = criterion(output_class_r, y_val.narrow(0, b, mini_batch_size)[:, 2])

                loss = loss_target + aux_loss_lambda * 0.5 * (loss_class_l + loss_class_r)

                val_loss += loss.item()

            # If the validation loss has stoped improving, we lower the learning rate
            scheduler.step(val_loss)

            # We compute the mean loss by batch
            train_loss /= (X_train.size(0) // mini_batch_size)
            val_loss /= (X_val.size(0) // mini_batch_size)

            # We save all the training and validation losses
            all_train_loss.append(train_loss)
            all_val_loss.append(val_loss)

        if verbose == 1 and e % 10 == 0:
            print('{} {}'.format(e, train_loss, val_loss))

        # If we don't have improved the validation loss since a moment, we stop the training (to mitigate the overfitting)
        if len(all_val_loss) - (torch.Tensor(all_val_loss).argmin() + 1) > early_stop_patience:
            if verbose == 1:
                print('Early stopping')
            break

    return all_train_loss, all_val_loss