import scipy.io
import numpy as np
import json  # 引入json模块


def load(name, path):
    # 处理UTD-MHAD的JSON文件（你的单帧数据）
    if name == 'UTD_MHAD':
        # 判断文件是否为JSON格式（通过后缀）
        if path.endswith('.json'):
            with open(path, 'r') as f:
                data = json.load(f)  # 加载JSON列表
            # 转换为NumPy数组，确保形状为[关节数, 3]
            data_np = np.array(data)
            # 校验数据格式（防止错误）
            assert data_np.ndim == 2 and data_np.shape[1] == 3, \
                f"UTD-MHAD JSON格式错误，应为[关节数, 3]，实际为{data_np.shape}"
            # 若需要单帧维度（[1, 关节数, 3]），可扩展维度
            return data_np[np.newaxis, ...]  # 输出形状：[1, N, 3]（1表示单帧）
        
        # 保留原UTD_MHAD的.mat文件支持（如果还需要处理.mat数据）
        elif path.endswith('.mat'):
            arr = scipy.io.loadmat(path)['d_skel']
            new_arr = np.zeros([arr.shape[2], arr.shape[0], arr.shape[1]])
            for i in range(arr.shape[2]):
                for j in range(arr.shape[0]):
                    for k in range(arr.shape[1]):
                        new_arr[i][j][k] = arr[j][k][i]
            return new_arr
        
        else:
            raise ValueError(f"UTD-MHAD不支持的文件格式：{path}")

    # 其他数据集的原有逻辑保持不变
    elif name == 'HumanAct12':
        return np.load(path, allow_pickle=True)
    elif name == "CMU_Mocap":
        return np.load(path, allow_pickle=True)
    elif name == "Human3.6M":
        return np.load(path, allow_pickle=True)[0::5]  # 下采样
    elif name == "NTU":
        return np.load(path, allow_pickle=True)[0::2]
    elif name == "HAA4D":
        return np.load(path, allow_pickle=True)
    else:
        raise ValueError(f"不支持的数据集名称：{name}")