backend/scripts/prepare_test_data.py

"""Prepare mock test data from a folder of dog-identity subfolders.

Input layout:
    input_dir/
      dog_a/  img1.jpg img2.jpg ...
      dog_b/  img1.jpg ...

Per dog folder, images are split into registration images (-> KnownDog) and holdout images
(-> UnknownDog + found case). Writes known_dogs.csv, found_dogs.csv, and pairs.csv (eval only)
to the output dir, with procedurally-generated owner/finder info. The found ZIP is set equal to
the dog's known ZIP so default-radius matching works.

Usage:
    python -m scripts.prepare_test_data --input-dir PATH [--output-dir PATH] \
        [--holdout N] [--seed N]
"""
from __future__ import annotations

import argparse
import csv
import random
from pathlib import Path

IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".gif"}

# Sample ZIPs that exist in data/zip_centroids.csv so radius matching is meaningful.
SAMPLE_ZIPS = ["10001", "20001", "60601", "94103", "98101", "30301", "78701", "85001"]
COLORS = ["black", "brown", "white", "tan", "golden", "brindle", "gray", "spotted"]
SIZES = ["small", "medium", "large"]
DOG_NAMES = ["Rex", "Bella", "Max", "Luna", "Charlie", "Daisy", "Cooper", "Lucy", "Buddy", "Sadie"]
FIRST_NAMES = ["Alex", "Sam", "Jordan", "Taylor", "Casey", "Morgan", "Riley", "Jamie", "Pat", "Quinn"]
LAST_NAMES = ["Smith", "Johnson", "Lee", "Garcia", "Brown", "Davis", "Martinez", "Clark", "Lewis", "Walker"]


def _images_in(folder: Path) -> list[Path]:
    return sorted(p for p in folder.iterdir() if p.is_file() and p.suffix.lower() in IMAGE_EXTS)


def _holdout_count(n_images: int, requested: int | None) -> int:
    """Default 1 holdout; 2 if 5+ images and not specified. Always leave >=1 registration image."""
    if requested is not None:
        h = requested
    else:
        h = 2 if n_images >= 5 else 1
    return max(0, min(h, n_images - 1))


def _fake_person(rng: random.Random) -> tuple[str, str, str]:
    name = f"{rng.choice(FIRST_NAMES)} {rng.choice(LAST_NAMES)}"
    handle = name.lower().replace(" ", ".")
    email = f"{handle}{rng.randint(1, 999)}@example.com"
    phone = f"{rng.randint(200, 989)}-{rng.randint(200, 989)}-{rng.randint(1000, 9999)}"
    return name, email, phone


def prepare(input_dir: Path, output_dir: Path, holdout: int | None, seed: int) -> dict:
    rng = random.Random(seed)
    subfolders = sorted(p for p in input_dir.iterdir() if p.is_dir())

    known_rows: list[dict] = []
    found_rows: list[dict] = []
    pair_rows: list[dict] = []
    skipped: list[str] = []

    for folder in subfolders:
        images = _images_in(folder)
        if not images:
            skipped.append(folder.name)
            continue

        # Per-dog attributes (deterministic given seed + folder order).
        dog_name = rng.choice(DOG_NAMES)
        color = rng.choice(COLORS)
        size = rng.choice(SIZES)
        zip_code = rng.choice(SAMPLE_ZIPS)
        owner_name, owner_email, owner_phone = _fake_person(rng)
        finder_name, finder_email, finder_phone = _fake_person(rng)
        description = f"{size} {color} dog"

        n_hold = _holdout_count(len(images), holdout)
        shuffled = images[:]
        rng.shuffle(shuffled)
        holdouts = shuffled[:n_hold]
        registrations = shuffled[n_hold:]

        for img in registrations:
            known_rows.append(
                {
                    "folder": folder.name,
                    "image_file": img.name,
                    "dog_name": dog_name,
                    "color": color,
                    "size": size,
                    "description": description,
                    "zip": zip_code,
                    "owner_name": owner_name,
                    "owner_email": owner_email,
                    "owner_phone": owner_phone,
                }
            )
        for img in holdouts:
            found_rows.append(
                {
                    "folder": folder.name,
                    "image_file": img.name,
                    "description": f"found {description}",
                    "color": color,
                    "size": size,
                    "found_zip": zip_code,  # == known zip so default-radius matching works
                    "current_location": "Local shelter",
                    "finder_name": finder_name,
                    "finder_email": finder_email,
                    "finder_phone": finder_phone,
                }
            )
        pair_rows.append(
            {
                "dog_folder": folder.name,
                "known_images": ";".join(i.name for i in registrations),
                "found_images": ";".join(i.name for i in holdouts),
            }
        )

    output_dir.mkdir(parents=True, exist_ok=True)
    _write_csv(output_dir / "known_dogs.csv", known_rows,
               ["folder", "image_file", "dog_name", "color", "size", "description",
                "zip", "owner_name", "owner_email", "owner_phone"])
    _write_csv(output_dir / "found_dogs.csv", found_rows,
               ["folder", "image_file", "description", "color", "size", "found_zip",
                "current_location", "finder_name", "finder_email", "finder_phone"])
    _write_csv(output_dir / "pairs.csv", pair_rows,
               ["dog_folder", "known_images", "found_images"])

    return {
        "dogs": len(pair_rows),
        "known_images": len(known_rows),
        "found_images": len(found_rows),
        "skipped_folders": skipped,
    }


def _write_csv(path: Path, rows: list[dict], fieldnames: list[str]) -> None:
    with path.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    parser = argparse.ArgumentParser(description="Prepare mock test data from dog-identity folders.")
    parser.add_argument("--input-dir", required=True, type=Path)
    parser.add_argument("--output-dir", type=Path, default=None)
    parser.add_argument("--holdout", type=int, default=None, help="Holdout images per dog (default 1; 2 if 5+ images)")
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    output_dir = args.output_dir or args.input_dir
    result = prepare(args.input_dir, output_dir, args.holdout, args.seed)
    print(f"Prepared {result['dogs']} dog(s): "
          f"{result['known_images']} registration + {result['found_images']} holdout image(s).")
    if result["skipped_folders"]:
        print(f"Skipped (no images): {', '.join(result['skipped_folders'])}")
    print(f"Wrote known_dogs.csv, found_dogs.csv, pairs.csv to {output_dir}")


if __name__ == "__main__":
    main()