"""DataLoader для GTA-UAV с текстовыми описаниями для ОБЕИХ сторон.

И дрон, и спутник имеют собственное текстовое описание (сгенерированное
одной и той же VLM по своей картинке). Модель сливает (картинка+текст)
СИММЕТРИЧНО на обеих сторонах через два отдельных GatedFusion, и только
затем сравнивает два слитых вектора.

Задача: symmetric fusion drone↔satellite retrieval (NGCG), где обе стороны
представлены парой (картинка, текст), слитой в единый вектор.
"""
from __future__ import annotations

import json
import logging
import random
import re
from pathlib import Path

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

from src.models.dgtrs.model import tokenize_dgtrs

LOGGER = logging.getLogger("cvgl.data.gta_uav")


# ---------------------------------------------------------------------------
# Загрузка и очистка текстовых описаний
# ---------------------------------------------------------------------------

# Зонная метка v1: позиционное слово (upper/lower/center...) + опциональное
# горизонтальное уточнение + двоеточие. Матчится ГДЕ УГОДНО в строке (в т.ч.
# в середине после `;`), регистронезависимо. Опциональный ведущий `;`/`,`
# съедается вместе с меткой, но пробельные границы — только [ \t], чтобы НЕ
# разрушать разрыв абзацев `\n\n` (иначе level'ы слипаются).
_ZONE_LABEL_RE = re.compile(
    r'[;,]?[ \t]*'
    r'\b(?:upper|lower|middle|top|bottom|center|central)'
    r'(?:[-‑ ]?(?:left|right|center|centre|middle))?'
    r'[ \t]*:[ \t]*',
    re.IGNORECASE,
)


def _clean_caption_text(raw_output: str) -> list[str]:
    """Очистить output от markdown-разметки (v1) и разбить на абзацы.

    V1 формат содержит markdown-заголовки (**P1 — ...**) и зонные метки
    (upper-left:, upper-right:, lower-center: ...), причём метки встречаются
    и в начале строки, и в середине после `;`. V2 формат — чистая проза.
    Функция обрабатывает оба и применяется одинаково к drone- и satellite-
    описаниям.
    """
    text = re.sub(r'\*\*P\d[^*]*\*\*\s*', '', raw_output)
    text = re.sub(r'\*\*([^*]*)\*\*', r'\1', text)
    text = _ZONE_LABEL_RE.sub(' ', text)
    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
    cleaned = []
    for p in paragraphs:
        merged = " ".join(line.strip() for line in p.split("\n") if line.strip())
        merged = re.sub(r'\s{2,}', ' ', merged).strip()
        if merged:
            cleaned.append(merged)
    return cleaned


def load_text_descriptions(
    descriptions_path: str | Path,
    view_type: str = "drone",
) -> dict[str, dict[str, str]]:
    """Загрузить текстовые описания сцен заданного типа (drone или satellite).

    Формат: *_caption.json файлы рекурсивно в descriptions_path, каждый
    содержит поле "view_type" ("drone" или "satellite") и поле "output"
    с 1-3 абзацами, разделёнными \\n\\n, которые маппятся на level1/2/3.

    Вызывается ДВАЖДЫ в Dataset.__init__: один раз с view_type="drone",
    один раз с view_type="satellite" — обе стороны фьюзятся одинаковым
    образом, поэтому обеим нужны свои описания.

    Args:
        descriptions_path: Путь к директории с *_caption.json.
        view_type: "drone" или "satellite" — какие описания оставить.

    Returns:
        dict: {"image_name.png": {"level1": "...", "level2": "...", "level3": "..."}}
    """
    path = Path(descriptions_path)
    caption_files = sorted(path.glob("**/*_caption.json"))

    if not caption_files:
        raise FileNotFoundError(
            f"No *_caption.json files found in {path}. "
            f"Expected structure: {path}/drone/images/*_caption.json "
            f"and {path}/satellite/*_caption.json"
        )

    data = {}
    n_single = 0
    n_skipped = 0

    for caption_file in caption_files:
        with open(caption_file, "r", encoding="utf-8") as f:
            entry = json.load(f)

        if entry.get("view_type") != view_type:
            n_skipped += 1
            continue

        if "rel_path" in entry:
            img_name = Path(entry["rel_path"]).name
        else:
            img_name = caption_file.name.replace("_caption.json", ".png")

        output_text = entry.get("output", "")
        paragraphs = _clean_caption_text(output_text)

        levels = {
            "level1": paragraphs[0] if len(paragraphs) >= 1 else "",
            "level2": paragraphs[1] if len(paragraphs) >= 2 else "",
            "level3": paragraphs[2] if len(paragraphs) >= 3 else "",
        }

        if len(paragraphs) < 2:
            n_single += 1

        data[img_name] = levels

    LOGGER.info(
        "📝 Loaded %d %s descriptions from %s (%d single-paragraph, %d other-view skipped)",
        len(data), view_type, path.name, n_single, n_skipped,
    )

    return data


