# src/cp4icl/oneshot.py
"""
One-shot conformal methods for in-context learning.

This module currently exposes two methods:

1) SCOS (Split Conformal for One-Shot predictors)
   - Calibrate a separate conformal threshold for each one-shot predictor (source/reference example).
   - Produce prediction sets for each predictor on each evaluation target.

2) CAOS (Conformal Aggregation of One-Shot predictors)
   - Aggregate a pool of one-shot predictors by taking (by default) the mean of the smallest-k
     nonconformity scores across predictors (smaller is better).
   - Calibrate a *single* conformal threshold on aggregated calibration scores.
   - Produce a single aggregated prediction set per evaluation target.

Conventions
-----------
- All scores are treated as *nonconformity* scores: lower is better.
  A prediction set is {y : score(y) <= qhat}.

- Input arrays are expected to have the following shapes:

    calib_scores: (N, N)
        Nonconformity scores for calibration, indexed by (source i, target j).
        In one-shot ICL, "source i" corresponds to the i-th labeled example used as a prompt exemplar.
        "target j" corresponds to the j-th labeled example treated as a pseudo-test point.
        Typically you will want to exclude the diagonal i==j (self-score) from calibration.

    eval_scores: (N, M, K)
        Nonconformity scores for evaluation, indexed by (source i, eval target t, label k).

    y_eval: (M,) optional
        True labels for the evaluation targets, integer in [0, K-1].
        If provided, coverage and oracle set size diagnostics are computed.

Output
------
The public functions return a PredictionSetResult dataclass with:
- pred_mask: boolean array indicating membership in the prediction set(s)
- qhat: the conformal threshold(s)
- set_size: per-example set sizes
- summary metrics (avg set size, coverage if y_eval provided, etc.)

Notes / Future improvements
---------------------------
- This implementation is intentionally NumPy-only for lightweight dependencies.
- For large N, vectorization is used but there are still improvement opportunities:
  * Memory: CAOS currently materializes some intermediate arrays; could be optimized.
  * Speed: could JIT (numba) or use torch for GPU if desired.
- Theoretical subtleties (ties / exact finite-sample conventions) are localized in _conformal_quantile().
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Literal, Optional, Tuple

import numpy as np


QuantileMethod = Literal["higher"]


@dataclass(frozen=True)
class PredictionSetResult:
    """
    Container for conformal prediction set outputs.

    Attributes
    ----------
    pred_mask:
        Boolean membership mask for prediction sets.

        - For SCOS: shape (N, M, K), one set per (source, eval target).
        - For CAOS: shape (M, K), one aggregated set per eval target.

    qhat:
        Conformal threshold(s).
        - For SCOS: shape (N,) (one threshold per source predictor).
        - For CAOS: scalar float (single threshold).

    set_size:
        Size of prediction sets.
        - For SCOS: shape (N, M) sizes per (source, eval target).
        - For CAOS: shape (M,) sizes per eval target.

    coverage:
        Empirical marginal coverage on eval data if y_eval was provided, else None.

        - For SCOS: coverage pooled over all sources and eval targets
        - For CAOS: coverage over eval targets

    avg_set_size:
        Mean of set_size entries (pooled appropriately).

    oracle_avg_set_size:
        Optional diagnostic: average *oracle* set size on eval targets.

        For SCOS: for each eval target, oracle chooses the smallest set among sources that contains the
        true label; then we average that size across eval targets.
        For CAOS: oracle set size is not meaningful in the same way (single aggregated set), so we set None.

    extra:
        Optional dictionary for additional diagnostics we may want to add later
        (e.g., per-source coverage, per-target coverage, etc.).
    """
    pred_mask: np.ndarray
    qhat: np.ndarray | float
    set_size: np.ndarray
    coverage: Optional[float]
    avg_set_size: float
    oracle_avg_set_size: Optional[float]
    extra: dict


def _validate_alpha(alpha: float) -> None:
    """Validate alpha is in (0, 1)."""
    if not (0.0 < alpha < 1.0):
        raise ValueError(f"alpha must be in (0,1). Got {alpha}.")


def _conformal_quantile(values: np.ndarray, alpha: float, method: QuantileMethod = "higher") -> float:
    """
    Compute the conformal quantile used for thresholding.

    This wraps NumPy's quantile API while keeping the "higher" convention stable across NumPy versions.

    We use the common finite-sample conformal choice:
        qhat = Quantile(values, tau) with tau = (1 - alpha) * (n + 1) / n
    and with 'higher' tie handling (conservative).

    Parameters
    ----------
    values:
        1D array of calibration nonconformity scores.
    alpha:
        Miscoverage level in (0,1).
    method:
        Quantile tie-handling. Currently only "higher" is supported.

    Returns
    -------
    qhat:
        Scalar threshold.
    """
    _validate_alpha(alpha)

    vals = np.asarray(values, dtype=float)
    if vals.ndim != 1:
        raise ValueError(f"values must be 1D, got shape {vals.shape}.")
    n = vals.shape[0]
    if n == 0:
        raise ValueError("values must be non-empty for calibration.")

    # tau corresponds to the (ceil((n+1)(1-alpha)))-th order statistic in the usual conformal construction.
    # Using 'higher' makes this conservative under ties.
    tau = (1.0 - alpha) * (n + 1.0) / n

    # NumPy expects q in [0,1]. tau can exceed 1.0 when alpha is extremely small or n is tiny.
    # In conformal prediction, tau>1 should just select the maximum score.
    q = min(max(tau, 0.0), 1.0)

    # NumPy API changed: 'interpolation' (old) -> 'method' (new).
    # We try the new API first, then fallback.
    try:
        return float(np.quantile(vals, q, method=method))
    except TypeError:
        # Older NumPy
        return float(np.quantile(vals, q, interpolation=method))


def _min_k_mean(arr: np.ndarray, k: int, axis: int = 0) -> np.ndarray:
    """
    Compute mean of the smallest-k elements along a given axis.

    This is the core aggregation primitive used in CAOS: https://arxiv.org/abs/2601.05219.

    Parameters
    ----------
    arr:
        Array of scores. Typical shapes:
        - (N, N) calibration matrix
        - (N, M, K) evaluation tensor
    k:
        Number of smallest elements to average.
    axis:
        Axis along which to take the smallest-k.

    Returns
    -------
    aggregated:
        Array with the same shape as arr, but with `axis` removed (reduced),
        since we aggregate across that axis.
    """
    x = np.asarray(arr, dtype=float)
    if k <= 0:
        raise ValueError(f"k must be >= 1, got {k}.")
    if k > x.shape[axis]:
        raise ValueError(f"k cannot exceed size along axis {axis} ({x.shape[axis]}), got k={k}.")

    # np.partition partially sorts so that the smallest-k elements are in the first k slots
    # along the chosen axis; ordering inside those k slots is arbitrary, which is fine.
    part = np.partition(x, kth=k - 1, axis=axis)

    # Slice out the first k elements along `axis`.
    slc = [slice(None)] * part.ndim
    slc[axis] = slice(0, k)
    smallest_k = part[tuple(slc)]

    # Average the k elements, removing that axis.
    return smallest_k.mean(axis=axis)


def _compute_coverage_from_mask(pred_mask: np.ndarray, y_eval: np.ndarray) -> float:
    """
    Compute empirical marginal coverage given a prediction-set mask and true labels.

    This function supports both:
    - pred_mask shape (M, K)  [CAOS]
    - pred_mask shape (N, M, K)  [SCOS pooled over sources]

    Returns pooled coverage (over sources as well if present).
    """
    y = np.asarray(y_eval, dtype=int)
    if y.ndim != 1:
        raise ValueError(f"y_eval must be 1D, got shape {y.shape}.")

    if pred_mask.ndim == 2:
        # CAOS: (M, K)
        M, K = pred_mask.shape
        if y.shape[0] != M:
            raise ValueError(f"y_eval has length {y.shape[0]} but pred_mask has M={M}.")
        return float(pred_mask[np.arange(M), y].mean())

    if pred_mask.ndim == 3:
        # SCOS: (N, M, K)
        N, M, K = pred_mask.shape
        if y.shape[0] != M:
            raise ValueError(f"y_eval has length {y.shape[0]} but pred_mask has M={M}.")
        # Broadcast indexing across sources:
        # pred_mask[:, np.arange(M), y] -> shape (N, M), then mean pools everything.
        return float(pred_mask[:, np.arange(M), y].mean())

    raise ValueError(f"pred_mask must be 2D or 3D, got ndim={pred_mask.ndim}.")


def scos(
    alpha: float,
    calib_scores: np.ndarray,
    eval_scores: np.ndarray,
    *,
    y_eval: Optional[np.ndarray] = None,
    method: QuantileMethod = "higher",
) -> PredictionSetResult:
    """
    SCOS: Split Conformal for One-Shot predictors.

    Intuition
    ---------
    We have a pool of one-shot predictors, indexed by i=0..N-1. Each predictor i is formed by
    prompting (conditioning on) the i-th labeled example.

    For each predictor i, we:
      1) Use its calibration nonconformity scores on the calibration targets to compute a threshold qhat_i.
      2) Form prediction sets on evaluation targets: {y : score_i(t, y) <= qhat_i}.

    Data (nonconformity scores) is assumed to be arranged into arrays:
      - calib_scores[i, j] = nonconformity score from source i on calibration target j
      - eval_scores[i, t, k] = nonconformity score from source i on eval target t for label k

    Parameters
    ----------
    alpha:
        Miscoverage level in (0,1).
    calib_scores:
        Array of shape (N, Nc). Calibration scores for each source predictor.
        NOTE: This is *not required* to be square for SCOS.
    eval_scores:
        Array of shape (N, M, K). Evaluation scores for each source predictor.
    y_eval:
        Optional array of shape (M,) with true labels in [0, K-1].
        If provided, we compute pooled coverage and an oracle average set size diagnostic.
    method:
        Quantile tie-handling; "higher" is conservative under ties.

    Returns
    -------
    PredictionSetResult
        pred_mask has shape (N, M, K) (prediction sets per source and eval target).
        qhat has shape (N,) (threshold per source).
        set_size has shape (N, M).
    """
    _validate_alpha(alpha)

    cal = np.asarray(calib_scores, dtype=float)
    ev = np.asarray(eval_scores, dtype=float)

    if cal.ndim != 2:
        raise ValueError(f"calib_scores must be 2D (N, Nc), got shape {cal.shape}.")
    if ev.ndim != 3:
        raise ValueError(f"eval_scores must be 3D (N, M, K), got shape {ev.shape}.")

    N, Nc = cal.shape
    N2, M, K = ev.shape
    if N2 != N:
        raise ValueError(f"Mismatch: calib_scores has N={N} but eval_scores has N={N2}.")

    # Compute one conformal threshold per source i.
    # This is the main SCOS calibration step.
    qhat = np.empty(N, dtype=float)
    for i in range(N):
        # NOTE: In a faster implementation, we could vectorize quantile computation,
        # but loop over sources is usually fine since it's O(N) calls.
        qhat[i] = _conformal_quantile(cal[i, :], alpha=alpha, method=method)

    # Form prediction sets: score <= qhat_i for each source i.
    # Broadcasting: qhat[:, None, None] broadcasts over (M, K).
    pred_mask = ev <= qhat[:, None, None]  # shape (N, M, K)

    # Compute set sizes for convenience.
    set_size = pred_mask.sum(axis=-1)  # (N, M)

    # Pooled summary stats.
    avg_set_size = float(set_size.mean())

    coverage = None
    oracle_avg_set_size = None
    extra: dict = {}

    if y_eval is not None:
        # Compute pooled coverage across all sources and eval targets.
        coverage = _compute_coverage_from_mask(pred_mask, y_eval)

        # Oracle diagnostic: for each eval target t, among sources that contain the true label,
        # take the smallest set size. If none contain it, oracle size defaults to K.
        y = np.asarray(y_eval, dtype=int)
        if y.shape != (M,):
            raise ValueError(f"y_eval must have shape (M,)={(M,)}, got {y.shape}.")

        # contains_true[i, t] indicates whether source i's prediction set contains the true label at t.
        contains_true = pred_mask[:, np.arange(M), y]  # (N, M)

        # sizes[i, t] is set size for source i at target t.
        sizes = set_size  # (N, M)

        # Compute oracle sizes by scanning targets.
        # NOTE: This loop is easy to read and fast enough for moderate M.
        # Improvement: could be vectorized with masked min, but that can be less readable.
        oracle_sizes = np.empty(M, dtype=float)
        for t in range(M):
            valid_sizes = sizes[contains_true[:, t], t]  # sizes of sources that cover true label
            if valid_sizes.size == 0:
                oracle_sizes[t] = float(K)  # worst-case if nobody covers true label
            else:
                oracle_sizes[t] = float(valid_sizes.min())

        oracle_avg_set_size = float(oracle_sizes.mean())

        # Provide extra optional diagnostics.
        extra["per_source_coverage"] = pred_mask[:, np.arange(M), y].mean(axis=1)  # (N,)
        extra["per_source_avg_size"] = set_size.mean(axis=1)  # (N,)

    # NOTE: Future improvement: we might want an option to return *lists of label indices*
    # for each prediction set. For now, boolean masks are composable and fast.
    return PredictionSetResult(
        pred_mask=pred_mask,
        qhat=qhat,
        set_size=set_size,
        coverage=coverage,
        avg_set_size=avg_set_size,
        oracle_avg_set_size=oracle_avg_set_size,
        extra=extra,
    )


def caos(
    alpha: float,
    calib_scores: np.ndarray,
    eval_scores: np.ndarray,
    *,
    k: int = 3,
    remove_self_score: bool = True,
    y_eval: Optional[np.ndarray] = None,
    method: QuantileMethod = "higher",
) -> PredictionSetResult:
    """
    CAOS: Conformal Aggregation of One-Shot predictors.

    Intuition
    ---------
    CAOS aggregates a pool of one-shot predictors (sources) in an instance-adaptive way.
    At a high level, for each target x and label y, we consider the scores produced by all
    one-shot predictors and aggregate them by averaging the smallest-k scores (lower is better).
    This favors the predictors that are most "compatible" with the target.

    Concretely:
      - Calibration aggregation:
            cal_agg[j] = mean_k( { calib_scores[i, j] : i in sources } )
        where mean_k takes the mean of the smallest k elements across sources i, for each target j.

      - Evaluation aggregation:
            eval_agg[t, y] = mean_k( { eval_scores[i, t, y] : i in sources } )

      - Conformal calibration:
            qhat = Quantile(cal_agg, tau) with tau = (1-alpha) * (n+1)/n

      - Prediction set:
            C(t) = { y : eval_agg[t, y] <= qhat }

    Self-score removal
    ------------------
    In many one-shot settings, calib_scores is square (N,N) with aligned indexing
    between sources and calibration targets. In that case, calib_scores[i,i] is a
    "self-score" that should be excluded (set to +inf) so it cannot artificially
    improve calibration.

    This is controlled by `remove_self_score`. If True, we require calib_scores to be square.

    Parameters
    ----------
    alpha:
        Miscoverage level in (0,1).
    calib_scores:
        Array of shape (N, N). Calibration scores (source, target).
        Must be square if remove_self_score is True.
    eval_scores:
        Array of shape (N, M, K). Evaluation scores (source, eval target, label).
    k:
        Number of smallest scores to average in the aggregation. Must satisfy 1 <= k <= N.
    remove_self_score:
        If True, set diagonal of calib_scores to +inf before aggregation. Requires square calib_scores.
    y_eval:
        Optional true labels of shape (M,). If provided, compute empirical coverage.
    method:
        Quantile tie-handling; "higher" is conservative under ties.

    Returns
    -------
    PredictionSetResult
        pred_mask has shape (M, K).
        qhat is a scalar float.
        set_size has shape (M,).
    """
    _validate_alpha(alpha)

    cal = np.asarray(calib_scores, dtype=float)
    ev = np.asarray(eval_scores, dtype=float)

    if cal.ndim != 2:
        raise ValueError(f"calib_scores must be 2D, got shape {cal.shape}.")
    if ev.ndim != 3:
        raise ValueError(f"eval_scores must be 3D (N, M, K), got shape {ev.shape}.")

    N, Nc = cal.shape
    N2, M, K = ev.shape
    if N2 != N:
        raise ValueError(f"Mismatch: calib_scores has N={N} but eval_scores has N={N2}.")

    if remove_self_score:
        # We only know what "self" means if calib_scores is square and aligned.
        if N != Nc:
            raise ValueError(
                "remove_self_score=True requires calib_scores to be square (N,N) with aligned indices. "
                f"Got shape {cal.shape}. Set remove_self_score=False if this is not applicable."
            )
        # Copy to avoid mutating user input.
        cal = cal.copy()
        # Set diagonal to +inf so it never appears among the smallest-k.
        np.fill_diagonal(cal, np.inf)

    # Aggregate calibration scores across sources: for each calibration target j,
    # compute mean of smallest-k across i.
    #
    # cal has shape (N, Nc); we aggregate along axis=0 (sources), yielding shape (Nc,).
    cal_agg = _min_k_mean(cal, k=k, axis=0)  # (Nc,)

    # Aggregate evaluation scores across sources: for each eval target t and label y,
    # compute mean of smallest-k across sources i.
    #
    # ev has shape (N, M, K); aggregate along axis=0, yielding shape (M, K).
    ev_agg = _min_k_mean(ev, k=k, axis=0)  # (M, K)

    # Calibrate a single conformal threshold on the aggregated calibration scores.
    qhat = _conformal_quantile(cal_agg, alpha=alpha, method=method)

    # Build prediction sets for each eval target: include label y if aggregated score <= qhat.
    pred_mask = ev_agg <= qhat  # (M, K)
    set_size = pred_mask.sum(axis=1)  # (M,)

    avg_set_size = float(set_size.mean())

    coverage = None
    extra: dict = {}
    if y_eval is not None:
        coverage = _compute_coverage_from_mask(pred_mask, y_eval)
        # Additional helpful diagnostics:
        extra["avg_score_margin"] = float((qhat - ev_agg.min(axis=1)).mean())

    # CAOS has a single prediction set per eval target, so oracle set size is not applicable.
    return PredictionSetResult(
        pred_mask=pred_mask,
        qhat=float(qhat),
        set_size=set_size,
        coverage=coverage,
        avg_set_size=avg_set_size,
        oracle_avg_set_size=None,
        extra=extra,
    )


# -----------------------------------------------------------------------------
# TODO / future improvements (intentionally left as comments for maintainability)
# -----------------------------------------------------------------------------
# - Add optional conversion utilities:
#     dict_scores_to_arrays(calib_scores_dict, eval_scores_dict, ...)
#   to ease migration from notebook code.
#
# - Add a unified "predict" interface that supports both methods:
#     predict(method="caos"|"scos", ...)
#
# - Add better typing for numpy arrays using numpy.typing.NDArray.
#
# - Add optional return format:
#     return_list=True -> list-of-arrays of label indices per prediction set
#
# - Add tests under tests/:
#     * quantile edge cases
#     * SCOS loop vs vectorized equivalence
#     * CAOS k=1 sanity checks
#
# - For huge N, consider memory optimizations or torch implementation.

