"""Train a tiny word-level transition table. Not a neural network."""
import argparse
import json
import re
from collections import Counter, defaultdict
from pathlib import Path


def tokenize(text):
    return re.findall(r"[A-Za-z0-9']+|[.,!?;:]", text)


def train(text):
    counts = defaultdict(Counter)
    for line in text.splitlines():
        words = tokenize(line.strip())
        if not words:
            continue
        prior = '<START>'
        for word in words:
            counts[prior][word] += 1
            prior = word
        counts[prior]['<END>'] += 1
    return {word: dict(next_words) for word, next_words in sorted(counts.items())}


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    p.add_argument('--input', default='corpus.txt')
    p.add_argument('--output', default='transitions.json')
    args = p.parse_args()
    model = train(Path(args.input).read_text(encoding='utf-8'))
    Path(args.output).write_text(json.dumps(model, indent=2, ensure_ascii=False) + '\n', encoding='utf-8')
    print(f'Wrote {len(model)} transition states to {args.output}')