"""Model 101: train and inspect a tiny bigram model, not an LLM.

Python 3.11+ standard library. Original fictional text only; no downloads,
network, provider account, GPU, persistent writes or equipment commands.
"""
import argparse
from collections import Counter, defaultdict
import hashlib
import json
import math
import random


CORPUS = [
    "the pump needs inspection",
    "the pump needs inspection",
    "the pump needs inspection",
    "the pump needs repair",
    "the valve needs repair",
]
SAMPLE_SEED = 0
EXTRA_REPAIR_EXAMPLES = 4
END = "<end>"
CASES = ("greedy", "sampled", "cooler", "earlier-context", "output-limit", "unseen", "retrained")


def train(sentences):
    counts = defaultdict(Counter)
    for sentence in sentences:
        tokens = sentence.split() + [END]
        for previous, following in zip(tokens, tokens[1:]):
            counts[previous][following] += 1
    return {previous: dict(sorted(row.items())) for previous, row in sorted(counts.items())}


def fingerprint(model):
    return hashlib.sha256(json.dumps(model, sort_keys=True).encode()).hexdigest()[:12]


def distribution(model, previous, temperature):
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("Temperature must be finite and greater than zero; greedy is a separate mode.")
    row = model.get(previous, {})
    if not row:
        return {}
    # log(count) is a toy score. Softmax(score / T) normalizes it.
    scores = {token: math.log(count) / temperature for token, count in row.items()}
    peak = max(scores.values())
    masses = {token: math.exp(score - peak) for token, score in scores.items()}
    total = sum(masses.values())
    return {token: mass / total for token, mass in masses.items()}


def generate(model, prompt, mode="greedy", temperature=1.0, max_steps=4):
    if mode not in ("greedy", "sample") or max_steps < 1:
        raise ValueError("Use greedy or sample with a positive token-step limit.")
    context = prompt.split()
    if not context:
        raise ValueError("The teaching prompt needs at least one word.")
    before = fingerprint(model)
    rng = random.Random(SAMPLE_SEED)
    trace, generated = [], []
    reason = "token-limit"
    for _ in range(max_steps):
        previous = context[-1]  # Deliberately only one word of context.
        probabilities = distribution(model, previous, temperature)
        if not probabilities:
            reason = "unseen-context"
            break
        draw = None
        if mode == "greedy":
            chosen = max(probabilities, key=probabilities.get)
        else:
            draw = rng.random()
            cumulative = 0.0
            chosen = next(reversed(probabilities))
            for token, probability in probabilities.items():
                cumulative += probability
                if draw < cumulative:
                    chosen = token
                    break
        trace.append({"previous": previous, "probabilities": probabilities, "draw": draw, "chosen": chosen})
        if chosen == END:
            reason = "end-token"
            break
        generated.append(chosen)
        context.append(chosen)
    return {"prompt": prompt, "mode": mode, "temperature": temperature,
            "continuation": " ".join(generated), "stop_reason": reason, "trace": trace,
            "parameters_before": before, "parameters_after": fingerprint(model),
            "parameters_unchanged": before == fingerprint(model)}


def experiment(name, base):
    model, prompt, settings = base, "the pump needs", {}
    if name == "sampled":
        settings = {"mode": "sample"}
    elif name == "cooler":
        settings = {"mode": "sample", "temperature": 0.5}
    elif name == "earlier-context":
        prompt = "the valve needs"
    elif name == "output-limit":
        settings = {"max_steps": 1}
    elif name == "unseen":
        prompt = "the compressor"
    elif name == "retrained":
        model = train(CORPUS + ["the pump needs repair"] * EXTRA_REPAIR_EXAMPLES)
    return {"case": name, "model_changed_from_base": fingerprint(model) != fingerprint(base),
            "needs_counts": model["needs"], **generate(model, prompt, **settings)}


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("case", nargs="?", choices=("compare", *CASES), default="compare")
    args = parser.parse_args()
    base = train(CORPUS)
    print(json.dumps({"model_kind": "toy word-bigram counts; not a neural LLM",
                      "training_sentences": CORPUS, "sample_seed": SAMPLE_SEED,
                      "learned_counts": base,
                      "results": [experiment(name, base) for name in (CASES if args.case == "compare" else [args.case])]}, indent=2))