"""
多进程并行统计：总时间序列条数和总时间戳数目。
支持自定义进程数，分进度条显示每个进程的读取速度。
"""

import argparse
import os
from os import path
import multiprocessing as mp

from tqdm import tqdm
from datasets import Dataset as HFDataset


def process_subsets(subset_paths, process_idx, result_queue):
    """
    单个进程的处理函数：处理分配到的子数据集，统计结果并放入队列。

    Args:
        subset_paths: 分配给当前进程的子数据集路径列表
        process_idx: 进程编号（用于设置进度条位置）
        result_queue: 用于传递统计结果的队列
    """
    process_series_count = 0
    process_timestamp_count = 0

    # 为当前进程创建独立的进度条，position 确保进度条不重叠
    pbar = tqdm(
        subset_paths,
        position=process_idx,
        desc=f"Process {process_idx}",
        leave=True,  # 进度条完成后保留在控制台
    )

    for subset_path in pbar:
        try:
            dataset = HFDataset.load_from_disk(subset_path)
            for sample in dataset:
                seq_len = len(sample["target"])
                process_series_count += 1
                process_timestamp_count += seq_len
            # 更新进度条后缀，显示当前进程已统计的条数
            pbar.set_postfix({"count": process_series_count})
        except Exception as e:
            print(
                f"\nProcess {process_idx} error processing {path.basename(subset_path)}: {str(e)}"
            )

    # 将当前进程的统计结果放入队列
    result_queue.put((process_series_count, process_timestamp_count))


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Multi-process statistics: total time series and timestamps."
    )
    parser.add_argument(
        "--subset_dir",
        type=str,
        required=True,
        help="Directory containing the split subsets (e.g., ./split_subsets).",
    )
    parser.add_argument(
        "--num_processes",
        type=int,
        default=mp.cpu_count(),
        help=f"Number of processes to use (default: CPU count = {mp.cpu_count()}).",
    )

    args = parser.parse_args()

    # ---------------------- 1. 收集并排序子数据集路径 ----------------------
    if not path.exists(args.subset_dir):
        raise ValueError(f"Subset directory {args.subset_dir} does not exist.")

    subset_folders = []
    for folder_name in os.listdir(args.subset_dir):
        folder_path = path.join(args.subset_dir, folder_name)
        if path.isdir(folder_path) and folder_name.startswith("dataset_subset_"):
            subset_folders.append(folder_path)

    if not subset_folders:
        raise ValueError(f"No valid subset folders found in {args.subset_dir}")

    subset_folders.sort()
    num_subsets = len(subset_folders)
    num_processes = min(args.num_processes, num_subsets)  # 进程数不超过子数据集数

    print(f"\nFound {num_subsets} subsets, using {num_processes} processes.\n")

    # ---------------------- 2. 分配任务给每个进程 ----------------------
    # 将子数据集路径均匀分配给各个进程
    subsets_per_process = [[] for _ in range(num_processes)]
    for i, subset_path in enumerate(subset_folders):
        subsets_per_process[i % num_processes].append(subset_path)

    # ---------------------- 3. 启动多进程处理 ----------------------
    result_queue = mp.Queue()  # 用于收集各进程的统计结果
    processes = []

    for i in range(num_processes):
        p = mp.Process(
            target=process_subsets, args=(subsets_per_process[i], i, result_queue)
        )
        processes.append(p)
        p.start()

    # ---------------------- 4. 收集并汇总结果 ----------------------
    total_series_count = 0
    total_timestamp_count = 0

    # 等待所有进程完成并收集结果
    for _ in range(num_processes):
        series_count, timestamp_count = result_queue.get()
        total_series_count += series_count
        total_timestamp_count += timestamp_count

    # 等待所有进程结束
    for p in processes:
        p.join()

    # ---------------------- 5. 打印最终结果 ----------------------
    print("\n" + "=" * 50)
    print(f"总时间序列条数: {total_series_count}")
    print(f"总时间戳数目:   {total_timestamp_count}")
    print("=" * 50)
