"""Markovo lab 评分器：runs.json + truth.json + 输出 md → scores.json / scores.md。

用法：
  tools/lab-venv/bin/python tools/lab_score.py --lab <RUN>/lab \
      --out <RUN>/lab/scores.json [--markdown <RUN>/lab/scores.md]

归一化：lower + 全部空白折成单空格 + 去掉 markdown 管道/星号/反引号/井号。
不适用于该 truth.kind 的指标写 null。
"""

import argparse
import json
import re
import sys
from pathlib import Path

METRICS_DOC = {
    "prose_recall": "truth 段落按 '. ' 切句，归一化后 ≥4 词的句子在输出中出现的比例",
    "reading_order": "每段前 6 词作锚点在输出中的位置，相邻段对位置严格递增的比例；锚点缺失计 0",
    "table_rows_intact": "truth 每行，存在某条输出行包含该行所有单元格字符串的比例",
    "table_header_intact": "leaf_header（无则 header）全部单元格出现在同一输出行：1 或 0；只测叶子单元格同行，不测跨列分组语义；doc 类样本有多张表/工作表时为满足该条件的表所占比例",
    "table_pipe_detected": "输出是否含 markdown 表格（| 开头结尾行且下一行为分隔行）：true/false",
    "pages_covered": "mixed 的 page_anchors 各自出现在输出中的比例；第三个锚点在扫描页，无 OCR 时上限 2/3",
    "scan_text_recall": "scan 文本按 prose_recall 口径的召回；无 OCR 预期为 0",
    "equation_recall": "expressions 经数学标记归一化（去 $ / \\[ \\] 定界、\\mathsf{} 等字体包裹、\\{ → {、^{\\wedge} → ^、\\_ → _、去全部空白）后子串命中的比例；对所有工具同一口径，衡量数学内容是否可恢复而非输出格式",
    "equation_recall_exact": "expressions 仅做大小写/空白归一化的纯文本精确子串命中比例（旧口径，保留作对照：输出 LaTeX 标记的路径在此列必然偏低）",
    "heading_recall": "doc 类：truth 标题文字出现在某条以 # 开头的输出行里的比例（不论级别）",
    "heading_level_exact": "doc 类：truth 标题以恰好 level 个 # 输出的比例；truth level 为空的标题（幻灯片标题、EPUB 章名）不计入",
    "list_items_intact": "doc 类：truth 列表项（含幻灯片要点）出现在以 - * + 或 数字. 数字) 开头的输出行里的比例",
    "code_block_intact": "doc 类：truth 代码块的每一行都出现在 ``` / ~~~ 围栏内的比例（按块计）",
    "boilerplate_excluded": "doc 类：导航、标语、Cookie 横幅、页脚等样板字符串没有出现在输出中的比例；越高说明正文抽取越干净",
    "slide_order": "doc 类：幻灯片标题按 reading_order 口径（前 6 词锚点）相邻对位置严格递增的比例",
    "notes_recall": "doc 类：幻灯片演讲者备注全文出现在输出中的比例",
    "sheet_coverage": "doc 类：至少有一行数据按 table_rows_intact 口径完整出现的工作表所占比例",
    "elapsed_s": "lab_run 实测耗时（秒），透传",
    "error": "lab_run 记录的异常文本，透传；无异常为 null",
}

_STRIP_RE = re.compile(r"[|*`#]")
_WS_RE = re.compile(r"\s+")
_PIPE_SEP_RE = re.compile(r"^\|[\s:\-|]+\|$")


def norm(s: str) -> str:
    s = _STRIP_RE.sub(" ", s.lower())
    return _WS_RE.sub(" ", s).strip()


_MATH_WRAP_RE = re.compile(r"\\(?:mathsf|mathrm|mathit|mathbf|text|textrm|operatorname)\{([^{}]*)\}")


def math_norm(s: str) -> str:
    """数学标记归一化：剥掉定界与字体包裹，还原转义，去全部空白。对 truth 与输出同样施加。"""
    s = s.lower()
    s = re.sub(r"\$+|\\\[|\\\]|\\\(|\\\)", " ", s)
    for _ in range(3):
        s = _MATH_WRAP_RE.sub(r"\1", s)
    s = re.sub(r"\\textsubscript\{([^{}]*)\}", r"_\1", s)
    s = re.sub(r"\\textsuperscript\{([^{}]*)\}", r"^\1", s)
    s = s.replace("{\\wedge}", "^").replace("\\wedge", "^")
    s = s.replace("\\{", "{").replace("\\}", "}").replace("\\_", "_")
    s = re.sub(r"[|*`#]", " ", s)
    return re.sub(r"\s+", "", s)


