from __future__ import annotations

"""Evaluation for caption quality test.

Recall@K for query(drone+text) -> gallery(satellite).
delta_r_at_1 compares caption-aware vs baseline runs.
"""

import json
import logging
from pathlib import Path

import gin
import torch
from torch.utils.data import DataLoader

from src.models.dual_encoder import DualEncoderCaptionTest

LOGGER = logging.getLogger("caption_test.eval")


def _recall_at_k(
    similarity: torch.Tensor,
    k_values: tuple[int, ...] = (1, 5, 10),
) -> dict[int, float]:
    """Recall@K assuming positives on the diagonal."""
    n_query = similarity.size(0)
    targets = torch.arange(n_query, device=similarity.device)
    sorted_idx = similarity.argsort(dim=1, descending=True)
    result: dict[int, float] = {}
    for k in k_values:
        top_k = sorted_idx[:, :k]
        hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
        result[k] = float(hit.mean().item())
    return result


@torch.no_grad()
def _encode_dataset(
    model: DualEncoderCaptionTest,
    loader: DataLoader,
    device: str,
) -> dict[str, torch.Tensor]:
    """Encode all samples into query and gallery embeddings."""
    model.eval()
    all_query: list[torch.Tensor] = []
    all_gallery: list[torch.Tensor] = []

    for batch in loader:
        drone_img = batch["drone_img"].to(device, non_blocking=True)
        sat_img = batch["sat_img"].to(device, non_blocking=True)
        caption_drone = batch["caption_drone"]

        embeddings = model(
            drone_img=drone_img,
            sat_img=sat_img,
            caption_drone=caption_drone,
        )
        all_query.append(embeddings["query"].cpu())
        all_gallery.append(embeddings["gallery"].cpu())

    return {
        "query": torch.cat(all_query, dim=0),
        "gallery": torch.cat(all_gallery, dim=0),
    }


def evaluate_retrieval(
    model: DualEncoderCaptionTest,
    loader: DataLoader,
    device: str,
    k_values: tuple[int, ...] = (1, 5, 10),
) -> dict[str, float]:
    """Compute R@K for query->gallery and gallery->query.

    Returns:
        Flat dict: r@1_query_to_gallery, r@5_query_to_gallery, etc.
    """
    feats = _encode_dataset(model=model, loader=loader, device=device)

    metrics: dict[str, float] = {}

    sim_q2g = feats["query"] @ feats["gallery"].t()

    for k, val in _recall_at_k(sim_q2g, k_values).items():
        metrics[f"r@{k}_query_to_gallery"] = val
    for k, val in _recall_at_k(sim_q2g.t(), k_values).items():
        metrics[f"r@{k}_gallery_to_query"] = val

    # Gate value for diagnostics.
    metrics["gate"] = model.fusion.gate_value

    return metrics


def delta_r_at_1(
    full_metrics: dict[str, float],
    baseline_metrics: dict[str, float],
) -> float:
    """R@1 gain from adding captions: full - baseline."""
    key = "r@1_query_to_gallery"
    return full_metrics[key] - baseline_metrics[key]


@gin.configurable
def run_evaluation_from_checkpoint(
    checkpoint_path: str,
    test_query_file: str,
    data_root: str,
    output_path: str = "eval_report.json",
    batch_size: int = 128,
    num_workers: int = 4,
    device: str = "cuda",
) -> dict[str, float]:
    """Standalone evaluation from checkpoint."""
    from src.datasets.visloc_with_captions import (
        GeoLocCaptionDataset,
        collate_caption_batch,
    )

    model = DualEncoderCaptionTest().to(device)
    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
    model.load_state_dict(ckpt["model_state"])
    model.eval()

    test_ds = GeoLocCaptionDataset(
        query_file=test_query_file,
        data_root=data_root,
        image_transform=model.preprocess,
    )
    test_loader = DataLoader(
        test_ds,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        collate_fn=collate_caption_batch,
        pin_memory=True,
    )

    metrics = evaluate_retrieval(model=model, loader=test_loader, device=device)

    report = {
        "checkpoint": checkpoint_path,
        "test_query_file": test_query_file,
        "metrics": metrics,
    }
    out = Path(output_path)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)

    LOGGER.info("evaluation report saved to %s", out)
    return metrics
