"""Weighted conformal prediction under test-dependent importance weights."""
import numpy as np

from .base import ConformalResult


def _weighted_quantile_with_test_mass(
    values: np.ndarray,
    weights: np.ndarray,
    alpha: float,
    test_weight: float,
) -> float:
    """Compute the weighted conformal threshold for one test point.

    The test point contributes mass at +inf, matching the weighted split
    conformal recipe used under covariate shift.
    """
    values = np.asarray(values, dtype=float)
    weights = np.asarray(weights, dtype=float)

    order = np.argsort(values)
    sorted_values = values[order]
    sorted_weights = weights[order]

    total_weight = sorted_weights.sum() + test_weight
    cutoff = (1.0 - alpha) * total_weight
    csum = np.cumsum(sorted_weights)
    idx = np.searchsorted(csum, cutoff, side="left")
    if idx >= len(sorted_values):
        return np.inf
    return float(sorted_values[idx])


def weighted_conformal(
    R_cal: np.ndarray,
    R_test: np.ndarray,
    alpha: float,
    weights_cal: np.ndarray,
    weights_test: np.ndarray | None = None,
) -> ConformalResult:
    """Weighted split conformal.

    Args:
        R_cal: calibration residuals (n_cal,)
        R_test: test residuals (n_test,)
        alpha: miscoverage level
        weights_cal: nonnegative calibration weights
        weights_test: nonnegative test weights. If omitted, uses ones.

    Returns:
        ConformalResult with test-specific weighted thresholds.
    """
    weights_cal = np.asarray(weights_cal, dtype=float)
    if weights_cal.shape != np.asarray(R_cal).shape:
        raise ValueError("weights_cal must have the same shape as R_cal")
    if not np.all(np.isfinite(weights_cal)):
        raise ValueError("weights_cal must be finite")
    if np.any(weights_cal < 0):
        raise ValueError("weights_cal must be nonnegative")
    if float(weights_cal.sum()) <= 0.0:
        raise ValueError("weights_cal must have positive total mass")

    if weights_test is None:
        weights_test = np.ones(len(R_test), dtype=float)
    else:
        weights_test = np.asarray(weights_test, dtype=float)
        if weights_test.shape != np.asarray(R_test).shape:
            raise ValueError("weights_test must have the same shape as R_test")
        if not np.all(np.isfinite(weights_test)):
            raise ValueError("weights_test must be finite")
        if np.any(weights_test < 0):
            raise ValueError("weights_test must be nonnegative")

    radius = np.array([
        _weighted_quantile_with_test_mass(R_cal, weights_cal, alpha, w_t)
        for w_t in weights_test
    ])
    covered = R_test <= radius
    return ConformalResult(covered=covered, radius=radius, threshold=radius.copy())