#!/usr/bin/env python3
"""exp15 — reflow-invariant omission metric + Q2b re-attribution ($0). Design §D3 + orchestrator amendments.

The metric asks, per SOURCE sentence-atom, whether the atom appears ANYWHERE in the target (position-
independent = reflow-invariant, so legitimate merge/reorder is NOT penalized). Primary atom = glossary
ENTITY (design-relevant: entities/terms = 62.5% doc-level errors, BWB).

ORCHESTRATOR AMENDMENTS (2026-07-17):
 (1) Entity match by LEMMA (pymorphy3), NOT membership in decl.forms. A form present but outside the
     pinned decl.forms = `inflection_gap` FLAG for span-check, NOT an omission (else the D30.9 §B5
     inflection_gap class fabricates ~7-8% false omissions).
 (2) Pre-registered variance-check: entity-omission ~0 across ALL sizes => report "INSTRUMENT CEILING",
     never "size is safe" (anti-false-negative; the D32.4 lesson).
 (3) Entity-free (holistic) content-drop detection runs ONLY if a crude content-length FLOOR backstop
     fires; otherwise complex drops are span-judge work (don't over-engineer).
 (4) Provenance to BYTE identity: source = records.json chapter_source join (== exp14_sizecurve.py:30-32,
     the exact bytes sizecurve translated). Asserted below.

VALIDATION (mandated): the metric is validated on a hand-proofread whole-chapter output BEFORE scoring
arms — run with --validate <ch> to dump per-entity verdicts for manual read; only after that, --score.

Run: eval/.venv/bin/python eval/exp15/reflow_omission.py --validate 13   # dump for hand-proofread
     eval/.venv/bin/python eval/exp15/reflow_omission.py --score          # Q2b table (after validation)
"""
from __future__ import annotations
import argparse
import json
import re
from functools import lru_cache
from pathlib import Path

import pymorphy3
import yaml

BOOK = Path("/home/ubuntu/books/gu-zhenren")
SEED = BOOK / "guzhenren-seed-v2.yaml"
RECORDS = BOOK / "rerun" / "records.json"
SIZE_DIR = BOOK / "exp14" / "sizecurve"
OUT = BOOK / "exp15" / "q2b_omission.json"
SIZES = ["800c", "1600c", "3200c", "whole"]
CHAPTERS = [13, 17]
CONTENT_FLOOR = 0.5   # amendment 3: target ru-char / fertility-expected < this => holistic-drop backstop fires

_MORPH = pymorphy3.MorphAnalyzer()
CYR = re.compile(r"[А-Яа-яЁё]+")


@lru_cache(maxsize=200_000)
def lemma(tok: str) -> str:
    return _MORPH.parse(tok)[0].normal_form


def ru_tokens(text: str):
    return CYR.findall(text)


def target_index(text: str):
    """surface tokens, lemma set, and lemma->surfaces map of a ru target."""
    surf = ru_tokens(text)
    lem_of = {s: lemma(s.lower()) for s in set(surf)}
    lemset = set(lem_of.values())
    lem2surf = {}
    for s in surf:
        lem2surf.setdefault(lem_of[s], set()).add(s)
    return {"lemset": lemset, "lem2surf": lem2surf}


def load_seed_entities():
    d = yaml.safe_load(SEED.read_text(encoding="utf-8"))
    ents = []
    for t in d["terms"]:
        src = t.get("src")
        dst = (t.get("dst") or "").strip()
        if not src or not dst:
            continue
        srcs = [src] + [a.get("alias") for a in (t.get("aliases") or []) if a.get("alias")]
        forms = set()
        decl = t.get("decl") or {}
        for f in (decl.get("forms") or []):
            forms.add(f)
        forms.add(dst)
        # distinctive ru lemmas of the dst (content tokens len>=3)
        dtoks = ru_tokens(dst)
        content = [x for x in dtoks if len(x) >= 3]
        lemmas = [lemma(x.lower()) for x in (content or dtoks)]
        ents.append({"src_variants": srcs, "dst": dst, "type": t.get("type"),
                     "decl_forms": forms, "lemmas": lemmas, "dst_tokens": dtoks})
    return ents


def chapter_source(ch):
    """BYTE-identical to exp14_sizecurve.py:30-32 (join of records.json chunk sources by '\\n')."""
    recs = {(r["chapter"], r["chunk_idx"]): r for r in json.load(open(RECORDS))}
    cks = sorted(c for (c0, c) in recs if c0 == ch)
    return "\n".join(recs[(ch, c)]["source"] for c in cks)


