# Copyright 1996-2024 Cyberbotics Ltd.
# ... (license details) ...

"""
This controller primarily tests if required Python packages
(numpy, matplotlib, scipy, gymnasium, stable_baselines3, torch)
are installed and accessible by the Python environment used by Webots.
It incorporates the basic structure of the original 'driver' example
but focuses on the import checks first.
"""

# --- Add these import statements to test your environment ---
print("--- Starting Check: Importing Required Packages ---")
import sys  # Needed for sys.exit() on failure

all_imports_successful = True # Flag to track success

try:
   import numpy as np
   print("[ OK ] Successfully imported numpy")
   # Optional basic check:
   # _ = np.array([1, 2])
   # print("[ OK ] numpy basic functionality check passed")

   # Import matplotlib.pyplot as requested
   import matplotlib.pyplot as plt
   print("[ OK ] Successfully imported matplotlib.pyplot")
   # Note: Using plt functions directly might be complex here,
   # but the import itself tests installation.

   import scipy
   print("[ OK ] Successfully imported scipy")
   # Optional basic check:
   # from scipy import constants
   # _ = constants.pi
   # print("[ OK ] scipy basic functionality check passed")

   import gymnasium as gym
   print("[ OK ] Successfully imported gymnasium")
   # Optional basic check:
   # _ = gym.spaces.Discrete(2)
   # print("[ OK ] gymnasium basic functionality check passed")

   import stable_baselines3
   print("[ OK ] Successfully imported stable_baselines3")
   # Optional basic check:
   # _ = stable_baselines3.A2C
   # print("[ OK ] stable_baselines3 basic functionality check passed")

   import torch
   print("[ OK ] Successfully imported torch")
   # Optional basic check:
   # _ = torch.tensor([1.0, 2.0])
   # print("[ OK ] torch basic functionality check passed")

   print("\n--- Check Complete: All specified packages imported successfully! ---")

except ImportError as e:
   all_imports_successful = False # Mark failure
   print(f"\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
   print(f"[FAIL] FAILED TO IMPORT A REQUIRED PACKAGE: {e}")
   print(f"       Error Type: {type(e)}")
   # e.name often holds the name of the failed module
   module_name = getattr(e, 'name', 'Unknown')
   print(f"       Failed Module: '{module_name}'")
   print(f"       -> Check if '{module_name}' is installed in the Python environment")
   print(f"          that Webots is configured to use.")
   print(f"       -> Verify Path: Webots > Tools > Preferences > Python command")
   print(f"!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n")
   # Exit the controller script if an import fails - prevents running Supervisor code
   sys.exit(f"ImportError: Required package '{module_name}' not found or import failed.")
# --- End of added import test section ---


# --- Original Webots example script structure (optional, runs only if imports succeed) ---
# --- This part might produce warnings if the world doesn't match the example exactly ---
if all_imports_successful:
   print("\n[INFO] Imports successful, proceeding with Supervisor initialization (if applicable).")
   from controller import Supervisor
   # Make sure 'common.py' exists if you are using the original example structure, otherwise define dummy
   try:
       # This specific example might not need common.py, adjust if needed
       pass # from common import common_print
   except ImportError:
       # Define a dummy function if needed
       # def common_print(msg): print(f"Dummy common_print: {msg}")
       print("[WARN] 'common.py' not found or not needed for this basic test.")


   class Driver (Supervisor):
       timeStep = 128 # Or get from world basicTimeStep

       def __init__(self):
           super(Driver, self).__init__()
           self.keyboard = None # Initialize to None
           self.timeStep = int(self.getBasicTimeStep()) # Use world timestep

           # Safely try to get keyboard (may not exist)
           try:
               self.keyboard = self.getKeyboard() # Get keyboard device
               if self.keyboard: # Check if keyboard was found
                   self.keyboard.enable(self.timeStep) # Enable keyboard reading
                   print("[INFO] Keyboard enabled for optional interaction.")
               else:
                   print("[WARN] Keyboard device not found in this world.")

           except Exception as e:
               print(f"[WARN] Could not initialize keyboard device: {e}")

       def run(self):
           print("[INFO] Controller 'run' method started.")
           print("[INFO] Primary goal (import tests) already passed.")
           print("[INFO] This controller will now run passively.")
           if self.keyboard:
                print("[INFO] You can press keys (if defined below) but main focus was imports.")
           # Passive loop to keep controller alive and show simulation running
           while self.step(self.timeStep) != -1:
               # Optionally add key checks here if needed for further testing
               # k = self.keyboard.getKey() if self.keyboard else -1
               # if k == ord('I'): print("Info key pressed")
               pass # Just keep stepping

           print("[INFO] Driver run loop finished or simulation stopped.")

   # --- Script execution ---
   if all_imports_successful:
       print("[INFO] Initializing Driver controller...")
       controller = Driver()
       print("[INFO] Starting controller run loop (passive)...")
       controller.run()
       print("[INFO] Controller run loop exited.")
   else:
       print("[INFO] Script halted due to import errors. Supervisor code was not executed.")

else:
    print("[INFO] Script halted before Supervisor initialization due to import errors.")