"""
PyTorch DataLoader for UAV-GeoLoc dataset (Cross-View Geo-Localization).

Dataset structure (discovered empirically):
  - 372 scenes: 171 Country + 200 Terrain + 1 Rot
  - 927K images: 652K query (drone, 512x512 JPEG) + 275K DB (satellite, PNG)
  - Satellite crops: stride = crop_size / 2 (50% overlap), 11 unique sizes (100-1000 px)
  - Query variants: 3 heights (100/125/150m) x 8 azimuths (0-315, step 45) = 24 per scene
  - Camera: 30 deg vertical FOV, top-down, 76 frames per trajectory

Supports:
  - Triplet training with random / semi-positive-aware negative mining
  - Pair training for contrastive learning
  - Evaluation with separate query / DB sets
  - Camera metadata (height, azimuth, GPS) per query
  - GPS-based localization error computation
  - Satellite tiling utility for new data
"""

import json
import logging
import math
import os
import re
import random
from pathlib import Path
from typing import Optional

import numpy as np
import torch
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as T


# ── Index file parsing ──────────────────────────────────────────────────────

def _parse_query_line(line: str):
    """Parse a query index line. Handles paths with spaces.

    Format: <query_path> <label_int> <pos_db_1> [pos_db_2 ...]
    DB paths always contain /DB/img/crop_.
    """
    line = line.strip()
    if not line:
        return None
    db_pattern = re.compile(r'\S*DB/img/crop_\S+')
    db_matches = list(db_pattern.finditer(line))
    if not db_matches:
        return None
    before_db = line[:db_matches[0].start()].rstrip()
    label_match = re.search(r'\s(\d+)\s*$', before_db)
    if not label_match:
        return None
    label = int(label_match.group(1))
    query_path = before_db[:label_match.start()].strip()
    positives = [m.group() for m in db_matches]
    return query_path, label, positives


def _parse_height_rot(query_path: str):
    """Extract height and rotation from query path.

    e.g. '.../height125_rot270/footage/...' -> (125, 270)
    Handles typos like 'eight150', '125_rot315', 'ght100'.
    """
    m = re.search(r'(?:h(?:eigh?t?)?)?(\d{3})_rot(\d+)', query_path)
    if m:
        return int(m.group(1)), int(m.group(2))
    return None, None


# ── GPS utilities ───────────────────────────────────────────────────────────

def load_db_positions(db_pos_path: str) -> dict:
    """Load db_postion.txt -> {crop_filename: (lon, lat, res_x, res_y)}."""
    positions = {}
    if not os.path.isfile(db_pos_path):
        return positions
    with open(db_pos_path) as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) >= 3:
                name = parts[0]
                lon, lat = float(parts[1]), float(parts[2])
                res_x = float(parts[3]) if len(parts) > 3 else 0.0
                res_y = float(parts[4]) if len(parts) > 4 else 0.0
                positions[name] = (lon, lat, res_x, res_y)
    return positions


def haversine_m(lon1, lat1, lon2, lat2):
    """Haversine distance in meters between two GPS points."""
    R = 6_371_000
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlam = math.radians(lon2 - lon1)
    a = math.sin(dphi / 2) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlam / 2) ** 2
    return 2 * R * math.atan2(math.sqrt(a), math.sqrt(1 - a))


# ── Satellite tiling utility ────────────────────────────────────────────────

def tile_satellite_image(
    image: np.ndarray,
    crop_size: int = 200,
    stride: Optional[int] = None,
) -> list:
    """Tile a satellite image following the UAV-GeoLoc convention.

    Args:
        image: HxWx3 numpy array (the merge.tif).
        crop_size: Size of each square crop in pixels.
        stride: Step between crops. Default = crop_size // 2 (50% overlap).

    Returns:
        List of (crop_array, x_idx, y_idx, pixel_x, pixel_y) tuples.
        crop_X_Y.png naming: X = col index, Y = row index.
    """
    if stride is None:
        stride = crop_size // 2
    h, w = image.shape[:2]
    crops = []
    for x_idx, px in enumerate(range(0, w - crop_size + 1, stride)):
        for y_idx, py in enumerate(range(0, h - crop_size + 1, stride)):
            crop = image[py:py + crop_size, px:px + crop_size]
            crops.append((crop, x_idx, y_idx, px, py))
    return crops


