"""Измерение процента обрезаемых описаний по комбинациям уровней.

Обрезка (truncation до 248 токенов) зависит ТОЛЬКО от текста и токенизатора,
не от обучения. Поэтому процент обрезанных сэмплов можно посчитать по готовым
данным — результат идентичен тому, что было во время обучения.

Считается ОТДЕЛЬНО по каждому split (cross-area train/test) и по каждой стороне
(drone queries / satellite gallery), чтобы процент отражал ровно те сэмплы, что
реально идут в обучение и в eval, а не весь датасет целиком (ср. check_coverage.py).

Для каждой из 6 комбинаций уровней считает:
    - сколько сэмплов превышает 248 токенов (обрезается)
    - средняя/максимальная длина в токенах
    - средний «перебор» у обрезанных (на сколько токенов текст длиннее 248)

Запуск:
    python measure_truncation.py \
        --data_root "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR" \
        --descriptions_path "/media/servml/SSD_2_2TB/datasets/cvgl_datasets/GTA-UAV-LR-captions_ v1" \
        --version v1
"""
from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

sys.path.insert(0, ".")

from src.data.gta_uav import load_text_descriptions, combine_text_levels
from src.data.gta_uav_eval import collect_gallery_names
from src.models.dgtrs.model import tokenize_dgtrs

CONTEXT_LENGTH = 248

# те же 6 комбинаций, что в экспериментах
COMBINATIONS = [
    ["level1"],
    ["level2"],
    ["level3"],
    ["level1", "level2"],
    ["level1", "level3"],
    ["level1", "level2", "level3"],
]


def count_tokens(text: str) -> int:
    """Реальное число ненулевых токенов ДО обрезки.

    Токенизируем с большим запасом (context_length=1024), чтобы увидеть полную
    длину без потери токенов на лимите 248.
    """
    if not text.strip():
        return 0
    toks = tokenize_dgtrs(text, context_length=1024, truncate=True)
    return int((toks != 0).sum())


def measure(names: list[str], descriptions: dict, combo: list[str]) -> dict:
    """Посчитать статистику обрезки для одной комбинации уровней.

    Учитываются только имена из `names` (сэмплы конкретного split/стороны),
    у которых есть непустой склеенный текст.
    """
    n_total = 0
    n_truncated = 0
    lengths = []
    overflows = []

    for name in names:
        desc = descriptions.get(name)
        if not desc:
            continue
        text = combine_text_levels(desc, combo)
        if not text.strip():
            continue
        n = count_tokens(text)
        n_total += 1
        lengths.append(n)
        if n > CONTEXT_LENGTH:
            n_truncated += 1
            overflows.append(n - CONTEXT_LENGTH)

    pct = 100.0 * n_truncated / n_total if n_total else 0.0
    avg_len = sum(lengths) / len(lengths) if lengths else 0
    max_len = max(lengths) if lengths else 0
    avg_overflow = sum(overflows) / len(overflows) if overflows else 0

    return {
        "combo": " + ".join(combo),
        "n_total": n_total,
        "n_truncated": n_truncated,
        "pct_truncated": pct,
        "avg_len": avg_len,
        "max_len": max_len,
        "avg_overflow": avg_overflow,
    }


def _report_side(title: str, names: list[str], descriptions: dict) -> None:
    """Напечатать таблицу обрезки по всем комбинациям для одной стороны."""
    print(f"\n── {title} (всего уникальных: {len(names)}) " + "─" * 20)
    print(f"{'Комбинация':<26} {'Всего':>7} {'Обрезано':>9} {'%':>7} "
          f"{'СрДлина':>8} {'МаксДлина':>10} {'СрПеребор':>10}")
    print("-" * 82)
    for combo in COMBINATIONS:
        r = measure(names, descriptions, combo)
        print(f"{r['combo']:<26} {r['n_total']:>7} {r['n_truncated']:>9} "
              f"{r['pct_truncated']:>6.1f}% {r['avg_len']:>8.0f} "
              f"{r['max_len']:>10} {r['avg_overflow']:>10.0f}")
    print("-" * 82)


def _report_split(
    data_root: Path,
    meta_name: str,
    desc_by_view: dict[str, dict],
) -> None:
    """Обрезка v для одного 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)):
        _report_side(
            f"{view.upper()} {'queries' if view == 'drone' else 'gallery'}",
            names, desc_by_view[view],
        )


def main():
    args = parse_args()
    data_root = Path(args.data_root)

    print("=" * 64)
    print(f"TRUNCATION CHECK  |  version={args.version}  |  лимит={CONTEXT_LENGTH} токенов")
    print("=" * 64)

    # Описания не зависят от split — грузим один раз на каждую сторону.
    desc_by_view = {
        view: load_text_descriptions(args.descriptions_path, view_type=view)
        for view in ("drone", "satellite")
    }
    for view, desc in desc_by_view.items():
        print(f"Загружено {len(desc)} описаний ({view}, {args.version})")

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

    print("\nОбрезано = число сэмплов, где склеенный текст длиннее лимита "
          "(хвост, включая level3-якорь, отсекается).")


def parse_args():
    p = argparse.ArgumentParser(description="Измерение обрезки по комбинациям уровней и split")
    p.add_argument("--data_root", type=str, required=True)
    p.add_argument("--descriptions_path", type=str, required=True)
    p.add_argument("--version", type=str, default="v1")
    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")
    return p.parse_args()


if __name__ == "__main__":
    main()