from examples.introduction.core.ground import LetterWorld
from examples.introduction.core.label import LetterWorldLabellingFunction
from examples.introduction.core.machine import LetterWorldRewardMachine
from examples.introduction.core.crossproduct import LetterWorldCrossProduct

# 1. Create the ground environment
ground_env = LetterWorld()

# 2. Create the labelling function
lf = LetterWorldLabellingFunction()

# 3. Create the Reward Machine
rm = LetterWorldRewardMachine()

# 4. Create the cross-product MDP
env = LetterWorldCrossProduct(
    ground_env=ground_env,
    machine=rm,
    lf=lf,
    max_steps=100,
)

# Use like a standard Gymnasium environment
obs, _ = env.reset()
action = env.action_space.sample()
next_obs, reward, terminated, truncated, info = env.step(action)

import numpy as np
from collections import defaultdict

q_table = defaultdict(lambda: np.zeros(env.action_space.n))

for episode in range(100):
    obs, _ = env.reset()
    done = False
    
    while not done:
        # Epsilon-greedy action selection
        if np.random.rand() < 0.1:
            action = env.action_space.sample()
        else:
            action = np.argmax(q_table[obs])
            
        next_obs, reward, terminated, truncated, _ = env.step(action)
        done = terminated or truncated
        
        # Q-learning update
        q_table[obs][action] += 0.1 * (
            reward + 0.99 * np.max(q_table[next_obs]) - q_table[obs][action]
        )
        
        obs = next_obs