def sentences(text: str):
    parts = [p.strip() for p in re.split(r"\.\s+", text) if p.strip()]
    return [p for p in parts if len(norm(p).split()) >= 4]


def has_pipe_table(text: str) -> bool:
    lines = text.splitlines()
    for i in range(len(lines) - 1):
        a, b = lines[i].strip(), lines[i + 1].strip()
        if a.startswith("|") and a.endswith("|") and _PIPE_SEP_RE.match(b):
            return True
    return False


def prose_recall(truth_paragraphs, out_norm: str) -> float:
    sents = []
    for p in truth_paragraphs:
        sents.extend(sentences(p))
    if not sents:
        return None
    hit = sum(1 for s in sents if norm(s) in out_norm)
    return round(hit / len(sents), 4)


def reading_order(truth_paragraphs, out_norm: str) -> float:
    anchors = [norm(" ".join(norm(p).split()[:6])) for p in truth_paragraphs]
    pos = [out_norm.find(a) for a in anchors]
    pairs = len(anchors) - 1
    if pairs <= 0:
        return None
    good = sum(1 for i in range(pairs)
               if pos[i] >= 0 and pos[i + 1] > pos[i])
    return round(good / pairs, 4)


def table_rows_intact(rows, out_lines_norm) -> float:
    if not rows:
        return None
    hit = 0
    for row in rows:
        cells = [norm(str(c)) for c in row if str(c).strip()]
        if not cells:
            continue
        if any(all(c in ln for c in cells) for ln in out_lines_norm):
            hit += 1
    return round(hit / len(rows), 4)


def table_header_intact(truth, out_lines_norm):
    header = truth.get("leaf_header") or truth.get("header")
    if not header:
        return None
    cells = [norm(str(c)) for c in header if str(c).strip()]
    return 1.0 if any(all(c in ln for c in cells) for ln in out_lines_norm) else 0.0


_HEADING_RE = re.compile(r"^(#{1,6})\s+(.*)$")
_LIST_RE = re.compile(r"^\s*(?:[-*+]|\d+[.)])\s+(.*)$")
_FENCE_RE = re.compile(r"^\s*(```|~~~)")


def heading_metrics(headings, out_text: str):
    lines = [ln.strip() for ln in out_text.splitlines()]
    heading_lines = [m for m in (_HEADING_RE.match(ln) for ln in lines) if m]
    recall = sum(1 for h in headings
                 if any(norm(h["text"]) in norm(m.group(2)) for m in heading_lines))
    leveled = [h for h in headings if h.get("level")]
    exact = sum(1 for h in leveled
                if any(len(m.group(1)) == h["level"] and norm(h["text"]) in norm(m.group(2))
                       for m in heading_lines))
    return (round(recall / len(headings), 4),
            round(exact / len(leveled), 4) if leveled else None)


def list_items_intact(items, out_text: str):
    list_lines = [norm(m.group(1)) for m in (_LIST_RE.match(ln) for ln in out_text.splitlines()) if m]
    hit = sum(1 for it in items if any(norm(it) in ln for ln in list_lines))
    return round(hit / len(items), 4)


def code_block_intact(blocks, out_text: str):
    fenced, inside = [], False
    for ln in out_text.splitlines():
        if _FENCE_RE.match(ln):
            inside = not inside
            continue
        if inside:
            fenced.append(norm(ln))
    hit = sum(1 for block in blocks
              if all(any(norm(line) in f for f in fenced) for line in block if line.strip()))
    return round(hit / len(blocks), 4)


def score_doc(truth, out_text: str, out_norm: str, out_lines_norm, m: dict):
    if truth.get("headings"):
        m["heading_recall"], m["heading_level_exact"] = heading_metrics(truth["headings"], out_text)
    paragraphs = truth.get("paragraphs") or []
    if paragraphs:
        m["prose_recall"] = prose_recall(paragraphs, out_norm)
        m["reading_order"] = reading_order(paragraphs, out_norm)
    items = [it for lst in truth.get("lists") or [] for it in lst["items"]]
    if items:
        m["list_items_intact"] = list_items_intact(items, out_text)
    tables = [*(truth.get("tables") or []), *(truth.get("sheets") or [])]
    if tables:
        m["table_rows_intact"] = table_rows_intact(
            [r for t in tables for r in t["rows"]], out_lines_norm)
        m["table_header_intact"] = round(
            sum(table_header_intact(t, out_lines_norm) or 0 for t in tables) / len(tables), 4)
        m["table_pipe_detected"] = has_pipe_table(out_text)
    if truth.get("code_blocks"):
        m["code_block_intact"] = code_block_intact(truth["code_blocks"], out_text)
    if truth.get("boilerplate"):
        bp = truth["boilerplate"]
        m["boilerplate_excluded"] = round(sum(1 for s in bp if norm(s) not in out_norm) / len(bp), 4)
    slides = truth.get("slides") or []
    if slides:
        m["slide_order"] = reading_order([s["title"] for s in slides], out_norm)
        notes = [s["notes"] for s in slides if s.get("notes")]
        if notes:
            m["notes_recall"] = round(sum(1 for n in notes if norm(n) in out_norm) / len(notes), 4)
    sheets = truth.get("sheets") or []
    if sheets:
        m["sheet_coverage"] = round(
            sum(1 for s in sheets if (table_rows_intact(s["rows"], out_lines_norm) or 0) > 0)
            / len(sheets), 4)


