#!/usr/bin/env python3
"""Prepare UAV-VisLoc dataset for retrieval training.

Pipeline:
  1. Resize drone images -> 256x256
  2. Stitch satellite tiles for route 09
  3. Crop satellite maps -> 512x512 patches with stride 256, resize -> 256x256
  4. Compute GPS for each crop center
  5. Match drone -> crops via GPS (positive, semi-positive)
  6. Generate metadata: positive.json, semi_positive.json, db_postion.txt
  7. Generate Index files (train_query.txt, train_db.txt, etc.)

Usage:
    python scripts/prepare_dataset.py \
        --src /path/to/UAV_VisLoc_dataset \
        --dst /path/to/UAV_VisLoc_processed \
        --crop-size 512 \
        --stride 256 \
        --target-size 256
"""

import argparse
import csv
import json
import math
import os
import warnings
from pathlib import Path

import numpy as np
from PIL import Image

warnings.filterwarnings("ignore", category=Image.DecompressionBombWarning)
Image.MAX_IMAGE_PIXELS = None

# Route 07 excluded: satellite map 3000x170 — too narrow for 512x512 crops
EXCLUDED_ROUTES = {"07"}

# Route 09 has satellite split into 4 tiles
SPLIT_TILE_ROUTE = "09"
TILE_LAYOUT_09 = {
    # (col, row): filename
    (0, 0): "satellite09_01-01.tif",
    (1, 0): "satellite09_01-02.tif",
    (0, 1): "satellite09_02-01.tif",
    (1, 1): "satellite09_02-02.tif",
}


def parse_args():
    parser = argparse.ArgumentParser(description="Prepare UAV-VisLoc for retrieval.")
    parser.add_argument("--src", required=True, help="Path to raw UAV_VisLoc_dataset")
    parser.add_argument("--dst", required=True, help="Path to output processed dataset")
    parser.add_argument("--crop-size", type=int, default=512, help="Satellite crop size in px")
    parser.add_argument("--stride", type=int, default=256, help="Crop stride in px")
    parser.add_argument("--target-size", type=int, default=256, help="Final resize for both drone and crops")
    parser.add_argument("--routes", nargs="*", default=None, help="Process only these routes (e.g. 01 02)")
    return parser.parse_args()


# ---------------------------------------------------------------------------
# 1. Read metadata
# ---------------------------------------------------------------------------

def read_satellite_bbox(src: Path) -> dict:
    """Read satellite GPS bounding boxes from CSV."""
    bbox = {}
    csv_path = src / "satellite_ coordinates_range.csv"
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            name = row["mapname"]
            route = name.replace("satellite", "").replace(".tif", "")
            bbox[route] = {
                "lt_lat": float(row["LT_lat_map"]),
                "lt_lon": float(row["LT_lon_map"]),
                "rb_lat": float(row["RB_lat_map"]),
                "rb_lon": float(row["RB_lon_map"]),
            }
    return bbox


def read_drone_metadata(src: Path, route: str) -> list[dict]:
    """Read drone GPS + pose from per-route CSV."""
    csv_path = src / route / f"{route}.csv"
    entries = []
    with open(csv_path) as f:
        reader = csv.DictReader(f)
        for row in reader:
            entries.append({
                "filename": row["filename"],
                "lat": float(row["lat"]),
                "lon": float(row["lon"]),
                "height": float(row["height"]),
            })
    return entries


def read_split(src: Path, split: str) -> set[str]:
    """Read train/test split CSV, return set of drone filenames."""
    csv_path = src / f"visloc_{split}.csv"
    filenames = set()
    with open(csv_path) as f:
        reader = csv.DictReader(f, delimiter="\t")
        for row in reader:
            fn = row["filename"]
            basename = os.path.basename(fn)
            filenames.add(basename)
    return filenames


# ---------------------------------------------------------------------------
# 2. Resize drone images
# ---------------------------------------------------------------------------

