import torch


def compute_args(args):
    # DataLoader
    if not hasattr(args, "dataset"):  # fix for previous version
        args.dataset = "MOSEI"

    if args.dataset == "MOSEI":
        args.dataloader = "Mosei_Dataset"
    if args.dataset == "MELD":
        args.dataloader = "Meld_Dataset"

    # Loss function to use
    if args.dataset == "MOSEI" and args.task == "sentiment":
        args.loss_fn = torch.nn.CrossEntropyLoss(reduction="sum")
    if args.dataset == "MOSEI" and args.task == "emotion":
        args.loss_fn = torch.nn.BCEWithLogitsLoss(reduction="sum")
    if args.dataset == "MELD":
        args.loss_fn = torch.nn.CrossEntropyLoss(reduction="sum")

    # Answer size
    if args.dataset == "MOSEI" and args.task == "sentiment":
        args.ans_size = 7
    if args.dataset == "MOSEI" and args.task == "sentiment" and args.task_binary:
        args.ans_size = 2
    if args.dataset == "MOSEI" and args.task == "emotion":
        args.ans_size = 6
    if args.dataset == "MELD" and args.task == "emotion":
        args.ans_size = 7
    if args.dataset == "MELD" and args.task == "sentiment":
        args.ans_size = 3

    if args.dataset == "MOSEI":
        args.pred_func = "amax"
    if args.dataset == "MOSEI" and args.task == "emotion":
        args.pred_func = "multi_label"
    if args.dataset == "MELD":
        args.pred_func = "amax"

    return args