#!/usr/bin/env python3
# SPDX-License-Identifier: CC-BY-4.0
# Creator: nlp-editor; source: AI Encyclopedia CN, Chapter 14.
"""Deterministically generate or check the Chapter 14 tokenization/TF-IDF fixture."""

import argparse
import json
import math
from pathlib import Path

DOCUMENTS = (
    ("d1", "猫喜欢鱼"),
    ("d2", "狗喜欢鱼"),
    ("d3", "猫喜欢狗"),
)
LEXICON = ("猫", "狗", "喜欢", "鱼")
VOCABULARY = ("猫", "喜欢", "鱼", "狗")


def longest_match(text, lexicon=LEXICON):
    """Tokenize left-to-right; prefer longest match, then lexicon order."""
    if not text:
        return []
    if not lexicon or any(not term for term in lexicon):
        raise ValueError("lexicon must contain only non-empty terms")
    tokens = []
    position = 0
    while position < len(text):
        matches = [term for term in lexicon if text.startswith(term, position)]
        if not matches:
            raise ValueError(f"no lexicon term matches at character offset {position}")
        longest = max(len(term) for term in matches)
        token = next(term for term in lexicon if len(term) == longest and term in matches)
        tokens.append(token)
        position += len(token)
    return tokens


def count_vector(tokens, vocabulary=VOCABULARY):
    """Return raw term-frequency counts in the declared vocabulary order."""
    unknown = set(tokens).difference(vocabulary)
    if unknown:
        raise ValueError(f"tokens outside vocabulary: {sorted(unknown)!r}")
    return [tokens.count(term) for term in vocabulary]


def build_payload():
    tokenized = [(doc_id, text, longest_match(text)) for doc_id, text in DOCUMENTS]
    tf_vectors = [count_vector(tokens) for _, _, tokens in tokenized]
    document_count = len(tokenized)
    document_frequency = [sum(count > 0 for count in column) for column in zip(*tf_vectors)]
    idf_values = [math.log(document_count / df) for df in document_frequency]
    if any(not math.isfinite(value) for value in idf_values):
        raise ValueError("all IDF values must be finite")
    tfidf_vectors = [
        [count * idf for count, idf in zip(vector, idf_values)]
        for vector in tf_vectors
    ]
    expected_documents = [
        {
            "id": doc_id,
            "text": text,
            "tokens": tokens,
            "term_frequency_vector": tf,
            "tfidf_vector": tfidf,
        }
        for (doc_id, text, tokens), tf, tfidf in zip(tokenized, tf_vectors, tfidf_vectors)
    ]
    expected_terms = [
        {
            "term": term,
            "document_frequency": df,
            "idf_expression": f"ln({document_count}/{df})",
            "idf_value": value,
            "idf_binary64_hex": value.hex(),
        }
        for term, df, value in zip(VOCABULARY, document_frequency, idf_values)
    ]
    return {
        "schema": 1,
        "creator": "nlp-editor",
        "source": "AI Encyclopedia CN, Chapter 14",
        "license": "CC BY 4.0",
        "license_url": "https://creativecommons.org/licenses/by/4.0/",
        "provenance": "Original fictional three-sentence teaching corpus; no people, external corpus, or deployment data.",
        "conventions": {
            "tokenization": "left-to-right longest lexicon match; ties follow lexicon order; unmatched input is an error",
            "vocabulary_order": "first appearance across documents, then first appearance within each document",
            "term_frequency": "raw token count",
            "document_frequency": "number of documents whose raw term frequency is greater than zero",
            "idf": "natural logarithm ln(N/df), with N=3 and no smoothing",
            "tfidf": "raw term frequency multiplied by unrounded binary64 IDF; no normalization or output rounding",
        },
        "primitives": {
            "documents": [{"id": doc_id, "text": text} for doc_id, text in DOCUMENTS],
            "lexicon_order": list(LEXICON),
            "vocabulary": list(VOCABULARY),
            "document_count": document_count,
            "idf_numerator": document_count,
            "idf_logarithm_base": "e",
            "idf_smoothing": False,
        },
        "expected": {
            "documents": expected_documents,
            "terms": expected_terms,
            "document_frequency_vector": document_frequency,
            "idf_vector": idf_values,
        },
    }


def encode():
    return (json.dumps(build_payload(), ensure_ascii=False, indent=2, sort_keys=True) + "\n").encode("utf-8")


def default_output_dir():
    here = Path(__file__).resolve()
    if here.parent.name == "data" and here.parent.parent.name == "public":
        return here.parent
    return here.parents[1] / "public" / "data"


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    mode = parser.add_mutually_exclusive_group(required=True)
    mode.add_argument("--check", action="store_true", help="compare without writing")
    mode.add_argument("--write", action="store_true", help="write the deterministic fixture")
    parser.add_argument("--output-dir", type=Path)
    args = parser.parse_args()
    output = (args.output_dir or default_output_dir()) / "ch14-tiny-corpus.json"
    expected = encode()
    if args.check:
        if not output.is_file() or output.read_bytes() != expected:
            raise SystemExit(f"FAIL: {output} missing or differs")
        print(f"PASS: {output} is deterministic; exact source terms, segmentation, counts, and unsmoothed TF-IDF verified")
        return
    output.parent.mkdir(parents=True, exist_ok=True)
    output.write_bytes(expected)
    print(f"WROTE: {output}")


if __name__ == "__main__":
    main()