import os
import torch
import torch.distributed as dist
import argparse
from utils.autodataloader_v3 import get_mesh_dataloader

def count_dataset(args):
    # 1. 初始化分布式环境
    # 使用环境变量自动获取 local_rank
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    torch.cuda.set_device(local_rank)
    dist.init_process_group(backend="nccl")
    rank = dist.get_rank()
    world_size = dist.get_world_size()

    # 2. 初始化数据加载器
    # 确保 get_mesh_dataloader 在分布式环境下已配置好对应的采样分片 (DistributedSampler)
    dataloader = get_mesh_dataloader(batch_size=args.batch_size, num_workers=8)

    local_batch_count = 0
    local_sample_count = 0

    # 3. 遍历统计
    print(f"[Rank {rank}] 开始统计数据...")
    for batch in dataloader:
        # 统计当前 batch 的样本数
        batch_size = len(batch['vertices'])
        local_batch_count += 1
        local_sample_count += batch_size

    # 4. 汇总各进程统计信息
    # 将统计结果转换为 tensor 用于通信
    batch_tensor = torch.tensor([local_batch_count], device=f"cuda:{local_rank}")
    sample_tensor = torch.tensor([local_sample_count], device=f"cuda:{local_rank}")
    
    # 使用 all_reduce 将所有卡的统计数据求和至所有卡
    dist.all_reduce(batch_tensor, op=dist.ReduceOp.SUM)
    dist.all_reduce(sample_tensor, op=dist.ReduceOp.SUM)

    # 5. 输出汇总结果 (仅在 rank 0 执行)
    if rank == 0:
        print("\n" + "="*40)
        print("数据集并行加载统计结果:")
        print(f"总参与卡数 (World Size): {world_size}")
        print(f"全局总 Batch 数: {batch_tensor.item()}")
        print(f"全局总数据样本数: {sample_tensor.item()}")
        print("="*40 + "\n")

    dist.destroy_process_group()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch_size", type=int, default=4)
    args = parser.parse_args()
    count_dataset(args)