from typing import Dict, List

import numpy as np
import tensorflow as tf


def compute_retrieval_metrics(logits: tf.Tensor, k: int | List[int] = 5) -> Dict[str, tf.Tensor]:
    """Computes metrics for a retrieval task using in-batch negatives.

    Supports a single k or a list of k values.

    Args:
        logits (tf.Tensor): The result of a dot product (Batch Size, Batch Size).
        k (Union[int, List[int]]): Number of top-K results to evaluate.

    Returns:
        Dict[str, tf.Tensor]: A dictionary containing recall_k, mrr_k, and ndcg_k.
    """
    try:
        k_list = [k] if isinstance(k, int) else list(k)

        batch_size = tf.shape(logits)[0]
        # Ground truth is on the diagonal (in-batch negatives)
        true_indices = tf.expand_dims(tf.range(batch_size), -1)
        sorted_indices = tf.argsort(logits, direction="DESCENDING")

        # Find the rank of the true positive (1-indexed)
        # where() returns (batch_idx, rank_idx), we take [:, 1] to get rank_idx and +1
        ranks = tf.where(tf.equal(sorted_indices, tf.cast(true_indices, tf.int32)))[:, 1] + 1
        ranks = tf.cast(ranks, tf.float32)

        results = {}

        for current_k in k_list:
            # 1. Recall@K (Is the true item's rank <= K?)
            is_in_top_k = ranks <= float(current_k)
            recall_at_k = tf.reduce_mean(tf.cast(is_in_top_k, tf.float32))

            # 2. MRR@K
            # If rank > K, RR is 0. Else RR is 1/rank
            rr = tf.where(is_in_top_k, 1.0 / ranks, 0.0)
            mrr_at_k = tf.reduce_mean(rr)

            # 3. NDCG@K
            # If rank > K, DCG is 0. Else DCG is 1 / log2(rank + 1)
            dcg = tf.where(is_in_top_k, 1.0 / (tf.math.log(ranks + 1.0) / tf.math.log(2.0)), 0.0)
            ndcg_at_k = tf.reduce_mean(dcg)

            results[f"recall_{current_k}"] = recall_at_k
            results[f"mrr_{current_k}"] = mrr_at_k
            results[f"ndcg_{current_k}"] = ndcg_at_k

        return results
    except Exception as e:
        print(f"Error computing retrieval metrics: {e}")
        results = {}
        k_list = [k] if isinstance(k, int) else list(k)
        for current_k in k_list:
            results[f"recall_{current_k}"] = tf.constant(0.0)
            results[f"mrr_{current_k}"] = tf.constant(0.0)
            results[f"ndcg_{current_k}"] = tf.constant(0.0)
        return results


def compute_retrieval_metrics_numpy(logits: np.ndarray, k: int | List[int] = 5) -> Dict[str, float]:
    """Computes metrics for a retrieval task using in-batch negatives in pure NumPy.

    Supports a single k or a list of k values. Frame-work agnostic.

    Args:
        logits (np.ndarray): The result of a dot product (Batch Size, Batch Size).
        k (Union[int, List[int]]): Number of top-K results to evaluate.

    Returns:
        Dict[str, float]: A dictionary containing recall_k, mrr_k, and ndcg_k.
    """
    try:
        k_list = [k] if isinstance(k, int) else list(k)

        batch_size = logits.shape[0]

        # Ground truth is on the diagonal (in-batch negatives)
        true_indices = np.arange(batch_size).reshape(-1, 1)

        # Sort indices descending
        sorted_indices = np.argsort(-logits, axis=1)

        # Find the rank of the true positive (1-indexed)
        # np.where returns a tuple of arrays representing indices.
        # For a 2D array, element 0 is row_idx, element 1 is col_idx (which acts as the 0-indexed rank).
        ranks = np.where(sorted_indices == true_indices)[1] + 1.0

        results = {}

        for current_k in k_list:
            # 1. Recall@K (Is the true item's rank <= K?)
            is_in_top_k = ranks <= float(current_k)
            recall_at_k = float(np.mean(is_in_top_k))

            # 2. MRR@K
            # If rank > K, RR is 0. Else RR is 1/rank
            rr = np.where(is_in_top_k, 1.0 / ranks, 0.0)
            mrr_at_k = float(np.mean(rr))

            # 3. NDCG@K
            # If rank > K, DCG is 0. Else DCG is 1 / log2(rank + 1)
            dcg = np.where(is_in_top_k, 1.0 / (np.log(ranks + 1.0) / np.log(2.0)), 0.0)
            ndcg_at_k = float(np.mean(dcg))

            results[f"recall_{current_k}"] = recall_at_k
            results[f"mrr_{current_k}"] = mrr_at_k
            results[f"ndcg_{current_k}"] = ndcg_at_k

        return results
    except Exception as e:
        print(f"Error computing numpy retrieval metrics: {e}")
        results = {}
        k_list = [k] if isinstance(k, int) else list(k)
        for current_k in k_list:
            results[f"recall_{current_k}"] = 0.0
            results[f"mrr_{current_k}"] = 0.0
            results[f"ndcg_{current_k}"] = 0.0
        return results
