import importlib
from typing import Dict, List, Tuple

import torch

def load_loss_weight(dataset_name: str) -> Tuple[List, List, List]:
    """Loss weight to balance the categories (obj, attr, aff)
    weight = k * -log(frequency)"""
    
    try:
        Weight = importlib.import_module('utils.aux_data.%s_weight'%dataset_name)
        obj_wgt  = torch.tensor(Weight.obj_weight, dtype=torch.float32)
        attr_wgt = torch.tensor(Weight.attr_weight, dtype=torch.float32)
        aff_wgt  = torch.tensor(Weight.aff_weight, dtype=torch.float32)
        return obj_wgt, attr_wgt, aff_wgt
        
    except ImportError:
        raise NotImplementedError("Loss weight for %s is not implemented yet"%dataset_name)


def load_wordvec_dict(dataset_name: str, vec_type: str) -> Tuple[Dict, Dict]:
    raise NotImplementedError()
    dsname_mapping = {
        "MITg": "MIT",
        "UTg": "UT",
    }
    if dataset_name in dsname_mapping:
        dataset_name = dsname_mapping[dataset_name]
        
    try:
        Wordvec = importlib.import_module('utils.aux_data.%s_%s'%(vec_type, dataset_name))
        return Wordvec.attrs_dict, Wordvec.objs_dict

    except ImportError:
        raise NotImplementedError("%s vector for %s is not ready yet"%(vec_type, dataset_name))
