"""Model-free #103 oracle receipts. Set PYTHONPATH to the immutable MTPLX checkout.

Run in the existing MLX 0.32.2 reference environment, passing the unchanged
Bare-Speed artifact folder. Only tokenizer assets are read; no model is loaded.
"""
import hashlib
import importlib.metadata
import json
import math
import sys
from pathlib import Path

import mlx.core as mx
import numpy as np
from mtplx.generation import _mx_lazy_sample, _sample_from_logits
from mtplx.sampling import SamplerConfig, distribution_from_logits
from mtplx.runtime import _load_tokenizer_resilient


def emit(**receipt):
    print(json.dumps(receipt, ensure_ascii=False), flush=True)


root = Path(sys.argv[1])
emit(case="environment", mlx=importlib.metadata.version("mlx"),
     device=mx.device_info(),
     hashes={n: hashlib.sha256((root/n).read_bytes()).hexdigest()
             for n in ["config.json", "tokenizer.json", "merges.txt", "chat_template.jinja"]})
config = SamplerConfig(temperature=1., top_p=0.5, top_k=2)
values = np.array([math.log(.4), math.log(.3)] + [math.log(.3/62)]*62, dtype=np.float32)
p = distribution_from_logits(values.astype(np.float64), config)
emit(case="sampling_law", probabilities=p[:2].tolist())
counts = [0]*64
for seed in range(128):
    token = _mx_lazy_sample(mx.array(values), config, mx.random.key(seed))
    mx.eval(token)
    counts[int(token.item())] += 1
emit(case="sampling_law_lazy", counts=counts[:2])
for temperature, top_k, top_p in [(.6,20,.95),(.2,8,.8),(1.,2,1.)]:
    row = values * np.float32(temperature)
    config = SamplerConfig(temperature=temperature, top_k=top_k, top_p=top_p)
    probabilities = distribution_from_logits(row.astype(np.float64), config)
    counts = [0]*64
    for seed in range(512):
        token = _mx_lazy_sample(mx.array(row), config, mx.random.key(seed))
        mx.eval(token)
        counts[int(token.item())] += 1
    emit(case="sampling_matrix", temperature=temperature, top_k=top_k, top_p=top_p,
         ordinary_probabilities=probabilities.tolist(), lazy_counts=counts)
for name, values in [
    ("nan", [float("nan")]*64),
    ("positive_inf", [float("inf")]+[0.]*63),
    ("all_negative_inf", [float("-inf")]*64),
    ("valid_mask", [float("-inf"),2.,float("-inf"),1.]+[float("-inf")]*60),
]:
    for temperature in [0.,.6,1.]:
        config = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
        try:
            token, _ = _sample_from_logits(mx.array(values), config, np.random.default_rng(0))
            emit(case="invalid_logits", row=name, temperature=temperature, token=int(token))
        except (ValueError, RuntimeError) as error:
            emit(case="invalid_logits", row=name, temperature=temperature, error=type(error).__name__)
    token = _mx_lazy_sample(mx.array(values), SamplerConfig(temperature=1., top_k=20, top_p=.95), mx.random.key(7))
    mx.eval(token)
    emit(case="invalid_logits_lazy", row=name, token=int(token.item()))

tokenizer = _load_tokenizer_resilient(root, json.loads((root/"config.json").read_text()))
for text in [
    "भारत एक विशाल देश है जिसकी संस्कृति बहुत पुरानी है।",
    "ประเทศไทยมีวัฒนธรรมที่เก่าแก่และสวยงามมาก",
    "إِنَّ اللُّغَةَ الْعَرَبِيَّةَ جَمِيلَةٌ",
    "The quick brown fox jumps over the lazy dog near the river bank.",
    "def foo(x):\n    return x + 1\n",
    "中国的文化历史悠久。",
    "a🙂b ©c ∑d —e …f €g ™h\n",
    "x\u0301y १२३ ١٢٣ ①Ⅷ²\u200dZ",
    " ①", " Ⅷ", " ²", " 𝟘",
]:
    emit(case="tokenizer", text=text, ids=tokenizer.encode(text, add_special_tokens=False))

try:
    from mtplx.attention_context import attention_phase
    from mtplx.attention_math import attention_gate
except ImportError:
    emit(case="bf16_verify_gate", available=False)
    sys.exit(0)
gate = mx.array([-6.84375]*4, dtype=mx.bfloat16)
value = mx.ones((4,), dtype=mx.bfloat16)
with attention_phase("decode_verify"):
    output = attention_gate(value, gate)
    mx.eval(output)
    emit(case="bf16_verify_gate", values=output.astype(mx.float32).tolist())