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

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

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

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

import argparse
import sys

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

from src.data.gta_uav import load_text_descriptions, combine_text_levels
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:
    """Реальное число ненулевых токенов ДО обрезки.

    tokenize с truncate=False даёт полную длину; если функция не поддерживает
    truncate=False, считаем через увеличенный context_length.
    """
    if not text.strip():
        return 0
    # токенизируем с большим запасом, чтобы увидеть полную длину без обрезки
    toks = tokenize_dgtrs(text, context_length=1024, truncate=True)
    return int((toks != 0).sum())


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

    for img_name, desc in descriptions.items():
        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 main():
    args = parse_args()

    descriptions = load_text_descriptions(args.descriptions_path, view_type="drone")
    print(f"Загружено {len(descriptions)} описаний ({args.version})\n")

    print(f"{'Комбинация':<26} {'Всего':>7} {'Обрезано':>9} {'%':>7} "
          f"{'СрДлина':>8} {'МаксДлина':>10} {'СрПеребор':>10}")
    print("-" * 82)

    rows = []
    for combo in COMBINATIONS:
        r = measure(descriptions, combo)
        rows.append(r)
        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)
    print(f"\nЛимит контекста: {CONTEXT_LENGTH} токенов")
    print("Обрезано = число сэмплов, где склеенный текст длиннее лимита "
          "(хвост, включая level3-якорь, отсекается).")


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


if __name__ == "__main__":
    main()