import os
import shutil
from collections import defaultdict
from pprint import pprint
import matplotlib.pyplot as plt

vid_dir_list = [
    "/Users/baizechen/Downloads/2025_10_18_1traj_sr15",
    "/Users/baizechen/Downloads/2025_10_19_step29_sr22",
    "/Users/baizechen/Downloads/2025_10_19_step59_sr17"
]

# for each vid_dir, split the videos in it into 2 parts
# one is name includes "success=True", the other is name includes "success=False"
# for vid_dir in vid_dir_list:
#     for file in os.listdir(vid_dir):
#         if "success=True" in file:
#             success_True_dir = os.path.join(vid_dir, "success=True")
#             if not os.path.exists(success_True_dir):
#                 os.makedirs(success_True_dir)
#             shutil.move(os.path.join(vid_dir, file), os.path.join(success_True_dir, file))
#         elif "success=False" in file:
#             success_False_dir = os.path.join(vid_dir, "success=False")
#             if not os.path.exists(success_False_dir):
#                 os.makedirs(success_False_dir)
#             shutil.move(os.path.join(vid_dir, file), os.path.join(success_False_dir, file))

# for each video dir, after spliting into success=True and success=False, get the success rate for each task
task_success_rate_across_dirs = defaultdict(list)
for vid_dir in vid_dir_list:
    success_dict = {}
    total_dict = {}
    for sub_dir in os.listdir(vid_dir):
        if "success=True" in sub_dir:
            success_flag = 1
        elif "success=False" in sub_dir:
            success_flag = 0
        else:
            continue
        for video_file in os.listdir(os.path.join(vid_dir, sub_dir)):
            task_name = video_file.split("task=")[1].split(".")[0]
            success_dict[task_name] = success_dict.get(task_name, 0) + success_flag
            total_dict[task_name] = total_dict.get(task_name, 0) + 1
    for task_name in success_dict.keys():
        success_rate = success_dict[task_name] / total_dict[task_name]
        task_success_rate_across_dirs[task_name].append(success_rate)

# pprint(task_success_rate_across_dirs)
# plot the success rate across trials into sub-plots, each sub-plot is a task
fig, axes = plt.subplots(4, 3, figsize=(16, 12))
for i, task_name in enumerate(task_success_rate_across_dirs.keys()):
    plt.subplot(4, 3, i + 1)
    plt.plot(task_success_rate_across_dirs[task_name])
    plt.title(task_name)
    plt.xlabel("Trial")
    plt.ylabel("Success Rate")
plt.tight_layout()
plt.savefig("success_rate_across_trials.png")
plt.close()