scripts/validate_sanitizer_data_bearing.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import csv
import json
import os
import sys
import time
from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
from typing import Any

from lxml import etree

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))
os.chdir(ROOT)

from modules.slot_layout_planner.chart_sanitizer import (  # noqa: E402
    DATA_BEARING_ATTRS,
    DATA_BEARING_ROLE_VALUES,
    sanitize_chart_svg,
)
from scripts.generate_template_samples import (  # noqa: E402
    configure_chart_type_filter,
    nested_chart_svg,
    run_sample,
    safe_slug,
)


GEOMETRY_ATTRS = (
    "x",
    "y",
    "x1",
    "y1",
    "x2",
    "y2",
    "cx",
    "cy",
    "r",
    "rx",
    "ry",
    "width",
    "height",
    "d",
)


def read_jsonl(path: Path) -> list[dict[str, Any]]:
    with open(path, "r", encoding="utf-8") as f:
        return [json.loads(line) for line in f if line.strip()]


def append_jsonl(path: Path, record: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "a", encoding="utf-8") as f:
        f.write(json.dumps(record, ensure_ascii=False) + "\n")


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        json.dump(payload, f, indent=2, ensure_ascii=False)


def local_name(elem: etree._Element) -> str:
    return etree.QName(elem).localname if isinstance(elem.tag, str) else ""


def normalized_role(elem: etree._Element) -> str:
    role = str(elem.get("data-role") or "").strip().lower().replace("_", " ").replace("-", " ")
    return " ".join(role.split())


def is_data_bearing(elem: etree._Element) -> bool:
    for attr in DATA_BEARING_ATTRS:
        value = elem.get(attr)
        if value is not None and str(value).strip():
            return True
    role = normalized_role(elem)
    return role in DATA_BEARING_ROLE_VALUES


def data_signature(elem: etree._Element) -> str:
    attrs: dict[str, str] = {
        "tag": local_name(elem),
        "class": str(elem.get("class") or ""),
        "id": str(elem.get("id") or ""),
    }
    for name, value in sorted(elem.attrib.items()):
        if name.startswith("data-") or name in GEOMETRY_ATTRS:
            attrs[name] = str(value)
    text = "".join(elem.itertext()).strip()
    if text:
        attrs["text"] = " ".join(text.split())
    return json.dumps(attrs, sort_keys=True, ensure_ascii=False)


def data_signatures(svg_path: Path) -> Counter[str]:
    parser = etree.XMLParser(remove_blank_text=False, recover=True, huge_tree=True)
    root = etree.parse(str(svg_path), parser).getroot()
    signatures: Counter[str] = Counter()
    for elem in root.xpath(".//*"):
        if is_data_bearing(elem):
            signatures[data_signature(elem)] += 1
    return signatures


def metrics(svg_path: Path) -> dict[str, int]:
    text = svg_path.read_text(encoding="utf-8", errors="ignore")
    return {
        "rect": text.count("<rect"),
        "path": text.count("<path"),
        "circle": text.count("<circle"),
        "line": text.count("<line"),
        "text": text.count("<text"),
        "image": text.count("<image"),
    }


def resolve_existing_chart_svg(record: dict[str, Any]) -> Path | None:
    chart_svg = record.get("chart_svg")
    if chart_svg:
        path = Path(chart_svg)
        if path.is_file():
            return path
    sample_dir = record.get("sample_dir")
    if sample_dir:
        nested = nested_chart_svg(Path(sample_dir))
        if nested and nested.is_file():
            return nested
    return None


def render_chart(record: dict[str, Any], sample_dir: Path, output_png: bool) -> dict[str, Any]:
    result = run_sample(
        template_key=record["template_key"],
        data_path=Path(record["data_source"]),
        sample_dir=sample_dir,
        output_png=output_png,
        chart_only=True,
    )
    return result


