import os
import subprocess

def main():
    base_dir = r"D:\Users\roy\CV_final_project\champ-master"
    ref_smpl_dir = os.path.join(base_dir, "reference_imgs", "smpl_results")
    driving_videos_dir = os.path.join(base_dir, "driving_videos")
    
    if not os.path.exists(ref_smpl_dir):
        print(f"錯誤：找不到參考骨架資料夾 {ref_smpl_dir}")
        return
        
    if not os.path.exists(driving_videos_dir):
        print(f"錯誤：找不到動作影片資料夾 {driving_videos_dir}")
        return

    # 指定只使用 ref-01.npy
    ref_file = "ref-01.npy"
    ref_path = os.path.join("reference_imgs", "smpl_results", ref_file)
    
    if not os.path.exists(os.path.join(base_dir, ref_path)):
        print(f"錯誤：找不到指定的參考骨架 {ref_file}")
        return
    
    # 取得所有的動作影片資料夾名稱
    video_names = [d for d in os.listdir(driving_videos_dir) if os.path.isdir(os.path.join(driving_videos_dir, d))]
    video_names.sort()
    
    print(f"只鎖定單一參考圖: {ref_file}")
    print(f"總共需要為 {len(video_names)} 支動作影片進行 SMPL 參數融合轉換...\n")

    for video_name in video_names:
        video_smpl_path = os.path.join("driving_videos", video_name)
        
        # 檢查該影片的 3D 骨架 npz 是否已成功產生
        if not os.path.exists(os.path.join(base_dir, video_smpl_path, "smpl_results", "smpls_group.npz")):
            print(f"[{video_name}] 尚未生成完整的 smpls_group.npz，略過此影片。")
            continue

        # 依照您的要求，不再多包一層 ref-01 資料夾
        output_folder = os.path.join("transferd_result", video_name)
        
        # 加入斷點續傳設計：檢查是否已經輸出過
        check_out_dir = os.path.join(base_dir, output_folder, "smpl_results")
        if os.path.exists(check_out_dir):
            output_smpls = [f for f in os.listdir(check_out_dir) if f.endswith('.npy')]
            # 如果資料夾存在且裡面已經有 npy 檔案，就認定為處理過了
            if len(output_smpls) > 0:
                print(f"[{video_name}] 已經轉換過，自動跳過...")
                continue
        
        print(f"========================================")
        print(f"正在轉換: 影片 [{video_name}] + 參考圖 [{ref_file}]")
        print(f"========================================")
        
        cmd = [
            "python", "-m", "scripts.data_processors.smpl.smpl_transfer",
            "--reference_path", ref_path,
            "--driving_path", video_smpl_path,
            "--output_folder", output_folder,
            "--figure_transfer",
            "--view_transfer"
        ]
        
        result = subprocess.run(cmd, cwd=base_dir)
        
        if result.returncode != 0:
            print(f"[{video_name}] 轉換時發生錯誤！")
        else:
            print(f"[{video_name}] 轉換完成！\n")

if __name__ == "__main__":
    main()