"""Sampling from the fixed five-token model; Python standard library only.

First run gpt_model_interface_demo.py in the same directory.
Run: python gpt_sampling_demo.py
Writes gpt_sampling_demo.json beside this script.
"""

import json
import math
from bisect import bisect_right
from collections import Counter
from itertools import accumulate
from pathlib import Path
from random import Random


def softmax(logits, temperature=1.0):
    """Positive temperature, with max subtraction for numerical stability."""
    if not math.isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be finite and positive; use argmax for greedy")
    maximum = max(logits)
    weights = [math.exp((z - maximum) / temperature) for z in logits]
    total = sum(weights)
    return [w / total for w in weights]


def sample_with_u(probabilities, u):
    """Locate uniform u in half-open cumulative intervals; return an original ID."""
    if not 0 <= u < 1:
        raise ValueError("u must be in [0, 1)")
    if any(p < 0 for p in probabilities) or not math.isclose(sum(probabilities), 1):
        raise ValueError("probabilities must be nonnegative and sum to 1")
    boundaries = list(accumulate(probabilities))
    boundaries[-1] = 1.0
    return bisect_right(boundaries, u)


def ranked_ids(probabilities):
    """Break probability ties by original ID; the main example has no ties."""
    return sorted(range(len(probabilities)), key=lambda i: (-probabilities[i], i))


def top_p_ids(probabilities, threshold):
    """Keep the shortest ranked prefix reaching the threshold, including the crossing token."""
    if not 0 < threshold <= 1:
        raise ValueError("top-p must be in (0, 1]")
    selected, mass = [], 0.0
    for token_id in ranked_ids(probabilities):
        selected.append(token_id)
        mass += probabilities[token_id]
        if mass >= threshold:
            break
    return selected


def gpt2_legacy_top_p_ids(probabilities, threshold):
    """Reproduce the old GPT-2 src/sample.py cutoff without running TensorFlow.

    Count cumulative_probs <= p, subtract 1, and clamp the index to 0.
    Keep candidates at or above that score. The main example has no ties.
    """
    order = ranked_ids(probabilities)
    cumulative = list(accumulate(probabilities[i] for i in order))
    index = max(sum(value <= threshold for value in cumulative) - 1, 0)
    cutoff = probabilities[order[index]]
    return [i for i in order if probabilities[i] >= cutoff]


def renormalize(probabilities, selected):
    """Preserve original ID positions; set excluded entries to zero and normalize."""
    allowed = set(selected)
    mass = sum(probabilities[i] for i in allowed)
    return [p / mass if i in allowed else 0.0 for i, p in enumerate(probabilities)]


def main():
    directory = Path(__file__).resolve().parent
    source = json.loads((directory / "gpt_model_interface_demo.json").read_text(encoding="utf-8"))
    logits = source["trace"]["last_logits"][0]
    probabilities = softmax(logits)
    labels = "abcde"
    temperatures = {str(t): softmax(logits, t) for t in (0.5, 1.0, 2.0)}
    top_k = {str(k): renormalize(probabilities, ranked_ids(probabilities)[:k]) for k in (1, 2, 3, 5)}
    top_p = {
        str(p): {"kept_ids": top_p_ids(probabilities, p), "probabilities": renormalize(probabilities, top_p_ids(probabilities, p))}
        for p in (0.6, 0.9, 0.92, 0.95, 1.0)
    }

    # Use the displayed, rounded probabilities for intervals and 100,000 draws only.
    displayed_probabilities = [0.2203, 0.6903, 0.0292, 0.0267, 0.0335]
    examples = [{"u": u, "token_id": sample_with_u(displayed_probabilities, u)} for u in (0.05, 0.5, 0.93, 0.95)]
    rng = Random(12)
    draws = Counter(sample_with_u(displayed_probabilities, rng.random()) for _ in range(100_000))

    assert [item["token_id"] for item in examples] == [0, 1, 2, 3]
    assert top_p["0.9"]["kept_ids"] == [1, 0]
    assert top_p["0.92"]["kept_ids"] == [1, 0, 4]
    assert top_p["0.95"]["kept_ids"] == [1, 0, 4, 2]
    assert gpt2_legacy_top_p_ids(probabilities, 0.9) == [1]
    assert all(max(range(5), key=p.__getitem__) == 1 for p in temperatures.values())
    for p in list(temperatures.values()) + list(top_k.values()) + [entry["probabilities"] for entry in top_p.values()]:
        assert math.isclose(sum(p), 1, abs_tol=1e-12)

    record = {
        "description": "Sampling from fixed teaching logits; top-p includes the token reaching or crossing the threshold.",
        "logits": logits,
        "base_probabilities": probabilities,
        "greedy_id": 1,
        "temperatures": temperatures,
        "top_k": top_k,
        "top_p": top_p,
        "legacy_gpt2_top_p_0.9": {"method": "Independent Python reproduction of the old GPT-2 cutoff; TensorFlow was not run", "kept_ids": [1]},
        "interval_example": {"probabilities": displayed_probabilities, "boundaries": list(accumulate(displayed_probabilities)), "draws": examples},
        "repeated_sampling": {"count": 100_000, "seed": 12, "probabilities": displayed_probabilities, "counts": {labels[i]: draws[i] for i in range(5)}},
        "checks_passed": True,
    }
    destination = directory / "gpt_sampling_demo.json"
    destination.write_text(json.dumps(record, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    print("Checked normalization, temperature ranking, cumulative intervals and filtering.")
    print("top-k=2", [round(100*p, 4) for p in top_k["2"]])
    print("top-p=0.90", top_p["0.9"]["kept_ids"])
    print("Old GPT-2 cutoff reproduced", gpt2_legacy_top_p_ids(probabilities, 0.9))
    print("100000 fixed-distribution draws", record["repeated_sampling"]["counts"])
    print("saved", destination)


if __name__ == "__main__":
    main()