import os
import torch
from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import CheckpointCallback # <--- NEW IMPORT
from kuka_env import KukaReachingEnv
from utils import MetricsCallback

if __name__ == "__main__":
    os.makedirs("models", exist_ok=True)
    os.makedirs("models/checkpoints", exist_ok=True) # <--- NEW FOLDER
    os.makedirs("logs", exist_ok=True)
    
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"Training on device: {device.upper()}")

    print("Initializing Environment...")
    env = KukaReachingEnv(render_mode=None)

    print("Setting up PPO with Convolutional Neural Network (CNN)...")
    model = PPO(
        "CnnPolicy", 
        env, 
        verbose=1, 
        tensorboard_log="./logs/",
        device=device,
        batch_size=128,      
        n_steps=2048         
    )

    # 1. Our custom metric tracker
    metrics_callback = MetricsCallback()

    # 2. NEW: The Checkpoint Saver!
    # Saves a backup of the model every 50,000 steps
    checkpoint_callback = CheckpointCallback(
        save_freq=50_000,
        save_path="./models/checkpoints/",
        name_prefix="ppo_vision"
    )

    print("Starting Deep Vision Training (3,000,000 steps)...")
    print("Backups will be saved automatically to models/checkpoints/")
    
    # Pass BOTH callbacks as a list
    model.learn(
        total_timesteps=3_000_000, 
        callback=[metrics_callback, checkpoint_callback]
    )

    print("Training complete! Saving the final model...")
    model.save("models/ppo_vision_kuka")
    
    env.close()