"""
Evaluation metrics for SAD.
"""

from collections import Counter
from typing import Dict, List, Optional

import torch
import numpy as np


def compute_exit_depth_histogram(
    exit_levels: torch.Tensor,
    num_levels: int,
) -> Dict[str, float]:
    """
    Compute histogram of token exit depths from adaptive decoding.

    Args:
        exit_levels: [B, S] int64 tensor of per-token exit levels
        num_levels:  total number of levels

    Returns:
        dict mapping "level_{l}_frac" to fraction of tokens that exited at level l.
    """
    flat = exit_levels.reshape(-1)
    total = flat.numel()
    hist = {}
    for l in range(num_levels):
        count = (flat == l).sum().item()
        hist[f"level_{l}_frac"] = count / max(total, 1)
    return hist


def compute_unresolved_over_steps(resolved_over_steps: List[float]) -> Dict[str, float]:
    """
    Compute statistics about the unresolved-token fraction over decoding steps.

    Args:
        resolved_over_steps: list of floats (fraction resolved at each step)

    Returns:
        dict with step statistics
    """
    if not resolved_over_steps:
        return {}
    arr = np.array(resolved_over_steps)
    return {
        "resolved_final": float(arr[-1]),
        "resolved_step_50pct": int(np.searchsorted(arr, 0.5)),
        "resolved_step_90pct": int(np.searchsorted(arr, 0.9)),
    }


def compute_diversity(token_ids: torch.Tensor) -> Dict[str, float]:
    """
    Compute simple text diversity metrics on a batch of generated sequences.

    Args:
        token_ids: [B, S]

    Returns:
        dict with dist-1, dist-2, unique_sequences fraction
    """
    B, S = token_ids.shape

    # dist-1: fraction of unique unigrams
    flat = token_ids.reshape(-1).tolist()
    uniq_1 = len(set(flat)) / max(len(flat), 1)

    # dist-2: fraction of unique bigrams
    bigrams = [(flat[i], flat[i + 1]) for i in range(len(flat) - 1)]
    uniq_2 = len(set(bigrams)) / max(len(bigrams), 1)

    # unique sequences
    seqs = [tuple(token_ids[i].tolist()) for i in range(B)]
    uniq_seqs = len(set(seqs)) / max(B, 1)

    return {
        "dist_1": uniq_1,
        "dist_2": uniq_2,
        "unique_sequences": uniq_seqs,
    }