"""Deterministic teaching byte BPE, not a production-model tokenizer.

Python 3.10+ standard library. See README.md before consulting this reference.
No network, model download or package installation. Default: print JSON report.
"""
import argparse
from collections import Counter
from dataclasses import dataclass
import hashlib
import json
import math
from pathlib import Path
import re
import sys

BOS, EOS, MERGE_START = 256, 257, 258
MAX_BYTES, MAX_MERGES = 262144, 128
FORMAT = "intuitiveai-byte-bpe-v1"
POLICY = "strict-utf8; no-normalization; document-wide; literal-special-text"


def pair_counts(sequences):
    counts = Counter()
    for sequence in sequences:
        counts.update(zip(sequence, sequence[1:]))
    return counts


def merge_pair(ids, pair, new_id):
    result, index = [], 0
    while index < len(ids):
        if index + 1 < len(ids) and (ids[index], ids[index + 1]) == pair:
            result.append(new_id)
            index += 2
        else:
            result.append(ids[index])
            index += 1
    return result


def apply_merges(ids, merges):
    result = list(ids)
    for rank, pair in enumerate(merges):
        result = merge_pair(result, pair, MERGE_START + rank)
    return result


def byte_input(text):
    if not isinstance(text, str):
        raise ValueError("Text must be a string")
    raw = text.encode("utf-8", errors="strict")
    if len(raw) > MAX_BYTES:
        raise ValueError("Teaching limit: input exceeds 256 KiB")
    return raw


@dataclass(frozen=True)
class Tokenizer:
    merges: tuple = ()
    corpus_sha256: str = ""
    training_ids: tuple = ()

    def __post_init__(self):
        if (not isinstance(self.merges, tuple) or len(self.merges) > MAX_MERGES
                or not isinstance(self.training_ids, tuple)
                or any(not isinstance(x, str) or not x for x in self.training_ids)
                or len(set(self.training_ids)) != len(self.training_ids)
                or not isinstance(self.corpus_sha256, str)
                or (self.corpus_sha256 and not re.fullmatch(r"[0-9a-f]{64}", self.corpus_sha256))):
            raise ValueError("Invalid tokenizer metadata")
        seen = set()
        lengths = {i: 1 for i in range(256)}
        expanded_bytes = 256
        for rank, pair in enumerate(self.merges):
            if (not isinstance(pair, tuple) or len(pair) != 2
                    or any(type(x) is not int or x < 0 or x >= MERGE_START + rank
                           or x in (BOS, EOS) for x in pair) or pair in seen):
                raise ValueError("Invalid, duplicate or forward-referencing merge")
            seen.add(pair)
            length = lengths[pair[0]] + lengths[pair[1]]
            expanded_bytes += length
            if length > MAX_BYTES or expanded_bytes > 1048576:
                raise ValueError("Expanded vocabulary exceeds teaching limit")
            lengths[MERGE_START + rank] = length

    @property
    def content_ids(self):
        return (*range(256), *range(MERGE_START, MERGE_START + len(self.merges)))

    def encode_bytes(self, raw, apply_fn=apply_merges):
        if not isinstance(raw, bytes) or len(raw) > MAX_BYTES:
            raise ValueError("Provide bytes within the 256 KiB teaching limit")
        return apply_fn(list(raw), self.merges)

    def encode(self, text, apply_fn=apply_merges):
        return self.encode_bytes(byte_input(text), apply_fn)

    def decode_bytes(self, ids):
        vocabulary = {i: bytes([i]) for i in range(256)}
        for rank, (left, right) in enumerate(self.merges):
            vocabulary[MERGE_START + rank] = vocabulary[left] + vocabulary[right]
        parts, size = [], 0
        for token in ids:
            if type(token) is not int or token not in vocabulary:
                raise ValueError("Content decode rejects unknown and reserved IDs")
            size += len(vocabulary[token])
            if size > MAX_BYTES:
                raise ValueError("Teaching limit: decoded output exceeds 256 KiB")
            parts.append(vocabulary[token])
        return b"".join(parts)

    def decode(self, ids):
        return self.decode_bytes(ids).decode("utf-8", errors="strict")

    def payload(self):
        return {"format": FORMAT, "policy": POLICY, "merges": [list(p) for p in self.merges],
                "corpus_sha256": self.corpus_sha256, "training_ids": list(self.training_ids)}

    def canonical_json(self):
        return json.dumps(self.payload(), sort_keys=True, separators=(",", ":"), ensure_ascii=True)

    def digest(self):
        return hashlib.sha256(self.canonical_json().encode("utf-8")).hexdigest()

    def save(self, path):
        with Path(path).open("x", encoding="utf-8") as handle:
            handle.write(self.canonical_json() + "\n")

    @classmethod
    def from_payload(cls, payload):
        expected = {"format", "policy", "merges", "corpus_sha256", "training_ids"}
        if (not isinstance(payload, dict) or set(payload) != expected
                or payload["format"] != FORMAT or payload["policy"] != POLICY
                or not isinstance(payload["merges"], list)
                or any(not isinstance(p, list) for p in payload["merges"])
                or not isinstance(payload["training_ids"], list)):
            raise ValueError("Unsupported or malformed tokenizer format/policy")
        return cls(tuple(tuple(p) for p in payload["merges"]), payload["corpus_sha256"],
                   tuple(payload["training_ids"]))

    @classmethod
    def load(cls, path):
        return cls.from_payload(json.loads(read_bounded(path)))


