"""Compare raw and CPT evaluation JSON and enforce the post-CPT gate."""
import argparse
import json
import re
import sys
from pathlib import Path


def repetition_4gram_rate(report: dict) -> float:
    rates = []
    for lang in ("ru", "en"):
        for item in report["generation"].get(lang, []):
            words = re.findall(r"\w+", item["output"].lower())
            grams = [tuple(words[i:i + 4]) for i in range(max(0, len(words) - 3))]
            rates.append(1 - len(set(grams)) / len(grams) if grams else 0.0)
    return sum(rates) / len(rates) if rates else 1.0


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw", required=True)
    ap.add_argument("--cpt", required=True)
    ap.add_argument("--max-ru-ppl-ratio", type=float, default=0.98)
    ap.add_argument("--max-en-ppl-ratio", type=float, default=1.10)
    ap.add_argument("--max-cyrillic-drop", type=float, default=0.0)
    args = ap.parse_args()
    raw = json.loads(Path(args.raw).read_text())
    cpt = json.loads(Path(args.cpt).read_text())
    raw_repetition = repetition_4gram_rate(raw)
    cpt_repetition = repetition_4gram_rate(cpt)

    checks = {
        "ru_ppl_improved_by_2pct": cpt["perplexity"]["ru"] <= raw["perplexity"]["ru"] * args.max_ru_ppl_ratio,
        "en_ppl_regression_within_10pct": cpt["perplexity"]["en"] <= raw["perplexity"]["en"] * args.max_en_ppl_ratio,
        "cyrillic_validity_not_worse": cpt["generation"]["ru_valid_pct"] >= raw["generation"]["ru_valid_pct"] - args.max_cyrillic_drop,
        "english_retained": bool(cpt["generation"]["en_retained"]),
        "generation_repetition_not_degraded": (
            cpt_repetition <= max(0.15, raw_repetition + 0.05)
        ),
    }
    report = {
        "pass": all(checks.values()),
        "checks": checks,
        "raw": raw.get("model"),
        "cpt": cpt.get("model"),
        "ratios": {
            "ru_ppl": cpt["perplexity"]["ru"] / raw["perplexity"]["ru"],
            "en_ppl": cpt["perplexity"]["en"] / raw["perplexity"]["en"],
        },
        "generation_repetition_4gram_rate": {
            "raw": raw_repetition,
            "cpt": cpt_repetition,
            "maximum_allowed": max(0.15, raw_repetition + 0.05),
        },
    }
    print(json.dumps(report, ensure_ascii=False, indent=2))
    sys.exit(0 if report["pass"] else 1)


if __name__ == "__main__":
    main()