"""Метрики для оценки качества перекрёстной геолокализации (CVGL).

    R@K (один positive на запрос):
        R@K = (1/N) Σ_i  𝟙[rank(i) ≤ K]
        где rank(i) — позиция (1-indexed) истинного соответствия.

    R@K (несколько positive, конвенция «hit-if-any», Sample4Geo):
        достаточно, чтобы ХОТЯ БЫ ОДИН positive попал в top-K.

    AP для одного запроса — доля позитивных предсказаний среди top-K;
    mAP — среднее AP по всем запросам.

    q2g (query-to-gallery): БПЛА → спутник (основная метрика);
    g2q (gallery-to-query): спутник → БПЛА (диагностика симметричности).
"""
from __future__ import annotations

import numpy as np
import torch


@torch.no_grad()
def compute_retrieval_metrics(
    query_embeddings: torch.Tensor,
    gallery_embeddings: torch.Tensor,
    ks: tuple[int, ...] = (1, 5, 10),
    positives: list[list[int]] | None = None,
    compute_map: bool = True,
    return_ranks: bool = False,
) -> dict[str, float]:
    """Вычислить метрики поиска q2g.

    Args:
        query_embeddings:   [N, D] нормализованные слитые drone-эмбеддинги.
        gallery_embeddings: [M, D] нормализованные слитые satellite-эмбеддинги.
        ks: Значения K для R@K.
        positives: Опционально — список positive-индексов галереи для каждого
            запроса (для конвенции «hit-if-any» при нескольких правильных
            ответах). Если None — предполагается диагональное соответствие
            query[i] ↔ gallery[i] (один positive на запрос).
        compute_map: Считать ли mAP.
        return_ranks: Если True — добавить в результат ключ "_best_ranks":
            np.ndarray [N] с 0-indexed рангом первого попавшего positive по
            каждому запросу (M, если ни один не попал). Из него выводится
            hit@k = (ranks < k) для любого k — нужно для bootstrap CI и
            paired-теста (протокол §8.3 / §9.2).

    Returns:
        dict: recall@1, recall@5, recall@10, recall@1%, mean_rank,
              median_rank, mAP (+ "_best_ranks", если return_ranks=True).
    """
    N = query_embeddings.shape[0]
    M = gallery_embeddings.shape[0]

    # Матрица косинусного сходства (эмбеддинги уже L2-нормализованы)
    sim = query_embeddings @ gallery_embeddings.T  # [N, M]
    sorted_indices = sim.argsort(dim=1, descending=True)  # [N, M]

    # positive-множество для каждого запроса
    if positives is None:
        # Диагональное соответствие: для запроса i правильный ответ — gallery i
        positives = [[i] for i in range(N)]

    # Ранг ПЕРВОГО найденного positive (hit-if-any) для каждого запроса
    best_ranks = torch.full((N,), M, dtype=torch.long)
    # AP по каждому запросу (для mAP)
    aps = torch.zeros(N, dtype=torch.float)

    sorted_np = sorted_indices.cpu().numpy()

    for i in range(N):
        pos_set = set(positives[i])
        if not pos_set:
            continue
        order = sorted_np[i]  # индексы галереи по убыванию сходства

        # Позиции всех positive в ранжированном списке (0-indexed)
        hit_positions = [r for r, idx in enumerate(order) if idx in pos_set]
        if not hit_positions:
            continue

        # Ранг первого попадания (1-indexed) — для R@K по «hit-if-any»
        best_ranks[i] = hit_positions[0]

        # AP запроса: среднее precision@(каждое попадание)
        # precision на k-м попадании = (номер попадания) / (позиция+1)
        if compute_map:
            ap = 0.0
            for hit_num, pos in enumerate(hit_positions, start=1):
                precision_at_hit = hit_num / (pos + 1)
                ap += precision_at_hit
            aps[i] = ap / len(pos_set)

    results: dict[str, float] = {}

    # R@K: доля запросов, где первый positive попал в top-K
    for k in ks:
        results[f"recall@{k}"] = (best_ranks < k).float().mean().item()

    # R@1%
    k_1pct = max(1, M // 100)
    results["recall@1%"] = (best_ranks < k_1pct).float().mean().item()

    # Mean / Median rank (1-indexed для читаемости)
    results["mean_rank"] = (best_ranks.float() + 1).mean().item()
    results["median_rank"] = (best_ranks.float() + 1).median().item()

    # mAP
    if compute_map:
        results["mAP"] = aps.mean().item()

    if return_ranks:
        results["_best_ranks"] = best_ranks.cpu().numpy()

    return results


@torch.no_grad()
def compute_bidirectional_metrics(
    query_embeddings: torch.Tensor,
    gallery_embeddings: torch.Tensor,
    ks: tuple[int, ...] = (1, 5, 10),
    positives_q2g: list[list[int]] | None = None,
    positives_g2q: list[list[int]] | None = None,
    return_ranks: bool = False,
) -> dict[str, float]:
    """Метрики в обе стороны.

    q2g (query-to-gallery): БПЛА → спутник — основная метрика.
    g2q (gallery-to-query): спутник → БПЛА — диагностика симметричности
        признакового пространства.

    Args:
        return_ranks: Если True — добавить ключи "q2g_best_ranks" и
            "g2q_best_ranks" (np.ndarray, 0-indexed ранги первого попадания)
            для последующего bootstrap CI / paired-теста. Числовые метрики
            остаются float, per-query массивы лежат под отдельными ключами.
    """
    q2g = compute_retrieval_metrics(
        query_embeddings, gallery_embeddings, ks=ks, positives=positives_q2g,
        return_ranks=return_ranks,
    )
    g2q = compute_retrieval_metrics(
        gallery_embeddings, query_embeddings, ks=ks, positives=positives_g2q,
        return_ranks=return_ranks,
    )

    q2g_ranks = q2g.pop("_best_ranks", None)
    g2q_ranks = g2q.pop("_best_ranks", None)

    results = {}
    for key, val in q2g.items():
        results[f"q2g_{key}"] = val
    for key, val in g2q.items():
        results[f"g2q_{key}"] = val

    # Средний R@K по обоим направлениям
    for k in ks:
        results[f"avg_recall@{k}"] = (q2g[f"recall@{k}"] + g2q[f"recall@{k}"]) / 2

    if return_ranks:
        results["q2g_best_ranks"] = q2g_ranks
        results["g2q_best_ranks"] = g2q_ranks

    return results


def compute_meter_distance(
    predicted_indices: np.ndarray,
    query_coords: np.ndarray,
    gallery_coords: np.ndarray,
) -> dict[str, float]:
    """Расстояние в метрах между предсказанной и истинной позицией.

    Берёт координаты top-1 предсказанного спутникового тайла и сравнивает
    с истинными координатами запроса.
    """
    predicted_coords = gallery_coords[predicted_indices]
    diffs = predicted_coords - query_coords
    distances = np.linalg.norm(diffs, axis=1)

    results = {
        "mean_distance": float(distances.mean()),
        "median_distance": float(np.median(distances)),
    }
    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 "mAP" 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)