def train(documents, budget=16, checksum="", count_fn=pair_counts, merge_fn=merge_pair):
    if type(budget) is not int or not 0 <= budget <= MAX_MERGES:
        raise ValueError("Merge budget must be an integer from 0 to 128")
    if (not documents or any(not isinstance(d, dict) or d.get("split") != "train"
                            or not isinstance(d.get("id"), str) or not d["id"]
                            or not isinstance(d.get("text"), str) for d in documents)):
        raise ValueError("Fit accepts labeled training documents only")
    sequences = [list(byte_input(d["text"])) for d in documents]
    if sum(map(len, sequences)) > MAX_BYTES or not any(sequences):
        raise ValueError("Training text must be nonempty and within 256 KiB")
    merges = []
    for _ in range(budget):
        counts = count_fn(sequences)
        if not counts:
            break
        pair = min(counts, key=lambda p: (-counts[p], p))
        if counts[pair] < 2:
            break
        new_id = MERGE_START + len(merges)
        sequences = [merge_fn(s, pair, new_id) for s in sequences]
        merges.append(pair)
    return Tokenizer(tuple(merges), checksum, tuple(d["id"] for d in documents))


def read_bounded(path):
    with Path(path).open("rb") as handle:
        raw = handle.read(1048577)
    if len(raw) > 1048576:
        raise ValueError("JSON exceeds 1 MiB teaching limit")
    return raw


def load_corpus(path):
    raw = read_bounded(path)
    corpus = json.loads(raw)
    if (not isinstance(corpus, dict) or not isinstance(corpus.get("version"), str)
            or not isinstance(corpus.get("license"), str)
            or not isinstance(corpus.get("documents"), list)):
        raise ValueError("Corpus needs metadata and documents")
    seen_ids, seen_text = set(), {}
    splits = {name: [] for name in ("train", "validation", "test")}
    for doc in corpus["documents"]:
        if not isinstance(doc, dict):
            raise ValueError("Malformed document")
        doc_id, split, text = (doc.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 not isinstance(split, str) or split not in splits
                or not isinstance(text, str) or not text.strip()):
            raise ValueError("Invalid document, split or repeated ID")
        byte_input(text)
        normalized = " ".join(text.casefold().split())
        if normalized in seen_text and seen_text[normalized] != split:
            raise ValueError("Exact duplicate text crosses protected splits")
        seen_ids.add(doc_id)
        seen_text[normalized] = split
        splits[split].append(doc)
    if any(not docs for docs in splits.values()):
        raise ValueError("Every split must be nonempty")
    return corpus, splits, hashlib.sha256(raw).hexdigest()


def representation_stats(tokenizer, documents, apply_fn=apply_merges):
    raw_bytes = sum(len(byte_input(d["text"])) for d in documents)
    characters = sum(len(d["text"]) for d in documents)
    tokens = sum(len(tokenizer.encode(d["text"], apply_fn)) for d in documents)
    return {"utf8_bytes": raw_bytes, "code_points": characters, "content_tokens": tokens,
            "bytes_per_content_token": raw_bytes / tokens if tokens else None,
            "tokens_per_code_point": tokens / characters if characters else None,
            "round_trip": all(tokenizer.decode(tokenizer.encode(d["text"], apply_fn)) == d["text"]
                              for d in documents)}