# ── Default transforms ──────────────────────────────────────────────────────

def _default_train_query_transform(img_size=224):
    return T.Compose([
        T.Resize((img_size, img_size)),
        T.RandomHorizontalFlip(),
        T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])


def _default_train_db_transform(img_size=224):
    return T.Compose([
        T.Resize((img_size, img_size)),
        T.RandomHorizontalFlip(),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])


def _default_eval_transform(img_size=224):
    return T.Compose([
        T.Resize((img_size, img_size)),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
    ])


# ── Training dataset: triplets ──────────────────────────────────────────────

class UAVGeoLocTrain(Dataset):
    """Training dataset returning (query, positive, negative) triplets.

    Supports semi-positive-aware negative mining: negatives are guaranteed
    to NOT be in the positive or semi-positive set for the query's scene.
    """

    def __init__(
        self,
        root: str,
        query_file: str = "Index/train_query.txt",
        db_file: str = "Index/train_db.txt",
        img_size: int = 224,
        transform_query=None,
        transform_db=None,
    ):
        self.root = root

        self.queries = []
        with open(os.path.join(root, query_file)) as f:
            for line in f:
                parsed = _parse_query_line(line)
                if parsed is None:
                    continue
                query_path, label, positives = parsed
                height, rot = _parse_height_rot(query_path)
                self.queries.append({
                    "path": query_path,
                    "label": label,
                    "positives": positives,
                    "height": height,
                    "rotation": rot,
                })

        self.db_paths = []
        with open(os.path.join(root, db_file)) as f:
            for line in f:
                p = line.strip()
                if p:
                    self.db_paths.append(p)

        # Build label -> all positive+semi-positive DB paths for clean negative mining
        self._label_positives = {}
        for q in self.queries:
            lbl = q["label"]
            if lbl not in self._label_positives:
                self._label_positives[lbl] = set()
            self._label_positives[lbl].update(q["positives"])

        self.transform_query = transform_query or _default_train_query_transform(img_size)
        self.transform_db = transform_db or _default_train_db_transform(img_size)

    def __len__(self):
        return len(self.queries)

    def _load(self, rel_path):
        return Image.open(os.path.join(self.root, rel_path)).convert("RGB")

    def __getitem__(self, idx):
        q = self.queries[idx]

        query_img = self.transform_query(self._load(q["path"]))
        pos_img = self.transform_db(self._load(random.choice(q["positives"])))

        # Negative: random DB image not in this scene's positive set
        pos_set = self._label_positives.get(q["label"], set())
        while True:
            neg_path = random.choice(self.db_paths)
            if neg_path not in pos_set:
                break
        neg_img = self.transform_db(self._load(neg_path))

        return {
            "query": query_img,
            "positive": pos_img,
            "negative": neg_img,
            "label": q["label"],
            "height": q["height"] or 0,
            "rotation": q["rotation"] or 0,
        }


# ── Training dataset: pairs ─────────────────────────────────────────────────

class UAVGeoLocPair(Dataset):
    """Training dataset returning (query, positive) pairs for contrastive learning."""

    def __init__(
        self,
        root: str,
        query_file: str = "Index/train_query.txt",
        img_size: int = 224,
        transform_query=None,
        transform_db=None,
    ):
        self.root = root

        self.queries = []
        with open(os.path.join(root, query_file)) as f:
            for line in f:
                parsed = _parse_query_line(line)
                if parsed is None:
                    continue
                query_path, label, positives = parsed
                height, rot = _parse_height_rot(query_path)
                self.queries.append({
                    "path": query_path,
                    "label": label,
                    "positives": positives,
                    "height": height,
                    "rotation": rot,
                })

        self.transform_query = transform_query or _default_train_query_transform(img_size)
        self.transform_db = transform_db or _default_train_db_transform(img_size)

    def __len__(self):
        return len(self.queries)

    def __getitem__(self, idx):
        q = self.queries[idx]
        query_img = self.transform_query(
            Image.open(os.path.join(self.root, q["path"])).convert("RGB")
        )
        pos_img = self.transform_db(
            Image.open(os.path.join(self.root, random.choice(q["positives"]))).convert("RGB")
        )
        return {
            "query": query_img,
            "positive": pos_img,
            "label": q["label"],
            "height": q["height"] or 0,
            "rotation": q["rotation"] or 0,
        }


