# SPDX-License-Identifier: BSD-3-Clause

"""Build a frozen evaluation set of 360-degree panorama tasks.

Discovery is city-seeded and adaptive, which measured far better than the
road-point probing in `build_pano_tasks.py`:

    45 m probe boxes at Street View coordinates ..........  5% of seeds hit
    adaptive box, jittered 28 km around cities ...........  3%
    adaptive box, ±2 km around cities .................... 21%
    adaptive box, city centres ........................... 26%, 16 sequences/hit

Two things drive that. Panoramas are extremely clustered — one hit typically
yields a dozen or more sequences from the same contributor's drive — so the
useful unit of discovery is a neighbourhood, not a point. And `is_pano=true` is
a real server-side filter, so one request answers "are there panoramas here"
instead of returning a hundred phone photographs to sift.

An eval set is not a training pool, so this applies three extra rules:

- **balance over size**: capped per country, because at n=100 balance decides
  what the number means
- **disjoint from training**: any sequence already in the training index is
  rejected, and so is anything within a spatial buffer of it, since frames sit
  about 3.3 m apart and holding out an image holds out nothing
- **mirrored and checksummed**: the imagery is downloaded and hashed, because
  Mapillary thumbnail URLs expire and uploads get deleted, and an eval whose
  images can change is not an eval

Usage:
    python scripts/build_eval_set.py --tasks 100 --per-country 2
    python scripts/build_eval_set.py --tasks 100 --exclude tasks/pano_v1.jsonl
"""

from __future__ import annotations

import argparse
import concurrent.futures
import functools
import hashlib
import json
import logging
import math
import pathlib
import random
import sys
import urllib.request

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
sys.path.insert(0, str(pathlib.Path(__file__).resolve().parents[1]))

from build_pano_tasks import (  # noqa: E402
    _get,
    _month,
    _token,
    frame_detail,
    sequence_frames,
    TIMEOUT_S,
)
from matplotlib.path import Path as MplPath  # noqa: E402
from server.render.minimap import locate  # noqa: E402


# Boxes are tried widest first; a dense city rejects the wide ones with a 500
# because the scan happens before the filter, so the ladder walks down.
BBOX_LADDER = (0.005, 0.002, 0.0006)

logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("build_eval_set")


def _box(lat: float, lon: float, half: float) -> str:
    return f"{lon - half},{lat - half},{lon + half},{lat + half}"


def _haversine_km(lat_a: float, lon_a: float, lat_b: float, lon_b: float) -> float:
    phi_a, phi_b = math.radians(lat_a), math.radians(lat_b)
    d_phi, d_lambda = phi_b - phi_a, math.radians(lon_b - lon_a)
    h = (
        math.sin(d_phi / 2) ** 2
        + math.cos(phi_a) * math.cos(phi_b) * math.sin(d_lambda / 2) ** 2
    )
    return 2 * 6371.0088 * math.asin(math.sqrt(min(1.0, h)))


@functools.lru_cache(maxsize=1)
def _country_polygons() -> list[tuple[str, MplPath]]:
    """Country name and outline for every admin-0 polygon, loaded once."""
    root = pathlib.Path(__file__).resolve().parents[1]
    features = json.loads(
        (root / "data" / "geo" / "ne_110m_admin_0_countries.geojson").read_text()
    )["features"]
    out = []
    for feature in features:
        geometry = feature["geometry"]
        rings = (
            [geometry["coordinates"][0]]
            if geometry["type"] == "Polygon"
            else [poly[0] for poly in geometry["coordinates"]]
        )
        for ring in rings:
            out.append((feature["properties"]["ADMIN"], MplPath(ring)))
    return out


def _country_of(lat: float, lon: float) -> str | None:
    """Country containing a coordinate, or `None` over water.

    Cheaper than `locate()`, which also scans every populated place for the
    nearest city — needless when only the country is wanted for seeding.
    """
    for name, path in _country_polygons():
        if path.contains_point((lon, lat)):
            return name
    return None


