"""Stage 00 reference engine. Python 3.10+ standard library; no network or GPU.

Run the learner starter before reading this reference. This count-based model is
not GPT. An output path is optional; without it the report prints to stdout.
"""
import argparse
from collections import Counter, defaultdict
import hashlib
import json
import math
from pathlib import Path
import sys

BOS, EOS, UNK = "<bos>", "<eos>", "<unk>"


def mean_nll(probabilities):
    values = list(probabilities)
    if not values:
        raise ValueError("At least one scored target is required")
    if any(not math.isfinite(p) or not 0 <= p <= 1 for p in values):
        raise ValueError("Each target probability must be finite and in [0,1]")
    if 0 in values:
        return math.inf
    return -math.fsum(math.log(p) for p in values) / len(values)


def conditional_probability(count, context_total, vocabulary_size, alpha):
    if (not isinstance(count, int) or not isinstance(context_total, int)
            or not isinstance(vocabulary_size, int) or count < 0
            or context_total < count or vocabulary_size < 1
            or not math.isfinite(alpha) or alpha <= 0):
        raise ValueError("Counts, vocabulary size and smoothing must be valid")
    return (count + alpha) / (context_total + alpha * vocabulary_size)


def load_corpus(path):
    raw = Path(path).read_bytes()
    corpus = json.loads(raw)
    if (not isinstance(corpus, dict) or not isinstance(corpus.get("version"), str)
            or not corpus["version"] or not isinstance(corpus.get("license"), str)):
        raise ValueError("Corpus needs version and license metadata")
    documents = corpus.get("documents", [])
    if not documents:
        raise ValueError("Corpus has no documents")
    seen_ids, seen_text, splits = set(), {}, defaultdict(list)
    for document in documents:
        doc_id, split, text = (document.get(k) for k in ("id", "split", "text"))
        if (not isinstance(doc_id, str) or not doc_id or doc_id in seen_ids
                or split not in ("train", "validation", "test")
                or not isinstance(text, str) or not text.strip()):
            raise ValueError("Invalid document, split or duplicate ID")
        normalized = " ".join(text.casefold().split())
        if normalized in seen_text and seen_text[normalized] != split:
            raise ValueError("Duplicate text crosses a protected split")
        seen_ids.add(doc_id)
        seen_text[normalized] = split
        splits[split].append(document)
    if any(not splits[s] for s in ("train", "validation", "test")):
        raise ValueError("All three document splits must be nonempty")
    return corpus, dict(splits), hashlib.sha256(raw).hexdigest()


def fit(train_documents):
    if not train_documents or any(d.get("split") != "train" for d in train_documents):
        raise ValueError("Fit accepts nonempty training documents only")
    vocabulary = (EOS, UNK, *sorted(set("".join(d["text"] for d in train_documents))))
    unigrams, bigrams, totals = Counter(), Counter(), Counter()
    for document in train_documents:
        previous = BOS
        for target in [*document["text"], EOS]:
            unigrams[target] += 1
            bigrams[previous, target] += 1
            totals[previous] += 1
            previous = target
    return {"vocabulary": vocabulary, "unigrams": unigrams,
            "bigrams": bigrams, "totals": totals}


def score(model, documents, kind, alpha=1.0, loss_fn=mean_nll,
          probability_fn=conditional_probability):
    if kind not in ("uniform", "unigram", "bigram"):
        raise ValueError("Unknown baseline")
    vocabulary = model["vocabulary"]
    probabilities, unknown = [], 0
    for document in documents:
        previous = BOS
        for character in [*document["text"], EOS]:
            target = character if character in vocabulary else UNK
            unknown += int(target == UNK)
            if kind == "uniform":
                probability = 1 / len(vocabulary)
            elif kind == "unigram":
                probability = probability_fn(model["unigrams"][target],
                    sum(model["unigrams"].values()), len(vocabulary), alpha)
            else:
                probability = probability_fn(model["bigrams"][previous, target],
                    model["totals"][previous], len(vocabulary), alpha)
            probabilities.append(probability)
            previous = target
    loss = loss_fn(probabilities)
    return {"nll_nats": loss, "perplexity": math.exp(loss),
            "targets": len(probabilities), "unknown_targets": unknown}


def run_experiment(corpus_path, loss_fn=mean_nll,
                   probability_fn=conditional_probability):
    corpus, splits, checksum = load_corpus(corpus_path)
    model = fit(splits["train"])
    # Predeclared candidates; select only on validation, not test.
    candidates = (0.1, 0.5, 1.0, 2.0)
    validation = {alpha: score(model, splits["validation"], "bigram", alpha,
                              loss_fn, probability_fn) for alpha in candidates}
    selected = min(candidates, key=lambda a: validation[a]["nll_nats"])
    metrics = {split: {kind: score(model, docs, kind, selected, loss_fn, probability_fn)
                       for kind in ("uniform", "unigram", "bigram")}
               for split, docs in splits.items()}
    return {"stage": "00", "implementation_version": "stage00-v1",
            "python_version": sys.version.split()[0], "corpus_version": corpus["version"],
            "corpus_sha256": checksum, "license": corpus["license"],
            "split_ids": {s: [d["id"] for d in docs] for s, docs in splits.items()},
            "vocabulary": list(model["vocabulary"]), "vocabulary_source": "train only",
            "normalization": "none; Unicode code points; BOS context and EOS target per document",
            "loss_unit": "nats per target, including EOS and UNK",
            "smoothing_candidates": list(candidates), "selected_alpha": selected,
            "selection_split": "validation", "test_used_for_selection": False,
            "validation_candidates": validation, "metrics": metrics,
            "limitations": ["Tiny authored corpus; not general-language evidence.",
                            "Count estimates, not a neural model or trained GPT.",
                            "Do not compare perplexities across different tokenization contracts."]}


def main(loss_fn=mean_nll, probability_fn=conditional_probability):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--corpus", type=Path, default=Path(__file__).with_name("corpus.json"))
    parser.add_argument("--output", type=Path, help="Optional new report path; existing files are not overwritten")
    args = parser.parse_args()
    report = run_experiment(args.corpus, loss_fn, probability_fn)
    rendered = json.dumps(report, indent=2, allow_nan=False) + "\n"
    if args.output:
        with args.output.open("x", encoding="utf-8") as handle:
            handle.write(rendered)
    else:
        print(rendered, end="")


if __name__ == "__main__":
    try:
        main()
    except (ValueError, OSError, NotImplementedError) as error:
        print(str(error), file=sys.stderr)
        raise SystemExit(1) from None