# ── Evaluation dataset ──────────────────────────────────────────────────────

class UAVGeoLocEval(Dataset):
    """Evaluation dataset for retrieval. Returns single images with metadata.

    mode="query": loads UAV query images with labels, positives, height, rotation.
    mode="db": loads satellite DB images.
    """

    def __init__(
        self,
        root: str,
        index_file: str,
        mode: str = "query",
        img_size: int = 224,
        transform=None,
    ):
        self.root = root
        self.mode = mode

        self.images = []
        self.labels = []
        self.positives = []
        self.heights = []
        self.rotations = []

        with open(os.path.join(root, index_file)) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                if mode == "db":
                    self.images.append(line)
                else:
                    parsed = _parse_query_line(line)
                    if parsed is None:
                        continue
                    query_path, label, positives = parsed
                    self.images.append(query_path)
                    self.labels.append(label)
                    self.positives.append(positives)
                    h, r = _parse_height_rot(query_path)
                    self.heights.append(h or 0)
                    self.rotations.append(r or 0)

        self.transform = transform or _default_eval_transform(img_size)

        if mode == "db":
            self.path_to_idx = {p: i for i, p in enumerate(self.images)}

    def __len__(self):
        return len(self.images)

    def __getitem__(self, idx):
        path = self.images[idx]
        img = Image.open(os.path.join(self.root, path)).convert("RGB")
        img = self.transform(img)

        out = {"image": img, "path": path, "index": idx}
        if self.mode == "query":
            out["label"] = self.labels[idx]
            out["positives"] = self.positives[idx]
            out["height"] = self.heights[idx]
            out["rotation"] = self.rotations[idx]
        return out


# ── Scene-based dataset (direct from directory, no index files) ─────────────

class UAVGeoLocScene(Dataset):
    """Load data directly from a scene directory. Useful for custom splits
    or scenes not covered by Index files (e.g., Rot subset).

    Returns (query_img, db_positive_img, metadata_dict) pairs.

    Args:
        scene_dir: Path to scene (e.g., .../Country/Australia/Adelaide/AdelaideCBD)
        heights: List of heights to include. Default: [100, 125, 150].
        rotations: List of rotations to include. Default: all 8.
        frames: List of frame indices or None for all.
    """

    def __init__(
        self,
        scene_dir: str,
        heights: Optional[list] = None,
        rotations: Optional[list] = None,
        frames: Optional[list] = None,
        img_size: int = 224,
        transform_query=None,
        transform_db=None,
    ):
        self.scene_dir = scene_dir
        heights = heights or [100, 125, 150]
        rotations = rotations or [0, 45, 90, 135, 180, 225, 270, 315]

        # Check for incomplete scene (missing DB crops or annotations)
        pos_path = os.path.join(scene_dir, "positive.json")
        db_img_dir = os.path.join(scene_dir, "DB", "img")
        missing = []
        if not os.path.isfile(pos_path):
            missing.append("positive.json")
        if not os.path.isdir(db_img_dir):
            missing.append("DB/img/")
        if missing:
            scene_name = os.path.basename(scene_dir)
            raise FileNotFoundError(
                f"Incomplete scene '{scene_name}': missing {', '.join(missing)}. "
                f"17 known incomplete scenes (Edinburgh, London, Manchester, "
                f"Birmingham/JewelleryQuarter, Chicago/__MACOSX) lack DB crops "
                f"and annotations — they cannot be used for training or evaluation."
            )

        # Load positive.json
        with open(pos_path) as f:
            self.positive_map = json.load(f)

        # Load semi_positive.json if available
        semi_path = os.path.join(scene_dir, "semi_positive.json")
        self.semi_positive_map = {}
        if os.path.isfile(semi_path):
            with open(semi_path) as f:
                self.semi_positive_map = json.load(f)

        # Load DB GPS positions
        db_dir = os.path.join(scene_dir, "DB")
        self.db_positions = load_db_positions(os.path.join(db_dir, "db_postion.txt"))

        # Enumerate valid (variant, frame) pairs
        self.samples = []
        query_dir = os.path.join(scene_dir, "query")
        for h in heights:
            for r in rotations:
                variant_name = f"height{h}_rot{r}"
                footage_dir = os.path.join(query_dir, variant_name, "footage")
                if not os.path.isdir(footage_dir):
                    continue

                available_frames = sorted([
                    f for f in os.listdir(footage_dir)
                    if f.endswith((".jpeg", ".jpg"))
                ])

                for frame_file in available_frames:
                    # Extract frame index (e.g. "height100_rot0_38.jpeg" -> "38")
                    m = re.search(r'_(\d+)\.jpe?g$', frame_file)
                    if m is None:
                        continue
                    frame_idx = m.group(1)

                    if frames is not None and int(frame_idx) not in frames:
                        continue

                    # Get positive DB crop(s)
                    pos_crops = self.positive_map.get(frame_idx, [])
                    if not pos_crops:
                        continue

                    self.samples.append({
                        "query_path": os.path.join(footage_dir, frame_file),
                        "frame_idx": frame_idx,
                        "height": h,
                        "rotation": r,
                        "positives": pos_crops,
                        "semi_positives": self.semi_positive_map.get(frame_idx, []),
                    })

        self.db_img_dir = os.path.join(db_dir, "img")

        self.transform_query = transform_query or _default_eval_transform(img_size)
        self.transform_db = transform_db or _default_eval_transform(img_size)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        s = self.samples[idx]

        query_img = self.transform_query(
            Image.open(s["query_path"]).convert("RGB")
        )
        pos_crop_name = random.choice(s["positives"])
        pos_img = self.transform_db(
            Image.open(os.path.join(self.db_img_dir, pos_crop_name)).convert("RGB")
        )

        # GPS of positive crop centroid
        gps = self.db_positions.get(pos_crop_name, (0.0, 0.0, 0.0, 0.0))

        return {
            "query": query_img,
            "positive": pos_img,
            "height": s["height"],
            "rotation": s["rotation"],
            "frame_idx": int(s["frame_idx"]),
            "positive_name": pos_crop_name,
            "positive_lon": gps[0],
            "positive_lat": gps[1],
            "semi_positives": s["semi_positives"],
        }


