import os
from stable_baselines3 import SAC
from stable_baselines3.common.vec_env import DummyVecEnv, VecFrameStack
from stable_baselines3.common.monitor import Monitor

# rc_car_envs.py dosyasından sınıfları çekiyoruz
from rc_car_envs import RcCarTargetEnv, RcCarDriveEnv, RcCarComplexEnv

# ==========================================
# AYARLAR: HANGİ MODU EĞİTECEKSİN?
# ==========================================
# Seçenekler: "target", "drive", "complex"
MODE = "complex" 

if __name__ == "__main__":
    
    log_dir = f"./logs_{MODE}/"
    model_dir = f"./models_{MODE}/"
    os.makedirs(log_dir, exist_ok=True)
    os.makedirs(model_dir, exist_ok=True)

    # 1. Ortamı Seç
    if MODE == "target":
        print("🎯 TARGET Mode Seçildi: Sadece hedefe gitmeyi öğrenecek.")
        env = RcCarTargetEnv()
    elif MODE == "drive":
        print("🏎️ DRIVE Mode Seçildi: Sadece hızlı sürmeyi öğrenecek.")
        env = RcCarDriveEnv()
    elif MODE == "complex":
        print("🧠 COMPLEX Mode Seçildi: Hedef + Engel + Hız.")
        env = RcCarComplexEnv()
    else:
        raise ValueError("Geçersiz Mod! 'target', 'drive' veya 'complex' seçin.")

    # 2. Wrapperlar (İzleme ve Stack)
    env = Monitor(env, filename=os.path.join(log_dir, "monitor"))
    env = DummyVecEnv([lambda: env])
    env = VecFrameStack(env, n_stack=4)

    # 3. Model Kurulumu
    print(f"🚀 {MODE.upper()} eğitimi başlıyor...")
    
    model = SAC.load("./models_target/sac_rc_car_target_interrupted", env=env, tensorboard_log=log_dir)

    # 4. Eğitim
    try:
        # İstersen adım sayısını artırabilirsin
        model.learn(total_timesteps=500_000, progress_bar=True)
        
        # Kaydet
        model.save(f"{model_dir}/sac_rc_car_{MODE}_final")
        print("✅ Eğitim Başarıyla Tamamlandı.")
        
    except KeyboardInterrupt:
        print("🛑 Eğitim manuel durduruldu. Kaydediliyor...")
        model.save(f"{model_dir}/sac_rc_car_{MODE}_interrupted")