scripts/run_d3_regression_from_baseline.py

#!/usr/bin/env python3
import argparse
import csv
import json
import logging
import math
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path

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

from scripts.generate_template_samples import (  # noqa: E402
    configure_chart_type_filter,
    newest_root_png,
    newest_root_svg,
    run_sample_job,
    safe_slug,
)
from modules.chart_engine.template.template_registry import scan_templates  # noqa: E402


DEFAULT_BASELINE = (
    "output/gpt_image_2_polish_bundle_20260528_0242/"
    "chart_template_samples_20260527_083945"
)

SHAPE_TAGS = ("path", "rect", "circle", "line", "polygon", "polyline", "ellipse", "use")
IMAGE_TAGS = ("image",)
FALLBACK_MARKER = "This is a fallback SVG using a PNG screenshot"


def read_jsonl(path: Path):
    if not path.is_file():
        return []
    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):
    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):
    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 parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--baseline-root", default=DEFAULT_BASELINE)
    parser.add_argument("--output-root", default=None)
    parser.add_argument("--samples-per-template", type=int, default=5)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--png-longest-side", type=int, default=1600)
    parser.add_argument(
        "--no-output-png",
        action="store_true",
        help="Skip PNG generation. PNG output is enabled by default.",
    )
    parser.add_argument("--resume", action="store_true")
    parser.add_argument("--rerun-failures", action="store_true")
    parser.add_argument("--compare-only", action="store_true")
    parser.add_argument("--chart-only", action="store_true")
    parser.add_argument(
        "--include-non-d3",
        action="store_true",
        help="Include matching non-d3 template records from the baseline manifest.",
    )
    parser.add_argument(
        "--remap-to-current-template-key",
        action="store_true",
        help="Map baseline records to the current registry template key with the same chart name.",
    )
    parser.add_argument(
        "--slot-polish-after-chart",
        action="store_true",
        help="Render each template chart-only, then run full_image_polisher in slot mode.",
    )
    parser.add_argument(
        "--slot-polish-dry-run",
        action="store_true",
        help="Generate slot masks/prompts/manifests without calling the image model.",
    )
    parser.add_argument(
        "--planned-slot-polish",
        action="store_true",
        help=(
            "Render chart-only, sanitize template title/image artifacts, plan editable slots, "
            "then run full_image_polisher in slot mode."
        ),
    )
    parser.add_argument(
        "--planned-slot-dry-run",
        action="store_true",
        help="Generate planned slot canvas/masks/prompts/manifests without calling the image model.",
    )
    parser.add_argument(
        "--planned-slot-disallow-chart-overlap",
        action="store_true",
        help="Do not allow planned slots to overlap the chart bbox. Overlap is allowed by default.",
    )
    parser.add_argument(
        "--planned-slot-polisher-base-url",
        default=None,
        help=(
            "Override base_url for planned-slot full_image_polisher. "
            "Use 'openai_default' to ignore config.base_url and call the official OpenAI endpoint."
        ),
    )
    parser.add_argument("--slot-polisher-backend", choices=("auto", "openai", "pinco"), default="auto")
    parser.add_argument("--planned-slot-polisher-backend", choices=("auto", "openai", "pinco"), default="auto")
    parser.add_argument(
        "--pinco-command",
        default=None,
        help=(
            "Command template for local Pinco inference. Placeholders: {input}, {mask}, "
            "{foreground}, {prompt_file}, {output}, {model}, {width}, {height}."
        ),
    )
    parser.add_argument("--pinco-url", default=None, help="HTTP endpoint for a Pinco inpainting service.")
    parser.add_argument("--pinco-timeout", type=int, default=600)
    parser.add_argument("--limit", type=int, default=0)
    parser.add_argument("--template", action="append", default=None)
    parser.add_argument("--log-file", default=None)
    return parser.parse_args()


def current_template_key_by_chart_name() -> dict[str, str]:
    priority = {"d3-js": 0, "echarts-js": 1, "echarts_py": 2}
    candidates: dict[str, tuple[int, str]] = {}
    templates = scan_templates(force=True)
    for engine, chart_types in templates.items():
        for chart_type, chart_names in chart_types.items():
            for chart_name in chart_names:
                key = f"{engine}/{chart_type}/{chart_name}"
                rank = priority.get(engine, 99)
                current = candidates.get(chart_name)
                if current is None or rank < current[0]:
                    candidates[chart_name] = (rank, key)
    return {chart_name: key for chart_name, (_rank, key) in candidates.items()}