def score_cell(truth, out_text: str) -> dict:
    out_norm = norm(out_text)
    out_lines_norm = [norm(ln) for ln in out_text.splitlines() if ln.strip()]
    kind = truth["kind"]
    m = {k: None for k in METRICS_DOC}

    prose_paras, scan_texts, tables, order_paras = [], [], [], None
    if kind == "doc":
        score_doc(truth, out_text, out_norm, out_lines_norm, m)
        return m
    if kind == "prose":
        prose_paras = truth["paragraphs"]
        order_paras = truth["paragraphs"]
    elif kind == "columns":
        prose_paras = truth["paragraphs"]
        order_paras = truth["paragraphs"]
    elif kind == "scan":
        scan_texts = [truth["text"]]
    elif kind == "table":
        tables = [truth]
    elif kind == "equation":
        out_math = math_norm(out_text)
        m["equation_recall"] = round(
            sum(1 for e in truth["expressions"] if math_norm(e) in out_math)
            / len(truth["expressions"]), 4)
        m["equation_recall_exact"] = round(
            sum(1 for e in truth["expressions"] if norm(e) in out_norm)
            / len(truth["expressions"]), 4)
    elif kind == "mixed":
        for pg in truth["pages"]:
            if pg["kind"] == "prose":
                prose_paras += pg["paragraphs"]
            elif pg["kind"] == "table":
                tables.append(pg)
            elif pg["kind"] == "scan":
                scan_texts.append(pg["text"])
        anchors = truth.get("page_anchors") or []
        if anchors:
            m["pages_covered"] = round(
                sum(1 for a in anchors if norm(a) in out_norm) / len(anchors), 4)

    if prose_paras:
        m["prose_recall"] = prose_recall(prose_paras, out_norm)
    if order_paras:
        m["reading_order"] = reading_order(order_paras, out_norm)
    if scan_texts:
        m["scan_text_recall"] = prose_recall(scan_texts, out_norm)
    if tables:
        all_rows = [r for t in tables for r in t["rows"]]
        m["table_rows_intact"] = table_rows_intact(all_rows, out_lines_norm)
        m["table_header_intact"] = table_header_intact(tables[0], out_lines_norm)
        m["table_pipe_detected"] = has_pipe_table(out_text)
    return m


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--lab", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--markdown", default=None)
    args = ap.parse_args(argv)

    lab = Path(args.lab)
    runs = json.loads((lab / "runs.json").read_text(encoding="utf-8"))
    corpus = lab / "corpus"

    cells = []
    for r in runs["results"]:
        truth = json.loads(
            (corpus / f"{r['sample']}.truth.json").read_text(encoding="utf-8"))
        out_path = lab / r["output"] if r.get("output") else None
        out_text = out_path.read_text(encoding="utf-8") if out_path and out_path.is_file() else ""
        m = score_cell(truth, out_text)
        m["elapsed_s"] = r.get("elapsed_s")
        m["error"] = r.get("error")
        cells.append({"tool": r["tool"], "sample": r["sample"],
                      "kind": truth["kind"], **m})

    scores = {
        "metrics_doc": METRICS_DOC,
        "tools": runs.get("tools"),
        "python": runs.get("python"),
        "date": runs.get("date"),
        "cells": cells,
    }
    Path(args.out).write_text(
        json.dumps(scores, ensure_ascii=False, indent=1) + "\n", encoding="utf-8")

    if args.markdown:
        keys = [k for k in METRICS_DOC if any(c[k] is not None for c in cells)]
        lines = ["| tool | sample | " + " | ".join(keys) + " |",
                 "| --- | --- |" + " --- |" * len(keys)]
        for c in cells:
            vals = []
            for k in keys:
                v = c[k]
                vals.append("" if v is None else
                            ("yes" if v is True else "no" if v is False else str(v)))
            lines.append(f"| {c['tool']} | {c['sample']} | " + " | ".join(vals) + " |")
        Path(args.markdown).write_text("\n".join(lines) + "\n", encoding="utf-8")
    return 0


if __name__ == "__main__":
    sys.exit(main())