"""Сравнение ПОКРЫТИЯ captions двух наборов (v1 vs v2) — отдельно от обучения.

Отвечает на один вопрос: одинаковый ли набор картинок покрыт текстом в v1 и v2?
Если да — сравнение R@1 между версиями чистое. Если нет — разница R@1 может
объясняться покрытием, а не качеством текста (см. протокол §5.4, §0.4).

Использует ТЕ ЖЕ функции чтения, что и обучение (load_text_descriptions,
combine_text_levels), поэтому «есть текст» здесь означает ровно то же, что
видит модель: непустой текст для выбранных text_levels, а не заглушку
"No description available.".

Запуск:
    python3 check_coverage.py \
    --data_root "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR" \
    --captions_v1 "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1" \
    --captions_v2 "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_v2" \
    --text_levels level1 level2 level3

"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

from src.data.gta_uav import combine_text_levels, load_text_descriptions
from src.data.gta_uav_eval import collect_gallery_names


def _has_real_text(
    name: str,
    descriptions: dict[str, dict],
    text_levels: list[str],
) -> bool:
    """True, если для картинки есть непустой текст на выбранных уровнях.

    Повторяет семантику обучения: пустой результат combine_text_levels →
    модель получит заглушку, т.е. текста фактически нет.
    """
    desc = descriptions.get(name)
    if not desc:
        return False
    # combine_text_levels возвращает заглушку, когда все уровни пустые →
    # это и есть «текста фактически нет».
    combined = combine_text_levels(desc, text_levels)
    return combined != "No description available."


def _covered_set(
    names: list[str],
    descriptions: dict[str, dict],
    text_levels: list[str],
) -> set[str]:
    """Подмножество имён, у которых есть реальный текст."""
    return {n for n in names if _has_real_text(n, descriptions, text_levels)}


def _report_side(
    title: str,
    names: list[str],
    desc_v1: dict[str, dict],
    desc_v2: dict[str, dict],
    text_levels: list[str],
) -> None:
    """Напечатать покрытие v1/v2 и пересечение для одной стороны."""
    total = len(names)
    cov_v1 = _covered_set(names, desc_v1, text_levels)
    cov_v2 = _covered_set(names, desc_v2, text_levels)
    both = cov_v1 & cov_v2
    only_v1 = cov_v1 - cov_v2
    only_v2 = cov_v2 - cov_v1

    def pct(n: int) -> str:
        return f"{n:>6d} / {total} ({100.0 * n / max(total, 1):5.1f}%)"

    print(f"\n── {title} (всего уникальных: {total}) " + "─" * 20)
    print(f"  покрытие v1:            {pct(len(cov_v1))}")
    print(f"  покрытие v2:            {pct(len(cov_v2))}")
    print(f"  есть в ОБОИХ (пересеч): {pct(len(both))}")
    print(f"  только v1:              {len(only_v1)}")
    print(f"  только v2:              {len(only_v2)}")

    if len(cov_v1) == len(cov_v2) == len(both):
        print("  ✅ наборы совпадают → покрытие не влияет на сравнение")
    else:
        diff = abs(len(cov_v1) - len(cov_v2))
        print(
            f"  ⚠️  расхождение: |v1−v2| = {diff} картинок; "
            f"для честного сравнения оценивайте на пересечении ({len(both)})"
        )


def _report_split(
    data_root: Path,
    meta_name: str,
    desc_by_view: dict[str, tuple[dict, dict]],
    text_levels: list[str],
) -> None:
    """Покрытие v1/v2 для одного split-файла (train или test)."""
    meta_path = data_root / meta_name
    if not meta_path.exists():
        print(f"\n⚠️  Пропущен (нет файла): {meta_path}")
        return

    with open(meta_path) as f:
        pairs_meta = json.load(f)

    # Дрон-запросы: уникальные имена картинок дрона в split.
    drone_names = sorted({e["drone_img_name"] for e in pairs_meta})
    # Спутниковая галерея: те же уникальные тайлы, что и в eval.
    gallery_names, _ = collect_gallery_names(pairs_meta)

    print("\n" + "=" * 64)
    print(f"SPLIT: {meta_name}  ({len(pairs_meta)} пар)")
    print("=" * 64)

    for view, names in (("drone", drone_names), ("satellite", gallery_names)):
        desc_v1, desc_v2 = desc_by_view[view]
        _report_side(
            f"{view.upper()} {'queries' if view == 'drone' else 'gallery'}",
            names, desc_v1, desc_v2, text_levels,
        )


def main(args: argparse.Namespace) -> None:
    data_root = Path(args.data_root)
    text_levels = list(args.text_levels)

    print("=" * 64)
    print(f"COVERAGE CHECK  |  text_levels={text_levels}")
    print("=" * 64)

    # Описания не зависят от split — грузим один раз на каждую сторону.
    desc_by_view = {
        view: (
            load_text_descriptions(args.captions_v1, view_type=view),
            load_text_descriptions(args.captions_v2, view_type=view),
        )
        for view in ("drone", "satellite")
    }

    # Покрытие важно и на train (что модель УЧИТ), и на test (что видит eval).
    for meta_name in (args.train_meta, args.test_meta):
        _report_split(data_root, meta_name, desc_by_view, text_levels)


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Сравнить покрытие captions v1 vs v2")
    p.add_argument("--data_root", type=str, required=True)
    p.add_argument("--train_meta", type=str, default="cross-area-drone2sate-train.json")
    p.add_argument("--test_meta", type=str, default="cross-area-drone2sate-test.json")
    p.add_argument("--captions_v1", type=str, required=True)
    p.add_argument("--captions_v2", type=str, required=True)
    p.add_argument("--text_levels", nargs="+", default=["level1"])
    return p.parse_args()


if __name__ == "__main__":
    main(parse_args())