class BaseBagGenerator(ABC, BaseEstimator):
  """Abstract base class for all dataset-shift bag generators.

  A bag generator draws repeated "bags" (labeled subsets) of a fixed
  size from a labeled dataset `(X, y)`, simulating a specific kind of
  dataset shift between training and test-time distributions. Bags
  produced this way follow the standard evaluation protocol in the
  quantification literature (the "Artificial Prevalence Protocol", APP),
  letting quantifiers be benchmarked across the full spectrum of
  possible test-time class prevalences (or covariate shifts) rather than
  relying on a single fixed train/test split.

  Parameters
  ----------
  n_bags : int, default = 100
    Number of bags to generate.
  bag_size : int, default = None
    Number of instances per bag. If None, defaults to `len(y)` (the size
    of the original dataset).
  random_state : int, RandomState instance or None, default = None
    Controls the randomness of the bag sampling process. Pass an int for
    reproducible bags across repeated calls to `generate` — particularly
    useful when comparing multiple quantifiers on the exact same
    sequence of bags (e.g. for `quack.visualization.prevalence_plot`).
  """

  def __init__(self, n_bags: int = 100, bag_size: int = None, random_state=None):
    self.n_bags = n_bags
    self.bag_size = bag_size
    self.random_state = random_state

  @staticmethod
  def _group_indices_by_class(y: np.ndarray, classes: np.ndarray) -> dict:
    """Maps each class label to the array of dataset indices holding it."""
    return {c: np.flatnonzero(y == c) for c in classes}

  def _validate(self, X: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    X, y = check_X_y(X, y, accept_sparse=True)
    if self.n_bags <= 0:
      raise ValueError(f"n_bags must be a positive integer, got {self.n_bags}.")
    if self.bag_size is not None and self.bag_size <= 0:
      raise ValueError(f"bag_size must be a positive integer, got {self.bag_size}.")
    return X, y

  @abstractmethod
  def generate(self, X: np.ndarray, y: np.ndarray) -> Generator[tuple, None, None]:
    """Lazily yields `n_bags` labeled bags `(X_bag, y_bag)` drawn from `(X, y)`.

    Parameters
    ----------
    X : {array-like, sparse matrix} of shape (n_samples, n_features)
      The pool of features to draw bags from.
    y : array-like of shape (n_samples,)
      The corresponding labels.

    Yields
    ------
    X_bag : ndarray of shape (bag_size, n_features)
      Feature matrix of a single generated bag.
    y_bag : ndarray of shape (bag_size,)
      Corresponding labels for the generated bag.
    """
    pass

  def to_list(self, X: np.ndarray, y: np.ndarray) -> list[tuple]:
    """Eagerly materializes `generate(X, y)` into a list of `(X_bag, y_bag)`.

    Useful when the same set of bags needs to be iterated multiple times
    (e.g. once per quantifier being benchmarked), since generators can
    otherwise only be consumed once.

    Returns
    -------
    bags : list[tuple[np.ndarray, np.ndarray]]
      List of `(X_bag, y_bag)` pairs, of length `n_bags`.
    """
    return list(self.generate(X, y))

class PriorShiftBagGenerator(BaseBagGenerator):
  """Simulates Prior Probability Shift by resampling bags across the
  class-prevalence simplex, preserving `P(X|y)`.

  For each bag, a target class-prevalence vector `p` is sampled (see
  `sampling_strategy`) and then, independently for each class `c`,
  `round(p[c] * bag_size)` instances are drawn from the pool of original
  instances of class `c` — so the class-conditional feature distribution
  `P(X|y=c)` is left untouched and only the marginal `P(y)` is shifted.
  This is the standard "Artificial Prevalence Protocol" (APP) used to
  benchmark quantifiers.

  Parameters
  ----------
  n_bags : int, default = 100
    Number of bags to generate.
  bag_size : int, default = None
    Number of instances per bag. If None, defaults to `len(y)`.
  sampling_strategy : {'uniform', 'dirichlet'}, default = 'uniform'
    Strategy used to sample each bag's target prevalence vector:
    - `'uniform'`: samples uniformly over the full probability simplex
      via the standard Kraemer algorithm (sorting `n_classes - 1`
      independent `Uniform(0, 1)` cut points).
    - `'dirichlet'`: samples from a `Dirichlet(dirichlet_alpha)`
      distribution, allowing control over how extreme/skewed the shifts
      are via `dirichlet_alpha` (values < 1 favor prevalences
      concentrated in a single class; values > 1 favor prevalences
      closer to uniform).
  dirichlet_alpha : float | array-like of shape (n_classes,), default = 1.0
    Concentration parameter(s) for the Dirichlet distribution. Only used
    when `sampling_strategy='dirichlet'`. A scalar is broadcast to all
    classes; `alpha=1.0` for every class is equivalent to `'uniform'`.
  with_replacement : bool, default = True
    Whether instances are drawn with replacement from each class pool.
    Automatically forced to True for a given class/bag whenever the
    requested count exceeds the number of available instances of that
    class, regardless of this setting.
  random_state : int, RandomState instance or None, default = None
    Controls the randomness of both the prevalence sampling and the
    instance resampling.

  Attributes
  ----------
  classes_ : ndarray of shape (n_classes,)
    The distinct class labels found in `y` the last time `generate` was
    called.
  sampled_prevalences_ : ndarray of shape (n_bags, n_classes)
    The realized class prevalence of each generated bag's `y_bag` (i.e.
    the target prevalence after the largest-remainder integer rounding).

  References
  ----------
  George Forman. Quantifying counts and costs via classification.
  Data Mining and Knowledge Discovery, 17(2):164-206, 2008.

  Examples
  --------
  >>> from sklearn.datasets import make_classification
  >>> from quack.bag_generator import PriorShiftBagGenerator
  >>> X, y = make_classification(n_samples=500, n_classes=2, random_state=0)
  >>> generator = PriorShiftBagGenerator(n_bags=5, bag_size=100, random_state=0)
  >>> bags = generator.to_list(X, y)
  >>> len(bags)
  5
  >>> generator.sampled_prevalences_.shape
  (5, 2)
  """

  def __init__(self,
               n_bags: int = 100,
               bag_size: int = None,
               sampling_strategy: str = 'uniform',
               dirichlet_alpha: float = 1.0,
               with_replacement: bool = True,
               random_state=None):
    super().__init__(n_bags=n_bags, bag_size=bag_size, random_state=random_state)
    self.sampling_strategy = sampling_strategy
    self.dirichlet_alpha = dirichlet_alpha
    self.with_replacement = with_replacement

  def _sample_prevalence(self, n_classes: int, rng) -> np.ndarray:
    """Draws a single target prevalence vector summing to 1.0."""
    if self.sampling_strategy == 'uniform':
      if n_classes == 1:
        return np.ones(1)
      cuts = np.sort(rng.uniform(0.0, 1.0, size=n_classes - 1))
      cuts = np.concatenate(([0.0], cuts, [1.0]))
      return np.diff(cuts)

    if self.sampling_strategy == 'dirichlet':
      alpha = self.dirichlet_alpha
      alpha = np.full(n_classes, alpha, dtype=float) if np.isscalar(alpha) else np.asarray(alpha, dtype=float)
      if alpha.shape[0] != n_classes:
        raise ValueError(
          f"dirichlet_alpha must be a scalar or have length n_classes={n_classes}, "
          f"got length {alpha.shape[0]}."
        )
      return rng.dirichlet(alpha)

    raise ValueError(
      f"Unknown sampling_strategy '{self.sampling_strategy}'. "
      "Supported options are 'uniform' and 'dirichlet'."
    )

  @staticmethod
  def _prevalence_to_counts(prevalence: np.ndarray, bag_size: int) -> np.ndarray:
    """Converts a real-valued prevalence vector into integer per-class
    counts summing to exactly `bag_size`, via the largest-remainder method."""
    raw_counts = prevalence * bag_size
    counts = np.floor(raw_counts).astype(int)

    remainder = bag_size - counts.sum()
    if remainder > 0:
      fractional_parts = raw_counts - counts
      top_indices = np.argsort(fractional_parts)[::-1][:remainder]
      counts[top_indices] += 1

    return counts

  def generate(self, X: np.ndarray, y: np.ndarray) -> Generator[tuple, None, None]:
    X, y = self._validate(X, y)
    rng = check_random_state(self.random_state)

    self.classes_ = np.unique(y)
    n_classes = len(self.classes_)
    class_pools = self._group_indices_by_class(y, self.classes_)

    bag_size = self.bag_size if self.bag_size is not None else len(y)
    self.sampled_prevalences_ = np.zeros((self.n_bags, n_classes))

    for i in range(self.n_bags):
      prevalence = self._sample_prevalence(n_classes, rng)
      counts = self._prevalence_to_counts(prevalence, bag_size)
      self.sampled_prevalences_[i] = counts / bag_size

      bag_indices = []
      for c_idx, count in enumerate(counts):
        if count == 0:
          continue
        pool = class_pools[self.classes_[c_idx]]
        replace = self.with_replacement or count > len(pool)
        bag_indices.append(rng.choice(pool, size=count, replace=replace))

      bag_indices = np.concatenate(bag_indices)
      rng.shuffle(bag_indices)

      yield X[bag_indices], y[bag_indices]

class CovariateShiftBagGenerator(BaseBagGenerator):
  """Simulates Covariate Shift by resampling bags biased towards random
  regions of the feature space, preserving `P(y|X)`.

  For each bag, a random "pivot" instance is drawn from the dataset and
  every instance's RBF kernel similarity to that pivot is computed:

    k(x, x_pivot) = exp( -gamma * ||x - x_pivot||^2 )

  Instances are then resampled with probability proportional to their
  similarity to the pivot, concentrating the bag around a random region
  of the feature space. Since instances — and their original labels —
  are drawn as-is (no label is ever altered), the conditional
  distribution `P(y|X)` is left untouched; only the marginal feature
  distribution `P(X)` (and, as a natural consequence in most real
  datasets, the marginal `P(y)` too) is shifted.

  Parameters
  ----------
  n_bags : int, default = 100
    Number of bags to generate.
  bag_size : int, default = None
    Number of instances per bag. If None, defaults to `len(y)`.
  gamma : float, default = None
    RBF kernel coefficient. Controls how concentrated each bag is around
    its pivot: larger values produce bags tightly clustered in feature
    space (stronger shift); smaller values approach the original,
    unshifted distribution. If None, defaults to `1 / n_features`
    (scikit-learn's `rbf_kernel` default).
  with_replacement : bool, default = True
    Whether instances are drawn with replacement. Automatically forced
    to True whenever `bag_size` exceeds the number of available
    instances, regardless of this setting.
  random_state : int, RandomState instance or None, default = None
    Controls the randomness of both the pivot selection and the
    instance resampling.

  Attributes
  ----------
  classes_ : ndarray of shape (n_classes,)
    The distinct class labels found in `y` the last time `generate` was
    called.
  pivot_indices_ : ndarray of shape (n_bags,)
    The dataset index of the pivot instance used to build each bag.
  sampled_prevalences_ : ndarray of shape (n_bags, n_classes)
    The realized class prevalence of each generated bag's `y_bag` — a
    side effect of the covariate shift, not directly controlled.

  References
  ----------
  Bickel, S., Brückner, M., & Scheffer, T. (2009). Discriminative
  learning under covariate shift. Journal of Machine Learning Research,
  10, 2137-2155.

  Examples
  --------
  >>> from sklearn.datasets import make_classification
  >>> from quack.bag_generator import CovariateShiftBagGenerator
  >>> X, y = make_classification(n_samples=500, n_classes=2, random_state=0)
  >>> generator = CovariateShiftBagGenerator(n_bags=5, bag_size=100, gamma=0.5, random_state=0)
  >>> bags = generator.to_list(X, y)
  >>> len(bags)
  5
  """

  def __init__(self,
               n_bags: int = 100,
               bag_size: int = None,
               gamma: float = None,
               with_replacement: bool = True,
               random_state=None):
    super().__init__(n_bags=n_bags, bag_size=bag_size, random_state=random_state)
    self.gamma = gamma
    self.with_replacement = with_replacement

  def generate(self, X: np.ndarray, y: np.ndarray) -> Generator[tuple, None, None]:
    X, y = self._validate(X, y)
    rng = check_random_state(self.random_state)

    self.classes_ = np.unique(y)
    n_classes = len(self.classes_)
    n_samples = X.shape[0]

    bag_size = self.bag_size if self.bag_size is not None else n_samples
    replace = self.with_replacement or bag_size > n_samples

    self.pivot_indices_ = np.zeros(self.n_bags, dtype=int)
    self.sampled_prevalences_ = np.zeros((self.n_bags, n_classes))

    for i in range(self.n_bags):
      pivot_idx = rng.randint(n_samples)
      self.pivot_indices_[i] = pivot_idx

      similarities = rbf_kernel(X, X[pivot_idx].reshape(1, -1), gamma=self.gamma).ravel()
      total_similarity = similarities.sum()
      weights = (similarities / total_similarity if total_similarity > 0
                else np.full(n_samples, 1.0 / n_samples))

      bag_indices = rng.choice(n_samples, size=bag_size, replace=replace, p=weights)
      self.sampled_prevalences_[i] = (y[bag_indices][:, None] == self.classes_[None, :]).mean(axis=0)

      yield X[bag_indices], y[bag_indices]