def get_words():
    with open("../txt/e098.txt", "r") as f:
        words = list(map(lambda w: w.strip().lower(), f.readlines()))
    words.sort(key=len, reverse=True)
    return words


def get_pair_subsets(xs):
    """ Returns all subsets of size two. """
    r = [(xs[i], xs[j])
            for i in range(len(xs))
            for j in range(i + 1, len(xs))]
    return r


def find_anagrams(xs):
    anagrams = {}
    for x in xs:
        x_tuple = tuple(sorted(x))
        try:
            anagrams[x_tuple].append(x)
        except KeyError:
            anagrams[x_tuple] = [x]
    anagrams = [a for a in anagrams.values() if len(a) > 1]
    return anagrams


def find_anagram_pairs(xs):
    anagrams = find_anagrams(xs)
    r = [pair
         for anagram in anagrams
         for pair in get_pair_subsets(anagram)]
    return r


def words_match_numbers(words, numbers):
    letters_a, letters_b = words
    digits_a, digits_b = numbers

    digit_to_letter = {}
    for l, d in zip(letters_a, digits_a):
        if d in digit_to_letter and digit_to_letter[d] != l:
            # Each digit can only be assigned to one letter
            return False
        digit_to_letter[d] = l


    letter_to_digit = {l: d for l, d in zip(letters_a, digits_a)}
    for l, d in zip(letters_b, digits_b):
        if letter_to_digit[l] != d:
            return False
    return True


def find_pairs(words, squares):
    pairs = []
    anagram_pairs = find_anagram_pairs(words)
    square_pairs = find_anagram_pairs(squares)
    for word_pair in anagram_pairs:
        for number_pair in square_pairs:
            if words_match_numbers(word_pair, number_pair):
                # pairs.append([word_pair, number_pair])
                pairs.append(number_pair)
    return pairs


def euler_098():
    squares = [str(i * i) for i in range(200000)]
    words = get_words()

    largest_number = 0
    for i in range(3, 9):
        words_i = list(filter(lambda w: len(w) == i, words))
        squares_i = list(filter(lambda w: len(w) == i, squares))
        pairs = find_pairs(words_i, squares_i)
        for pair in pairs:
            n = max(map(int, pair))
            if n > largest_number:
                largest_number = n
    return largest_number


if __name__ == "__main__":
    solution = euler_098()
    print("e098.py: " + str(solution))
    assert(solution == 18769)