scripts/_common.py

"""Shared helpers for seed scripts."""
import logging
import math
import random
from datetime import datetime, timedelta, timezone

from PIL import Image
from sqlalchemy.orm import Session

from app.config import get_settings
from app.db import SessionLocal
from app.models import Sighting
from app.services.detector import NoDogDetectedError
from app.services.pipeline import process_and_store
from app.services.storage import ensure_bucket

log = logging.getLogger("seeder")

_METERS_PER_DEG_LAT = 111_000.0


def random_geo(rng: random.Random) -> tuple[float, float]:
    """Filler-style scatter: gaussian around the configured city center."""
    s = get_settings()
    return (
        rng.gauss(s.seed_center_lat, s.seed_geo_sigma),
        rng.gauss(s.seed_center_lng, s.seed_geo_sigma),
    )


def random_city_point(rng: random.Random) -> tuple[float, float]:
    """One uniform point inside the city scatter — used as a per-identity 'home'."""
    s = get_settings()
    return (
        rng.gauss(s.seed_center_lat, s.seed_geo_sigma),
        rng.gauss(s.seed_center_lng, s.seed_geo_sigma),
    )


def offset_within_meters(
    rng: random.Random,
    center_lat: float,
    center_lng: float,
    max_meters: float,
) -> tuple[float, float]:
    """Uniform random point inside a disc of `max_meters` radius around (lat, lng)."""
    # sqrt(uniform) keeps the distribution uniform over area, not radius.
    radius_m = max_meters * math.sqrt(rng.random())
    angle = rng.uniform(0.0, 2.0 * math.pi)
    dlat = (radius_m * math.cos(angle)) / _METERS_PER_DEG_LAT
    dlng = (radius_m * math.sin(angle)) / (
        _METERS_PER_DEG_LAT * math.cos(math.radians(center_lat))
    )
    return (center_lat + dlat, center_lng + dlng)


def random_recent_timestamp(rng: random.Random, *, max_days: int = 30) -> datetime:
    seconds = rng.uniform(0, max_days * 24 * 3600)
    return datetime.now(timezone.utc) - timedelta(seconds=seconds)


def insert_sighting(
    session: Session,
    image: Image.Image,
    *,
    source: str,
    identity: str | None,
    rng: random.Random,
    prefix: str,
    coords: tuple[float, float] | None = None,
    lost_dog_id=None,
) -> Sighting | None:
    try:
        stored = process_and_store(image, prefix=prefix)
    except NoDogDetectedError:
        log.warning("Skipping image — no dog detected.")
        return None

    lat, lng = coords if coords is not None else random_geo(rng)
    sighting = Sighting(
        image_url=stored.image_url,
        cropped_url=stored.cropped_url,
        embedding=stored.embedding.tolist(),
        latitude=lat,
        longitude=lng,
        sighted_at=random_recent_timestamp(rng),
        notes=None,
        source=source,
        identity=identity,
        lost_dog_id=lost_dog_id,
    )
    session.add(sighting)
    return sighting


def open_session() -> Session:
    ensure_bucket()
    return SessionLocal()