import torch


def update_train_state(args, model, train_state):
    """
    Handle the training states update.
    Components: Early stopping and model checkpoint saving.
    :param args:
    :param model:
    :param train_state:
    :return: a new train state
    """

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

    # Save model if performance improved
    if train_state["epoch_index"] >= 1:
        loss_tml, loss_t = train_state["dev_loss"][-2:]
        if loss_t >= loss_tml:
            train_state["early_stopping_step"] += 1
        else:
            if loss_t < train_state["early_stopping_best_val"]:
                torch.save(model.state_dict(), train_state["model_filename"])
                train_state["early_stopping_best_val"] = loss_t
            train_state["early_stopping_step"] = 0

        train_state["early_stopping"] = \
            train_state["early_stopping_step"] >= args.early_stopping_criteria

    return train_state