# ── Collate functions ───────────────────────────────────────────────────────

def eval_collate_fn(batch):
    """Collate for eval datasets (handles variable-length positives)."""
    images = torch.stack([item["image"] for item in batch])
    indices = torch.tensor([item["index"] for item in batch])
    paths = [item["path"] for item in batch]
    out = {"image": images, "index": indices, "path": paths}
    if "label" in batch[0]:
        out["label"] = torch.tensor([item["label"] for item in batch])
        out["positives"] = [item["positives"] for item in batch]
        out["height"] = torch.tensor([item["height"] for item in batch])
        out["rotation"] = torch.tensor([item["rotation"] for item in batch])
    return out


def scene_collate_fn(batch):
    """Collate for UAVGeoLocScene (handles variable-length semi_positives)."""
    return {
        "query": torch.stack([b["query"] for b in batch]),
        "positive": torch.stack([b["positive"] for b in batch]),
        "height": torch.tensor([b["height"] for b in batch]),
        "rotation": torch.tensor([b["rotation"] for b in batch]),
        "frame_idx": torch.tensor([b["frame_idx"] for b in batch]),
        "positive_name": [b["positive_name"] for b in batch],
        "positive_lon": torch.tensor([b["positive_lon"] for b in batch], dtype=torch.float64),
        "positive_lat": torch.tensor([b["positive_lat"] for b in batch], dtype=torch.float64),
        "semi_positives": [b["semi_positives"] for b in batch],
    }


# ── Localization error evaluation ───────────────────────────────────────────

