import ast, subprocess, os

# 从文件中获取特定变量信息，返回变量直和类型
def get_var_from_file(filename, var_name="norm_cfg"):
    with open(filename, 'r', encoding='utf-8') as file:
        code = file.read()

    # 解析文件的 AST 树
    tree = ast.parse(code)

    # 遍历 AST 树，查找指定变量
    for node in ast.walk(tree):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                if isinstance(target, ast.Name) and target.id == var_name:
                    # 将 AST 节点转换为 Python 对象
                    var_value = ast.literal_eval(node.value)
                    # 返回变量的值和类型
                    return var_value, type(var_value)

    # 如果没有找到指定的变量，返回 None 和 None
    return None, None

def get_var_from_py_file(file_path, var_name="auxiliary_head"):
    # 定义一个字典用于保存文件中的变量
    context = {}
    
    # 读取并执行文件
    with open(file_path, 'r') as file:
        exec(file.read(), context)
    
    # 获取 auxiliary_head 变量
    if var_name in context:
        return context[var_name]
    else:
        raise AttributeError(f"文件中没有找到 {var_name} 变量")

# 更新list的dict变量
def update_list_dict_var(var, var_new):
    # 判断 var 和 var_new 的长度是否相等
    if len(var) != len(var_new):
        raise ValueError(f"var {len(var)} 和 var_new {len(var_new)} 的大小不相等，无法更新")
    
    # 遍历 var_new，按索引更新 var
    for i, new_entry in enumerate(var_new):
        # 将 new_entry 中的键值覆盖 var[i] 中的相同键
        var[i].update(new_entry)
    return var

if __name__ == '__main__':
    # 示例调用
    var_value, var_type = get_var_from_file('./configs/_base_/models', 'norm_cfg')
    if var_value is not None:
        print(f"变量名: norm_cfg\n值: {var_value}\n类型: {var_type}")
    else:
        print("未找到指定变量")

# 将文件以dict格式输出
def format_dict(d, indent_level=1):
    formatted_items = []
    indent = '    ' * indent_level  # 根据缩进级别生成空格

    for key, value in d.items():
        if isinstance(value, dict):
            # 如果值是字典，递归调用 format_dict_as_func 并增加缩进级别
            formatted_value = format_dict(value, indent_level + 1)
            formatted_items.append(f"{indent}{key}={formatted_value}")
        elif isinstance(value, str):
            # 如果是字符串，格式化时加引号
            formatted_items.append(f"{indent}{key}='{value}'")
        else:
            # 其他类型（如数值等）不加引号
            formatted_items.append(f"{indent}{key}={value}")

    # 将所有键值对合成为多行的格式，并加上结尾逗号
    formatted_str = ",\n".join(formatted_items)
    
    # 返回更美观的 dict() 的格式，保留缩进和换行
    return f"dict(\n{formatted_str},\n{'    ' * (indent_level - 1)})"

# 所有格式正确输出
def format_all_data_old(data, indent_level=0):
    indent = '    ' * indent_level  # 根据缩进级别生成空格
    if isinstance(data, dict):
        formatted_items = []
        for key, value in data.items():
            formatted_value = format_all_data(value, indent_level + 1)
            formatted_items.append(f"{indent}    {key}={formatted_value}")
        formatted_str = ",\n".join(formatted_items)
        return f"dict(\n{formatted_str},\n{indent})"
    
    elif isinstance(data, list):
        formatted_items = [f"{indent}    {format_all_data(item, indent_level + 1)}" for item in data]
        formatted_str = ",\n".join(formatted_items)
        return f"[\n{formatted_str},\n{indent}]"

    elif isinstance(data, tuple):
        return f"{data}"

    elif isinstance(data, str):
        # 如果字符串包含单引号，则使用双引号，否则使用单引号
        if "'" in data:
            return f'"{data}"'
        else:
            return f"'{data}'"
    
    else:
        return str(data)

# 所有格式正确输出，加入键值中带有"."的处理
def format_all_data(data, indent_level=0):
    indent = '    ' * indent_level  # 根据缩进级别生成空格
    dot_key_items = []  # 用于存储带点号的键
    regular_items = []  # 用于存储普通键

    if isinstance(data, dict):
        for key, value in data.items():
            formatted_value = format_all_data(value, indent_level + 1)
            # 如果键包含点号，收集到 dot_key_items 中
            if '.' in key:
                dot_key_items.append(f"('{key}', {formatted_value})")
            else:
                regular_items.append(f"{indent}    {key}={formatted_value}")
        
        # 生成带点号键的部分，并放在最前面
        if dot_key_items:
            dot_key_str = f"{indent}    [{', '.join(dot_key_items)}]"
        else:
            dot_key_str = ""

        # 生成普通键的部分
        regular_str = ",\n".join(regular_items)

        # 返回最终的组合字符串，带点号键放在普通键前面
        if dot_key_str:
            return f"dict(\n{dot_key_str},\n{regular_str},\n{indent})"
        else:
            return f"dict(\n{regular_str},\n{indent})"
    
    elif isinstance(data, list):
        formatted_items = [f"{indent}    {format_all_data(item, indent_level + 1)}" for item in data]
        formatted_str = ",\n".join(formatted_items)
        return f"[\n{formatted_str},\n{indent}]"

    elif isinstance(data, tuple):
        return f"{data}"

    elif isinstance(data, str):
        # 如果字符串包含单引号，则使用双引号，否则使用单引号
        if "'" in data:
            return f'"{data}"'
        else:
            return f"'{data}'"
    
    else:
        return str(data)
    
