"""Repeated-trial robustness check: single samples are noisy, so run each
config N times with independent seeds and measure a crash rate, not just
read one lucky/unlucky sample.

Flags a trial as "degenerate" if it exhibits the failure modes seen in
tune_alpha.py's fixed-alpha sweep: blank-template exploitation (runs of
underscores) or markdown-bold meta-commentary artifacts.
"""
import re

import numpy as np

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

N_TRIALS = 5

CONFIGS = [
    ("sine p16 a0.6 off-0.2", lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=-0.2)),
    ("sine p16 a0.7 off-0.3", lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.7, offset=-0.3)),
    ("sine p12 a0.8 off-0.2", lambda: sine_alpha_schedule(period_tokens=12, amplitude=0.8, offset=-0.2)),
    ("sine p16 a0.6 off 0.0", lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=0.0)),
    ("fixed alpha = 0.0", lambda: 0.0),
    ("fixed alpha = +0.3", lambda: 0.3),
]


# Patterns seen in manual reads that the old blank/markdown-only check missed:
# <think> leaks, reading-comprehension/quiz framing, and reasoning scaffolding.
# Topic hijacks and POV breaks (also observed manually) aren't reliably
# catchable by keyword matching without false-flagging legitimate creative
# swerves - those remain a manual-read gap, not something this function claims
# to cover.
_META_PATTERNS = [
    r"<think",
    r"thinking process",
    r"thinking step",
    r"analyze the request",
    r"question:\s",
    r"most likely outcome",
    r"what does (the word|this) .* mean",
    r"which (ending|option|answer)",
    r"\ba\.\s.{5,60}\bb\.\s",  # "A. ... B. ..." multiple-choice framing
]


def _has_repetition_loop(text, n=6):
    words = text.split()
    seen = set()
    for i in range(len(words) - n + 1):
        gram = tuple(words[i : i + n])
        if gram in seen:
            return True
        seen.add(gram)
    return False


def _has_meta_commentary(text):
    low = text.lower()
    return any(re.search(p, low) for p in _META_PATTERNS)


def is_degenerate(text):
    blank_run = re.search(r"_{2,}", text) is not None or text.count("_") >= 3
    markdown_bold = text.count("**") >= 2
    return (
        blank_run
        or markdown_bold
        or _has_repetition_loop(text)
        or _has_meta_commentary(text)
    )


def run_trial(sampler, alpha_factory, 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
    text = sampler.generate(
        PROMPT, max_new_tokens=MAX_NEW_TOKENS, alpha=alpha_factory(), verbose=False
    )
    sampler.step = orig_step
    stats = summarize(text, log)
    stats["degenerate"] = is_degenerate(text)
    stats["text"] = text
    return stats


if __name__ == "__main__":
    sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096)
    try:
        results = {}
        for label, alpha_factory in CONFIGS:
            trials = []
            for seed in range(N_TRIALS):
                trials.append(run_trial(sampler, alpha_factory, seed=100 + seed))
            results[label] = trials

            print(f"=== {label} ===")
            for i, t in enumerate(trials):
                flag = "DEGENERATE" if t["degenerate"] else "ok"
                print(
                    f"  seed={100+i} [{flag:10s}] mean_p={t['mean_p']:.3f} "
                    f"mean_h={t['mean_h']:.3f}"
                )
                print(f"    {t['text'][:160]!r}")
            crash_rate = sum(t["degenerate"] for t in trials) / len(trials)
            mean_p = sum(t["mean_p"] for t in trials) / len(trials)
            mean_h = sum(t["mean_h"] for t in trials) / len(trials)
            print(f"  -> crash_rate={crash_rate:.0%}  avg_mean_p={mean_p:.3f}  avg_mean_h={mean_h:.3f}\n")

        print("=== summary ===")
        for label, trials in results.items():
            crash_rate = sum(t["degenerate"] for t in trials) / len(trials)
            mean_p = sum(t["mean_p"] for t in trials) / len(trials)
            mean_h = sum(t["mean_h"] for t in trials) / len(trials)
            print(f"{label:24s} crash_rate={crash_rate:.0%}  avg_mean_p={mean_p:.3f}  avg_mean_h={mean_h:.3f}")
    finally:
        sampler.close()