from __future__ import annotations

"""GTA-UAV-LR dataset loader with L1/L2/L3 hierarchical captions for CVGL.

Reads cross-area pair JSONs and VLM-generated caption JSONs.
Produces (drone_img, sat_img, caption_l1, caption_l2, caption_l3) tuples.

Caption levels:
    L1 (overview):    First sentence of P1 (land cover summary).
    L2 (full):        Complete P1 + P2 text.
    L3 (fingerprint): P3 unique signature section.

For short captions (pure water, no P markers): all levels get the same short text.
"""

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

import coloredlogs
import torch
from PIL import Image
from torch.utils.data import Dataset
from tqdm import tqdm

LOGGER = logging.getLogger("caption_test.gtauav_dataset")
coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")

# Default paths.
_RGB_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR")
_CAPTION_ROOT = Path("/home/servml/Документы/datasets/GTA-UAV-LR-captions")
_EMPTY_CAPTION = ""

# Regex to split P1/P2/P3 sections.
_P_SPLIT = re.compile(r"\*\*P[123][^*]*\*\*\s*:?\s*")


def _parse_caption_levels(output: str) -> tuple[str, str, str]:
    """Split VLM caption output into L1, L2, L3 levels.

    Returns:
        (l1_overview, l2_full, l3_fingerprint)
    """
    sections = _P_SPLIT.split(output)
    # sections[0] is empty (before **P1**), sections[1]=P1, [2]=P2, [3]=P3
    sections = [s.strip() for s in sections if s.strip()]

    if len(sections) >= 3:
        p1, p2, p3 = sections[0], sections[1], sections[2]
    elif len(sections) == 2:
        p1, p2, p3 = sections[0], sections[1], sections[0]
    elif len(sections) == 1:
        p1 = p2 = p3 = sections[0]
    else:
        p1 = p2 = p3 = output.strip()

    # L1: first sentence of P1.
    first_dot = p1.find(". ")
    l1 = p1[:first_dot + 1] if first_dot > 0 else p1

    # L2: full P1 + P2.
    l2 = p1 + " " + p2

    # L3: P3 (fingerprint / unique signature).
    l3 = p3

    return l1, l2, l3


def _load_caption_index(caption_root: Path) -> dict[str, dict]:
    """Build index: image_name -> caption JSON data.

    Scans drone/images/ and satellite/ directories.
    """
    index: dict[str, dict] = {}

    for subdir in ["drone/images", "satellite"]:
        cap_dir = caption_root / subdir
        if not cap_dir.exists():
            continue
        cap_files = sorted(cap_dir.glob("*_caption.json"))
        for cap_file in tqdm(cap_files, desc=f"  📄 Loading {subdir} captions", unit="file", leave=False):
            with open(cap_file) as f:
                data = json.load(f)
            # Key by the image name (without _caption suffix).
            img_name = cap_file.name.replace("_caption.json", ".png")
            index[img_name] = data

    return index