def combine_text_levels(
    descriptions: dict[str, str],
    levels: list[str],
    separator: str = " ",
) -> str:
    """Объединить указанные уровни описаний в один текст.

    Используется одинаково для drone- и satellite-описаний — выбранные
    text_levels (level1/level2/level3) применяются СИММЕТРИЧНО к обеим
    сторонам в рамках одного эксперимента.
    """
    parts = []
    for level in levels:
        text = descriptions.get(level, "")
        if text:
            parts.append(text.strip())

    combined = separator.join(parts)

    if not combined:
        LOGGER.debug("Empty text after combining levels %s", levels)
        combined = "No description available."

    return combined


# ---------------------------------------------------------------------------
# Image transforms (общие для drone и satellite изображений)
# ---------------------------------------------------------------------------

def get_image_transforms(
    image_size: int = 256,
    is_train: bool = True,
) -> transforms.Compose:
    """Трансформации для изображений (используются и для drone, и для satellite).

    Train: augmentations (flip, color jitter, resize+crop).
    Test: только resize + normalize.
    """
    normalize = transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225],
    )

    if is_train:
        return transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.RandomVerticalFlip(p=0.3),
            transforms.ColorJitter(
                brightness=0.2, contrast=0.2, saturation=0.1, hue=0.05,
            ),
            transforms.ToTensor(),
            normalize,
        ])
    else:
        return transforms.Compose([
            transforms.Resize((image_size, image_size)),
            transforms.ToTensor(),
            normalize,
        ])


# ---------------------------------------------------------------------------
# Dataset
# ---------------------------------------------------------------------------

