"""Метрики для оценки retrieval качества.

После обучения мы прогоняем все тексты и изображения через энкодеры,
строим матрицу сходства и для каждого запроса ранжируем кандидатов.

Recall@K — основная метрика:
    «Какая доля запросов нашла правильный ответ в топ-K?»
    Recall@1 = 0.65 значит: в 65% случаев правильное изображение — первое в списке.

AP (Average Precision):
    Усреднённая точность по всем позициям ранжирования.
    Учитывает не только «попал ли в топ-K», но и на какой именно позиции.

Meter-level distance (специфика GTA-UAV):
    Каждое изображение имеет GPS-координаты. Мы можем посчитать
    расстояние в метрах между предсказанной и реальной позицией.
"""
from __future__ import annotations

import numpy as np
import torch


@torch.no_grad()
def compute_retrieval_metrics(
    text_embeddings: torch.Tensor,
    image_embeddings: torch.Tensor,
    ks: list[int] = (1, 5, 10),
    compute_ap: bool = True,
) -> dict[str, float]:
    """Вычислить retrieval метрики.
    
    Предполагается что text_embeddings[i] соответствует image_embeddings[i]
    (i-й текст описывает i-е изображение).
    
    Args:
        text_embeddings:  [N, D] нормализованные текстовые эмбеддинги.
        image_embeddings: [N, D] нормализованные визуальные эмбеддинги.
        ks: Значения K для Recall@K.
        compute_ap: Считать ли Average Precision (медленнее).
    
    Returns:
        dict с метриками: recall@1, recall@5, recall@10, AP, ...
    """
    N = text_embeddings.shape[0]
    
    # Матрица сходства: [N, N]. sim[i][j] = косинусное сходство текста i и изображения j.
    sim = text_embeddings @ image_embeddings.T  # уже L2-нормализованы → cosine sim
    
    # Для каждого текста ранжируем все изображения по убыванию сходства.
    # ranks[i] = на какой позиции (0-indexed) стоит правильное изображение i.
    # Правильное изображение для текста i — это image i (диагональ).
    sorted_indices = sim.argsort(dim=1, descending=True)  # [N, N]
    
    # Найти позицию правильного ответа (i-е изображение для i-го текста)
    ranks = torch.zeros(N, dtype=torch.long, device=sim.device)
    for i in range(N):
        rank = (sorted_indices[i] == i).nonzero(as_tuple=True)[0].item()
        ranks[i] = rank
    
    results = {}
    
    # Recall@K: доля запросов, где правильный ответ в топ-K
    for k in ks:
        recall = (ranks < k).float().mean().item()
        results[f"recall@{k}"] = recall

    # Recall@1%: правильный ответ в топ-1% базы
    k_1pct = max(1, N // 100)
    results["recall@1%"] = (ranks < k_1pct).float().mean().item()

    # Mean Rank и Median Rank (для анализа — чем меньше, тем лучше)
    results["mean_rank"] = ranks.float().mean().item()
    results["median_rank"] = ranks.float().median().item()
    
    # Average Precision
    if compute_ap:
        # AP = среднее по всем запросам от 1/(rank+1)
        # Это упрощённый AP для случая одного правильного ответа на запрос.
        ap = (1.0 / (ranks.float() + 1)).mean().item()
        results["AP"] = ap
    
    return results


@torch.no_grad()
def compute_bidirectional_metrics(
    text_embeddings: torch.Tensor,
    image_embeddings: torch.Tensor,
    ks: list[int] = (1, 5, 10),
) -> dict[str, float]:
    """Метрики в обе стороны: text→image и image→text.
    
    Returns:
        dict с префиксами t2i_ и i2t_ для каждого направления.
    """
    t2i = compute_retrieval_metrics(text_embeddings, image_embeddings, ks=ks)
    # Для image→text: переворачиваем — ищем текст по изображению
    i2t = compute_retrieval_metrics(image_embeddings, text_embeddings, ks=ks)

    results = {}
    for key, val in t2i.items():
        results[f"t2i_{key}"] = val
    for key, val in i2t.items():
        results[f"i2t_{key}"] = val
    
    # Средний Recall@K по обоим направлениям
    for k in ks:
        results[f"avg_recall@{k}"] = (t2i[f"recall@{k}"] + i2t[f"recall@{k}"]) / 2
    
    return results


def compute_meter_distance(
    predicted_indices: np.ndarray,
    query_coords: np.ndarray,
    gallery_coords: np.ndarray,
) -> dict[str, float]:
    """Расстояние в метрах между предсказанной и реальной позицией.
    
    Специфика GTA-UAV: каждое изображение имеет GPS-координаты.
    Мы берём координаты top-1 предсказания и считаем расстояние до GT.
    
    Args:
        predicted_indices: [N] индексы top-1 предсказанных изображений.
        query_coords:      [N, 2] GPS-координаты запросов (lat, lon).
        gallery_coords:    [M, 2] GPS-координаты базы изображений.
    
    Returns:
        dict: mean_distance_m, median_distance_m, recall_at_Xm.
    """
    predicted_coords = gallery_coords[predicted_indices]  # [N, 2]
    
    # Haversine distance (приблизительно, для малых расстояний ≈ Euclidean × scale)
    # Для GTA-UAV координаты в игровых юнитах, не GPS → используем Euclidean
    diffs = predicted_coords - query_coords  # [N, 2]
    distances = np.linalg.norm(diffs, axis=1)  # [N] в метрах (или юнитах)
    
    results = {
        "mean_distance": float(distances.mean()),
        "median_distance": float(np.median(distances)),
    }
    
    # Recall at distance thresholds
    for threshold in (5, 10, 25, 50, 100):
        results[f"recall@{threshold}m"] = float((distances < threshold).mean())
    
    return results


def format_metrics(metrics: dict[str, float], prefix: str = "") -> str:
    """Форматировать метрики для логирования."""
    lines = []
    for key, val in metrics.items():
        if "recall" in key or "AP" in key or "acc" in key:
            lines.append(f"{prefix}{key}: {val:.4f} ({val*100:.1f}%)")
        elif "rank" in key:
            lines.append(f"{prefix}{key}: {val:.1f}")
        elif "distance" in key:
            lines.append(f"{prefix}{key}: {val:.1f}m")
        else:
            lines.append(f"{prefix}{key}: {val:.4f}")
    return "\n".join(lines)