class GTAUAVDataset(Dataset):
    """GTA-UAV-LR dataset with hierarchical L1/L2/L3 captions.

    Args:
        pair_json: Path to cross-area-drone2sate-{train,test}.json.
        rgb_root: Root of GTA-UAV-LR RGB images.
        caption_root: Root of GTA-UAV-LR-captions.
        drone_transform: Transform for drone images (can include augmentations).
        sat_transform: Transform for satellite images (can include augmentations).
        image_transform: Fallback single transform for both (used if drone/sat not set).
        filter_meta: Path to seg_filter.json (exclude 90%+ bg/water).
        drop_caption_prob: Probability of dropping captions (ablation).
        seed: Random seed.
    """

    def __init__(
        self,
        pair_json: str,
        rgb_root: str = str(_RGB_ROOT),
        caption_root: str = str(_CAPTION_ROOT),
        drone_transform: Callable[[Image.Image], torch.Tensor] | None = None,
        sat_transform: Callable[[Image.Image], torch.Tensor] | None = None,
        image_transform: Callable[[Image.Image], torch.Tensor] | None = None,
        filter_meta: str | None = None,
        drop_caption_prob: float = 0.0,
        seed: int = 0,
    ) -> None:
        self.rgb_root = Path(rgb_root)
        self.caption_root = Path(caption_root)
        self.drone_transform = drone_transform or image_transform
        self.sat_transform = sat_transform or image_transform
        self.drop_caption_prob = drop_caption_prob
        self._rng = random.Random(seed)

        # Load exclusion set from segmentation filter.
        self.excluded: set[str] = set()
        if filter_meta is not None:
            self._load_filter(Path(filter_meta))

        # Load caption index.
        LOGGER.info("📚 Loading caption index from %s", caption_root)
        self.caption_index = _load_caption_index(self.caption_root)
        LOGGER.info("📚 Caption index: %d entries", len(self.caption_index))

        # Load pairs.
        self.entries: list[dict[str, Any]] = []
        self._load_pairs(Path(pair_json))
        LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json)

    def _load_filter(self, path: Path) -> None:
        with open(path) as f:
            meta = json.load(f)
        # excluded list contains paths like "drone/images/xxx.png"
        for exc in meta.get("excluded", []):
            # Extract image name from segm path.
            self.excluded.add(Path(exc).name)
        LOGGER.info("🔻 Filter loaded: %d excluded images", len(self.excluded))

    def _load_pairs(self, pair_json: Path) -> None:
        with open(pair_json) as f:
            raw_pairs = json.load(f)

        for pair in raw_pairs:
            drone_name = pair["drone_img_name"]

            # Skip excluded images.
            if drone_name in self.excluded:
                continue

            # Get positive/semi-positive satellite images.
            pos_list = pair.get("pair_pos_sate_img_list", [])
            semipos_list = pair.get("pair_pos_semipos_sate_img_list", [])
            semipos_weights = pair.get("pair_pos_semipos_sate_weight_list", [])

            # Use positives if available, else semi-positives.
            if pos_list:
                sat_candidates = pos_list
                sat_weights = None
            elif semipos_list:
                sat_candidates = semipos_list
                sat_weights = semipos_weights if semipos_weights else None
            else:
                continue  # No match, skip.

            # Get drone captions.
            cap_data = self.caption_index.get(drone_name)
            if cap_data is not None:
                l1, l2, l3 = _parse_caption_levels(cap_data["output"])
            else:
                l1 = l2 = l3 = _EMPTY_CAPTION

            # Pre-parse satellite captions for all candidates.
            sat_captions: dict[str, tuple[str, str, str]] = {}
            for sat_name in sat_candidates:
                sat_cap = self.caption_index.get(sat_name)
                if sat_cap is not None:
                    sat_captions[sat_name] = _parse_caption_levels(sat_cap["output"])

            self.entries.append({
                "drone_name": drone_name,
                "drone_dir": pair["drone_img_dir"],
                "sat_dir": pair["sate_img_dir"],
                "sat_candidates": sat_candidates,
                "sat_weights": sat_weights,
                "caption_l1": l1,
                "caption_l2": l2,
                "caption_l3": l3,
                "sat_captions": sat_captions,
            })

    def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor:
        path = self.rgb_root / directory / filename
        with Image.open(path) as img:
            rgb = img.convert("RGB")
        if transform is not None:
            return transform(rgb)
        return torch.tensor(0)  # placeholder if no transform

    def get_all_valid_sat_names(self) -> list[list[str]]:
        """Return all valid satellite matches per drone query (for evaluation).

        In GTA-UAV, each drone has multiple valid satellite crops (partial IoU).
        Standard diagonal R@K is wrong — must check if ANY valid match is in top-K.
        """
        return [entry["sat_candidates"] for entry in self.entries]

    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["drone_dir"], entry["drone_name"], self.drone_transform)

        # Sample satellite match (weighted if semi-positive).
        if entry["sat_weights"] is not None:
            sat_idx = self._rng.choices(
                range(len(entry["sat_candidates"])),
                weights=entry["sat_weights"],
                k=1,
            )[0]
            sat_name = entry["sat_candidates"][sat_idx]
            pos_weight = entry["sat_weights"][sat_idx]
        else:
            sat_idx = self._rng.randrange(len(entry["sat_candidates"]))
            sat_name = entry["sat_candidates"][sat_idx]
            pos_weight = 1.0  # strict positive

        sat_img = self._load_image(entry["sat_dir"], sat_name, self.sat_transform)

        # Drone captions with optional dropout.
        if self.drop_caption_prob > 0 and self._rng.random() < self.drop_caption_prob:
            l1 = l2 = l3 = _EMPTY_CAPTION
        else:
            l1, l2, l3 = entry["caption_l1"], entry["caption_l2"], entry["caption_l3"]

        # Satellite captions (empty string if not available).
        sat_caps = entry["sat_captions"].get(sat_name)
        if sat_caps is not None:
            sat_l1, sat_l2, sat_l3 = sat_caps
        else:
            sat_l1 = sat_l2 = sat_l3 = _EMPTY_CAPTION

        return {
            "drone_img": drone_img,
            "sat_img": sat_img,
            "caption_l1": l1,
            "caption_l2": l2,
            "caption_l3": l3,
            "sat_caption_l1": sat_l1,
            "sat_caption_l2": sat_l2,
            "sat_caption_l3": sat_l3,
            "pair_id": entry["drone_name"],
            "sat_name": sat_name,
            "positive_weight": pos_weight,
        }


def collate_gtauav_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_l1": [b["caption_l1"] for b in batch],
        "caption_l2": [b["caption_l2"] for b in batch],
        "caption_l3": [b["caption_l3"] for b in batch],
        "sat_caption_l1": [b["sat_caption_l1"] for b in batch],
        "sat_caption_l2": [b["sat_caption_l2"] for b in batch],
        "sat_caption_l3": [b["sat_caption_l3"] for b in batch],
        "pair_ids": [b["pair_id"] for b in batch],
        "sat_names": [b["sat_name"] for b in batch],
        "positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32),
    }