class GTAUAVTextDataset(Dataset):
    """GTA-UAV dataset для symmetric drone(image+text) ↔ satellite(image+text).

    Каждый сэмпл: (drone_image, drone_tokens, satellite_image,
    satellite_tokens, weight, metadata).

    Обе стороны несут пару (картинка, текст), которая сливается моделью
    (GatedFusion) в единый вектор перед сравнением.

    Args:
        data_root:          Корневая папка GTA-UAV.
        pairs_meta_file:    JSON-файл с парами.
        descriptions_path:  Путь к текстовым описаниям (drone И satellite).
        text_levels:        Какие уровни текста использовать — одинаково
                             применяются к обеим сторонам.
        image_size:         Размер изображения (drone и satellite).
        is_train:           Train/test режим.
        use_semipos:        Использовать semi-positive пары (с весами).
        context_length:     Максимальная длина токенов.
    """

    def __init__(
        self,
        data_root: str | Path,
        pairs_meta_file: str,
        descriptions_path: str | Path,
        text_levels: list[str] = ("level1",),
        image_size: int = 256,
        is_train: bool = True,
        use_semipos: bool = True,
        context_length: int = 248,
    ) -> None:
        super().__init__()
        self.data_root = Path(data_root)
        self.text_levels = list(text_levels)
        self.is_train = is_train
        self.use_semipos = use_semipos
        self.context_length = context_length
        self.image_size = image_size

        meta_path = self.data_root / pairs_meta_file
        with open(meta_path, "r") as f:
            self.pairs_meta = json.load(f)
        LOGGER.info(
            "📦 GTA-UAV %s: %d entries from %s",
            "train" if is_train else "test",
            len(self.pairs_meta),
            pairs_meta_file,
        )

        # Описания загружаются ДВАЖДЫ — для дрона и для спутника отдельно,
        # т.к. это разные картинки с разными описаниями.
        self.drone_descriptions = load_text_descriptions(descriptions_path, view_type="drone")
        self.satellite_descriptions = load_text_descriptions(descriptions_path, view_type="satellite")

        covered_drone = sum(
            1 for entry in self.pairs_meta
            if entry["drone_img_name"] in self.drone_descriptions
        )
        LOGGER.info(
            "📝 Drone text coverage: %d/%d entries (%.1f%%)",
            covered_drone, len(self.pairs_meta),
            100.0 * covered_drone / max(len(self.pairs_meta), 1),
        )
        LOGGER.info(
            "📝 Satellite descriptions available: %d images",
            len(self.satellite_descriptions),
        )

        self.drone_transform = get_image_transforms(image_size, is_train)
        self.satellite_transform = get_image_transforms(image_size, is_train)

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

    def _load_drone_image(self, entry: dict) -> torch.Tensor:
        """Загрузить картинку дрона."""
        drone_img_dir = entry.get("drone_img_dir", "drone/images")
        drone_img_name = entry["drone_img_name"]
        drone_path = self.data_root / drone_img_dir / drone_img_name

        try:
            image = Image.open(drone_path).convert("RGB")
            return self.drone_transform(image)
        except (FileNotFoundError, OSError) as e:
            LOGGER.debug("Failed to load drone image %s: %s", drone_path, e)
            return torch.zeros(3, self.image_size, self.image_size)

    def _load_satellite_image(self, entry: dict) -> tuple[torch.Tensor, float, str]:
        """Загрузить спутниковый снимок (цель поиска, галерея).

        Returns:
            (image_tensor, weight, sate_img_name) — имя нужно для того,
            чтобы найти соответствующее текстовое описание спутника.
        """
        key = "pair_pos_semipos_sate_img_list" if self.use_semipos else "pair_pos_sate_img_list"
        weight_key = "pair_pos_semipos_sate_weight_list" if self.use_semipos else "pair_pos_sate_weight_list"

        sate_list = entry.get(key, entry.get("pair_pos_sate_img_list", []))
        weight_list = entry.get(weight_key, entry.get("pair_pos_sate_weight_list", []))

        if not sate_list:
            LOGGER.debug("No satellite matches for %s, using placeholder", entry.get("drone_img_name"))
            return torch.zeros(3, self.image_size, self.image_size), 0.0, ""

        if self.is_train:
            if weight_list and sum(weight_list) > 0:
                probs = torch.tensor(weight_list, dtype=torch.float)
                probs = probs / probs.sum()
                chosen_idx = torch.multinomial(probs, 1).item()
            else:
                chosen_idx = random.randint(0, len(sate_list) - 1)
        else:
            chosen_idx = 0

        sate_img_name = sate_list[chosen_idx]
        weight = weight_list[chosen_idx] if weight_list else 1.0
        sate_img_dir = entry.get("sate_img_dir", "satellite")
        sate_path = self.data_root / sate_img_dir / sate_img_name

        try:
            image = Image.open(sate_path).convert("RGB")
            image = self.satellite_transform(image)
        except (FileNotFoundError, OSError) as e:
            LOGGER.debug("Failed to load %s: %s", sate_path, e)
            image = torch.zeros(3, self.image_size, self.image_size)
            weight = 0.0

        return image, weight, sate_img_name

    def __getitem__(self, idx: int) -> dict:
        entry = self.pairs_meta[idx]
        drone_img_name = entry["drone_img_name"]

        # --- Дрон: картинка + текст ---
        drone_image = self._load_drone_image(entry)
        drone_desc = self.drone_descriptions.get(drone_img_name, {})
        drone_text = combine_text_levels(drone_desc, self.text_levels)
        drone_tokens = tokenize_dgtrs(drone_text, context_length=self.context_length, truncate=True)
        drone_tokens = drone_tokens.squeeze(0)  # [context_length]

        # --- Спутник: картинка + текст (НОВОЕ: текст для спутника тоже) ---
        satellite_image, weight, sate_img_name = self._load_satellite_image(entry)
        sat_desc = self.satellite_descriptions.get(sate_img_name, {})
        sat_text = combine_text_levels(sat_desc, self.text_levels)
        satellite_tokens = tokenize_dgtrs(sat_text, context_length=self.context_length, truncate=True)
        satellite_tokens = satellite_tokens.squeeze(0)  # [context_length]

        # В meta координаты дрона называются drone_loc_x_y (игровые X/Y), а не lat/lon.
        drone_loc = entry.get("drone_loc_x_y", [0.0, 0.0])

        return {
            "drone_image": drone_image,                # [3, H, W]
            "drone_tokens": drone_tokens,                # [248]
            "satellite_image": satellite_image,          # [3, H, W]
            "satellite_tokens": satellite_tokens,         # [248] — НОВОЕ ПОЛЕ
            "weight": torch.tensor(weight, dtype=torch.float32),
            "drone_img_name": drone_img_name,
            "drone_x": drone_loc[0],
            "drone_y": drone_loc[1],
        }


