# Experiment E1: Text generation with Markov chains (random only, offline)
# Requires: random (standard library) | no dataset, no internet
# ===========================================================================

import random

random.seed(42)

# 1. Tiny training text typed directly in the code (no dataset)
text = ("the cat sat on the mat and the dog sat on the rug . "
        "the cat saw the dog and the dog saw the cat . "
        "the dog ran to the mat and the cat ran to the rug .")
words = text.split()
print("Number of words:", len(words))

# 2. Build the Markov chain: word -> list of words that follow it
chain = {}
for i in range(len(words) - 1):
    current_word = words[i]
    next_word = words[i + 1]
    if current_word not in chain:
        chain[current_word] = []
    chain[current_word].append(next_word)

print("\nTransition table:")
for w in chain:
    print(f"  {w:>4} -> {chain[w]}")

# 3. Next-word probabilities (counts / total)
def probabilities(word):
    followers = chain[word]
    return {w: followers.count(w) / len(followers) for w in set(followers)}

print("\nP(next | 'the') =", {w: round(p, 2) for w, p in probabilities("the").items()})

# 4. Generate text: start word -> repeatedly pick the next word at random from the chain
def generate(start="the", length=20):
    word = start
    output = [word]
    for _ in range(length - 1):
        if word not in chain:              # dead end -> stop
            break
        word = random.choice(chain[word])  # sampling = probability of each follower
        output.append(word)
    return " ".join(output)

print("\nGenerated text (word level):")
for start in ["the", "dog", "cat"]:
    print(" ", generate(start, 15))

# 5. Character-level Markov chain of order k (uses the previous k characters)
def build_char_chain(text, k=3):
    model = {}
    for i in range(len(text) - k):
        state = text[i:i + k]
        model.setdefault(state, []).append(text[i + k])
    return model

def generate_chars(model, k=3, length=120, seed_text=None):
    state = seed_text or random.choice(list(model.keys()))
    out = state
    for _ in range(length):
        if state not in model:
            break
        ch = random.choice(model[state])
        out += ch
        state = out[-k:]
    return out

for k in [2, 4]:
    model = build_char_chain(text, k)
    print(f"\nCharacter-level, order {k} ({len(model)} states):")
    print(" ", generate_chars(model, k, 100, seed_text=text[:k]))