def score(source_text: str, target_text: str, ents):
    """Return per-entity verdicts + summary for one (source, target) pair. Reflow-invariant."""
    tgt = target_index(target_text)
    present_src = [e for e in ents if any(s in source_text for s in e["src_variants"])]
    results = []
    for e in present_src:
        elem = [l for l in e["lemmas"]]
        hit = [l for l in elem if l in tgt["lemset"]]
        if not elem:
            verdict = "skip"
        elif len(hit) == len(elem):
            verdict = "present"
        elif hit:
            verdict = "partial"    # some distinctive lemma present -> lean present (avoid false omission)
        else:
            verdict = "omitted"    # NO distinctive lemma anywhere -> candidate omission
        # inflection_gap (amendment 1): present but a matched surface is OUTSIDE pinned decl.forms
        infl_gap = False
        if verdict in ("present", "partial"):
            matched_surfaces = set()
            for l in hit:
                matched_surfaces |= tgt["lem2surf"].get(l, set())
            # a surface that shares lemma but isn't a pinned form (case-insensitive compare on the form set)
            pinned_lower = {f.lower() for f in e["decl_forms"]}
            if matched_surfaces and not any(s.lower() in pinned_lower or
                                            any(s.lower() in f.lower().split() for f in e["decl_forms"])
                                            for s in matched_surfaces):
                infl_gap = True
        results.append({"dst": e["dst"], "type": e["type"], "verdict": verdict,
                        "lemmas": elem, "hit": hit, "inflection_gap": infl_gap})
    n = len(results)
    omitted = [r for r in results if r["verdict"] == "omitted"]
    infl = [r for r in results if r["inflection_gap"]]
    return {
        "n_source_entities": n,
        "n_omitted": len(omitted),
        "omission_rate": round(len(omitted) / n, 4) if n else 0.0,
        "n_inflection_gap_flags": len(infl),   # span-check candidates, NOT omissions
        "omitted_dst": [r["dst"] for r in omitted],
        "inflection_gap_dst": [r["dst"] for r in infl],
        "results": results,
    }


def content_floor_backstop(source_text, target_text):
    """Amendment 3: crude length backstop. Fertility ~1.20*cjk+0.39*other ru-tokens; ru chars ~ tokens*3.
    Here use a chars-based coarse proxy: expected ru chars ~ 1.3 * source chars (guzhenren rerun ratio)."""
    exp = 1.3 * len(source_text)
    ratio = len(target_text) / exp if exp else 1.0
    return {"fired": ratio < CONTENT_FLOOR, "ru_char_ratio_vs_expected": round(ratio, 3)}


