import subprocess
import sys
import time

# Participants đại diện: p01 (Cơ bản), p08 (Thử thách cho Pooling), p11 (Thử thách cho Loss)
# p11 đã hoàn thành (fold 00 theo cách gọi của người dùng)
PARTICIPANTS = ["p01", "p08"] 
IDS = range(1, 9)

print(f"{'='*60}")
print(f" STARTING ABLATION BATCH FOR {PARTICIPANTS} ")
print(f" IDs: {list(IDS)} ")
print(f"{'='*60}")

start_total = time.time()

for participant in PARTICIPANTS:
    print(f"\n{'#'*40}")
    print(f" PROCESSING PARTICIPANT: {participant} ")
    print(f"{'#'*40}")
    
    for id_val in IDS:
        print(f"\n>>> Running Ablation ID #{id_val} for {participant}...")
        cmd = [
            sys.executable, "src/train_ablation.py",
            "--id", str(id_val),
            "--participant", participant,
            "--epochs", "30", 
            "--workers", "4"
        ]
        
        start_id = time.time()
        process = subprocess.Popen(cmd)
        process.wait()
        end_id = time.time()
        
        if process.returncode == 0:
            print(f"SUCCESS: {participant} ID #{id_val} completed in {(end_id - start_id)/60:.2f} mins.")
        else:
            print(f"FAILED: {participant} ID #{id_val} with exit code {process.returncode}")

end_total = time.time()
print(f"\n{'='*60}")
print(f" BATCH COMPLETED in {(end_total - start_total)/3600:.2f} hours ")
print(f"{'='*60}")