"""Model-free #105 receipt from the pinned MTPLX/MLX oracle, never production."""
import json
import subprocess
from pathlib import Path

import mlx.core as mx
import numpy as np
from mtplx import generation
from mtplx.sampling import NonFiniteLogitsError, SamplerConfig

reference = Path(generation.__file__).resolve().parent.parent
assert subprocess.check_output(["git", "-C", str(reference), "rev-parse", "HEAD"], text=True).strip() == "9882703f3105363ddc37eca9f97aa09a1d387112"
assert mx.__version__ == "0.32.2"
rows = {
    "nan": np.full(64, np.nan, np.float32),
    "positive_inf": np.array([np.inf] + [0.] * 63, np.float32),
    "all_negative_inf": np.full(64, -np.inf, np.float32),
    "valid_mask": np.array([-np.inf, 2., -np.inf, 1.] + [-np.inf] * 60, np.float32),
}
for name, values in rows.items():
    row = mx.array(values)
    for temperature in (0., .6, 1.):
        config = SamplerConfig(temperature=temperature, top_p=.95, top_k=20)
        try:
            token, _ = generation._sample_from_logits(row, config, np.random.default_rng(0))
            result = {"token": token}
        except NonFiniteLogitsError as error:
            result = {"error": str(error), "nan": error.nan_count, "inf": error.inf_count, "vocab": error.vocab_size}
        print(json.dumps(dict(case="ordinary", row=name, temperature=temperature, **result)), flush=True)
    config = SamplerConfig(temperature=1., top_p=.95, top_k=20)
    token = int(generation._mx_lazy_sample(row, config, mx.random.key(7)).item())
    print(json.dumps(dict(case="lazy", row=name, token=token)), flush=True)
    # The specialized MTP greedy helper still lacks the ordinary guard at this
    # revision. Keep the conflict with #105 visible; never label this parity.
    token, _ = generation._greedy_draft_token_and_top_values(row.reshape(1,1,-1), topk=2)
    print(json.dumps(dict(case="mtp_greedy_specialized", row=name, token=token)), flush=True)