class QuantificationMetric(ABC):
  """
  Abstract class to all quantification metrics. It uses the strategy design pattern.

  Parameters
  ----------
  name : str
    Human-readable name of the metric (used in reports/plots).
  lower_is_better : bool, default = True
    Whether lower values of this metric indicate better quantification
    performance. All metrics currently shipped with `quack` are error/
    divergence measures, so this defaults to True.
  """
  def __init__(self, name: str, lower_is_better: bool = True):
    self.name = name
    self.lower_is_better = lower_is_better

  def _validate_inputs(self, p_true: np.ndarray, p_pred: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """ Ensure that all inputs are in the same format, shape and type.

    Parameters
    ----------
    p_true: np.ndarray
      Array-like with all true prevalences.
    p_pred: np.ndarray
      Array-like with all predicted prevalences.

    Raises
    ------
    ValueError
      thrown when shapes are incompatible or inputs aren't 1D.

    Returns
    -------
    (p_true, p_pred): tuple[np.ndarray, np.ndarray]
      Returns the two arrays in a numpy compatible format.
    """
    p_true = np.asarray(p_true, dtype=np.float64)
    p_pred = np.asarray(p_pred, dtype=np.float64)

    if p_true.shape != p_pred.shape:
      raise ValueError(
        f"Shape Error: p_true has a shape {p_true.shape} "
        f"and p_pred have shape {p_pred.shape}."
      )
    if p_true.ndim != 1:
      raise ValueError(
        f"Quantification metrics expect 1D prevalence vectors (one value "
        f"per class), got arrays with shape {p_true.shape}."
      )

    for label, p in (("p_true", p_true), ("p_pred", p_pred)):
      if not np.isclose(p.sum(), 1.0, atol=1e-6):
        warnings.warn(
          f"'{label}' does not sum to 1.0 (got {p.sum():.6f}); it may not "
          "represent a valid prevalence distribution and the metric result "
          "may be misleading.",
          stacklevel=3,
        )

    return p_true, p_pred

  @staticmethod
  def _smooth(p: np.ndarray, epsilon: float) -> np.ndarray:
    """Additive smoothing that keeps `p` a valid probability distribution.

    Following the convention adopted in the quantification literature
    (Forman, 2008; Esuli & Sebastiani, 2015), this both avoids division-
    by-zero / `log(0)` issues for classes absent from a bag and, unlike a
    naive `p + epsilon`, renormalizes so the smoothed vector still sums
    to exactly 1.0.

      p_s(c) = (p(c) + epsilon) / (1 + n_classes * epsilon)

    Parameters
    ----------
    p: np.ndarray
      Prevalence vector of shape `(n_classes,)`.
    epsilon: float
      Smoothing factor.

    Returns
    -------
      p_smoothed: np.ndarray
        Smoothed prevalence vector, still summing to 1.0.
    """
    n_classes = p.shape[0]
    return (p + epsilon) / (1.0 + n_classes * epsilon)

  @abstractmethod
  def compute(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    """ Each metric implements its own logic and mathematics."""
    pass

  def __call__(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    """ Call and perform the input validation and metric computation.

    Parameters
    ----------
    p_true: np.ndarray
      Array-like with all true prevalences.
    p_pred: np.ndarray
      Array-like with all predicted prevalences.

    Returns
    -------
    result: float
      The computed metric value.
    """
    p_true_clean, p_pred_clean = self._validate_inputs(p_true, p_pred)
    return self.compute(p_true_clean, p_pred_clean)

  def __repr__(self) -> str:
    return f"{self.__class__.__name__}(name={self.name!r}, lower_is_better={self.lower_is_better})"

class AbsoluteError(QuantificationMetric):
  """Mean Absolute Error (AE) between true and predicted prevalence vectors.

  Averages the absolute per-class deviation across all classes, bounding
  the metric to `[0, 1]` regardless of the number of classes (0 = perfect
  quantification, 1 = maximally wrong).

    AE(p, p_hat) = (1 / n_classes) * sum_c |p(c) - p_hat(c)|

  References
  ----------
  George Forman. Quantifying counts and costs via classification.
  Data Mining and Knowledge Discovery, 17(2):164-206, 2008.
  """
  def __init__(self):
    super().__init__(name="Absolute Error", lower_is_better=True)

  def compute(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    return float(np.mean(np.abs(p_true - p_pred)))

class RelativeAbsoluteError(QuantificationMetric):
  """Relative Absolute Error (RAE) between true and predicted prevalence vectors.

  Averages the per-class absolute deviation relative to the true
  prevalence. Both `p_true` and `p_pred` are additively smoothed (see
  `QuantificationMetric._smooth`) before the ratio is computed, since
  dividing by a true prevalence of exactly 0 would otherwise make the
  metric undefined for classes absent from the test bag.

    p_s(c)        = (p(c) + eps) / (1 + n_classes * eps)
    RAE(p, p_hat) = (1 / n_classes) * sum_c |p_s(c) - p_hat_s(c)| / p_s(c)

  Parameters
  ----------
  epsilon : float, default = 1e-5
    Smoothing factor applied to both `p_true` and `p_pred` before
    computing the ratio.

  References
  ----------
  George Forman. Quantifying counts and costs via classification.
  Data Mining and Knowledge Discovery, 17(2):164-206, 2008.
  """
  def __init__(self, epsilon: float = 1e-5):
    super().__init__(name="Relative Absolute Error", lower_is_better=True)
    self.epsilon = epsilon

  def compute(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    p_true_smoothed = self._smooth(p_true, self.epsilon)
    p_pred_smoothed = self._smooth(p_pred, self.epsilon)
    return float(np.mean(np.abs(p_true_smoothed - p_pred_smoothed) / p_true_smoothed))

class KullbackLeiblerDivergence(QuantificationMetric):
  """Kullback-Leibler Divergence (KLD) between true and predicted prevalence vectors.

  Both `p_true` and `p_pred` are additively smoothed and renormalized
  (see `QuantificationMetric._smooth`) so they remain valid probability
  distributions before computing the divergence, avoiding `log(0)` /
  division issues for classes with zero true or estimated prevalence.

    p_s(c)          = (p(c) + eps) / (1 + n_classes * eps)
    KLD(p, p_hat)   = sum_c p_s(c) * log( p_s(c) / p_hat_s(c) )

  Parameters
  ----------
  epsilon : float, default = 1e-5
    Smoothing factor applied to both `p_true` and `p_pred`.

  References
  ----------
  Esuli, A. & Sebastiani, F. (2015). Optimizing text quantifiers for
  multivariate loss functions. ACM Transactions on Knowledge Discovery
  from Data, 9(4), 1-27.
  """
  def __init__(self, epsilon: float = 1e-5):
    super().__init__(name="Kullback-Leibler Divergence", lower_is_better=True)
    self.epsilon = epsilon

  def compute(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    p_true_smoothed = self._smooth(p_true, self.epsilon)
    p_pred_smoothed = self._smooth(p_pred, self.epsilon)
    return float(np.sum(p_true_smoothed * np.log(p_true_smoothed / p_pred_smoothed)))

class NormalizedKullbackLeiblerDivergence(QuantificationMetric):
  """Normalized Kullback-Leibler Divergence (NKLD).

  Squashes the unbounded KLD into the `[0, 1)` range via a logistic-style
  transform, making it comparable across experiments/datasets:

    NKLD(p, p_hat) = max(0, 2 * exp(KLD(p, p_hat)) / (1 + exp(KLD(p, p_hat))) - 1)

  Parameters
  ----------
  epsilon : float, default = 1e-5
    Smoothing factor forwarded to the internal `KullbackLeiblerDivergence`
    instance (`self.kld`).

  References
  ----------
  Esuli, A. & Sebastiani, F. (2015). Optimizing text quantifiers for
  multivariate loss functions. ACM Transactions on Knowledge Discovery
  from Data, 9(4), 1-27.
  """
  def __init__(self, epsilon: float = 1e-5):
    super().__init__(name="Normalized Kullback-Leibler Divergence", lower_is_better=True)
    self.epsilon = epsilon
    self.kld = KullbackLeiblerDivergence(epsilon=epsilon)

  def compute(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    # inputs are already validated by __call__, so we go straight to
    # .compute() on the internal KLD instance instead of re-validating
    # through its own __call__ (which does not accept an `eps` kwarg).
    exp_kld = math.exp(self.kld.compute(p_true, p_pred))
    return max(0.0, 2.0 * exp_kld / (1.0 + exp_kld) - 1.0)

class NormalizedAbsoluteError(QuantificationMetric):
  """Normalized Absolute Error (NAE) between true and predicted prevalence vectors.

  Normalizes the (unaveraged) Absolute Error by its theoretical maximum
  given `p_true`, bounding the metric to `[0, 1]` regardless of how
  skewed the true prevalence is. This makes NAE more comparable across
  experiments/datasets with very different training or test prevalences
  than the plain `AbsoluteError`.

    NAE(p, p_hat) = sum_c |p(c) - p_hat(c)| / (2 * (1 - min_c p(c)))

  References
  ----------
  Esuli, A. & Sebastiani, F. (2015). Optimizing text quantifiers for
  multivariate loss functions. ACM Transactions on Knowledge Discovery
  from Data, 9(4), 1-27.
  """
  def __init__(self):
    super().__init__(name="Normalized Absolute Error", lower_is_better=True)

  def compute(self, p_true: np.ndarray, p_pred: np.ndarray) -> float:
    max_ae = 2.0 * (1.0 - np.min(p_true))
    if max_ae <= 0:
      # only possible when n_classes == 1 (a single class holds all the mass
      # across every class simultaneously, i.e. a degenerate 1-class problem)
      return 0.0
    return float(np.sum(np.abs(p_true - p_pred)) / max_ae)