import torch


def update_train_state(args, model, train_state):
    """ Handle the training state updates.

    Components:
    - Early Stopping: Prevent overfitting.
    - Model Checkpoint: Model is saved if the model is better

    :param args: main arguments
    :param model: model to train
    :param train_state: a dictionary representing the training state values
    :returns: a new train_state
    """

    # Save one model at least
    if train_state["epoch_index"] == 0:
        torch.save(model.state_dict(), train_state["model_filename"])
        train_state["stop_early"] = False

    # Save model if performance improved
    elif train_state["epoch_index"] >= 1:
        loss_tm1, loss_t = train_state["val_loss"][-2:]

        # If loss worsened
        if loss_t >= train_state["early_stopping_best_val"]:
            # Update step
            train_state["early_stopping_step"] += 1
        # Loss decreased
        else:
            # Save the best model
            if loss_t < train_state["early_stopping_best_val"]:
                torch.save(model.state_dict(), train_state["model_filename"])

            # Rest early stopping step
            train_state["early_stopping_step"] = 0

            # Update early_stopping_best_val
            train_state["early_stopping_best_val"] = loss_t

        # Stop early ?
        train_state["stop_early"] = \
            train_state["early_stopping_step"] >= args.early_stopping_criteria

    return train_state