import os
import orjson
import torch
import numpy as np
from model import TMR_textencoder

EMBS = "data/unit_motion_embs"


def load_json(path):
    with open(path, "rb") as ff:
        return orjson.loads(ff.read())


def load_keyids(split):
    path = os.path.join(EMBS, f"{split}.keyids")
    with open(path) as ff:
        keyids = np.array([x.strip() for x in ff.readlines()])
    return keyids


def load_keyids_splits(splits):
    return {
        split: load_keyids(split)
        for split in splits
    }


def load_unit_motion_embs(split, device):
    path = os.path.join(EMBS, f"{split}_motion_embs_unit.npy")
    tensor = torch.from_numpy(np.load(path)).to(device)
    return tensor


def load_unit_motion_embs_splits(splits, device):
    return {
        split: load_unit_motion_embs(split, device)
        for split in splits
    }


def load_model(device):
    text_params = {
        'latent_dim': 256, 'ff_size': 1024, 'num_layers': 6, 'num_heads': 4,
        'activation': 'gelu', 'modelpath': 'distilbert-base-uncased'
    }
    "unit_motion_embs"
    model = TMR_textencoder(**text_params)
    state_dict = torch.load("data/textencoder.pt", map_location=device)
    # load values for the transformer only
    model.load_state_dict(state_dict, strict=False)
    model = model.eval()
    return model