"""Choose compile settings and the confidence policy by simulated typing on held-out prompts.

The weights are always produced by the same closed-form compile; this script only
compares a handful of compile settings (context length, smoothing, source weights,
vocabulary, pruning) and policy thresholds, and reports which one saves the most
keystrokes while keeping precision >= TARGET_PRECISION.

Simulation: each held-out prompt (10% of data/prompts.txt, never compiled in) is
"typed" one character at a time. After every keystroke the policy shows ghost text
or nothing. A shown suggestion is correct only if the user's next characters start
with the whole ghost text. Keystrokes saved replays the prompt accepting every
correct suggestion with one Tab press.

    python scripts/sweep.py            # writes docs/sweep.json
"""

import json
import sys
import time
import zlib
from dataclasses import asdict, replace
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))
import build_weights as bw  # noqa: E402
import suggest as sg  # noqa: E402

TARGET_PRECISION = 0.90
MAX_EXPORT_MB = 50
TAUS = [round(0.5 + 0.01 * i, 2) for i in range(46)]  # 0.50 .. 0.95
SUPPORTS = [0, 2, 5, 10]
BASE = bw.Settings(n_prev=3, smoothing="wb", phrase_weight=5, prompt_weight=1, reply_weight=1, max_vocab=8000, min_count=2)


def heldout_texts() -> list[str]:
    seen, out = set(), []
    for line in (bw.DATA_DIR / "prompts.txt").read_text().splitlines():
        toks = bw.pieces(line)
        if toks and zlib.crc32(" ".join(toks).encode()) % 100 < 10:
            text = bw.NORMALIZER.normalize_str(line.strip())
            if text not in seen:
                seen.add(text)
                out.append(text)
    return out


def train_sources() -> dict:
    src = bw.load_sources()
    src["prompts"], _ = bw.heldout_split(src["prompts"])
    return src


def record(sugg: sg.Suggester, texts: list[str]) -> list[tuple[str, list]]:
    """For every prompt and every keystroke position, the unthresholded suggestion chain."""
    return [(t, [None] + [sugg.steps(t[:i]) for i in range(1, len(t))]) for t in texts]


def evaluate(records, policy: sg.Policy) -> dict:
    shown = correct = positions = chars = strokes = 0
    for text, steps in records:
        ghosts = [None] + [sg.Suggester.apply(st, policy) for st in steps[1:]]
        for i in range(1, len(text)):
            positions += 1
            if ghosts[i]:
                shown += 1
                correct += text.startswith(ghosts[i], i)
        pos, n = 0, 0
        while pos < len(text):
            pos, n = pos + 1, n + 1
            g = ghosts[pos] if pos < len(text) else ""
            if g and text.startswith(g, pos):
                pos, n = pos + len(g), n + 1
        chars, strokes = chars + len(text), strokes + n
    return {
        "precision": round(correct / shown, 4) if shown else None,
        "coverage": round(shown / positions, 4),
        "keystrokes_saved": round(1 - strokes / chars, 4),
        "shown": shown,
    }


def best_policy(records) -> dict:
    results = []
    for tau in TAUS:
        for sup in SUPPORTS:
            m = evaluate(records, sg.Policy(tau=tau, min_support=sup))
            results.append({"tau": tau, "min_support": sup, **m})
    ok = [r for r in results if r["precision"] is not None and r["precision"] >= TARGET_PRECISION]
    best = max(ok, key=lambda r: r["keystrokes_saved"]) if ok else None
    return {"best": best, "grid": results}


def export_mb(ngram: bw.NGram, vocab: list[str], oov: dict) -> float:
    """Size of web/model.bin for these tables (see export_web.py)."""
    size = sum(len(p.encode()) + 1 for p in vocab) + 8 * len(vocab) + sum(len(p.encode()) + 5 for p in oov)
    for k in range(1, ngram.order + 1):
        for table in ngram.tables[k].values():
            size += 2 * k + 8 + 6 * len(table)
    return round(size / 1e6, 1)


def run(s: bw.Settings, src: dict, texts: list[str]) -> dict:
    t0 = time.time()
    vocab, ngram, stats = bw.compile_tables(s, src)
    oov = sg.oov_prefix_rates(s, src, vocab)
    recs = record(sg.Suggester(vocab, ngram, oov), texts)
    pol = best_policy(recs)
    out = {
        "settings": asdict(s),
        **stats,
        "export_mb": export_mb(ngram, vocab, oov),
        "best": pol["best"],
        "seconds": round(time.time() - t0, 1),
        "grid": pol["grid"],
    }
    b = out["best"] or {}
    print(
        f"n_prev={s.n_prev} {s.smoothing} w=({s.phrase_weight},{s.prompt_weight},{s.reply_weight}) "
        f"V={s.max_vocab} min={s.min_count} | {out['export_mb']} MB | "
        f"saved={b.get('keystrokes_saved')} prec={b.get('precision')} cov={b.get('coverage')} "
        f"tau={b.get('tau')} sup={b.get('min_support')} | {out['seconds']}s",
        flush=True,
    )
    return out


def score(r: dict) -> float:
    if r["best"] is None or r["export_mb"] > MAX_EXPORT_MB:
        return -1
    return r["best"]["keystrokes_saved"]


def main() -> None:
    src, texts = train_sources(), heldout_texts()
    print(f"held-out prompts: {len(texts)}, characters: {sum(map(len, texts))}", flush=True)
    runs = []

    print("== stage A: context length x smoothing", flush=True)
    stage = [run(replace(BASE, n_prev=n, smoothing=sm), src, texts) for sm in ["wb", "kn"] for n in [3, 4, 5, 7, 9]]
    runs += stage
    best = max(stage, key=score)

    print("== stage B: source weights", flush=True)
    base = bw.Settings(**best["settings"])
    stage = [
        run(replace(base, phrase_weight=pw, reply_weight=rw), src, texts)
        for pw in [1.0, 5.0]
        for rw in [0.0, 0.3, 1.0]
        if (pw, rw) != (base.phrase_weight, base.reply_weight)
    ]
    runs += stage
    best = max(runs, key=score)

    print("== stage C: vocabulary x pruning", flush=True)
    base = bw.Settings(**best["settings"])
    stage = [
        run(replace(base, max_vocab=v, min_count=mc), src, texts)
        for v in [3000, 8000, 16000]
        for mc in [2, 3]
        if (v, mc) != (base.max_vocab, base.min_count)
    ]
    runs += stage
    best = max(runs, key=score)

    print("== stage D: context length again, with the final settings", flush=True)
    base = bw.Settings(**best["settings"])
    stage = [run(replace(base, n_prev=n), src, texts) for n in [3, 4, 5, 7, 9] if n != base.n_prev]
    runs += stage
    best = max(runs, key=score)

    out = ROOT / "docs" / "sweep.json"
    out.write_text(json.dumps({
        "target_precision": TARGET_PRECISION,
        "max_export_mb": MAX_EXPORT_MB,
        "heldout_prompts": len(texts),
        "best": {k: v for k, v in best.items() if k != "grid"},
        "runs": [{k: v for k, v in r.items() if k != "grid"} for r in runs],
        "best_grid": best["grid"],
    }, indent=2) + "\n")
    print("BEST", json.dumps({k: v for k, v in best.items() if k != "grid"}, indent=2))


if __name__ == "__main__":
    main()