import os
import json
import numpy as np

def find_subfolders_and_pth_files(root_folder):
    """
    查找指定文件夹下的所有子文件夹（不包括下一级的子文件夹），
    以及子文件夹下所有后缀为 '.pth' 的文件，并将它们对应存储。
    同时，如果根文件夹下也有 '.pth' 文件，也会将其添加到结果中。
    结果会根据子文件夹的名字进行排序。

    :param root_folder: 要查找的根文件夹路径
    :return: 一个字典，键为文件夹路径，值为该文件夹下所有 '.pth' 文件的路径列表
    """
    result = {}

    # 首先检查根文件夹下的 .pth 文件
    root_pth_files = []
    for file in os.listdir(root_folder):
        file_path = os.path.join(root_folder, file)
        if os.path.isfile(file_path) and file.endswith('.pth'):  # 确保是文件且后缀为 '.pth'
            root_pth_files.append(file_path)

    # 如果根文件夹下有 .pth 文件，则将其添加到结果中
    if root_pth_files:
        result[root_folder] = root_pth_files

    # 遍历根文件夹下的所有子文件夹
    for subfolder in sorted(os.listdir(root_folder)):  # 按子文件夹名字排序
        subfolder_path = os.path.join(root_folder, subfolder)
        if os.path.isdir(subfolder_path):  # 确保是文件夹
            pth_files = []

            # 遍历子文件夹下的所有文件
            for file in os.listdir(subfolder_path):
                file_path = os.path.join(subfolder_path, file)
                if os.path.isfile(file_path) and file.endswith('.pth'):  # 确保是文件且后缀为 '.pth'
                    pth_files.append(file)

            # 如果子文件夹下有 .pth 文件，则将其添加到结果中
            if pth_files:
                result[subfolder_path] = pth_files

    return result

def save_dict_to_json(dictionary, file_path, file_name):
    """
    将字典保存为 JSON 文件。

    :param dictionary: 要保存的字典
    :param file_path: 文件保存路径
    :param file_name: 文件名（不需要后缀）
    """
    # 确保路径存在
    os.makedirs(file_path, exist_ok=True)

    # 拼接完整文件路径
    full_path = os.path.join(file_path, f"{file_name}.json")

    # 将字典保存为 JSON 文件
    with open(full_path, 'w', encoding='utf-8') as f:
        json.dump(dictionary, f, indent=4, ensure_ascii=False)

    print(f"字典已保存到: {full_path}")

def save_list_of_dicts_to_json(list_of_dicts, file_path, file_name="all_result_under_this_path"):
    """
    将一个包含多个字典的列表保存为 JSON 文件。

    :param list_of_dicts: 包含多个字典的列表
    :param file_path: 文件保存路径
    :param file_name: 文件名（不需要后缀）
    """
    # 确保路径存在
    os.makedirs(file_path, exist_ok=True)

    # 拼接完整文件路径
    full_path = os.path.join(file_path, f"{file_name}.json")

    list_of_dicts_tmp=convert_values_to_float(list_of_dicts)
    # 将列表保存为 JSON 文件
    with open(full_path, 'w', encoding='utf-8') as f:
        json.dump(list_of_dicts_tmp, f, indent=4, ensure_ascii=False)

    print(f"列表已保存到: {full_path}")

def convert_values_to_float(obj):
    """
    递归地将字典或列表中的所有数值类型转换为 float 格式。
    支持处理 ndarray，将其转换为列表。

    :param obj: 输入的字典、列表或其他数据类型
    :return: 转换后的字典或列表
    """
    if isinstance(obj, dict):
        # 如果是字典，递归处理每个值
        return {key: convert_values_to_float(value) for key, value in obj.items()}
    elif isinstance(obj, list):
        # 如果是列表，递归处理每个元素
        return [convert_values_to_float(item) for item in obj]
    elif isinstance(obj, np.ndarray):
        # 如果是 ndarray，转换为列表并递归处理
        return convert_values_to_float(obj.tolist())
    elif isinstance(obj, (int, float, np.float32, np.float64, np.int32, np.int64)):
        # 如果是数值类型，转换为 float
        return float(obj)
    else:
        # 其他类型（如字符串）保持不变
        return obj



