import torch


def add_dics(base, new, is_first):
    """
    Add a new dictionary to an old dictionary, where the values in each dictionary
    are lists that should be concatenated, and the keys in the new dictionary might
    not be the same as those in the old one.
    Args:
        base (dict): base dictionary to be added to
        new (dict): new dictionary adding on
        is_first (bool): whether this is the first batch we've loaded
    Returns:
        base (dict): updated base dictionary
    """

    for key, val in new.items():
        if is_first:
            base[key] = val
        else:
            if key in base:
                base[key] += val
    if is_first:
        return base

    # any keys that are new to the dictionary despite this not being
    # the first batch added (i.e. they're in this batch but weren't
    # in previous batches)
    extra_keys = [key for key in new.keys() if key not in
                  base.keys()]

    for key in extra_keys:
        dim = len(base["smiles"])
        old_val = [torch.tensor(float("nan"))
                   for _ in range(dim)]
        new_val = old_val + new[key]
        base[key] = new_val

    # same idea for keys that were here before and aren't now
    missing_keys = [key for key in base.keys() if key not in 
                    new.keys()]

    for key in missing_keys:
        dim = len(new["smiles"])
        base[key] += [torch.tensor(float('nan')) for _ in range(dim)]

    return base