def validate_record(
    record: dict[str, Any],
    output_root: Path,
    render: bool,
    output_png: bool,
) -> dict[str, Any]:
    template_key = record["template_key"]
    sample_index = int(record.get("sample_index", 0))
    sample_dir = output_root / "charts" / safe_slug(template_key) / f"sample_{sample_index:02d}"
    render_result: dict[str, Any] = {}

    if render:
        render_result = render_chart(record, sample_dir, output_png)
        chart_svg = Path(render_result.get("chart_svg") or "") if render_result.get("chart_svg") else None
    else:
        chart_svg = resolve_existing_chart_svg(record)

    if chart_svg is None or not chart_svg.is_file():
        return {
            "template_key": template_key,
            "sample_index": sample_index,
            "data_source": record.get("data_source", ""),
            "success": False,
            "error": "chart_svg_missing",
        }

    sanitize_dir = output_root / "sanitized" / safe_slug(template_key) / f"sample_{sample_index:02d}"
    sanitized_svg = sanitize_dir / "sanitized_chart.svg"
    report_path = sanitize_dir / "sanitized_chart.report.json"

    before = data_signatures(chart_svg)
    before_metrics = metrics(chart_svg)
    report = sanitize_chart_svg(chart_svg, sanitized_svg, report_path=report_path)
    after = data_signatures(sanitized_svg)
    after_metrics = metrics(sanitized_svg)
    lost = before - after
    gained = after - before

    row = {
        "template_key": template_key,
        "sample_index": sample_index,
        "data_source": record.get("data_source", ""),
        "success": True,
        "render_success": render_result.get("success") if render else "",
        "chart_svg": str(chart_svg),
        "sanitized_svg": str(sanitized_svg),
        "sanitizer_report": str(report_path),
        "removed_count": report.removed_count,
        "removed_by_reason": report.removed_by_reason,
        "data_bearing_before": sum(before.values()),
        "data_bearing_after": sum(after.values()),
        "data_bearing_lost": sum(lost.values()),
        "data_bearing_gained": sum(gained.values()),
        "lost_examples": list(lost.keys())[:5],
        "rect_before": before_metrics["rect"],
        "rect_after": after_metrics["rect"],
        "path_before": before_metrics["path"],
        "path_after": after_metrics["path"],
        "circle_before": before_metrics["circle"],
        "circle_after": after_metrics["circle"],
        "text_before": before_metrics["text"],
        "text_after": after_metrics["text"],
        "image_before": before_metrics["image"],
        "image_after": after_metrics["image"],
        "error": "",
    }
    return row


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Batch validate that SVG sanitization does not remove data-bearing chart elements."
    )
    parser.add_argument("--records", type=Path, required=True, help="JSONL records with template_key/data_source/sample_index.")
    parser.add_argument("--output-root", type=Path, required=True)
    parser.add_argument("--render-chart", action="store_true", help="Render current chart SVGs before validation.")
    parser.add_argument("--output-png", action="store_true", help="Generate chart PNGs during chart-only render.")
    parser.add_argument("--workers", type=int, default=8)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--resume", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    configure_chart_type_filter(False)
    args.output_root.mkdir(parents=True, exist_ok=True)
    records = read_jsonl(args.records)
    if args.limit:
        records = records[: args.limit]

    manifest_path = args.output_root / "manifest.jsonl"
    done: set[tuple[str, int]] = set()
    if args.resume and manifest_path.exists():
        for row in read_jsonl(manifest_path):
            done.add((row.get("template_key", ""), int(row.get("sample_index", -1))))
    elif manifest_path.exists():
        manifest_path.unlink()

    pending = [
        record
        for record in records
        if (record.get("template_key", ""), int(record.get("sample_index", -1))) not in done
    ]
    write_json(
        args.output_root / "run_config.json",
        {
            "records": str(args.records),
            "output_root": str(args.output_root),
            "render_chart": args.render_chart,
            "output_png": args.output_png,
            "workers": args.workers,
            "limit": args.limit,
            "selected_records": len(records),
            "pending_records": len(pending),
        },
    )

    started = time.time()
    print(f"validating records={len(records)} pending={len(pending)} workers={args.workers}", flush=True)
    if pending and args.workers > 1:
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = [
                executor.submit(validate_record, record, args.output_root, args.render_chart, args.output_png)
                for record in pending
            ]
            for index, future in enumerate(as_completed(futures), 1):
                row = future.result()
                append_jsonl(manifest_path, row)
                elapsed = time.time() - started
                rate = index / elapsed if elapsed else 0
                remaining = (len(pending) - index) / rate if rate else 0
                print(
                    f"[{index}/{len(pending)}] {row.get('template_key')} "
                    f"sample_{int(row.get('sample_index', 0)):02d} "
                    f"success={row.get('success')} lost={row.get('data_bearing_lost', '')} "
                    f"eta={remaining/60:.1f}m",
                    flush=True,
                )
    else:
        for index, record in enumerate(pending, 1):
            row = validate_record(record, args.output_root, args.render_chart, args.output_png)
            append_jsonl(manifest_path, row)
            elapsed = time.time() - started
            rate = index / elapsed if elapsed else 0
            remaining = (len(pending) - index) / rate if rate else 0
            print(
                f"[{index}/{len(pending)}] {row.get('template_key')} "
                f"sample_{int(row.get('sample_index', 0)):02d} "
                f"success={row.get('success')} lost={row.get('data_bearing_lost', '')} "
                f"eta={remaining/60:.1f}m",
                flush=True,
            )

    rows = read_jsonl(manifest_path) if manifest_path.exists() else []
    successes = [row for row in rows if row.get("success")]
    failures = [row for row in rows if not row.get("success")]
    suspicious = [row for row in successes if int(row.get("data_bearing_lost") or 0) > 0]
    removed_reasons: Counter[str] = Counter()
    for row in successes:
        for reason, count in (row.get("removed_by_reason") or {}).items():
            removed_reasons[reason] += int(count)

    summary = {
        "total_records": len(rows),
        "success": len(successes),
        "failed": len(failures),
        "templates": len({row.get("template_key") for row in rows}),
        "suspicious_data_bearing_loss": len(suspicious),
        "total_data_bearing_lost": sum(int(row.get("data_bearing_lost") or 0) for row in successes),
        "total_removed_nodes": sum(int(row.get("removed_count") or 0) for row in successes),
        "removed_by_reason": dict(removed_reasons.most_common()),
        "failure_examples": failures[:20],
        "suspicious_examples": suspicious[:20],
        "manifest": str(manifest_path),
        "csv": str(args.output_root / "summary.csv"),
    }
    write_json(args.output_root / "summary.json", summary)

    fieldnames = [
        "template_key",
        "sample_index",
        "success",
        "render_success",
        "removed_count",
        "removed_by_reason",
        "data_bearing_before",
        "data_bearing_after",
        "data_bearing_lost",
        "data_bearing_gained",
        "rect_before",
        "rect_after",
        "path_before",
        "path_after",
        "circle_before",
        "circle_after",
        "text_before",
        "text_after",
        "image_before",
        "image_after",
        "chart_svg",
        "sanitized_svg",
        "sanitizer_report",
        "data_source",
        "error",
    ]
    with open(args.output_root / "summary.csv", "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames, extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)

    print(json.dumps(summary, indent=2, ensure_ascii=False), flush=True)
    return 0 if not failures and not suspicious else 1


if __name__ == "__main__":
    raise SystemExit(main())