import gymnasium as gym
import torch

episodes = 100
env = gym.make("MountainCar-v0", render_mode="human") # MountainCar environment with visualization

def policy_random(state):
	return env.action_space.sample()  # Random action

def policy_custom(state):
	position, velocity = state # hole Position und Geschwindigkeit aus state vector/tensor
	position = position + 0.5 # => Mitte 'Tal' ist jetzt bei 0 (ursprünglich -0.5 blöde environment;)
	# Aufgabe: car soll den Berg hochfahren zur Fahne, finde eine gute Policy
	# wähle aktion: # 0 = left, 1 = no action, 2 = right
	links = 0
	rechts = 2
	action = 1 # do nothing
	if position > 0:
		action = links
	# return links # immer nach links
	return action

def policy(state):
	# return policy_random(state)
	return policy_custom(state)

for step in range(episodes): # Schleife noch ohne Training, nur zum Testen der Policy
	state = torch.tensor(env.reset()[0]) # unser Zustand ist ein Tensor (bei uns: [position, velocity] vector)
	done = False
	while not done:
		action = policy(state)
		next_state, reward, terminated, truncated, info = env.step(action) # p
		state = torch.tensor(next_state) # umwandeln in Tensor / Vector
		done = terminated or truncated # Ziel erreicht oder abgebrochen wegen Zeit