from pathlib import Path
import json

def main():
    file_content = Path('words_en.txt').read_text(encoding='utf-8')
    training_data = file_content.splitlines()

    weights = {}
    for word in training_data:
        extend_weights(weights, word)
    
    json_output = json.dumps(weights, indent=2)
    Path('weights.json').write_text(json_output, encoding='utf-8')
    
def extend_weights(weights, word):
    for i in range(len(word)):
        left = word[:i]
        right = word[i:]
        context = left[-1:]
        next_letter = right[0]

        update_weights(context, next_letter, weights)


def update_weights(context, next_letter, weights):
    if context not in weights:
        weights[context] = {}
    
    vekter = weights[context]

    if next_letter not in vekter:
        vekter[next_letter] = 1
    else:
        vekter[next_letter] += 1


if __name__ == '__main__':
    main()