def select_baseline_records(
    baseline_root: Path,
    samples_per_template: int,
    templates: set[str] | None,
    include_non_d3: bool = False,
    remap_to_current_template_key: bool = False,
):
    records = read_jsonl(baseline_root / "manifest.jsonl")
    current_key_by_name = (
        current_template_key_by_chart_name()
        if remap_to_current_template_key
        else {}
    )
    grouped = {}
    for record in records:
        baseline_template_key = record.get("template_key") or ""
        chart_name = baseline_template_key.split("/")[-1]
        template_key = current_key_by_name.get(chart_name, baseline_template_key)
        if not include_non_d3 and not template_key.startswith("d3-js/"):
            continue
        if (
            templates
            and chart_name not in templates
            and template_key not in templates
            and baseline_template_key not in templates
        ):
            continue
        if template_key != baseline_template_key:
            record = {
                **record,
                "baseline_template_key": baseline_template_key,
                "template_key": template_key,
            }
        grouped.setdefault(template_key, []).append(record)

    selected = []
    for template_key in sorted(grouped):
        template_records = sorted(
            grouped[template_key],
            key=lambda item: (item.get("sample_index", 10**9), item.get("data_source", "")),
        )
        selected.extend(template_records[:samples_per_template])
    return selected


def sample_key(record: dict):
    return (record.get("template_key"), int(record.get("sample_index", -1)))


def latest_manifest_records(manifest_path: Path):
    latest = {}
    for record in read_jsonl(manifest_path):
        latest[sample_key(record)] = record
    return latest


def completed_keys(manifest_path: Path, successful_only: bool = False):
    latest = latest_manifest_records(manifest_path)
    if successful_only:
        return {key for key, record in latest.items() if record.get("success")}
    return set(latest)


def resolve_sample_dir(root: Path, record: dict, use_template_key: bool = False):
    template_key = record.get("template_key") or ""
    slug_source = template_key if use_template_key else template_key.split("/")[-1]
    return root / safe_slug(slug_source) / f"sample_{int(record.get('sample_index', 0)):02d}"


def nested_chart_svg(sample_dir: Path):
    candidates = sorted(
        [p for p in sample_dir.glob("*/chart.svg") if p.is_file()],
        key=lambda p: p.stat().st_mtime,
    )
    return candidates[-1] if candidates else None


def file_metrics(path: Path | None):
    if path is not None:
        path = Path(path)
    if path is None or not path.is_file():
        return {
            "exists": False,
            "bytes": 0,
            "fallback": False,
            "n_shapes": 0,
            "n_text": 0,
            "n_images": 0,
            "n_visible": 0,
            "n_elements": 0,
            "empty": True,
        }
    text = path.read_text(encoding="utf-8", errors="ignore")
    n_shapes = sum(text.count(f"<{tag}") for tag in SHAPE_TAGS)
    n_text = text.count("<text")
    n_images = sum(text.count(f"<{tag}") for tag in IMAGE_TAGS)
    n_visible = n_shapes + n_text + n_images
    return {
        "exists": True,
        "bytes": path.stat().st_size,
        "fallback": FALLBACK_MARKER in text,
        "n_shapes": n_shapes,
        "n_text": n_text,
        "n_images": n_images,
        "n_visible": n_visible,
        "n_elements": text.count("<"),
        "empty": n_visible == 0,
    }