def compute_localization_error(
    query_dataset: UAVGeoLocEval,
    db_dataset: UAVGeoLocEval,
    predictions: np.ndarray,
    db_positions_cache: Optional[dict] = None,
) -> dict:
    """Compute localization error in meters given retrieval predictions.

    Args:
        query_dataset: UAVGeoLocEval in "query" mode.
        db_dataset: UAVGeoLocEval in "db" mode.
        predictions: (N_query,) array of predicted DB indices.
        db_positions_cache: Pre-loaded {db_path: (lon, lat, ...)} dict.
            If None, loads from db_postion.txt files on the fly.

    Returns:
        dict with 'mean_error_m', 'median_error_m', 'errors' (per-query list).
    """
    if db_positions_cache is None:
        db_positions_cache = {}
        # Discover all unique scene DB dirs from db paths
        scene_dirs = set()
        for p in db_dataset.images:
            # e.g. "Terrain/Mountain/Andes/DB/img/crop_0_0.png" -> "Terrain/Mountain/Andes/DB"
            db_dir = str(Path(p).parent.parent)
            scene_dirs.add(db_dir)
        for sd in scene_dirs:
            pos_file = os.path.join(db_dataset.root, sd, "db_postion.txt")
            positions = load_db_positions(pos_file)
            for crop_name, coords in positions.items():
                full_path = os.path.join(sd, "img", crop_name)
                db_positions_cache[full_path] = coords

    errors = []
    for i, pred_idx in enumerate(predictions):
        # Get GT positive crops for this query
        gt_positives = query_dataset.positives[i]
        # Get predicted DB path
        pred_path = db_dataset.images[int(pred_idx)]

        # GPS of prediction
        pred_gps = db_positions_cache.get(pred_path)
        if pred_gps is None:
            continue

        # GPS of nearest GT positive
        min_dist = float("inf")
        for gt_path in gt_positives:
            gt_gps = db_positions_cache.get(gt_path)
            if gt_gps is None:
                continue
            dist = haversine_m(pred_gps[0], pred_gps[1], gt_gps[0], gt_gps[1])
            min_dist = min(min_dist, dist)

        if min_dist < float("inf"):
            errors.append(min_dist)

    errors_arr = np.array(errors)
    return {
        "mean_error_m": float(np.mean(errors_arr)) if len(errors_arr) else 0.0,
        "median_error_m": float(np.median(errors_arr)) if len(errors_arr) else 0.0,
        "errors": errors,
        "num_evaluated": len(errors),
    }


# ── Convenience builder ─────────────────────────────────────────────────────

def build_dataloaders(
    root: str,
    split: str = "terrain",
    batch_size: int = 32,
    img_size: int = 224,
    num_workers: int = 4,
    mode: str = "triplet",
):
    """Build train/val/test dataloaders.

    Args:
        root: Path to UAV-GeoLoc dataset root.
        split: "terrain" (default), "country", or "all".
        batch_size: Batch size for all loaders.
        img_size: Resize images to (img_size, img_size).
        num_workers: DataLoader workers.
        mode: "triplet" for (query, pos, neg) or "pair" for (query, pos).

    Returns:
        dict with keys: "train", "val_query", "val_db", "test_query", "test_db"
    """
    suffix = {"terrain": "", "country": "_country", "all": "_all"}[split]

    if mode == "triplet":
        train_ds = UAVGeoLocTrain(
            root,
            query_file=f"Index/train_query{suffix}.txt",
            db_file=f"Index/train_db{suffix}.txt",
            img_size=img_size,
        )
    else:
        train_ds = UAVGeoLocPair(
            root,
            query_file=f"Index/train_query{suffix}.txt",
            img_size=img_size,
        )

    train_loader = DataLoader(
        train_ds,
        batch_size=batch_size,
        shuffle=True,
        num_workers=num_workers,
        pin_memory=True,
        drop_last=True,
    )

    loaders = {"train": train_loader}

    for phase in ["val", "test"]:
        q_file = f"Index/{phase}_query{suffix}.txt"
        d_file = f"Index/{phase}_db{suffix}.txt"
        # Fall back to unsuffixed if suffixed file doesn't exist
        if not os.path.isfile(os.path.join(root, q_file)):
            q_file = f"Index/{phase}_query.txt"
        if not os.path.isfile(os.path.join(root, d_file)):
            d_file = f"Index/{phase}_db.txt"

        q_ds = UAVGeoLocEval(root, q_file, mode="query", img_size=img_size)
        d_ds = UAVGeoLocEval(root, d_file, mode="db", img_size=img_size)

        loaders[f"{phase}_query"] = DataLoader(
            q_ds, batch_size=batch_size, shuffle=False,
            num_workers=num_workers, pin_memory=True, collate_fn=eval_collate_fn,
        )
        loaders[f"{phase}_db"] = DataLoader(
            d_ds, batch_size=batch_size, shuffle=False,
            num_workers=num_workers, pin_memory=True, collate_fn=eval_collate_fn,
        )

    return loaders


