#!/usr/bin/env python3
"""
Setup and run FSD-Level5-CoT training on SADC dataset.
This script:
1. Clones the model repo
2. Installs dependencies
3. Runs end-to-end training
"""
import subprocess
import sys
import os

def run(cmd, **kwargs):
    print(f"$ {cmd}")
    result = subprocess.run(cmd, shell=True, capture_output=True, text=True, **kwargs)
    if result.stdout:
        print(result.stdout[-2000:])
    if result.stderr:
        print(result.stderr[-2000:])
    return result.returncode == 0

# Install dependencies
print("=" * 60)
print("  Installing dependencies...")
print("=" * 60)
run("pip install torch torchvision datasets huggingface_hub trackio Pillow numpy --quiet")

# Clone model repo
print("\n" + "=" * 60)
print("  Cloning model repo...")
print("=" * 60)
if not os.path.exists("/app/fsd_repo"):
    run("git clone https://huggingface.co/Reality123b/FSD-Level5-CoT /app/fsd_repo")
else:
    run("cd /app/fsd_repo && git pull")

# Run training with small subset for CPU feasibility
print("\n" + "=" * 60)
print("  Starting training...")
print("=" * 60)
os.chdir("/app/fsd_repo")

cmd = [
    sys.executable, "train_sadc_e2e.py",
    "--train_samples", "200",
    "--val_samples", "50",
    "--epochs", "2",
    "--batch_size", "2",
    "--grad_accum", "4",
    "--lr", "3e-4",
    "--num_workers", "0",
    "--log_every", "5",
    "--eval_every", "50",
    "--push_to_hub",
    "--hub_model_id", "Reality123b/FSD-Level5-CoT",
    "--save_dir", "/app/checkpoints",
    "--data_dir", "/app/sadc_subset",
]

print(" ".join(cmd))
result = subprocess.run(cmd, cwd="/app/fsd_repo")
print(f"\nTraining exited with code: {result.returncode}")