import os
import joblib
from huggingface_hub import hf_hub_download
access_token_smpl = os.environ.get('HF_SMPL_TOKEN')

def get_smpl_models():
    REPO_ID = 'athn-nik/smpl_models'
    from huggingface_hub import snapshot_download
    print(access_token_smpl)
    # aaa = hf_hub_download(REPO_ID, filename="smplh/SMPLH_NEUTRAL.npz")
    # print(aaa)
    return snapshot_download(repo_id=REPO_ID, allow_patterns="smplh*",
                             token=access_token_smpl)

def download_models():
    REPO_ID = 'athn-nik/example-model'
    return hf_hub_download(REPO_ID, filename="tmed_compressed.ckpt")

def download_model_config():
    REPO_ID = 'athn-nik/example-model'
    path_to_config = hf_hub_download(REPO_ID, filename="tmed/.hydra/config.yaml")
    from omegaconf import OmegaConf
    model_cfg = OmegaConf.load(path_to_config)
    return model_cfg.model.input_feats


def download_motion_from_dataset(key_to_dl):
    REPO_ID = 'athn-nik/example-model'
    from huggingface_hub import snapshot_download
    keytodl = key_to_dl
    keytodl = '000008'
    path_for_ds = snapshot_download(repo_id=REPO_ID, 
                                           allow_patterns=f"dataset_inputs/{keytodl}",
                                           token=access_token_smpl)
    path_for_ds_sample = path_for_ds + f'/dataset_inputs/{keytodl}.pth.tar'
    return path_for_ds_sample

def download_tmr():
    REPO_ID = 'athn-nik/example-model'
    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
    from huggingface_hub import snapshot_download
    return snapshot_download(repo_id=REPO_ID, allow_patterns="tmr*",
                             token=access_token_smpl)

def download_motionfix():
    REPO_ID = 'athn-nik/example-model'
    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
    from huggingface_hub import snapshot_download
    return snapshot_download(repo_id=REPO_ID, allow_patterns="motionfix*",
                             token=access_token_smpl)

def download_motionfix_dataset():
    REPO_ID = 'athn-nik/example-model'
    dataset_downloaded_path = hf_hub_download(REPO_ID, filename="motionfix.pth.tar")
    dataset_dict = joblib.load(dataset_downloaded_path)
    return dataset_dict

def download_embeddings():
    REPO_ID = 'athn-nik/example-model'
    # return hf_hub_download(REPO_ID, filename="min_checkpoint.ckpt")
    from huggingface_hub import snapshot_download
    return snapshot_download(repo_id=REPO_ID, allow_patterns="embeddings*",
                             token=access_token_smpl)