def build_rot_loader(
    root: str,
    batch_size: int = 32,
    img_size: int = 224,
    num_workers: int = 4,
    heights: Optional[list] = None,
    rotations: Optional[list] = None,
):
    """Build a DataLoader for the Rot evaluation subset.

    The Rot subset has 88 query variants (72 at h100 every 5deg + 16 at h125/h150)
    over a single scene (SouthernSuburbs), useful for rotation robustness evaluation.

    Returns:
        DataLoader yielding scene_collate_fn batches.
    """
    scene_dir = os.path.join(root, "Rot", "SouthernSuburbs")
    # Rot has fine-grained rotations at height100
    if rotations is None and heights is None:
        # Include all: h100 has 72 rotations (0-355 step 5), h125/h150 have 8 each
        heights_rot = [(100, r) for r in range(0, 360, 5)]
        heights_rot += [(h, r) for h in [125, 150] for r in range(0, 360, 45)]
        all_heights = list(set(h for h, _ in heights_rot))
        all_rotations = list(set(r for _, r in heights_rot))
    else:
        all_heights = heights or [100, 125, 150]
        all_rotations = rotations or list(range(0, 360, 5))

    ds = UAVGeoLocScene(
        scene_dir,
        heights=all_heights,
        rotations=all_rotations,
        img_size=img_size,
    )
    return DataLoader(
        ds, batch_size=batch_size, shuffle=False,
        num_workers=num_workers, pin_memory=True, collate_fn=scene_collate_fn,
    )


# ── Quick test ──────────────────────────────────────────────────────────────

if __name__ == "__main__":
    ROOT = "/mnt/data1tb/cvgl_datasets/UAV-GeoLoc"

    print("=" * 60)
    print("Building dataloaders (terrain split, img_size=224)...")
    loaders = build_dataloaders(ROOT, split="terrain", batch_size=4, num_workers=0)

    for name, loader in loaders.items():
        print(f"  {name}: {len(loader.dataset)} samples")

    batch = next(iter(loaders["train"]))
    print(f"\nTrain batch:")
    print(f"  query:    {batch['query'].shape}")
    print(f"  positive: {batch['positive'].shape}")
    print(f"  negative: {batch['negative'].shape}")
    print(f"  labels:   {batch['label']}")
    print(f"  heights:  {batch['height']}")
    print(f"  rotations:{batch['rotation']}")

    batch = next(iter(loaders["val_query"]))
    print(f"\nVal query batch:")
    print(f"  image:    {batch['image'].shape}")
    print(f"  heights:  {batch['height']}")
    print(f"  rotations:{batch['rotation']}")

    print("\n" + "=" * 60)
    print("Scene-based loader (AdelaideCBD, h100 only)...")
    scene_ds = UAVGeoLocScene(
        os.path.join(ROOT, "Country/Australia/Adelaide/AdelaideCBD"),
        heights=[100],
        rotations=[0, 90, 180, 270],
    )
    print(f"  Samples: {len(scene_ds)}")
    s = scene_ds[0]
    print(f"  query: {s['query'].shape}, height={s['height']}, rot={s['rotation']}")
    print(f"  positive: {s['positive_name']}, GPS=({s['positive_lon']:.4f}, {s['positive_lat']:.4f})")

    print("\n" + "=" * 60)
    print("Rot evaluation loader...")
    rot_loader = build_rot_loader(ROOT, batch_size=4, num_workers=0, heights=[100], rotations=[0, 45, 90])
    print(f"  Samples: {len(rot_loader.dataset)}")
    batch = next(iter(rot_loader))
    print(f"  query: {batch['query'].shape}")
    print(f"  rotations: {batch['rotation']}")

    print("\n" + "=" * 60)
    print("Tiling utility test...")
    dummy = np.random.randint(0, 255, (500, 600, 3), dtype=np.uint8)
    crops = tile_satellite_image(dummy, crop_size=200, stride=100)
    print(f"  Input: 600x500, crop=200, stride=100")
    print(f"  Crops generated: {len(crops)}")
    print(f"  Grid: {max(c[1] for c in crops)+1} x {max(c[2] for c in crops)+1}")

    print("\nAll tests passed.")