Blame for scripts/_common.py

from __future__ import annotations

import binascii
import csv
import json
import math
import shutil
import struct
import zlib
from pathlib import Path
from typing import Any, Iterable


REPO_ROOT = Path(__file__).resolve().parents[1]


def ensure_dir(path: str | Path) -> Path:
    out = Path(path)
    out.mkdir(parents=True, exist_ok=True)
    return out


def read_json(path: str | Path) -> Any:
    with Path(path).open("r", encoding="utf-8") as f:
        return json.load(f)


def write_json(path: str | Path, obj: Any) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as f:
        json.dump(obj, f, indent=2, sort_keys=True)
        f.write("\n")


def read_jsonl(path: str | Path) -> list[dict[str, Any]]:
    rows: list[dict[str, Any]] = []
    with Path(path).open("r", encoding="utf-8") as f:
        for line_no, line in enumerate(f, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSONL at {path}:{line_no}: {exc}") from exc
            if not isinstance(row, dict):
                raise ValueError(f"Expected a JSON object at {path}:{line_no}")
            rows.append(row)
    return rows


def write_jsonl(path: str | Path, rows: Iterable[dict[str, Any]]) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    with path.open("w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, sort_keys=True))
            f.write("\n")


def write_csv(path: str | Path, rows: list[dict[str, Any]], fieldnames: list[str] | None = None) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    if fieldnames is None:
        keys: list[str] = []
        for row in rows:
            for key in row:
                if key not in keys:
                    keys.append(key)
        fieldnames = keys
    with path.open("w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for row in rows:
            writer.writerow(row)


def _table_fieldnames(rows: list[dict[str, Any]], fieldnames: list[str] | None = None) -> list[str]:
    if fieldnames is not None:
        return fieldnames
    keys: list[str] = []
    for row in rows:
        for key in row:
            if key not in keys:
                keys.append(key)
    return keys


def write_markdown_table(path: str | Path, rows: list[dict[str, Any]], fieldnames: list[str] | None = None) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    fieldnames = _table_fieldnames(rows, fieldnames)
    with path.open("w", encoding="utf-8") as f:
        f.write("| " + " | ".join(fieldnames) + " |\n")
        f.write("| " + " | ".join(["---"] * len(fieldnames)) + " |\n")
        for row in rows:
            f.write("| " + " | ".join(str(row.get(name, "")) for name in fieldnames) + " |\n")


def _latex_escape(value: Any) -> str:
    text = str(value)
    return (
        text.replace("\\", "\\textbackslash{}")
        .replace("&", "\\&")
        .replace("%", "\\%")
        .replace("$", "\\$")
        .replace("#", "\\#")
        .replace("_", "\\_")
        .replace("{", "\\{")
        .replace("}", "\\}")
    )


def write_latex_table(
    path: str | Path,
    rows: list[dict[str, Any]],
    fieldnames: list[str] | None = None,
    caption: str = "Table-ready experiment results.",
    label: str = "tab:cmevs_results",
) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    fieldnames = _table_fieldnames(rows, fieldnames)
    align = "l" * len(fieldnames)
    with path.open("w", encoding="utf-8") as f:
        f.write("\\begin{table}[t]\n")
        f.write("\\centering\n")
        f.write(f"\\caption{{{_latex_escape(caption)}}}\n")
        safe_label = str(label).replace("{", "").replace("}", "")
        f.write(f"\\label{{{safe_label}}}\n")
        f.write(f"\\begin{{tabular}}{{{align}}}\n")
        f.write("\\toprule\n")
        f.write(" & ".join(_latex_escape(name) for name in fieldnames) + " \\\\\n")
        f.write("\\midrule\n")
        for row in rows:
            f.write(" & ".join(_latex_escape(row.get(name, "")) for name in fieldnames) + " \\\\\n")
        f.write("\\bottomrule\n")
        f.write("\\end{tabular}\n")
        f.write("\\end{table}\n")


def copy_file(src: str | Path, dst: str | Path) -> None:
    dst = Path(dst)
    ensure_dir(dst.parent)
    shutil.copy2(src, dst)


def candidate_by_id(candidates: Iterable[dict[str, Any]]) -> dict[str, dict[str, Any]]:
    result = {}
    for row in candidates:
        cid = row.get("candidate_id")
        if not isinstance(cid, str) or not cid.strip():
            raise ValueError("candidate_id must be a nonempty string")
        if cid in result:
            raise ValueError(f"Duplicate candidate_id: {cid}")
        result[cid] = row
    return result


def validate_candidates(candidates: list[dict[str, Any]]) -> None:
    candidate_by_id(candidates)
    if not candidates:
        raise ValueError("Candidate input is empty")
    scenes = set()
    for row in candidates:
        if not isinstance(row.get("valid"), bool):
            raise ValueError("Candidate valid must be a JSON boolean")
        for key in ("source", "scene_id"):
            if not isinstance(row.get(key), str) or not row[key]:
                raise ValueError(f"Candidate {key} must be a nonempty string")
        scenes.add((row["source"], row["scene_id"]))
        position = row.get("position")
        if not isinstance(position, list) or len(position) != 3 or not all(
            isinstance(x, (int, float)) and not isinstance(x, bool) and math.isfinite(x)
            for x in position
        ):
            raise ValueError("Candidate position must contain three finite numbers")
        for key in ("single_view_probe_coverage", "conflict_prior"):
            value = row.get(key)
            if isinstance(value, bool) or not isinstance(value, (int, float)) or not math.isfinite(value) or not 0 <= value <= 1:
                raise ValueError(f"Candidate {key} must be finite and in [0, 1]")
        for key in ("yaw_deg", "oracle_gain"):
            if key in row and (isinstance(row[key], bool) or not isinstance(row[key], (int, float)) or not math.isfinite(row[key])):
                raise ValueError(f"Candidate {key} must be finite")
        cells = row.get("covered_cells")
        if not isinstance(cells, list) or not all(isinstance(x, str) for x in cells):
            raise ValueError("covered_cells must be an array of strings")
    if len(scenes) != 1:
        raise ValueError("A normalized candidate file must describe one source and scene")


def validate_selection(selected_doc: dict[str, Any], by_id: dict[str, dict[str, Any]]) -> None:
    ids = selected_ids(selected_doc)
    for cid in ids:
        if cid not in by_id:
            raise ValueError(f"Unknown selected candidate: {cid}")
        if not by_id[cid].get("valid", True):
            raise ValueError(f"Selected candidate is invalid: {cid}")
        if selected_doc.get("scene_id") != by_id[cid].get("scene_id"):
            raise ValueError(f"Selected candidate belongs to another scene: {cid}")


def valid_candidates(candidates: Iterable[dict[str, Any]]) -> list[dict[str, Any]]:
    return [row for row in candidates if bool(row.get("valid", True))]


def cell_set(candidate: dict[str, Any]) -> set[str]:
    return {str(cell) for cell in candidate.get("covered_cells", [])}


def universe_cells(candidates: Iterable[dict[str, Any]]) -> set[str]:
    cells: set[str] = set()
    for candidate in candidates:
        if bool(candidate.get("valid", True)):
            cells.update(cell_set(candidate))
    return cells


def selected_ids(selected_doc: dict[str, Any]) -> list[str]:
    if not isinstance(selected_doc, dict) or not isinstance(selected_doc.get("selected_viewpoints"), list):
        raise ValueError("Selection must contain a selected_viewpoints array")
    ids = []
    for row in selected_doc["selected_viewpoints"]:
        cid = row.get("candidate_id") if isinstance(row, dict) else None
        if not isinstance(cid, str) or not cid.strip():
            raise ValueError("Selected candidate_id must be a nonempty string")
        if cid in ids:
            raise ValueError(f"Duplicate selected candidate_id: {cid}")
        ids.append(cid)
    return ids


def safe_div(num: float, den: float) -> float:
    return 0.0 if den == 0 else num / den


def pearson(xs: list[float], ys: list[float]) -> float:
    if len(xs) != len(ys) or len(xs) < 2:
        return float("nan")
    mx = sum(xs) / len(xs)
    my = sum(ys) / len(ys)
    num = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    vx = sum((x - mx) ** 2 for x in xs)
    vy = sum((y - my) ** 2 for y in ys)
    if vx <= 0.0 or vy <= 0.0:
        return float("nan")
    return num / math.sqrt(vx * vy)


def _png_chunk(kind: bytes, payload: bytes) -> bytes:
    return (
        struct.pack(">I", len(payload))
        + kind
        + payload
        + struct.pack(">I", binascii.crc32(kind + payload) & 0xFFFFFFFF)
    )


def write_solid_png(path: str | Path, width: int, height: int, rgb: tuple[int, int, int]) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    raw = bytearray()
    row = bytes(rgb) * width
    for _ in range(height):
        raw.append(0)
        raw.extend(row)
    ihdr = struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0)
    data = zlib.compress(bytes(raw), level=9)
    with path.open("wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(_png_chunk(b"IHDR", ihdr))
        f.write(_png_chunk(b"IDAT", data))
        f.write(_png_chunk(b"IEND", b""))


def write_npy_f4(path: str | Path, height: int, width: int, value: float) -> None:
    path = Path(path)
    ensure_dir(path.parent)
    header = "{'descr': '<f4', 'fortran_order': False, 'shape': (%d, %d), }" % (height, width)
    header_bytes = header.encode("latin1")
    prefix_len = 6 + 2 + 2
    padding = 16 - ((prefix_len + len(header_bytes) + 1) % 16)
    header_bytes += b" " * padding + b"\n"
    row = struct.pack("<" + "f" * width, *([float(value)] * width))
    with path.open("wb") as f:
        f.write(b"\x93NUMPY")
        f.write(b"\x01\x00")
        f.write(struct.pack("<H", len(header_bytes)))
        f.write(header_bytes)
        for _ in range(height):
            f.write(row)