#!/usr/bin/env python3
"""What a task set's answers cost a model in tokens, written four ways.

    python3 -I token_cost.py answers.json --tokenizer NAME=MODEL.gguf [...] --tool PATH [--out FILE]

answers.json is what `wai-quantum eval corpus` writes: each task's answer as
plain OpenQASM (the program alone, in a fenced block), in the answer view
layer (whole, and with the request alone), and as a job's JSON. For each tokenizer, the vocabulary of a model read
from its GGUF file by a tokenizer tool (one that takes `-m MODEL -f FILE --ids
--no-bos --log-disable` and prints the ids as a list), this counts the tokens of
every answer in each form, and reports the totals, the per-answer means, each
form against plain, and the bytes. A tokenizer's size is its vocabulary's, read
from the GGUF file's tokenizer.ggml.tokens. Standard library only; it reads the
model files and writes nothing but its report.
"""
import argparse
import json
import os
import struct
import subprocess
import sys
import tempfile

FORMS = ["plain", "views", "views_request", "job"]


def vocabulary(path):
    """The number of tokens in a GGUF file's tokenizer.ggml.tokens, or None."""
    sizes = {0: 1, 1: 1, 2: 2, 3: 2, 4: 4, 5: 4, 6: 4, 7: 1, 10: 8, 11: 8, 12: 8}
    with open(path, "rb") as f:
        if f.read(4) != b"GGUF":
            return None
        (version,) = struct.unpack("<I", f.read(4))
        if version < 2:
            return None
        _tensors, kvs = struct.unpack("<QQ", f.read(16))

        def string():
            (n,) = struct.unpack("<Q", f.read(8))
            return f.read(n)

        def skip(t):
            if t in sizes:
                f.seek(sizes[t], 1)
            elif t == 8:
                string()
            elif t == 9:
                et, n = struct.unpack("<IQ", f.read(12))
                if et in sizes:
                    f.seek(sizes[et] * n, 1)
                else:
                    for _ in range(n):
                        skip(et)
            else:
                raise ValueError(f"GGUF value type {t}")

        for _ in range(kvs):
            key = string().decode("utf-8", "replace")
            (t,) = struct.unpack("<I", f.read(4))
            if key == "tokenizer.ggml.tokens" and t == 9:
                _et, n = struct.unpack("<IQ", f.read(12))
                return n
            skip(t)
    return None


def count(tool, model, text):
    with tempfile.NamedTemporaryFile("w", suffix=".txt", delete=False, encoding="utf-8") as f:
        f.write(text)
        name = f.name
    try:
        out = subprocess.run([tool, "-m", model, "-f", name, "--ids", "--no-bos", "--log-disable"], capture_output=True, text=True, check=True).stdout
    finally:
        os.unlink(name)
    line = [l for l in out.splitlines() if l.strip().startswith("[")][-1]
    return len(json.loads(line))


def main():
    p = argparse.ArgumentParser(description=__doc__.split("\n")[1])
    p.add_argument("answers")
    p.add_argument("--tokenizer", action="append", required=True, metavar="NAME=MODEL.gguf")
    p.add_argument("--tool", required=True, metavar="PATH", help="the tokenizer tool")
    p.add_argument("--out")
    a = p.parse_args()
    answers = json.load(open(a.answers, encoding="utf-8"))["answers"]
    report = {
        "format": "wai.quantum.token-cost/1",
        "about": "The tokens a task set's answers take, each written four ways: plain (the program alone, in a fenced block), views (the answer view layer: the program, its circuit, a sample when it measures, a verdict for each intent, a seal), views_request (the layer with the sample and the verdicts alone) and job (the request an agent sends a run tool, as JSON). views_request and job carry the same program and the same request; views adds the circuit and the seal; plain carries the program alone.",
        "answers": len(answers),
        "bytes": {form: sum(len(x[form].encode("utf-8")) for x in answers) for form in FORMS},
        "tokenizers": [],
    }
    for spec in a.tokenizer:
        name, model = spec.split("=", 1)
        counts = {form: [count(a.tool, model, x[form]) for x in answers] for form in FORMS}
        totals = {form: sum(v) for form, v in counts.items()}
        report["tokenizers"].append({
            "name": name,
            "vocabulary": vocabulary(model),
            "totals": totals,
            "mean": {form: round(totals[form] / len(answers), 1) for form in FORMS},
            "against_plain": {form: round(totals[form] / totals["plain"], 3) for form in FORMS},
            "views_request_against_job": round(totals["views_request"] / totals["job"], 3),
            "per_answer": [{"name": x["name"], **{form: counts[form][i] for form in FORMS}} for i, x in enumerate(answers)],
        })
        print(f"{name} ({report['tokenizers'][-1]['vocabulary']} tokens): " + ", ".join(f"{f} {totals[f]}" for f in FORMS), file=sys.stderr)
    text = json.dumps(report, indent=2) + "\n"
    if a.out:
        open(a.out, "w", encoding="utf-8").write(text)
    else:
        sys.stdout.write(text)


if __name__ == "__main__":
    main()