"""Measure the actual per-token overhead of the entropy sampler vs a plain
single-token-decode baseline (greedy, no forking) on this model/hardware.
"""
import time

import numpy as np

import llama_capi as C
from entropy_sampler import EntropySampler, softmax

PROMPT = "The old lighthouse keeper had seen many storms, but none like"
N_TOKENS = 60


def time_baseline(sampler, n_tokens):
    base_logits = sampler.prime(PROMPT)
    t0 = time.perf_counter()
    for _ in range(n_tokens):
        probs = softmax(base_logits)
        tok = int(np.argmax(probs))  # greedy - cheapest possible comparison point
        sampler._decode([tok], [0], [sampler.n_past], [True])
        sampler.n_past += 1
        base_logits = sampler._logits_row(-1)
        if C.lib.llama_vocab_is_eog(sampler.vocab, tok):
            break
    return time.perf_counter() - t0


def time_entropy_sampler(sampler, n_tokens):
    fork_time = 0.0
    orig_score = sampler._score_candidates

    def timed_score(base_logits):
        nonlocal fork_time
        t0 = time.perf_counter()
        result = orig_score(base_logits)
        fork_time += time.perf_counter() - t0
        return result

    sampler._score_candidates = timed_score
    t0 = time.perf_counter()
    sampler.generate(PROMPT, max_new_tokens=n_tokens, alpha=0.0, verbose=False)
    total_time = time.perf_counter() - t0
    sampler._score_candidates = orig_score
    return total_time, fork_time


if __name__ == "__main__":
    for top_k in (4, 8, 12, 20):
        sampler = EntropySampler(top_k=top_k, top_n_future=20, n_ctx=4096)
        try:
            t_base = time_baseline(sampler, N_TOKENS)
            t_entropy, t_fork = time_entropy_sampler(sampler, N_TOKENS)

            print(f"=== top_k={top_k} ===")
            print(f"  baseline (no forking):   {t_base:6.2f}s  ({N_TOKENS/t_base:5.2f} tok/s)")
            print(f"  entropy sampler:         {t_entropy:6.2f}s  ({N_TOKENS/t_entropy:5.2f} tok/s)")
            print(f"    of which fork+score:   {t_fork:6.2f}s ({100*t_fork/t_entropy:.0f}% of total)")
            print(f"    of which commit decode:{t_entropy - t_fork:6.2f}s")
            print(f"  overhead: {t_entropy/t_base:.2f}x\n")
        finally:
            sampler.close()