"""Sweep alpha (and a couple of tamed sine schedules) to find the sweet spot
between the pure-probability baseline (generic but coherent) and the
pure-entropy extreme (degenerates into blank-template exploitation).
"""
from entropy_sampler import EntropySampler, sine_alpha_schedule

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

FIXED_ALPHAS = [-1.0, -0.5, 0.0, 0.3, 0.5, 0.7, 1.0]
SINE_CONFIGS = [
    {"period_tokens": 10, "amplitude": 0.4},
    {"period_tokens": 16, "amplitude": 0.6},
]


def summarize(text, diagnostics_log):
    ps = [d["candidates"][d["chosen"]]["p"] for d in diagnostics_log]
    hs = [d["candidates"][d["chosen"]]["h_hat"] for d in diagnostics_log]
    degenerate = sum(1 for d in diagnostics_log if "_" in d["candidates"][d["chosen"]]["token"])
    return {
        "mean_p": sum(ps) / len(ps),
        "mean_h": sum(hs) / len(hs),
        "degenerate_tokens": degenerate,
    }


def run_labeled(sampler, label, alpha):
    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
    text = sampler.generate(PROMPT, max_new_tokens=MAX_NEW_TOKENS, alpha=alpha, verbose=False)
    sampler.step = orig_step

    stats = summarize(text, log)
    print(f"=== {label} ===")
    print(text)
    print(
        f"[mean p={stats['mean_p']:.3f}  mean H_hat={stats['mean_h']:.3f}  "
        f"degenerate(_)={stats['degenerate_tokens']}/{MAX_NEW_TOKENS}]\n"
    )


if __name__ == "__main__":
    sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096)
    try:
        for a in FIXED_ALPHAS:
            run_labeled(sampler, f"alpha = {a:+.1f}", a)

        for cfg in SINE_CONFIGS:
            schedule = sine_alpha_schedule(**cfg)
            label = f"sine period={cfg['period_tokens']} amp={cfg['amplitude']}"
            run_labeled(sampler, label, schedule)
    finally:
        sampler.close()