#!/usr/bin/env python3
"""Sampling sweep against the llama.cpp router's marco-nano-q8: which settings stop degeneration?

    python -I scripts/router_sampling_sweep.py --out reports/router_sampling/sweep.jsonl

Same prompts and seed per config. Counts outputs that hit max_tokens and loops (16-gram repeat share
of the last 600 tokens >= 0.5, the T1.4 check's rule, computed on characters-as-tokens proxies via the
Marco tokenizer).
"""
from __future__ import annotations

import argparse
import json
import random
import sys
import time
import urllib.error
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
from fp8_check import loop_score  # noqa: E402

URL = "http://127.0.0.1:8081/v1/chat/completions"
DATA = Path("/path/to/data")
CONFIGS = {
    "spec_t0.7_p0.9": {"temperature": 0.7, "top_p": 0.9},
    "t0.3_p0.9": {"temperature": 0.3, "top_p": 0.9},
    "t1.0_p0.9": {"temperature": 1.0, "top_p": 0.9},
    "t0.7_minp0.05": {"temperature": 0.7, "top_p": 1.0, "min_p": 0.05},
    "t1.0_minp0.1": {"temperature": 1.0, "top_p": 1.0, "min_p": 0.1},
    "t0.7_rep1.1": {"temperature": 0.7, "top_p": 0.9, "repeat_penalty": 1.1},
    "t0.7_presence0.5": {"temperature": 0.7, "top_p": 0.9, "presence_penalty": 0.5},
    "t0.7_dry0.8": {"temperature": 0.7, "top_p": 0.9, "dry_multiplier": 0.8, "dry_base": 1.75, "dry_allowed_length": 2},
}


def prompts() -> list[dict]:
    loopy = [json.loads(l) for l in open(DATA / "fp8check/aya_loopy4.jsonl")]
    rng = random.Random(0)
    out = []
    for lang in ("ur", "tr", "bn", "ne"):
        out += rng.sample([r for r in loopy if r["language"] == lang], 10)
    other = [json.loads(l) for l in open(DATA / "bench/aya_marco29_prompts.jsonl")]
    for lang in ("en", "zh", "de", "ar", "ja", "ru", "es", "ko"):
        out += rng.sample([r for r in other if r["language"] == lang], 2)
    out += [json.loads(l) for l in open(DATA / "bench/json_prompts.jsonl")]
    return out


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--out", required=True)
    p.add_argument("--max-tokens", type=int, default=1536)
    p.add_argument("--only", nargs="*")
    a = p.parse_args()
    from tokenizers import Tokenizer
    tok = Tokenizer.from_file(str(DATA / "models/Marco-Nano-Instruct/tokenizer.json"))
    ps = prompts()
    done = set()
    if Path(a.out).exists():
        done = {(r["config"], r["i"]) for r in map(json.loads, open(a.out))}
    with open(a.out, "a") as f:
        for name, cfg in CONFIGS.items():
            if a.only and name not in a.only:
                continue
            t0 = time.time()
            for i, r in enumerate(ps):
                if (name, i) in done:
                    continue
                body = {"model": "marco-nano-q8", "messages": [{"role": "user", "content": r["text"]}],
                        "max_tokens": a.max_tokens, "seed": 1000 + i, **cfg}
                req = urllib.request.Request(URL, json.dumps(body).encode(), {"Content-Type": "application/json"})
                for attempt in range(5):  # the router may be swapping models (one loaded at a time)
                    try:
                        resp = json.loads(urllib.request.urlopen(req, timeout=600).read())
                        break
                    except urllib.error.HTTPError:
                        time.sleep(15 * (attempt + 1))
                else:
                    raise SystemExit(f"router kept failing on {name} prompt {i}")
                text = resp["choices"][0]["message"]["content"] or ""
                ids = tok.encode(text, add_special_tokens=False).ids
                f.write(json.dumps({"config": name, "i": i, "lang": r["language"], "finish": resp["choices"][0]["finish_reason"],
                                    "n_tokens": len(ids), "loop_score": round(loop_score(ids), 3), "text": text},
                                   ensure_ascii=False) + "\n")
                f.flush()
            print(name, round(time.time() - t0), "s", flush=True)


if __name__ == "__main__":
    main()