#!/usr/bin/env python3
"""An inspectable word n-gram language model, using only Python 3.10+.

Run: python3 tiny_language_model.py --context 1 --seed 7
     python3 tiny_language_model.py --train my_train.txt --test my_test.txt
Each input line is one sentence. Tokenization is deliberately simple English:
lowercase words (including internal apostrophes) and . ! ? punctuation.
Training counts and vocabulary never include held-out test sentences.
Uses add-one smoothing and backs off to a shorter context when unseen.
NLL is measured in bits per token, including the sentence-end token.
This is a count model, not a neural network or a Transformer.
"""
import argparse
from collections import Counter, defaultdict
import json
import math
from pathlib import Path
import random
import re

TRAIN = [
    'the cat naps .', 'the cat naps .', 'the cat plays .',
    'the dog plays .', 'the dog runs .', 'the dog naps .',
    'a small cat runs .', 'a small dog plays .',
    'the small cat plays .', 'the small dog naps .',
    'a cat naps .', 'a dog runs .',
]
TEST = ['a small cat naps .', 'a dog naps .', 'the small dog runs .', 'a bird plays .']
START, END, UNKNOWN = '<s>', '</s>', '<unk>'


def tokenize(text):
    return re.findall(r"[a-z]+(?:'[a-z]+)*|[.!?]", text.lower())


def read_sentences(path, fallback):
    if path is None:
        lines = fallback
    else:
        source = Path(path)
        if source.stat().st_size > 1_000_000:
            raise ValueError('Use a classroom text file smaller than 1 MB.')
        lines = source.read_text(encoding='utf-8').splitlines()
    sentences = [tokenize(line) for line in lines if line.strip()]
    if not sentences or any(not line for line in sentences):
        raise ValueError('Provide nonempty English sentences, one per line.')
    return sentences


class NGram:
    def __init__(self, sentences, context=1):
        if context not in (0, 1, 2):
            raise ValueError('Context length must be 0, 1 or 2.')
        if not sentences or any(not sentence for sentence in sentences):
            raise ValueError('Training needs at least one nonempty sentence.')
        self.context = context
        self.vocabulary = sorted({token for sentence in sentences for token in sentence} | {END, UNKNOWN})
        self.counts = defaultdict(Counter)
        for sentence in sentences:
            history = [START] * context
            for token in sentence + [END]:
                for n in range(context + 1):
                    self.counts[tuple(history[-n:]) if n else ()][token] += 1
                history.append(token)

    def known(self, token):
        return token if token in self.vocabulary or token == START else UNKNOWN

    def distribution(self, history):
        history = [self.known(token) for token in history]
        for n in range(min(self.context, len(history)), -1, -1):
            key = tuple(history[-n:]) if n else ()
            counter = self.counts.get(key)
            if counter:
                denominator = sum(counter.values()) + len(self.vocabulary)
                return [(word, (counter[word] + 1) / denominator) for word in self.vocabulary]
        raise ValueError('No training counts available.')

    def nll(self, sentences):
        loss, count, unknown_count = 0.0, 0, 0
        for sentence in sentences:
            history = [START] * self.context
            for raw in sentence + [END]:
                token = self.known(raw)
                probability = dict(self.distribution(history))[token]
                loss -= math.log2(probability)
                count += 1
                unknown_count += token == UNKNOWN
                history.append(token)
        if not count:
            raise ValueError('Evaluation needs nonempty sentences.')
        return loss / count, count, unknown_count

    def generate(self, prompt, seed, max_new_tokens=25):
        rng = random.Random(seed)
        words = [self.known(token) for token in tokenize(prompt)]
        history = [START] * self.context + words
        for _ in range(max_new_tokens):
            distribution = self.distribution(history)
            word = rng.choices([w for w, _ in distribution], [p for _, p in distribution], k=1)[0]
            if word == END:
                break
            words.append(word)
            history.append(word)
        return ' '.join(words)


def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('--context', type=int, choices=(0, 1, 2), default=1)
    parser.add_argument('--seed', type=int, default=7)
    parser.add_argument('--train', help='UTF-8 training text: one English sentence per line')
    parser.add_argument('--test', help='Separate held-out test text: one English sentence per line')
    parser.add_argument('--prompt', default='the cat')
    parser.add_argument('--json', action='store_true', help='Print inspectable results as JSON')
    args = parser.parse_args()
    if bool(args.train) != bool(args.test):
        parser.error('Provide both --train and --test, or neither for the built-in split.')
    if args.train and Path(args.train).resolve() == Path(args.test).resolve():
        parser.error('Training and test paths must be different.')
    try:
        training = read_sentences(args.train, TRAIN)
        testing = read_sentences(args.test, TEST)
        if {tuple(s) for s in training} & {tuple(s) for s in testing}:
            raise ValueError('Exact duplicate sentences cross training and test. Reserve different test sentences.')
        model = NGram(training, args.context)
        train_nll, train_tokens, _ = model.nll(training)
        test_nll, test_tokens, unknown = model.nll(testing)
        distribution = sorted(model.distribution(tokenize(args.prompt)), key=lambda pair: (-pair[1], pair[0]))
        results = {'context': args.context, 'seed': args.seed, 'vocabulary': model.vocabulary,
                   'train_sentences': len(training), 'test_sentences': len(testing),
                   'train_tokens': train_tokens, 'test_tokens': test_tokens,
                   'train_nll_bits': train_nll, 'test_nll_bits': test_nll,
                   'test_unknown_tokens': unknown, 'prompt': args.prompt,
                   'continuations': distribution,
                   'sample': model.generate(args.prompt, args.seed)}
        if args.json:
            print(json.dumps(results, ensure_ascii=False, indent=2))
        else:
            print(f'Word n-gram baseline | context={args.context} | seed={args.seed}')
            print(f'Train/test sentences: {len(training)}/{len(testing)}; vocabulary: {len(model.vocabulary)}')
            print('Vocabulary:', ', '.join(model.vocabulary))
            print(f'Train NLL: {train_nll:.4f} bits/token; test NLL: {test_nll:.4f} bits/token')
            print(f'Test unknown tokens: {unknown}; end-of-sentence tokens are included in NLL.')
            print(f'Continuations after {args.prompt!r}:')
            for word, probability in distribution:
                print(f'  {word:>10} {probability:.4f}')
            print('Sample:', results['sample'])
            print('Fluency is not factual truth. Keep test data out of training and model selection.')
    except (OSError, UnicodeError, ValueError) as error:
        parser.error(str(error))


if __name__ == '__main__':
    main()