scripts/seed_targets.py

"""Seed target identity sightings from `seed_data/targets/{identity}/*.jpg`.

For each identity folder, the script holds out one photo into
`seed_data/holdout/{identity}.{ext}` (so you can use it as a query) and
inserts the rest into the database, marked source='target' identity='{name}'.

The identity column is for evaluation only — the search API never reads it.

Run inside the backend container:
    python -m scripts.seed_targets
"""
import argparse
import logging
import random
import shutil
from pathlib import Path

from PIL import Image
from sqlalchemy import select

from app.config import get_settings
from app.models import LostDog, Sighting
from scripts._common import (
    insert_sighting,
    log,
    offset_within_meters,
    open_session,
    random_city_point,
    random_recent_timestamp,
)

VALID_EXTS = {".jpg", ".jpeg", ".png", ".webp"}


def _is_image(p: Path) -> bool:
    return p.is_file() and p.suffix.lower() in VALID_EXTS


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--force",
        action="store_true",
        help="Re-seed even if target rows already exist.",
    )
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument(
        "--root",
        type=Path,
        default=Path("/seed_data"),
        help="Root containing targets/ and holdout/ subdirs.",
    )
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(name)s: %(message)s")
    rng = random.Random(args.seed)

    targets_root: Path = args.root / "targets"
    holdout_root: Path = args.root / "holdout"
    holdout_root.mkdir(parents=True, exist_ok=True)

    if not targets_root.exists():
        log.info(
            "No targets directory at %s — drop dogs into seed_data/targets/{name}/ "
            "and re-run. Skipping.",
            targets_root,
        )
        return

    identity_dirs = [d for d in sorted(targets_root.iterdir()) if d.is_dir()]
    if not identity_dirs:
        log.info(
            "targets/ is empty — drop ~5 dogs with 2-3 photos each into "
            "seed_data/targets/{name}/ and re-run. Skipping."
        )
        return

    session = open_session()
    try:
        existing = session.scalar(
            select(Sighting).where(Sighting.source == "target").limit(1)
        )
        if existing and not args.force:
            log.info("Target rows already present — skipping (pass --force to re-seed).")
            return
        if existing and args.force:
            # Cascade-delete the LostDog rows; FK ondelete=CASCADE will remove
            # the linked sightings too.
            session.query(LostDog).filter(
                LostDog.id.in_(
                    select(Sighting.lost_dog_id)
                    .where(Sighting.source == "target")
                    .where(Sighting.lost_dog_id.is_not(None))
                )
            ).delete(synchronize_session=False)
            # Catch any orphan target sightings (no lost_dog_id) too.
            session.query(Sighting).filter(Sighting.source == "target").delete()
            session.commit()

        settings = get_settings()
        total_dogs = 0
        total_photos = 0
        for ident_dir in identity_dirs:
            identity = ident_dir.name
            photos = sorted([p for p in ident_dir.iterdir() if _is_image(p)])
            if len(photos) < 2:
                log.warning(
                    "Identity '%s' has %d photo(s); need at least 2 (one for index, "
                    "one as holdout query). Skipping.",
                    identity,
                    len(photos),
                )
                continue

            # One "home turf" point per identity — every photo of this dog is placed
            # within seed_target_radius_m (default 100 m) of this center, so the
            # matched sightings cluster tightly on the map.
            home_lat, home_lng = random_city_point(rng)

            # Hold out one photo deterministically (last one alphabetically).
            holdout_src = photos[-1]
            index_photos = photos[:-1]
            holdout_dst = holdout_root / f"{identity}{holdout_src.suffix.lower()}"
            shutil.copyfile(holdout_src, holdout_dst)
            log.info(
                "Identity '%s': home=(%.5f, %.5f), indexing %d photo(s), holdout=%s",
                identity,
                home_lat,
                home_lng,
                len(index_photos),
                holdout_dst.name,
            )

            # Create the LostDog cluster row first so FKs are valid.
            lost_dog = LostDog(
                name=identity.title(),
                last_seen_at=random_recent_timestamp(rng),
                last_seen_lat=home_lat,
                last_seen_lng=home_lng,
                contact_name=f"{identity.title()}'s Owner",
                contact_email="koenig.andres@gmail.com",
                contact_phone=None,
            )
            session.add(lost_dog)
            session.flush()  # populate lost_dog.id
            total_dogs += 1

            for photo in index_photos:
                try:
                    image = Image.open(photo)
                    image.load()
                except Exception as exc:  # noqa: BLE001
                    log.warning("Could not open %s: %s", photo, exc)
                    continue
                coords = offset_within_meters(
                    rng, home_lat, home_lng, settings.seed_target_radius_m
                )
                row = insert_sighting(
                    session,
                    image,
                    source="target",
                    identity=identity,
                    coords=coords,
                    rng=rng,
                    prefix="target",
                    lost_dog_id=lost_dog.id,
                )
                if row is not None:
                    total_photos += 1

            session.commit()

        log.info(
            "Seeded %d lost-dog clusters with %d total reference photos.",
            total_dogs,
            total_photos,
        )
    finally:
        session.close()


if __name__ == "__main__":
    main()