"""Multi-positive eval для GTA-UAV cross-view retrieval (протокол §6.2).

В отличие от обучающего загрузчика (одна спутниковая пара на дрон,
диагональная оценка), здесь строится КОРРЕКТНАЯ retrieval-постановка:

    * галерея = УНИКАЛЬНЫЕ спутниковые тайлы теста (без дублей);
    * для каждого дрона-запроса — список ВСЕХ валидных positive-тайлов
      (positive + semi-positive, IoU ≥ 0.14) как индексы в галерее;
    * метрика hit-if-any: запрос засчитан, если ХОТЯ БЫ ОДИН из его
      positive-тайлов попал в top-K.

Галерея собирается как объединение всех pos/semi-pos тайлов по всем
дронам теста — это все достижимые reference-тайлы данной области.

ВАЖНО (задокументировано): отдельного официального списка галереи
для этого набора НЕТ (в папке набора — только split описаний
drone/satellite и pipeline.log). Поэтому объединение pos/semi-pos —
канонический и воспроизводимый источник галереи здесь. Единственное
следствие: тайлы тест-области, не являющиеся positive НИ для одного
дрона (чистые дистракторы), в галерею не попадают → абсолютный recall
слегка оптимистичен. Смещение ОДИНАКОВО для всех вариантов, поэтому
ΔR@1 между вариантами (основа decision rule) остаётся честным.
"""
from __future__ import annotations

import json
import logging
from pathlib import Path

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

from src.data.gta_uav import (
    combine_text_levels,
    get_image_transforms,
    load_text_descriptions,
)
from src.models.dgtrs.model import tokenize_dgtrs

LOGGER = logging.getLogger("cvgl.eval")

# Поля meta с positive/semi-positive списками (pos + semi-pos, IoU ≥ 0.14).
_SEMIPOS_KEY = "pair_pos_semipos_sate_img_list"
_POS_KEY = "pair_pos_sate_img_list"


def _entry_positive_names(entry: dict) -> list[str]:
    """Список positive+semi-positive спутниковых тайлов для одного дрона."""
    names = entry.get(_SEMIPOS_KEY) or entry.get(_POS_KEY) or []
    # Уникализируем, сохраняя порядок (у одного дрона тайл не должен дублиться).
    seen: dict[str, None] = {}
    for name in names:
        seen.setdefault(name, None)
    return list(seen.keys())


def collect_gallery_names(pairs_meta: list[dict]) -> tuple[list[str], dict[str, str]]:
    """Собрать уникальную галерею тайлов из meta теста.

    Returns:
        gallery_names: отсортированный список уникальных имён тайлов.
        name_to_dir:   имя тайла → каталог (sate_img_dir) для загрузки.
    """
    name_to_dir: dict[str, str] = {}
    for entry in pairs_meta:
        sate_dir = entry.get("sate_img_dir", "satellite")
        for name in _entry_positive_names(entry):
            name_to_dir.setdefault(name, sate_dir)
    gallery_names = sorted(name_to_dir.keys())
    return gallery_names, name_to_dir


class _SatelliteGalleryDataset(Dataset):
    """Уникальные спутниковые тайлы (картинка + текст) — галерея."""

    def __init__(
        self,
        data_root: Path,
        gallery_names: list[str],
        name_to_dir: dict[str, str],
        satellite_descriptions: dict[str, dict],
        text_levels: list[str],
        image_size: int,
        context_length: int,
    ) -> None:
        self.data_root = data_root
        self.gallery_names = gallery_names
        self.name_to_dir = name_to_dir
        self.satellite_descriptions = satellite_descriptions
        self.text_levels = text_levels
        self.image_size = image_size
        self.context_length = context_length
        self.transform = get_image_transforms(image_size, is_train=False)

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

    def __getitem__(self, idx: int) -> dict:
        name = self.gallery_names[idx]
        path = self.data_root / self.name_to_dir[name] / name
        try:
            image = self.transform(Image.open(path).convert("RGB"))
        except (FileNotFoundError, OSError) as e:
            LOGGER.debug("Failed to load gallery tile %s: %s", path, e)
            image = torch.zeros(3, self.image_size, self.image_size)

        desc = self.satellite_descriptions.get(name, {})
        text = combine_text_levels(desc, self.text_levels)
        tokens = tokenize_dgtrs(text, context_length=self.context_length, truncate=True).squeeze(0)
        return {"image": image, "tokens": tokens, "name": name}


