import pytest

from . import state, wordle

TESTWORDS = [
    "APPAA",
    "APPAB",
    "APPAC",
    "APAPD",

    "BPPAB",
    "PPAPB",
    "PPBBA",
    "BPABB",
    "PPPAC"
]


@pytest.fixture
def wordleEnv():
    env = wordle.WordleEnvBase(
        words=TESTWORDS,
        max_turns=6,
    )
    return env


def test_reset(wordleEnv):
    wordleEnv.reset()


def test_guess_win(wordleEnv):
    wordleEnv.reset()
    goal = wordleEnv.goal_word
    new_state, reward, done, _ = wordleEnv.step(goal)
    assert done
    assert wordleEnv.done
    assert reward == 0

    try:
        wordleEnv.step(goal)
        raise ValueError("Shouldn't reach here!")
    except ValueError:
        pass


def test_win_reward(wordleEnv):
    wordleEnv.reset()
    goal = wordleEnv.goal_word
    word_index = (goal + 1) % len(wordleEnv.words)
    new_state, reward, done, _ = wordleEnv.step(word_index)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-1
    assert not done
    assert not wordleEnv.done
    assert reward < 1

    new_state, reward, done, _ = wordleEnv.step(goal)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-2
    assert done
    assert wordleEnv.done
    assert reward == wordle.REWARD

    try:
        wordleEnv.step(goal)
        raise ValueError("Shouldn't reach here!")
    except ValueError:
        pass


def test_win_reward_6(wordleEnv):
    wordleEnv.reset()
    goal = wordleEnv.goal_word
    random_word = (goal + 1) % len(wordleEnv.words)

    for i in range(5):
        new_state, reward, done, _ = wordleEnv.step(random_word)

    new_state, reward, done, _ = wordleEnv.step(goal)

    assert wordleEnv.max_turns - state.remaining_steps(new_state) == 6
    assert done
    assert wordleEnv.done
    assert reward == wordle.REWARD


def test_lose_reward(wordleEnv):
    wordleEnv.reset()
    goal = wordleEnv.goal_word
    random_word = (goal + 1) % len(wordleEnv.words)
    for i in range(1, wordleEnv.max_turns):
        new_state, reward, done, _ = wordleEnv.step(random_word)
        assert state.remaining_steps(new_state) == wordleEnv.max_turns-i
        assert not done
        assert not wordleEnv.done
        assert reward < 1

    word_index = (goal + wordleEnv.max_turns) % len(wordleEnv.words)
    random_word = word_index
    new_state, reward, done, _ = wordleEnv.step(random_word)
    assert state.remaining_steps(new_state) == 0
    assert done
    assert wordleEnv.done
    assert reward < -wordle.REWARD + 1

    try:
        wordleEnv.step(goal)
        raise ValueError("Shouldn't reach here!")
    except ValueError:
        pass


def letter_test(char, state, letter_state):
    offset = 1+3*5*(ord(char)-ord('A'))
    assert tuple(state[offset:offset+15]) == letter_state


def test_step(wordleEnv):
    wordleEnv.reset()
    wordleEnv.set_goal_encoded(0)

    cur_state = wordleEnv.state
    new_state, reward, done, _ = wordleEnv.step(1)
    assert state.remaining_steps(cur_state) == wordleEnv.max_turns
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-1
    # Expect B to be all 1,0,0
    letter_test('B', new_state, tuple([1, 0, 0]*5))

    # Expect A to be right in position 0 3, no in 1 2 and 0 otherwise
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 0, 0)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2, no in 0 2 and 0 otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    0, 0, 0)
    letter_test('P', new_state, letter_state)

    # Expect C to be no everywhere except in 4
    letter_state = (1, 0, 0,
                    1, 0, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 0)
    letter_test('C', new_state, letter_state)
    cur_state = wordleEnv.state
    new_state, reward, done, _ = wordleEnv.step(3)
    assert state.remaining_steps(cur_state) == wordleEnv.max_turns-1
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-2
    # Expect D to be all 1,0,0
    letter_state = tuple([1, 0, 0]*5)
    letter_test('D', new_state, letter_state)

    # Expect A to be right in position 0 3 and Maybe in 4
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 1, 0)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2 no in 3 and maybe in 4
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    0, 1, 0)
    letter_test('P', new_state, letter_state)

    new_state, reward, done, _ = wordleEnv.step(2)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-3
    # Expect B to be all 1,0,0
    letter_state = tuple([1, 0, 0]*5)
    letter_test('C', new_state, letter_state)

    # Expect C to be all 1,0,0
    letter_state = tuple([1, 0, 0]*5)
    letter_test('C', new_state, letter_state)

    # Expect A to be right in position 0 3 and, no 3 and maybe otherwise
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 1, 0)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2, no in 0 3 and maybe otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    0, 1, 0)
    letter_test('P', new_state, letter_state)

    new_state, reward, done, _ = wordleEnv.step(0)
    # Expect A to be right in position 0 3 and 4
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 0, 1)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2 and not otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)
    assert state.remaining_steps(new_state) == wordleEnv.max_turns-4
    assert done
    assert wordleEnv.done
    assert reward == wordle.REWARD


