dataset/find_task_samples.py

"""
Find representative task samples at three difficulty levels from the dataset.

Easy   : 3–4  elements, noise=0.05
Medium : 5–7  elements, noise=0.10
Hard   : 7–10 elements, noise=0.15

Writes dataset/task_samples.json for inference.py, copies referenced PNGs into
dataset/sample_images/, and rewrites image_path / layer_image_path in the JSON
to point at sample_images/ (paths relative to the dataset directory).

For each sample background, copies the matching saliency array from
dataset/saliency_images/ into dataset/sample_saliency_images/ and sets
sample["saliency_image_path"] (e.g. sample_saliency_images/3960_bg.npy). Run
preprocess_saliency.py first so those .npy files exist.
"""

from __future__ import annotations

import copy
import json
import random
import shutil
import sys
from pathlib import Path

SCRIPT_DIR = Path(__file__).resolve().parent
DATASET_JSON = SCRIPT_DIR / "genposter_5000_images.json"
OUTPUT_JSON = SCRIPT_DIR / "task_samples.json"
SAMPLE_IMAGES_DIR = SCRIPT_DIR / "sample_images"
SAMPLE_IMAGES_PREFIX = "sample_images"
SAMPLE_SALIENCY_DIR = SCRIPT_DIR / "sample_saliency_images"
SAMPLE_SALIENCY_PREFIX = "sample_saliency_images"
SALIENCY_SOURCE_DIR = SCRIPT_DIR / "saliency_images"

TASKS = [
    {
        "task_id": "easy",
        "min_elements": 3,
        "max_elements": 4,
        "noise": 0.05,
        "max_steps": 50,
    },
    {
        "task_id": "medium",
        "min_elements": 5,
        "max_elements": 7,
        "noise": 0.10,
        "max_steps": 100,
    },
    {
        "task_id": "hard",
        "min_elements": 7,
        "max_elements": 10,
        "noise": 0.15,
        "max_steps": 200,
    },
]


def _copy_and_remap_path(
    rel: str,
    sample_images_dir: Path,
    dataset_root: Path,
    missing: list[str],
) -> str:
    name = Path(rel).name
    src = (dataset_root / rel).resolve()
    if not src.is_file():
        missing.append(rel)
        return f"{SAMPLE_IMAGES_PREFIX}/{name}"

    dest = sample_images_dir / name
    shutil.copy2(src, dest)
    return f"{SAMPLE_IMAGES_PREFIX}/{name}"


def _copy_saliency_for_background(
    sample: dict,
    *,
    sample_saliency_dir: Path,
    saliency_source_dir: Path,
    missing: list[str],
) -> None:
    """Match preprocess_saliency output: saliency_images/<bg_stem>.npy"""
    sample.pop("saliency_image_path", None)
    ip = sample.get("image_path")
    if not ip:
        return
    sal_name = f"{Path(ip).stem}.npy"
    src = (saliency_source_dir / sal_name).resolve()
    if not src.is_file():
        missing.append(f"saliency_images/{sal_name}")
        return
    sample_saliency_dir.mkdir(parents=True, exist_ok=True)
    dest = sample_saliency_dir / sal_name
    shutil.copy2(src, dest)
    sample["saliency_image_path"] = f"{SAMPLE_SALIENCY_PREFIX}/{sal_name}"


def copy_media_and_rewrite_paths(
    output: list[dict],
    *,
    sample_images_dir: Path = SAMPLE_IMAGES_DIR,
    sample_saliency_dir: Path = SAMPLE_SALIENCY_DIR,
    saliency_source_dir: Path = SALIENCY_SOURCE_DIR,
    dataset_root: Path = SCRIPT_DIR,
) -> None:
    if sample_images_dir.exists():
        shutil.rmtree(sample_images_dir)
    sample_images_dir.mkdir(parents=True)

    if sample_saliency_dir.exists():
        shutil.rmtree(sample_saliency_dir)
    sample_saliency_dir.mkdir(parents=True)

    missing: list[str] = []
    saliency_missing: list[str] = []
    for entry in output:
        sample = entry["sample"]
        ip = sample.get("image_path")
        if ip:
            sample["image_path"] = _copy_and_remap_path(
                ip, sample_images_dir, dataset_root, missing
            )
        for el in sample.get("elements") or []:
            lp = el.get("layer_image_path")
            if lp:
                el["layer_image_path"] = _copy_and_remap_path(
                    lp, sample_images_dir, dataset_root, missing
                )
        _copy_saliency_for_background(
            sample,
            sample_saliency_dir=sample_saliency_dir,
            saliency_source_dir=saliency_source_dir,
            missing=saliency_missing,
        )

    if missing:
        for m in missing:
            print(f"WARNING: missing source file: {dataset_root / m}", file=sys.stderr)
    if saliency_missing:
        for m in saliency_missing:
            print(
                f"WARNING: missing saliency map (run preprocess_saliency.py): "
                f"{dataset_root / m}",
                file=sys.stderr,
            )


def main() -> None:
    with open(DATASET_JSON, "r", encoding="utf-8") as f:
        dataset = json.load(f)

    rng = random.Random(42)

    output: list[dict] = []
    for task in TASKS:
        candidates = [
            s
            for s in dataset
            if task["min_elements"] <= len(s.get("elements", [])) <= task["max_elements"]
            and s.get("image_path")
        ]
        if not candidates:
            candidates = [
                s
                for s in dataset
                if task["min_elements"] <= len(s.get("elements", [])) <= task["max_elements"]
            ]
        if not candidates:
            print(
                f"WARNING: No samples found for task '{task['task_id']}' "
                f"({task['min_elements']}–{task['max_elements']} elements). Skipping."
            )
            continue

        chosen = rng.choice(candidates)
        output.append(
            {
                "task_id": task["task_id"],
                "noise": task["noise"],
                "max_steps": task["max_steps"],
                "sample": copy.deepcopy(chosen),
            }
        )
        print(
            f"  {task['task_id']:8s}  id={chosen['id']:<6}  "
            f"elements={len(chosen['elements'])}  noise={task['noise']}"
        )

    copy_media_and_rewrite_paths(output)

    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(output, f, indent=2)

    print(f"\nCopied media to {SAMPLE_IMAGES_DIR}")
    print(f"Copied saliency arrays (.npy) to {SAMPLE_SALIENCY_DIR}")
    print(f"Saved {len(output)} task samples to {OUTPUT_JSON}")


if __name__ == "__main__":
    main()