# 在 mmdet3d/utils/profiler_hook.py 添加以下内容
import os
import torch
from mmengine.hooks import Hook
from mmengine.registry import HOOKS
from mmengine.runner import Runner

@HOOKS.register_module()
class SimpleProfilerHook(Hook):
    """简化版性能分析钩子

    Args:
        interval (int): 执行分析的迭代间隔，默认50
        output_dir (str): 分析结果输出目录，默认'./profiler'
    """

    def __init__(self, interval: int = 50, output_dir: str = './profiler') -> None:
        self.interval = interval
        self.output_dir = output_dir
        self.profiler = None
        self.step_count = 0

    def after_train_iter(self,
                         runner: Runner,
                         batch_idx: int,
                         data_batch: dict = None,
                         outputs: dict = None) -> None:
        """在训练迭代后执行性能分析"""
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)

        # 首次迭代初始化分析器
        if self.profiler is None:
            self.profiler = torch.profiler.profile(
                activities=[
                    torch.profiler.ProfilerActivity.CPU,
                    torch.profiler.ProfilerActivity.CUDA
                ],
                schedule=torch.profiler.schedule(
                    wait=30,       # 跳过前1次迭代
                    warmup=1,     # 预热1次迭代
                    active=5,     # 分析3次迭代
                    repeat=1      # 只执行1轮
                ),
                on_trace_ready=torch.profiler.tensorboard_trace_handler(self.output_dir),
                record_shapes=True,
                profile_memory=False,
                with_stack=True
            )
            self.profiler.start()
            runner.logger.info(f"Profiler started. Results will be saved to {self.output_dir}")

        # 执行分析器步骤
        self.profiler.step()
        self.step_count += 1

        # 每间隔打印摘要
        if self.step_count == self.interval and self.profiler is not None:
            runner.logger.info("Profiler summary:")
            runner.logger.info(
                self.profiler.key_averages().table(
                    sort_by="cuda_time_total",
                    row_limit=-1
                )
            )
            # self.total_steps = 30 + 1 + 5
            # results = self.profiler.key_averages().table(sort_by="cuda_time_total")
            # log_file = "./BW_log_step{}.txt".format(self.total_steps)
            # with open(log_file, mode='w') as file:
            #     file.write(str(results))

            # 保存当前跟踪
            # trace_path = os.path.join(self.output_dir, f"trace_step_{self.step_count}.json")
            # self.profiler.export_chrome_trace(trace_path)
            # runner.logger.info(f"Trace saved to {trace_path}")
