# Download word list from https://www-cs-faculty.stanford.edu/~knuth/sgb-words.txt

import random
import sys
from collections import defaultdict

def get_wordle_hints(guess, answer):
    result = ''
    answer_minus_green = ''
    for i in range(len(guess)):
        if guess[i] == answer[i]:
            answer_minus_green += ' '
        else:
            answer_minus_green += answer[i]
    
    for i in range(len(guess)):
        if guess[i] == answer[i]:
            result += 'g'
        elif guess[i] in answer_minus_green:
            result += 'y'
            pos = answer_minus_green.find(guess[i])
            answer_minus_green = answer_minus_green[0:pos] + ' ' + answer_minus_green[pos+1:]
        else:
            result += 'k'
    return result

def get_best_guesses(words, guesses_so_far):
    valid_words = []
    for word in words:
        success = True
        for guess in guesses_so_far:
            if get_wordle_hints(guess, word) != guesses_so_far[guess]:
                success = False
        if success:
            valid_words.append(word)

    if len(valid_words) == 0: # Can only happen with invalid hints
        return [('xxxxx', 0)]
    if len(valid_words) == 1:
        return [(valid_words[0], 0)]

    guess_ratings = dict()
    # for guess in valid_words: # HARD MODE
    counter = 0
    for guess in words:
        counter += 1
        # if (counter % 100 == 0):
        #     print(counter)
        hints_dict = defaultdict(int)
        for answer in valid_words:
            hints_dict[get_wordle_hints(guess, answer)] += 1
        guess_ratings[guess] = max(hints_dict.values())
    
    sorted_guess_ratings = sorted(guess_ratings.items(), key=lambda item: item[1])
    best_rating = sorted_guess_ratings[0][1]
    best_guesses = [x for x in sorted_guess_ratings if x[1] == best_rating]
    best_guesses_in_valid_words = [x for x in best_guesses if x[0] in valid_words]
    if len(best_guesses_in_valid_words) > 0:
        return best_guesses_in_valid_words
    else:
        return best_guesses

words = list()
with open('sgb-words.txt') as f:
    for line in f:
        if len(line.strip()) == 5:
            words.append(line.strip())

first_guess = 'aloes'

guess = first_guess
guesses_so_far = dict()
while True:
    print('Next guess: ' + guess)
    hint = input('Enter result (k for black, y for yellow, g for green): ')
    guesses_so_far[guess] = hint
    best_guesses = get_best_guesses(words, guesses_so_far)
    guess = best_guesses[0][0]
    if guess == 'xxxxx':
        print('Word is not in dictionary.')
        break
    elif best_guesses[0][1] == 0:
        print('The word is: ' + guess)
        break
    
sys.exit(0)

# The below code is used to test the solver on all the words in the dictionary and see how it does

print("Building second_guess_dict...")

second_guess_dict = dict()
for c1 in ['k', 'y', 'g']:
    for c2 in ['k', 'y', 'g']:
        for c3 in ['k', 'y', 'g']:
            for c4 in ['k', 'y', 'g']:
                for c5 in ['k', 'y', 'g']:
                    hint = c1 + c2 + c3 + c4 + c5
                    best_guess = get_best_guesses(words, {first_guess: hint})[0][0]
                    second_guess_dict[hint] = best_guess
                    print(hint + ' ' + best_guess)

num_guesses_dict = defaultdict(int)

counter = 0
for word in words:
    counter += 1
    if (counter % 100) == 0:
        print(counter)
    guess = first_guess
    guesses_so_far = dict()
    # print(word)
    if guess != word:
        hint = get_wordle_hints(guess, word)
        guesses_so_far[guess] = hint
        # print(guess + ' ' + hint)
        guess = second_guess_dict[hint]
    while guess != word:
        hint = get_wordle_hints(guess, word)
        guesses_so_far[guess] = hint
        # print(guess + ' ' + hint)
        best_guesses = get_best_guesses(words, guesses_so_far)
        # guess = best_guesses[random.randint(0,len(best_guesses)-1)][0]
        guess = best_guesses[0][0]
    num_guesses = len(guesses_so_far) + 1
    num_guesses_dict[num_guesses] += 1
    # print(word + ' ' + str(num_guesses))
    if num_guesses > 6:
        print(word + ' ' + str(num_guesses))

for k, v in sorted(num_guesses_dict.items(), key=lambda item: item[0]):
    print(str(k) + ': ' + str(v))
print(float(sum([k*v for k, v in num_guesses_dict.items()])) / sum([v for k, v in num_guesses_dict.items()]))