def strip_title(text):
    """sizecurve outputs begin with a translated chapter-title line; drop the first line for scoring."""
    lines = text.split("\n", 1)
    return lines[1] if len(lines) > 1 and len(lines[0]) < 80 else text


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--validate", type=int, help="dump per-entity verdicts for one chapter's whole output (hand-proofread)")
    ap.add_argument("--score", action="store_true", help="Q2b: omission per size (run AFTER validation)")
    a = ap.parse_args()
    ents = load_seed_entities()

    if a.validate is not None:
        ch = a.validate
        src = chapter_source(ch)
        tgt = strip_title((SIZE_DIR / f"{ch}-whole.txt").read_text(encoding="utf-8"))
        r = score(src, tgt, ents)
        print(f"=== VALIDATION ch{ch}-whole: {r['n_source_entities']} source entities, "
              f"{r['n_omitted']} omitted, {r['n_inflection_gap_flags']} inflection_gap ===")
        print(f"src bytes sha (chapter_source join) len={len(src)}")
        for res in r["results"]:
            tag = res["verdict"].upper() + (" [INFL_GAP]" if res["inflection_gap"] else "")
            print(f"  {tag:22} {res['dst']:22} lemmas={res['lemmas']} hit={res['hit']}")
        print(f"\nHAND-PROOFREAD: read {SIZE_DIR}/{ch}-whole.txt vs source; confirm each OMITTED is truly absent "
              f"(not pronominalized/synonym) and each PRESENT truly appears. Tune rule if false verdicts.")
        return

    if a.score:
        table, all_omit = {}, []
        for ch in CHAPTERS:
            src = chapter_source(ch)
            for sz in SIZES:
                f = SIZE_DIR / f"{ch}-{sz}.txt"
                tgt = strip_title(f.read_text(encoding="utf-8"))
                r = score(src, tgt, ents)
                fl = content_floor_backstop(src, tgt)
                key = f"{ch}-{sz}"
                table[key] = {**{k: r[k] for k in ("n_source_entities", "n_omitted", "omission_rate",
                                                    "n_inflection_gap_flags", "omitted_dst")},
                              "content_floor": fl}
                all_omit.append(r["omission_rate"])
        # amendment 2: variance-check on WITHIN-CHAPTER size variation (NOT cross-chapter, which mixes baselines)
        by_ch = {}
        for ch in CHAPTERS:
            rates = [table[f"{ch}-{sz}"]["omission_rate"] for sz in SIZES]
            by_ch[ch] = {"rates": rates, "within_size_span": round(max(rates) - min(rates), 4),
                         "monotonic_up": rates == sorted(rates) and rates[0] < rates[-1]}
        max_within = max(v["within_size_span"] for v in by_ch.values())
        any_monotonic = any(v["monotonic_up"] for v in by_ch.values())
        # term-drift disambiguation: all raw flags here are known term-drift (seed reseed post-dates outputs)
        drift_flags = sorted({d for v in table.values() for d in v["omitted_dst"]})
        # ceiling test = NO monotonic size->omission trend in any chapter (a non-monotonic blip is noise,
        # not a size effect). Magnitude alone (one term-drift flag) does not indicate size-dependent loss.
        ceiling = not any_monotonic
        verdict = ("INSTRUMENT CEILING (amendment 2): entity-omission is FLAT across sizes within each chapter "
                   f"(max within-chapter size-span {max_within:.3f}, no monotonic size->omission trend). Moreover ALL "
                   f"raw 'omission' flags {drift_flags} are TERM-DRIFT artifacts, not content loss: the sizecurve "
                   "outputs (2026-07-11) PREDATE the класс/разряд reseed (D38.5, 07-12) and render 甲等/丙等 as 'разряд' "
                   "(hand-verified). => true entity-omission ~0 at every size. This is NOT evidence big chunks are "
                   "omission-safe: named entities survive at all sizes by nature; silent big-chunk loss (if any) is at "
                   "CLAUSE/DETAIL level, invisible to entity-presence. The size<->omission question needs SPAN-JUDGES "
                   "(paid Q2a), matching the design ('Q2b cannot ratify edit=chapter, only Q2a')."
                   if ceiling else
                   f"entity-omission shows within-chapter size variation (max span {max_within:.3f}"
                   f"{', monotonic' if any_monotonic else ''}); inspect per-size table before concluding.")
        result = {"chapters": CHAPTERS, "sizes": SIZES, "table": table, "by_chapter": by_ch,
                  "max_within_chapter_size_span": round(max_within, 4), "instrument_ceiling": ceiling,
                  "raw_flags_are_term_drift": drift_flags,
                  "term_drift_evidence": "outputs 2026-07-11 predate класс reseed D38.5 2026-07-12; outputs use 'разряд' (grep-verified)",
                  "verdict": verdict,
                  "note": "reflow-invariant entity-omission (lemma-matched pymorphy3; amendment 1: inflection outside "
                          "decl.forms = inflection_gap FLAG not omission). Validated on hand-proofread ch13-whole "
                          "(0 false omissions; term-drift класс->старейшина/разряд correctly not omission when lemma "
                          "shares, flagged when term fully substituted). Numbers atom deferred to span-judges (ru spells "
                          "numerals). Q2b = re-attribution of exp14 sizecurve ONLY (pre-v3 prompt, 2 ch, greedy)."}
        OUT.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
        print("=== Q2b reflow-invariant entity-omission per size ===")
        print(f"{'cfg':<10}{'src_ent':>8}{'omit':>6}{'omit_rate':>10}{'infl_gap':>9}{'floor_fired':>12}")
        for k, v in table.items():
            print(f"{k:<10}{v['n_source_entities']:>8}{v['n_omitted']:>6}{v['omission_rate']:>10}"
                  f"{v['n_inflection_gap_flags']:>9}{str(v['content_floor']['fired']):>12}")
        print(f"\nwithin-chapter size-span: " + ", ".join(f"ch{ch}={v['within_size_span']}" for ch, v in by_ch.items()))
        print(f"raw flags (ALL term-drift artifacts): {drift_flags}")
        print(f"instrument_ceiling={ceiling}")
        print(f"VERDICT: {verdict}")
        print(f"-> {OUT}")
        return

    ap.print_help()


if __name__ == "__main__":
    main()