def load_seeds(
    geo_dir: pathlib.Path, limit: int, seed: int, stratify: bool = True
) -> list[tuple[str, float, float]]:
    """
    City seeds, ordered to spread across countries.

    Panorama coverage clusters hard: a first run probed 900 cities and found
    1,300 sequences, but from only 30 countries — so a per-country cap ran out
    of countries long before it ran out of sequences, and the set stopped at 58
    of 100. Round-robin ordering by country means the first N probes touch N
    different countries instead of hammering whichever metros the shuffle picked.

    Args:
        geo_dir (`pathlib.Path`):
            Directory holding the bundled Natural Earth files.
        limit (`int`):
            How many seeds to return.
        seed (`int`):
            Seed for the shuffle within each country.
        stratify (`bool`, *optional*, defaults to `True`):
            Order round-robin by country rather than at random.

    Returns:
        `list[tuple[str, float, float]]`: `(name, lat, lon)` seeds.
    """
    detail = geo_dir / "detail" / "places.json"
    if detail.exists():
        rows = json.loads(detail.read_text())
        points = [(r.get("n") or "", r["c"][1], r["c"][0]) for r in rows]
    else:
        features = json.loads((geo_dir / "ne_50m_populated_places.geojson").read_text())
        points = [
            (
                f["properties"]["name"],
                f["geometry"]["coordinates"][1],
                f["geometry"]["coordinates"][0],
            )
            for f in features["features"]
        ]
    rng = random.Random(seed)
    rng.shuffle(points)
    if not stratify:
        return points[:limit]

    by_country: dict[str | None, list] = {}
    for name, lat, lon in points:
        by_country.setdefault(_country_of(lat, lon), []).append((name, lat, lon))
    by_country.pop(None, None)

    buckets = list(by_country.values())
    rng.shuffle(buckets)
    ordered: list[tuple[str, float, float]] = []
    depth = 0
    while len(ordered) < limit and any(len(b) > depth for b in buckets):
        for bucket in buckets:
            if len(bucket) > depth:
                ordered.append(bucket[depth])
                if len(ordered) >= limit:
                    break
        depth += 1
    logger.info(
        "seeded %d cities round-robin across %d countries", len(ordered), len(buckets)
    )
    return ordered


def discover(token: str, seeds, workers: int) -> dict[str, dict]:
    """Probe each seed and return one anchor image per panorama sequence."""

    def probe(seed):
        name, lat, lon = seed
        for half in BBOX_LADDER:
            result = _get(
                token,
                "images",
                bbox=_box(lat, lon, half),
                fields="id,sequence,camera_type,quality_score,computed_geometry",
                is_pano="true",
            )
            if "__error" not in result:
                return name, result.get("data", [])
        return name, []

    found: dict[str, dict] = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
        for index, (name, images) in enumerate(pool.map(probe, seeds), 1):
            for image in images:
                sequence = image.get("sequence")
                if not sequence or sequence in found:
                    continue
                if image.get("camera_type") != "spherical":
                    continue
                found[sequence] = {**image, "__seed": name}
            if index % 50 == 0:
                logger.info(
                    "  %d/%d seeds probed, %d sequences", index, len(seeds), len(found)
                )
    return found


def mirror_frames(
    token: str, image_ids: list[str], cache_dir: pathlib.Path, workers: int
) -> dict[str, str]:
    """
    Download and hash a task's frames, in parallel.

    An eval that prefetches all 24 frames so `move()` can be evaluated offline
    was spending 24 sequential round trips per task — about 30 s of the 34 s a
    task cost. Fetching them concurrently makes mirroring roughly free next to
    frame assembly.

    Args:
        token (`str`):
            Mapillary access token.
        image_ids (`list[str]`):
            Frames to mirror.
        cache_dir (`pathlib.Path`):
            Where the JPEGs live.
        workers (`int`):
            Concurrent downloads.

    Returns:
        `dict[str, str]`: sha256 by image id, for every frame that landed.
    """

    def fetch_one(image_id: str) -> tuple[str, str | None]:
        path = cache_dir / f"{image_id}.jpg"
        if not path.exists():
            meta = _get(token, image_id, fields="thumb_2048_url")
            url = meta.get("thumb_2048_url") if "__error" not in meta else None
            if not url:
                return image_id, None
            try:
                with urllib.request.urlopen(url, timeout=TIMEOUT_S) as response:
                    path.write_bytes(response.read())
            except Exception as exc:  # noqa: BLE001
                logger.warning("  fetch failed for %s: %r", image_id, exc)
                return image_id, None
        return image_id, hashlib.sha256(path.read_bytes()).hexdigest()

    checksums: dict[str, str] = {}
    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as pool:
        for image_id, digest in pool.map(fetch_one, image_ids):
            if digest:
                checksums[image_id] = digest
    return checksums