def image_metrics_and_diff(before: Path | None, after: Path | None, max_side: int = 256):
    if before is not None:
        before = Path(before)
    if after is not None:
        after = Path(after)
    result = {
        "before_png_exists": bool(before and before.is_file()),
        "after_png_exists": bool(after and after.is_file()),
        "before_size": "",
        "after_size": "",
        "png_mae": "",
        "png_rmse": "",
    }
    if not result["before_png_exists"] or not result["after_png_exists"]:
        return result
    try:
        from PIL import Image
    except Exception as exc:
        result["png_error"] = f"PIL import failed: {exc}"
        return result

    try:
        with Image.open(before) as img_a, Image.open(after) as img_b:
            img_a = img_a.convert("RGB")
            img_b = img_b.convert("RGB")
            result["before_size"] = f"{img_a.width}x{img_a.height}"
            result["after_size"] = f"{img_b.width}x{img_b.height}"
            target_w = max(img_a.width, img_b.width)
            target_h = max(img_a.height, img_b.height)
            scale = min(1.0, max_side / max(target_w, target_h))
            size = (max(1, int(target_w * scale)), max(1, int(target_h * scale)))
            img_a = img_a.resize(size)
            img_b = img_b.resize(size)
            pixels_a = img_a.tobytes()
            pixels_b = img_b.tobytes()
            total = len(pixels_a)
            abs_sum = 0
            sq_sum = 0
            for a, b in zip(pixels_a, pixels_b):
                delta = abs(a - b)
                abs_sum += delta
                sq_sum += delta * delta
            result["png_mae"] = round(abs_sum / total / 255, 6)
            result["png_rmse"] = round(math.sqrt(sq_sum / total) / 255, 6)
    except Exception as exc:
        result["png_error"] = str(exc)
    return result


def compare_records(baseline_root: Path, output_root: Path, selected_records: list[dict]):
    current_by_key = {
        sample_key(record): record
        for record in read_jsonl(output_root / "manifest.jsonl")
    }
    rows = []
    for before_record in selected_records:
        key = sample_key(before_record)
        after_record = current_by_key.get(key, {})
        before_dir = resolve_sample_dir(baseline_root, before_record)
        after_dir = resolve_sample_dir(output_root, before_record, use_template_key=True)
        before_svg = newest_root_svg(before_dir)
        after_svg = newest_root_svg(after_dir)
        before_png = newest_root_png(before_dir)
        after_png = newest_root_png(after_dir)
        before_chart_svg = nested_chart_svg(before_dir)
        after_chart_svg = nested_chart_svg(after_dir)
        before_chart = file_metrics(before_chart_svg)
        after_chart = file_metrics(after_chart_svg)
        before_final = file_metrics(before_svg)
        after_final = file_metrics(after_svg)
        image_diff = image_metrics_and_diff(before_png, after_png)
        row = {
            "template_key": before_record.get("template_key", ""),
            "chart_name": (before_record.get("template_key") or "").split("/")[-1],
            "sample_index": before_record.get("sample_index"),
            "data_source": before_record.get("data_source", ""),
            "baseline_success": bool(before_record.get("success")),
            "current_success": bool(after_record.get("success")),
            "baseline_final_svg": str(before_svg or ""),
            "current_final_svg": str(after_svg or ""),
            "baseline_final_png": str(before_png or ""),
            "current_final_png": str(after_png or ""),
            "baseline_chart_svg": str(before_chart_svg or ""),
            "current_chart_svg": str(after_chart_svg or ""),
            "baseline_chart_fallback": before_chart["fallback"],
            "current_chart_fallback": after_chart["fallback"],
            "baseline_chart_shapes": before_chart["n_shapes"],
            "current_chart_shapes": after_chart["n_shapes"],
            "baseline_chart_text": before_chart["n_text"],
            "current_chart_text": after_chart["n_text"],
            "baseline_chart_images": before_chart["n_images"],
            "current_chart_images": after_chart["n_images"],
            "baseline_chart_visible": before_chart["n_visible"],
            "current_chart_visible": after_chart["n_visible"],
            "baseline_chart_empty": before_chart["empty"],
            "current_chart_empty": after_chart["empty"],
            "baseline_final_bytes": before_final["bytes"],
            "current_final_bytes": after_final["bytes"],
            "baseline_final_empty": before_final["empty"],
            "current_final_empty": after_final["empty"],
            **image_diff,
        }
        row["status_change"] = (
            "same_success" if row["baseline_success"] and row["current_success"]
            else "same_failure" if not row["baseline_success"] and not row["current_success"]
            else "regressed_failure" if row["baseline_success"] and not row["current_success"]
            else "fixed_from_failure"
        )
        row["fallback_change"] = (
            "same"
            if row["baseline_chart_fallback"] == row["current_chart_fallback"]
            else "regressed_fallback"
            if row["current_chart_fallback"]
            else "fixed_fallback"
        )
        row["empty_chart_change"] = (
            "same_empty"
            if row["baseline_chart_empty"] and row["current_chart_empty"]
            else "regressed_empty"
            if row["current_chart_empty"]
            else "fixed_empty"
            if row["baseline_chart_empty"]
            else "same_visible"
        )
        rows.append(row)
    return rows


