"""Larger-N robustness check on confidence_threshold=0.5: the 5-seed sweep
looked clean, but 5 seeds isn't enough to trust a 0% crash rate. Run 20 fresh
seeds (not overlapping the earlier 5-seed batch) against both threshold=0.5
and the always-fork baseline, for a fair speed + crash-rate comparison.
"""
import time

import numpy as np

from entropy_sampler import EntropySampler, sine_alpha_schedule
from tune_alpha import PROMPT, MAX_NEW_TOKENS, summarize
from robustness_test import is_degenerate

N_TRIALS = 20
TOP_K = 12
SEED_START = 300  # fresh seeds, no overlap with the earlier 200-204 batch

DEFAULT_SCHEDULE = lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=-0.2)


def run_trial(sampler, threshold, seed):
    sampler.rng = np.random.default_rng(seed)
    log = []
    orig_step = sampler.step

    def step_and_log(base_logits, a, confidence_threshold=None):
        token, next_logits, diag = orig_step(base_logits, a, confidence_threshold)
        log.append(diag)
        return token, next_logits, diag

    sampler.step = step_and_log
    t0 = time.perf_counter()
    text = sampler.generate(
        PROMPT, max_new_tokens=MAX_NEW_TOKENS, alpha=DEFAULT_SCHEDULE(),
        verbose=False, confidence_threshold=threshold,
    )
    elapsed = time.perf_counter() - t0
    sampler.step = orig_step

    forked_log = [d for d in log if d["forked"]]
    stats = summarize(text, forked_log) if forked_log else {"mean_p": float("nan"), "mean_h": float("nan")}
    stats["degenerate"] = is_degenerate(text)
    stats["text"] = text
    stats["elapsed"] = elapsed
    stats["skip_fraction"] = 1 - len(forked_log) / len(log)
    return stats


def run_batch(sampler, threshold, label):
    trials = [run_trial(sampler, threshold, seed=SEED_START + i) for i in range(N_TRIALS)]
    print(f"=== {label} ===")
    for i, t in enumerate(trials):
        flag = "DEGENERATE" if t["degenerate"] else "ok"
        print(
            f"  seed={SEED_START+i} [{flag:10s}] skip={t['skip_fraction']:.0%} "
            f"time={t['elapsed']:5.1f}s ({MAX_NEW_TOKENS/t['elapsed']:.2f} tok/s) mean_p={t['mean_p']:.3f}"
        )
        print(f"    {t['text']!r}")
    crash_rate = sum(t["degenerate"] for t in trials) / len(trials)
    avg_skip = sum(t["skip_fraction"] for t in trials) / len(trials)
    avg_time = sum(t["elapsed"] for t in trials) / len(trials)
    avg_tokps = MAX_NEW_TOKENS / avg_time
    print(
        f"  -> avg_skip={avg_skip:.0%}  avg_time={avg_time:.1f}s  "
        f"avg_tok/s={avg_tokps:.2f}  crash_rate={crash_rate:.0%} ({sum(t['degenerate'] for t in trials)}/{N_TRIALS})\n"
    )
    return label, avg_skip, avg_time, avg_tokps, crash_rate


if __name__ == "__main__":
    sampler = EntropySampler(top_k=TOP_K, top_n_future=20, n_ctx=4096)
    try:
        rows = []
        rows.append(run_batch(sampler, None, "always fork (baseline), n=20"))
        rows.append(run_batch(sampler, 0.5, "threshold=0.5, n=20"))

        print("=== summary ===")
        baseline_time = rows[0][2]
        for label, avg_skip, avg_time, avg_tokps, crash_rate in rows:
            speedup = baseline_time / avg_time
            print(
                f"{label:32s} skip={avg_skip:5.0%}  {avg_tokps:5.2f} tok/s  "
                f"speedup={speedup:.2f}x  crash_rate={crash_rate:.0%}"
            )
    finally:
        sampler.close()