def baseline_comparison(tokenizer, splits, apply_fn=apply_merges):
    """Compare canonical encodings; not normalized probabilities over byte strings."""
    results = {}
    for name, encoder in (("raw_bytes", Tokenizer()), ("bpe", tokenizer)):
        vocab = (*encoder.content_ids, EOS)  # BOS is context-only, never a target.
        unigrams, bigrams, totals = Counter(), Counter(), Counter()
        for doc in splits["train"]:
            previous = BOS
            for target in [*encoder.encode(doc["text"], apply_fn), EOS]:
                unigrams[target] += 1
                bigrams[previous, target] += 1
                totals[previous] += 1
                previous = target
        tables = {}
        for split, docs in splits.items():
            metrics = {}
            for kind in ("uniform", "unigram", "bigram"):
                losses = []
                for doc in docs:
                    previous = BOS
                    for target in [*encoder.encode(doc["text"], apply_fn), EOS]:
                        if kind == "uniform":
                            probability = 1 / len(vocab)
                        elif kind == "unigram":
                            probability = (unigrams[target] + 0.1) / (sum(unigrams.values()) + 0.1 * len(vocab))
                        else:
                            probability = (bigrams[previous, target] + 0.1) / (totals[previous] + 0.1 * len(vocab))
                        losses.append(-math.log(probability))
                        previous = target
                nll = math.fsum(losses) / len(losses)
                nbytes = sum(len(byte_input(d["text"])) for d in docs)
                metrics[kind] = {"targets_including_eos": len(losses), "nll_nats_per_target": nll,
                                 "perplexity_per_target": math.exp(nll),
                                 "canonical_path_nats_per_utf8_byte": math.fsum(losses) / nbytes}
            tables[split] = metrics
        results[name] = {"output_vocabulary_size_including_eos": len(vocab), "metrics": tables}
    return results


def default_corpus():
    local = Path(__file__).with_name("corpus.json")
    return local if local.exists() else Path(__file__).parent.parent / "gpt-stage-00" / "corpus.json"


def run_experiment(path, count_fn=pair_counts, merge_fn=merge_pair, apply_fn=apply_merges):
    corpus, splits, checksum = load_corpus(path)
    model = train(splits["train"], checksum=checksum, count_fn=count_fn, merge_fn=merge_fn)
    probes = ["नमस्ते", "東京", "👩🏽‍💻", "e\u0301", "<eos>", "x = 1\n"]
    return {"stage": "01", "implementation_version": FORMAT, "python_version": sys.version.split()[0],
            "corpus_version": corpus["version"], "corpus_sha256": checksum, "corpus_license": corpus["license"],
            "split_ids": {s: [d["id"] for d in docs] for s, docs in splits.items()},
            "fit_split": "train", "merge_budget": 16, "minimum_pair_frequency": 2,
            "tokenizer": model.payload(), "tokenizer_sha256": model.digest(),
            "representation": {s: representation_stats(model, docs, apply_fn) for s, docs in splits.items()},
            "transfer_probes": [{"text": text, "ids": model.encode(text, apply_fn),
                                 **representation_stats(model, [{"text": text}], apply_fn)} for text in probes],
            "baseline_contract": {"alpha": 0.1, "selection": "predeclared, no test tuning",
                                  "bos": "context only", "eos": "scored once per document",
                                  "cross_tokenizer_token_perplexity_comparable": False,
                                  "byte_score": "canonical path including EOS; NOT normalized string likelihood"},
            "baselines": baseline_comparison(model, splits, apply_fn),
            "limitations": ["Tiny English fixture, not a multilingual efficiency benchmark.",
                            "Document-wide naive BPE, no regex pre-tokenization or model compatibility.",
                            "Lossless encoding does not establish language competence or security."]}


def main(count_fn=pair_counts, merge_fn=merge_pair, apply_fn=apply_merges):
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--corpus", type=Path, default=default_corpus())
    parser.add_argument("--output", type=Path, help="New JSON report path; existing files are not overwritten")
    args = parser.parse_args()
    report = run_experiment(args.corpus, count_fn, merge_fn, apply_fn)
    rendered = json.dumps(report, indent=2, ensure_ascii=True, 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