def write_comparison(output_root: Path, rows: list[dict]):
    jsonl_path = output_root / "comparison.jsonl"
    if jsonl_path.exists():
        jsonl_path.unlink()
    for row in rows:
        append_jsonl(jsonl_path, row)

    csv_path = output_root / "comparison.csv"
    fieldnames = list(rows[0].keys()) if rows else ["template_key"]
    with open(csv_path, "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)

    by_status = {}
    by_fallback = {}
    by_empty_chart = {}
    high_diff = {"mae_gt_0_01": 0, "mae_gt_0_05": 0, "mae_gt_0_15": 0}
    for row in rows:
        by_status[row["status_change"]] = by_status.get(row["status_change"], 0) + 1
        by_fallback[row["fallback_change"]] = by_fallback.get(row["fallback_change"], 0) + 1
        by_empty_chart[row["empty_chart_change"]] = by_empty_chart.get(row["empty_chart_change"], 0) + 1
        mae = row.get("png_mae")
        if isinstance(mae, (int, float)):
            if mae > 0.01:
                high_diff["mae_gt_0_01"] += 1
            if mae > 0.05:
                high_diff["mae_gt_0_05"] += 1
            if mae > 0.15:
                high_diff["mae_gt_0_15"] += 1

    top_diffs = sorted(
        [row for row in rows if isinstance(row.get("png_mae"), (int, float))],
        key=lambda row: row["png_mae"],
        reverse=True,
    )[:50]
    problem_rows = [
        row for row in rows
        if row["status_change"] == "regressed_failure"
        or row["fallback_change"] == "regressed_fallback"
        or row["empty_chart_change"] in {"regressed_empty", "same_empty"}
        or not row["current_success"]
    ][:100]
    summary = {
        "total_samples": len(rows),
        "total_templates": len({row["template_key"] for row in rows}),
        "status_counts": by_status,
        "fallback_counts": by_fallback,
        "empty_chart_counts": by_empty_chart,
        "png_diff_counts": high_diff,
        "top_png_diffs": top_diffs,
        "problem_rows": problem_rows,
        "comparison_jsonl": str(jsonl_path),
        "comparison_csv": str(csv_path),
    }
    write_json(output_root / "comparison_summary.json", summary)
    return summary


def redirect_output(log_path: Path):
    log_path.parent.mkdir(parents=True, exist_ok=True)
    log_file = open(log_path, "a", encoding="utf-8", buffering=1)
    sys.stdout = log_file
    sys.stderr = log_file
    for handler in logging.getLogger().handlers:
        if hasattr(handler, "stream"):
            handler.stream = log_file
    return log_file


def main():
    args = parse_args()
    configure_chart_type_filter(False)
    os.environ["RENDER_LONGEST_SIDE"] = str(args.png_longest_side)
    output_png = not args.no_output_png

    baseline_root = Path(args.baseline_root)
    if not baseline_root.is_dir():
        raise SystemExit(f"baseline root is not a directory: {baseline_root}")

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    output_root = Path(args.output_root or f"output/d3_regression_from_baseline_{timestamp}")
    output_root.mkdir(parents=True, exist_ok=True)
    log_handle = None
    if args.log_file:
        log_handle = redirect_output(Path(args.log_file))

    templates = set(args.template) if args.template else None
    selected = select_baseline_records(
        baseline_root,
        args.samples_per_template,
        templates,
        include_non_d3=args.include_non_d3,
        remap_to_current_template_key=args.remap_to_current_template_key,
    )
    if args.limit:
        selected = selected[:args.limit]
    write_json(output_root / "run_config.json", {
        "baseline_root": str(baseline_root),
        "output_root": str(output_root),
        "samples_per_template": args.samples_per_template,
        "workers": args.workers,
        "png_longest_side": args.png_longest_side,
        "output_png": output_png,
        "rerun_failures": args.rerun_failures,
        "chart_only": args.chart_only,
        "include_non_d3": args.include_non_d3,
        "remap_to_current_template_key": args.remap_to_current_template_key,
        "slot_polish_after_chart": args.slot_polish_after_chart,
        "slot_polish_dry_run": args.slot_polish_dry_run,
        "planned_slot_polish": args.planned_slot_polish,
        "planned_slot_dry_run": args.planned_slot_dry_run,
        "planned_slot_allow_chart_overlap": not args.planned_slot_disallow_chart_overlap,
        "planned_slot_polisher_base_url": args.planned_slot_polisher_base_url,
        "slot_polisher_backend": args.slot_polisher_backend,
        "planned_slot_polisher_backend": args.planned_slot_polisher_backend,
        "pinco_command": args.pinco_command or "",
        "pinco_url": args.pinco_url or "",
        "pinco_timeout": args.pinco_timeout,
        "limit": args.limit,
        "template": args.template,
        "selected_samples": len(selected),
        "selected_templates": len({record.get("template_key") for record in selected}),
    })
    selected_path = output_root / "selected_baseline_records.jsonl"
    if not selected_path.exists():
        for record in selected:
            append_jsonl(selected_path, record)

    manifest_path = output_root / "manifest.jsonl"
    done = (
        completed_keys(manifest_path, successful_only=args.rerun_failures)
        if args.resume or args.compare_only or args.rerun_failures
        else set()
    )
    jobs = []
    if not args.compare_only:
        for record in selected:
            key = sample_key(record)
            if key in done:
                continue
            chart_name = (record.get("template_key") or "").split("/")[-1]
            sample_dir = resolve_sample_dir(output_root, record, use_template_key=True)
            jobs.append((
                record.get("template_key"),
                int(record.get("sample_index", 0)),
                Path(record.get("data_source")),
                sample_dir,
                output_png,
                args.chart_only,
                args.slot_polish_after_chart,
                args.slot_polish_dry_run,
                args.planned_slot_polish,
                args.planned_slot_dry_run,
                not args.planned_slot_disallow_chart_overlap,
                args.planned_slot_polisher_base_url,
                args.slot_polisher_backend,
                args.planned_slot_polisher_backend,
                args.pinco_command,
                args.pinco_url,
                args.pinco_timeout,
            ))

    started = time.time()
    print(
        f"selected_templates={len({r.get('template_key') for r in selected})} "
        f"selected_samples={len(selected)} to_run={len(jobs)} workers={args.workers}",
        flush=True,
    )
    if jobs and args.workers > 1:
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = [executor.submit(run_sample_job, job) for job in jobs]
            for index, future in enumerate(as_completed(futures), 1):
                record = future.result()
                append_jsonl(manifest_path, record)
                elapsed = time.time() - started
                rate = index / elapsed if elapsed else 0
                remaining = (len(jobs) - index) / rate if rate else 0
                print(
                    f"[{index}/{len(jobs)}] {record['template_key']} "
                    f"sample_{record['sample_index']:02d} success={record['success']} "
                    f"elapsed={elapsed/60:.1f}m eta={remaining/60:.1f}m",
                    flush=True,
                )
    else:
        for index, job in enumerate(jobs, 1):
            record = run_sample_job(job)
            append_jsonl(manifest_path, record)
            elapsed = time.time() - started
            rate = index / elapsed if elapsed else 0
            remaining = (len(jobs) - index) / rate if rate else 0
            print(
                f"[{index}/{len(jobs)}] {record['template_key']} "
                f"sample_{record['sample_index']:02d} success={record['success']} "
                f"elapsed={elapsed/60:.1f}m eta={remaining/60:.1f}m",
                flush=True,
            )

    rows = compare_records(baseline_root, output_root, selected)
    summary = write_comparison(output_root, rows)
    print(json.dumps({
        "output_root": str(output_root),
        "total_samples": summary["total_samples"],
        "total_templates": summary["total_templates"],
        "status_counts": summary["status_counts"],
        "fallback_counts": summary["fallback_counts"],
        "empty_chart_counts": summary["empty_chart_counts"],
        "png_diff_counts": summary["png_diff_counts"],
    }, indent=2, ensure_ascii=False), flush=True)
    if log_handle:
        log_handle.close()


if __name__ == "__main__":
    main()