def resize_drone_images(src: Path, dst: Path, route: str, target_size: int) -> int:
    """Resize all drone images for a route to target_size x target_size."""
    drone_src = src / route / "drone"
    drone_dst = dst / route / "drone"
    drone_dst.mkdir(parents=True, exist_ok=True)

    count = 0
    for img_name in sorted(os.listdir(drone_src)):
        if not img_name.upper().endswith(".JPG"):
            continue
        src_path = drone_src / img_name
        dst_path = drone_dst / img_name

        if dst_path.exists():
            count += 1
            continue

        with Image.open(src_path) as img:
            img_resized = img.resize((target_size, target_size), Image.LANCZOS)
            img_resized.save(dst_path, "JPEG", quality=95)
        count += 1

    return count


# ---------------------------------------------------------------------------
# 3. Stitch satellite tiles (route 09)
# ---------------------------------------------------------------------------

def stitch_route09(src: Path) -> Image.Image:
    """Stitch 4 satellite tiles for route 09 into a single image."""
    tiles = {}
    for (col, row), fname in TILE_LAYOUT_09.items():
        tile_path = src / SPLIT_TILE_ROUTE / fname
        tiles[(col, row)] = Image.open(tile_path)

    # Compute full image dimensions
    # Row 0: tiles (0,0) and (1,0) side by side
    # Row 1: tiles (0,1) and (1,1) side by side
    w0 = tiles[(0, 0)].width + tiles[(1, 0)].width
    w1 = tiles[(0, 1)].width + tiles[(1, 1)].width
    full_w = max(w0, w1)

    h0 = max(tiles[(0, 0)].height, tiles[(1, 0)].height)
    h1 = max(tiles[(0, 1)].height, tiles[(1, 1)].height)
    full_h = h0 + h1

    merged = Image.new("RGB", (full_w, full_h))
    merged.paste(tiles[(0, 0)], (0, 0))
    merged.paste(tiles[(1, 0)], (tiles[(0, 0)].width, 0))
    merged.paste(tiles[(0, 1)], (0, h0))
    merged.paste(tiles[(1, 1)], (tiles[(0, 1)].width, h0))

    # Close tiles to free memory (~4.3 GB).
    for tile in tiles.values():
        tile.close()

    return merged


def load_satellite(src: Path, route: str) -> Image.Image:
    """Load satellite map for a route (handles route 09 stitching)."""
    if route == SPLIT_TILE_ROUTE:
        return stitch_route09(src)
    sat_path = src / route / f"satellite{route}.tif"
    return Image.open(sat_path)


# ---------------------------------------------------------------------------
# 4. Crop satellite map
# ---------------------------------------------------------------------------

