import time
import pybullet as p
from stable_baselines3 import PPO
from kuka_env import KukaReachingEnv

if __name__ == "__main__":
    print("Loading Environment in GUI mode...")
    env = KukaReachingEnv(render_mode="human")
    
    print("Loading Vision-Trained Model...")
    try:
        model = PPO.load("models/ppo_vision_kuka", env=env)
    except FileNotFoundError:
        print("Error: Could not find the model.")
        env.close()
        exit()

    obs, info = env.reset()
    
    print("Watch the AI control the arm using only raw pixels!")
    try:
        for i in range(5000):
            if not p.isConnected():
                break
                
            action, _states = model.predict(obs, deterministic=True)
            obs, reward, terminated, truncated, info = env.step(action)
            
            if terminated:
                print(f"Target Hit! (Distance: {info['distance']:.3f})")
                time.sleep(0.5) 
                obs, info = env.reset()
            elif truncated:
                obs, info = env.reset()
                
    except KeyboardInterrupt:
        pass
    finally:
        env.close()