# # 批量将参数内容写入文件 # V1 传统版
# def write_config_to_file(output_file, **kwargs):
#     """
#     将传入的任意数量的参数写入指定文件，并格式化输出。
    
#     :param output_file: 要写入的文件路径
#     :param kwargs: 任意数量的配置项，格式为 key=value
#     """
#     with open(output_file, 'w', encoding='utf-8') as file:
#         # 遍历 kwargs，将每个 key, value 写入文件
#         for key, value in kwargs.items():
#             file.write(f"{key} = {format_all_data(value)}\n\n")
    
#     # 打印成功信息
#     print(f"\033[93m{output_file} file generated successfully\033[0m")

# 批量将参数内容写入文件 # V2 加入训练过程可视化 TODO
def format_all_data(value):
    """
    一个辅助函数，用于将 Python 对象格式化为适合写入配置文件的字符串。
    使用 repr() 可以确保字符串、列表、字典等都保持其 Python 语法格式。
    """
    return repr(value)
def write_config_to_file(output_file, **kwargs):
    """
    将传入的任意数量的参数以及一个自动生成的 visualizer 配置写入指定文件。

    :param output_file: 要写入的文件路径 (e.g., 'configs/exp1.py')
    :param kwargs: 任意数量的配置项，格式为 key=value
    """
    # --- 1. 从 output_file 路径中提取实验名称 ---
    # 首先获取基本文件名 (e.g., 'exp1.py')
    base_name = os.path.basename(output_file)
    # 然后去掉文件扩展名，得到纯净的实验名 (e.g., 'exp1')
    experiment_name, _ = os.path.splitext(base_name)

    # --- 2. 构建 visualizer 配置字典 ---
    vis_backends = [
        dict(type='LocalVisBackend'),
        dict(type='TensorboardVisBackend'),
        dict(
            type='WandbVisBackend',
            init_kwargs=dict(
                project='Seg_MMSeg_Test',  # 你的 wandb 项目名称
                name=experiment_name  # 使用上面提取的文件名作为实验名
            )
        )
    ]
    visualizer = dict(
        name='visualizer',
        type='SegLocalVisualizer',
        vis_backends=vis_backends
    )

    # --- 3. 将自动生成的 visualizer 添加到要写入的内容中 ---
    # 如果 kwargs 中已经有 'visualizer'，它将会被新的配置覆盖
    kwargs['vis_backends'] = vis_backends
    kwargs['visualizer'] = visualizer

    # --- 4. 将所有配置项写入文件 ---
    with open(output_file, 'w', encoding='utf-8') as file:
        # 遍历所有配置项（包括新加入的 visualizer），写入文件
        for key, value in kwargs.items():
            # 使用 format_all_data 保证格式正确
            file.write(f"{key} = {format_all_data(value)}\n\n")

    # 打印成功信息
    print(f"\033[93mConfiguration saved to '{output_file}' successfully.\033[0m")
    print(f"\033[96mWandB experiment name will be: '{experiment_name}'\033[0m")

# 获取系统中所有 GPU 的可用显存信息。
def get_gpu_info():
    """
    获取系统中所有 GPU 的可用显存信息。
    :return: GPU 显存信息的列表，列表中的每一项是一个 (GPU编号, 剩余显存) 元组
    """
    try:
        # 使用 nvidia-smi 命令获取 GPU 显存信息
        result = subprocess.run(
            ['nvidia-smi', '--query-gpu=index,memory.free', '--format=csv,noheader,nounits'],
            stdout=subprocess.PIPE,
            encoding='utf-8'
        )
        
        # 解析结果
        gpu_info = []
        lines = result.stdout.strip().split('\n')
        for line in lines:
            index, memory_free = line.split(',')
            gpu_info.append((int(index.strip()), int(memory_free.strip())))
        
        return gpu_info
    except FileNotFoundError:
        print("\033[91mError: nvidia-smi 命令未找到，请确保 NVIDIA 驱动正确安装。\033[0m")
        return []

# 批量生成dict
def create_dict_by_kwargs(**kwargs):
    # 批量生成字典，kwargs 会自动收集所有传入的命名参数
    return {key: value for key, value in kwargs.items() if value is not None}

if __name__ == '__main__':
    # 示例字典
    my_dict = {'pretrained': './My_Local_Model/open_mmlab/resnet50_v1c.pth'}

    # 打印成 dict() 的格式
    print(format_dict(my_dict))