def crop_satellite(
    sat_img: Image.Image,
    dst: Path,
    route: str,
    crop_size: int,
    stride: int,
    target_size: int,
) -> list[dict]:
    """Crop satellite image into patches and save resized versions.

    Returns list of crop metadata: [{"name": "crop_X_Y.png", "x": X, "y": Y, "px_x": ..., "px_y": ...}]
    """
    crop_dir = dst / route / "DB" / "img"
    crop_dir.mkdir(parents=True, exist_ok=True)

    w, h = sat_img.size
    cols = max(0, (w - crop_size) // stride + 1)
    rows = max(0, (h - crop_size) // stride + 1)

    crops_meta = []

    for cx in range(cols):
        for cy in range(rows):
            px_x = cx * stride
            px_y = cy * stride

            crop_name = f"crop_{cx}_{cy}.png"
            crop_path = crop_dir / crop_name

            if not crop_path.exists():
                box = (px_x, px_y, px_x + crop_size, px_y + crop_size)
                patch = sat_img.crop(box)
                patch_resized = patch.resize((target_size, target_size), Image.LANCZOS)
                patch_resized.save(crop_path, "PNG")
                patch.close()
                patch_resized.close()

            crops_meta.append({
                "name": crop_name,
                "x": cx,
                "y": cy,
                "px_x": px_x,
                "px_y": px_y,
            })

    return crops_meta


# ---------------------------------------------------------------------------
# 5. Compute GPS for crops + match drone -> crops
# ---------------------------------------------------------------------------

def compute_crop_gps(crops_meta: list[dict], bbox: dict, sat_w: int, sat_h: int, crop_size: int) -> list[dict]:
    """Add GPS center coordinates to each crop metadata entry."""
    lt_lat = bbox["lt_lat"]
    lt_lon = bbox["lt_lon"]
    rb_lat = bbox["rb_lat"]
    rb_lon = bbox["rb_lon"]

    for crop in crops_meta:
        center_px_x = crop["px_x"] + crop_size / 2
        center_px_y = crop["px_y"] + crop_size / 2

        crop["lon"] = lt_lon + (center_px_x / sat_w) * (rb_lon - lt_lon)
        crop["lat"] = lt_lat + (center_px_y / sat_h) * (rb_lat - lt_lat)

    return crops_meta


def _frame_id(filename: str) -> str:
    """Extract frame ID from drone filename: '01_0001.JPG' -> '0001'."""
    stem = os.path.splitext(filename)[0]  # '01_0001'
    return stem.split("_", 1)[1]          # '0001'


def match_drone_to_crops(
    drone_entries: list[dict],
    crops_meta: list[dict],
) -> tuple[dict, dict]:
    """Match each drone image to its positive and semi-positive crops.

    positive: crop with minimum GPS distance to drone center.
    semi-positive: all crops within ±1 grid step from the positive crop.

    Returns:
        positive_map: {frame_id: [crop_name]}
        semi_positive_map: {frame_id: [crop_name, ...]}
    """
    # Build arrays for vectorized distance computation
    crop_lats = np.array([c["lat"] for c in crops_meta])
    crop_lons = np.array([c["lon"] for c in crops_meta])

    # Grid lookup: (x, y) -> crop_name for O(1) neighbor search
    grid = {}
    for crop in crops_meta:
        grid[(crop["x"], crop["y"])] = crop["name"]

    positive_map = {}
    semi_positive_map = {}

    for drone in drone_entries:
        d_lat, d_lon = drone["lat"], drone["lon"]
        fid = _frame_id(drone["filename"])

        # Compute distances to all crops (vectorized haversine)
        dlat = np.radians(crop_lats - d_lat)
        dlon = np.radians(crop_lons - d_lon)
        a = np.sin(dlat / 2) ** 2 + np.cos(math.radians(d_lat)) * np.cos(np.radians(crop_lats)) * np.sin(dlon / 2) ** 2
        dists = 6_371_000 * 2 * np.arcsin(np.sqrt(a))

        # Positive: closest crop
        best_idx = int(np.argmin(dists))
        best_crop = crops_meta[best_idx]
        positive_map[fid] = [best_crop["name"]]

        # Semi-positives: ±1 neighbors in grid via dict lookup
        bx, by = best_crop["x"], best_crop["y"]
        semi = []
        for dx in range(-1, 2):
            for dy in range(-1, 2):
                if dx == 0 and dy == 0:
                    continue
                neighbor = grid.get((bx + dx, by + dy))
                if neighbor is not None:
                    semi.append(neighbor)

        semi_positive_map[fid] = semi

    return positive_map, semi_positive_map


# ---------------------------------------------------------------------------
# 6. Write metadata files
# ---------------------------------------------------------------------------

def write_positive_json(dst: Path, route: str, positive_map: dict):
    out_path = dst / route / "positive.json"
    with open(out_path, "w") as f:
        json.dump(positive_map, f, indent=2)


def write_semi_positive_json(dst: Path, route: str, semi_positive_map: dict):
    out_path = dst / route / "semi_positive.json"
    with open(out_path, "w") as f:
        json.dump(semi_positive_map, f, indent=2)


def write_db_position(dst: Path, route: str, crops_meta: list[dict],
                      sat_w: int, sat_h: int, bbox: dict):
    """Write db_postion.txt with GPS coordinates and scale for each crop.

    Format matches UAV-GeoLoc: name\tlon\tlat\tscale_lon\tscale_lat (tab-separated).
    scale_lon/scale_lat = degrees per pixel in the satellite map.
    Note: filename uses original UAV-GeoLoc spelling 'postion' (sic).
    """
    scale_lon = (bbox["rb_lon"] - bbox["lt_lon"]) / sat_w
    scale_lat = (bbox["rb_lat"] - bbox["lt_lat"]) / sat_h

    out_path = dst / route / "DB" / "db_postion.txt"
    with open(out_path, "w") as f:
        for crop in crops_meta:
            f.write(f"{crop['name']}\t{crop['lon']:.8f}\t{crop['lat']:.8f}"
                    f"\t{scale_lon:.2e}\t{scale_lat:.2e}\n")


# ---------------------------------------------------------------------------
# 7. Generate Index files
# ---------------------------------------------------------------------------

def generate_index_files(
    dst: Path,
    all_routes: list[str],
    route_data: dict,
    train_files: set[str],
    test_files: set[str],
):
    """Generate train/test query and DB index files in UAV-GeoLoc format.

    train_query.txt format:
        route/drone/filename label positive_crop1 positive_crop2 ...
    train_db.txt format:
        route/DB/img/crop_name
    """
    index_dir = dst / "Index"
    index_dir.mkdir(parents=True, exist_ok=True)

    train_query_lines = []
    test_query_lines = []
    all_db_set = set()

    for route in all_routes:
        data = route_data[route]
        positive_map = data["positive_map"]
        semi_positive_map = data["semi_positive_map"]
        crops_meta = data["crops_meta"]

        # All DB crops for this route go into the gallery
        for crop in crops_meta:
            db_path = f"{route}/DB/img/{crop['name']}"
            all_db_set.add(db_path)

        for drone_entry in data["drone_entries"]:
            fname = drone_entry["filename"]
            fid = _frame_id(fname)
            drone_path = f"{route}/drone/{fname}"

            positives = positive_map.get(fid, [])
            semi_pos = semi_positive_map.get(fid, [])
            all_pos = positives + semi_pos

            pos_paths = [f"{route}/DB/img/{p}" for p in all_pos]
            line = f"{drone_path} 0 {' '.join(pos_paths)}"

            if fname in train_files:
                train_query_lines.append(line)
            elif fname in test_files:
                test_query_lines.append(line)

    # DB gallery is the same for train and test (split is by query, not by route)
    sorted_db = sorted(all_db_set)

    # Write files
    _write_lines(index_dir / "train_query.txt", train_query_lines)
    _write_lines(index_dir / "test_query.txt", test_query_lines)
    _write_lines(index_dir / "train_db.txt", sorted_db)
    _write_lines(index_dir / "test_db.txt", sorted_db)
    _write_lines(index_dir / "all_db.txt", sorted_db)

    # Route lists
    _write_lines(index_dir / "train.txt", all_routes)
    _write_lines(index_dir / "test.txt", all_routes)

    return {
        "train_queries": len(train_query_lines),
        "test_queries": len(test_query_lines),
        "total_db": len(all_db_set),
    }


def _write_lines(path: Path, lines: list[str]):
    with open(path, "w") as f:
        for line in lines:
            f.write(line + "\n")


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def process_route(
    src: Path,
    dst: Path,
    route: str,
    bbox: dict,
    crop_size: int,
    stride: int,
    target_size: int,
) -> dict:
    """Process a single route: resize drone, crop satellite, match pairs."""
    print(f"\n{'='*60}")
    print(f"Route {route}")
    print(f"{'='*60}")

    # 1. Read drone metadata
    drone_entries = read_drone_metadata(src, route)
    print(f"  Drone images: {len(drone_entries)}")

    # 2. Resize drone images
    n_drone = resize_drone_images(src, dst, route, target_size)
    print(f"  Drone resized: {n_drone}")

    # 3. Load and crop satellite
    print(f"  Loading satellite map...")
    sat_img = load_satellite(src, route)
    sat_w, sat_h = sat_img.size
    print(f"  Satellite size: {sat_w}x{sat_h}")

    crops_meta = crop_satellite(sat_img, dst, route, crop_size, stride, target_size)
    print(f"  Satellite crops: {len(crops_meta)}")

    # Free memory
    del sat_img

    # 4. Compute GPS for crops
    route_bbox = bbox[route]
    crops_meta = compute_crop_gps(crops_meta, route_bbox, sat_w, sat_h, crop_size)

    # 5. Match drone -> crops
    positive_map, semi_positive_map = match_drone_to_crops(drone_entries, crops_meta)
    print(f"  Positives matched: {len(positive_map)}")

    # Stats
    avg_semi = sum(len(v) for v in semi_positive_map.values()) / max(1, len(semi_positive_map))
    print(f"  Avg semi-positives per drone: {avg_semi:.1f}")

    # 6. Write metadata
    write_positive_json(dst, route, positive_map)
    write_semi_positive_json(dst, route, semi_positive_map)
    write_db_position(dst, route, crops_meta, sat_w, sat_h, route_bbox)
    print(f"  Metadata written.")

    return {
        "drone_entries": drone_entries,
        "crops_meta": crops_meta,
        "positive_map": positive_map,
        "semi_positive_map": semi_positive_map,
    }


def main():
    args = parse_args()
    src = Path(args.src)
    dst = Path(args.dst)

    print(f"Source: {src}")
    print(f"Destination: {dst}")
    print(f"Crop: {args.crop_size}x{args.crop_size}, stride: {args.stride}")
    print(f"Target size: {args.target_size}x{args.target_size}")
    print(f"Excluded routes: {EXCLUDED_ROUTES}")

    dst.mkdir(parents=True, exist_ok=True)

    # Read satellite bounding boxes
    bbox = read_satellite_bbox(src)

    # Determine routes to process
    all_routes = sorted([
        d for d in os.listdir(src)
        if os.path.isdir(src / d) and d.isdigit() and d not in EXCLUDED_ROUTES
    ])
    if args.routes:
        all_routes = [r for r in all_routes if r in args.routes]

    print(f"Routes to process: {all_routes}")

    # Read train/test splits
    train_files = read_split(src, "train")
    test_files = read_split(src, "test")
    print(f"Train split: {len(train_files)} files")
    print(f"Test split: {len(test_files)} files")

    # Process each route
    route_data = {}
    total_drone = 0
    total_crops = 0

    for route in all_routes:
        data = process_route(src, dst, route, bbox, args.crop_size, args.stride, args.target_size)
        route_data[route] = data
        total_drone += len(data["drone_entries"])
        total_crops += len(data["crops_meta"])

    # Generate Index files
    print(f"\n{'='*60}")
    print("Generating Index files...")
    idx_stats = generate_index_files(dst, all_routes, route_data, train_files, test_files)
    print(f"  Train queries: {idx_stats['train_queries']}")
    print(f"  Test queries: {idx_stats['test_queries']}")
    print(f"  DB gallery (all crops): {idx_stats['total_db']}")

    # Summary
    print(f"\n{'='*60}")
    print("SUMMARY")
    print(f"{'='*60}")
    print(f"  Routes processed: {len(all_routes)}")
    print(f"  Total drone images (resized to {args.target_size}): {total_drone}")
    print(f"  Total satellite crops ({args.crop_size}→{args.target_size}): {total_crops}")
    print(f"  Output: {dst}")


if __name__ == "__main__":
    main()