class _DroneQueryDataset(Dataset):
    """Дрон-запросы (картинка + текст)."""

    def __init__(
        self,
        data_root: Path,
        entries: list[dict],
        drone_descriptions: dict[str, dict],
        text_levels: list[str],
        image_size: int,
        context_length: int,
    ) -> None:
        self.data_root = data_root
        self.entries = entries
        self.drone_descriptions = drone_descriptions
        self.text_levels = text_levels
        self.image_size = image_size
        self.context_length = context_length
        self.transform = get_image_transforms(image_size, is_train=False)

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

    def __getitem__(self, idx: int) -> dict:
        entry = self.entries[idx]
        name = entry["drone_img_name"]
        img_dir = entry.get("drone_img_dir", "drone/images")
        path = self.data_root / img_dir / name
        try:
            image = self.transform(Image.open(path).convert("RGB"))
        except (FileNotFoundError, OSError) as e:
            LOGGER.debug("Failed to load drone query %s: %s", path, e)
            image = torch.zeros(3, self.image_size, self.image_size)

        desc = self.drone_descriptions.get(name, {})
        text = combine_text_levels(desc, self.text_levels)
        tokens = tokenize_dgtrs(text, context_length=self.context_length, truncate=True).squeeze(0)
        return {"image": image, "tokens": tokens, "name": name}


def _collate(batch: list[dict]) -> dict:
    return {
        "image": torch.stack([b["image"] for b in batch]),
        "tokens": torch.stack([b["tokens"] for b in batch]),
        "name": [b["name"] for b in batch],
    }


def invert_positives(
    positives_q2g: list[list[int]],
    n_gallery: int,
) -> list[list[int]]:
    """Инвертировать positive-карту: для каждого тайла — список дрон-запросов."""
    positives_g2q: list[list[int]] = [[] for _ in range(n_gallery)]
    for query_idx, gallery_indices in enumerate(positives_q2g):
        for g in gallery_indices:
            positives_g2q[g].append(query_idx)
    return positives_g2q


def build_multipos_eval(
    data_root: str | Path,
    test_meta: str,
    descriptions_path: str | Path,
    text_levels: list[str],
    image_size: int = 256,
    batch_size: int = 64,
    num_workers: int = 4,
    context_length: int = 248,
) -> dict:
    """Собрать multi-positive eval: query-loader, gallery-loader, positive-карты.

    Дроны без единого валидного positive-тайла в галерее исключаются из
    запросов (recall для них не определён), с логированием количества.

    Returns:
        dict с ключами:
            drone_loader:    DataLoader дрон-запросов.
            gallery_loader:  DataLoader уникальной спутниковой галереи.
            positives_q2g:   list[list[int]] — дрон → индексы тайлов галереи.
            positives_g2q:   list[list[int]] — тайл → индексы дрон-запросов.
            gallery_names:   list[str] — имена тайлов галереи (порядок = индекс).
            n_queries / n_gallery / n_skipped: счётчики.
    """
    data_root = Path(data_root)
    meta_path = data_root / test_meta
    with open(meta_path, "r") as f:
        pairs_meta = json.load(f)

    drone_descriptions = load_text_descriptions(descriptions_path, view_type="drone")
    satellite_descriptions = load_text_descriptions(descriptions_path, view_type="satellite")

    # 1. Уникальная галерея тайлов + индекс.
    gallery_names, name_to_dir = collect_gallery_names(pairs_meta)
    name_to_idx = {name: i for i, name in enumerate(gallery_names)}

    # 2. Запросы + positive-карта (только дроны с ≥1 positive в галерее).
    kept_entries: list[dict] = []
    positives_q2g: list[list[int]] = []
    n_skipped = 0
    for entry in pairs_meta:
        pos_idx = [name_to_idx[n] for n in _entry_positive_names(entry) if n in name_to_idx]
        if not pos_idx:
            n_skipped += 1
            continue
        kept_entries.append(entry)
        positives_q2g.append(pos_idx)

    positives_g2q = invert_positives(positives_q2g, len(gallery_names))

    LOGGER.info(
        "🔎 Multi-positive eval: %d queries (%d skipped, no positive), "
        "%d unique gallery tiles",
        len(kept_entries), n_skipped, len(gallery_names),
    )

    gallery_ds = _SatelliteGalleryDataset(
        data_root, gallery_names, name_to_dir, satellite_descriptions,
        text_levels, image_size, context_length,
    )
    drone_ds = _DroneQueryDataset(
        data_root, kept_entries, drone_descriptions,
        text_levels, image_size, context_length,
    )

    drone_loader = DataLoader(
        drone_ds, batch_size=batch_size, shuffle=False,
        num_workers=num_workers, collate_fn=_collate,
    )
    gallery_loader = DataLoader(
        gallery_ds, batch_size=batch_size, shuffle=False,
        num_workers=num_workers, collate_fn=_collate,
    )

    return {
        "drone_loader": drone_loader,
        "gallery_loader": gallery_loader,
        "positives_q2g": positives_q2g,
        "positives_g2q": positives_g2q,
        "gallery_names": gallery_names,
        "n_queries": len(kept_entries),
        "n_gallery": len(gallery_names),
        "n_skipped": n_skipped,
    }