# ---------------------------------------------------------------------------
# Collate и DataLoader factory
# ---------------------------------------------------------------------------

def collate_fn(batch: list[dict]) -> dict:
    """Custom collate — теперь включает satellite_tokens."""
    return {
        "drone_image": torch.stack([b["drone_image"] for b in batch]),
        "drone_tokens": torch.stack([b["drone_tokens"] for b in batch]),
        "satellite_image": torch.stack([b["satellite_image"] for b in batch]),
        "satellite_tokens": torch.stack([b["satellite_tokens"] for b in batch]),
        "weight": torch.stack([b["weight"] for b in batch]),
        "drone_img_name": [b["drone_img_name"] for b in batch],
        "drone_x": torch.tensor([b["drone_x"] for b in batch]),
        "drone_y": torch.tensor([b["drone_y"] for b in batch]),
    }


def build_dataloaders(
    data_root: str | Path,
    descriptions_path: str | Path,
    text_levels: list[str],
    train_meta: str = "cross-area-drone2sate-train.json",
    test_meta: str = "cross-area-drone2sate-test.json",
    batch_size: int = 64,
    num_workers: int = 4,
    image_size: int = 256,
    build_test: bool = True,
    mutually_exclusive: bool = True,
    seed: int = 42,
) -> tuple[DataLoader, DataLoader | None]:
    """Создать train и (опционально) диагональный test DataLoader.

    build_test=False — не строить test-набор (оценка идёт через
    multi-positive eval из gta_uav_eval, который сам грузит галерею и
    описания; иначе описания читаются дважды).

    mutually_exclusive=True — формировать train-батчи так, чтобы в одном
    батче не было двух дронов с общим позитивным тайлом (протокол §6.4);
    убирает false-negatives на диагонали InfoNCE. seed синхронизирует
    порядок с train seed. False → обычный random shuffle.
    """
    train_dataset = GTAUAVTextDataset(
        data_root=data_root,
        pairs_meta_file=train_meta,
        descriptions_path=descriptions_path,
        text_levels=text_levels,
        image_size=image_size,
        is_train=True,
        use_semipos=True,
    )

    if mutually_exclusive:
        from src.data.mutually_exclusive_sampler import MutuallyExclusiveBatchSampler

        # Позитивные тайлы (pos + semi-pos) каждого дрона → множества int-id.
        tile_to_id: dict[str, int] = {}
        positive_sets: list[set[int]] = []
        for entry in train_dataset.pairs_meta:
            names = (
                entry.get("pair_pos_semipos_sate_img_list")
                or entry.get("pair_pos_sate_img_list")
                or []
            )
            ids = set()
            for name in names:
                ids.add(tile_to_id.setdefault(name, len(tile_to_id)))
            positive_sets.append(ids)

        batch_sampler = MutuallyExclusiveBatchSampler(
            positive_sets, batch_size=batch_size, seed=seed, drop_last=True,
        )
        train_loader = DataLoader(
            train_dataset,
            batch_sampler=batch_sampler,
            num_workers=num_workers,
            collate_fn=collate_fn,
            pin_memory=True,
        )
        LOGGER.info(
            "🧩 Mutually-exclusive sampling ON (seed=%d): нет общих позитивных "
            "тайлов внутри батча (§6.4)", seed,
        )
    else:
        train_loader = DataLoader(
            train_dataset,
            batch_size=batch_size,
            shuffle=True,
            num_workers=num_workers,
            collate_fn=collate_fn,
            pin_memory=True,
            drop_last=True,
        )
        LOGGER.info("🔀 Random shuffle sampling (mutually_exclusive=False)")

    if not build_test:
        LOGGER.info(
            "📦 DataLoader ready: train=%d batches (bs=%d, levels=%s); test пропущен",
            len(train_loader), batch_size, text_levels,
        )
        return train_loader, None

    test_dataset = GTAUAVTextDataset(
        data_root=data_root,
        pairs_meta_file=test_meta,
        descriptions_path=descriptions_path,
        text_levels=text_levels,
        image_size=image_size,
        is_train=False,
        use_semipos=False,
    )

    test_loader = DataLoader(
        test_dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collate_fn,
        pin_memory=True,
        drop_last=False,
    )

    LOGGER.info(
        "📦 DataLoaders ready: train=%d batches, test=%d batches (bs=%d, levels=%s)",
        len(train_loader), len(test_loader), batch_size, text_levels,
    )

    return train_loader, test_loader