def load_exclusions(
    paths: list[pathlib.Path],
) -> tuple[set[str], list[tuple[float, float]]]:
    """Sequences and coordinates that a new task must stay away from."""
    sequences: set[str] = set()
    points: list[tuple[float, float]] = []
    for path in paths:
        if not path.exists():
            logger.warning("exclusion file missing, skipping: %s", path)
            continue
        for line in path.read_text().splitlines():
            if not line.strip():
                continue
            row = json.loads(line)
            if row.get("sequence_id"):
                sequences.add(row["sequence_id"])
            for frame in row.get("frames", []):
                points.append((frame["lat"], frame["lon"]))
    return sequences, points


def build(args: argparse.Namespace) -> None:
    """Discover, filter, mirror and write the eval set."""
    token = _token()
    root = pathlib.Path(__file__).resolve().parents[1]
    out_dir = args.out
    cache_dir = args.cache
    out_dir.mkdir(parents=True, exist_ok=True)
    cache_dir.mkdir(parents=True, exist_ok=True)

    excluded_sequences, excluded_points = load_exclusions(args.exclude)
    logger.info(
        "excluding %d training sequences and staying %.1f km from %d of their frames",
        len(excluded_sequences),
        args.buffer_km,
        len(excluded_points),
    )

    seeds = load_seeds(
        root / "data" / "geo", args.seeds, args.seed, stratify=not args.no_stratify
    )
    logger.info("probing %d city seeds with %d workers", len(seeds), args.workers)
    candidates = discover(token, seeds, args.workers)
    logger.info("found %d panorama sequences", len(candidates))

    order = sorted(
        candidates.items(),
        key=lambda kv: -(kv[1].get("quality_score") or 0.0),
    )
    per_country: dict[str, int] = {}
    tasks: list[dict] = []

    rejected_early = 0
    for sequence_id, anchor in order:
        if len(tasks) >= args.tasks:
            break
        if sequence_id in excluded_sequences:
            continue

        # Reject before assembling, not after. Discovery already returned the
        # anchor's coordinates, so the country cap and the training buffer can
        # both be checked for free — and with coverage as clustered as it is,
        # most candidates fall in an already-full country. Assembling 25
        # requests' worth of frames first made a task cost 14.4 s instead of 5.
        geometry = anchor.get("computed_geometry") or anchor.get("geometry")
        if geometry:
            anchor_lon, anchor_lat = geometry["coordinates"]
            country_guess = _country_of(anchor_lat, anchor_lon)
            if (
                country_guess is None
                or per_country.get(country_guess, 0) >= args.per_country
            ):
                rejected_early += 1
                continue
            if any(
                _haversine_km(anchor_lat, anchor_lon, lat, lon) < args.buffer_km
                for lat, lon in excluded_points
            ):
                rejected_early += 1
                continue

        ids = sequence_frames(token, sequence_id)
        if len(ids) < args.min_frames:
            continue
        anchor_id = str(anchor["id"])
        position = ids.index(anchor_id) if anchor_id in ids else 0
        low = max(0, position - args.frames // 2)
        window = ids[low : low + args.frames]

        with concurrent.futures.ThreadPoolExecutor(max_workers=args.workers) as pool:
            details = list(pool.map(lambda i: frame_detail(token, i), window))

        frames = []
        for detail in details:
            if not detail or not detail.get("is_pano"):
                continue
            geometry = detail.get("computed_geometry") or detail.get("geometry")
            if not geometry:
                continue
            lon, lat = geometry["coordinates"]
            frames.append(
                {
                    "image_id": str(detail["id"]),
                    "lat": lat,
                    "lon": lon,
                    "compass_angle": detail.get("computed_compass_angle")
                    or detail.get("compass_angle")
                    or 0.0,
                    "captured_at": _month(detail.get("captured_at")),
                    "is_pano": True,
                }
            )
        if len(frames) < args.min_frames:
            continue

        # A missing captured_at renders as an epoch date, which is harmless for
        # scoring but breaks a temporal split and looks wrong in the UI.
        # Mapillary launched in 2013.
        frames = [
            f for f in frames if not f["captured_at"] or f["captured_at"][:4] >= "2010"
        ]
        if len(frames) < args.min_frames:
            continue

        start = min(len(frames) // 2, len(frames) - 1)
        start_frame = frames[start]

        too_close = any(
            _haversine_km(start_frame["lat"], start_frame["lon"], lat, lon)
            < args.buffer_km
            for lat, lon in excluded_points
        )
        if too_close:
            continue

        place = locate(start_frame["lat"], start_frame["lon"])
        country = place.country or "unknown"
        if country == "unknown" or per_country.get(country, 0) >= args.per_country:
            continue

        anchor_detail = next(
            (d for d in details if d and str(d["id"]) == start_frame["image_id"]), {}
        )
        creator = anchor_detail.get("creator", {}) or {}

        # Mirror the imagery. An eval whose pictures can change upstream is not
        # an eval, and Mapillary thumbnail URLs expire.
        wanted = [f["image_id"] for f in frames[: args.prefetch]]
        if start_frame["image_id"] not in wanted:
            wanted.append(start_frame["image_id"])
        checksums = mirror_frames(token, wanted, cache_dir, args.workers)

        if start_frame["image_id"] not in checksums:
            continue

        per_country[country] = per_country.get(country, 0) + 1
        tasks.append(
            {
                "task_index": len(tasks),
                "task_id": f"eval-{len(tasks):04d}",
                "country": country,
                "sequence_id": sequence_id,
                "provider": "mapillary",
                "start_frame": start,
                "frames": frames,
                "attribution": {
                    "creator_username": creator.get("username", ""),
                    "creator_id": creator.get("id", ""),
                    "licence": "CC-BY-SA-4.0",
                    "source": "Mapillary",
                },
                "meta": {
                    "seed_name": anchor.get("__seed", ""),
                    "camera_make": anchor_detail.get("make", ""),
                    "camera_model": anchor_detail.get("model", ""),
                    "quality_score": anchor_detail.get("quality_score"),
                    "continent": place.continent,
                    "subregion": place.subregion,
                    "nearest_city": place.nearest_city,
                    "sha256": checksums,
                    "mirrored_frames": len(checksums),
                },
            }
        )
        logger.info(
            "  eval %3d  %-24s %-8s %2d frames  %d mirrored",
            len(tasks) - 1,
            country[:24],
            start_frame["captured_at"],
            len(frames),
            len(checksums),
        )

    index_path = out_dir / args.name
    with index_path.open("w") as handle:
        for task in tasks:
            handle.write(json.dumps(task) + "\n")

    logger.info(
        "rejected %d candidates before assembly, on country cap or buffer",
        rejected_early,
    )
    countries = sorted(per_country.items(), key=lambda kv: -kv[1])
    logger.info(
        "\nwrote %d eval tasks to %s (%.0f KB)",
        len(tasks),
        index_path,
        index_path.stat().st_size / 1024,
    )
    logger.info("countries: %d -> %s", len(countries), dict(countries[:12]))
    mirrored = sum(t["meta"]["mirrored_frames"] for t in tasks)
    logger.info("mirrored %d images into %s", mirrored, cache_dir)


def main() -> None:
    """Command-line entry point."""
    root = pathlib.Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--tasks", type=int, default=100)
    parser.add_argument("--per-country", type=int, default=2)
    parser.add_argument("--frames", type=int, default=24)
    parser.add_argument("--min-frames", type=int, default=8)
    parser.add_argument("--prefetch", type=int, default=1)
    parser.add_argument("--seeds", type=int, default=900)
    parser.add_argument("--workers", type=int, default=10)
    parser.add_argument("--seed", type=int, default=1234)
    parser.add_argument(
        "--no-stratify",
        action="store_true",
        help="Probe cities in random order instead of round-robin by country.",
    )
    parser.add_argument("--buffer-km", type=float, default=1.0)
    parser.add_argument("--name", default="eval_pano_v1.jsonl")
    parser.add_argument(
        "--exclude",
        nargs="*",
        type=pathlib.Path,
        default=[root / "tasks" / "pano_v1.jsonl"],
    )
    parser.add_argument("--out", type=pathlib.Path, default=root / "tasks")
    parser.add_argument("--cache", type=pathlib.Path, default=root / "data" / "panos")
    build(parser.parse_args())


if __name__ == "__main__":
    main()