Fix complex case test and update state updater

wordle_env/state.py

@@ -171,7 +171,11 @@ def update(state: WordleState, word: str, goal_word: str) -> Tuple[WordleState,
                 # Char at all positions = no
                 state[offset:offset + 3 * WORDLE_N] = [1, 0, 0] * WORDLE_N
             else: # goal_word.count(c) <= processed_letters.count(c) and goal in word
-                # Only information at this point is that char at position i = no
+                # At i and in every position which is not set = no
                 state[offset + 3 * i:offset + 3 * i + 3] = [1, 0, 0]
+                for char_idx in range(0, WORDLE_N * 3, 3):
+                    char_offset = offset + char_idx
+                    if tuple(state[char_offset: char_offset + 3]) == (0, 0, 0):
+                        state[char_offset: char_offset + 3] = [1, 0, 0]
             processed_letters.append(c)
     return state, reward

wordle_env/test_wordle.py

@@ -161,7 +161,6 @@ def test_step(wordleEnv):
                     1, 0, 0,
                     0, 0, 1,
                     0, 1, 0)
-    print(new_state, letter_state)
     letter_test('A', new_state, letter_state)
 
     # Expect P to be right in position 1 2 no in 3 and maybe in 4
@@ -288,6 +287,6 @@ def test_special_step_cases(wordleEnv):
     letter_state = (1, 0, 0,
                     0, 0, 1,
                     1, 0, 0,
-                    0, 0, 0,
-                    0, 0, 0)
+                    1, 0, 0,
+                    1, 0, 0)
     letter_test('P', new_state, letter_state)