def test_special_step_cases(wordleEnv):
    wordleEnv.reset()
    wordleEnv.set_goal_encoded(4)
    # BPPAB - goal
    # PPAPB - 1st guess
    new_state, _, _, _ = wordleEnv.step(5)
    # Expect A to be all maybe except 2, 1 and 4 that are no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 1, 0,
                    1, 0, 0)
    letter_test('A', new_state, letter_state)
    # Expect B to be all 0 except 4 that is yes and 1 that is no
    letter_state = (0, 0, 0,
                    1, 0, 0,
                    0, 0, 0,
                    0, 0, 0,
                    0, 0, 1)
    letter_test('B', new_state, letter_state)
    # Expect P to be yes at 1, maybe at 2 and 4, and no in 0, 3 and 4
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 1, 0,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)
    # BPPAB - goal
    # PPAPB - 1st guess
    # PPBBA - 2nd guess
    new_state, _, _, _ = wordleEnv.step(6)
    # Expect A to be all maybe except 2, 1 and 4 that are no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 1, 0,
                    1, 0, 0)
    letter_test('A', new_state, letter_state)
    # Expect B to be maybe at 0, yes at 4, and 1 2 and 3 are no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1)
    letter_test('B', new_state, letter_state)
    # Expect P to be yes at 1, maybe at 2, and no in 0, 3 and 4
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 1, 0,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)

    wordleEnv.reset()
    wordleEnv.set_goal_encoded(7)
    # BPABB - goal
    # PPPAC - 1st guess
    new_state, _, _, _ = wordleEnv.step(8)
    # Expect A to be all maybe except 1 and 3 that is no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    0, 1, 0,
                    1, 0, 0,
                    0, 1, 0)
    letter_test('A', new_state, letter_state)
    # Expect P to be yes at 1, maybe at 3 and 4 and no otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)


def test_mask_update(wordleEnv):
    wordleEnv.reset()
    wordleEnv.set_goal_encoded(0)

    cur_state = wordleEnv.state
    # "APPAA"
    # "APPAB"
    new_state = state.update_from_mask(
        cur_state, wordleEnv.words[1], [2, 2, 2, 2, 0])
    # Expect B to be all 1,0,0
    letter_test('B', new_state, tuple([1, 0, 0]*5))

    # Expect A to be right in position 0 3, no in 1 2 and 0 otherwise
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 0, 0)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2, no in 0 2 and 0 otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    0, 0, 0)
    letter_test('P', new_state, letter_state)

    # Expect C to be no everywhere except in 4
    letter_state = (1, 0, 0,
                    1, 0, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 0)
    letter_test('C', new_state, letter_state)
    # "APPAA",
    # "APPAB",
    # "APAPD",
    new_state = state.update_from_mask(
        new_state, wordleEnv.words[3], [2, 2, 1, 1, 0])
    # Expect D to be all 1,0,0
    letter_state = tuple([1, 0, 0]*5)
    letter_test('D', new_state, letter_state)

    # Expect A to be right in position 0 3 and Maybe in 4
    letter_state = (0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1,
                    0, 1, 0)
    letter_test('A', new_state, letter_state)

    # Expect P to be right in position 1 2 no in 3 and maybe in 4
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 0, 1,
                    1, 0, 0,
                    0, 1, 0)
    letter_test('P', new_state, letter_state)

    wordleEnv.reset()
    cur_state = wordleEnv.state
    wordleEnv.set_goal_encoded(4)
    # BPPAB - goal
    # PPAPB - 1st guess
    new_state = state.update_from_mask(
        cur_state, wordleEnv.words[5], [1, 2, 1, 0, 2])
    # Expect A to be all maybe except 2, 1 and 4 that are no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 1, 0,
                    1, 0, 0)
    letter_test('A', new_state, letter_state)
    # Expect B to be all 0 except 4 that is yes and 1 that is no
    letter_state = (0, 0, 0,
                    1, 0, 0,
                    0, 0, 0,
                    0, 0, 0,
                    0, 0, 1)
    letter_test('B', new_state, letter_state)
    # Expect P to be yes at 1, maybe at 2 and 4, and no in 0, 3 and 4
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 1, 0,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)
    # BPPAB - goal
    # PPAPB - 1st guess
    # PPBBA - 2nd guess
    new_state = state.update_from_mask(
        new_state, wordleEnv.words[6], [1, 2, 1, 1, 1])
    # Expect A to be all maybe except 2, 1 and 4 that are no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 1, 0,
                    1, 0, 0)
    letter_test('A', new_state, letter_state)
    # Expect B to be maybe at 0, yes at 4, and 1 2 and 3 are no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    1, 0, 0,
                    1, 0, 0,
                    0, 0, 1)
    letter_test('B', new_state, letter_state)
    # Expect P to be yes at 1, maybe at 2, and no in 0, 3 and 4
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    0, 1, 0,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)
    wordleEnv.reset()
    wordleEnv.set_goal_encoded(7)
    # BPABB - goal
    # PPPAC - 1st guess
    new_state = state.update_from_mask(
        new_state, wordleEnv.words[8], [0, 2, 0, 1, 0])
    new_state, _, _, _ = wordleEnv.step(8)
    # Expect A to be all maybe except 1 and 3 that is no
    letter_state = (0, 1, 0,
                    1, 0, 0,
                    0, 1, 0,
                    1, 0, 0,
                    0, 1, 0)
    letter_test('A', new_state, letter_state)
    # Expect P to be yes at 1, maybe at 3 and 4 and no otherwise
    letter_state = (1, 0, 0,
                    0, 0, 1,
                    1, 0, 0,
                    1, 0, 0,
                    1, 0, 0)
    letter_test('P', new_state, letter_state)