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

Supports:
  - Training with (query, positive_db, negative_db) triplets for metric learning
  - Evaluation with separate query and DB sets for retrieval

Index file format (train_query.txt / val_query.txt / test_query.txt):
  <query_path> <scene_label> <positive_db_1> [positive_db_2 ...]

Index file format (train_db.txt / val_db.txt / test_db.txt):
  <db_image_path>
"""

import os
import re
import random
from typing import Optional

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


def _parse_query_line(line: str):
    """Parse a query index line that may contain spaces in file paths.

    Format: <query_path> <label_int> <pos_db_1> [pos_db_2 ...]
    Paths can contain spaces (e.g. 'height150_rot180 (1)').
    DB paths always match */DB/img/crop_*.png.
    """
    line = line.strip()
    if not line:
        return None

    # Find all DB positive paths (they always contain /DB/img/crop_)
    db_pattern = re.compile(r'\S*DB/img/crop_\S+')
    db_matches = list(db_pattern.finditer(line))
    if not db_matches:
        return None

    # Everything before the first DB match is "query_path label"
    before_db = line[:db_matches[0].start()].rstrip()
    # Label is the last whitespace-separated token before DB paths
    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


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

    def __init__(
        self,
        root: str,
        query_file: str = "Index/train_query.txt",
        db_file: str = "Index/train_db.txt",
        img_size: int = 512,
        transform_query: Optional[T.Compose] = None,
        transform_db: Optional[T.Compose] = None,
        mining: str = "random",  # "random" or "hard" (hard requires external update)
    ):
        self.root = root
        self.mining = mining

        # Parse query file: each line = query_path label pos_db_1 [pos_db_2 ...]
        self.queries = []  # list of (query_path, label, [positive_db_paths])
        with open(os.path.join(root, query_file)) as f:
            for line in f:
                parsed = _parse_query_line(line)
                if parsed is None:
                    continue
                self.queries.append(parsed)

        # Parse DB file and build label -> db_paths mapping for negative mining
        self.db_paths = []
        self.label_to_db = {}
        with open(os.path.join(root, db_file)) as f:
            for line in f:
                path = line.strip()
                if path:
                    self.db_paths.append(path)

        # Build label -> set of positive db paths for efficient negative sampling
        self._label_positives = {}
        for _, label, positives in self.queries:
            if label not in self._label_positives:
                self._label_positives[label] = set()
            self._label_positives[label].update(positives)

        # Default transforms
        self.transform_query = transform_query or 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]),
        ])
        self.transform_db = transform_db or 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 __len__(self):
        return len(self.queries)

    def _load_image(self, rel_path: str) -> Image.Image:
        return Image.open(os.path.join(self.root, rel_path)).convert("RGB")

    def __getitem__(self, idx):
        query_path, label, positives = self.queries[idx]

        # Load query
        query_img = self.transform_query(self._load_image(query_path))

        # Load random positive
        pos_path = random.choice(positives)
        pos_img = self.transform_db(self._load_image(pos_path))

        # Mine a negative (random DB image not in this query's positive set)
        pos_set = self._label_positives.get(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_image(neg_path))

        return {
            "query": query_img,
            "positive": pos_img,
            "negative": neg_img,
            "label": label,
        }


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

    Use mode="query" for UAV query images, mode="db" for satellite DB images.
    """

    def __init__(
        self,
        root: str,
        index_file: str,  # e.g. "Index/val_query.txt" or "Index/val_db.txt"
        mode: str = "query",  # "query" or "db"
        img_size: int = 512,
        transform: Optional[T.Compose] = None,
    ):
        self.root = root
        self.mode = mode

        self.images = []    # list of image paths
        self.labels = []    # scene labels (query only)
        self.positives = [] # positive db paths per query (query only)

        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:  # query
                    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)

        self.transform = transform or 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]),
        ])

        # Build path-to-index mapping for DB (used in retrieval evaluation)
        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]
        return out


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 = 512,
        transform_query: Optional[T.Compose] = None,
        transform_db: Optional[T.Compose] = 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
                self.queries.append(parsed)

        self.transform_query = transform_query or T.Compose([
            T.Resize((img_size, img_size)),
            T.RandomHorizontalFlip(),
            T.RandomRotation(15),
            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]),
        ])
        self.transform_db = transform_db or 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 __len__(self):
        return len(self.queries)

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


# ── Collate function for eval (handles variable-length positives list) ──────

def eval_collate_fn(batch):
    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]
    return out


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

def build_dataloaders(
    root: str,
    split: str = "terrain",  # "terrain", "country", or "all"
    batch_size: int = 32,
    img_size: int = 512,
    num_workers: int = 4,
    mode: str = "triplet",  # "triplet" or "pair"
):
    """Build train/val/test dataloaders.

    Args:
        root: Path to UAV-GeoLoc dataset root.
        split: Which subset - "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]

    # Train loader
    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,
    )

    # Val/Test loaders (separate query and db for retrieval evaluation)
    loaders = {"train": train_loader}

    for phase in ["val", "test"]:
        q_suffix = suffix if os.path.exists(
            os.path.join(root, f"Index/{phase}_query{suffix}.txt")
        ) else ""
        d_suffix = suffix if os.path.exists(
            os.path.join(root, f"Index/{phase}_db{suffix}.txt")
        ) else ""

        q_ds = UAVGeoLocEval(
            root, f"Index/{phase}_query{q_suffix}.txt", mode="query", img_size=img_size
        )
        d_ds = UAVGeoLocEval(
            root, f"Index/{phase}_db{d_suffix}.txt", 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


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

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

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

    print(f"Train: {len(loaders['train'].dataset)} samples")
    print(f"Val query: {len(loaders['val_query'].dataset)} samples")
    print(f"Val DB: {len(loaders['val_db'].dataset)} samples")
    print(f"Test query: {len(loaders['test_query'].dataset)} samples")
    print(f"Test DB: {len(loaders['test_db'].dataset)} samples")

    # Smoke test one batch
    batch = next(iter(loaders["train"]))
    print(f"\nTrain batch shapes:")
    print(f"  query:    {batch['query'].shape}")
    print(f"  positive: {batch['positive'].shape}")
    print(f"  negative: {batch['negative'].shape}")
    print(f"  labels:   {batch['label']}")

    batch = next(iter(loaders["val_query"]))
    print(f"\nVal query batch:")
    print(f"  image: {batch['image'].shape}")
    print(f"  labels: {batch['label']}")
    print(f"  positives[0]: {batch['positives'][0]}")