from __future__ import annotations

"""UAV-GeoLoc dataset loader with template captions for CVGL caption test.

Reads UAV-GeoLoc Index format (train_query.txt / train_db.txt) and generates
template captions from path metadata (terrain type, scene, altitude, heading).

train_query.txt format:
    Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg 0 .../DB/img/crop_X_Y.png ...

Template caption (P3-style fingerprint for cross-view matching):
    "Aerial view at 100m facing northwest over volcanic terrain near KilaueaVolcano.
     Plan-view features: dark lava flows, crater edges, sparse vegetation patches."
"""

import random
import re
from pathlib import Path
from typing import Any, Callable

import gin
import torch
from PIL import Image
from torch.utils.data import Dataset

# Compass lookup: heading_deg -> direction name.
_COMPASS = ["north", "northeast", "east", "southeast",
            "south", "southwest", "west", "northwest"]

# Simple terrain-to-features mapping for template captions.
_TERRAIN_FEATURES: dict[str, str] = {
    "Volcano": "lava flows, crater edges, volcanic rock",
    "Mountain": "ridgelines, steep slopes, rocky terrain",
    "Hill": "rolling hills, gentle slopes, scattered trees",
    "Desert": "sand dunes, arid ground, sparse scrub",
    "Plain": "flat open fields, agricultural plots, dirt roads",
    "Plateau": "flat elevated terrain, cliffs, mesa edges",
    "Basin": "lowland depression, dry lake bed, sediment patterns",
    "Delta": "river channels, sediment fans, wetland patches",
    "Gorge": "deep canyon, exposed rock layers, narrow valley",
    "Island": "shoreline, coastal vegetation, water boundary",
    "Wetland": "marshes, water channels, aquatic vegetation",
    "Glacier": "ice formations, crevasses, glacial moraine",
    "Forest": "dense canopy, tree shadows, clearings",
    "Farm": "crop fields, irrigation lines, farm buildings",
    "Prairie": "grassland, open meadow, fence lines",
    "Finca": "rural estate, orchards, scattered structures",
    "Calcification": "mineral deposits, white crusts, barren soil",
    "StoneForest": "karst pillars, eroded limestone, sparse vegetation",
    "Oasis": "palm clusters, water pool, surrounding desert",
    "Flowers": "colorful ground cover, floral patterns, open field",
    "Terrace": "terraced hillside, stepped fields, retaining walls",
    "Snow": "snow cover, tracks, exposed rock patches",
    "Pasture": "grazing land, fenced paddocks, grass",
    "Danxia": "red sandstone, layered cliffs, erosion patterns",
    "Hylare": "sparse woodland, rocky outcrops",
    "Karst": "sinkholes, limestone towers, caves",
    "Fall": "autumn foliage, colored canopy, leaf litter",
}

# Country/city features.
_COUNTRY_FEATURES: str = "buildings, roads, urban blocks, rooftops, intersections"


def _parse_query_path(query_path: str) -> dict[str, str]:
    """Extract metadata from UAV-GeoLoc query path.

    Example: Terrain/Volcano/KilaueaVolcano/query/height100_rot315/footage/file.jpeg
    Returns: {category, terrain_type, scene, height_m, heading_deg, heading_dir}
    """
    parts = query_path.split("/")

    category = parts[0] if parts else "Unknown"

    if category == "Terrain":
        terrain_type = parts[1] if len(parts) > 1 else "Unknown"
        scene = parts[2] if len(parts) > 2 else "Unknown"
    elif category == "Country":
        terrain_type = "Urban"
        scene = "/".join(parts[1:3]) if len(parts) > 2 else "Unknown"
    else:
        terrain_type = "Unknown"
        scene = parts[1] if len(parts) > 1 else "Unknown"

    # Parse height and rotation from trajectory folder name.
    height_m = "100"
    heading_deg = "0"
    for part in parts:
        m = re.match(r"height(\d+)_rot(\d+)", part)
        if m:
            height_m = m.group(1)
            heading_deg = m.group(2)
            break

    heading_idx = round(int(heading_deg) / 45) % 8
    heading_dir = _COMPASS[heading_idx]

    return {
        "category": category,
        "terrain_type": terrain_type,
        "scene": scene,
        "height_m": height_m,
        "heading_deg": heading_deg,
        "heading_dir": heading_dir,
    }


def _make_template_caption(meta: dict[str, str]) -> str:
    """Generate a template caption from parsed metadata."""
    terrain = meta["terrain_type"]
    features = _TERRAIN_FEATURES.get(terrain, _COUNTRY_FEATURES)

    return (
        f"Aerial view at {meta['height_m']}m facing {meta['heading_dir']} "
        f"over {terrain.lower()} terrain near {meta['scene']}. "
        f"Plan-view features: {features}."
    )


@gin.configurable
class GeoLocCaptionDataset(Dataset):
    """UAV-GeoLoc pairs with template captions.

    Reads train_query.txt, randomly samples one positive crop per query.

    Args:
        query_file: Path to train_query.txt (or test_query.txt).
        data_root: Root directory of UAV-GeoLoc dataset.
        image_transform: Callable applied to PIL images.
        drop_caption_prob: Probability of dropping caption (for ablation).
        seed: Random seed.
    """

    def __init__(
        self,
        query_file: str,
        data_root: str,
        image_transform: Callable[[Image.Image], torch.Tensor],
        drop_caption_prob: float = 0.0,
        seed: int = 0,
    ) -> None:
        self.data_root = Path(data_root)
        self.image_transform = image_transform
        self.drop_caption_prob = drop_caption_prob
        self._rng = random.Random(seed)

        self.entries: list[dict[str, Any]] = []
        self._load_query_file(Path(query_file))

    def _load_query_file(self, query_file: Path) -> None:
        """Parse train_query.txt into list of entries."""
        with open(query_file) as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                parts = line.split()
                query_path = parts[0]
                # parts[1] is label (always 0), parts[2:] are positive crop paths.
                positive_crops = parts[2:]
                if not positive_crops:
                    continue

                meta = _parse_query_path(query_path)
                caption = _make_template_caption(meta)

                self.entries.append({
                    "query_path": query_path,
                    "positive_crops": positive_crops,
                    "caption": caption,
                    "meta": meta,
                })

    def _load_image(self, relative_path: str) -> torch.Tensor:
        path = self.data_root / relative_path
        with Image.open(path) as img:
            rgb = img.convert("RGB")
        return self.image_transform(rgb)

    def __len__(self) -> int:
        return len(self.entries)

    def __getitem__(self, idx: int) -> dict[str, Any]:
        entry = self.entries[idx]

        drone_img = self._load_image(entry["query_path"])

        # Randomly sample one positive crop.
        crop_path = self._rng.choice(entry["positive_crops"])
        sat_img = self._load_image(crop_path)

        caption = entry["caption"]
        if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
            caption = ""

        return {
            "drone_img": drone_img,
            "sat_img": sat_img,
            "caption_drone": caption,
            "pair_id": entry["query_path"],
        }


def collate_caption_batch(
    batch: list[dict[str, Any]],
) -> dict[str, Any]:
    """Collate into batched dict. Captions stay as string lists."""
    return {
        "drone_img": torch.stack([b["drone_img"] for b in batch], dim=0),
        "sat_img": torch.stack([b["sat_img"] for b in batch], dim=0),
        "caption_drone": [b["caption_drone"] for b in batch],
        "pair_ids": [b["pair_id"] for b in batch],
    }
