# Words by the Numbers 5, Exercise 1: a first-order chain on a short invented passage
import random
from collections import Counter, defaultdict

text = "the old lamp lit the old road and the road led to the old gate"   # invented
words = text.split()
followers = defaultdict(list)
for a, b in zip(words, words[1:]):
    followers[a].append(b)

print(len(words), "words,", len(set(zip(words, words[1:]))), "distinct bigrams")
for w in ("the", "old", "road"):
    c = Counter(followers[w])
    print(f"after {w!r}:", {k: f"{v}/{len(followers[w])}" for k, v in c.items()})

# chance of the exact path the -> old -> road
p = 1.0
path = ["the", "old", "road"]
for a, b in zip(path, path[1:]):
    p *= Counter(followers[a])[b] / len(followers[a])
print("chance of", " ".join(path), "starting from 'the':", round(p, 3))

random.seed(2)
word, out = "the", ["the"]
for _ in range(9):
    if word not in followers:
        break
    word = random.choice(followers[word])
    out.append(word)
print(" ".join(out))

# Output when run (Python 3.14; the documentation-text figures depend on the Python version):
# 15 words, 12 distinct bigrams
# after 'the': {'old': '3/4', 'road': '1/4'}
# after 'old': {'lamp': '1/3', 'road': '1/3', 'gate': '1/3'}
# after 'road': {'and': '1/2', 'led': '1/2'}
# chance of the old road starting from 'the': 0.25
# the old lamp lit the old gate