seed
        Random seed value used for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import feature_selection
    >>> from river import stream

    >>> selector = feature_selection.PoissonInclusion(p=0.1, seed=42)

    >>> dataset = iter(datasets.TrumpApproval())

    >>> feature_names = next(dataset)[0].keys()
    >>> n = 0

    >>> while True:
    ...     x, y = next(dataset)
    ...     xt = selector.transform_one(x)
    ...     if xt.keys() == feature_names:
    ...         break
    ...     n += 1

    >>> n
    12

    References
    ----------
    [^1]: [McMahan, H.B., Holt, G., Sculley, D., Young, M., Ebner, D., Grady, J., Nie, L., Phillips, T., Davydov, E., Golovin, D. and Chikkerur, S., 2013, August. Ad click prediction: a view from the trenches. In Proceedings of the 19th ACM SIGKDD international conference on Knowledge discovery and data mining (pp. 1222-1230)](https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf)

    """

    def __init__(self, p: float, seed: int | None = None):
        self.p = p
        self.seed = seed
        self.rng = random.Random(seed)
        self.included: set[base.typing.FeatureName] = set()

    def transform_one(self, x):
        xt = {}

        for i, xi in x.items():
            if i in self.included:
                xt[i] = xi
            elif self.rng.random() < self.p:
                self.included.add(i)
                xt[i] = xi

        return xt


# river/feature_selection/variance.py

from __future__ import annotations

import collections

from river import base, stats


class VarianceThreshold(base.Transformer):
    """Removes low-variance features.

    Parameters
    ----------
    threshold
        Only features with a variance above the threshold will be kept.
    min_samples
        The minimum number of samples required to perform selection.

    Attributes
    ----------
    variances : dict
        The variance of each feature.

    Examples
    --------

    >>> from river import feature_selection
    >>> from river import stream

    >>> X = [
    ...     [0, 2, 0, 3],
    ...     [0, 1, 4, 3],
    ...     [0, 1, 1, 3]
    ... ]

    >>> selector = feature_selection.VarianceThreshold()

    >>> for x, _ in stream.iter_array(X):
    ...     print(selector.learn_one(x).transform_one(x))
    {0: 0, 1: 2, 2: 0, 3: 3}
    {1: 1, 2: 4}
    {1: 1, 2: 1}

    """

    def __init__(self, threshold=0, min_samples=2):
        self.threshold = threshold
        self.min_samples = min_samples
        self.variances = collections.defaultdict(stats.Var)

    def learn_one(self, x):
        for i, xi in x.items():
            self.variances[i].update(xi)

        return self

    def check_feature(self, feature):
        if feature not in self.variances:
            return True
        if self.variances[feature].mean.n < self.min_samples:
            return True
        return self.variances[feature].get() > self.threshold

    def transform_one(self, x):
        return {i: xi for i, xi in x.items() if self.check_feature(i)}


# river/forest/__init__.py

"""This module implements forest-based classifiers and regressors."""
from __future__ import annotations

from .adaptive_random_forest import ARFClassifier, ARFRegressor
from .aggregated_mondrian_forest import AMFClassifier
from .online_extra_trees import OXTRegressor

__all__ = [
    "ARFClassifier",
    "ARFRegressor",
    "AMFClassifier",
    "OXTRegressor",
]


# river/forest/adaptive_random_forest.py

from __future__ import annotations

import abc
import collections
import copy
import math
import random
import typing

import numpy as np

from river import base, metrics, stats
from river.drift import ADWIN
from river.tree.hoeffding_tree_classifier import HoeffdingTreeClassifier
from river.tree.hoeffding_tree_regressor import HoeffdingTreeRegressor
from river.tree.nodes.arf_htc_nodes import (
    RandomLeafMajorityClass,
    RandomLeafNaiveBayes,
    RandomLeafNaiveBayesAdaptive,
)
from river.tree.nodes.arf_htr_nodes import RandomLeafAdaptive, RandomLeafMean, RandomLeafModel
from river.tree.splitter import Splitter
from river.utils.random import poisson


class BaseForest(base.Ensemble):
    _FEATURES_SQRT = "sqrt"
    _FEATURES_LOG2 = "log2"

    def __init__(
        self,
        n_models: int,
        max_features: bool | str | int,
        lambda_value: int,
        drift_detector: base.DriftDetector | None,
        warning_detector: base.DriftDetector | None,
        metric: metrics.base.MultiClassMetric | metrics.base.RegressionMetric,
        disable_weighted_vote,
        seed,
    ):
        super().__init__([])  # type: ignore
        self.n_models = n_models
        self.max_features = max_features
        self.lambda_value = lambda_value
        self.metric = metric
        self.disable_weighted_vote = disable_weighted_vote
        self.drift_detector = drift_detector
        self.warning_detector = warning_detector
        self.seed = seed
        self._rng = random.Random(self.seed)

        # Internal parameters
        self._n_samples_seen = 0
        self._base_member_class: ForestMemberClassifier | ForestMemberRegressor | None = None

    @property
    def _min_number_of_models(self):
        return 0

    @classmethod
    def _unit_test_params(cls):
        yield {"n_models": 3}

    def _unit_test_skips(self):
        return {"check_shuffle_features_no_impact"}

    def learn_one(self, x: dict, y: base.typing.Target, **kwargs):
        self._n_samples_seen += 1

        if len(self) == 0:
            self._init_ensemble(sorted(x.keys()))

        for model in self:
            # Get prediction for instance
            y_pred = model.predict_one(x)

            # Update performance evaluator
            model.metric.update(y_true=y, y_pred=y_pred)

            k = poisson(rate=self.lambda_value, rng=self._rng)
            if k > 0:
                model.learn_one(x=x, y=y, sample_weight=k, n_samples_seen=self._n_samples_seen)

        return self

    def _init_ensemble(self, features: list):
        self._set_max_features(len(features))

        self.data = [
            self._base_member_class(  # type: ignore
                index_original=i,
                model=self._new_base_model(),
                created_on=self._n_samples_seen,
                drift_detector=self.drift_detector,
                warning_detector=self.warning_detector,
                is_background_learner=False,
                metric=self.metric,
            )
            for i in range(self.n_models)
        ]

    @abc.abstractmethod
    def _new_base_model(self):
        raise NotImplementedError

    def _set_max_features(self, n_features):
        if self.max_features == "sqrt":
            self.max_features = round(math.sqrt(n_features))
        elif self.max_features == "log2":
            self.max_features = round(math.log2(n_features))
        elif isinstance(self.max_features, int):
            # Consider 'max_features' features at each split.
            pass
        elif isinstance(self.max_features, float):
            # Consider 'max_features' as a percentage
            self.max_features = int(self.max_features * n_features)
        elif self.max_features is None:
            self.max_features = n_features
        else:
            raise AttributeError(
                f"Invalid max_features: {self.max_features}.\n"
                f"Valid options are: int [2, M], float (0., 1.],"
                f" {self._FEATURES_SQRT}, {self._FEATURES_LOG2}"
            )
        # Sanity checks
        # max_features is negative, use max_features + n
        if self.max_features < 0:
            self.max_features += n_features
        # max_features <= 0
        # (m can be negative if max_features is negative and abs(max_features) > n),
        # use max_features = 1
        if self.max_features <= 0:
            self.max_features = 1
        # max_features > n, then use n
        if self.max_features > n_features:
            self.max_features = n_features


class BaseTreeClassifier(HoeffdingTreeClassifier):
    """Adaptive Random Forest Hoeffding Tree Classifier.

    This is the base-estimator of the Adaptive Random Forest classifier.
    This variant of the Hoeffding Tree classifier includes the `max_features`
    parameter, which defines the number of randomly selected features to be
    considered at each split.

    """

    def __init__(
        self,
        max_features: int = 2,
        grace_period: int = 200,
        max_depth: int | None = None,
        split_criterion: str = "info_gain",
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "nba",
        nb_threshold: int = 0,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
        rng: random.Random | None = None,
    ):
        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            split_criterion=split_criterion,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            nb_threshold=nb_threshold,
            nominal_attributes=nominal_attributes,
            splitter=splitter,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self.max_features = max_features
        self.rng = rng

    def _new_leaf(self, initial_stats=None, parent=None):
        if initial_stats is None:
            initial_stats = {}

        if parent is None:
            depth = 0
        else:
            depth = parent.depth + 1

        if self._leaf_prediction == self._MAJORITY_CLASS:
            return RandomLeafMajorityClass(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                self.rng,
            )
        elif self._leaf_prediction == self._NAIVE_BAYES:
            return RandomLeafNaiveBayes(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                self.rng,
            )
        else:  # NAIVE BAYES ADAPTIVE (default)
            return RandomLeafNaiveBayesAdaptive(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                self.rng,
            )

    def new_instance(self):
        new_instance = self.clone()
        # Use existing rng to enforce a different model
        new_instance.rng = self.rng
        return new_instance


class BaseTreeRegressor(HoeffdingTreeRegressor):
    """ARF Hoeffding Tree regressor.

    This is the base-estimator of the Adaptive Random Forest regressor.
    This variant of the Hoeffding Tree regressor includes the `max_features`
    parameter, which defines the number of randomly selected features to be
    considered at each split.

    """

    def __init__(
        self,
        max_features: int = 2,
        grace_period: int = 200,
        max_depth: int | None = None,
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "adaptive",
        leaf_model: base.Regressor | None = None,
        model_selector_decay: float = 0.95,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        min_samples_split: int = 5,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
        rng: random.Random | None = None,
    ):
        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            leaf_model=leaf_model,
            model_selector_decay=model_selector_decay,
            nominal_attributes=nominal_attributes,
            splitter=splitter,
            min_samples_split=min_samples_split,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self.max_features = max_features
        self.rng = rng

    def _new_leaf(self, initial_stats=None, parent=None):  # noqa
        """Create a new learning node.

        The type of learning node depends on the tree configuration.
        """

        if parent is not None:
            depth = parent.depth + 1
        else:
            depth = 0

        leaf_model = None
        if self.leaf_prediction in {self._MODEL, self._ADAPTIVE}:
            if parent is None:
                leaf_model = copy.deepcopy(self.leaf_model)
            else:
                try:
                    leaf_model = copy.deepcopy(parent._leaf_model)  # noqa
                except AttributeError:
                    leaf_model = copy.deepcopy(self.leaf_model)

        if self.leaf_prediction == self._TARGET_MEAN:
            return RandomLeafMean(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                self.rng,
            )
        elif self.leaf_prediction == self._MODEL:
            return RandomLeafModel(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                self.rng,
                leaf_model=leaf_model,
            )
        else:  # adaptive learning node
            new_adaptive = RandomLeafAdaptive(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                self.rng,
                leaf_model=leaf_model,
            )
            if parent is not None and isinstance(parent, RandomLeafAdaptive):
                new_adaptive._fmse_mean = parent._fmse_mean  # noqa
                new_adaptive._fmse_model = parent._fmse_model  # noqa

            return new_adaptive

    def new_instance(self):
        new_instance = self.clone()
        # Use existing rng to enforce a different model
        new_instance.rng = self.rng
        return new_instance


class ARFClassifier(BaseForest, base.Classifier):
    """Adaptive Random Forest classifier.

    The 3 most important aspects of Adaptive Random Forest [^1] are:

    1. inducing diversity through re-sampling

    2. inducing diversity through randomly selecting subsets of features for
       node splits

    3. drift detectors per base tree, which cause selective resets in response
       to drifts

    It also allows training background trees, which start training if a
    warning is detected and replace the active tree if the warning escalates
    to a drift.

    Parameters
    ----------
    n_models
        Number of trees in the ensemble.
    max_features
        Max number of attributes for each node split.<br/>
        - If `int`, then consider `max_features` at each split.<br/>
        - If `float`, then `max_features` is a percentage and
          `int(max_features * n_features)` features are considered per split.<br/>
        - If "sqrt", then `max_features=sqrt(n_features)`.<br/>
        - If "log2", then `max_features=log2(n_features)`.<br/>
        - If None, then ``max_features=n_features``.
    lambda_value
        The lambda value for bagging (lambda=6 corresponds to Leveraging Bagging).
    metric
        Metric used to track trees performance within the ensemble.
        Defaults to `metrics.Accuracy()`.
    disable_weighted_vote
        If `True`, disables the weighted vote prediction.
    drift_detector
        Drift Detection method. Set to None to disable Drift detection.
        Defaults to `drift.ADWIN(delta=0.001)`.
    warning_detector
        Warning Detection method. Set to None to disable warning detection.
        Defaults to `drift.ADWIN(delta=0.01)`.
    grace_period
        [*Tree parameter*] Number of instances a leaf should observe between
        split attempts.
    max_depth
        [*Tree parameter*] The maximum depth a tree can reach. If `None`, the
        tree will grow indefinitely.
    split_criterion
        [*Tree parameter*] Split criterion to use.<br/>
        - 'gini' - Gini<br/>
        - 'info_gain' - Information Gain<br/>
        - 'hellinger' - Hellinger Distance
    delta
        [*Tree parameter*] Allowed error in split decision, a value closer to 0
        takes longer to decide.
    tau
        [*Tree parameter*] Threshold below which a split will be forced to break
        ties.
    leaf_prediction
        [*Tree parameter*] Prediction mechanism used at leafs.<br/>
        - 'mc' - Majority Class<br/>
        - 'nb' - Naive Bayes<br/>
        - 'nba' - Naive Bayes Adaptive
    nb_threshold
        [*Tree parameter*] Number of instances a leaf should observe before
        allowing Naive Bayes.
    nominal_attributes
        [*Tree parameter*] List of Nominal attributes. If empty, then assume that
        all attributes are numerical.
    splitter
        [*Tree parameter*] The Splitter or Attribute Observer (AO) used to monitor the class
        statistics of numeric features and perform splits. Splitters are available in the
        `tree.splitter` module. Different splitters are available for classification and
        regression tasks. Classification and regression splitters can be distinguished by their
        property `is_target_class`. This is an advanced option. Special care must be taken when
        choosing different splitters. By default, `tree.splitter.GaussianSplitter` is used
        if `splitter` is `None`.
    binary_split
        [*Tree parameter*] If True, only allow binary splits.
    max_size
        [*Tree parameter*] Maximum memory (MB) consumed by the tree.
    memory_estimate_period
        [*Tree parameter*] Number of instances between memory consumption checks.
    stop_mem_management
        [*Tree parameter*] If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        [*Tree parameter*] If True, disable poor attributes to reduce memory usage.
    merit_preprune
        [*Tree parameter*] If True, enable merit-based tree pre-pruning.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> from river import evaluate
    >>> from river import forest
    >>> from river import metrics
    >>> from river.datasets import synth

    >>> dataset = synth.ConceptDriftStream(
    ...     seed=42,
    ...     position=500,
    ...     width=40
    ... ).take(1000)

    >>> model = forest.ARFClassifier(seed=8, leaf_prediction="mc")

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 71.07%

    References
    ----------
    [^1]: Heitor Murilo Gomes, Albert Bifet, Jesse Read, Jean Paul Barddal,
         Fabricio Enembreck, Bernhard Pfharinger, Geoff Holmes, Talel Abdessalem.
         Adaptive random forests for evolving data stream classification.
         In Machine Learning, DOI: 10.1007/s10994-017-5642-8, Springer, 2017.

    """

    def __init__(
        self,
        n_models: int = 10,
        max_features: bool | str | int = "sqrt",
        lambda_value: int = 6,
        metric: metrics.base.MultiClassMetric | None = None,
        disable_weighted_vote=False,
        drift_detector: base.DriftDetector | None = None,
        warning_detector: base.DriftDetector | None = None,
        # Tree parameters
        grace_period: int = 50,
        max_depth: int | None = None,
        split_criterion: str = "info_gain",
        delta: float = 0.01,
        tau: float = 0.05,
        leaf_prediction: str = "nba",
        nb_threshold: int = 0,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 2_000_000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
        seed: int | None = None,
    ):
        super().__init__(
            n_models=n_models,
            max_features=max_features,
            lambda_value=lambda_value,
            metric=metric or metrics.Accuracy(),
            disable_weighted_vote=disable_weighted_vote,
            drift_detector=drift_detector or ADWIN(delta=0.001),
            warning_detector=warning_detector or ADWIN(delta=0.01),
            seed=seed,
        )

        self._n_samples_seen = 0
        self._base_member_class = ForestMemberClassifier  # type: ignore

        # Tree parameters
        self.grace_period = grace_period
        self.max_depth = max_depth
        self.split_criterion = split_criterion
        self.delta = delta
        self.tau = tau
        self.leaf_prediction = leaf_prediction
        self.nb_threshold = nb_threshold
        self.nominal_attributes = nominal_attributes
        self.splitter = splitter
        self.binary_split = binary_split
        self.max_size = max_size
        self.memory_estimate_period = memory_estimate_period
        self.stop_mem_management = stop_mem_management
        self.remove_poor_attrs = remove_poor_attrs
        self.merit_preprune = merit_preprune

    @property
    def _mutable_attributes(self):
        return {
            "max_features",
            "lambda_value",
            "grace_period",
            "delta",
            "tau",
        }

    @property
    def _multiclass(self):
        return True

    def predict_proba_one(self, x: dict) -> dict[base.typing.ClfTarget, float]:
        y_pred: typing.Counter = collections.Counter()

        if len(self) == 0:
            self._init_ensemble(sorted(x.keys()))
            return y_pred  # type: ignore

        for model in self:
            y_proba_temp = model.predict_proba_one(x)
            metric_value = model.metric.get()
            if not self.disable_weighted_vote and metric_value > 0.0:
                y_proba_temp = {k: val * metric_value for k, val in y_proba_temp.items()}
            y_pred.update(y_proba_temp)

        total = sum(y_pred.values())
        if total > 0:
            return {label: proba / total for label, proba in y_pred.items()}
        return y_pred  # type: ignore

    def _new_base_model(self):
        return BaseTreeClassifier(
            max_features=self.max_features,
            grace_period=self.grace_period,
            split_criterion=self.split_criterion,
            delta=self.delta,
            tau=self.tau,
            leaf_prediction=self.leaf_prediction,
            nb_threshold=self.nb_threshold,
            nominal_attributes=self.nominal_attributes,
            splitter=self.splitter,
            max_depth=self.max_depth,
            binary_split=self.binary_split,
            max_size=self.max_size,
            memory_estimate_period=self.memory_estimate_period,
            stop_mem_management=self.stop_mem_management,
            remove_poor_attrs=self.remove_poor_attrs,
            merit_preprune=self.merit_preprune,
            rng=self._rng,
        )


class ARFRegressor(BaseForest, base.Regressor):
    r"""Adaptive Random Forest regressor.

    The 3 most important aspects of Adaptive Random Forest [^1] are:

    1. inducing diversity through re-sampling

    2. inducing diversity through randomly selecting subsets of features for
       node splits

    3. drift detectors per base tree, which cause selective resets in response
       to drifts

    Notice that this implementation is slightly different from the original
    algorithm proposed in [^2]. The `HoeffdingTreeRegressor` is used as base
    learner, instead of `FIMT-DD`. It also adds a new strategy to monitor the
    predictions and check for concept drifts. The deviations of the predictions
    to the target are monitored and normalized in the [0, 1] range to fulfill ADWIN's
    requirements. We assume that the data subjected to the normalization follows
    a normal distribution, and thus, lies within the interval of the mean $\pm3\sigma$.

    Parameters
    ----------
    n_models
        Number of trees in the ensemble.
    max_features
        Max number of attributes for each node split.<br/>
        - If `int`, then consider `max_features` at each split.<br/>
        - If `float`, then `max_features` is a percentage and
          `int(max_features * n_features)` features are considered per split.<br/>
        - If "sqrt", then `max_features=sqrt(n_features)`.<br/>
        - If "log2", then `max_features=log2(n_features)`.<br/>
        - If None, then ``max_features=n_features``.
    lambda_value
        The lambda value for bagging (lambda=6 corresponds to Leveraging Bagging).
    metric
        Metric used to track trees performance within the ensemble. Depending,
        on the configuration, this metric is also used to weight predictions
        from the members of the ensemble. Defaults to `metrics.MSE()`.
    aggregation_method
        The method to use to aggregate predictions in the ensemble.<br/>
        - 'mean'<br/>
        - 'median' - If selected will disable the weighted vote.
    disable_weighted_vote
        If `True`, disables the weighted vote prediction, i.e. does not assign
        weights to individual tree's predictions and uses the arithmetic mean
        instead. Otherwise will use the `metric` value to weight predictions.
    drift_detector
        Drift Detection method. Set to None to disable Drift detection.
        Defaults to `drift.ADWIN(0.001)`.
    warning_detector
        Warning Detection method. Set to None to disable warning detection.
        Defaults to `drift.ADWIN(0.01)`.
    grace_period
        [*Tree parameter*] Number of instances a leaf should observe between
        split attempts.
    max_depth
        [*Tree parameter*] The maximum depth a tree can reach. If `None`, the
        tree will grow indefinitely.
    delta
        [*Tree parameter*] Allowed error in split decision, a value closer to 0
        takes longer to decide.
    tau
        [*Tree parameter*] Threshold below which a split will be forced to break
        ties.
    leaf_prediction
        [*Tree parameter*] Prediction mechanism used at leaves.</br>
        - 'mean' - Target mean</br>
        - 'model' - Uses the model defined in `leaf_model`</br>
        - 'adaptive' - Chooses between 'mean' and 'model' dynamically</br>
    leaf_model
        [*Tree parameter*] The regression model used to provide responses if
        `leaf_prediction='model'`. If not provided, an instance of
        `river.linear_model.LinearRegression` with the default hyperparameters
         is used.
    model_selector_decay
        [*Tree parameter*] The exponential decaying factor applied to the learning models'
        squared errors, that are monitored if `leaf_prediction='adaptive'`. Must be
        between `0` and `1`. The closer to `1`, the more importance is going to
        be given to past observations. On the other hand, if its value
        approaches `0`, the recent observed errors are going to have more
        influence on the final decision.
    nominal_attributes
        [*Tree parameter*] List of Nominal attributes. If empty, then assume that
        all attributes are numerical.
    splitter
        [*Tree parameter*] The Splitter or Attribute Observer (AO) used to monitor the class
        statistics of numeric features and perform splits. Splitters are available in the
        `tree.splitter` module. Different splitters are available for classification and
        regression tasks. Classification and regression splitters can be distinguished by their
        property `is_target_class`. This is an advanced option. Special care must be taken when
        choosing different splitters.By default, `tree.splitter.EBSTSplitter` is used if
        `splitter` is `None`.
    min_samples_split
        [*Tree parameter*] The minimum number of samples every branch resulting from a split
        candidate must have to be considered valid.
    binary_split
        [*Tree parameter*] If True, only allow binary splits.
    max_size
        [*Tree parameter*] Maximum memory (MB) consumed by the tree.
    memory_estimate_period
        [*Tree parameter*] Number of instances between memory consumption checks.
    stop_mem_management
        [*Tree parameter*] If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        [*Tree parameter*] If True, disable poor attributes to reduce memory usage.
    merit_preprune
        [*Tree parameter*] If True, enable merit-based tree pre-pruning.
    seed
        Random seed for reproducibility.

    References
    ----------
    [^1]: Gomes, H.M., Bifet, A., Read, J., Barddal, J.P., Enembreck, F.,
          Pfharinger, B., Holmes, G. and Abdessalem, T., 2017. Adaptive random
          forests for evolving data stream classification. Machine Learning,
          106(9-10), pp.1469-1495.

    [^2]: Gomes, H.M., Barddal, J.P., Boiko, L.E., Bifet, A., 2018.
          Adaptive random forests for data stream regression. ESANN 2018.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import forest
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     forest.ARFRegressor(seed=42)
    ... )

    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 0.800649

    """

    _MEAN = "mean"
    _MEDIAN = "median"
    _VALID_AGGREGATION_METHOD = [_MEAN, _MEDIAN]

    def __init__(
        self,
        # Forest parameters
        n_models: int = 10,
        max_features="sqrt",
        aggregation_method: str = "median",
        lambda_value: int = 6,
        metric: metrics.base.RegressionMetric | None = None,
        disable_weighted_vote=True,
        drift_detector: base.DriftDetector | None = None,
        warning_detector: base.DriftDetector | None = None,
        # Tree parameters
        grace_period: int = 50,
        max_depth: int | None = None,
        delta: float = 0.01,
        tau: float = 0.05,
        leaf_prediction: str = "adaptive",
        leaf_model: base.Regressor | None = None,
        model_selector_decay: float = 0.95,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        min_samples_split: int = 5,
        binary_split: bool = False,
        max_size: float = 500.0,
        memory_estimate_period: int = 2_000_000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
        seed: int | None = None,
    ):
        super().__init__(
            n_models=n_models,
            max_features=max_features,
            lambda_value=lambda_value,
            metric=metric or metrics.MSE(),
            disable_weighted_vote=disable_weighted_vote,
            drift_detector=drift_detector or ADWIN(0.001),
            warning_detector=warning_detector or ADWIN(0.01),
            seed=seed,
        )

        self._n_samples_seen = 0
        self._base_member_class = ForestMemberRegressor  # type: ignore

        # Tree parameters
        self.grace_period = grace_period
        self.max_depth = max_depth
        self.delta = delta
        self.tau = tau
        self.leaf_prediction = leaf_prediction
        self.leaf_model = leaf_model
        self.model_selector_decay = model_selector_decay
        self.nominal_attributes = nominal_attributes
        self.splitter = splitter
        self.min_samples_split = min_samples_split
        self.binary_split = binary_split
        self.max_size = max_size
        self.memory_estimate_period = memory_estimate_period
        self.stop_mem_management = stop_mem_management
        self.remove_poor_attrs = remove_poor_attrs
        self.merit_preprune = merit_preprune

        if aggregation_method in self._VALID_AGGREGATION_METHOD:
            self.aggregation_method = aggregation_method
        else:
            raise ValueError(
                f"Invalid aggregation_method: {aggregation_method}.\n"
                f"Valid values are: {self._VALID_AGGREGATION_METHOD}"
            )

    @property
    def _mutable_attributes(self):
        return {
            "max_features",
            "aggregation_method",
            "lambda_value",
            "grace_period",
            "delta",
            "tau",
            "model_selector_decay",
        }

    def predict_one(self, x: dict) -> base.typing.RegTarget:
        if len(self) == 0:
            self._init_ensemble(sorted(x.keys()))
            return 0.0  # type: ignore

        y_pred = np.zeros(self.n_models)

        if not self.disable_weighted_vote and self.aggregation_method != self._MEDIAN:
            weights = np.zeros(self.n_models)
            sum_weights = 0.0
            for idx, model in enumerate(self):
                y_pred[idx] = model.predict_one(x)
                weights[idx] = model.metric.get()
                sum_weights += weights[idx]

            if sum_weights != 0:
                # The higher the error, the worse is the tree
                weights = sum_weights - weights
                # Normalize weights to sum up to 1
                weights /= weights.sum()
                y_pred *= weights
        else:
            for idx, model in enumerate(self):
                y_pred[idx] = model.predict_one(x)

        if self.aggregation_method == self._MEAN:
            y_pred = y_pred.mean()
        else:
            y_pred = float(np.median(y_pred))  # type: ignore

        return float(y_pred)  # type: ignore

    def _new_base_model(self):
        return BaseTreeRegressor(
            max_features=self.max_features,
            grace_period=self.grace_period,
            max_depth=self.max_depth,
            delta=self.delta,
            tau=self.tau,
            leaf_prediction=self.leaf_prediction,
            leaf_model=self.leaf_model,
            model_selector_decay=self.model_selector_decay,
            nominal_attributes=self.nominal_attributes,
            splitter=self.splitter,
            binary_split=self.binary_split,
            max_size=self.max_size,
            memory_estimate_period=self.memory_estimate_period,
            stop_mem_management=self.stop_mem_management,
            remove_poor_attrs=self.remove_poor_attrs,
            merit_preprune=self.merit_preprune,
            rng=self._rng,
        )

    @property
    def valid_aggregation_method(self):
        """Valid aggregation_method values."""
        return self._VALID_AGGREGATION_METHOD


class BaseForestMember:
    """Base forest member class.

    This class represents a tree member of the forest. It includes a
    base tree model, the background learner, drift detectors and performance
    tracking parameters.

    The main purpose of this class is to train the foreground model.
    Optionally, it monitors drift detection. Depending on the configuration,
    if drift is detected then the foreground model is reset or replaced by a
    background model.

    Parameters
    ----------
    index_original
        Tree index within the ensemble.
    model
        Tree learner.
    created_on
        Number of instances seen by the tree.
    drift_detector
        Drift Detection method.
    warning_detector
        Warning Detection method.
    is_background_learner
        True if the tree is a background learner.
    metric
        Metric to track performance.

    """

    def __init__(
        self,
        index_original: int,
        model: BaseTreeClassifier | BaseTreeRegressor,
        created_on: int,
        drift_detector: base.DriftDetector,
        warning_detector: base.DriftDetector,
        is_background_learner,
        metric: metrics.base.MultiClassMetric | metrics.base.RegressionMetric,
    ):
        self.index_original = index_original
        self.model = model
        self.created_on = created_on
        self.is_background_learner = is_background_learner
        self.metric = metric.clone()
        self.background_learner = None

        # Drift and warning detection
        self.last_drift_on = 0
        self.last_warning_on = 0
        self.n_drifts_detected = 0
        self.n_warnings_detected = 0

        # Initialize drift and warning detectors
        if drift_detector is not None:
            self._use_drift_detector = True
            self.drift_detector = drift_detector.clone()
        else:
            self._use_drift_detector = False
            self.drift_detector = None

        if warning_detector is not None:
            self._use_background_learner = True
            self.warning_detector = warning_detector.clone()
        else:
            self._use_background_learner = False
            self.warning_detector = None

    def reset(self, n_samples_seen):
        if self._use_background_learner and self.background_learner is not None:
            # Replace foreground model with background model
            self.model = self.background_learner.model
            self.warning_detector = self.background_learner.warning_detector
            self.drift_detector = self.background_learner.drift_detector
            self.metric = self.background_learner.metric
            self.created_on = self.background_learner.created_on
            self.background_learner = None
        else:
            # Reset model
            self.model = self.model.new_instance()
            self.metric = self.metric.clone()
            self.created_on = n_samples_seen
            self.drift_detector = self.drift_detector.clone()

    def learn_one(self, x: dict, y: base.typing.Target, *, sample_weight: int, n_samples_seen: int):
        self.model.learn_one(x, y, sample_weight=sample_weight)

        if self.background_learner:
            # Train the background learner
            self.background_learner.model.learn_one(x=x, y=y, sample_weight=sample_weight)

        if self._use_drift_detector and not self.is_background_learner:
            drift_detector_input = self._drift_detector_input(
                y_true=y, y_pred=self.model.predict_one(x)  # type: ignore
            )

            # Check for warning only if use_background_learner is set
            if self._use_background_learner:
                self.warning_detector.update(drift_detector_input)
                # Check if there was a (warning) change
                if self.warning_detector.drift_detected:
                    self.last_warning_on = n_samples_seen
                    self.n_warnings_detected += 1
                    # Create a new background learner object
                    self.background_learner = self.__class__(  # type: ignore
                        index_original=self.index_original,
                        model=self.model.new_instance(),
                        created_on=n_samples_seen,
                        drift_detector=self.drift_detector,
                        warning_detector=self.warning_detector,
                        is_background_learner=True,
                        metric=self.metric,
                    )
                    # Reset the warning detector for the current object
                    self.warning_detector = self.warning_detector.clone()

            # Update the drift detector
            self.drift_detector.update(drift_detector_input)

            # Check if there was a change
            if self.drift_detector.drift_detected:
                self.last_drift_on = n_samples_seen
                self.n_drifts_detected += 1
                self.reset(n_samples_seen)

    @abc.abstractmethod
    def _drift_detector_input(
        self,
        y_true: base.typing.ClfTarget | base.typing.RegTarget,
        y_pred: base.typing.ClfTarget | base.typing.RegTarget,
    ):
        raise NotImplementedError


class ForestMemberClassifier(BaseForestMember, base.Classifier):  # type: ignore
    """Forest member class for classification"""

    def __init__(
        self,
        index_original: int,
        model: BaseTreeClassifier,
        created_on: int,
        drift_detector: base.DriftDetector,
        warning_detector: base.DriftDetector,
        is_background_learner,
        metric: metrics.base.MultiClassMetric,
    ):
        super().__init__(
            index_original=index_original,
            model=model,
            created_on=created_on,
            drift_detector=drift_detector,
            warning_detector=warning_detector,
            is_background_learner=is_background_learner,
            metric=metric,
        )

    def _drift_detector_input(  # type: ignore
        self, y_true: base.typing.ClfTarget, y_pred: base.typing.ClfTarget
    ):
        return int(not y_true == y_pred)  # Not correctly_classifies

    def predict_one(self, x):
        return self.model.predict_one(x)

    def predict_proba_one(self, x):
        return self.model.predict_proba_one(x)


class ForestMemberRegressor(BaseForestMember, base.Regressor):  # type: ignore
    """Forest member class for regression"""

    def __init__(
        self,
        index_original: int,
        model: BaseTreeRegressor,
        created_on: int,
        drift_detector: base.DriftDetector,
        warning_detector: base.DriftDetector,
        is_background_learner,
        metric: metrics.base.RegressionMetric,
    ):
        super().__init__(
            index_original=index_original,
            model=model,
            created_on=created_on,
            drift_detector=drift_detector,
            warning_detector=warning_detector,
            is_background_learner=is_background_learner,
            metric=metric,
        )
        self._var = stats.Var()  # Used to track drift

    def _drift_detector_input(self, y_true: float, y_pred: float):  # type: ignore
        drift_input = y_true - y_pred
        self._var.update(drift_input)

        if self._var.mean.n == 1:
            return 0.5  # The expected error is the normalized mean error

        sd = math.sqrt(self._var.get())

        # We assume the error follows a normal distribution -> (empirical rule)
        # 99.73% of the values lie  between [mean - 3*sd, mean + 3*sd]. We
        # assume this range for the normalized data. Hence, we can apply the
        # min-max norm to cope with  ADWIN's requirements
        return (drift_input + 3 * sd) / (6 * sd) if sd > 0 else 0.5

    def reset(self, n_samples_seen):
        super().reset(n_samples_seen)
        # Reset the stats for the drift detector
        self._var = stats.Var()

    def predict_one(self, x):
        return self.model.predict_one(x)


# river/forest/aggregated_mondrian_forest.py

from __future__ import annotations

import abc
import random

from river import base
from river.tree.mondrian import MondrianTreeClassifier


class AMFLearner(base.Ensemble, abc.ABC):
    """Base class for Aggregated Mondrian Forest classifier and regressors for online learning.

    Parameters
    ----------
    n_estimators
        The number of trees in the forest.
    step
         Step-size for the aggregation weights.
    loss
        The loss used for the computation of the aggregation weights.
    use_aggregation
        Controls if aggregation is used in the trees. It is highly recommended to
        leave it as `True`.
    split_pure
        Controls if nodes that contains only sample of the same class should be
        split ("pure" nodes). Default is `False`, namely pure nodes are not split,
        but `True` can be sometimes better.
    seed
        Random seed for reproducibility.

    Note
    ----
    This class is not intended for end users but for development only.

    """

    def __init__(
        self,
        n_estimators: int = 10,
        step: float = 0.1,
        loss: str = "log",
        use_aggregation: bool = True,
        split_pure: bool = False,
        seed: int | None = None,
    ):
        super().__init__([])  # type: ignore

        self.n_estimators = n_estimators
        self.step = step
        self.loss = loss
        self.use_aggregation = use_aggregation
        self.split_pure = split_pure
        self.seed = seed

        self._rng = random.Random(self.seed)

    @property
    def _is_initialized(self) -> bool:
        """Indicate whether the model has been trained at least once before."""
        return len(self) > 0

    @abc.abstractmethod
    def _initialize_trees(self):
        """Initialize the forest members."""

    @property
    def _min_number_of_models(self):
        return 0


class AMFClassifier(AMFLearner, base.Classifier):
    """Aggregated Mondrian Forest classifier for online learning.

    This implementation is truly online, in the sense that a single pass is performed, and that
    predictions can be produced anytime.

    Each node in a tree predicts according to the distribution of the labels
    it contains. This distribution is regularized using a "Jeffreys" prior
    with parameter `dirichlet`. For each class with `count` labels in the
    node and `n_samples` samples in it, the prediction of a node is given by

    $\\frac{count + dirichlet}{n_{samples} + dirichlet \\times n_{classes}}$.

    The prediction for a sample is computed as the aggregated predictions of all the
    subtrees along the path leading to the leaf node containing the sample. The
    aggregation weights are exponential weights with learning rate `step` and log-loss
    when `use_aggregation` is `True`.

    This computation is performed exactly thanks to a context tree weighting algorithm.
    More details can be found in the paper cited in the references below.

    The final predictions are the average class probabilities predicted by each of the
    `n_estimators` trees in the forest.

    Parameters
    ----------
    n_estimators
        The number of trees in the forest.
    step
        Step-size for the aggregation weights. Default is 1 for classification with
        the log-loss, which is usually the best choice.
    use_aggregation
        Controls if aggregation is used in the trees. It is highly recommended to
        leave it as `True`.
    dirichlet
        Regularization level of the class frequencies used for predictions in each node. A rule of
        thumb is to set this to `1 / n_classes`, where `n_classes` is the expected number of
        classes which might appear. Default is `dirichlet = 0.5`, which works well for binary
        classification problems.
    split_pure
        Controls if nodes that contains only sample of the same class should be
        split ("pure" nodes). Default is `False`, namely pure nodes are not split,
        but `True` can be sometimes better.
    seed
        Random seed for reproducibility.

    Notes
    -----
    Only log_loss used for the computation of the aggregation weights is supported for now, namely
    the log-loss for multi-class classification.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import forest
    >>> from river import metrics

    >>> dataset = datasets.Bananas().take(500)

    >>> model = forest.AMFClassifier(
    ...     n_estimators=10,
    ...     use_aggregation=True,
    ...     dirichlet=0.5,
    ...     seed=1
    ... )

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 84.97%

    References
    ----------
    J. Mourtada, S. Gaiffas and E. Scornet, *AMF: Aggregated Mondrian Forests for Online Learning*, arXiv:1906.10529, 2019.

    """

    def __init__(
        self,
        n_estimators: int = 10,
        step: float = 1.0,
        use_aggregation: bool = True,
        dirichlet: float = 0.5,
        split_pure: bool = False,
        seed: int | None = None,
    ):
        super().__init__(
            n_estimators=n_estimators,
            step=step,
            loss="log",
            use_aggregation=use_aggregation,
            split_pure=split_pure,
            seed=seed,
        )
        self.dirichlet = dirichlet
        # memory of the classes
        self._classes: set[base.typing.ClfTarget] = set()

    def _initialize_trees(self):
        self.data: list[MondrianTreeClassifier] = []
        for _ in range(self.n_estimators):
            tree = MondrianTreeClassifier(
                self.step,
                self.use_aggregation,
                self.dirichlet,
                self.split_pure,
                iteration=0,
                # We don't want to have the same stochastic scheme for each tree, or it'll break the randomness
                # Hence we introduce a new seed for each, that is derived of the given seed by a deterministic process
                seed=self._rng.randint(0, 9999999),
            )
            self.data.append(tree)

    def learn_one(self, x, y):
        # Updating the previously seen classes with the new sample
        self._classes.add(y)

        # Checking if the forest has been created
        if not self._is_initialized:
            self._initialize_trees()

        # we fit all the trees using the new sample
        for tree in self:
            tree.learn_one(x, y)

        return self

    def predict_proba_one(self, x):
        # Checking that the model has been trained once at least
        # Otherwise return the default empty dict
        if not self._is_initialized:
            return {}

        # initialize the scores
        scores = {c: 0 for c in self._classes}

        # Simply computes the prediction for each tree and average it
        for tree in self:
            predictions = tree.predict_proba_one(x)
            for c in self._classes:
                scores[c] += predictions[c] / self.n_estimators

        return scores

    @property
    def _multiclass(self):
        return True


# river/forest/online_extra_trees.py

from __future__ import annotations

import abc
import collections
import math
import numbers
import random
import sys

from river import base, drift, metrics, tree, utils
from river.tree.nodes.et_nodes import ETLeafAdaptive, ETLeafMean, ETLeafModel
from river.tree.splitter.random_splitter import RegRandomSplitter


class Sampler(base.Base, abc.ABC):
    @abc.abstractmethod
    def __call__(self, rate, rng) -> int:
        pass


class ConstantSampler(Sampler):
    def __call__(self, rate, rng):
        return 1


class BaggingSampler(Sampler):
    def __call__(self, rate, rng):
        return utils.random.poisson(rate, rng)


class SubBaggingSampler(Sampler):
    def __call__(self, rate, rng) -> int:
        return 1 if rng.random() <= rate else 0


class ExtraTrees(base.Ensemble, metaclass=abc.ABCMeta):
    _FEATURES_SQRT = "sqrt"
    _FEATURES_LOG2 = "log2"
    _FEATURES_RANDOM = "random"

    _BAGGING = "bagging"
    _SUBBAGGING = "subbagging"

    _DETECTION_ALL = "all"
    _DETECTION_DROP = "drop"
    _DETECTION_OFF = "off"

    def __init__(
        self,
        n_models: int,
        max_features: bool | str | int,
        resampling_strategy: str | None,
        resampling_rate: int | float,
        detection_mode: str,
        warning_detector: base.DriftDetector | None,
        drift_detector: base.DriftDetector | None,
        max_depth: int | None,
        randomize_tree_depth: bool,
        track_metric: metrics.base.MultiClassMetric | metrics.base.RegressionMetric,
        disable_weighted_vote: bool,
        split_buffer_size: int,
        seed: int | None,
    ):
        self.data = []
        self.n_models = n_models
        self.max_features = max_features

        if resampling_strategy not in [None, self._BAGGING, self._SUBBAGGING]:
            raise ValueError(f"Invalid resampling strategy: {resampling_strategy}")
        self.resampling_strategy = resampling_strategy

        self.resampling_rate: int | float = 0
        if self.resampling_strategy is not None:
            if self.resampling_strategy == self._BAGGING:
                if resampling_rate < 1:
                    raise ValueError(
                        "'resampling_rate' must be an integer greater than or"
                        "equal to 1, when resample_strategy='bagging'."
                    )
                # Cast to integer (online bagging using poisson sampling)
                self.resampling_rate = int(resampling_rate)

            if self.resampling_strategy == self._SUBBAGGING:
                if not 0 < resampling_rate <= 1:
                    raise ValueError(
                        "resampling_rate must be a float in the interval (0, 1],"
                        "when resampling_strategy='subbagging'."
                    )
                self.resampling_rate = resampling_rate

        if detection_mode not in [
            self._DETECTION_ALL,
            self._DETECTION_DROP,
            self._DETECTION_OFF,
        ]:
            raise ValueError(
                f"Invalid drift detection mode. Valid values are: '{self._DETECTION_ALL}',"
                f" {self._DETECTION_DROP}, and '{self._DETECTION_OFF}'."
            )

        self.detection_mode = detection_mode
        self.warning_detector = warning_detector or drift.ADWIN(delta=0.01)
        self.drift_detector = drift_detector or drift.ADWIN(delta=0.001)

        self.max_depth = max_depth
        self.randomize_tree_depth = randomize_tree_depth
        self.track_metric = track_metric
        self.disable_weighted_vote = disable_weighted_vote
        self.split_buffer_size = split_buffer_size
        self.seed = seed

        # The predictive performance of each tree
        self._perfs: list = []
        # Keep a running estimate of the sum of performances
        self._perf_sum: float = 0

        # Number of times a tree will use each instance to learn from it
        self._weight_sampler = self.__weight_sampler_factory()

        # General statistics
        # Counter of the number of instances each ensemble member has processed (instance weights
        # are not accounted for, just the number of instances)
        self._sample_counter: collections.Counter = collections.Counter()
        # Total of samples processed by the Extra Trees ensemble
        self._total_instances: float = 0
        # Number of warnings triggered
        self._n_warnings: collections.Counter = collections.Counter()
        # Number of drifts detected
        self._n_drifts: collections.Counter = collections.Counter()
        # Number of tree swaps
        self._n_tree_swaps: collections.Counter = collections.Counter()

        self._background_trees: dict[int, tree.hoeffding_tree.HoeffdingTree] = {}
        # Initialize drift detectors and select the detection mode procedure
        if self.detection_mode == self._DETECTION_ALL:
            self._warn_detectors = {i: self.warning_detector.clone() for i in range(self.n_models)}
            self._drift_detectors = {i: self.drift_detector.clone() for i in range(self.n_models)}
        elif self.detection_mode == self._DETECTION_DROP:
            self._warn_detectors = {}
            self._drift_detectors = {i: self.drift_detector.clone() for i in range(self.n_models)}
        else:  # detection_mode: "off"
            self._warn_detectors = {}
            self._drift_detectors = {}
        self._detect = self.__detection_mode_factory()

        # Set the rng
        self._rng = random.Random(seed)

    @abc.abstractmethod
    def _new_member(self, max_features, max_depth, seed) -> base.Classifier | base.Regressor:
        pass

    @abc.abstractmethod
    def _drift_input(self, y, y_hat) -> int | float:
        pass

    def _calculate_tree_depth(self) -> int | float:
        max_depth = self.max_depth or math.inf
        if not self.randomize_tree_depth:
            return max_depth
        return self._rng.randint(1, max_depth if not math.isinf(max_depth) else 9999)  # type: ignore

    def _calculate_max_features(self, n_features) -> int:
        if self.max_features == self._FEATURES_RANDOM:
            # Generate a random integer
            return self._rng.randint(2, n_features)
        else:
            if self.max_features == self._FEATURES_SQRT:
                max_feat = round(math.sqrt(n_features))
            elif self.max_features == self._FEATURES_LOG2:
                max_feat = round(math.log2(n_features))
            elif isinstance(self.max_features, int):
                max_feat = n_features
            elif isinstance(self.max_features, float):
                # Consider 'max_features' as a percentage
                max_feat = int(self.max_features * n_features)
            elif self.max_features is None:
                max_feat = n_features
            else:
                raise AttributeError(
                    f"Invalid max_features: {self.max_features}.\n"
                    f"Valid options are: int [2, M], float (0., 1.],"
                    f" {self._FEATURES_SQRT}, {self._FEATURES_LOG2}"
                )

            # Sanity checks
            # max_feat is negative, use max_feat + n
            if max_feat < 0:
                max_feat += n_features
            # max_feat <= 0
            # (m can be negative if max_feat is negative and abs(max_feat) > n),
            # use max_features = 1
            if max_feat <= 0:
                max_feat = 1
            # max_feat > n, then use n
            if max_feat > n_features:
                max_feat = n_features

            return max_feat

    def _init_trees(self, n_features: int):
        for _ in range(self.n_models):
            self.data.append(
                self._new_member(
                    max_features=self._calculate_max_features(n_features),
                    max_depth=self._calculate_tree_depth(),
                    seed=self._rng.randint(0, sys.maxsize),  # randomly creates a new seed
                )
            )
            self._perfs.append(self.track_metric.clone())

    def __weight_sampler_factory(self) -> Sampler:
        if self.resampling_strategy == self._BAGGING:
            return BaggingSampler()
        elif self.resampling_strategy == self._SUBBAGGING:
            return SubBaggingSampler()
        else:
            return ConstantSampler()

    @staticmethod
    def _detection_mode_all(
        drift_detector: base.DriftDetector,
        warning_detector: base.DriftDetector,
        detector_input: numbers.Number,
    ) -> tuple[bool, bool]:
        in_warning = warning_detector.update(detector_input).drift_detected
        in_drift = drift_detector.update(detector_input).drift_detected

        return in_drift, in_warning

    @staticmethod
    def _detection_mode_drop(
        drift_detector: base.DriftDetector,
        warning_detector: base.DriftDetector,
        detector_input: numbers.Number,
    ) -> tuple[bool, bool]:
        in_drift = drift_detector.update(detector_input).drift_detected

        return in_drift, False

    @staticmethod
    def _detection_mode_off(
        drift_detector: base.DriftDetector,
        warning_detector: base.DriftDetector,
        detector_input: numbers.Number,
    ) -> tuple[bool, bool]:
        return False, False

    def __detection_mode_factory(self):
        if self.detection_mode == self._DETECTION_ALL:
            return self._detection_mode_all
        elif self.detection_mode == self._DETECTION_DROP:
            return self._detection_mode_drop
        else:
            return self._detection_mode_off

    def learn_one(self, x, y):
        if not self.models:
            self._init_trees(len(x))

        self._total_instances += 1
        trained = []
        for i, model in enumerate(self.models):
            y_hat = model.predict_one(x)
            in_drift, in_warning = self._detect(
                self._drift_detectors.get(i),
                self._warn_detectors.get(i),
                self._drift_input(y, y_hat),
            )

            if in_warning:
                self._background_trees[i] = self._new_member(
                    max_features=self._calculate_max_features(len(x)),
                    max_depth=self._calculate_tree_depth(),
                    seed=self._rng.randint(0, sys.maxsize),  # randomly creates a new seed
                )
                # Reset the warning detector
                self._warn_detectors[i] = self.warning_detector.clone()
                # Update statistics
                self._n_warnings.update([i])

            # Drift detected: time to change (swap or reset) the affected tree
            if in_drift:
                if i in self._background_trees:
                    self.data[i] = self._background_trees[i]
                    del self._background_trees[i]

                    self._n_tree_swaps.update([i])
                else:
                    self.data[i] = self._new_member(
                        max_features=self._calculate_max_features(len(x)),
                        max_depth=self._calculate_tree_depth(),
                        seed=self._rng.randint(0, sys.maxsize),  # randomly creates a new seed
                    )
                # Reset the drift detector
                self._drift_detectors[i] = self.drift_detector.clone()
                # Update statistics
                self._n_drifts.update([i])
                # Also reset tree's error estimates
                self._perf_sum -= self._perfs[i].get()
                self._perfs[i] = self.track_metric.clone()
                self._perf_sum += self._perfs[i].get()
                # And the number of observations of the new model
                self._sample_counter[i] = 0

            # Remove the old performance estimate
            self._perf_sum -= self._perfs[i].get()
            # Update metric
            self._perfs[i].update(y, y_hat)
            # Add the new performance estimate
            self._perf_sum += self._perfs[i].get()

            # Define the weight of the instance
            w = self._weight_sampler(self.resampling_rate, self._rng)
            if w == 0:  # Skip model update if w is zero
                continue

            model.learn_one(x, y, sample_weight=w)

            if i in self._background_trees:
                self._background_trees[i].learn_one(x, y, sample_weight=w)

            trained.append(i)

        # Increase by one the count of instances observed by each trained model
        self._sample_counter.update(trained)

        return self

    # Properties
    @property
    def n_warnings(self) -> collections.Counter:
        """The number of warnings detected per ensemble member."""
        return self._n_warnings

    @property
    def n_drifts(self) -> collections.Counter:
        """The number of concept drifts detected per ensemble member."""
        return self._n_drifts

    @property
    def n_tree_swaps(self) -> collections.Counter:
        """The number of performed alternate tree swaps.

        Not applicable if the warning detectors are disabled.
        """
        return self._n_tree_swaps

    @property
    def total_instances(self) -> float:
        """The total number of instances processed by the ensemble."""
        return self._total_instances

    @property
    def instances_per_tree(self) -> collections.Counter:
        """The number of instances processed by each one of the current forest members.

        Each time a concept drift is detected, the count corresponding to the affected tree is
        reset.
        """
        return self._sample_counter

    @classmethod
    def _unit_test_params(cls):
        yield {"n_models": 3}

    def _unit_test_skips(self):
        return {"check_shuffle_features_no_impact"}


class ETRegressor(tree.HoeffdingTreeRegressor):
    """Extra Tree regressor.

    This is the base-estimator of the Extra Trees regressor.
    This variant of the Hoeffding Tree regressor includes the `max_features` parameter,
    which defines the number of randomly selected features to be considered at each split.
    It also evaluates split candidates randomly.

    """

    def __init__(
        self,
        max_features,
        grace_period,
        max_depth,
        delta,
        tau,
        leaf_prediction,
        leaf_model,
        model_selector_decay,
        nominal_attributes,
        min_samples_split,
        binary_split,
        max_size,
        memory_estimate_period,
        stop_mem_management,
        remove_poor_attrs,
        merit_preprune,
        split_buffer_size,
        seed,
    ):
        self.max_features = max_features
        self.split_buffer_size = split_buffer_size
        self.seed = seed
        self._rng = random.Random(self.seed)

        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            leaf_model=leaf_model,
            model_selector_decay=model_selector_decay,
            nominal_attributes=nominal_attributes,
            splitter=RegRandomSplitter(
                seed=self._rng.randint(0, sys.maxsize),
                buffer_size=self.split_buffer_size,
            ),
            min_samples_split=min_samples_split,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

    def _new_learning_node(self, initial_stats=None, parent=None):  # noqa
        """Create a new learning node.
        The type of learning node depends on the tree configuration.
        """

        if parent is not None:
            depth = parent.depth + 1
        else:
            depth = 0

        # Generate a random seed for the new learning node
        seed = self._rng.randint(0, sys.maxsize)

        leaf_model = None
        if self.leaf_prediction in {self._MODEL, self._ADAPTIVE}:
            if parent is None:
                leaf_model = self.leaf_model.clone()
            else:
                leaf_model = parent._leaf_model.clone(include_attributes=True)  # noqa

        if self.leaf_prediction == self._TARGET_MEAN:
            return ETLeafMean(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                seed,
            )
        elif self.leaf_prediction == self._MODEL:
            return ETLeafModel(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                seed,
                leaf_model=leaf_model,
            )
        else:  # adaptive learning node
            new_adaptive = ETLeafAdaptive(
                initial_stats,
                depth,
                self.splitter,
                self.max_features,
                seed,
                leaf_model=leaf_model,
            )
            if parent is not None and isinstance(parent, ETLeafAdaptive):
                new_adaptive._fmse_mean = parent._fmse_mean  # noqa
                new_adaptive._fmse_model = parent._fmse_model  # noqa

            return new_adaptive


class OXTRegressor(ExtraTrees, base.Regressor):
    """Online Extra Trees regressor.

    The online Extra Trees[^1] ensemble takes some steps further into randomization when
    compared to Adaptive Random Forests (ARF). A subspace of the feature space is considered
    at each split attempt, as ARF does, and online bagging or subbagging can also be
    (optionally) used. Nonetheless, Extra Trees randomizes the split candidates evaluated by each
    leaf node (just a single split is tested by numerical feature, which brings significant
    speedups to the ensemble), and might also randomize the maximum depth of the forest members,
    as well as the size of the feature subspace processed by each of its trees' leaves.

    On the other hand, OXT suffers from a cold-start problem. As the splits are random, the
    predictive performance in small samples is usually worse than using a deterministic
    split approach, such as the one used by ARF.

    Parameters
    ----------
    n_models
        The number of trees in the ensemble.
    max_features
        Max number of attributes for each node split.</br>
        - If int, then consider `max_features` at each split.</br>
        - If float, then `max_features` is a percentage and `int(max_features * n_features)`
        features are considered per split.</br>
        - If "sqrt", then `max_features=sqrt(n_features)`.</br>
        - If "log2", then `max_features=log2(n_features)`.</br>
        - If "random", then `max_features` will assume a different random number in the interval
        `[2, n_features]` for each tree leaf.</br>
        - If None, then `max_features=n_features`.
    resampling_strategy
        The chosen instance resampling strategy:</br>
        - If `None`, no resampling will be done and the trees will process all instances.
        - If `'baggging'`, online bagging will be performed (sampling with replacement).
        - If `'subbagging'`, online subbagging will be performed (sampling without replacement).
    resampling_rate
        Only valid if `resampling_strategy` is not None. Controls the parameters of the resampling
        strategy.</br>.
        - If `resampling_strategy='bagging'`, must be an integer greater than or equal to 1 that
        parameterizes the poisson distribution used to simulate bagging in online learning
        settings. It acts as the lambda parameter of Oza Bagging and Leveraging Bagging.</br>
        - If `resampling_strategy='subbagging'`, must be a float in the interval $(0, 1]$ that
        controls the chance of each instance being used by a tree for learning.
    detection_mode
        The concept drift detection mode in which the forest operates. Valid values are:</br>
        - "all": creates both warning and concept drift detectors. If a warning is detected,
        an alternate tree starts being trained in the background. If the warning trigger escalates
        to a concept drift, the affected tree is replaced by the alternate tree.</br>
        - "drop": only the concept drift detectors are created. If a drift is detected, the
        affected tree is dropped and replaced by a new tree.</br>
        - "off": disables the concept drift adaptation capabilities. The forest will act as if
        the processed stream is stationary.
    warning_detector
        The detector that will be used to trigger concept drift warnings. Defaults to
        `drift.ADWIN(0.01)`.
    drift_detector
        The detector used to detect concept drifts. Defaults to `drift.ADWIN(0.001)`.
    max_depth
        The maximum depth the ensemble members might reach. If `None`, the trees will grow
        indefinitely.
    randomize_tree_depth
        Whether or not randomize the maximum depth of each tree in the ensemble. If `max_depth`
        is provided, it is going to act as an upper bound to generate the maximum depth for each
        tree.
    track_metric
        The performance metric used to weight predictions. Defaults to `metrics.MAE()`.
    disable_weighted_vote
        Defines whether or not to use predictions weighted by each trees' prediction performance.
    split_buffer_size
        Defines the size of the buffer used by the tree splitters when determining the feature
        range and a random split point in this interval.
    seed
        Random seed to support reproducibility.
    grace_period
        [*Tree parameter*] Number of instances a leaf should observe between
        split attempts.
    delta
        [*Tree parameter*] Allowed error in split decision, a value closer to 0
        takes longer to decide.
    tau
        [*Tree parameter*] Threshold below which a split will be forced to break
        ties.
    leaf_prediction
        [*Tree parameter*] Prediction mechanism used at leaves.</br>
        - 'mean' - Target mean</br>
        - 'model' - Uses the model defined in `leaf_model`</br>
        - 'adaptive' - Chooses between 'mean' and 'model' dynamically</br>
    leaf_model
        [*Tree parameter*] The regression model used to provide responses if
        `leaf_prediction='model'`. If not provided, an instance of
        `river.linear_model.LinearRegression` with the default hyperparameters
         is used.
    model_selector_decay
        [*Tree parameter*] The exponential decaying factor applied to the learning models'
        squared errors, that are monitored if `leaf_prediction='adaptive'`. Must be
        between `0` and `1`. The closer to `1`, the more importance is going to
        be given to past observations. On the other hand, if its value
        approaches `0`, the recent observed errors are going to have more
        influence on the final decision.
    nominal_attributes
        [*Tree parameter*] List of Nominal attributes. If empty, then assume that
        all attributes are numerical.
    min_samples_split
        [*Tree parameter*] The minimum number of samples every branch resulting from a split
        candidate must have to be considered valid.
    binary_split
        [*Tree parameter*] If True, only allow binary splits.
    max_size
        [*Tree parameter*] Maximum memory (MB) consumed by the tree.
    memory_estimate_period
        [*Tree parameter*] Number of instances between memory consumption checks.
    stop_mem_management
        [*Tree parameter*] If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        [*Tree parameter*] If True, disable poor attributes to reduce memory usage.
    merit_preprune
        [*Tree parameter*] If True, enable merit-based tree pre-pruning.

    Notes
    -----
    As the Online Extra Trees change the way in which Hoeffding Trees perform split attempts
    and monitor numerical input features, some of the parameters of the vanilla Hoeffding Tree
    algorithms are not available.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import forest

    >>> dataset = datasets.synth.Friedman(seed=42).take(5000)

    >>> model = forest.OXTRegressor(n_models=3, seed=42)

    >>> metric = metrics.RMSE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    RMSE: 3.127311

    References
    ----------
    [^1]: Mastelini, S. M., Nakano, F. K., Vens, C., & de Leon Ferreira, A. C. P. (2022).
    Online Extra Trees Regressor. IEEE Transactions on Neural Networks and Learning Systems.

    """

    def __init__(
        self,
        n_models: int = 10,
        max_features: bool | str | int = "sqrt",
        resampling_strategy: str | None = "subbagging",
        resampling_rate: int | float = 0.5,
        detection_mode: str = "all",
        warning_detector: base.DriftDetector | None = None,
        drift_detector: base.DriftDetector | None = None,
        max_depth: int | None = None,
        randomize_tree_depth: bool = False,
        track_metric: metrics.base.RegressionMetric | None = None,
        disable_weighted_vote: bool = True,
        split_buffer_size: int = 5,
        seed: int | None = None,
        grace_period: int = 50,
        delta: float = 0.01,
        tau: float = 0.05,
        leaf_prediction: str = "adaptive",
        leaf_model: base.Regressor | None = None,
        model_selector_decay: float = 0.95,
        nominal_attributes: list | None = None,
        min_samples_split: int = 5,
        binary_split: bool = False,
        max_size: int = 500,
        memory_estimate_period: int = 2_000_000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
    ):
        super().__init__(
            n_models=n_models,
            max_features=max_features,
            resampling_strategy=resampling_strategy,
            resampling_rate=resampling_rate,
            detection_mode=detection_mode,
            warning_detector=warning_detector,
            drift_detector=drift_detector,
            max_depth=max_depth,
            randomize_tree_depth=randomize_tree_depth,
            track_metric=track_metric or metrics.MAE(),
            disable_weighted_vote=disable_weighted_vote,
            split_buffer_size=split_buffer_size,
            seed=seed,
        )

        # Tree parameters
        self.grace_period = grace_period
        self.delta = delta
        self.tau = tau
        self.leaf_prediction = leaf_prediction
        self.leaf_model = leaf_model
        self.model_selector_decay = model_selector_decay
        self.nominal_attributes = nominal_attributes
        self.min_samples_split = min_samples_split
        self.binary_split = binary_split
        self.max_size = max_size
        self.memory_estimate_period = memory_estimate_period
        self.stop_mem_management = stop_mem_management
        self.remove_poor_attrs = remove_poor_attrs
        self.merit_preprune = merit_preprune

    def _new_member(self, max_features, max_depth, seed):
        return ETRegressor(
            max_features=max_features,
            grace_period=self.grace_period,
            max_depth=max_depth,
            delta=self.delta,
            tau=self.tau,
            leaf_prediction=self.leaf_prediction,
            leaf_model=self.leaf_model,
            model_selector_decay=self.model_selector_decay,
            nominal_attributes=self.nominal_attributes,
            min_samples_split=self.min_samples_split,
            binary_split=self.binary_split,
            max_size=self.max_size,
            memory_estimate_period=self.memory_estimate_period,
            stop_mem_management=self.stop_mem_management,
            remove_poor_attrs=self.remove_poor_attrs,
            merit_preprune=self.merit_preprune,
            split_buffer_size=self.split_buffer_size,
            seed=seed,
        )

    def _drift_input(self, y, y_hat) -> int | float:
        return abs(y - y_hat)

    def predict_one(self, x: dict) -> base.typing.RegTarget:
        if not self.models:
            self._init_trees(len(x))
            return 0.0  # type: ignore

        if not self.disable_weighted_vote:
            preds = []
            weights = []

            for perf, model in zip(self._perfs, self.models):
                preds.append(model.predict_one(x))
                weights.append(perf.get())

            sum_weights = sum(weights)
            if sum_weights != 0:
                if self.track_metric.bigger_is_better:
                    preds = [(w / sum_weights) * pred for w, pred in zip(weights, preds)]
                else:
                    weights = [(1 + 1e-8) / (w + 1e-8) for w in weights]
                    sum_weights = sum(weights)
                    preds = [(w / sum_weights) * pred for w, pred in zip(weights, preds)]
                return sum(preds)
        else:
            preds = [model.predict_one(x) for model in self.models]

        return sum(preds) / len(preds)


# river/imblearn/__init__.py

"""Sampling methods."""
from __future__ import annotations

from .chebyshev import ChebyshevOverSampler, ChebyshevUnderSampler
from .hard_sampling import HardSamplingClassifier, HardSamplingRegressor
from .random import RandomOverSampler, RandomSampler, RandomUnderSampler

__all__ = [
    "ChebyshevOverSampler",
    "ChebyshevUnderSampler",
    "HardSamplingClassifier",
    "HardSamplingRegressor",
    "RandomOverSampler",
    "RandomUnderSampler",
    "RandomSampler",
]


# river/imblearn/chebyshev.py

from __future__ import annotations

import math
import random

from river import base, linear_model, stats


class ChebyshevUnderSampler(base.Wrapper, base.Regressor):
    """Under-sampling for imbalanced regression using Chebyshev's inequality.

    Chebyshev's inequality can be used to define the probability of target observations
    being frequent values (w.r.t. the distribution mean).

    Let $Y$ be a random variable with finite expected value $\\overline{y}$ and non-zero variance
    $\\sigma^2$. For any real number $t > 0$, the Chebyshev's inequality states that, for a wide class of
    unimodal probability distributions: $Pr(|y-\\overline{y}| \\ge t\\sigma) \\le \\dfrac{1}{t^2}$.

    Taking $t=\\dfrac{|y-\\overline{y}|}{\\sigma}$, and assuming $t > 1$, the Chebyshev’s inequality
    for an observation $y$ becomes: $P(|y - \\overline{y}|=t) = \\dfrac{\\sigma^2}{|y-\\overline{y}|}$.
    The reciprocal of this probability is used for under-sampling[^1] the most frequent cases. Extreme
    valued or rare cases have higher probabilities of selection, whereas the most frequent cases are
    likely to be discarded. Still, frequent cases have a small chance of being selected (controlled via
    the `sp` parameter) in case few rare instances were observed.

    Parameters
    ----------
    regressor
        The regression model that will receive the biased sample.
    sp
        Second chance probability. Even if an example is not initially selected for training, it still has
        a small chance of being selected in case the number of rare case observed so far is small.
    seed
        Random seed to support reproducibility.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import metrics
    >>> from river import preprocessing
    >>> from river import rules

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     imblearn.ChebyshevUnderSampler(
    ...         regressor=rules.AMRules(
    ...             n_min=50, delta=0.01,
    ...         ),
    ...         seed=42
    ...     )
    ... )

    >>> evaluate.progressive_val_score(
    ...     datasets.TrumpApproval(),
    ...     model,
    ...     metrics.MAE(),
    ...     print_every=500
    ... )
    [500] MAE: 1.787162
    [1,000] MAE: 1.515711
    [1,001] MAE: 1.515236
    MAE: 1.515236

    References
    ----------
    [^1]: Aminian, Ehsan, Rita P. Ribeiro, and João Gama. "Chebyshev approaches for imbalanced data
    streams regression models." Data Mining and Knowledge Discovery 35.6 (2021): 2389-2466.

    """

    def __init__(self, regressor: base.Regressor, sp: float = 0.15, seed: int | None = None):
        self.regressor = regressor
        self.sp = sp
        self.seed = seed

        self._var = stats.Var()
        self._rng = random.Random(self.seed)

        self._freq_c = 0
        self._rare_c = 0

    @property
    def _wrapped_model(self):
        return self.regressor

    def predict_one(self, x, **kwargs):
        return self.regressor.predict_one(x, **kwargs)

    def learn_one(self, x, y, **kwargs):
        self._var.update(y)
        sd = self._var.get() ** 0.5

        if sd > 0:
            mean = self._var.mean.get()
            dev = abs(y - mean)  # noqa
            t = dev / sd

            # Small values for rare cases and 1 for frequent cases
            prob_threshold = 1 / (t * t) if t > 1 else 1
            p = self._rng.random()

            if p >= prob_threshold:
                self.regressor.learn_one(x, y, **kwargs)
                self._rare_c += 1
            elif self._freq_c < self._rare_c and p <= self.sp:
                self.regressor.learn_one(x, y, **kwargs)
                self._freq_c += 1
        else:
            self.regressor.learn_one(x, y, **kwargs)

        return self

    @classmethod
    def _unit_test_params(cls):
        yield {"regressor": linear_model.LinearRegression(), "seed": 42}


class ChebyshevOverSampler(base.Wrapper, base.Regressor):
    """Over-sampling for imbalanced regression using Chebyshev's inequality.

    Chebyshev's inequality can be used to define the probability of target observations
    being frequent values (w.r.t. the distribution mean).

    Let $Y$ be a random variable with finite expected value $\\overline{y}$ and non-zero variance
    $\\sigma^2$. For any real number $t > 0$, the Chebyshev's inequality states that, for a wide class of
    unimodal probability distributions: $Pr(|y-\\overline{y}| \\ge t\\sigma) \\le \\dfrac{1}{t^2}$.

    Taking $t=\\dfrac{|y-\\overline{y}|}{\\sigma}$, and assuming $t > 1$, the Chebyshev’s inequality
    for an observation $y$ becomes: $P(|y - \\overline{y}|=t) = \\dfrac{\\sigma^2}{|y-\\overline{y}|}$.

    Alternatively, one can use $t$ directly to estimate a frequency weight $\\kappa = \\lceil t\\rceil$
    and define an over-sampling strategy for extreme and rare target values[^1]. Each incoming instance is
    used $\\kappa$ times to update the underlying regressor. Frequent target values contribute only once
    to the underlying regressor, whereas rares cases are used multiple times for training.

    Parameters
    ----------
    regressor
        The regression model that will receive the biased sample.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import metrics
    >>> from river import preprocessing
    >>> from river import rules

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     imblearn.ChebyshevOverSampler(
    ...         regressor=rules.AMRules(
    ...             n_min=50, delta=0.01
    ...         )
    ...     )
    ... )

    >>> evaluate.progressive_val_score(
    ...     datasets.TrumpApproval(),
    ...     model,
    ...     metrics.MAE(),
    ...     print_every=500
    ... )
    [500] MAE: 1.682627
    [1,000] MAE: 1.761306
    [1,001] MAE: 1.759576
    MAE: 1.759576

    References
    ----------
    [^1]: Aminian, Ehsan, Rita P. Ribeiro, and João Gama. "Chebyshev approaches for imbalanced data
    streams regression models." Data Mining and Knowledge Discovery 35.6 (2021): 2389-2466.

    """

    def __init__(self, regressor: base.Regressor):
        self.regressor = regressor

        self._var = stats.Var()

    @property
    def _wrapped_model(self):
        return self.regressor

    def predict_one(self, x, **kwargs):
        return self.regressor.predict_one(x, **kwargs)

    def learn_one(self, x, y, **kwargs):
        self._var.update(y)
        var = self._var.get()
        sd = var**0.5

        if sd > 0:
            mean = self._var.mean.get()
            dev = abs(y - mean)  # noqa
            t = dev / sd

            kappa = int(math.ceil(t))

            for k in range(kappa):
                self.regressor.learn_one(x, y, **kwargs)
        else:
            self.regressor.learn_one(x, y, **kwargs)

        return self

    @classmethod
    def _unit_test_params(cls):
        yield {"regressor": linear_model.LinearRegression()}


# river/imblearn/hard_sampling.py

from __future__ import annotations

import collections
import random
import typing

from river import base, linear_model, optim, utils


class Triplet(collections.namedtuple("Triplet", "x y loss")):
    def __lt__(self, other):
        return self.loss < other.loss


class HardSampling(base.Wrapper):
    """Hard sampler."""

    def __init__(self, model, loss, size, p, seed: int | None = None):
        self.model = model
        self.loss = loss
        self.p = p
        self.size = size
        self.buffer = utils.SortedWindow(self.size)
        self.seed = seed
        self._rng = random.Random(seed)

    @property
    def _wrapped_model(self):
        return self.model

    def predict_one(self, x, **kwargs):
        return self.model.predict_one(x, **kwargs)

    @property
    def _model_pred_func(self) -> typing.Callable:
        if isinstance(self.model, base.Classifier):
            if not self.model._multiclass:
                return lambda x: self.model.predict_proba_one(x)[True]
            return self.model.predict_proba_one
        return self.model.predict_one

    def learn_one(self, x, y, **kwargs):
        loss = self.loss(y_true=y, y_pred=self._model_pred_func(x))

        if len(self.buffer) < self.size:
            self.buffer.append(Triplet(x=x, y=y, loss=loss))

        elif loss > self.buffer[0].loss:
            self.buffer.pop(0)

            self.buffer.append(Triplet(x=x, y=y, loss=loss))

        # Probability p
        if self._rng.uniform(0, 1) <= self.p:
            i = self._rng.randint(0, len(self.buffer) - 1)

            triplet = self.buffer.pop(i)

            self.model.learn_one(triplet.x, triplet.y, **kwargs)

            loss = self.loss(y_true=triplet.y, y_pred=self._model_pred_func(triplet.x))

            self.buffer.append(Triplet(x=triplet.x, y=triplet.y, loss=loss))

        # Probability (1 - p)
        else:
            self.model.learn_one(x, y)

        return self


class HardSamplingRegressor(HardSampling, base.Regressor):
    """Hard sampling regressor.

    This wrapper enables a model to retrain on past samples who's output was hard to predict.
    This works by storing the hardest samples in a buffer of a fixed size. When a new sample
    arrives, the wrapped model is either trained on one of the buffered samples with a probability
    p or on the new sample with a probability (1 - p).

    The hardness of an observation is evaluated with a loss function that compares the sample's
    ground truth with the wrapped model's prediction. If the buffer is not full, then the sample
    is added to the buffer. If the buffer is full and the new sample has a bigger loss than the
    lowest loss in the buffer, then the sample takes it's place.

    Parameters
    ----------
    regressor
    size
        Size of the buffer.
    p
        Probability of updating the model with a sample from the buffer instead of a new incoming
        sample.
    loss
        Criterion used to evaluate the hardness of a sample.
    seed
        Random seed.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     imblearn.HardSamplingRegressor(
    ...         regressor=linear_model.LinearRegression(),
    ...         p=.2,
    ...         size=30,
    ...         seed=42,
    ...     )
    ... )

    >>> evaluate.progressive_val_score(
    ...     datasets.TrumpApproval(),
    ...     model,
    ...     metrics.MAE(),
    ...     print_every=500
    ... )
    [500] MAE: 2.292501
    [1,000] MAE: 1.395797
    [1,001] MAE: 1.394693
    MAE: 1.394693

    """

    def __init__(
        self,
        regressor: base.Regressor,
        size: int,
        p: float,
        loss: optim.losses.RegressionLoss | None = None,
        seed: int | None = None,
    ):
        if loss is None:
            loss = optim.losses.Absolute()
        super().__init__(model=regressor, loss=loss, size=size, p=p, seed=seed)

    @property
    def regressor(self):
        return self.model

    @classmethod
    def _unit_test_params(cls):
        yield {"regressor": linear_model.LinearRegression(), "p": 0.1, "size": 40}


class HardSamplingClassifier(HardSampling, base.Classifier):
    """Hard sampling classifier.

    This wrapper enables a model to retrain on past samples who's output was hard to predict.
    This works by storing the hardest samples in a buffer of a fixed size. When a new sample
    arrives, the wrapped model is either trained on one of the buffered samples with a probability
    p or on the new sample with a probability (1 - p).

    The hardness of an observation is evaluated with a loss function that compares the sample's
    ground truth with the wrapped model's prediction. If the buffer is not full, then the sample
    is added to the buffer. If the buffer is full and the new sample has a bigger loss than the
    lowest loss in the buffer, then the sample takes it's place.

    Parameters
    ----------
    classifier
    size
        Size of the buffer.
    p
        Probability of updating the model with a sample from the buffer instead of a new
        incoming sample.
    loss
        Criterion used to evaluate the hardness of a sample.
    seed
        Random seed.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     imblearn.HardSamplingClassifier(
    ...         classifier=linear_model.LogisticRegression(),
    ...         p=0.1,
    ...         size=40,
    ...         seed=42,
    ...     )
    ... )

    >>> evaluate.progressive_val_score(
    ...     dataset=datasets.Phishing(),
    ...     model=model,
    ...     metric=metrics.ROCAUC(),
    ...     print_every=500,
    ... )
    [500] ROCAUC: 92.71%
    [1,000] ROCAUC: 94.75%
    [1,250] ROCAUC: 95.05%
    ROCAUC: 95.05%

    """

    def __init__(
        self,
        classifier: base.Classifier,
        size: int,
        p: float,
        loss: optim.losses.BinaryLoss | optim.losses.MultiClassLoss | None = None,
        seed: int | None = None,
    ):
        if loss is None:
            loss = optim.losses.CrossEntropy() if classifier._multiclass else optim.losses.Log()
        super().__init__(model=classifier, loss=loss, size=size, p=p, seed=seed)

    @property
    def classifier(self):
        return self.model

    @property
    def _multiclass(self):
        return self.model._multiclass

    def predict_proba_one(self, x, **kwargs):
        return self.model.predict_proba_one(x, **kwargs)

    @classmethod
    def _unit_test_params(cls):
        yield {"classifier": linear_model.LogisticRegression(), "p": 0.1, "size": 40}


# river/imblearn/random.py

from __future__ import annotations

import collections
import random
import typing

from river import base, utils


class ClassificationSampler(base.Wrapper, base.Classifier):
    def __init__(self, classifier, seed: int | None = None):
        self.classifier = classifier
        self.seed = seed
        self._rng = random.Random(seed)

    @property
    def _wrapped_model(self):
        return self.classifier

    def predict_proba_one(self, x, **kwargs):
        return self.classifier.predict_proba_one(x, **kwargs)

    def predict_one(self, x, **kwargs):
        return self.classifier.predict_one(x, **kwargs)


class RandomUnderSampler(ClassificationSampler):
    """Random under-sampling.

    This is a wrapper for classifiers. It will train the provided classifier by under-sampling the
    stream of given observations so that the class distribution seen by the classifier follows
    a given desired distribution. The implementation is a discrete version of rejection sampling.

    See [Working with imbalanced data](/user-guide/imbalanced-learning) for example usage.

    Parameters
    ----------
    classifier
    desired_dist
        The desired class distribution. The keys are the classes whilst the values are the desired
        class percentages. The values must sum up to 1.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> model = imblearn.RandomUnderSampler(
    ...     (
    ...         preprocessing.StandardScaler() |
    ...         linear_model.LogisticRegression()
    ...     ),
    ...     desired_dist={False: 0.4, True: 0.6},
    ...     seed=42
    ... )

    >>> dataset = datasets.CreditCard().take(3000)

    >>> metric = metrics.LogLoss()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    LogLoss: 0.0728

    References
    ----------
    [^1]: [Under-sampling a dataset with desired ratios](https://maxhalford.github.io/blog/undersampling-ratios/)
    [^2]: [Wikipedia article on rejection sampling](https://www.wikiwand.com/en/Rejection_sampling)

    """

    def __init__(self, classifier: base.Classifier, desired_dist: dict, seed: int | None = None):
        super().__init__(classifier=classifier, seed=seed)
        self.desired_dist = desired_dist
        self._actual_dist: typing.Counter = collections.Counter()
        self._pivot = None

    def learn_one(self, x, y, **kwargs):
        self._actual_dist[y] += 1
        f = self.desired_dist
        g = self._actual_dist

        # Check if the pivot needs to be changed
        if y != self._pivot:
            self._pivot = max(g.keys(), key=lambda y: f[y] / g[y])
        else:
            self.classifier.learn_one(x, y, **kwargs)
            return self

        # Determine the sampling ratio if the class is not the pivot
        M = f[self._pivot] / g[self._pivot]  # Likelihood ratio
        ratio = f[y] / (M * g[y])

        if ratio < 1 and self._rng.random() < ratio:
            self.classifier.learn_one(x, y, **kwargs)

        return self


class RandomOverSampler(ClassificationSampler):
    """Random over-sampling.

    This is a wrapper for classifiers. It will train the provided classifier by over-sampling the
    stream of given observations so that the class distribution seen by the classifier follows
    a given desired distribution. The implementation is a discrete version of reverse rejection
    sampling.

    See [Working with imbalanced data](/user-guide/imbalanced-learning) for example usage.

    Parameters
    ----------
    classifier
    desired_dist
        The desired class distribution. The keys are the classes whilst the values are the desired
        class percentages. The values must sum up to 1.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> model = imblearn.RandomOverSampler(
    ...     (
    ...         preprocessing.StandardScaler() |
    ...         linear_model.LogisticRegression()
    ...     ),
    ...     desired_dist={False: 0.4, True: 0.6},
    ...     seed=42
    ... )

    >>> dataset = datasets.CreditCard().take(3000)

    >>> metric = metrics.LogLoss()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    LogLoss: 0.054338

    """

    def __init__(self, classifier: base.Classifier, desired_dist: dict, seed: int | None = None):
        super().__init__(classifier=classifier, seed=seed)
        self.desired_dist = desired_dist
        self._actual_dist: typing.Counter = collections.Counter()
        self._pivot = None

    def learn_one(self, x, y, **kwargs):
        self._actual_dist[y] += 1
        f = self.desired_dist
        g = self._actual_dist

        # Check if the pivot needs to be changed
        if y != self._pivot:
            self._pivot = max(g.keys(), key=lambda y: g[y] / f[y])
        else:
            self.classifier.learn_one(x, y, **kwargs)
            return self

        M = g[self._pivot] / f[self._pivot]
        rate = M * f[y] / g[y]

        for _ in range(utils.random.poisson(rate, rng=self._rng)):
            self.classifier.learn_one(x, y, **kwargs)

        return self


class RandomSampler(ClassificationSampler):
    """Random sampling by mixing under-sampling and over-sampling.

    This is a wrapper for classifiers. It will train the provided classifier by both under-sampling
    and over-sampling the stream of given observations so that the class distribution seen by the
    classifier follows a given desired distribution.

    See [Working with imbalanced data](/user-guide/imbalanced-learning) for example usage.

    Parameters
    ----------
    classifier
    desired_dist
        The desired class distribution. The keys are the classes whilst the values are the desired
        class percentages. The values must sum up to 1. If set to `None`, then the observations
        will be sampled uniformly at random, which is stricly equivalent to using
        `ensemble.BaggingClassifier`.
    sampling_rate
        The desired ratio of data to sample.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import imblearn
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> model = imblearn.RandomSampler(
    ...     (
    ...         preprocessing.StandardScaler() |
    ...         linear_model.LogisticRegression()
    ...     ),
    ...     desired_dist={False: 0.4, True: 0.6},
    ...     sampling_rate=0.8,
    ...     seed=42
    ... )

    >>> dataset = datasets.CreditCard().take(3000)

    >>> metric = metrics.LogLoss()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    LogLoss: 0.131988

    """

    def __init__(
        self,
        classifier: base.Classifier,
        desired_dist: dict,
        sampling_rate=1.0,
        seed: int | None = None,
    ):
        super().__init__(classifier=classifier, seed=seed)
        self.sampling_rate = sampling_rate
        self._actual_dist: typing.Counter = collections.Counter()
        if desired_dist is None:
            desired_dist = self._actual_dist
        self.desired_dist = desired_dist
        self._n = 0

    def learn_one(self, x, y, **kwargs):
        self._actual_dist[y] += 1
        self._n += 1
        f = self.desired_dist
        g = self._actual_dist

        rate = self.sampling_rate * f[y] / (g[y] / self._n)

        for _ in range(utils.random.poisson(rate, rng=self._rng)):
            self.classifier.learn_one(x, y, **kwargs)

        return self


# river/linear_model/__init__.py

"""Linear models."""
from __future__ import annotations

from . import base
from .alma import ALMAClassifier
from .bayesian_lin_reg import BayesianLinearRegression
from .lin_reg import LinearRegression
from .log_reg import LogisticRegression
from .pa import PAClassifier, PARegressor
from .perceptron import Perceptron
from .softmax import SoftmaxRegression

__all__ = [
    "base",
    "ALMAClassifier",
    "BayesianLinearRegression",
    "LinearRegression",
    "LogisticRegression",
    "PAClassifier",
    "PARegressor",
    "Perceptron",
    "SoftmaxRegression",
]


# river/linear_model/alma.py

from __future__ import annotations

import collections
import math

from river import base, utils

__all__ = ["ALMAClassifier"]


class ALMAClassifier(base.Classifier):
    """Approximate Large Margin Algorithm (ALMA).

    Parameters
    ----------
    p
    alpha
    B
    C

    Attributes
    ----------
    w : collections.defaultdict
        The current weights.
    k : int
        The number of instances seen during training.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.ALMAClassifier()
    ... )

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 82.64%

    References
    ----------
    [^1]: [Gentile, Claudio. "A new approximate maximal margin classification algorithm." Journal of Machine Learning Research 2.Dec (2001): 213-242](http://www.jmlr.org/papers/volume2/gentile01a/gentile01a.pdf)

    """

    def __init__(self, p=2, alpha=0.9, B=1 / 0.9, C=2**0.5):
        self.p = p
        self.alpha = alpha
        self.B = B
        self.C = C
        self.w = collections.defaultdict(float)
        self.k = 1

    def _raw_dot(self, x):
        return utils.math.dot(x, self.w)

    def predict_proba_one(self, x):
        yp = utils.math.sigmoid(self._raw_dot(x))
        return {False: 1 - yp, True: yp}

    def learn_one(self, x, y):
        # Convert 0 to -1
        y = int(y or -1)

        gamma = self.B * math.sqrt(self.p - 1) / math.sqrt(self.k)

        if y * self._raw_dot(x) < (1 - self.alpha) * gamma:
            eta = self.C / (math.sqrt(self.p - 1) * math.sqrt(self.k))

            for i, xi in x.items():
                self.w[i] += eta * y * xi

            norm = utils.math.norm(self.w, order=self.p)

            for i in x:
                self.w[i] /= max(1, norm)

            self.k += 1

        return self


# river/linear_model/base.py

from __future__ import annotations

import contextlib
import numbers

import numpy as np
import pandas as pd

from river import optim, utils

__all__ = ["GLM"]


class GLM:
    """Generalized Linear Model.

    This serves as a base class for linear and logistic regression.

    Parameters
    ----------
    optimizer
        The sequential optimizer used for updating the weights. Note that the intercept updates are
        handled separately.
    loss
        The loss function to optimize for.
    l2
        Amount of L2 regularization used to push weights towards 0.
        For now, only one type of penalty can be used. The joint use of L1 and L2 is not explicitly supported.
    l1
        Amount of L1 regularization used to push weights towards 0.
        For now, only one type of penalty can be used. The joint use of L1 and L2 is not explicitly supported.
    intercept_init
        Initial intercept value.
    intercept_lr
        Learning rate scheduler used for updating the intercept. A `optim.schedulers.Constant` is
        used if a `float` is provided. The intercept is not updated when this is set to 0.
    clip_gradient
        Clips the absolute value of each gradient value.
    initializer
        Weights initialization scheme.

    """

    def __init__(
        self,
        optimizer,
        loss,
        l2,
        l1,
        intercept_init,
        intercept_lr,
        clip_gradient,
        initializer,
    ):
        self.optimizer = optimizer
        self.loss = loss
        self.l2 = l2
        self.l1 = l1
        self.intercept_init = intercept_init
        self.intercept = intercept_init
        self.intercept_lr = (
            optim.schedulers.Constant(intercept_lr)
            if isinstance(intercept_lr, numbers.Number)
            else intercept_lr
        )
        self.clip_gradient = clip_gradient
        self.initializer = initializer
        self._weights = utils.VectorDict(None)

        # The predict_many functions are going to return pandas.Series. We can name the series with
        # the name given to the y series seen during the last learn_many call.
        self._y_name = None

        if l1 != 0:
            if l2 != 0:
                raise NotImplementedError(
                    "The joint use of L1 and L2 penalties is currently not explicitly supported!"
                )

            # L1-specific parameters
            self.max_cum_l1 = 0
            self.cum_l1 = utils.VectorDict(None, optim.initializers.Zeros())

    @property
    def _mutable_attributes(self):
        return {"optimizer", "l2", "l1", "loss", "intercept_lr", "clip_gradient", "initializer"}

    @property
    def weights(self):
        return self._weights.to_dict()

    @contextlib.contextmanager
    def _learn_mode(self, mask=None):
        weights = self._weights
        try:
            # enable the initializer and set a mask
            self._weights = utils.VectorDict(weights, self.initializer, mask)
            yield
        finally:
            self._weights = weights

    def _get_intercept_update(self, loss_gradient):
        return self.intercept_lr.get(self.optimizer.n_iterations) * loss_gradient

    def _fit(self, x, y, w, get_grad):
        # Some optimizers need to do something before a prediction is made
        self.optimizer.look_ahead(w=self._weights)

        # Calculate the gradient
        gradient, loss_gradient = get_grad(x, y, w)

        # Update the intercept
        self.intercept -= self._get_intercept_update(loss_gradient)

        # Update the weights
        self.optimizer.step(w=self._weights, g=gradient)

        # Apply L1 cumulative penalty if applicable
        if self.l1 != 0.0:
            # This should be called after the learning_rate update in case of adaptive learning rate
            self.max_cum_l1 = self.max_cum_l1 + self.l1 * self.optimizer.learning_rate

            self._update_weights(x)

        return self

    def _update_weights(self, x):
        # L1 cumulative penalty helper

        # Apply penalty to each weight iteratively, with the potential of being parrallelized by using VectorDict
        for j, xj in x.items():
            wj_temp = self._weights[j]

            if wj_temp > 0:
                self._weights[j] = max(0, wj_temp - (self.max_cum_l1 + self.cum_l1[j]))
            elif wj_temp < 0:
                self._weights[j] = min(0, wj_temp + (self.max_cum_l1 - self.cum_l1[j]))
            else:
                self._weights[j] = wj_temp

            # Update the penalty state of the estimator
            self.cum_l1[j] = self.cum_l1[j] + (self._weights[j] - wj_temp)

    # Single instance methods

    def _raw_dot_one(self, x: dict) -> float:
        return self._weights @ utils.VectorDict(x) + self.intercept

    def _eval_gradient_one(self, x: dict, y: float, w: float) -> tuple[dict, float]:
        loss_gradient = self.loss.gradient(y_true=y, y_pred=self._raw_dot_one(x))
        loss_gradient *= w
        loss_gradient = float(
            utils.math.clamp(loss_gradient, -self.clip_gradient, self.clip_gradient)
        )

        if self.l2:
            return (
                loss_gradient * utils.VectorDict(x) + self.l2 * self._weights,
                loss_gradient,
            )

        return (loss_gradient * utils.VectorDict(x), loss_gradient)

    def learn_one(self, x, y, w=1.0):
        with self._learn_mode(x):
            return self._fit(x, y, w, get_grad=self._eval_gradient_one)

    # Mini-batch methods

    def _raw_dot_many(self, X: pd.DataFrame) -> np.ndarray:
        return X.values @ self._weights.to_numpy(X.columns) + self.intercept

    def _eval_gradient_many(
        self, X: pd.DataFrame, y: pd.Series, w: float | pd.Series
    ) -> tuple[dict, float]:
        loss_gradient = self.loss.gradient(y_true=y.values, y_pred=self._raw_dot_many(X))
        loss_gradient *= w
        loss_gradient = np.clip(loss_gradient, -self.clip_gradient, self.clip_gradient)

        # At this point we have a feature matrix X of shape (n, p). The loss gradient is a vector
        # of length p. We want to multiply each of X's rows by the corresponding value in the loss
        # gradient. When this is all done, we collapse X by computing the average of each column,
        # thereby obtaining the mean gradient of the batch. From thereon, the code reduces to the
        # single instance case.
        gradient = np.einsum("ij,i->ij", X.values, loss_gradient).mean(axis=0)
        if self.l2:
            gradient += self.l2 * self._weights.to_numpy(X.columns)

        return dict(zip(X.columns, gradient)), loss_gradient.mean()

    def learn_many(self, X: pd.DataFrame, y: pd.Series, w: float | pd.Series = 1):
        self._y_name = y.name
        with self._learn_mode(set(X)):
            return self._fit(X, y, w, get_grad=self._eval_gradient_many)


# river/linear_model/bayesian_lin_reg.py

from __future__ import annotations

import numpy as np

from river import base, proba, utils


class BayesianLinearRegression(base.Regressor):
    """Bayesian linear regression.

    An advantage of Bayesian linear regression over standard linear regression is that features
    do not have to scaled beforehand. Another attractive property is that this flavor of linear
    regression is somewhat insensitive to its hyperparameters. Finally, this model can output
    instead a predictive distribution rather than just a point estimate.

    The downside is that the learning step runs in `O(n^2)` time, whereas the learning step of
    standard linear regression takes `O(n)` time.

    Parameters
    ----------
    alpha
        Prior parameter.
    beta
        Noise parameter.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics

    >>> dataset = datasets.TrumpApproval()
    >>> model = linear_model.BayesianLinearRegression()
    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric).get()
    0.5818

    >>> x, _ = next(iter(dataset))
    >>> model.predict_one(x)
    43.61

    >>> model.predict_one(x, with_dist=True)
    𝒩(μ=43.616, σ=1.003)

    References
    ----------
    [^1]: [Pattern Recognition and Machine Learning, page 52 — Christopher M. Bishop](https://www.microsoft.com/en-us/research/uploads/prod/2006/01/Bishop-Pattern-Recognition-and-Machine-Learning-2006.pdf)
    [^2]: [Bayesian/Streaming Algorithms — Vincent Warmerdam](https://koaning.io/posts/bayesian-propto-streaming/)
    [^3]: [Bayesian linear regression for practitioners — Max Halford](https://maxhalford.github.io/blog/bayesian-linear-regression/)

    """

    def __init__(self, alpha=1, beta=1):
        self.alpha = alpha
        self.beta = beta
        self._ss = {}
        self._ss_inv = {}
        self._m = {}
        self._n = 1

    def _unit_test_skips(self):
        return {"check_shuffle_features_no_impact", "check_emerging_features"}

    def _get_arrays(self, features, m=True, ss=True, ss_inv=True):
        m_arr = np.array([self._m.get(i, 0.0) for i in features]) if m else None
        ss_arr = (
            np.array([[self._ss.get(min((i, j), (j, i)), 0.0) for j in features] for i in features])
            if ss
            else None
        )
        ss_inv_arr = (
            np.array(
                [
                    [
                        self._ss_inv.get(min((i, j), (j, i)), 1.0 / self.alpha if i == j else 0.0)
                        for j in features
                    ]
                    for i in features
                ],
                order="F",
            )
            if ss_inv
            else None
        )
        return m_arr, ss_arr, ss_inv_arr

    def _set_arrays(self, features, m_arr, ss_arr, ss_inv_arr):
        for i, fi in enumerate(features):
            self._m[fi] = m_arr[i]
            ss_row = ss_arr[i]
            ss_inv_row = ss_inv_arr[i]
            for j, fj in enumerate(features):
                self._ss[min((fi, fj), (fj, fi))] = ss_row[j]
                self._ss_inv[min((fi, fj), (fj, fi))] = ss_inv_row[j]

    def learn_one(self, x, y):
        x_arr = np.array(list(x.values()))
        m_arr, ss_arr, ss_inv_arr = self._get_arrays(x.keys())

        bx = self.beta * x_arr
        utils.math.sherman_morrison(A=ss_inv_arr, u=bx, v=x_arr)
        # Bishop equation 3.50
        m_arr = ss_inv_arr @ (ss_arr @ m_arr + bx * y)
        # Bishop equation 3.51
        ss_arr += np.outer(bx, x_arr)

        self._set_arrays(x.keys(), m_arr, ss_arr, ss_inv_arr)

        return self

    def predict_one(self, x, with_dist=False):
        """Predict the output of features `x`.

        Parameters
        ----------
        x
            A dictionary of features.
        with_dist
            Whether to return a predictive distribution, or instead just the most likely value.

        Returns
        -------
        The prediction.

        """

        # Bishop equation 3.58
        y_pred_mean = utils.math.dot(self._m, x)
        if not with_dist:
            return y_pred_mean

        x_arr = np.array(list(x.values()))
        _, _, ss_inv_arr = self._get_arrays(x.keys(), m=False, ss=False)
        # Bishop equation 3.59
        y_pred_var = 1 / self.beta + x_arr @ ss_inv_arr @ x_arr.T

        return proba.Gaussian._from_state(n=1, m=y_pred_mean, sig=y_pred_var**0.5, ddof=0)


# river/linear_model/lin_reg.py

from __future__ import annotations

import numpy as np
import pandas as pd

from river import base, linear_model, optim, utils


class LinearRegression(linear_model.base.GLM, base.MiniBatchRegressor):
    """Linear regression.

    This estimator supports learning with mini-batches. On top of the single instance methods, it
    provides the following methods: `learn_many`, `predict_many`, `predict_proba_many`. Each method
    takes as input a `pandas.DataFrame` where each column represents a feature.

    It is generally a good idea to scale the data beforehand in order for the optimizer to
    converge. You can do this online with a `preprocessing.StandardScaler`.

    Parameters
    ----------
    optimizer
        The sequential optimizer used for updating the weights. Note that the intercept updates are
        handled separately.
    loss
        The loss function to optimize for.
    l2
        Amount of L2 regularization used to push weights towards 0.
        For now, only one type of penalty can be used. The joint use of L1 and L2 is not explicitly supported.
    l1
        Amount of L1 regularization used to push weights towards 0.
        For now, only one type of penalty can be used. The joint use of L1 and L2 is not explicitly supported.
    intercept_init
        Initial intercept value.
    intercept_lr
        Learning rate scheduler used for updating the intercept. A `optim.schedulers.Constant` is
        used if a `float` is provided. The intercept is not updated when this is set to 0.
    clip_gradient
        Clips the absolute value of each gradient value.
    initializer
        Weights initialization scheme.

    Attributes
    ----------
    weights : dict
        The current weights.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LinearRegression(intercept_lr=.1)
    ... )
    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 0.555971

    >>> model['LinearRegression'].intercept
    35.617670

    You can call the `debug_one` method to break down a prediction. This works even if the
    linear regression is part of a pipeline.

    >>> x, y = next(iter(dataset))
    >>> report = model.debug_one(x)
    >>> print(report)
    0. Input
    --------
    gallup: 43.84321 (float)
    ipsos: 46.19925 (float)
    morning_consult: 48.31875 (float)
    ordinal_date: 736389 (int)
    rasmussen: 44.10469 (float)
    you_gov: 43.63691 (float)
    <BLANKLINE>
    1. StandardScaler
    -----------------
    gallup: 1.18810 (float)
    ipsos: 2.10348 (float)
    morning_consult: 2.73545 (float)
    ordinal_date: -1.73032 (float)
    rasmussen: 1.26872 (float)
    you_gov: 1.48391 (float)
    <BLANKLINE>
    2. LinearRegression
    -------------------
    Name              Value      Weight      Contribution
          Intercept    1.00000    35.61767       35.61767
              ipsos    2.10348     0.62689        1.31866
    morning_consult    2.73545     0.24180        0.66144
             gallup    1.18810     0.43568        0.51764
          rasmussen    1.26872     0.28118        0.35674
            you_gov    1.48391     0.03123        0.04634
       ordinal_date   -1.73032     3.45162       -5.97242
    <BLANKLINE>
    Prediction: 32.54607

    """

    def __init__(
        self,
        optimizer: optim.base.Optimizer | None = None,
        loss: optim.losses.RegressionLoss | None = None,
        l2=0.0,
        l1=0.0,
        intercept_init=0.0,
        intercept_lr: optim.base.Scheduler | float = 0.01,
        clip_gradient=1e12,
        initializer: optim.base.Initializer | None = None,
    ):
        super().__init__(
            optimizer=optim.SGD(0.01) if optimizer is None else optimizer,
            loss=optim.losses.Squared() if loss is None else loss,
            intercept_init=intercept_init,
            intercept_lr=intercept_lr,
            l2=l2,
            l1=l1,
            clip_gradient=clip_gradient,
            initializer=initializer if initializer else optim.initializers.Zeros(),
        )

    def predict_one(self, x):
        return self.loss.mean_func(self._raw_dot_one(x))

    def predict_many(self, X):
        return pd.Series(
            self.loss.mean_func(self._raw_dot_many(X)),
            index=X.index,
            name=self._y_name,
            copy=False,
        )

    def debug_one(self, x: dict, decimals: int = 5) -> str:
        """Debugs the output of the linear regression.

        Parameters
        ----------
        x
            A dictionary of features.
        decimals
            The number of decimals use for printing each numeric value.

        Returns
        -------
        A table which explains the output.

        """

        def fmt_float(x):
            return "{: ,.{prec}f}".format(x, prec=decimals)

        names = list(map(str, x.keys())) + ["Intercept"]
        values = list(map(fmt_float, list(x.values()) + [1]))
        weights = list(map(fmt_float, [self._weights.get(i, 0) for i in x] + [self.intercept]))
        contributions = [xi * self._weights.get(i, 0) for i, xi in x.items()] + [self.intercept]
        order = list(reversed(np.argsort(contributions)))
        contributions = list(map(fmt_float, contributions))

        table = utils.pretty.print_table(
            headers=["Name", "Value", "Weight", "Contribution"],
            columns=[names, values, weights, contributions],
            order=order,
        )

        return table


# river/linear_model/log_reg.py

from __future__ import annotations

import pandas as pd

from river import base, linear_model, optim


class LogisticRegression(linear_model.base.GLM, base.MiniBatchClassifier):
    """Logistic regression.

    This estimator supports learning with mini-batches. On top of the single instance methods, it
    provides the following methods: `learn_many`, `predict_many`, `predict_proba_many`. Each method
    takes as input a `pandas.DataFrame` where each column represents a feature.

    It is generally a good idea to scale the data beforehand in order for the optimizer to
    converge. You can do this online with a `preprocessing.StandardScaler`.

    Parameters
    ----------
    optimizer
        The sequential optimizer used for updating the weights. Note that the intercept is handled
        separately.
    loss
        The loss function to optimize for. Defaults to `optim.losses.Log`.
    l2
        Amount of L2 regularization used to push weights towards 0.
        For now, only one type of penalty can be used. The joint use of L1 and L2 is not explicitly supported.
    l1
        Amount of L1 regularization used to push weights towards 0.
        For now, only one type of penalty can be used. The joint use of L1 and L2 is not explicitly supported.
    intercept_init
        Initial intercept value.
    intercept_lr
        Learning rate scheduler used for updating the intercept. A `optim.schedulers.Constant` is
        used if a `float` is provided. The intercept is not updated when this is set to 0.
    clip_gradient
        Clips the absolute value of each gradient value.
    initializer
        Weights initialization scheme.

    Attributes
    ----------
    weights
        The current weights.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer=optim.SGD(.1))
    ... )

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 88.96%

    """

    def __init__(
        self,
        optimizer: optim.base.Optimizer | None = None,
        loss: optim.losses.BinaryLoss | None = None,
        l2=0.0,
        l1=0.0,
        intercept_init=0.0,
        intercept_lr: float | optim.base.Scheduler = 0.01,
        clip_gradient=1e12,
        initializer: optim.base.Initializer | None = None,
    ):
        super().__init__(
            optimizer=optim.SGD(0.01) if optimizer is None else optimizer,
            loss=optim.losses.Log() if loss is None else loss,
            intercept_init=intercept_init,
            intercept_lr=intercept_lr,
            l2=l2,
            l1=l1,
            clip_gradient=clip_gradient,
            initializer=initializer if initializer else optim.initializers.Zeros(),
        )

    def predict_proba_one(self, x):
        p = self.loss.mean_func(self._raw_dot_one(x))  # Convert logit to probability
        return {False: 1.0 - p, True: p}

    def predict_proba_many(self, X: pd.DataFrame) -> pd.DataFrame:
        p = self.loss.mean_func(self._raw_dot_many(X))  # Convert logits to probabilities
        return pd.DataFrame({False: 1.0 - p, True: p}, index=X.index, copy=False)


# river/linear_model/pa.py

from __future__ import annotations

import collections

import numpy as np

from river import base, optim, utils

__all__ = ["PAClassifier", "PARegressor"]


class BasePA:
    def __init__(self, C, mode, learn_intercept):
        self.C = C
        self.mode = mode
        self.calc_tau = {0: self._calc_tau_0, 1: self._calc_tau_1, 2: self._calc_tau_2}[mode]
        self.learn_intercept = learn_intercept
        self.weights = collections.defaultdict(float)
        self.intercept = 0.0

    @classmethod
    def _calc_tau_0(cls, x, loss):
        norm = utils.math.norm(x, order=2) ** 2
        if norm > 0:
            return loss / utils.math.norm(x, order=2) ** 2
        return 0

    def _calc_tau_1(self, x, loss):
        norm = utils.math.norm(x, order=2) ** 2
        if norm > 0:
            return min(self.C, loss / norm)
        return 0

    def _calc_tau_2(self, x, loss):
        return loss / (utils.math.norm(x, order=2) ** 2 + 0.5 / self.C)


class PARegressor(BasePA, base.Regressor):
    """Passive-aggressive learning for regression.

    Parameters
    ----------
    C
    mode
    eps
    learn_intercept

    Examples
    --------

    The following example is taken from [this blog post](https://www.bonaccorso.eu/2017/10/06/ml-algorithms-addendum-passive-aggressive-algorithms/).

    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import stream
    >>> import numpy as np
    >>> from sklearn import datasets

    >>> np.random.seed(1000)
    >>> X, y = datasets.make_regression(n_samples=500, n_features=4)

    >>> model = linear_model.PARegressor(
    ...     C=0.01,
    ...     mode=2,
    ...     eps=0.1,
    ...     learn_intercept=False
    ... )
    >>> metric = metrics.MAE() + metrics.MSE()

    >>> for xi, yi in stream.iter_array(X, y):
    ...     y_pred = model.predict_one(xi)
    ...     model = model.learn_one(xi, yi)
    ...     metric = metric.update(yi, y_pred)

    >>> print(metric)
    MAE: 9.809402, MSE: 472.393532

    References
    ----------
    [^1]: [Crammer, K., Dekel, O., Keshet, J., Shalev-Shwartz, S. and Singer, Y., 2006. Online passive-aggressive algorithms. Journal of Machine Learning Research, 7(Mar), pp.551-585.](http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf)

    """

    def __init__(self, C=1.0, mode=1, eps=0.1, learn_intercept=True):
        super().__init__(C=C, mode=mode, learn_intercept=learn_intercept)
        self.eps = eps
        self.loss = optim.losses.EpsilonInsensitiveHinge(eps=eps)

    def learn_one(self, x, y):
        y_pred = self.predict_one(x)
        tau = self.calc_tau(x, self.loss(y, y_pred))
        step = tau * np.sign(y - y_pred)

        for i, xi in x.items():
            self.weights[i] += step * xi
        if self.learn_intercept:
            self.intercept += step

        return self

    def predict_one(self, x):
        return utils.math.dot(x, self.weights) + self.intercept


class PAClassifier(BasePA, base.Classifier):
    """Passive-aggressive learning for classification.

    Parameters
    ----------
    C
    mode
    eps
    learn_intercept

    Examples
    --------

    The following example is taken from [this blog post](https://www.bonaccorso.eu/2017/10/06/ml-algorithms-addendum-passive-aggressive-algorithms/).

    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import stream
    >>> import numpy as np
    >>> from sklearn import datasets
    >>> from sklearn import model_selection

    >>> np.random.seed(1000)
    >>> X, y = datasets.make_classification(
    ...     n_samples=5000,
    ...     n_features=4,
    ...     n_informative=2,
    ...     n_redundant=0,
    ...     n_repeated=0,
    ...     n_classes=2,
    ...     n_clusters_per_class=2
    ... )

    >>> X_train, X_test, y_train, y_test = model_selection.train_test_split(
    ...     X,
    ...     y,
    ...     test_size=0.35,
    ...     random_state=1000
    ... )

    >>> model = linear_model.PAClassifier(
    ...     C=0.01,
    ...     mode=1
    ... )

    >>> for xi, yi in stream.iter_array(X_train, y_train):
    ...     y_pred = model.learn_one(xi, yi)

    >>> metric = metrics.Accuracy() + metrics.LogLoss()

    >>> for xi, yi in stream.iter_array(X_test, y_test):
    ...     metric = metric.update(yi, model.predict_proba_one(xi))

    >>> print(metric)
    Accuracy: 88.46%, LogLoss: 0.325727

    References
    ----------
    [^1]: [Crammer, K., Dekel, O., Keshet, J., Shalev-Shwartz, S. and Singer, Y., 2006. Online passive-aggressive algorithms. Journal of Machine Learning Research, 7(Mar), pp.551-585](http://jmlr.csail.mit.edu/papers/volume7/crammer06a/crammer06a.pdf)

    """

    def __init__(self, C=1.0, mode=1, learn_intercept=True):
        super().__init__(C=C, mode=mode, learn_intercept=learn_intercept)
        self.loss = optim.losses.Hinge()

    def learn_one(self, x, y):
        y_pred = utils.math.dot(x, self.weights) + self.intercept
        tau = self.calc_tau(x, self.loss(y, y_pred))
        step = tau * (y or -1)  # y == False becomes -1

        for i, xi in x.items():
            self.weights[i] += step * xi
        if self.learn_intercept:
            self.intercept += step

        return self

    def predict_proba_one(self, x):
        y_pred = utils.math.sigmoid(utils.math.dot(x, self.weights) + self.intercept)
        return {False: 1.0 - y_pred, True: y_pred}


# river/linear_model/perceptron.py

from __future__ import annotations

from river import optim

from .log_reg import LogisticRegression


class Perceptron(LogisticRegression):
    """Perceptron classifier.

    In this implementation, the Perceptron is viewed as a special case of the logistic regression.
    The loss function that is used is the Hinge loss with a threshold set to 0, whilst the learning
    rate of the stochastic gradient descent procedure is set to 1 for both the weights and the
    intercept.

    Parameters
    ----------
    l2
        Amount of L2 regularization used to push weights towards 0.
    clip_gradient
        Clips the absolute value of each gradient value.
    initializer
        Weights initialization scheme.

    Attributes
    ----------
    weights
        The current weights.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model as lm
    >>> from river import metrics
    >>> from river import preprocessing as pp

    >>> dataset = datasets.Phishing()

    >>> model = pp.StandardScaler() | lm.Perceptron()

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 85.84%

    """

    def __init__(
        self,
        l2=0.0,
        clip_gradient=1e12,
        initializer: optim.initializers.Initializer | None = None,
    ):
        super().__init__(
            optimizer=optim.SGD(1),
            intercept_lr=1,
            loss=optim.losses.Hinge(threshold=0.0),
            l2=l2,
            clip_gradient=clip_gradient,
            initializer=initializer,
        )


# river/linear_model/softmax.py

from __future__ import annotations

import collections
import copy
import functools

from river import base, optim, utils

__all__ = ["SoftmaxRegression"]


class SoftmaxRegression(base.Classifier):
    """Softmax regression is a generalization of logistic regression to multiple classes.

    Softmax regression is also known as "multinomial logistic regression". There are a set weights
    for each class, hence the `weights` attribute is a nested `collections.defaultdict`. The main
    advantage of using this instead of a one-vs-all logistic regression is that the probabilities
    will be calibrated. Moreover softmax regression is more robust to outliers.

    Parameters
    ----------
    optimizer
        The sequential optimizer used to tune the weights.
    loss
        The loss function to optimize for.
    l2
        Amount of L2 regularization used to push weights towards 0.

    Attributes
    ----------
    weights : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.ImageSegments()

    >>> model = preprocessing.StandardScaler()
    >>> model |= linear_model.SoftmaxRegression()

    >>> metric = metrics.MacroF1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MacroF1: 81.88%

    References
    ----------
    [^1]: [Course on classification stochastic gradient descent](https://www.inf.ed.ac.uk/teaching/courses/mlp/2016/mlp02-sln.pdf)
    [^2]: [Binary vs. Multi-Class Logistic Regression](https://chrisyeh96.github.io/2018/06/11/logistic-regression.html)

    """

    def __init__(
        self,
        optimizer: optim.base.Optimizer | None = None,
        loss: optim.losses.MultiClassLoss | None = None,
        l2=0,
    ):
        if optimizer is None:
            optimizer = optim.SGD(0.01)
        new_optimizer = functools.partial(copy.deepcopy, optimizer)
        self.optimizer = optimizer
        self.optimizers = collections.defaultdict(new_optimizer)  # type: ignore
        self.loss = optim.losses.CrossEntropy() if loss is None else loss
        self.l2 = l2
        self.weights = collections.defaultdict(
            functools.partial(collections.defaultdict, float)
        )  # type: ignore

    @property
    def _multiclass(self):
        return True

    def learn_one(self, x, y):
        # Some optimizers need to do something before a prediction is made
        for label, weights in self.weights.items():
            self.optimizers[label].look_ahead(w=weights)

        # Make a prediction for the given features
        y_pred = self.predict_proba_one(x)

        # Compute the gradient of the loss w.r.t. each label
        loss_gradients = self.loss.gradient(y_true=y, y_pred=y_pred)

        for label, loss in loss_gradients.items():
            # Compute the gradient w.r.t. each feature
            weights = self.weights[label]
            gradient = {i: xi * loss + self.l2 * weights.get(i, 0) for i, xi in x.items()}
            self.weights[label] = self.optimizers[label].step(w=weights, g=gradient)

        return self

    def predict_proba_one(self, x):
        return utils.math.softmax(
            {label: utils.math.dot(weights, x) for label, weights in self.weights.items()}
        )


# river/linear_model/test_glm.py

from __future__ import annotations

import copy
import itertools
import math
import random

import numpy as np
import pandas as pd
import pytest
from sklearn import linear_model as sklm
from sklearn.datasets import make_classification, make_regression
from sklearn.metrics import log_loss

from river import datasets, optim, preprocessing, stream, utils
from river import linear_model as lm


def iter_perturbations(keys, n=10):
    """Enumerate perturbations that will be applied to the weights."""

    # Enumerate unit vectors
    for i in keys:
        yield {j: int(i == j) for j in keys}

    # Enumerate samples from a spherical Gaussian
    for _ in range(n):
        p = {j: random.gauss(0, 1) for j in keys}
        norm = utils.math.norm(p, order=2)
        for j in p:
            p[j] /= norm
        yield p


@pytest.mark.parametrize(
    "lm, dataset",
    [
        pytest.param(
            lm(optimizer=copy.deepcopy(optimizer), initializer=initializer, l2=0),
            dataset,
            id=f"{lm.__name__} - {optimizer} - {initializer}",
        )
        for lm, dataset in [
            (lm.LinearRegression, datasets.TrumpApproval().take(100)),
            (lm.LogisticRegression, datasets.Bananas().take(100)),
        ]
        for optimizer, initializer in itertools.product(
            [
                optim.AdaBound(),
                optim.AdaDelta(),
                optim.AdaGrad(),
                optim.AdaMax(),
                optim.Adam(),
                optim.AMSGrad(),
                # TODO: check momentum optimizers
                # optim.Momentum(),
                # optim.NesterovMomentum(),
                optim.RMSProp(),
                optim.SGD(),
            ],
            [
                optim.initializers.Zeros(),
                optim.initializers.Normal(mu=0, sigma=1, seed=42),
            ],
        )
    ],
)
def test_finite_differences(lm, dataset):
    """Checks the gradient of a linear model via finite differences.

    References
    ----------
    [^1]: [How to test gradient implementations](https://timvieira.github.io/blog/post/2017/04/21/how-to-test-gradient-implementations/)
    [^2]: [Stochastic Gradient Descent Tricks](https://cilvr.cs.nyu.edu/diglib/lsml/bottou-sgd-tricks-2012.pdf)

    """

    scaler = preprocessing.StandardScaler()
    eps = 1e-6

    for x, y in dataset:
        x = scaler.learn_one(x).transform_one(x)

        # Store the current gradient and weights
        gradient, _ = lm._eval_gradient_one(x, y, 1)
        weights = copy.deepcopy(lm._weights)

        # d is a set of weight perturbations
        for d in iter_perturbations(weights.keys()):
            # Pertubate the weights and obtain the loss with the new weights
            lm._weights = utils.VectorDict({i: weights[i] + eps * di for i, di in d.items()})
            forward = lm.loss(y_true=y, y_pred=lm._raw_dot_one(x))
            lm._weights = utils.VectorDict({i: weights[i] - eps * di for i, di in d.items()})
            backward = lm.loss(y_true=y, y_pred=lm._raw_dot_one(x))

            # We expect g and h to be equal
            g = utils.math.dot(d, gradient)
            h = (forward - backward) / (2 * eps)

            # Compare signs
            # TODO: reactivate this check
            # assert np.sign(g) == np.sign(h)

            # Check absolute difference
            # TODO: decrease the tolerance
            assert abs(g - h) < 1e-5

        # Reset the weights to their original values in order not to influence
        # the training loop, even though it doesn't really matter.
        lm._weights = weights
        lm.learn_one(x, y)


def test_one_many_consistent():
    """Checks that using learn_one or learn_many produces the same result."""

    X = pd.read_csv(datasets.TrumpApproval().path)
    Y = X.pop("five_thirty_eight")

    one = lm.LinearRegression()
    for x, y in stream.iter_pandas(X[:100], Y[:100]):
        one.learn_one(x, y)

    many = lm.LinearRegression()
    for xb, yb in zip(np.array_split(X[:100], len(X[:100])), np.array_split(Y[:100], len(Y[:100]))):
        many.learn_many(xb, yb)

    for i in X:
        assert math.isclose(one.weights[i], many.weights[i])


def test_shuffle_columns():
    """Checks that learn_many works identically whether columns are shuffled or not."""

    X = pd.read_csv(datasets.TrumpApproval().path)
    Y = X.pop("five_thirty_eight")

    normal = lm.LinearRegression()
    for xb, yb in zip(np.array_split(X, 10), np.array_split(Y, 10)):
        normal.learn_many(xb, yb)

    shuffled = lm.LinearRegression()
    for xb, yb in zip(np.array_split(X, 10), np.array_split(Y, 10)):
        cols = np.random.permutation(X.columns)
        shuffled.learn_many(xb[cols], yb)

    for i in X:
        assert math.isclose(normal.weights[i], shuffled.weights[i])


def test_add_remove_columns():
    """Checks that no exceptions are raised whenever columns are dropped and/or added."""

    X = pd.read_csv(datasets.TrumpApproval().path)
    Y = X.pop("five_thirty_eight")

    lin_reg = lm.LinearRegression()
    for xb, yb in zip(np.array_split(X, 10), np.array_split(Y, 10)):
        # Pick half of the columns at random
        cols = np.random.choice(X.columns, len(X.columns) // 2, replace=False)
        lin_reg.learn_many(xb[cols], yb)


class ScikitLearnSquaredLoss:
    """sklearn removes the leading 2 from the gradient of the squared loss."""

    def gradient(self, y_true, y_pred):
        return y_pred - y_true


lin_reg_tests = {
    "Vanilla": (
        {"optimizer": optim.SGD(1e-2), "loss": ScikitLearnSquaredLoss()},
        {"learning_rate": "constant", "eta0": 1e-2, "alpha": 0},
    ),
    "Huber": (
        {"optimizer": optim.SGD(1e-2), "loss": optim.losses.Huber()},
        {"loss": "huber", "learning_rate": "constant", "eta0": 1e-2, "alpha": 0},
    ),
    "No intercept": (
        {
            "optimizer": optim.SGD(1e-2),
            "intercept_lr": 0,
            "loss": ScikitLearnSquaredLoss(),
        },
        {
            "learning_rate": "constant",
            "eta0": 1e-2,
            "alpha": 0,
            "fit_intercept": False,
        },
    ),
    "L2 regu": (
        {
            "optimizer": optim.SGD(1e-2),
            "loss": ScikitLearnSquaredLoss(),
            "l2": 1e-3,
        },
        {"learning_rate": "constant", "eta0": 1e-2, "alpha": 1e-3},
    ),
}


@pytest.mark.parametrize(
    "river_params, sklearn_params",
    lin_reg_tests.values(),
    ids=lin_reg_tests.keys(),
)
def test_lin_reg_sklearn_coherence(river_params, sklearn_params):
    """Checks that the sklearn and river implementations produce the same results."""

    ss = preprocessing.StandardScaler()
    rv = lm.LinearRegression(**river_params)
    sk = sklm.SGDRegressor(**sklearn_params)

    for x, y in datasets.TrumpApproval().take(100):
        x = ss.learn_one(x).transform_one(x)
        rv.learn_one(x, y)
        sk.partial_fit([list(x.values())], [y])

    for i, w in enumerate(rv.weights.values()):
        assert math.isclose(w, sk.coef_[i])

    assert math.isclose(rv.intercept, sk.intercept_[0])


@pytest.mark.parametrize(
    "river_params, sklearn_params",
    lin_reg_tests.values(),
    ids=lin_reg_tests.keys(),
)
def test_lin_reg_sklearn_learn_many_coherence(river_params, sklearn_params):
    """Checks that the sklearn and river implementations produce the same results
    when learn_many is used."""

    ss = preprocessing.StandardScaler()
    rv = lm.LinearRegression(**river_params)
    sk = sklm.SGDRegressor(**sklearn_params)

    for x, y in datasets.TrumpApproval().take(100):
        x = ss.learn_one(x).transform_one(x)
        rv.learn_many(pd.DataFrame([x]), pd.Series([y]))
        sk.partial_fit([list(x.values())], [y])

    for i, w in enumerate(rv.weights.values()):
        assert math.isclose(w, sk.coef_[i])

    assert math.isclose(rv.intercept, sk.intercept_[0])


log_reg_tests = {
    "Vanilla": (
        {"optimizer": optim.SGD(1e-2)},
        {"learning_rate": "constant", "eta0": 1e-2, "alpha": 0, "loss": "log_loss"},
    ),
    "Hinge": (
        {"optimizer": optim.SGD(1e-2), "loss": optim.losses.Hinge()},
        {"learning_rate": "constant", "eta0": 1e-2, "alpha": 0},
    ),
    "No intercept": (
        {"optimizer": optim.SGD(1e-2), "intercept_lr": 0},
        {
            "learning_rate": "constant",
            "eta0": 1e-2,
            "alpha": 0,
            "loss": "log_loss",
            "fit_intercept": False,
        },
    ),
    "L2 regu": (
        {
            "optimizer": optim.SGD(1e-2),
            "l2": 1e-3,
        },
        {
            "learning_rate": "constant",
            "eta0": 1e-2,
            "alpha": 1e-3,
            "loss": "log_loss",
        },
    ),
    "Inverse-scaling": (
        {
            "optimizer": optim.SGD(optim.schedulers.InverseScaling(1e-2)),
            "intercept_lr": optim.schedulers.InverseScaling(1e-2),
        },
        {"eta0": 1e-2, "alpha": 0, "learning_rate": "invscaling", "loss": "log_loss"},
    ),
    "Optimal": (
        {
            "optimizer": optim.SGD(optim.schedulers.Optimal(optim.losses.Hinge(), alpha=1e-3)),
            "loss": optim.losses.Hinge(),
            "intercept_lr": optim.schedulers.Optimal(optim.losses.Hinge(), alpha=1e-3),
            "l2": 1e-3,
        },
        {"learning_rate": "optimal", "alpha": 1e-3},
    ),
    "Optimal no intercept": (
        {
            "optimizer": optim.SGD(optim.schedulers.Optimal(optim.losses.Hinge(), alpha=1e-3)),
            "loss": optim.losses.Hinge(),
            "intercept_lr": 0,
            "l2": 1e-3,
        },
        {"learning_rate": "optimal", "alpha": 1e-3, "fit_intercept": False},
    ),
}


@pytest.mark.parametrize(
    "river_params, sklearn_params",
    log_reg_tests.values(),
    ids=log_reg_tests.keys(),
)
def test_log_reg_sklearn_coherence(river_params, sklearn_params):
    """Checks that the sklearn and river implementations produce the same results."""

    ss = preprocessing.StandardScaler()
    rv = lm.LogisticRegression(**river_params)
    sk = sklm.SGDClassifier(**sklearn_params)

    for x, y in datasets.Bananas().take(100):
        x = ss.learn_one(x).transform_one(x)
        rv.learn_one(x, y)
        sk.partial_fit([list(x.values())], [y], classes=[False, True])

    for i, w in enumerate(rv.weights.values()):
        assert math.isclose(w, sk.coef_[0][i])

    assert math.isclose(rv.intercept, sk.intercept_[0])


perceptron_tests = {
    "Vanilla": (
        {},
        {},
    ),
    "L2 regu": (
        {
            "l2": 1e-3,
        },
        {
            "alpha": 1e-3,
            "penalty": "l2",
        },
    ),
}


@pytest.mark.parametrize(
    "river_params, sklearn_params",
    perceptron_tests.values(),
    ids=perceptron_tests.keys(),
)
def test_perceptron_sklearn_coherence(river_params, sklearn_params):
    """Checks that the sklearn and river implementations produce the same results."""

    ss = preprocessing.StandardScaler()
    rv = lm.Perceptron(**river_params)
    sk = sklm.Perceptron(**sklearn_params)

    for x, y in datasets.Bananas().take(100):
        x = ss.learn_one(x).transform_one(x)
        rv.learn_one(x, y)
        sk.partial_fit([list(x.values())], [y], classes=[False, True])

    for i, w in enumerate(rv.weights.values()):
        assert math.isclose(w, sk.coef_[0][i])

    assert math.isclose(rv.intercept, sk.intercept_[0])


def test_lin_reg_sklearn_l1_non_regression():
    """Checks that the river L1 implementation results are no worse than sklearn L1."""

    X, y, true_coeffs = make_regression(
        n_samples=1000,
        n_features=20,
        n_informative=4,
        coef=True,
        random_state=273,
    )
    X = pd.DataFrame(X)
    y = pd.Series(y)

    ss = preprocessing.StandardScaler()

    rv = lm.LinearRegression(
        **{
            "optimizer": optim.SGD(1e-2),
            "loss": ScikitLearnSquaredLoss(),
            "l1": 1e-1,
        }
    )

    sk = sklm.SGDRegressor(
        **{
            "learning_rate": "constant",
            "eta0": 1e-2,
            "alpha": 1e-1,
            "penalty": "l1",
        }
    )

    for xi, yi in stream.iter_pandas(X, y):
        xi_tr = ss.learn_one(xi).transform_one(xi)
        rv.learn_one(xi_tr, yi)
        sk.partial_fit([list(xi_tr.values())], [yi])

    rv_coeffs = np.array(list(rv.weights.values()))
    sk_coeffs = sk.coef_

    # check that the shrinkage to zero is no worse than in sklearn
    assert np.sum(rv_coeffs > 0) <= np.sum(sk_coeffs > 0)

    # check that the river coefficients shrink to the true ones with 10% relative tolerance
    assert np.isclose(rv_coeffs, true_coeffs, rtol=5e-02, atol=0.0).all()


def test_log_reg_sklearn_l1_non_regression():
    """Checks that the river L1 implementation results are no worse than sklearn L1."""

    (
        X,
        y,
    ) = make_classification(
        n_samples=1000,
        n_features=20,
        n_informative=4,
        n_classes=2,
        random_state=273,
    )
    X = pd.DataFrame(X)
    y = pd.Series(y)

    ss = preprocessing.StandardScaler()

    rv = lm.LogisticRegression(
        **{
            "optimizer": optim.SGD(1e-2),
            "l1": 1e-3,
        }
    )

    sk = sklm.SGDClassifier(
        **{
            "learning_rate": "constant",
            "eta0": 1e-2,
            "alpha": 1e-3,
            "penalty": "l1",
            "loss": "log_loss",
        }
    )

    rv_pred = list()
    sk_pred = list()
    for xi, yi in stream.iter_pandas(X, y):
        xi_tr = ss.learn_one(xi).transform_one(xi)
        rv.learn_one(xi_tr, yi)
        sk.partial_fit([list(xi_tr.values())], [yi], classes=[False, True])

        rv_pred.append(rv.predict_one(xi_tr))
        sk_pred.append(sk.predict([list(xi_tr.values())])[0])

    rv_coeffs = np.array(list(rv.weights.values()))
    sk_coeffs = sk.coef_

    # check that the shrinkage to zero is no worse than in sklearn
    assert np.sum(rv_coeffs > 0) <= np.sum(sk_coeffs > 0)

    # since we can't access true coefficients from sklearn.make_classification, we compare losses
    assert math.isclose(log_loss(y, rv_pred), log_loss(y, sk_pred))


# river/metrics/__init__.py

"""Evaluation metrics.

All the metrics are updated one sample at a time. This way we can track performance of
predictive methods over time.

Note that all metrics have a `revert` method, enabling them to be wrapped in `utils.Rolling`.
This allows computirng rolling metrics:

```py
from river import metrics, utils

y_true = [True, False, True, True]
y_pred = [False, False, True, True]

metric = utils.Rolling(metrics.Accuracy(), window_size=3)

for yt, yp in zip(y_true, y_pred):
    print(metric.update(yt, yp))
```

```
Accuracy: 0.00%
Accuracy: 50.00%
Accuracy: 66.67%
Accuracy: 100.00%
```

"""
from __future__ import annotations

from . import base, multioutput
from .accuracy import Accuracy
from .balanced_accuracy import BalancedAccuracy
from .confusion import ConfusionMatrix
from .cross_entropy import CrossEntropy
from .fbeta import (
    F1,
    FBeta,
    MacroF1,
    MacroFBeta,
    MicroF1,
    MicroFBeta,
    MultiFBeta,
    WeightedF1,
    WeightedFBeta,
)
from .fowlkes_mallows import FowlkesMallows
from .geometric_mean import GeometricMean
from .jaccard import Jaccard, MacroJaccard, MicroJaccard, WeightedJaccard
from .kappa import CohenKappa
from .log_loss import LogLoss
from .mae import MAE
from .mape import MAPE
from .mcc import MCC
from .mse import MSE, RMSE, RMSLE
from .mutual_info import AdjustedMutualInfo, MutualInfo, NormalizedMutualInfo
from .precision import MacroPrecision, MicroPrecision, Precision, WeightedPrecision
from .r2 import R2
from .rand import AdjustedRand, Rand
from .recall import MacroRecall, MicroRecall, Recall, WeightedRecall
from .report import ClassificationReport
from .roc_auc import ROCAUC
from .rolling_roc_auc import RollingROCAUC
from .silhouette import Silhouette
from .smape import SMAPE
from .vbeta import Completeness, Homogeneity, VBeta

__all__ = [
    "Accuracy",
    "AdjustedMutualInfo",
    "AdjustedRand",
    "BalancedAccuracy",
    "base",
    "Completeness",
    "ClassificationReport",
    "ClusteringReport",
    "CohenKappa",
    "ConfusionMatrix",
    "CrossEntropy",
    "Jaccard",
    "MacroJaccard",
    "MicroJaccard",
    "F1",
    "FBeta",
    "GeometricMean",
    "Homogeneity",
    "LogLoss",
    "MAE",
    "MAPE",
    "MacroF1",
    "MacroFBeta",
    "MacroPrecision",
    "MacroRecall",
    "MCC",
    "MicroF1",
    "MicroFBeta",
    "MicroPrecision",
    "MicroRecall",
    "MultiFBeta",
    "multioutput",
    "MSE",
    "MutualInfo",
    "NormalizedMutualInfo",
    "Precision",
    "Rand",
    "Recall",
    "RMSE",
    "FowlkesMallows",
    "RMSLE",
    "ROCAUC",
    "RollingROCAUC",
    "R2",
    "Precision",
    "Silhouette",
    "SMAPE",
    "VBeta",
    "WeightedF1",
    "WeightedFBeta",
    "WeightedPrecision",
    "WeightedRecall",
    "WeightedJaccard",
]


# river/metrics/accuracy.py

from __future__ import annotations

from river import metrics

__all__ = ["Accuracy"]


class Accuracy(metrics.base.MultiClassMetric):
    """Accuracy score, which is the percentage of exact matches.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [True, False, True, True, True]
    >>> y_pred = [True, True, False, True, True]

    >>> metric = metrics.Accuracy()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    Accuracy: 60.00%

    """

    def get(self):
        try:
            return self.cm.total_true_positives / self.cm.total_weight
        except ZeroDivisionError:
            return 0.0


# river/metrics/balanced_accuracy.py

from __future__ import annotations

from river import metrics

__all__ = ["BalancedAccuracy"]


class BalancedAccuracy(metrics.base.MultiClassMetric):
    """Balanced accuracy.

    Balanced accuracy is the average of recall obtained on each class. It is used to
    deal with imbalanced datasets in binary and multi-class classification problems.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

        >>> from river import metrics
        >>> y_true = [True, False, True, True, False, True]
        >>> y_pred = [True, False, True, True, True, False]

        >>> metric = metrics.BalancedAccuracy()
        >>> for yt, yp in zip(y_true, y_pred):
        ...     metric = metric.update(yt, yp)

        >>> metric
        BalancedAccuracy: 62.50%

        >>> y_true = [0, 1, 0, 0, 1, 0]
        >>> y_pred = [0, 1, 0, 0, 0, 1]
        >>> metric = metrics.BalancedAccuracy()
        >>> for yt, yp in zip(y_true, y_pred):
        ...     metric = metric.update(yt, yp)

        >>> metric
        BalancedAccuracy: 62.50%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                total += self.cm[c][c] / self.cm.sum_row[c]
            except ZeroDivisionError:
                continue
        try:
            score = total / len(self.cm.classes)

            return score

        except ZeroDivisionError:
            return 0.0


# river/metrics/base.py

from __future__ import annotations

import abc
import collections
import numbers
import operator

from river import base, stats, utils

from . import confusion

__all__ = [
    "BinaryMetric",
    "ClassificationMetric",
    "Metric",
    "Metrics",
    "MultiClassMetric",
    "RegressionMetric",
    "WrapperMetric",
]


class Metric(base.Base, abc.ABC):
    """Mother class for all metrics."""

    @abc.abstractmethod
    def update(self, y_true, y_pred) -> Metric:
        """Update the metric."""

    @abc.abstractmethod
    def revert(self, y_true, y_pred) -> Metric:
        """Revert the metric."""

    @abc.abstractmethod
    def get(self) -> float:
        """Return the current value of the metric."""

    @property
    @abc.abstractmethod
    def bigger_is_better(self) -> bool:
        """Indicate if a high value is better than a low one or not."""

    @abc.abstractmethod
    def works_with(self, model: base.Estimator) -> bool:
        """Indicates whether or not a metric can work with a given model."""

    @property
    def works_with_weights(self) -> bool:
        """Indicate whether the model takes into consideration the effect of sample weights"""
        return True

    def is_better_than(self, other) -> bool:
        """Indicate if the current metric is better than another one."""
        op = operator.gt if self.bigger_is_better else operator.lt
        return op(self.get(), other.get())

    def __gt__(self, other):
        return self.is_better_than(other)

    def __repr__(self):
        """Return the class name along with the current value of the metric."""
        return f"{self.__class__.__name__}: {self.get():{self._fmt}}".rstrip("0")

    def __str__(self):
        return repr(self)


class ClassificationMetric(Metric):
    """Mother class for all classification metrics.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    """

    _fmt = ".2%"  # output a percentage, e.g. 0.427 becomes "42,7%"

    def __init__(self, cm: confusion.ConfusionMatrix | None = None):
        if cm is None:
            cm = confusion.ConfusionMatrix()
        self.cm = cm

    def update(self, y_true, y_pred, sample_weight=1.0):
        self.cm.update(
            y_true,
            y_pred,
            sample_weight=sample_weight,
        )
        return self

    def revert(self, y_true, y_pred, sample_weight=1.0):
        self.cm.revert(
            y_true,
            y_pred,
            sample_weight=sample_weight,
        )
        return self

    @property
    def bigger_is_better(self):
        return True

    def works_with(self, model) -> bool:
        return utils.inspect.isclassifier(model)

    @property
    def requires_labels(self):
        """Indicates if labels are required, rather than probabilities."""
        return True

    @staticmethod
    def _clamp_proba(p):
        """Clamp a number in between the (0, 1) interval."""
        return utils.math.clamp(p, minimum=1e-15, maximum=1 - 1e-15)

    def __add__(self, other) -> Metrics:
        if not isinstance(other, ClassificationMetric):
            raise ValueError(
                f"{self.__class__.__name__} and {other.__class__.__name__} metrics "
                "are not compatible"
            )
        return Metrics([self, other])


class BinaryMetric(ClassificationMetric):
    """Mother class for all binary classification metrics.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.
    pos_val
        Value to treat as "positive".

    """

    def __init__(self, cm=None, pos_val=True):
        super().__init__(cm)
        self.pos_val = pos_val

    def update(
        self,
        y_true: bool,
        y_pred: bool | float | dict[bool, float],
        sample_weight=1.0,
    ) -> BinaryMetric:
        if self.requires_labels:
            y_pred = y_pred == self.pos_val
        return super().update(y_true == self.pos_val, y_pred, sample_weight)

    def revert(
        self,
        y_true: bool,
        y_pred: bool | float | dict[bool, float],
        sample_weight=1.0,
    ) -> BinaryMetric:
        if self.requires_labels:
            y_pred = y_pred == self.pos_val
        return super().revert(y_true == self.pos_val, y_pred, sample_weight)


class MultiClassMetric(ClassificationMetric):
    """Mother class for all multi-class classification metrics.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    """

    def works_with(self, model) -> bool:
        return utils.inspect.isclassifier(model) or utils.inspect.isclusterer(model)


class RegressionMetric(Metric):
    """Mother class for all regression metrics."""

    _fmt = ",.6f"  # use commas to separate big numbers and show 6 decimals

    @abc.abstractmethod
    def update(self, y_true: numbers.Number, y_pred: numbers.Number) -> RegressionMetric:
        """Update the metric."""

    @abc.abstractmethod
    def revert(self, y_true: numbers.Number, y_pred: numbers.Number) -> RegressionMetric:
        """Revert the metric."""

    @property
    def bigger_is_better(self):
        return False

    def works_with(self, model) -> bool:
        return utils.inspect.isregressor(model)

    def __add__(self, other) -> Metrics:
        if not isinstance(other, RegressionMetric):
            raise ValueError(
                f"{self.__class__.__name__} and {other.__class__.__name__} metrics "
                "are not compatible"
            )
        return Metrics([self, other])


class Metrics(Metric, collections.UserList):
    """A container class for handling multiple metrics at once.

    Parameters
    ----------
    metrics
    str_sep

    """

    def __init__(self, metrics, str_sep=", "):
        super().__init__(metrics)
        self.str_sep = str_sep

    def update(self, y_true, y_pred, sample_weight=1.0):
        # If the metrics are classification metrics, then we have to handle the case where some
        # of the metrics require labels, whilst others need to be fed probabilities
        if hasattr(self, "requires_labels") and not self.requires_labels:
            for m in self:
                if m.requires_labels:
                    m.update(y_true, max(y_pred, key=y_pred.get))
                else:
                    m.update(y_true, y_pred)
            return self

        for m in self:
            m.update(y_true, y_pred)
        return self

    def revert(self, y_true, y_pred, sample_weight=1.0):
        # If the metrics are classification metrics, then we have to handle the case where some
        # of the metrics require labels, whilst others need to be fed probabilities
        if hasattr(self, "requires_labels") and not self.requires_labels:
            for m in self:
                if m.requires_labels:
                    m.revert(y_true, max(y_pred, key=y_pred.get), sample_weight)
                else:
                    m.revert(y_true, y_pred, sample_weight)
            return self

        for m in self:
            m.revert(y_true, y_pred, sample_weight)
        return self

    def get(self):
        return [m.get() for m in self]

    def works_with(self, model) -> bool:
        return all(m.works_with(model) for m in self)

    @property
    def bigger_is_better(self):
        raise NotImplementedError

    @property
    def works_with_weights(self):
        return all(m.works_with_weights for m in self)

    @property
    def requires_labels(self):
        return all(m.requires_labels for m in self)

    def __repr__(self):
        return self.str_sep.join(str(m) for m in self)

    def __add__(self, other):
        try:
            other + self[0]  # Will raise a ValueError if incompatible
        except IndexError:
            pass
        self.append(other)
        return self

    def clone(self):
        return self.__class__([m.clone() for m in self])


class WrapperMetric(Metric):
    @property
    def _fmt(self):
        return self.metric._fmt

    @property
    @abc.abstractmethod
    def metric(self):
        """Gives access to the wrapped metric."""

    def get(self):
        return self.metric.get()

    @property
    def bigger_is_better(self):
        return self.metric.bigger_is_better

    def works_with(self, model):
        return self.metric.works_with(model)

    @property
    def requires_labels(self):
        return self.metric.requires_labels

    @property
    def __metaclass__(self):
        return self.metric.__class__

    def __repr__(self):
        name = self.__class__.__name__
        return super().__repr__().replace(name, f"{name}({self.metric.__class__.__name__})")


class MeanMetric(abc.ABC):
    """Many metrics are just running averages. This is a utility class that avoids repeating
    tedious stuff throughout the module for such metrics.

    """

    def __init__(self):
        self._mean = stats.Mean()

    @abc.abstractmethod
    def _eval(self, y_true, y_pred):
        pass

    def update(self, y_true, y_pred, sample_weight=1.0):
        self._mean.update(x=self._eval(y_true, y_pred), w=sample_weight)
        return self

    def revert(self, y_true, y_pred, sample_weight=1.0):
        self._mean.revert(x=self._eval(y_true, y_pred), w=sample_weight)
        return self

    def get(self):
        return self._mean.get()


class ClusteringMetric(base.Base, abc.ABC):
    """
    Mother class of all internal clustering metrics.
    """

    # Define the format specification used for string representation.
    _fmt = ",.6f"  # Use commas to separate big numbers and show 6 decimals

    @abc.abstractmethod
    def update(self, x, y_pred, centers, sample_weight=1.0) -> ClusteringMetric:
        """Update the metric."""

    @abc.abstractmethod
    def revert(self, x, y_pred, centers, sample_weight=1.0) -> ClusteringMetric:
        """Revert the metric."""

    @abc.abstractmethod
    def get(self) -> float:
        """Return the current value of the metric."""

    @property
    @abc.abstractmethod
    def bigger_is_better(self) -> bool:
        """Indicates if a high value is better than a low one or not."""

    def works_with(self, model: base.Estimator) -> bool:
        """Indicates whether or not a metric can work with a given model."""
        return utils.inspect.isclusterer(model)

    def __repr__(self):
        """Returns the class name along with the current value of the metric."""
        return f"{self.__class__.__name__}: {self.get():{self._fmt}}".rstrip("0")


# river/metrics/confusion.py

from __future__ import annotations

import functools
from collections import defaultdict

from river import utils


class ConfusionMatrix:
    """Confusion Matrix for binary and multi-class classification.

    Parameters
    ----------
    classes
        The initial set of classes. This is optional and serves only for displaying purposes.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = ['cat', 'ant', 'cat', 'cat', 'ant', 'bird']
    >>> y_pred = ['ant', 'ant', 'cat', 'cat', 'ant', 'cat']

    >>> cm = metrics.ConfusionMatrix()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     cm = cm.update(yt, yp)

    >>> cm
           ant  bird   cat
     ant     2     0     0
    bird     0     0     1
     cat     1     0     2

    >>> cm['bird']['cat']
    1.0

    Notes
    -----
    This confusion matrix is a 2D matrix of shape `(n_classes, n_classes)`, corresponding
    to a single-target (binary and multi-class) classification task.

    Each row represents `true` (actual) class-labels, while each column corresponds
    to the `predicted` class-labels. For example, an entry in position `[1, 2]` means
    that the true class-label is 1, and the predicted class-label is 2 (incorrect prediction).

    This structure is used to keep updated statistics about a single-output classifier's
    performance and to compute multiple evaluation metrics.

    """

    def __init__(self, classes=None):
        self._init_classes = set(classes) if classes is not None else set()
        self.sum_row = defaultdict(float)
        self.sum_col = defaultdict(float)
        self.data = defaultdict(functools.partial(defaultdict, float))
        self.n_samples = 0
        self.total_weight = 0

    def __getitem__(self, key):
        """Syntactic sugar for accessing the counts directly."""
        return self.data[key]

    def update(self, y_true, y_pred, sample_weight=1.0):
        self.n_samples += 1
        self._update(y_true, y_pred, sample_weight)
        return self

    def revert(self, y_true, y_pred, sample_weight=1.0):
        self.n_samples -= 1
        # Revert is equal to subtracting so we pass the negative sample_weight
        self._update(y_true, y_pred, -sample_weight)
        return self

    def _update(self, y_true, y_pred, sample_weight):
        self.data[y_true][y_pred] += sample_weight
        self.total_weight += sample_weight
        self.sum_row[y_true] += sample_weight
        self.sum_col[y_pred] += sample_weight

    @property
    def classes(self):
        return list(
            {c for c, n in self.sum_row.items() if n} | {c for c, n in self.sum_col.items() if n}
        )

    def __repr__(self):
        classes = sorted(self.classes)
        if not classes:
            return ""

        headers = [""] + list(map(str, classes))
        columns = [headers[1:]]
        for col in classes:
            columns.append([f"{int(self.data[row][col]):,}" for row in classes])

        return utils.pretty.print_table(headers, columns)

    def support(self, label):
        return self.sum_row[label]

    def true_positives(self, label):
        return self.data[label][label]

    def true_negatives(self, label):
        return self.total_true_positives - self.data[label][label]

    def false_positives(self, label):
        return self.sum_col[label] - self.data[label][label]

    def false_negatives(self, label):
        return self.sum_row[label] - self.data[label][label]

    @property
    def total_true_positives(self):
        return sum(self.true_positives(label) for label in self.classes)

    @property
    def total_true_negatives(self):
        return sum(self.true_negatives(label) for label in self.classes)

    @property
    def total_false_positives(self):
        return sum(self.false_positives(label) for label in self.classes)

    @property
    def total_false_negatives(self):
        return sum(self.false_negatives(label) for label in self.classes)


# river/metrics/cross_entropy.py

from __future__ import annotations

import math

from river import metrics, utils

__all__ = ["CrossEntropy"]


class CrossEntropy(metrics.base.MeanMetric, metrics.base.MultiClassMetric):
    """Multiclass generalization of the logarithmic loss.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2]
    >>> y_pred = [
    ...     {0: 0.29450637, 1: 0.34216758, 2: 0.36332605},
    ...     {0: 0.21290077, 1: 0.32728332, 2: 0.45981591},
    ...     {0: 0.42860913, 1: 0.33380113, 2: 0.23758974},
    ...     {0: 0.44941979, 1: 0.32962558, 2: 0.22095463}
    ... ]

    >>> metric = metrics.CrossEntropy()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)
    ...     print(metric.get())
    1.222454
    1.169691
    1.258864
    1.321597

    >>> metric
    CrossEntropy: 1.321598

    """

    _fmt = ""

    @property
    def bigger_is_better(self):
        return False

    @property
    def requires_labels(self):
        return False

    def _eval(self, y_true, y_pred):
        total = 0

        for label, proba in y_pred.items():
            if y_true == label:
                total += math.log(utils.math.clamp(x=proba, minimum=1e-15, maximum=1 - 1e-15))

        return -total


# river/metrics/fbeta.py

from __future__ import annotations

import collections
import functools

from river import metrics

__all__ = [
    "F1",
    "FBeta",
    "MacroF1",
    "MacroFBeta",
    "MicroF1",
    "MicroFBeta",
    "MultiFBeta",
    "WeightedF1",
    "WeightedFBeta",
]


class FBeta(metrics.base.BinaryMetric):
    """Binary F-Beta score.

    The FBeta score is a weighted harmonic mean between precision and recall. The higher the
    `beta` value, the higher the recall will be taken into account. When `beta` equals 1,
    precision and recall and equivalently weighted, which results in the F1 score (see
    `metrics.F1`).

    Parameters
    ----------
    beta
        Weight of precision in the harmonic mean.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.
    pos_val
        Value to treat as "positive".

    Attributes
    ----------
    precision : metrics.Precision
    recall : metrics.Recall

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [False, False, False, True, True, True]
    >>> y_pred = [False, False, True, True, False, False]

    >>> metric = metrics.FBeta(beta=2)
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    FBeta: 35.71%

    """

    def __init__(self, beta: float, cm=None, pos_val=True):
        super().__init__(cm, pos_val)
        self.beta = beta
        self.precision = metrics.Precision(self.cm, self.pos_val)
        self.recall = metrics.Recall(self.cm, self.pos_val)

    def get(self):
        p = self.precision.get()
        r = self.recall.get()
        b2 = self.beta**2
        try:
            return (1 + b2) * p * r / (b2 * p + r)
        except ZeroDivisionError:
            return 0.0


class MacroFBeta(metrics.base.MultiClassMetric):
    """Macro-average F-Beta score.

    This works by computing the F-Beta score per class, and then performs a global average.

    Parameters
    ----------
    beta
        Weight of precision in harmonic mean.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MacroFBeta(beta=.8)

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MacroFBeta: 100.00%
    MacroFBeta: 31.06%
    MacroFBeta: 54.04%
    MacroFBeta: 54.04%
    MacroFBeta: 48.60%

    """

    def __init__(self, beta, cm=None):
        super().__init__(cm)
        self.beta = beta

    def get(self):
        total = 0
        b2 = self.beta**2

        for c in self.cm.classes:
            try:
                p = self.cm[c][c] / self.cm.sum_col[c]
            except ZeroDivisionError:
                p = 0

            try:
                r = self.cm[c][c] / self.cm.sum_row[c]
            except ZeroDivisionError:
                r = 0

            try:
                total += (1 + b2) * p * r / (b2 * p + r)
            except ZeroDivisionError:
                continue

        try:
            return total / len(self.cm.classes)
        except ZeroDivisionError:
            return 0.0


class MicroFBeta(metrics.base.MultiClassMetric):
    """Micro-average F-Beta score.

    This computes the F-Beta score by merging all the predictions and true labels, and then
    computes a global F-Beta score.

    Parameters
    ----------
    beta
        Weight of precision in the harmonic mean.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 0]
    >>> y_pred = [0, 1, 1, 2, 1]

    >>> metric = metrics.MicroFBeta(beta=2)
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    MicroFBeta: 60.00%

    References
    ----------
    1. [Why are precision, recall and F1 score equal when using micro averaging in a multi-class problem?](https://simonhessner.de/why-are-precision-recall-and-f1-score-equal-when-using-micro-averaging-in-a-multi-class-problem/)

    """

    def __init__(self, beta: float, cm=None):
        super().__init__(cm)
        self.beta = beta
        self.precision = metrics.MicroPrecision(self.cm)
        self.recall = metrics.MicroRecall(self.cm)

    def get(self):
        p = self.precision.get()
        r = self.recall.get()
        b2 = self.beta**2
        try:
            return (1 + b2) * p * r / (b2 * p + r)
        except ZeroDivisionError:
            return 0.0


class WeightedFBeta(metrics.base.MultiClassMetric):
    """Weighted-average F-Beta score.

    This works by computing the F-Beta score per class, and then performs a global weighted average
    according to the support of each class.

    Parameters
    ----------
    beta
        Weight of precision in the harmonic mean.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.WeightedFBeta(beta=0.8)

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    WeightedFBeta: 100.00%
    WeightedFBeta: 31.06%
    WeightedFBeta: 54.04%
    WeightedFBeta: 65.53%
    WeightedFBeta: 62.63%

    """

    def __init__(self, beta, cm=None):
        super().__init__(cm)
        self.beta = beta

    def get(self):
        total = 0
        b2 = self.beta**2

        for c in self.cm.classes:
            try:
                p = self.cm.sum_row[c] * self.cm[c][c] / self.cm.sum_col[c]
            except ZeroDivisionError:
                p = 0

            try:
                r = self.cm.sum_row[c] * self.cm[c][c] / self.cm.sum_row[c]
            except ZeroDivisionError:
                r = 0

            try:
                total += (1 + b2) * p * r / (b2 * p + r)
            except ZeroDivisionError:
                continue

        try:
            return total / self.cm.total_weight
        except ZeroDivisionError:
            return 0.0


class MultiFBeta(metrics.base.MultiClassMetric):
    """Multi-class F-Beta score with different betas per class.

    The multiclass F-Beta score is the arithmetic average of the binary F-Beta scores of each class.
    The mean can be weighted by providing class weights.

    Parameters
    ----------
    betas
        Weight of precision in the harmonic mean of each class.
    weights
        Class weights. If not provided then uniform weights will be used.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MultiFBeta(
    ...     betas={0: 0.25, 1: 1, 2: 4},
    ...     weights={0: 1, 1: 1, 2: 2}
    ... )

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MultiFBeta: 100.00%
    MultiFBeta: 25.76%
    MultiFBeta: 62.88%
    MultiFBeta: 62.88%
    MultiFBeta: 46.88%

    """

    def __init__(self, betas, weights, cm=None):
        super().__init__(cm)
        self.betas = betas
        self.weights = (
            collections.defaultdict(functools.partial(int, 1)) if weights is None else weights
        )

    def get(self):
        total = 0

        for c in self.cm.classes:
            b2 = self.betas[c] ** 2

            try:
                p = self.cm[c][c] / self.cm.sum_col[c]
            except ZeroDivisionError:
                p = 0

            try:
                r = self.cm[c][c] / self.cm.sum_row[c]
            except ZeroDivisionError:
                r = 0

            try:
                total += self.weights[c] * (1 + b2) * p * r / (b2 * p + r)
            except ZeroDivisionError:
                continue

        try:
            return total / sum(self.weights[c] for c in self.cm.classes)
        except ZeroDivisionError:
            return 0.0


class F1(FBeta):
    """Binary F1 score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [False, False, False, True, True, True]
    >>> y_pred = [False, False, True, True, False, False]

    >>> metric = metrics.F1()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    F1: 40.00%

    """

    def __init__(self, cm=None, pos_val=True):
        super().__init__(beta=1.0, cm=cm, pos_val=pos_val)


class MacroF1(MacroFBeta):
    """Macro-average F1 score.

    This works by computing the F1 score per class, and then performs a global average.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MacroF1()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MacroF1: 100.00%
    MacroF1: 33.33%
    MacroF1: 55.56%
    MacroF1: 55.56%
    MacroF1: 48.89%

    """

    def __init__(self, cm=None):
        super().__init__(beta=1.0, cm=cm)


class MicroF1(MicroFBeta):
    """Micro-average F1 score.

    This computes the F1 score by merging all the predictions and true labels, and then computes a
    global F1 score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 0]
    >>> y_pred = [0, 1, 1, 2, 1]

    >>> metric = metrics.MicroF1()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    MicroF1: 60.00%

    References
    ----------
    [^1]: [Why are precision, recall and F1 score equal when using micro averaging in a multi-class problem?](https://simonhessner.de/why-are-precision-recall-and-f1-score-equal-when-using-micro-averaging-in-a-multi-class-problem/)

    """

    def __init__(self, cm=None):
        super().__init__(beta=1.0, cm=cm)


class WeightedF1(WeightedFBeta):
    """Weighted-average F1 score.

    This works by computing the F1 score per class, and then performs a global weighted average by
    using the support of each class.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.WeightedF1()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    WeightedF1: 100.00%
    WeightedF1: 33.33%
    WeightedF1: 55.56%
    WeightedF1: 66.67%
    WeightedF1: 61.33%

    """

    def __init__(self, cm=None):
        super().__init__(beta=1.0, cm=cm)


# river/metrics/fowlkes_mallows.py

from __future__ import annotations

import math

from river import metrics

__all__ = ["FowlkesMallows"]


class FowlkesMallows(metrics.base.MultiClassMetric):
    r"""Fowlkes-Mallows Index.

    The Fowlkes-Mallows Index [^1] [^2] is an external evaluation method that is
    used to determine the similarity between two clusterings, and also a metric
    to measure confusion matrices. The measure of similarity could be either between
    two hierarchical clusterings or a clustering and a benchmark classification. A
    higher value for the Fowlkes-Mallows index indicates a greater similarity between
    the clusters and the benchmark classifications.

    The Fowlkes-Mallows Index, for two cluster algorithms, is defined as:

    $$
    FM = \sqrt{PPV \times TPR} = \sqrt{\frac{TP}{TP+FP} \times \frac{TP}{TP+FN}}
    $$

    where

    * TP, FP, FN are respectively the number of true positives, false positives and
    false negatives;

    * TPR is the True Positive Rate (or Sensitivity/Recall), PPV is the Positive Predictive
    Rate (or Precision).

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------
    >>> from river import metrics

    >>> y_true = [0, 0, 0, 1, 1, 1]
    >>> y_pred = [0, 0, 1, 1, 2, 2]

    >>> metric = metrics.FowlkesMallows()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    FowlkesMallows: 0.00%
    FowlkesMallows: 100.00%
    FowlkesMallows: 57.74%
    FowlkesMallows: 40.82%
    FowlkesMallows: 35.36%
    FowlkesMallows: 47.14%

    References
    ----------
    [^1]: Wikipedia contributors. (2020, December 22).
          Fowlkes–Mallows index. In Wikipedia, The Free Encyclopedia,
          from https://en.wikipedia.org/w/index.php?title=Fowlkes%E2%80%93Mallows_index&oldid=995714222
    [^2]: E. B. Fowkles and C. L. Mallows (1983).
          “A method for comparing two hierarchical clusterings”.
          Journal of the American Statistical Association

    """

    @property
    def works_with_weights(self):
        return False

    def get(self):
        n = self.cm.n_samples
        tk = sum(c * c for row in self.cm.data.values() for c in row.values()) - n
        pk = sum(sc * sc for sc in self.cm.sum_col.values()) - n
        qk = sum(sr * sr for sr in self.cm.sum_row.values()) - n

        try:
            return math.sqrt(tk / pk) * math.sqrt(tk / qk)
        except ZeroDivisionError:
            return 0.0


# river/metrics/geometric_mean.py

from __future__ import annotations

import numpy as np
from scipy import stats

from river import metrics

__all__ = ["GeometricMean"]


class GeometricMean(metrics.base.MultiClassMetric):
    r"""Geometric mean score.

    The geometric mean is a good indicator of a classifier's performance in the presence of class
    imbalance because it is independent of the distribution of examples between classes. This
    implementation computes the geometric mean of class-wise sensitivity (recall).

    $$
    gm = \sqrt[n]{s_1\cdot s_2\cdot s_3\cdot \ldots\cdot s_n}
    $$

    where $s_i$ is the sensitivity (recall) of class $i$ and $n$ is the
    number of classes.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing a
        confusion matrix reduces the amount of storage and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = ['cat', 'ant', 'cat', 'cat', 'ant', 'bird', 'bird']
    >>> y_pred = ['ant', 'ant', 'cat', 'cat', 'ant', 'cat', 'bird']

    >>> metric = metrics.GeometricMean()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    GeometricMean: 69.34%

    References
    ----------
    [^1]: Barandela, R. et al. “Strategies for learning in class imbalance problems”, Pattern Recognition, 36(3), (2003), pp 849-851.

    """

    def get(self):
        if len(self.cm.classes) > 0:
            sensitivity_per_class = np.zeros(len(self.cm.classes), float)
            for i, c in enumerate(self.cm.classes):
                try:
                    sensitivity_per_class[i] = self.cm[c][c] / self.cm.sum_row[c]
                except ZeroDivisionError:
                    continue
            with np.errstate(divide="ignore", invalid="ignore"):
                return stats.gmean(sensitivity_per_class)
        return 0.0


# river/metrics/jaccard.py

from __future__ import annotations

from river import metrics

__all__ = ["Jaccard", "MacroJaccard", "MicroJaccard", "WeightedJaccard"]


class Jaccard(metrics.base.BinaryMetric):
    """Jaccard score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing
        a confusion matrix reduces the amount of storage and computation time.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [False, True, True]
    >>> y_pred = [True, True, True]

    >>> metric = metrics.Jaccard()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    Jaccard: 0.00%
    Jaccard: 50.00%
    Jaccard: 66.67%

    References
    ----------
    [^1]: [Jaccard index](https://www.wikiwand.com/en/Jaccard_index)

    """

    def get(self):
        tp = self.cm.true_positives(self.pos_val)
        fp = self.cm.false_positives(self.pos_val)
        fn = self.cm.false_negatives(self.pos_val)
        try:
            return tp / (tp + fp + fn)
        except ZeroDivisionError:
            return 0.0


class MacroJaccard(metrics.base.MultiClassMetric):
    """Macro-average Jaccard score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MacroJaccard()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MacroJaccard: 100.00%
    MacroJaccard: 25.00%
    MacroJaccard: 50.00%
    MacroJaccard: 50.00%
    MacroJaccard: 38.89%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                tp = self.cm.true_positives(c)
                fp = self.cm.false_positives(c)
                fn = self.cm.false_negatives(c)
                total += tp / (tp + fp + fn)
            except ZeroDivisionError:
                continue
        try:
            return total / len(self.cm.classes)
        except ZeroDivisionError:
            return 0.0


class MicroJaccard(metrics.base.MultiClassMetric):
    """Micro-average Jaccard score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MicroJaccard()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MicroJaccard: 100.00%
    MicroJaccard: 33.33%
    MicroJaccard: 50.00%
    MicroJaccard: 60.00%
    MicroJaccard: 42.86%

    """

    def get(self):
        tp = self.cm.total_true_positives
        fp = self.cm.total_false_positives
        fn = self.cm.total_false_negatives

        try:
            return tp / (tp + fp + fn)
        except ZeroDivisionError:
            return 0.0


class WeightedJaccard(metrics.base.MultiClassMetric):
    """Weighted average Jaccard score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.WeightedJaccard()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    WeightedJaccard: 100.00%
    WeightedJaccard: 25.00%
    WeightedJaccard: 50.00%
    WeightedJaccard: 62.50%
    WeightedJaccard: 50.00%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                tp = self.cm.true_positives(c)
                fp = self.cm.false_positives(c)
                fn = self.cm.false_negatives(c)
                total += self.cm.support(c) * tp / (tp + fp + fn)
            except ZeroDivisionError:
                continue
        try:
            return total / self.cm.total_weight
        except ZeroDivisionError:
            return 0.0


# river/metrics/kappa.py

from __future__ import annotations

from river import metrics

__all__ = ["CohenKappa"]


class CohenKappa(metrics.base.MultiClassMetric):
    r"""Cohen's Kappa score.

    Cohen's Kappa expresses the level of agreement between two annotators on a classification
    problem. It is defined as

    $$
    \kappa = (p_o - p_e) / (1 - p_e)
    $$

    where $p_o$ is the empirical probability of agreement on the label
    assigned to any sample (prequential accuracy), and $p_e$ is
    the expected agreement when both annotators assign labels randomly.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = ['cat', 'ant', 'cat', 'cat', 'ant', 'bird']
    >>> y_pred = ['ant', 'ant', 'cat', 'cat', 'ant', 'cat']

    >>> metric = metrics.CohenKappa()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    CohenKappa: 42.86%

    References
    ----------
    [^1]: J. Cohen (1960). "A coefficient of agreement for nominal scales". Educational and Psychological Measurement 20(1):37-46. doi:10.1177/001316446002000104.

    """

    def get(self):
        try:
            p0 = self.cm.total_true_positives / self.cm.n_samples  # same as accuracy
        except ZeroDivisionError:
            p0 = 0

        pe = 0

        for c in self.cm.classes:
            estimation_row = self.cm.sum_row[c] / self.cm.n_samples
            estimation_col = self.cm.sum_col[c] / self.cm.n_samples
            pe += estimation_row * estimation_col

        try:
            return (p0 - pe) / (1 - pe)
        except ZeroDivisionError:
            return 0.0


# river/metrics/log_loss.py

from __future__ import annotations

import math

from river import metrics

__all__ = ["LogLoss"]


class LogLoss(metrics.base.MeanMetric, metrics.base.BinaryMetric):
    """Binary logarithmic loss.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [True, False, False, True]
    >>> y_pred = [0.9,  0.1,   0.2,   0.65]

    >>> metric = metrics.LogLoss()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)
    ...     print(metric.get())
    0.105360
    0.105360
    0.144621
    0.216161

    >>> metric
    LogLoss: 0.216162

    """

    _fmt = ""

    @property
    def bigger_is_better(self):
        return False

    @property
    def requires_labels(self):
        return False

    def _eval(self, y_true, y_pred):
        p_true = y_pred.get(True, 0.0) if isinstance(y_pred, dict) else y_pred
        p_true = self._clamp_proba(p_true)
        if y_true:
            return -math.log(p_true)
        return -math.log(1 - p_true)


# river/metrics/mae.py

from __future__ import annotations

from river import metrics

__all__ = ["MAE"]


class MAE(metrics.base.MeanMetric, metrics.base.RegressionMetric):
    """Mean absolute error.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]

    >>> metric = metrics.MAE()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    0.5
    0.5
    0.333
    0.5

    >>> metric
    MAE: 0.5

    """

    def _eval(self, y_true, y_pred):
        return abs(y_true - y_pred)


# river/metrics/mape.py

from __future__ import annotations

from river import metrics

__all__ = ["MAPE"]


class MAPE(metrics.base.MeanMetric, metrics.base.RegressionMetric):
    """Mean absolute percentage error.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]

    >>> metric = metrics.MAPE()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    MAPE: 32.738095

    """

    def _eval(self, y_true, y_pred):
        if y_true == 0:
            return 0.0
        return abs(y_true - y_pred) / abs(y_true)

    def get(self):
        return 100 * super().get()


# river/metrics/mcc.py

from __future__ import annotations

import math

from river import metrics

__all__ = ["MCC"]


class MCC(metrics.base.BinaryMetric):
    """Matthews correlation coefficient.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing a
        confusion matrix reduces the amount of storage and computation time.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [True, True, True, False]
    >>> y_pred = [True, False, True, True]

    >>> mcc = metrics.MCC()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     mcc = mcc.update(yt, yp)

    >>> mcc
    MCC: -0.333333

    References
    ----------
    [^1]: [Wikipedia article](https://www.wikiwand.com/en/Matthews_correlation_coefficient)

    """

    _fmt = ""

    def get(self):
        tp = self.cm.true_positives(self.pos_val)
        tn = self.cm.true_negatives(self.pos_val)
        fp = self.cm.false_positives(self.pos_val)
        fn = self.cm.false_negatives(self.pos_val)

        n = (tp + tn + fp + fn) or 1
        s = (tp + fn) / n
        p = (tp + fp) / n

        try:
            return (tp / n - s * p) / math.sqrt(p * s * (1 - s) * (1 - p))
        except ZeroDivisionError:
            return 0.0


# river/metrics/mse.py

from __future__ import annotations

import math

from river import metrics

__all__ = ["MSE", "RMSE", "RMSLE"]


class MSE(metrics.base.MeanMetric, metrics.base.RegressionMetric):
    """Mean squared error.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]

    >>> metric = metrics.MSE()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    0.25
    0.25
    0.1666
    0.375

    """

    def _eval(self, y_true, y_pred):
        return (y_true - y_pred) ** 2


class RMSE(MSE):
    """Root mean squared error.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]

    >>> metric = metrics.RMSE()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    0.5
    0.5
    0.408248
    0.612372

    >>> metric
    RMSE: 0.612372

    """

    def get(self):
        return super().get() ** 0.5


class RMSLE(RMSE):
    """Root mean squared logarithmic error.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]

    >>> metric = metrics.RMSLE()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    RMSLE: 0.357826

    """

    def update(self, y_true, y_pred, sample_weight=1.0):
        return super().update(math.log(y_true + 1), math.log(y_pred + 1), sample_weight)


# river/metrics/mutual_info.py

from __future__ import annotations

import math

import numpy as np

from river import metrics
from river.metrics.expected_mutual_info import expected_mutual_info

__all__ = [
    "AdjustedMutualInfo",
    "MutualInfo",
    "NormalizedMutualInfo",
]


class MutualInfo(metrics.base.MultiClassMetric):
    r"""Mutual Information between two clusterings.

    The Mutual Information [^1] is a measure of the similarity between two labels of
    the same data. Where $|U_i|$ is the number of samples in cluster $U_i$ and $|V_j|$
    is the number of the samples in cluster $V_j$, the Mutual Information between
    clusterings $U$ and $V$ can be calculated as:

    $$
    MI(U,V) = \sum_{i=1}^{|U|} \sum_{v=1}^{|V|} \frac{|U_i \cup V_j|}{N} \log \frac{N |U_i \cup V_j|}{|U_i| |V_j|}
    $$

    This metric is independent of the absolute values of the labels: a permutation
    of the class or cluster label values won't change the score.

    This metric is furthermore symmetric: switching `y_true` and `y_pred` will return
    the same score value. This can be useful to measure the agreement of two independent
    label assignments strategies on the same dataset when the real ground truth is
    not known.

    The Mutual Information can be equivalently expressed as:

    $$
    MI(U,V) = H(U) - H(U | V) = H(V) - H(V | U)
    $$

    where $H(U)$ and $H(V)$ are the marginal entropies, $H(U | V)$ and $H(V | U)$ are the
    conditional entropies.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [1, 1, 2, 2, 3, 3]
    >>> y_pred = [1, 1, 1, 2, 2, 2]

    >>> metric = metrics.MutualInfo()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    0.0
    0.0
    0.0
    0.215761
    0.395752
    0.462098

    >>> metric
    MutualInfo: 0.462098

    References
    ----------
    [^1]: Wikipedia contributors. (2021, March 17). Mutual information.
          In Wikipedia, The Free Encyclopedia,
          from https://en.wikipedia.org/w/index.php?title=Mutual_information&oldid=1012714929
    """

    _fmt = ""

    @property
    def works_with_weights(self):
        return False

    def get(self):
        mutual_info_score = 0.0

        for i in self.cm.classes:
            for j in self.cm.classes:
                try:
                    temp = (
                        self.cm[i][j]
                        / self.cm.n_samples
                        * (
                            math.log(self.cm.n_samples * self.cm[i][j])
                            - math.log(self.cm.sum_row[i] * self.cm.sum_col[j])
                        )
                    )
                except (ValueError, ZeroDivisionError):
                    continue
                temp = 0.0 if (abs(temp) < np.finfo("float64").eps) else temp
                # temp = 0.0 if temp < 0.0 else temp   # TODO confirm if we need to clip here
                mutual_info_score += temp
        return mutual_info_score


class NormalizedMutualInfo(metrics.base.MultiClassMetric):
    r"""Normalized Mutual Information between two clusterings.

    Normalized Mutual Information (NMI) is a normalized version of the Mutual Information (MI) score
    to scale the results between the range of 0 (no mutual information) and 1 (perfectly mutual
    information). In the formula, the mutual information will be normalized by a generalized mean of
    the entropy of true and predicted labels, defined by the `average_method`.

    We note that this measure is not adjusted for chance (i.e corrected the effect of result
    agreement solely due to chance); as a result, the Adjusted Mutual Info Score will mostly be preferred.
    However, this metric is still symmetric, which means that switching true and predicted labels will not
    alter the score value. This fact can be useful when the metric is used to measure the agreement between
    two indepedent label solutions on the same dataset, when the ground truth remains unknown.

    Another advantage of the metric is that as it is based on the calculation of entropy-related measures,
    it is independent of the permutation of class/cluster labels.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    average_method
        This parameter defines how to compute the normalizer in the denominator.
        Possible options include `min`, `max`, `arithmetic` and `geometric`.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [1, 1, 2, 2, 3, 3]
    >>> y_pred = [1, 1, 1, 2, 2, 2]

    >>> metric = metrics.NormalizedMutualInfo()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    1.0
    1.0
    0.0
    0.343711
    0.458065
    0.515803

    >>> metric
    NormalizedMutualInfo: 0.515804

    References
    ----------
    [^1]: Wikipedia contributors. (2021, March 17). Mutual information.
          In Wikipedia, The Free Encyclopedia,
          from https://en.wikipedia.org/w/index.php?title=Mutual_information&oldid=1012714929
    """
    _AVERAGE_MIN = "min"
    _AVERAGE_MAX = "max"
    _AVERAGE_GEOMETRIC = "geometric"
    _AVERAGE_ARITHMETIC = "arithmetic"
    _VALID_AVERAGE = [
        _AVERAGE_MIN,
        _AVERAGE_MAX,
        _AVERAGE_GEOMETRIC,
        _AVERAGE_ARITHMETIC,
    ]

    _fmt = ""

    def __init__(self, cm=None, average_method="arithmetic"):
        super().__init__(cm)
        if average_method not in self._VALID_AVERAGE:
            raise ValueError(
                f"Valid 'average_methods' are {self._VALID_AVERAGE}, "
                f"but {average_method} was passed."
            )
        self.average_method = average_method
        if average_method == self._AVERAGE_MIN:
            self._generalized_average = min
        elif average_method == self._AVERAGE_MAX:
            self._generalized_average = max
        elif average_method == self._AVERAGE_GEOMETRIC:
            self._generalized_average = _average_geometric
        else:  # average_method == self._AVERAGE_ARITHMETIC
            self._generalized_average = _average_arithmetic

    @property
    def works_with_weights(self):
        return False

    def get(self):
        n_classes = len([i for i in self.cm.sum_row.values() if i != 0])
        n_clusters = len([i for i in self.cm.sum_col.values() if i != 0])

        if (n_classes == n_clusters == 1) or (n_classes == n_clusters == 0):
            return 1.0

        mutual_info_score = metrics.MutualInfo(self.cm).get()

        entropy_true = _entropy(cm=self.cm, y_true=True)
        entropy_pred = _entropy(cm=self.cm, y_true=False)

        normalizer = self._generalized_average(entropy_true, entropy_pred)

        normalizer = max(normalizer, np.finfo("float64").eps)

        return mutual_info_score / normalizer


class AdjustedMutualInfo(metrics.base.MultiClassMetric):
    r"""Adjusted Mutual Information between two clusterings.

    Adjusted Mutual Information (AMI) is an adjustment of the Mutual Information score
    that accounts for chance. It corrects the effect of agreement solely due to chance
    between clusterings, similar to the way the Adjusted Rand Index corrects the Rand Index.
    It is closely related to variation of information. The adjusted measure, however, is
    no longer metrical.

    For two clusterings $U$ and $V$, the Adjusted Mutual Information is calculated as:

    $$
    AMI(U, V) = \frac{MI(U, V) - E(MI(U, V))}{avg(H(U), H(V)) - E(MI(U, V))}
    $$

    This metric is independent of the permutation of the class or cluster label values;
    furthermore, it is also symmetric. This can be useful to measure the agreement of
    two label assignments strategies on the same dataset, regardless of the ground truth.

    However, due to the complexity of the Expected Mutual Info Score, the computation of
    this metric is an order of magnitude slower than most other metrics, in general.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    average_method
        This parameter defines how to compute the normalizer in the denominator.
        Possible options include `min`, `max`, `arithmetic` and `geometric`.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [1, 1, 2, 2, 3, 3]
    >>> y_pred = [1, 1, 1, 2, 2, 2]

    >>> metric = metrics.AdjustedMutualInfo()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    1.0
    1.0
    0.0
    0.0
    0.105891
    0.298792

    >>> metric
    AdjustedMutualInfo: 0.298792

    References
    ----------
    [^1]: Wikipedia contributors. (2021, March 17). Mutual information.
          In Wikipedia, The Free Encyclopedia,
          from https://en.wikipedia.org/w/index.php?title=Mutual_information&oldid=1012714929
    """
    _AVERAGE_MIN = "min"
    _AVERAGE_MAX = "max"
    _AVERAGE_GEOMETRIC = "geometric"
    _AVERAGE_ARITHMETIC = "arithmetic"
    _VALID_AVERAGE = [
        _AVERAGE_MIN,
        _AVERAGE_MAX,
        _AVERAGE_GEOMETRIC,
        _AVERAGE_ARITHMETIC,
    ]

    _fmt = ""

    def __init__(self, cm=None, average_method="arithmetic"):
        super().__init__(cm)
        self.average_method = average_method
        if average_method not in self._VALID_AVERAGE:
            raise ValueError(
                f"Valid 'average_methods' are {self._VALID_AVERAGE}, "
                f"but {average_method} was passed."
            )
        self.average_method = average_method
        if average_method == self._AVERAGE_MIN:
            self._generalized_average = min
        elif average_method == self._AVERAGE_MAX:
            self._generalized_average = max
        elif average_method == self._AVERAGE_GEOMETRIC:
            self._generalized_average = _average_geometric
        else:  # average_method == self._AVERAGE_ARITHMETIC
            self._generalized_average = _average_arithmetic

    @property
    def works_with_weights(self):
        return False

    def get(self):
        n_classes = len([i for i in self.cm.sum_row.values() if i != 0])
        n_clusters = len([i for i in self.cm.sum_col.values() if i != 0])

        if (n_classes == n_clusters == 1) or (n_classes == n_clusters == 0):
            return 1.0

        mutual_info_score = metrics.MutualInfo(self.cm).get()

        expected_mutual_info_score = expected_mutual_info(self.cm)

        entropy_true = _entropy(cm=self.cm, y_true=True)
        entropy_pred = _entropy(cm=self.cm, y_true=False)

        normalizer = self._generalized_average(entropy_true, entropy_pred)

        denominator = normalizer - expected_mutual_info_score

        if denominator < 0:
            denominator = min(denominator, -np.finfo("float64").eps)
        else:
            denominator = max(denominator, np.finfo("float64").eps)

        adjusted_mutual_info_score = (mutual_info_score - expected_mutual_info_score) / denominator

        return adjusted_mutual_info_score


def _entropy(cm, y_true):
    n_samples = cm.n_samples
    if n_samples == 0:
        return 1.0

    if y_true:
        values = cm.sum_row
    else:
        values = cm.sum_col
    entropy = 0.0
    for i in cm.classes:
        if i in values and values[i] > 0:
            entropy -= (values[i] / n_samples) * (np.log(values[i]) - np.log(n_samples))
    return entropy


def _average_geometric(u, v):
    return math.sqrt(u * v)


def _average_arithmetic(u, v):
    return (u + v) / 2


# river/metrics/precision.py

from __future__ import annotations

from river import metrics

__all__ = ["MacroPrecision", "MicroPrecision", "Precision", "WeightedPrecision"]


class Precision(metrics.base.BinaryMetric):
    """Binary precision score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing
        a confusion matrix reduces the amount of storage and computation time.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [True, False, True, True, True]
    >>> y_pred = [True, True, False, True, True]

    >>> metric = metrics.Precision()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    Precision: 100.00%
    Precision: 50.00%
    Precision: 50.00%
    Precision: 66.67%
    Precision: 75.00%

    """

    def get(self):
        tp = self.cm.true_positives(self.pos_val)
        fp = self.cm.false_positives(self.pos_val)
        try:
            return tp / (tp + fp)
        except ZeroDivisionError:
            return 0.0


class MacroPrecision(metrics.base.MultiClassMetric):
    """Macro-average precision score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MacroPrecision()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MacroPrecision: 100.00%
    MacroPrecision: 25.00%
    MacroPrecision: 50.00%
    MacroPrecision: 50.00%
    MacroPrecision: 50.00%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                tp = self.cm.true_positives(c)
                fp = self.cm.false_positives(c)
                total += tp / (tp + fp)
            except ZeroDivisionError:
                continue
        try:
            return total / len(self.cm.classes)
        except ZeroDivisionError:
            return 0.0


class MicroPrecision(metrics.base.MultiClassMetric):
    """Micro-average precision score.

    The micro-average precision score is exactly equivalent to the micro-average recall as well as
    the micro-average F1 score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MicroPrecision()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MicroPrecision: 100.00%
    MicroPrecision: 50.00%
    MicroPrecision: 66.67%
    MicroPrecision: 75.00%
    MicroPrecision: 60.00%

    References
    ----------
    [^1]: [Why are precision, recall and F1 score equal when using micro averaging in a multi-class problem?](https://simonhessner.de/why-are-precision-recall-and-f1-score-equal-when-using-micro-averaging-in-a-multi-class-problem/)

    """

    def get(self):
        tp = self.cm.total_true_positives
        fp = self.cm.total_false_positives
        try:
            return tp / (tp + fp)
        except ZeroDivisionError:
            return 0.0


class WeightedPrecision(metrics.base.MultiClassMetric):
    """Weighted-average precision score.

    This uses the support of each label to compute an average score, whereas
    `metrics.MacroPrecision` ignores the support.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.WeightedPrecision()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    WeightedPrecision: 100.00%
    WeightedPrecision: 25.00%
    WeightedPrecision: 50.00%
    WeightedPrecision: 62.50%
    WeightedPrecision: 70.00%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                tp = self.cm.true_positives(c)
                fp = self.cm.false_positives(c)
                total += self.cm.support(c) * tp / (tp + fp)
            except ZeroDivisionError:
                continue
        try:
            return total / self.cm.total_weight
        except ZeroDivisionError:
            return 0.0


# river/metrics/r2.py

from __future__ import annotations

from river import metrics, stats

__all__ = ["R2"]


class R2(metrics.base.RegressionMetric):
    """Coefficient of determination ($R^2$) score

    The coefficient of determination, denoted $R^2$ or $r^2$, is the proportion
    of the variance in the dependent variable that is predictable from the
    independent variable(s). [^1]

    Best possible score is 1.0 and it can be negative (because the model can be
    arbitrarily worse). A constant model that always predicts the expected
    value of $y$, disregarding the input features, would get a $R^2$ score of
    0.0.

    $R^2$ is not defined when less than 2 samples have been observed. This
    implementation returns 0.0 in this case.

    Examples
    --------
    >>> from river import metrics

    >>> y_true = [3, -0.5, 2, 7]
    >>> y_pred = [2.5, 0.0, 2, 8]

    >>> metric = metrics.R2()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    0.0
    0.9183
    0.9230
    0.9486

    References
    ----------
    [^1]: [Coefficient of determination (Wikipedia)](https://en.wikipedia.org/wiki/Coefficient_of_determination)

    """

    def __init__(self):
        self._y_var = stats.Var()
        self._total_sum_of_squares = 0
        self._residual_sum_of_squares = 0

    @property
    def bigger_is_better(self):
        return True

    def update(self, y_true, y_pred, sample_weight=1.0):
        self._y_var.update(y_true, w=sample_weight)
        squared_error = (y_true - y_pred) * (y_true - y_pred) * sample_weight
        self._residual_sum_of_squares += squared_error
        return self

    def revert(self, y_true, y_pred, sample_weight=1.0):
        self._y_var.update(y_true, w=-sample_weight)
        self._residual_sum_of_squares -= (y_true - y_pred) * (y_true - y_pred) * sample_weight
        return self

    def get(self):
        if self._y_var.mean.n > 1:
            try:
                total_sum_of_squares = (self._y_var.mean.n - 1) * self._y_var.get()
                return 1 - (self._residual_sum_of_squares / total_sum_of_squares)
            except ZeroDivisionError:
                return 0.0

        # Not defined for n_samples < 2
        return 0.0


# river/metrics/rand.py

from __future__ import annotations

import collections

from river import metrics

__all__ = ["AdjustedRand", "Rand"]


def _pair_confusion(cm):
    pair_confusion_matrix = {i: collections.defaultdict(int) for i in range(2)}

    sum_squares = 0
    false_positives = 0
    false_negatives = 0

    for i in cm.classes:
        for j in cm.classes:
            sum_squares += cm[i][j] * cm[i][j]
            false_positives += cm[i][j] * cm.sum_col[j]
            false_negatives += cm[j][i] * cm.sum_row[j]

    true_positives = sum_squares - cm.n_samples

    false_positives -= sum_squares

    false_negatives -= sum_squares

    true_negatives = cm.n_samples * cm.n_samples - (false_positives + false_negatives) - sum_squares

    pair_confusion_matrix[0][0] = true_negatives
    pair_confusion_matrix[0][1] = false_positives
    pair_confusion_matrix[1][0] = false_negatives
    pair_confusion_matrix[1][1] = true_positives

    return pair_confusion_matrix


class Rand(metrics.base.MultiClassMetric):
    """Rand Index.

    The Rand Index [^1] [^2] is a measure of the similarity between two data clusterings.
    Given a set of elements `S` and two partitions of `S` to compare, `X` and `Y`,
    define the following:

    * a, the number of pairs of elements in `S` that are in the **same** subset in `X`
    and in the **same** subset in `Y`

    * b, the number of pairs of elements in `S` that are in the **different** subset in `X`
    and in **different** subsets in `Y`

    * c, the number of pairs of elements in `S` that are in the **same** subset in `X`
    and in **different** subsets in `Y`

    * d, the number of pairs of elements in `S` that are in the **different** subset in `X`
    and in the **same** subset in `Y`

    The Rand index, R, is

    $$
    R = \frac{a+b}{a+b+c+d} = \frac{a+b}{\frac{n(n-1)}{2}}.
    $$

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------
    >>> from river import metrics

    >>> y_true = [0, 0, 0, 1, 1, 1]
    >>> y_pred = [0, 0, 1, 1, 2, 2]

    >>> metric = metrics.Rand()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    Rand: 0.666667

    References
    ----------
    [^1]: Wikipedia contributors. (2021, January 13). Rand index.
          In Wikipedia, The Free Encyclopedia,
          from https://en.wikipedia.org/w/index.php?title=Rand_index&oldid=1000098911
    [^2]: W. M. Rand (1971). "Objective criteria for the evaluation of clustering methods".
          Journal of the American Statistical Association. American Statistical Association.
          66 (336): 846–850. arXiv:1704.01036. doi:10.2307/2284239. JSTOR 2284239.

    """

    _fmt = ""

    @property
    def works_with_weights(self):
        return False

    def get(self):
        pair_confusion_matrix = _pair_confusion(self.cm)

        true_positives = pair_confusion_matrix[1][1]
        true_negatives = pair_confusion_matrix[0][0]

        total_pairs = self.cm.n_samples * (self.cm.n_samples - 1)

        try:
            return (true_positives + true_negatives) / total_pairs
        except ZeroDivisionError:
            return 1.0


class AdjustedRand(metrics.base.MultiClassMetric):
    """Adjusted Rand Index.

    The Adjusted Rand Index is the corrected-for-chance version of the Rand Index [^1] [^2].
    Such a correction for chance establishes a baseline by using the expected similarity
    of all pair-wise comparisions between clusterings specified by a random model.

    Traditionally, the Rand Index was corrected using the Permutation Model for Clustering.
    However, the premises of the permutation model are frequently violated; in many
    clustering scenarios, either the number of clusters or the size distribution of those
    clusters vary drastically. Variations of the adjusted Rand Index account for different
    models of random clusterings.

    Though the Rand Index may only yield a value between 0 and 1, the Adjusted Rand index
    can yield negative values if the index is less than the expected index.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------
    >>> from river import metrics

    >>> y_true = [0, 0, 0, 1, 1, 1]
    >>> y_pred = [0, 0, 1, 1, 2, 2]

    >>> metric = metrics.AdjustedRand()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    1.0
    1.0
    0.0
    0.0
    0.09090909090909091
    0.24242424242424243

    >>> metric
    AdjustedRand: 0.242424

    References
    ----------
    [^1]: Wikipedia contributors. (2021, January 13). Rand index.
          In Wikipedia, The Free Encyclopedia,
          from https://en.wikipedia.org/w/index.php?title=Rand_index&oldid=1000098911
    [^2]: W. M. Rand (1971). "Objective criteria for the evaluation of clustering methods".
          Journal of the American Statistical Association. American Statistical Association.
          66 (336): 846–850. arXiv:1704.01036. doi:10.2307/2284239. JSTOR 2284239.

    """

    _fmt = ""

    @property
    def works_with_weights(self):
        return False

    def get(self):
        pair_confusion_matrix = _pair_confusion(self.cm)

        true_negatives, false_positives = pair_confusion_matrix[0].values()
        false_negatives, true_positives = pair_confusion_matrix[1].values()

        try:
            return (
                2.0
                * (true_positives * true_negatives - false_negatives * false_positives)
                / (
                    (true_positives + false_negatives) * (false_negatives + true_negatives)
                    + (true_positives + false_positives) * (false_positives + true_negatives)
                )
            )
        except ZeroDivisionError:
            return 1.0


# river/metrics/recall.py

from __future__ import annotations

from river import metrics

__all__ = ["MacroRecall", "MicroRecall", "Recall", "WeightedRecall"]


class Recall(metrics.base.BinaryMetric):
    """Binary recall score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing a
        confusion matrix reduces the amount of storage and computation time.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [True, False, True, True, True]
    >>> y_pred = [True, True, False, True, True]

    >>> metric = metrics.Recall()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    Recall: 100.00%
    Recall: 100.00%
    Recall: 50.00%
    Recall: 66.67%
    Recall: 75.00%

    """

    def get(self):
        tp = self.cm.true_positives(self.pos_val)
        fn = self.cm.false_negatives(self.pos_val)
        try:
            return tp / (tp + fn)
        except ZeroDivisionError:
            return 0.0


class MacroRecall(metrics.base.MultiClassMetric):
    """Macro-average recall score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing a
        confusion matrix reduces the amount of storage and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MacroRecall()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MacroRecall: 100.00%
    MacroRecall: 50.00%
    MacroRecall: 66.67%
    MacroRecall: 66.67%
    MacroRecall: 55.56%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                total += self.cm[c][c] / self.cm.sum_row[c]
            except ZeroDivisionError:
                continue
        try:
            return total / len(self.cm.classes)
        except ZeroDivisionError:
            return 0.0


class MicroRecall(metrics.MicroPrecision):
    """Micro-average recall score.

    The micro-average recall is exactly equivalent to the micro-average precision as well as the
    micro-average F1 score.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing a
        confusion matrix reduces the amount of storage and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.MicroRecall()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    MicroRecall: 100.00%
    MicroRecall: 50.00%
    MicroRecall: 66.67%
    MicroRecall: 75.00%
    MicroRecall: 60.00%

    References
    ----------
    [^1]: [Why are precision, recall and F1 score equal when using micro averaging in a multi-class problem?](https://simonhessner.de/why-are-precision-recall-and-f1-score-equal-when-using-micro-averaging-in-a-multi-class-problem/)

    """


class WeightedRecall(metrics.base.MultiClassMetric):
    """Weighted-average recall score.

    This uses the support of each label to compute an average score, whereas `MacroRecall`
    ignores the support.

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion matrix between multiple metrics. Sharing a
        confusion matrix reduces the amount of storage and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 1, 2, 2, 2]
    >>> y_pred = [0, 0, 2, 2, 1]

    >>> metric = metrics.WeightedRecall()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp))
    WeightedRecall: 100.00%
    WeightedRecall: 50.00%
    WeightedRecall: 66.67%
    WeightedRecall: 75.00%
    WeightedRecall: 60.00%

    """

    def get(self):
        total = 0
        for c in self.cm.classes:
            try:
                total += self.cm[c][c]
            except ZeroDivisionError:
                continue
        try:
            return total / self.cm.total_weight
        except ZeroDivisionError:
            return 0.0


# river/metrics/report.py

from __future__ import annotations

from river import metrics, utils


class ClassificationReport(metrics.base.MultiClassMetric):
    """A report for monitoring a classifier.

    This class maintains a set of metrics and updates each of them every time `update` is called.
    You can print this class at any time during a model's lifetime to get a tabular visualization
    of various metrics.

    You can wrap a `metrics.ClassificationReport` with `utils.Rolling` in order to obtain a
    classification report over a window of observations. You can also wrap it with
    `utils.TimeRolling` to obtain a report over a period of time.

    Parameters
    ----------
    decimals
        The number of decimals to display in each cell.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = ['pear', 'apple', 'banana', 'banana', 'banana']
    >>> y_pred = ['apple', 'pear', 'banana', 'banana', 'apple']

    >>> report = metrics.ClassificationReport()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     report = report.update(yt, yp)

    >>> print(report)
                   Precision   Recall   F1       Support
    <BLANKLINE>
       apple       0.00%    0.00%    0.00%         1
      banana     100.00%   66.67%   80.00%         3
        pear       0.00%    0.00%    0.00%         1
    <BLANKLINE>
       Macro      33.33%   22.22%   26.67%
       Micro      40.00%   40.00%   40.00%
    Weighted      60.00%   40.00%   48.00%
    <BLANKLINE>
                     40.00% accuracy

    """

    def __init__(self, decimals=2, cm=None):
        super().__init__(cm)
        self.decimals = decimals
        self._f1s = {}
        self._macro_precision = metrics.MacroPrecision(self.cm)
        self._macro_recall = metrics.MacroRecall(self.cm)
        self._macro_f1 = metrics.MacroF1(self.cm)
        self._micro_precision = metrics.MicroPrecision(self.cm)
        self._micro_recall = metrics.MicroRecall(self.cm)
        self._micro_f1 = metrics.MicroF1(self.cm)
        self._weighted_precision = metrics.WeightedPrecision(self.cm)
        self._weighted_recall = metrics.WeightedRecall(self.cm)
        self._weighted_f1 = metrics.WeightedF1(self.cm)
        self._accuracy = metrics.Accuracy(self.cm)

    def get(self):
        raise NotImplementedError

    def __repr__(self):
        def fmt_pct(x):
            return f"{x:.{self.decimals}%}"

        headers = ["", "Precision", "Recall", "F1", "Support"]
        classes = sorted(self.cm.classes)

        for c in classes:
            if c not in self._f1s:
                self._f1s[c] = metrics.F1(cm=self.cm, pos_val=c)

        columns = [
            # Row names
            ["", *map(str, classes), "", "Macro", "Micro", "Weighted"],
            # Precision values
            [
                "",
                *[fmt_pct(self._f1s[c].precision.get()) for c in classes],
                "",
                *map(
                    fmt_pct,
                    [
                        self._macro_precision.get(),
                        self._micro_precision.get(),
                        self._weighted_precision.get(),
                    ],
                ),
            ],
            # Recall values
            [
                "",
                *[fmt_pct(self._f1s[c].recall.get()) for c in classes],
                "",
                *map(
                    fmt_pct,
                    [
                        self._macro_recall.get(),
                        self._micro_recall.get(),
                        self._weighted_recall.get(),
                    ],
                ),
            ],
            # F1 values
            [
                "",
                *[fmt_pct(self._f1s[c].get()) for c in classes],
                "",
                *map(
                    fmt_pct,
                    [
                        self._macro_f1.get(),
                        self._micro_f1.get(),
                        self._weighted_f1.get(),
                    ],
                ),
            ],
            # Support
            ["", *[str(self.cm.sum_row[c]).rstrip("0").rstrip(".") for c in classes], *[""] * 4],
        ]

        # Build the table
        table = utils.pretty.print_table(headers, columns)

        # Write down the accuracy
        width = len(table.splitlines()[0])
        accuracy = f"{self._accuracy.get():.{self.decimals}%}" + " accuracy"
        table += "\n\n" + f"{accuracy:^{width}}"

        return table


# river/metrics/roc_auc.py

from __future__ import annotations

from scipy import integrate

from river import metrics, utils

__all__ = ["ROCAUC"]


class ROCAUC(metrics.base.BinaryMetric):
    """Receiving Operating Characteristic Area Under the Curve.

    This metric is an approximation of the true ROC AUC. Computing the true ROC AUC would
    require storing all the predictions and ground truths, which isn't desirable. The approximation
    error is not significant as long as the predicted probabilities are well calibrated. In any
    case, this metric can still be used to reliably compare models between each other.

    Parameters
    ----------
    n_thresholds
        The number of thresholds used for discretizing the ROC curve. A higher value will lead to
        more accurate results, but will also cost more time and memory.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [ 0,  0,   1,  1]
    >>> y_pred = [.1, .4, .35, .8]

    >>> metric = metrics.ROCAUC()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    ROCAUC: 87.50%

    The true ROC AUC is in fact 0.75. We can improve the accuracy by increasing the amount
    of thresholds. This comes at the cost more computation time and more memory usage.

    >>> metric = metrics.ROCAUC(n_thresholds=20)

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    ROCAUC: 75.00%

    """

    def __init__(self, n_thresholds=10, pos_val=True):
        self.n_thresholds = n_thresholds
        self.pos_val = pos_val
        self.thresholds = [i / (n_thresholds - 1) for i in range(n_thresholds)]
        self.thresholds[0] -= 1e-7
        self.thresholds[-1] += 1e-7
        self.cms = [metrics.ConfusionMatrix() for _ in range(n_thresholds)]

    def works_with(self, model) -> bool:
        return (
            super().works_with(model)
            or utils.inspect.isanomalydetector(model)
            or utils.inspect.isanomalyfilter(model)
        )

    def update(self, y_true, y_pred, sample_weight=1.0):
        p_true = y_pred.get(True, 0.0) if isinstance(y_pred, dict) else y_pred
        for t, cm in zip(self.thresholds, self.cms):
            cm.update(y_true == self.pos_val, p_true > t, sample_weight)
        return self

    def revert(self, y_true, y_pred, sample_weight=1.0):
        p_true = y_pred.get(True, 0.0) if isinstance(y_pred, dict) else y_pred
        for t, cm in zip(self.thresholds, self.cms):
            cm.revert(y_true == self.pos_val, p_true > t, sample_weight)
        return self

    @property
    def requires_labels(self):
        return False

    def get(self):
        tprs = [0] * self.n_thresholds
        fprs = [0] * self.n_thresholds

        def safe_div(a, b):
            try:
                return a / b
            except ZeroDivisionError:
                return 0.0

        for i, cm in enumerate(self.cms):
            tp = cm.true_positives(self.pos_val)
            tn = cm.true_negatives(self.pos_val)
            fp = cm.false_positives(self.pos_val)
            fn = cm.false_negatives(self.pos_val)

            tprs[i] = safe_div(a=tp, b=tp + fn)
            fprs[i] = safe_div(a=fp, b=fp + tn)

        return -integrate.trapz(x=fprs, y=tprs)


# river/metrics/rolling_roc_auc.py

from __future__ import annotations

from river import metrics, utils

from .efficient_rollingrocauc import EfficientRollingROCAUC

__all__ = ["RollingROCAUC"]


class RollingROCAUC(metrics.base.BinaryMetric):
    """Rolling version of the Receiving Operating Characteristic Area Under the Curve.

    The RollingROCAUC calculates the metric using the instances in its window
    of size S. It keeps a queue of the instances, when an instance is added and
    the queue length is equal to S, the last instance is removed. The metric has
    a tree with ordered instances, in order to calculate the AUC efficiently.
    It was implemented based on the algorithm presented in Brzezinski and
    Stefanowski, 2017.

    The difference between this metric and the standard ROCAUC is that the latter
    calculates an approximation of the real metric considering all data from the
    beginning of the stream, while the RollingROCAUC calculates the exact value
    considering only the last S instances. This approach may be beneficial if
    it's necessary to evaluate the model's performance over time, since
    calculating the metric using the entire stream may hide the current
    performance of the classifier.

    Parameters
    ----------
    window_size
        The max length of the window.
    pos_val
        Value to treat as "positive".

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [ 0,  1,  0,  1,  0,  1,  0,  0,   1,  1]
    >>> y_pred = [.3, .5, .5, .7, .1, .3, .1, .4, .35, .8]

    >>> metric = metrics.RollingROCAUC(window_size=4)

    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    RollingROCAUC: 75.00%

    """

    def __init__(self, window_size=1000, pos_val=True):
        self.window_size = window_size
        self.pos_val = pos_val
        self.__metric = EfficientRollingROCAUC(pos_val, window_size)

    def works_with(self, model) -> bool:
        return (
            super().works_with(model)
            or utils.inspect.isanomalydetector(model)
            or utils.inspect.isanomalyfilter(model)
        )

    def update(self, y_true, y_pred):
        p_true = y_pred.get(True, 0.0) if isinstance(y_pred, dict) else y_pred
        self.__metric.update(y_true, p_true)
        return self

    def revert(self, y_true, y_pred):
        p_true = y_pred.get(True, 0.0) if isinstance(y_pred, dict) else y_pred
        self.__metric.revert(y_true, p_true)
        return self

    @property
    def requires_labels(self) -> bool:
        return False

    @property
    def works_with_weights(self) -> bool:
        return False

    def get(self):
        return self.__metric.get()


# river/metrics/silhouette.py

from __future__ import annotations

import math

from river import metrics, utils


class Silhouette(metrics.base.ClusteringMetric):
    """
    Silhouette coefficient [^1], roughly speaking, is the ratio between cohesion and the average distances
    from the points to their second-closest centroid. It rewards the clustering algorithm where
    points are very close to their assigned centroids and far from any other centroids,
    that is, clustering results with good cohesion and good separation.

    It rewards clusterings where points are very close to their assigned centroids and far from any other
    centroids, that is clusterings with good cohesion and good separation. [^2]

    The definition of Silhouette coefficient for online clustering evaluation is different from that of
    batch learning. It does not store information and calculate pairwise distances between all points at the
    same time, since the practice is too expensive for an incremental metric.

    Examples
    --------

    >>> from river import cluster
    >>> from river import stream
    >>> from river import metrics

    >>> X = [
    ...     [1, 2],
    ...     [1, 4],
    ...     [1, 0],
    ...     [4, 2],
    ...     [4, 4],
    ...     [4, 0],
    ...     [-2, 2],
    ...     [-2, 4],
    ...     [-2, 0]
    ... ]

    >>> k_means = cluster.KMeans(n_clusters=3, halflife=0.4, sigma=3, seed=0)
    >>> metric = metrics.Silhouette()

    >>> for x, _ in stream.iter_array(X):
    ...     k_means = k_means.learn_one(x)
    ...     y_pred = k_means.predict_one(x)
    ...     metric = metric.update(x, y_pred, k_means.centers)

    >>> metric
    Silhouette: 0.32145

    References
    ----------
    [^1]: Rousseeuw, P. (1987). Silhouettes: a graphical aid to the intepretation and validation
          of cluster analysis 20, 53 - 65. DOI: 10.1016/0377-0427(87)90125-7
    [^2]: Bifet, A. et al. (2018). "Machine Learning for Data Streams".
          DOI: 10.7551/mitpress/10654.001.0001.

    """

    def __init__(self):
        super().__init__()
        self._sum_distance_closest_centroid = 0
        self._sum_distance_second_closest_centroid = 0

    @staticmethod
    def _find_distance_second_closest_center(centers, x):
        distances = {i: math.sqrt(utils.math.minkowski_distance(centers[i], x, 2)) for i in centers}
        return sorted(distances.values())[-2]

    def update(self, x, y_pred, centers, sample_weight=1.0):
        distance_closest_centroid = math.sqrt(utils.math.minkowski_distance(centers[y_pred], x, 2))
        self._sum_distance_closest_centroid += distance_closest_centroid

        distance_second_closest_centroid = self._find_distance_second_closest_center(centers, x)
        self._sum_distance_second_closest_centroid += distance_second_closest_centroid

        return self

    def revert(self, x, y_pred, centers, sample_weight=1.0):
        distance_closest_centroid = math.sqrt(utils.math.minkowski_distance(centers[y_pred], x, 2))
        self._sum_distance_closest_centroid -= distance_closest_centroid

        distance_second_closest_centroid = self._find_distance_second_closest_center(centers, x)
        self._sum_distance_second_closest_centroid -= distance_second_closest_centroid

        return self

    def get(self):
        try:
            return self._sum_distance_closest_centroid / self._sum_distance_second_closest_centroid
        except ZeroDivisionError:
            return math.inf

    @property
    def bigger_is_better(self):
        return False


# river/metrics/smape.py

from __future__ import annotations

from river import metrics

__all__ = ["SMAPE"]


class SMAPE(metrics.base.MeanMetric, metrics.base.RegressionMetric):
    """Symmetric mean absolute percentage error.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [0, 0.07533, 0.07533, 0.07533, 0.07533, 0.07533, 0.07533, 0.0672, 0.0672]
    >>> y_pred = [0, 0.102, 0.107, 0.047, 0.1, 0.032, 0.047, 0.108, 0.089]

    >>> metric = metrics.SMAPE()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     metric = metric.update(yt, yp)

    >>> metric
    SMAPE: 37.869392

    """

    def _eval(self, y_true, y_pred):
        den = abs(y_true) + abs(y_pred)
        if den == 0:
            return 0.0
        return 2.0 * abs(y_true - y_pred) / den

    def get(self):
        return 100 * super().get()


# river/metrics/test_cross_entropy.py

from __future__ import annotations

import math

from sklearn import metrics as sk_metrics

from river import metrics


def test_cross_entropy():
    metric = metrics.CrossEntropy()

    y_true = [0, 1, 2, 2]
    y_pred = [
        [0.29450637, 0.34216758, 0.36332605],
        [0.21290077, 0.32728332, 0.45981591],
        [0.42860913, 0.33380113, 0.23758974],
        [0.44941979, 0.32962558, 0.22095463],
    ]

    for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
        yp = dict(enumerate(yp))
        metric.update(yt, yp)

        if i >= 1:
            assert math.isclose(
                metric.get(),
                sk_metrics.log_loss(y_true[: i + 1], y_pred[: i + 1], labels=[0, 1, 2]),
            )

    metric.revert(y_true[-1], dict(enumerate(y_pred[-1])))
    assert math.isclose(metric.get(), sk_metrics.log_loss(y_true[:-1], y_pred[:-1]))


# river/metrics/test_fbeta.py

from __future__ import annotations

import collections
import math

from sklearn import metrics as sk_metrics

from river import metrics, utils


def test_multi_fbeta():
    fbeta = metrics.MultiFBeta(betas={0: 0.25, 1: 1, 2: 4}, weights={0: 1, 1: 1, 2: 2})
    sk_fbeta = sk_metrics.fbeta_score
    y_true = [0, 1, 2, 2, 2]
    y_pred = [0, 1, 0, 2, 1]

    for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
        fbeta.update(yt, yp)

        if i >= 2:
            fbeta_0, _, _ = sk_fbeta(y_true[: i + 1], y_pred[: i + 1], beta=0.25, average=None)
            _, fbeta_1, _ = sk_fbeta(y_true[: i + 1], y_pred[: i + 1], beta=1, average=None)
            _, _, fbeta_2 = sk_fbeta(y_true[: i + 1], y_pred[: i + 1], beta=4, average=None)

            multi_fbeta = fbeta_0 * 1 + fbeta_1 * 1 + fbeta_2 * 2
            multi_fbeta /= 1 + 1 + 2

            assert math.isclose(fbeta.get(), multi_fbeta)


def test_rolling_multi_fbeta():
    def tail(iterable, n):
        return collections.deque(iterable, maxlen=n)

    fbeta = utils.Rolling(
        metrics.MultiFBeta(betas={0: 0.25, 1: 1, 2: 4}, weights={0: 1, 1: 1, 2: 2}),
        window_size=3,
    )
    n = fbeta.window_size
    sk_fbeta = sk_metrics.fbeta_score
    y_true = [0, 1, 2, 2, 2]
    y_pred = [0, 1, 0, 2, 1]

    for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
        fbeta.update(yt, yp)

        if i >= 2:
            sk_y_true, sk_y_pred = tail(y_true[: i + 1], n), tail(y_pred[: i + 1], n)
            fbeta_0, _, _ = sk_fbeta(sk_y_true, sk_y_pred, beta=0.25, average=None)
            _, fbeta_1, _ = sk_fbeta(sk_y_true, sk_y_pred, beta=1, average=None)
            _, _, fbeta_2 = sk_fbeta(sk_y_true, sk_y_pred, beta=4, average=None)

            multi_fbeta = fbeta_0 * 1 + fbeta_1 * 1 + fbeta_2 * 2
            multi_fbeta /= 1 + 1 + 2

            assert math.isclose(fbeta.get(), multi_fbeta)


# river/metrics/test_log_loss.py

from __future__ import annotations

import math

from sklearn import metrics as sk_metrics

from river import metrics


def test_log_loss():
    metric = metrics.LogLoss()

    y_true = [True, False, False, True]
    y_pred = [0.9, 0.1, 0.2, 0.65]

    for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
        metric.update(yt, yp)

        if i >= 1:
            assert math.isclose(metric.get(), sk_metrics.log_loss(y_true[: i + 1], y_pred[: i + 1]))

    metric.revert(y_true[-1], y_pred[-1])
    assert math.isclose(metric.get(), sk_metrics.log_loss(y_true[:-1], y_pred[:-1]))


# river/metrics/test_metrics.py

from __future__ import annotations

import collections
import copy
import functools
import importlib
import inspect
import pickle
import platform
import random

import numpy as np
import pytest
from sklearn import metrics as sk_metrics

from river import metrics, utils


def load_metrics():
    """Yields all the metrics."""

    for name, obj in inspect.getmembers(importlib.import_module("river.metrics"), inspect.isclass):
        if name == "Metrics":
            continue

        if inspect.isabstract(obj):
            continue

        elif name == "RegressionMultiOutput":
            yield obj(metric=metrics.MSE())
            continue

        try:
            sig = inspect.signature(obj)
            yield obj(
                **{
                    param.name: param.default if param.default != param.empty else 5
                    for param in sig.parameters.values()
                }
            )
        except ValueError:
            yield obj()


@pytest.mark.parametrize(
    "metric",
    [pytest.param(metric, id=type(metric).__name__) for metric in load_metrics()],
)
def test_pickling(metric):
    assert isinstance(pickle.loads(pickle.dumps(metric)), metric.__class__)
    assert isinstance(copy.deepcopy(metric), metric.__class__)


@pytest.mark.parametrize(
    "metric",
    [pytest.param(metric, id=type(metric).__name__) for metric in load_metrics()],
)
def test_repr(metric):
    repr(metric)

    for y_true, y_pred, _ in generate_test_cases(metric=metric, n=10):
        m = copy.deepcopy(metric)

        for yt, yp in zip(y_true, y_pred):
            m.update(yt, yp)
            repr(metric)


def generate_test_cases(metric, n):
    """Yields n (y_true, y_pred, sample_weight) triplets for a given metric."""

    sample_weights = [random.random() for _ in range(n)]

    if isinstance(metric, metrics.base.ClassificationMetric):
        y_true = [random.choice([False, True]) for _ in range(n)]
        if metric.requires_labels:
            y_pred = [random.choice([False, True]) for _ in range(n)]
        else:
            y_pred = [dict(zip([False, True], np.random.dirichlet(np.ones(2)))) for _ in range(n)]
        yield y_true, y_pred, sample_weights

    if isinstance(metric, metrics.base.MultiClassMetric):
        y_true = [random.choice([0, 1, 2]) for _ in range(n)]
        if metric.requires_labels:
            y_pred = [random.choice([0, 1, 2]) for _ in range(n)]
        else:
            y_pred = [dict(zip([0, 1, 2], np.random.dirichlet(np.ones(3)))) for _ in range(n)]
        yield y_true, y_pred, sample_weights

    if isinstance(metric, metrics.base.RegressionMetric):
        yield (
            [random.random() for _ in range(n)],
            [random.random() for _ in range(n)],
            sample_weights,
        )


def partial(f, **kwargs):
    return functools.update_wrapper(functools.partial(f, **kwargs), f)


def roc_auc_score(y_true, y_score):
    """
    This functions is a wrapper to the scikit-learn roc_auc_score function.
    It was created because the scikit version utilizes array of scores and
    may raise a ValueError if there is only one class present in y_true.
    This wrapper returns 0 if y_true has only one class and
    deals with the scores.
    """
    nonzero = np.count_nonzero(y_true)
    if nonzero == 0 or nonzero == len(y_true):
        return 0

    scores = [s[True] for s in y_score]

    return sk_metrics.roc_auc_score(y_true, scores)


TEST_CASES = [
    (metrics.Accuracy(), sk_metrics.accuracy_score),
    (metrics.Precision(), partial(sk_metrics.precision_score, zero_division=0)),
    (
        metrics.MacroPrecision(),
        partial(sk_metrics.precision_score, average="macro", zero_division=0),
    ),
    (
        metrics.MicroPrecision(),
        partial(sk_metrics.precision_score, average="micro", zero_division=0),
    ),
    (
        metrics.WeightedPrecision(),
        partial(sk_metrics.precision_score, average="weighted", zero_division=0),
    ),
    (metrics.Recall(), partial(sk_metrics.recall_score, zero_division=0)),
    (
        metrics.MacroRecall(),
        partial(sk_metrics.recall_score, average="macro", zero_division=0),
    ),
    (
        metrics.MicroRecall(),
        partial(sk_metrics.recall_score, average="micro", zero_division=0),
    ),
    (
        metrics.WeightedRecall(),
        partial(sk_metrics.recall_score, average="weighted", zero_division=0),
    ),
    (
        metrics.FBeta(beta=0.5),
        partial(sk_metrics.fbeta_score, beta=0.5, zero_division=0),
    ),
    (
        metrics.MacroFBeta(beta=0.5),
        partial(sk_metrics.fbeta_score, beta=0.5, average="macro", zero_division=0),
    ),
    (
        metrics.MicroFBeta(beta=0.5),
        partial(sk_metrics.fbeta_score, beta=0.5, average="micro", zero_division=0),
    ),
    (
        metrics.WeightedFBeta(beta=0.5),
        partial(sk_metrics.fbeta_score, beta=0.5, average="weighted", zero_division=0),
    ),
    (metrics.F1(), partial(sk_metrics.f1_score, zero_division=0)),
    (metrics.MacroF1(), partial(sk_metrics.f1_score, average="macro", zero_division=0)),
    (metrics.MicroF1(), partial(sk_metrics.f1_score, average="micro", zero_division=0)),
    (
        metrics.WeightedF1(),
        partial(sk_metrics.f1_score, average="weighted", zero_division=0),
    ),
    (metrics.MCC(), sk_metrics.matthews_corrcoef),
    (metrics.MAE(), sk_metrics.mean_absolute_error),
    (metrics.MSE(), sk_metrics.mean_squared_error),
    (metrics.Homogeneity(), sk_metrics.homogeneity_score),
    (metrics.Completeness(), sk_metrics.completeness_score),
    (metrics.VBeta(beta=0.5), partial(sk_metrics.v_measure_score, beta=0.5)),
    (metrics.FowlkesMallows(), sk_metrics.fowlkes_mallows_score),
    (metrics.Rand(), sk_metrics.rand_score),
    (metrics.AdjustedRand(), sk_metrics.adjusted_rand_score),
    (metrics.MutualInfo(), sk_metrics.mutual_info_score),
    (
        metrics.NormalizedMutualInfo(average_method="min"),
        partial(sk_metrics.normalized_mutual_info_score, average_method="min"),
    ),
    (
        metrics.NormalizedMutualInfo(average_method="max"),
        partial(sk_metrics.normalized_mutual_info_score, average_method="max"),
    ),
    (
        metrics.NormalizedMutualInfo(average_method="arithmetic"),
        partial(sk_metrics.normalized_mutual_info_score, average_method="arithmetic"),
    ),
    (
        metrics.NormalizedMutualInfo(average_method="geometric"),
        partial(sk_metrics.normalized_mutual_info_score, average_method="geometric"),
    ),
    (
        metrics.AdjustedMutualInfo(average_method="max"),
        partial(sk_metrics.adjusted_mutual_info_score, average_method="max"),
    ),
    (
        metrics.AdjustedMutualInfo(average_method="arithmetic"),
        partial(sk_metrics.adjusted_mutual_info_score, average_method="arithmetic"),
    ),
    (
        metrics.AdjustedMutualInfo(average_method="geometric"),
        partial(sk_metrics.adjusted_mutual_info_score, average_method="geometric"),
    ),
    (metrics.Jaccard(), partial(sk_metrics.jaccard_score, average="binary")),
    (metrics.MacroJaccard(), partial(sk_metrics.jaccard_score, average="macro")),
    (metrics.MicroJaccard(), partial(sk_metrics.jaccard_score, average="micro")),
    (metrics.WeightedJaccard(), partial(sk_metrics.jaccard_score, average="weighted")),
    (metrics.RollingROCAUC(), roc_auc_score),
]

# HACK: not sure why this is needed, see this CI run https://github.com/online-ml/river/runs/7992357532?check_suite_focus=true
if platform.system() != "Linux":
    TEST_CASES.append(
        (
            metrics.AdjustedMutualInfo(average_method="min"),
            partial(sk_metrics.adjusted_mutual_info_score, average_method="min"),
        )
    )


@pytest.mark.parametrize(
    "metric, sk_metric",
    [
        pytest.param(metric, sk_metric, id=f"{metric.__class__.__name__}")
        for metric, sk_metric in TEST_CASES
    ],
)
@pytest.mark.filterwarnings("ignore::RuntimeWarning")
@pytest.mark.filterwarnings("ignore::sklearn.exceptions.UndefinedMetricWarning")
def test_metric(metric, sk_metric):
    for y_true, y_pred, sample_weights in generate_test_cases(metric=metric, n=30):
        m = copy.deepcopy(metric)
        for i, (yt, yp, w) in enumerate(zip(y_true, y_pred, sample_weights)):
            if metric.works_with_weights:
                m.update(y_true=yt, y_pred=yp, sample_weight=w)
            else:
                m.update(y_true=yt, y_pred=yp)

            if i >= 1:
                if metric.works_with_weights:
                    kwargs = {"sample_weight": sample_weights[: i + 1]}
                else:
                    kwargs = {}
                assert abs(m.get() - sk_metric(y_true[: i + 1], y_pred[: i + 1], **kwargs)) < 1e-6


@pytest.mark.parametrize(
    "metric, sk_metric",
    [
        pytest.param(metric, sk_metric, id=f"{metric.__class__.__name__}")
        for metric, sk_metric in TEST_CASES
    ],
)
@pytest.mark.filterwarnings("ignore::RuntimeWarning")
@pytest.mark.filterwarnings("ignore::sklearn.exceptions.UndefinedMetricWarning")
def test_rolling_metric(metric, sk_metric):
    def tail(iterable, n):
        return collections.deque(iterable, maxlen=n)

    for n in (1, 2, 5, 10):
        for y_true, y_pred, _ in generate_test_cases(metric=metric, n=30):
            m = utils.Rolling(copy.deepcopy(metric), window_size=n)

            # Check str works
            str(m)

            for i, (yt, yp) in enumerate(zip(y_true, y_pred)):
                m.update(y_true=yt, y_pred=yp)

                if i >= 1:
                    assert (
                        abs(
                            m.get()
                            - sk_metric(
                                tail(y_true[: i + 1], n),
                                tail(y_pred[: i + 1], n),
                            )
                        )
                        < 1e-10
                    )


def test_compose():
    metrics.MAE() + metrics.MSE()
    metrics.Accuracy() + metrics.LogLoss()

    with pytest.raises(ValueError):
        _ = metrics.MSE() + metrics.LogLoss()

    with pytest.raises(ValueError):
        _ = metrics.MSE() + metrics.MAE() + metrics.LogLoss()


# river/metrics/vbeta.py

from __future__ import annotations

import math

from river import metrics

__all__ = ["Completeness", "Homogeneity", "VBeta"]


class Homogeneity(metrics.base.MultiClassMetric):
    r"""Homogeneity Score.

    Homogeneity metric [^1] of a cluster labeling given a ground truth.

    In order to satisfy the homogeneity criteria, a clustering must assign only
    those data points that are members of a single class to a single cluster. That
    is, the class distribution within each cluster should be skewed to a single
    class, that is, zero entropy. We determine how close a given clustering is to
    this ideal by examining the conditional entropy of the class distribution given
    the proposed clustering.

    However, in an imperfect situation, the size of this value is dependent on the
    size of the dataset and the distribution of class sizes. Therefore, instead of
    taking the raw conditional entropy, we normalize by the maximum reduction in
    entropy the clustering information could provide.

    As such, we define homogeneity as:

    $$
    h = \begin{cases}
    1 if H(C) = 0, \\
    1 - \frac{H(C|K)}{H(C)} otherwise.
    \end{cases}.
    $$

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [1, 1, 2, 2, 3, 3]
    >>> y_pred = [1, 1, 1, 2, 2, 2]

    >>> metric = metrics.Homogeneity()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    1.0
    1.0
    0.0
    0.311278
    0.37515
    0.42062

    >>> metric
    Homogeneity: 42.06%

    References
    ----------
    [^1]: Andrew Rosenberg and Julia Hirschberg (2007).
          V-Measure: A conditional entropy-based external cluster evaluation measure.
          Proceedings of the 2007 Joing Conference on Empirical Methods in Natural Language
          Processing and Computational Natural Language Learning, pp. 410 - 420,
          Prague, June 2007.

    """

    @property
    def works_with_weights(self):
        return False

    def get(self):
        raw_conditional_entropy = 0.0
        max_reduction_entropy = 0.0

        for i in self.cm.classes:
            for j in self.cm.classes:
                try:
                    raw_conditional_entropy -= (
                        self.cm[j][i]
                        / self.cm.n_samples
                        * math.log(self.cm[j][i] / self.cm.sum_col[i], 2)
                    )
                except (ValueError, ZeroDivisionError):
                    continue

            try:
                max_reduction_entropy -= (
                    self.cm.sum_row[i]
                    / self.cm.n_samples
                    * math.log(self.cm.sum_row[i] / self.cm.n_samples, 2)
                )
            except (ValueError, ZeroDivisionError):
                continue

        try:
            return 1.0 - raw_conditional_entropy / max_reduction_entropy
        except ZeroDivisionError:
            return 1.0


class Completeness(metrics.base.MultiClassMetric):
    r"""Completeness Score.

    Completeness [^1] is symmetrical to homogeneity. In order to satisfy the
    completeness criteria, a clustering must assign all of those datapoints
    that are members of a single class to a single cluster. To evaluate completeness,
    we examine the distribution cluster assignments within each class. In a
    perfectly complete clustering solution, each of these distributions will be
    completely skewed to a single cluster.

    We can evaluate this degree of skew by calculating the conditional entropy of
    the proposed cluster distribution given the class of the component data points.
    However, in the worst case scenario, each class is represented by every cluster
    with a distribution equal to the distribution of cluster sizes. Therefore,
    symmetric to the claculation above, we define completeness as:

    $$
    c = \begin{cases}
    1 if H(K) = 0, \\
    1 - \frac{H(K|C)}{H(K)} otherwise.
    \end{cases}.
    $$

    Parameters
    ----------
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [1, 1, 2, 2, 3, 3]
    >>> y_pred = [1, 1, 1, 2, 2, 2]

    >>> metric = metrics.Completeness()
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    1.0
    1.0
    1.0
    0.3836885465963443
    0.5880325916843805
    0.6666666666666667

    >>> metric
    Completeness: 66.67%

    References
    ----------
    [^1]: Andrew Rosenberg and Julia Hirschberg (2007).
          V-Measure: A conditional entropy-based external cluster evaluation measure.
          Proceedings of the 2007 Joing Conference on Empirical Methods in Natural Language
          Processing and Computational Natural Language Learning, pp. 410 - 420,
          Prague, June 2007.

    """

    @property
    def works_with_weights(self):
        return False

    def get(self):
        raw_conditional_entropy = 0
        max_reduction_entropy = 0

        for i in self.cm.classes:
            for j in self.cm.classes:
                try:
                    raw_conditional_entropy -= (
                        self.cm[i][j]
                        / self.cm.n_samples
                        * math.log(self.cm[i][j] / self.cm.sum_row[i])
                    )
                except (ValueError, ZeroDivisionError):
                    continue

            try:
                max_reduction_entropy -= (
                    self.cm.sum_col[i]
                    / self.cm.n_samples
                    * math.log(self.cm.sum_col[i] / self.cm.n_samples)
                )
            except (ValueError, ZeroDivisionError):
                continue

        try:
            return 1.0 - raw_conditional_entropy / max_reduction_entropy
        except ZeroDivisionError:
            return 1.0


class VBeta(metrics.base.MultiClassMetric):
    r"""VBeta.

    VBeta (or V-Measure) [^1] is an external entropy-based cluster evaluation measure.
    It provides an elegant solution to many problems that affect previously defined
    cluster evaluation measures including

    * Dependance of clustering algorithm or dataset,

    * The "problem of matching", where the clustering of only a portion of data
    points are evaluated, and

    * Accurate evaluation and combination of two desirable aspects of clustering,
    homogeneity and completeness.

    Based upon the calculations of homogeneity and completeness, a clustering
    solution's V-measure is calculated by computing the weighted harmonic mean
    of homogeneity and completeness,

    $$
    V_{\beta} = \frac{(1 + \beta) \times h \times c}{\beta \times h + c}.
    $$

    Parameters
    ----------
    beta
        Weight of Homogeneity in the harmonic mean.
    cm
        This parameter allows sharing the same confusion
        matrix between multiple metrics. Sharing a confusion matrix reduces the amount of storage
        and computation time.

    Examples
    --------

    >>> from river import metrics

    >>> y_true = [1, 1, 2, 2, 3, 3]
    >>> y_pred = [1, 1, 1, 2, 2, 2]

    >>> metric = metrics.VBeta(beta=1.0)
    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(metric.update(yt, yp).get())
    1.0
    1.0
    0.0
    0.3437110184854507
    0.4580652856440158
    0.5158037429793888

    >>> metric
    VBeta: 51.58%

    References
    ----------
    [^1]: Andrew Rosenberg and Julia Hirschberg (2007).
          V-Measure: A conditional entropy-based external cluster evaluation measure.
          Proceedings of the 2007 Joing Conference on Empirical Methods in Natural Language
          Processing and Computational Natural Language Learning, pp. 410 - 420,
          Prague, June 2007.

    """

    def __init__(self, beta: float = 1.0, cm=None):
        super().__init__(cm)
        self.beta = beta
        self.homogeneity = metrics.Homogeneity(self.cm)
        self.completeness = metrics.Completeness(self.cm)

    @property
    def works_with_weights(self):
        return False

    def get(self):
        h = self.homogeneity.get()
        c = self.completeness.get()

        try:
            return (1 + self.beta) * h * c / (self.beta * h + c)
        except ZeroDivisionError:
            return 0.0


# river/misc/__init__.py

"""Miscellaneous.

This module essentially regroups some implementations that have nowhere else to go.

"""
from __future__ import annotations

from .sdft import SDFT
from .skyline import Skyline

__all__ = ["SDFT", "Skyline"]


# river/misc/sdft.py

from __future__ import annotations

import collections

import numpy as np

from river import base


class SDFT(base.Base):
    """Sliding Discrete Fourier Transform (SDFT).

    Initially, the coefficients are all equal to 0, up until enough values have been seen. A call
    to `numpy.fft.fft` is triggered once `window_size` values have been seen. Subsequent values
    will update the coefficients online. This is much faster than recomputing an FFT from scratch
    for every new value.

    Parameters
    ----------
    window_size
        The size of the window.

    Examples
    --------

    >>> import numpy as np
    >>> from river import misc

    >>> X = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

    >>> window_size = 5
    >>> sdft = misc.SDFT(window_size)

    >>> for i, x in enumerate(X):
    ...     sdft = sdft.update(x)
    ...
    ...     if i + 1 >= window_size:
    ...         assert np.allclose(sdft.coefficients, np.fft.fft(X[i+1 - window_size:i+1]))

    References
    ----------
    [^1]: [Jacobsen, E. and Lyons, R., 2003. The sliding DFT. IEEE Signal Processing Magazine, 20(2), pp.74-80.](https://www.comm.utoronto.ca/~dimitris/ece431/slidingdft.pdf)
    [^2]: [Understanding and Implementing the Sliding DFT](https://www.dsprelated.com/showarticle/776.php)

    """

    def __init__(self, window_size):
        self.coefficients = collections.deque(maxlen=window_size)
        self.window = collections.deque(maxlen=window_size)

    @property
    def window_size(self):
        return self.coefficients.maxlen

    def update(self, x):
        # Simply append the new value if the window isn't full yet
        if len(self.window) < self.window.maxlen - 1:
            self.window.append(x)

        # Compute an initial FFT the first time the window is full
        elif len(self.window) == self.window.maxlen - 1:
            self.window.append(x)
            self.coefficients.extend(np.fft.fft(self.window))

        # Update the coefficients for subsequent values
        else:
            diff = x - self.window[0]
            for i, c in enumerate(self.coefficients):
                self.coefficients[i] = (c + diff) * np.exp(2j * np.pi * i / self.window_size)
            self.window.append(x)

        return self


# river/misc/skyline.py

from __future__ import annotations

import collections

from river import base


class Skyline(collections.UserList, base.Base):
    """A skyline is set of points which is not dominated by any other point.

    This implementation uses a block nested loop. Identical observations are all part of the
    skyline if applicable.

    Parameters
    ----------
    minimize
        A list of features for which the values need to be minimized. Can be omitted as
        long as `maximize` is specified.
    maximize
        A list of features for which the values need to be maximized. Can be omitted as
        long as `minimize` is specified.

    Examples
    --------

    Here is an example taken from [this](https://maxhalford.github.io/blog/skyline-queries/) blog post.

    >>> import random
    >>> from river import misc

    >>> city_prices = {
    ...     'Bordeaux': 4045,
    ...     'Lyon': 4547,
    ...     'Toulouse': 3278
    ... }

    >>> def random_house():
    ...     city = random.choice(['Bordeaux', 'Lyon', 'Toulouse'])
    ...     size = round(random.gauss(200, 50))
    ...     price = round(random.uniform(0.8, 1.2) * city_prices[city] * size)
    ...     return {'city': city, 'size': size, 'price': price}

    >>> skyline = misc.Skyline(minimize=['price'], maximize=['size'])

    >>> random.seed(42)

    >>> for _ in range(100):
    ...     house = random_house()
    ...     skyline = skyline.update(house)

    >>> print(len(skyline))
    13

    >>> print(skyline[0])
    {'city': 'Toulouse', 'size': 280, 'price': 763202}

    Here is another example using the kart data from *Mario Kart: Double Dash!!*.

    >>> import collections
    >>> from river import misc

    >>> Kart = collections.namedtuple(
    ...      'Kart',
    ...      'name speed off_road acceleration weight turbo'
    ... )

    >>> karts = [
    ...     Kart('Red Fire', 5, 4, 4, 5, 2),
    ...     Kart('Green Fire', 7, 3, 3, 4, 2),
    ...     Kart('Heart Coach', 4, 6, 6, 5, 2),
    ...     Kart('Bloom Coach', 6, 4, 5, 3, 2),
    ...     Kart('Turbo Yoshi', 4, 5, 6, 6, 2),
    ...     Kart('Turbo Birdo', 6, 4, 4, 7, 2),
    ...     Kart('Goo-Goo Buggy', 1, 9, 9, 2, 3),
    ...     Kart('Rattle Buggy', 2, 9, 8, 2, 3),
    ...     Kart('Toad Kart', 3, 9, 7, 2, 3),
    ...     Kart('Toadette Kart', 1, 9, 9, 2, 3),
    ...     Kart('Koopa Dasher', 2, 8, 8, 3, 3),
    ...     Kart('Para-Wing', 1, 8, 9, 3, 3),
    ...     Kart('DK Jumbo', 8, 2, 2, 8, 1),
    ...     Kart('Barrel Train', 8, 7, 3, 5, 3),
    ...     Kart('Koopa King', 9, 1, 1, 9, 1),
    ...     Kart('Bullet Blaster', 8, 1, 4, 1, 3),
    ...     Kart('Wario Car', 7, 3, 3, 7, 1),
    ...     Kart('Waluigi Racer', 5, 9, 5, 6, 2),
    ...     Kart('Piranha Pipes', 8, 7, 2, 9, 1),
    ...     Kart('Boo Pipes', 2, 9, 8, 9, 1),
    ...     Kart('Parade Kart', 7, 3, 4, 7, 3)
    ... ]

    >>> skyline = misc.Skyline(
    ...     maximize=['speed', 'off_road', 'acceleration', 'turbo'],
    ...     minimize=['weight']
    ... )

    >>> for kart in karts:
    ...     skyline = skyline.update(kart._asdict())

    >>> best_cart_names = [kart['name'] for kart in skyline]
    >>> for name in best_cart_names:
    ...     print(f'- {name}')
    - Green Fire
    - Heart Coach
    - Bloom Coach
    - Goo-Goo Buggy
    - Rattle Buggy
    - Toad Kart
    - Toadette Kart
    - Barrel Train
    - Koopa King
    - Bullet Blaster
    - Waluigi Racer
    - Parade Kart

    >>> for name in sorted(set(kart.name for kart in karts) - set(best_cart_names)):
    ...     print(f'- {name}')
    - Boo Pipes
    - DK Jumbo
    - Koopa Dasher
    - Para-Wing
    - Piranha Pipes
    - Red Fire
    - Turbo Birdo
    - Turbo Yoshi
    - Wario Car

    References
    ----------
    [^1]: [Skyline queries in Python](https://maxhalford.github.io/blog/skyline-queries-in-python)
    [^2]: [Borzsony, S., Kossmann, D. and Stocker, K., 2001, April. The skyline operator. In Proceedings 17th international conference on data engineering (pp. 421-430). IEEE.](https://infolab.usc.edu/csci599/Fall2007/papers/e-1.pdf)
    [^3]: [Tao, Y. and Papadias, D., 2006. Maintaining sliding window skylines on data streams. IEEE Transactions on Knowledge and Data Engineering, 18(3), pp.377-391.](http://www.cs.ust.hk/~dimitris/PAPERS/TKDE06-Sky.pdf)

    """

    def __init__(self, minimize: list | None = None, maximize: list | None = None):
        super().__init__()

        self.minimize = [] if minimize is None else minimize
        self.maximize = [] if maximize is None else maximize

        if len(self.minimize) + len(self.maximize) == 0:
            raise ValueError("At least one name has to be specified")

    def _count_diffs(self, a, b):
        n_better = 0
        n_worse = 0

        for f in self.minimize:
            n_better += a[f] < b[f]
            n_worse += a[f] > b[f]

        for f in self.maximize:
            n_better += a[f] > b[f]
            n_worse += a[f] < b[f]

        return n_better, n_worse

    def update(self, x):
        # If the skyline is empty then the first element is part of the skyline
        if not self:
            self.append(x)
            return self

        to_drop = []
        is_dominated = False

        for i, _ in enumerate(self):
            n_better, n_worse = self._count_diffs(x, self[i])

            # x is dominated by one element in the skyline
            if n_worse > 0 and n_better == 0:
                is_dominated = True
                break

            # x dominates one element in the skyline
            if n_better > 0 and n_worse == 0:
                to_drop.append(i)

        if is_dominated:
            return self

        # Remove dominated elements
        if to_drop:
            for i in sorted(to_drop, reverse=True):
                del self[i]

        # Add x to the skyline
        self.append(x)

        return self


# river/model_selection/__init__.py

"""Model selection.

This module regroups a variety of methods that may be used for performing model selection. An
model selector is provided with a list of models. These are called "experts" in the expert learning
literature. The model selector's goal is to perform at least as well as the best model. Indeed,
initially, the best model is not known. The performance of each model becomes more apparent as time
goes by. Different strategies are possible, each one offering a different tradeoff in terms of
accuracy and computational performance.

Model selection can be used for tuning the hyperparameters of a model. This may be done by creating
a copy of the model for each set of hyperparameters, and treating each copy as a separate model.
The `utils.expand_param_grid` function can be used for this purpose.

"""
from __future__ import annotations

from . import base
from .bandit import BanditClassifier, BanditRegressor
from .greedy import GreedyRegressor
from .sh import SuccessiveHalvingClassifier, SuccessiveHalvingRegressor

__all__ = [
    "base",
    "BanditClassifier",
    "BanditRegressor",
    "GreedyRegressor",
    "SuccessiveHalvingClassifier",
    "SuccessiveHalvingRegressor",
]


# river/model_selection/bandit.py

from __future__ import annotations

from river import bandit, metrics, model_selection


class BanditModelSection:
    def _pick_arms(self):
        """Decide which models (arms) to pull.

        We want to pull every model at least burn_in times. This provides a significant accuracy
        boost, albeit at the cost of a longer training time in the beginning.

        """
        arm_ids = [
            arm_id
            for arm_id in range(len(self.models))
            if self.policy._counts[arm_id] < self.policy.burn_in
        ]
        if arm_ids:
            return arm_ids
        return [self.policy.pull(range(len(self.models)))]

    @property
    def best_model(self):
        if not self.policy._rewards:
            return None
        best_arm_id = max(self.policy._rewards, key=lambda arm: self.policy._rewards[arm].get())
        return self[best_arm_id]

    @classmethod
    def _unit_test_params(cls):
        for params in super()._unit_test_params():
            yield {**params, "policy": bandit.EpsilonGreedy(0.2)}

    def _unit_test_skips(self):
        return (
            {"check_model_selection_order_does_not_matter"} if hasattr(self.policy, "seed") else {}
        )


class BanditRegressor(BanditModelSection, model_selection.base.ModelSelectionRegressor):
    """Bandit-based model selection for regression.

    Each model is associated with an arm. At each `learn_one` call, the policy decides which
    arm/model to pull. The reward is the performance of the model on the provided sample. The
    `predict_one` method uses the current best model.

    Parameters
    ----------
    models
        The models to select from.
    metric
        The metric that is used to measure the performance of each model.
    policy
        The bandit policy to use.

    Examples
    --------

    >>> from river import bandit
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import model_selection
    >>> from river import optim
    >>> from river import preprocessing

    >>> models = [
    ...     linear_model.LinearRegression(optimizer=optim.SGD(lr=lr))
    ...     for lr in [0.0001, 0.001, 1e-05, 0.01]
    ... ]

    >>> dataset = datasets.TrumpApproval()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     model_selection.BanditRegressor(
    ...         models,
    ...         metric=metrics.MAE(),
    ...         policy=bandit.EpsilonGreedy(
    ...             epsilon=0.1,
    ...             decay=0.001,
    ...             burn_in=100,
    ...             seed=42
    ...         )
    ...     )
    ... )
    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 3.13815

    Here's another example using the UCB policy. The latter is more sensitive to the target scale,
    and usually works better when the target is rescaled.

    >>> models = [
    ...     linear_model.LinearRegression(optimizer=optim.SGD(lr=lr))
    ...     for lr in [0.0001, 0.001, 1e-05, 0.01]
    ... ]

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     preprocessing.TargetStandardScaler(
    ...         model_selection.BanditRegressor(
    ...             models,
    ...             metric=metrics.MAE(),
    ...             policy=bandit.UCB(
    ...                 delta=1,
    ...                 burn_in=100
    ...             )
    ...         )
    ...     )
    ... )
    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 0.875457

    """

    def __init__(
        self,
        models,
        metric: metrics.base.RegressionMetric,
        policy: bandit.base.Policy,
    ):
        super().__init__(models, metric)
        self.policy = policy.clone({"reward_obj": self.metric})

    def learn_one(self, x, y):
        for arm_id in self._pick_arms():
            model = self[arm_id]
            y_pred = model.predict_one(x)
            self.policy.update(arm_id, y_true=y, y_pred=y_pred)
            model.learn_one(x, y)

        return self


class BanditClassifier(BanditModelSection, model_selection.base.ModelSelectionClassifier):
    """Bandit-based model selection for classification.

    Each model is associated with an arm. At each `learn_one` call, the policy decides which
    arm/model to pull. The reward is the performance of the model on the provided sample. The
    `predict_one` and `predict_proba_one` methods use the current best model.

    Parameters
    ----------
    models
        The models to select from.
    metric
        The metric that is used to measure the performance of each model.
    policy
        The bandit policy to use.

    Examples
    --------

    >>> from river import bandit
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import model_selection
    >>> from river import optim
    >>> from river import preprocessing

    >>> models = [
    ...     linear_model.LogisticRegression(optimizer=optim.SGD(lr=lr))
    ...     for lr in [0.0001, 0.001, 1e-05, 0.01]
    ... ]

    >>> dataset = datasets.Phishing()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     model_selection.BanditClassifier(
    ...         models,
    ...         metric=metrics.Accuracy(),
    ...         policy=bandit.EpsilonGreedy(
    ...             epsilon=0.1,
    ...             decay=0.001,
    ...             burn_in=20,
    ...             seed=42
    ...         )
    ...     )
    ... )
    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 88.96%

    """

    def __init__(
        self,
        models,
        metric: metrics.base.ClassificationMetric,
        policy: bandit.base.Policy,
    ):
        super().__init__(models, metric)
        self.policy = policy.clone({"reward_obj": self.metric})

    def learn_one(self, x, y):
        for arm_id in self._pick_arms():
            model = self[arm_id]
            y_pred = (
                model.predict_one(x) if self.metric.requires_labels else model.predict_proba_one(x)
            )
            self.policy.update(arm_id, y_true=y, y_pred=y_pred)
            model.learn_one(x, y)

        return self


# river/model_selection/base.py

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Iterator

from river import base, metrics

__all__ = ["ModelSelectionRegressor", "ModelSelectionClassifier"]


class ModelSelector(base.Ensemble, ABC):
    """A generic model selector.

    Parameters
    ----------
    models
    metric

    """

    def __init__(self, models: Iterator[base.Estimator], metric: metrics.base.Metric):
        super().__init__(models)
        for model in models:
            if not metric.works_with(model):
                raise ValueError(
                    f"{metric.__class__.__name__} metric can't be used to evaluate a "
                    + f"{model.__class__.__name__}"
                )
        self.metric = metric

    @property
    @abstractmethod
    def best_model(self):
        """The current best model."""


class ModelSelectionRegressor(ModelSelector, base.Regressor):
    """A model selector for regression.

    Parameters
    ----------
    models
    metric

    """

    def predict_one(self, x):
        if self.best_model is None:
            return self.models[0].predict_one(x)
        return self.best_model.predict_one(x)

    @classmethod
    def _unit_test_params(cls):
        from river import compose, linear_model, optim, preprocessing

        yield {
            "models": [
                compose.Pipeline(
                    preprocessing.StandardScaler(),
                    linear_model.LinearRegression(optimizer=optim.SGD(lr=1e-2)),
                ),
                compose.Pipeline(
                    preprocessing.StandardScaler(),
                    linear_model.LinearRegression(optimizer=optim.SGD(lr=1e-1)),
                ),
            ],
            "metric": metrics.MAE(),
        }
        yield {
            "models": [
                compose.Pipeline(
                    preprocessing.StandardScaler(),
                    linear_model.LinearRegression(optimizer=optim.SGD(lr=lr)),
                )
                for lr in [1e-4, 1e-3, 1e-2, 1e-1]
            ],
            "metric": metrics.MAE(),
        }


class ModelSelectionClassifier(ModelSelector, base.Classifier):
    """A model selector for classification.

    Parameters
    ----------
    models
    metric

    """

    def predict_proba_one(self, x):
        if self.best_model is None:
            return self.models[0].predict_proba_one(x)
        return self.best_model.predict_proba_one(x)

    @classmethod
    def _unit_test_params(cls):
        from river import compose, linear_model, optim, preprocessing

        yield {
            "models": [
                compose.Pipeline(
                    preprocessing.StandardScaler(),
                    linear_model.LogisticRegression(optimizer=optim.SGD(lr=1e-2)),
                ),
                compose.Pipeline(
                    preprocessing.StandardScaler(),
                    linear_model.LogisticRegression(optimizer=optim.SGD(lr=1e-1)),
                ),
            ],
            "metric": metrics.Accuracy(),
        }
        yield {
            "models": [
                compose.Pipeline(
                    preprocessing.StandardScaler(),
                    linear_model.LogisticRegression(optimizer=optim.SGD(lr=lr)),
                )
                for lr in [1e-4, 1e-3, 1e-2, 1e-1]
            ],
            "metric": metrics.Accuracy(),
        }


# river/model_selection/greedy.py

from __future__ import annotations

from copy import deepcopy

from river import base, metrics
from river.model_selection.base import ModelSelectionRegressor


class GreedyRegressor(ModelSelectionRegressor):
    """Greedy selection regressor.

    This selection method simply updates each model at each time step. The current best model is
    used to make predictions. It's greedy in the sense that updating each model can be costly. On
    the other hand, bandit-like algorithms are more temperate in that only update a subset of the
    models at each step.

    Parameters
    ----------
    models
        The models to select from.
    metric
        The metric that is used to measure the performance of each model.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import model_selection
    >>> from river import optim
    >>> from river import preprocessing

    >>> models = [
    ...     linear_model.LinearRegression(optimizer=optim.SGD(lr=lr))
    ...     for lr in [1e-5, 1e-4, 1e-3, 1e-2]
    ... ]

    >>> dataset = datasets.TrumpApproval()
    >>> metric = metrics.MAE()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     model_selection.GreedyRegressor(models, metric)
    ... )

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 1.35

    """

    def __init__(
        self, models: list[base.Regressor], metric: metrics.base.RegressionMetric | None = None
    ):
        if metric is None:
            metric = metrics.MAE()
        super().__init__(models, metric)  # type: ignore
        self.metrics = [deepcopy(metric) for _ in range(len(self))]
        self._best_model = None
        self._best_metric = self.metrics[0]

    def learn_one(self, x, y):
        for model, metric in zip(self, self.metrics):
            y_pred = model.predict_one(x)
            metric.update(y, y_pred)
            model.learn_one(x, y)

            if metric.is_better_than(self._best_metric):
                self._best_model = model
                self._best_metric = metric

        return self

    @property
    def best_model(self):
        return self._best_model


# river/model_selection/sh.py

from __future__ import annotations

import abc
import math

from river import metrics
from river.model_selection.base import ModelSelectionClassifier, ModelSelectionRegressor

__all__ = ["SuccessiveHalvingClassifier", "SuccessiveHalvingRegressor"]


class SuccessiveHalving(abc.ABC):
    def __init__(
        self,
        models,
        metric: metrics.base.Metric,
        budget: int,
        eta=2,
        verbose=False,
        **print_kwargs,
    ):
        super().__init__(models=models, metric=metric)  # type: ignore
        self.budget = budget
        self.eta = eta
        self.verbose = verbose
        self.print_kwargs = print_kwargs

        self._n = len(models)
        self._metrics = [metric.clone() for _ in range(self._n)]
        self._s = self._n
        self._n_rungs = 0
        self._rankings = list(range(self._n))
        self._r = math.floor(budget / (self._s * math.ceil(math.log(self._n, eta))))
        self._budget_used = 0
        self._n_iterations = 0
        self._best_model_idx = 0

    @abc.abstractmethod
    def _pred_func(self, model):
        ...

    @property
    def best_model(self):
        """The current best model."""
        return self.models[self._best_model_idx]

    def learn_one(self, x, y):
        for i in self._rankings[: self._s]:
            model = self.models[i]
            metric = self._metrics[i]
            y_pred = self._pred_func(model)(x)
            metric.update(y_true=y, y_pred=y_pred)
            model.learn_one(x, y)

            # Check for a new best model
            if metric.is_better_than(self._metrics[self._best_model_idx]):
                self._best_model_idx = i

        self._n_iterations += 1

        if self._s > 1 and self._n_iterations == self._r:
            self._n_rungs += 1
            self._budget_used += self._s * self._r

            # Update the rankings of the current models based on their respective metric values
            self._rankings[: self._s] = sorted(
                self._rankings[: self._s],
                key=lambda i: self._metrics[i].get(),
                reverse=self.metric.bigger_is_better,
            )

            # Determine how many models to keep for the current rung
            cutoff = math.ceil(self._s / self.eta)

            if self.verbose:
                print(
                    "\t".join(
                        (
                            f"[{self._n_rungs}]",
                            f"{self._s - cutoff} removed",
                            f"{cutoff} left",
                            f"{self._r} iterations",
                            f"budget used: {self._budget_used}",
                            f"budget left: {self.budget - self._budget_used}",
                            f"best {self._metrics[self._rankings[0]]}",
                        )
                    ),
                    **self.print_kwargs,
                )

            # Determine where the next rung is located
            self._s = cutoff
            self._r = math.floor(self.budget / (self._s * math.ceil(math.log(self._n, self.eta))))

        return self


class SuccessiveHalvingRegressor(SuccessiveHalving, ModelSelectionRegressor):
    r"""Successive halving algorithm for regression.

    Successive halving is a method for performing model selection without having to train each
    model on all the dataset. At certain points in time (called "rungs"), the worst performing will
    be discarded and the best ones will keep competing between each other. The rung values are
    designed so that at most `budget` model updates will be performed in total.

    If you have `k` combinations of hyperparameters and that your dataset contains `n`
    observations, then the maximal budget you can allocate is:

    $$\frac{2kn}{eta}$$

    It is recommended that you check this beforehand. This bound can't be checked by the function
    because the size of the dataset is not known. In fact it is potentially infinite, in which case
    the algorithm will terminate once all the budget has been spent.

    If you have a budget of `B`, and that your dataset contains `n` observations, then the
    number of hyperparameter combinations that will spend all the budget and go through all the
    data is:

    $$\left\lceil\left\lfloor\frac{B}{2n}\right\rfloor \times eta \right\rceil$$

    Parameters
    ----------
    models
        The models to compare.
    metric
        Metric used for comparing models with.
    budget
        Total number of model updates you wish to allocate.
    eta
        Rate of elimination. At every rung, `math.ceil(k / eta)` models are kept, where `k` is the
        number of models that have reached the rung. A higher `eta` value will focus on less models
        but will allocate more iterations to the best models.
    verbose
        Whether to display progress or not.
    print_kwargs
        Extra keyword arguments are passed to the `print` function. For instance, this allows
        providing a `file` argument, which indicates where to output progress.

    Examples
    --------

    As an example, let's use successive halving to tune the optimizer of a linear regression.
    We'll first define the model.

    >>> from river import linear_model
    >>> from river import preprocessing

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LinearRegression(intercept_lr=.1)
    ... )

    Let's now define a grid of parameters which we would like to compare. We'll try
    different optimizers with various learning rates.

    >>> from river import optim
    >>> from river import utils

    >>> models = utils.expand_param_grid(model, {
    ...     'LinearRegression': {
    ...         'optimizer': [
    ...             (optim.SGD, {'lr': [.1, .01, .005]}),
    ...             (optim.Adam, {'beta_1': [.01, .001], 'lr': [.1, .01, .001]}),
    ...             (optim.Adam, {'beta_1': [.1], 'lr': [.001]}),
    ...         ]
    ...     }
    ... })

    We can check how many models we've created.

    >>> len(models)
    10

    We can now pass these models to a `SuccessiveHalvingRegressor`. We also need to pick a
    metric to compare the models, and a budget which indicates how many iterations to run
    before picking the best model and discarding the rest.

    >>> from river import model_selection

    >>> sh = model_selection.SuccessiveHalvingRegressor(
    ...     models,
    ...     metric=metrics.MAE(),
    ...     budget=2000,
    ...     eta=2,
    ...     verbose=True
    ... )

    A `SuccessiveHalvingRegressor` is also a regressor with a `learn_one` and a `predict_one`
    method. We can therefore evaluate it like any other classifier with
    `evaluate.progressive_val_score`.

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics

    >>> evaluate.progressive_val_score(
    ...     dataset=datasets.TrumpApproval(),
    ...     model=sh,
    ...     metric=metrics.MAE()
    ... )
    [1]	5 removed	5 left	50 iterations	budget used: 500	budget left: 1500	best MAE: 4.540491
    [2]	2 removed	3 left	100 iterations	budget used: 1000	budget left: 1000	best MAE: 2.458765
    [3]	1 removed	2 left	166 iterations	budget used: 1498	budget left: 502	best MAE: 1.583751
    [4]	1 removed	1 left	250 iterations	budget used: 1998	budget left: 2	best MAE: 1.147296
    MAE: 0.488387

    We can now view the best model.

    >>> sh.best_model
    Pipeline (
      StandardScaler (
        with_std=True
      ),
      LinearRegression (
        optimizer=Adam (
          lr=Constant (
            learning_rate=0.1
          )
          beta_1=0.01
          beta_2=0.999
          eps=1e-08
        )
        loss=Squared ()
        l2=0.
        l1=0.
        intercept_init=0.
        intercept_lr=Constant (
          learning_rate=0.1
        )
        clip_gradient=1e+12
        initializer=Zeros ()
      )
    )

    References
    ----------
    [^1]: [Jamieson, K. and Talwalkar, A., 2016, May. Non-stochastic best arm identification and hyperparameter optimization. In Artificial Intelligence and Statistics (pp. 240-248).](http://proceedings.mlr.press/v51/jamieson16.pdf)
    [^2]: [Li, L., Jamieson, K., Rostamizadeh, A., Gonina, E., Hardt, M., Recht, B. and Talwalkar, A., 2018. Massively parallel hyperparameter tuning. arXiv preprint arXiv:1810.05934.](https://arxiv.org/pdf/1810.05934.pdf)
    [^3]: [Li, L., Jamieson, K., DeSalvo, G., Rostamizadeh, A. and Talwalkar, A., 2017. Hyperband: A novel bandit-based approach to hyperparameter optimization. The Journal of Machine Learning Research, 18(1), pp.6765-6816.](https://arxiv.org/pdf/1603.06560.pdf)

    """

    def _pred_func(self, model):
        return model.predict_one

    def predict_one(self, x):
        return self.best_model.predict_one(x)

    @classmethod
    def _unit_test_params(cls):
        for params in super()._unit_test_params():
            yield {**params, "budget": 500}


class SuccessiveHalvingClassifier(SuccessiveHalving, ModelSelectionClassifier):
    r"""Successive halving algorithm for classification.

    Successive halving is a method for performing model selection without having to train each
    model on all the dataset. At certain points in time (called "rungs"), the worst performing will
    be discarded and the best ones will keep competing between each other. The rung values are
    designed so that at most `budget` model updates will be performed in total.

    If you have `k` combinations of hyperparameters and that your dataset contains `n`
    observations, then the maximal budget you can allocate is:

    $$\frac{2kn}{eta}$$

    It is recommended that you check this beforehand. This bound can't be checked by the function
    because the size of the dataset is not known. In fact it is potentially infinite, in which case
    the algorithm will terminate once all the budget has been spent.

    If you have a budget of `B`, and that your dataset contains `n` observations, then the
    number of hyperparameter combinations that will spend all the budget and go through all the
    data is:

    $$\left\lceil\left\lfloor\frac{B}{2n}\right\rfloor \times eta \right\rceil$$

    Parameters
    ----------
    models
        The models to compare.
    metric
        Metric used for comparing models with.
    budget
        Total number of model updates you wish to allocate.
    eta
        Rate of elimination. At every rung, `math.ceil(k / eta)` models are kept, where `k` is the
        number of models that have reached the rung. A higher `eta` value will focus on less models
        but will allocate more iterations to the best models.
    verbose
        Whether to display progress or not.
    print_kwargs
        Extra keyword arguments are passed to the `print` function. For instance, this allows
        providing a `file` argument, which indicates where to output progress.

    Examples
    --------

    As an example, let's use successive halving to tune the optimizer of a logistic regression.
    We'll first define the model.

    >>> from river import linear_model
    >>> from river import preprocessing

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression()
    ... )

    Let's now define a grid of parameters which we would like to compare. We'll try
    different optimizers with various learning rates.

    >>> from river import utils
    >>> from river import optim

    >>> models = utils.expand_param_grid(model, {
    ...     'LogisticRegression': {
    ...         'optimizer': [
    ...             (optim.SGD, {'lr': [.1, .01, .005]}),
    ...             (optim.Adam, {'beta_1': [.01, .001], 'lr': [.1, .01, .001]}),
    ...             (optim.Adam, {'beta_1': [.1], 'lr': [.001]}),
    ...         ]
    ...     }
    ... })

    We can check how many models we've created.

    >>> len(models)
    10

    We can now pass these models to a `SuccessiveHalvingClassifier`. We also need to pick a
    metric to compare the models, and a budget which indicates how many iterations to run
    before picking the best model and discarding the rest.

    >>> from river import model_selection

    >>> sh = model_selection.SuccessiveHalvingClassifier(
    ...     models,
    ...     metric=metrics.Accuracy(),
    ...     budget=2000,
    ...     eta=2,
    ...     verbose=True
    ... )

    A `SuccessiveHalvingClassifier` is also a classifier with a `learn_one` and a
    `predict_proba_one` method. We can therefore evaluate it like any other classifier with
    `evaluate.progressive_val_score`.

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics

    >>> evaluate.progressive_val_score(
    ...     dataset=datasets.Phishing(),
    ...     model=sh,
    ...     metric=metrics.ROCAUC()
    ... )
    [1] 5 removed       5 left  50 iterations   budget used: 500        budget left: 1500       best Accuracy: 80.00%
    [2] 2 removed       3 left  100 iterations  budget used: 1000       budget left: 1000       best Accuracy: 84.00%
    [3] 1 removed       2 left  166 iterations  budget used: 1498       budget left: 502        best Accuracy: 86.14%
    [4] 1 removed       1 left  250 iterations  budget used: 1998       budget left: 2  best Accuracy: 84.80%
    ROCAUC: 95.29%

    We can now view the best model.

    >>> sh.best_model
    Pipeline (
      StandardScaler (
        with_std=True
      ),
      LogisticRegression (
        optimizer=Adam (
          lr=Constant (
            learning_rate=0.01
          )
          beta_1=0.01
          beta_2=0.999
          eps=1e-08
        )
        loss=Log (
          weight_pos=1.
          weight_neg=1.
        )
        l2=0.
        l1=0.
        intercept_init=0.
        intercept_lr=Constant (
          learning_rate=0.01
        )
        clip_gradient=1e+12
        initializer=Zeros ()
      )
    )

    References
    ----------
    [^1]: [Jamieson, K. and Talwalkar, A., 2016, May. Non-stochastic best arm identification and hyperparameter optimization. In Artificial Intelligence and Statistics (pp. 240-248).](http://proceedings.mlr.press/v51/jamieson16.pdf)
    [^2]: [Li, L., Jamieson, K., Rostamizadeh, A., Gonina, E., Hardt, M., Recht, B. and Talwalkar, A., 2018. Massively parallel hyperparameter tuning. arXiv preprint arXiv:1810.05934.](https://arxiv.org/pdf/1810.05934.pdf)
    [^3]: [Li, L., Jamieson, K., DeSalvo, G., Rostamizadeh, A. and Talwalkar, A., 2017. Hyperband: A novel bandit-based approach to hyperparameter optimization. The Journal of Machine Learning Research, 18(1), pp.6765-6816.](https://arxiv.org/pdf/1603.06560.pdf)

    """

    def _pred_func(self, model):
        if self.metric.requires_labels:
            return model.predict_one
        return model.predict_proba_one

    def predict_proba_one(self, x):
        return self.best_model.predict_proba_one(x)

    def _multiclass(self):
        return all(model._multiclass for model in self.models)


# river/multiclass/__init__.py

"""Multi-class classification."""
from __future__ import annotations

from .occ import OutputCodeClassifier
from .ovo import OneVsOneClassifier
from .ovr import OneVsRestClassifier

__all__ = ["OutputCodeClassifier", "OneVsOneClassifier", "OneVsRestClassifier"]


# river/multiclass/occ.py

from __future__ import annotations

import collections
import random

from river import base, linear_model

__all__ = ["OutputCodeClassifier"]


def l1_dist(a, b):
    return sum(abs(ai - bi) for ai, bi in zip(a, b))


class OutputCodeClassifier(base.Wrapper, base.Classifier):
    """Output-code multiclass strategy.

    This also referred to as "error-correcting output codes".

    This class allows to learn a multi-class classification problem with a binary classifier. Each
    class is converted to a code of 0s and 1s. The length of the code is called  the code size. A
    copy of the classifier made for code. The codes associated with the classes are stored in a
    code book.

    When a new sample arrives, the label's code is retrieved from the code book. Then, each
    classifier is trained on the relevant part of code, which is either a 0 or a 1.

    For predicting, each classifier outputs a probability. These are then compared to each code in
    the code book, and the label which is the "closest" is chosen as the most likely class.
    Closeness is determined in terms of Manhattan distance.

    One specificity of online learning is that we don't how many classes there are initially.
    Therefore, a random procedure generates random codes on the fly whenever a previously unseed
    label appears.

    Parameters
    ----------
    classifier
        A binary classifier, although a multi-class classifier will work too.
    code_size
        The code size, which dictates how many copies of the provided classifiers to train. Must be
        strictly positive.
    coding_method
        The method used to generate the codes. Can be either 'exact' or 'random'. The 'exact'
        method generates all possible codes of a given size in memory, and streams them in a random
        order. The 'random' method generates random codes of a given size on the fly. The 'exact'
        method necessarily generates different codes for each class, but requires more memory. The
        'random' method can generate duplicate codes for different classes, but requires less
        memory.
    seed
        A random seed number that can be set for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multiclass
    >>> from river import preprocessing

    >>> dataset = datasets.ImageSegments()

    >>> scaler = preprocessing.StandardScaler()
    >>> ooc = multiclass.OutputCodeClassifier(
    ...     classifier=linear_model.LogisticRegression(),
    ...     code_size=10,
    ...     coding_method='random',
    ...     seed=1
    ... )
    >>> model = scaler | ooc

    >>> metric = metrics.MacroF1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MacroF1: 79.32%

    References
    ----------
    [^1]: [Dietterich, T.G. and Bakiri, G., 1994. Solving multiclass learning problems via error-correcting output codes. Journal of artificial intelligence research, 2, pp.263-286.](https://arxiv.org/pdf/cs/9501101.pdf)
    [^2]: [James, G. and Hastie, T., 1998. The error coding method and PICTs. Journal of Computational and Graphical statistics, 7(3), pp.377-387.](https://www.semanticscholar.org/paper/The-Error-Coding-Method-and-PICTs-James-Hastie/73abc83ed720921ed912a709aab6b734915b1012)

    """

    def __init__(
        self,
        classifier: base.Classifier,
        code_size: int,
        coding_method: str = "random",
        seed: int | None = None,
    ):
        self.classifier = classifier
        self.code_size = code_size
        self.coding_method = coding_method
        self.seed = seed
        self._rng = random.Random(seed)

        self.classifiers = {i: classifier.clone() for i in range(code_size)}

        # We don't know how many classes there are, therefore we can't generate the code book
        # from the start. Therefore, we define a random queue of integers. When a new class
        # appears, we get the next integer and convert it to a code. There are different ways to do
        # this.
        if self.coding_method == "exact":
            integers = list(range(2**code_size))
            self._rng.shuffle(integers)
            self._integers = iter(integers)
        self.code_book: collections.defaultdict = collections.defaultdict(self._next_code)

    def _next_code(self):
        if self.coding_method == "random":
            return tuple(self._rng.randint(0, 1) for _ in range(self.code_size))
        elif self.coding_method == "exact":
            i = next(self._integers)
            b = bin(i)[2:]  # convert to a string of 0s and 1s
            b = b.zfill(self.code_size)  # ensure the code is of length code_size
            return tuple(int(c) for c in b)

    @property
    def _multiclass(self):
        return True

    @property
    def _wrapped_model(self):
        return self.classifier

    @classmethod
    def _unit_test_params(cls):
        yield {
            "classifier": linear_model.LogisticRegression(),
            "code_size": 6,
            "coding_method": "exact",
        }
        # A code size of 30 would overload RAM with the exact method
        yield {
            "classifier": linear_model.LogisticRegression(),
            "code_size": 30,
            "coding_method": "random",
        }

    def learn_one(self, x, y, **kwargs):
        code = self.code_book[y]

        for i, c in enumerate(code):
            self.classifiers[i].learn_one(x, c, **kwargs)

        return self

    def predict_one(self, x, **kwargs):
        if not self.code_book:  # it's empty
            return None

        output = [None for _ in range(self.code_size)]

        for i, clf in self.classifiers.items():
            output[i] = clf.predict_proba_one(x, **kwargs).get(True, 0.0)

        return min(self.code_book, key=lambda c: l1_dist(self.code_book[c], output))


# river/multiclass/ovo.py

from __future__ import annotations

import collections
import copy
import functools

from river import base, linear_model

__all__ = ["OneVsOneClassifier"]


class OneVsOneClassifier(base.Wrapper, base.Classifier):
    """One-vs-One (OvO) multiclass strategy.

    This strategy consists in fitting one binary classifier for each pair of classes. Because we
    are in a streaming context, the number of classes isn't known from the start, hence new
    classifiers are instantiated on the fly.

    The number of classifiers is `k * (k - 1) / 2`, where `k` is the number of classes. However,
    each call to `learn_one` only requires training `k - 1` models. Indeed, only the models that
    pertain to the given label have to be trained. Meanwhile, making a prediction requires going
    through each and every model.

    Parameters
    ----------
    classifier
        A binary classifier, although a multi-class classifier will work too.

    Attributes
    ----------
    classifiers : dict
        A mapping between pairs of classes and classifiers. The keys are tuples which contain a
        pair of classes. Each pair is sorted in lexicographical order.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multiclass
    >>> from river import preprocessing

    >>> dataset = datasets.ImageSegments()

    >>> scaler = preprocessing.StandardScaler()
    >>> ovo = multiclass.OneVsOneClassifier(linear_model.LogisticRegression())
    >>> model = scaler | ovo

    >>> metric = metrics.MacroF1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MacroF1: 80.76%

    """

    def __init__(self, classifier):
        self.classifier = classifier
        new_clf = functools.partial(copy.deepcopy, classifier)
        self.classifiers = collections.defaultdict(new_clf)
        self.classes = set()

    @property
    def _wrapped_model(self):
        return self.classifier

    @property
    def _multiclass(self):
        return True

    @classmethod
    def _unit_test_params(cls):
        yield {"classifier": linear_model.LogisticRegression()}

    def learn_one(self, x, y, **kwargs):
        self.classes.add(y)

        for c in self.classes - {y}:
            pair = (c, y) if c < y else (y, c)
            self.classifiers[pair].learn_one(x, y=c < y, **kwargs)

        return self

    def predict_one(self, x, **kwargs):
        if not self.classifiers:  # is empty
            return None

        votes = collections.defaultdict(int)

        for pair, clf in self.classifiers.items():
            if clf.predict_one(x, **kwargs):
                votes[pair[1]] += 1
            else:
                votes[pair[0]] += 1

        return max(votes, key=votes.get)


# river/multiclass/ovr.py

from __future__ import annotations

import pandas as pd

from river import base, linear_model

__all__ = ["OneVsRestClassifier"]


class OneVsRestClassifier(base.Wrapper, base.Classifier):
    """One-vs-the-rest (OvR) multiclass strategy.

    This strategy consists in fitting one binary classifier per class. Because we are in a
    streaming context, the number of classes isn't known from the start. Hence, new classifiers are
    instantiated on the fly. Likewise, the predicted probabilities will only include the classes
    seen up to a given point in time.

    Note that this classifier supports mini-batches as well as single instances.

    The computational complexity for both learning and predicting grows linearly with the number of
    classes. If you have a very large number of classes, then you might want to consider using an
    `multiclass.OutputCodeClassifier` instead.

    Parameters
    ----------
    classifier
        A binary classifier, although a multi-class classifier will work too.

    Attributes
    ----------
    classifiers : dict
        A mapping between classes and classifiers.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multiclass
    >>> from river import preprocessing

    >>> dataset = datasets.ImageSegments()

    >>> scaler = preprocessing.StandardScaler()
    >>> ovr = multiclass.OneVsRestClassifier(linear_model.LogisticRegression())
    >>> model = scaler | ovr

    >>> metric = metrics.MacroF1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MacroF1: 77.46%

    This estimator also also supports mini-batching.

    >>> for X in pd.read_csv(dataset.path, chunksize=64):
    ...     y = X.pop('category')
    ...     y_pred = model.predict_many(X)
    ...     model = model.learn_many(X, y)

    """

    def __init__(self, classifier: base.Classifier):
        self.classifier = classifier
        self.classifiers: dict[base.typing.ClfTarget, base.Classifier] = {}
        self._y_name = None

    @property
    def _wrapped_model(self):
        return self.classifier

    @property
    def _multiclass(self):
        return True

    @classmethod
    def _unit_test_params(cls):
        yield {"classifier": linear_model.LogisticRegression()}

    def learn_one(self, x, y, **kwargs):
        # Instantiate a new binary classifier if the class is new
        if y not in self.classifiers:
            self.classifiers[y] = self.classifier.clone()

        # Train each label's associated classifier
        for label, model in self.classifiers.items():
            model.learn_one(x, y == label, **kwargs)

        return self

    def predict_proba_one(self, x, **kwargs):
        y_pred = {}
        total = 0.0

        for label, model in self.classifiers.items():
            yp = model.predict_proba_one(x, **kwargs)[True]
            y_pred[label] = yp
            total += yp

        if total:
            return {label: votes / total for label, votes in y_pred.items()}
        return {label: 1 / len(y_pred) for label in y_pred}

    def learn_many(self, X, y, **kwargs):
        self._y_name = y.name

        # Instantiate a new binary classifier for the classes that have not yet been seen
        for label in y.unique():
            if label not in self.classifiers:
                self.classifiers[label] = self.classifier.clone()

        # Train each label's associated classifier
        for label, model in self.classifiers.items():
            model.learn_many(X, y == label, **kwargs)

        return self

    def predict_proba_many(self, X, **kwargs):
        y_pred = pd.DataFrame(columns=self.classifiers.keys(), index=X.index)

        for label, clf in self.classifiers.items():
            y_pred[label] = clf.predict_proba_many(X, **kwargs)[True]

        return y_pred.div(y_pred.sum(axis="columns"), axis="rows")

    def predict_many(self, X, **kwargs):
        if not self.classifiers:
            return pd.Series([None] * len(X), index=X.index, dtype="object")
        return self.predict_proba_many(X, **kwargs).idxmax(axis="columns").rename(self._y_name)


# river/multiclass/test_ovr.py

from __future__ import annotations

import pandas as pd

from river import datasets, linear_model, metrics, multiclass, preprocessing, stream


def test_online_batch_consistent():
    # Batch

    batch = preprocessing.StandardScaler() | multiclass.OneVsRestClassifier(
        linear_model.LogisticRegression()
    )

    dataset = datasets.ImageSegments()

    batch_metric = metrics.MacroF1()

    for i, x in enumerate(pd.read_csv(dataset.path, chunksize=1)):
        y = x.pop("category")
        y_pred = batch.predict_many(x)
        batch.learn_many(x, y)

        for yt, yp in zip(y, y_pred):
            if yp is not None:
                batch_metric.update(yt, yp)

        if i == 30:
            break

    # Online

    online = preprocessing.StandardScaler() | multiclass.OneVsRestClassifier(
        linear_model.LogisticRegression()
    )

    online_metric = metrics.MacroF1()

    X = pd.read_csv(dataset.path)
    Y = X.pop("category")

    for i, (x, y) in enumerate(stream.iter_pandas(X, Y)):
        y_pred = online.predict_one(x)
        online.learn_one(x, y)

        if y_pred is not None:
            online_metric.update(y, y_pred)

        if i == 30:
            break

    assert online_metric.get() == batch_metric.get()


# river/multioutput/__init__.py

"""Multi-output models."""
from __future__ import annotations

from .chain import (
    ClassifierChain,
    MonteCarloClassifierChain,
    ProbabilisticClassifierChain,
    RegressorChain,
)
from .encoder import MultiClassEncoder

__all__ = [
    "ClassifierChain",
    "MonteCarloClassifierChain",
    "MultiClassEncoder",
    "ProbabilisticClassifierChain",
    "RegressorChain",
]


# river/multioutput/chain.py

from __future__ import annotations

import collections
import copy
import random

from river import base, linear_model
from river.utils.math import prod

__all__ = [
    "ClassifierChain",
    "RegressorChain",
    "ProbabilisticClassifierChain",
    "MonteCarloClassifierChain",
]


class BaseChain(base.Wrapper, collections.UserDict):
    def __init__(self, model, order: list | None = None):
        super().__init__()
        self.model = model
        self.order = order or []

    @property
    def _wrapped_model(self):
        return self.model

    def __getitem__(self, key):
        try:
            return collections.UserDict.__getitem__(self, key)
        except KeyError:
            collections.UserDict.__setitem__(self, key, copy.deepcopy(self.model))
            return self[key]


class ClassifierChain(BaseChain, base.MultiLabelClassifier):
    """A multi-output model that arranges classifiers into a chain.

    This will create one model per output. The prediction of the first output will be used as a
    feature in the second model. The prediction for the second output will be used as a feature
    for the third model, etc. This "chain model" is therefore capable of capturing dependencies
    between outputs.

    Parameters
    ----------
    model
        A classifier model used for each label.
    order
        A list with the targets order in which to construct the chain. If `None` then the order
        will be inferred from the order of the keys in the target.

    Examples
    --------

    >>> from river import feature_selection
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multioutput
    >>> from river import preprocessing
    >>> from river import stream
    >>> from sklearn import datasets

    >>> dataset = stream.iter_sklearn_dataset(
    ...     dataset=datasets.fetch_openml('yeast', version=4, parser='auto', as_frame=False),
    ...     shuffle=True,
    ...     seed=42
    ... )

    >>> model = feature_selection.VarianceThreshold(threshold=0.01)
    >>> model |= preprocessing.StandardScaler()
    >>> model |= multioutput.ClassifierChain(
    ...     model=linear_model.LogisticRegression(),
    ...     order=list(range(14))
    ... )

    >>> metric = metrics.multioutput.MicroAverage(metrics.Jaccard())

    >>> for x, y in dataset:
    ...     # Convert y values to booleans
    ...     y = {i: yi == 'TRUE' for i, yi in y.items()}
    ...     y_pred = model.predict_one(x)
    ...     metric = metric.update(y, y_pred)
    ...     model = model.learn_one(x, y)

    >>> metric
    MicroAverage(Jaccard): 41.95%

    References
    ----------
    [^1]: [Multi-Output Chain Models and their Application in Data Streams](https://jmread.github.io/talks/2019_03_08-Imperial_Stats_Seminar.pdf)

    """

    def __init__(self, model: base.Classifier, order: list | None = None):
        super().__init__(model, order)

    @classmethod
    def _unit_test_params(cls):
        from river import neighbors

        yield {"model": linear_model.LogisticRegression()}  # binary classifier
        yield {"model": linear_model.SoftmaxRegression()}  # multi-class classifier
        yield {
            "model": neighbors.KNNClassifier(n_neighbors=1, window_size=5)
        }  # multi-class classifier

    @property
    def _multiclass(self):
        return self.model._multiclass

    def learn_one(self, x, y, **kwargs):
        x = copy.copy(x)
        n_seen = 0

        for o in self.order:
            clf = self[o]

            # Make predictions before the model is updated to avoid leakage
            y_pred = clf.predict_proba_one(x, **kwargs)

            # We handle the case where an output has been seen in the past but is missing now
            try:
                y_o = y[o]
                n_seen += 1
                clf.learn_one(x, y_o, **kwargs)
            except KeyError:
                pass

            # The predictions are stored as features for the next label
            if clf._multiclass:
                for label, proba in y_pred.items():
                    x[f"{o}_{label}"] = proba
            else:
                x[o] = y_pred[True]

        # Now we check if there are any new outputs
        n_unseen = len(y) - n_seen
        if n_unseen:
            for o in y:
                if o not in self.order:
                    self.order.append(o)

        return self

    def predict_proba_one(self, x, **kwargs):
        x = copy.copy(x)
        y_pred = {}

        for o in self.order:
            clf = self[o]

            y_pred[o] = clf.predict_proba_one(x, **kwargs)

            # The predictions are stored as features for the next label
            if clf._multiclass:
                for label, proba in y_pred[o].items():
                    x[f"{o}_{label}"] = proba
            else:
                x[o] = y_pred[o][True]

        return y_pred


class RegressorChain(BaseChain, base.MultiTargetRegressor):
    """A multi-output model that arranges regressors into a chain.

    This will create one model per output. The prediction of the first output will be used as a
    feature in the second output. The prediction for the second output will be used as a feature
    for the third, etc. This "chain model" is therefore capable of capturing dependencies between
    outputs.

    Parameters
    ----------
    model
        The regression model used to make predictions for each target.
    order
        A list with the targets order in which to construct the chain. If `None` then the order
        will be inferred from the order of the keys in the target.

    Examples
    --------

    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multioutput
    >>> from river import preprocessing
    >>> from river import stream

    >>> from sklearn import datasets

    >>> dataset = stream.iter_sklearn_dataset(
    ...     dataset=datasets.load_linnerud(),
    ...     shuffle=True,
    ...     seed=42
    ... )

    >>> model = multioutput.RegressorChain(
    ...     model=(
    ...         preprocessing.StandardScaler() |
    ...         linear_model.LinearRegression(intercept_lr=0.3)
    ...     ),
    ...     order=[0, 1, 2]
    ... )

    >>> metric = metrics.multioutput.MicroAverage(metrics.MAE())

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MicroAverage(MAE): 12.649592

    """

    def __init__(self, model: base.Regressor, order: list | None = None):
        super().__init__(model, order)

    @classmethod
    def _unit_test_params(cls):
        yield {"model": linear_model.LinearRegression()}

    def learn_one(self, x, y, **kwargs):
        x = copy.copy(x)
        n_seen = 0

        for o in self.order:
            reg = self[o]

            # Make predictions before the model is updated to avoid leakage
            y_pred = reg.predict_one(x)

            # We handle the case where an output has been seen in the past but is missing now
            try:
                y_o = y[o]
                n_seen += 1
                reg.learn_one(x, y_o, **kwargs)
            except KeyError:
                pass

            # The predictions are stored as features for the next label
            x[o] = y_pred

        # Now we check if there are any new outputs
        n_unseen = len(y) - n_seen
        if n_unseen:
            for o in y:
                if o not in self.order:
                    self.order.append(o)

        return self

    def predict_one(self, x, **kwargs):
        x = copy.copy(x)
        y_pred = {}

        if not isinstance(self.order, list):
            return y_pred

        for o, clf in self.items():
            y_pred[o] = clf.predict_one(x, **kwargs)
            x[o] = y_pred[o]

        return y_pred


class ProbabilisticClassifierChain(ClassifierChain):
    r"""Probabilistic Classifier Chains.

    The Probabilistic Classifier Chains (PCC) [^1] is a Bayes-optimal method
    based on the Classifier Chains (CC).

    Consider the concept of chaining classifiers as searching a path in a
    binary tree whose leaf nodes are associated with a label $y \in Y$. While
    CC searches only a single path in the aforementioned binary tree, PCC looks
    at each of the $2^l$ paths, where $l$ is the number of labels. This limits
    the applicability of the method to data sets with a small to moderate
    number of labels. The authors recommend no more than about 15 labels for
    real-world applications.

    Parameters
    ----------
    model

    Examples
    --------

    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multioutput
    >>> from river.datasets import synth

    >>> dataset = synth.Logical(seed=42, n_tiles=100)

    >>> model = multioutput.ProbabilisticClassifierChain(
    ...     model=linear_model.LogisticRegression()
    ... )

    >>> metric = metrics.multioutput.MicroAverage(metrics.Jaccard())

    >>> for x, y in dataset:
    ...    y_pred = model.predict_one(x)
    ...    y_pred = {k: y_pred.get(k, 0) for k in y}
    ...    metric = metric.update(y, y_pred)
    ...    model = model.learn_one(x, y)

    >>> metric
    MicroAverage(Jaccard): 51.84%

    References
    ----------
    [^1]: Cheng, W., Hüllermeier, E., & Dembczynski, K. J. (2010).
          Bayes optimal multilabel classification via probabilistic classifier
          chains. In Proceedings of the 27th international conference on
          machine learning (ICML-10) (pp. 279-286).

    """

    def __init__(self, model: base.Classifier):
        super().__init__(model)

    def predict_one(self, x, **kwargs):
        y_pred = {}

        if not isinstance(self.order, list):
            return y_pred

        max_payoff = 0.0
        n_labels = len(self.order)
        # for each and every possible label combination
        for label in range(2**n_labels):
            # put together a binary label vector
            y_gen = {i: int(v) for i, v in zip(self.order, list(bin(label)[2:].zfill(n_labels)))}
            # ... and gauge a probability for it (given x)
            payoff = self._payoff(x=x, y=y_gen, **kwargs)
            # if it performs well, keep it, and record the max
            if payoff > max_payoff:
                y_pred = copy.copy(y_gen)
                max_payoff = payoff
        return y_pred

    def _payoff(self, x, y, **kwargs):
        # Calculate payoff for predicting y | x, under the chains model.
        p = {}

        x = copy.copy(x)

        for label in self.order:
            clf = self[label]

            y_pred = clf.predict_proba_one(x, **kwargs)
            # Extend features
            x[label] = y.get(label, 0)
            p[label] = y_pred.get(y[label], 0) if label in y else 0

        return prod(p.values()) if p else 0


class MonteCarloClassifierChain(ProbabilisticClassifierChain):
    """Monte Carlo Sampling Classifier Chains.

    Probabilistic Classifier Chains using Monte Carlo sampling, as
    described in [^1].

    m samples are taken from the posterior distribution. Therefore we
    need a probabilistic interpretation of the output, and thus, this is a
    particular variety of ProbabilisticClassifierChain.

    Parameters
    ----------
    model
    m
        Number of samples to take from the posterior distribution.
    seed
        Random number generator seed for reproducibility.

    Examples
    --------
    >>> from river import feature_selection
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import multioutput
    >>> from river import preprocessing
    >>> from river.datasets import synth

    >>> dataset = synth.Logical(seed=42, n_tiles=100)

    >>> model = multioutput.MonteCarloClassifierChain(
    ...     model=linear_model.LogisticRegression(),
    ...     m=10,
    ...     seed=42
    ... )

    >>> metric = metrics.multioutput.MicroAverage(metrics.Jaccard())

    >>> for x, y in dataset:
    ...    y_pred = model.predict_one(x)
    ...    y_pred = {k: y_pred.get(k, 0) for k in y}
    ...    metric = metric.update(y, y_pred)
    ...    model = model.learn_one(x, y)

    >>> metric
    MicroAverage(Jaccard): 51.79%

    References
    ----------
    [^1]: Read, J., Martino, L., & Luengo, D. (2014). Efficient monte carlo
          methods for multi-dimensional learning with classifier chains.
          Pattern Recognition, 47(3), 1535-1546.

    """

    def __init__(self, model: base.Classifier, m: int = 10, seed: int | None = None):
        ClassifierChain.__init__(self, model=model, order=None)
        self.seed = seed
        self._rng = random.Random(seed)
        self.m = m

    def _sample(self, x):
        # Sample y ~ P(y|x)
        p = {}
        y = {}
        x = copy.copy(x)

        for label in self.order:
            clf = self[label]

            y_pred = clf.predict_proba_one(x)
            if y_pred:
                y_val = self._rng.choices(
                    range(len(y_pred)), k=1, weights=[v for v in y_pred.values()]
                )[0]
                # Extend features
                x[label] = y_val
                y[label] = y_val
                p[label] = list(y_pred.keys())[y_val]

        return y, p

    def predict_one(self, x, **kwargs):
        y_pred = {}

        if not isinstance(self.order, list):
            return y_pred

        y_pred = ClassifierChain.predict_one(self, x, **kwargs)
        max_payoff = self._payoff(x=x, y=y_pred)
        # for M times
        for m in range(self.m):
            y_, p_ = self._sample(x)  # N.B. in fact, the calculation p_ is done again in P.
            payoff = self._payoff(x=x, y=y_)
            # if it performs well, keep it, and record the max
            if payoff > max_payoff:
                y_pred = copy.copy(y_)
                max_payoff = payoff
        return y_pred


# river/multioutput/encoder.py

from __future__ import annotations

import typing

from river import base


class MultiClassEncoder(base.MultiLabelClassifier):
    """Convert a multi-label task into multiclass.

    Assigns a class to each unique combination of labels, and proceeds
    with training the supplied multi-class classifier.

    The transformation is done by converting the label set, which could be seen
    as a binary number, into an integer representing a class. At prediction time,
    the predicted integer is converted back to a binary number which is the
    predicted label set.

    Parameters
    ----------
    model
        The classifier used for learning.

    Examples
    --------

    >>> from river import forest
    >>> from river import metrics
    >>> from river import multioutput
    >>> from river.datasets import synth

    >>> dataset = synth.Logical(seed=42, n_tiles=100)

    >>> model = multioutput.MultiClassEncoder(
    ...     model=forest.ARFClassifier(seed=7)
    ... )

    >>> metric = metrics.multioutput.MicroAverage(metrics.Jaccard())

    >>> for x, y in dataset:
    ...    y_pred = model.predict_one(x)
    ...    y_pred = {k: y_pred.get(k, 0) for k in y}
    ...    metric = metric.update(y, y_pred)
    ...    model = model.learn_one(x, y)

    >>> metric
    MicroAverage(Jaccard): 95.10%

    """

    def __init__(self, model: base.Classifier):
        super().__init__()
        self.model = model

        self._label_map: dict[tuple, int] = {}
        self._r_label_map: dict[int, tuple] = {}
        self._labels: set[typing.Hashable] = set()

    def learn_one(self, x, y):
        self._labels.update(y.keys())

        aux = tuple(sorted(y.items()))
        if aux not in self._label_map:
            code = len(self._label_map)
            self._label_map[aux] = code
            self._r_label_map[code] = aux

        # Encode
        y_encoded = self._label_map[aux]

        # Update the classifier
        self.model.learn_one(x, y_encoded)

        return self

    def predict_proba_one(self, x, **kwargs):
        enc_probas = self.model.predict_proba_one(x, **kwargs)

        if not enc_probas:
            return {}

        enc_class = max(enc_probas, key=enc_probas.get)

        return {
            label_id: {
                bool(label_val): enc_probas[enc_class],
                not bool(label_val): 1 - enc_probas[enc_class],
            }
            for label_id, label_val in self._r_label_map[enc_class]
        }

    @classmethod
    def _unit_test_params(cls):
        from river import tree

        yield {"model": tree.HoeffdingTreeClassifier()}


# river/naive_bayes/__init__.py

"""Naive Bayes algorithms."""
from __future__ import annotations

from .bernoulli import BernoulliNB
from .complement import ComplementNB
from .gaussian import GaussianNB
from .multinomial import MultinomialNB

__all__ = ["BernoulliNB", "ComplementNB", "GaussianNB", "MultinomialNB"]


# river/naive_bayes/base.py

from __future__ import annotations

import abc
import math

import numpy as np
import pandas as pd
from scipy import sparse, special

from river import base


class BaseNB(base.MiniBatchClassifier):
    """Base Naive Bayes class."""

    @abc.abstractmethod
    def joint_log_likelihood(self, x: dict) -> float:
        """Compute the unnormalized posterior log-likelihood of x.

        The log-likelihood is `log P(c) + log P(x|c)`.

        """

    @abc.abstractmethod
    def joint_log_likelihood_many(self, X: pd.DataFrame) -> pd.DataFrame:
        """Compute the unnormalized posterior log-likelihood of x in mini-batches.

        The log-likelihood is `log P(c) + log P(x|c)`.

        """

    def predict_proba_one(self, x):
        """Return probabilities using the log-likelihoods."""
        jll = self.joint_log_likelihood(x)
        if not jll:
            return {}
        lse = special.logsumexp(list(jll.values()))
        return {label: math.exp(ll - lse) for label, ll in jll.items()}

    def predict_proba_many(self, X: pd.DataFrame) -> pd.DataFrame:
        """Return probabilities using the log-likelihoods in mini-batchs setting."""
        jll = self.joint_log_likelihood_many(X)
        if jll.empty:
            return jll
        lse = pd.Series(special.logsumexp(jll, axis=1))
        return np.exp(jll.subtract(lse.values, axis="rows"))

    @property
    def _multiclass(self):
        return True


def from_dict(data: dict) -> pd.DataFrame:
    """Convert a dict into a pandas dataframe.
    This function is faster than pd.from_dict (01/02/2021).

    Parameters
    ----------
    data
        Input data as dict.

    Returns
    --------
        Dict to pandas dataframe.

    """
    dict_data, index = list(data.values()), list(data.keys())
    return pd.DataFrame(data=dict_data, index=index, dtype="float32")


def one_hot_encode(y: pd.Series) -> pd.DataFrame:
    """One hot encode input pandas series into sparse pandas DataFrame.

    Parameters
    ----------
    y
        Pandas Series of strings.

    Returns
    --------
    One hot encoded sparse dataframe.

    """
    classes = np.unique(y)
    indices = np.searchsorted(classes, y)
    indptr = np.hstack((0, np.cumsum(np.in1d(y, classes))))
    data = np.empty_like(indices)
    data.fill(1)
    return pd.DataFrame.sparse.from_spmatrix(
        sparse.csr_matrix((data, indices, indptr), shape=(y.shape[0], classes.shape[0])),
        index=y.index,
        columns=[str(c) for c in classes],
    )


# river/naive_bayes/bernoulli.py

from __future__ import annotations

import collections
import math

import numpy as np
import pandas as pd
from scipy import sparse

from . import base

__all__ = ["BernoulliNB"]


class BernoulliNB(base.BaseNB):
    """Bernoulli Naive Bayes.

    Bernoulli Naive Bayes model learns from occurrences between features such as word
    counts and discrete classes. The input vector must contain positive values, such as
    counts or TF-IDF values.

    Parameters
    ----------
    alpha
        Additive (Laplace/Lidstone) smoothing parameter (use 0 for no smoothing).
    true_threshold
        Threshold for binarizing (mapping to booleans) features.

    Attributes
    ----------
    class_counts : collections.Counter
        Number of times each class has been seen.
    feature_counts : collections.defaultdict
        Total frequencies per feature and class.

    Examples
    --------

    >>> import pandas as pd
    >>> from river import compose
    >>> from river import feature_extraction
    >>> from river import naive_bayes

    >>> docs = [
    ...     ("Chinese Beijing Chinese", "yes"),
    ...     ("Chinese Chinese Shanghai", "yes"),
    ...     ("Chinese Macao", "yes"),
    ...     ("Tokyo Japan Chinese", "no")
    ... ]

    >>> model = compose.Pipeline(
    ...     ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
    ...     ("nb", naive_bayes.BernoulliNB(alpha=1))
    ... )

    >>> for sentence, label in docs:
    ...     model = model.learn_one(sentence, label)

    >>> model["nb"].p_class("yes")
    0.75
    >>> model["nb"].p_class("no")
    0.25

    >>> model.predict_proba_one("test")
    {'yes': 0.8831539823829913, 'no': 0.11684601761700895}

    >>> model.predict_one("test")
    'yes'

    You can train the model and make predictions in mini-batch mode using the class methods
    `learn_many` and `predict_many`.

    >>> df_docs = pd.DataFrame(docs, columns = ["docs", "y"])

    >>> X = pd.Series([
    ...    "Chinese Beijing Chinese",
    ...    "Chinese Chinese Shanghai",
    ...    "Chinese Macao",
    ...    "Tokyo Japan Chinese"
    ... ])

    >>> y = pd.Series(["yes", "yes", "yes", "no"])

    >>> model = compose.Pipeline(
    ...     ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
    ...     ("nb", naive_bayes.BernoulliNB(alpha=1))
    ... )

    >>> model = model.learn_many(X, y)

    >>> unseen = pd.Series(["Taiwanese Taipei", "Chinese Shanghai"])

    >>> model.predict_proba_many(unseen)
             no       yes
    0  0.116846  0.883154
    1  0.047269  0.952731

    >>> model.predict_many(unseen)
    0    yes
    1    yes
    dtype: object

    References
    ----------
    [^1]: [The Bernoulli model](https://nlp.stanford.edu/IR-book/html/htmledition/the-bernoulli-model-1.html)

    """

    def __init__(self, alpha=1.0, true_threshold=0.0):
        self.alpha = alpha
        self.true_threshold = true_threshold
        self.class_counts = collections.Counter()
        self.feature_counts = collections.defaultdict(collections.Counter)

    def learn_one(self, x, y):
        """Updates the model with a single observation.

        Parameters
        ----------
        x
            Dictionary of term frequencies.
        y
            Target class.

        Returns
        -------
        self

        """
        self.class_counts.update((y,))

        for i, xi in x.items():
            self.feature_counts[i].update({y: xi > self.true_threshold})

        return self

    def p_feature_given_class(self, f: str, c: str) -> float:
        num = self.feature_counts.get(f, {}).get(c, 0.0) + self.alpha
        den = self.class_counts[c] + self.alpha * 2
        return num / den

    def p_class(self, c: str) -> float:
        return self.class_counts[c] / sum(self.class_counts.values())

    def p_class_many(self) -> pd.DataFrame:
        return base.from_dict(self.class_counts).T[list(self.class_counts.keys())] / sum(
            self.class_counts.values()
        )

    def joint_log_likelihood(self, x):
        """Computes the joint log likelihood of input features.

        Parameters
        ----------
        x
            Dictionary of term frequencies.

        Returns
        -------
        Mapping between classes and joint log likelihood.

        """
        return {
            c: math.log(self.p_class(c))
            + sum(
                map(
                    math.log,
                    (
                        10e-10 + self.p_feature_given_class(f, c)
                        if f in x and x[f] > self.true_threshold
                        else 10e-10 + (1.0 - self.p_feature_given_class(f, c))
                        for f in self.feature_counts
                    ),
                )
            )
            for c in self.class_counts
        }

    def learn_many(self, X: pd.DataFrame, y: pd.Series):
        """Learn from a batch of count vectors.

        Parameters
        ----------
        X
            Count vectors.
        y
            Target classes.

        Returns
        -------
        self

        """
        # One hot encode y and convert it into sparse matrix
        y = base.one_hot_encode(y)
        columns, classes = X.columns, y.columns
        y = sparse.csc_matrix(y.sparse.to_coo()).T

        self.class_counts.update({c: count.item() for c, count in zip(classes, y.sum(axis=1))})

        if hasattr(X, "sparse"):
            X = sparse.csr_matrix(X.sparse.to_coo())
            X.data = X.data > self.true_threshold
        else:
            X = X > self.true_threshold

        fc = y @ X

        # Update the defaultdict feature counts by slicing the sparse matrix per row.
        # Each row of y @ X correspond to a class and each column to a token.
        for c, i in zip(classes, range(fc.shape[0])):
            if sparse.issparse(fc):
                counts = {c: {columns[f]: count for f, count in zip(fc[i].indices, fc[i].data)}}

            else:
                counts = {c: {f: count for f, count in zip(columns, fc[i])}}

            # Transform {classe_i: {token_1: f_1, ... token_n: f_n}} into:
            # [{token_1: {classe_i: f_1}},.. {token_n: {class_i: f_n}}]
            for dict_count in [
                {token: {c: f} for token, f in frequencies.items()}
                for c, frequencies in counts.items()
            ]:
                for f, count in dict_count.items():
                    self.feature_counts[f].update(count)

        return self

    def _feature_log_prob(self, columns: list) -> pd.DataFrame:
        """Compute log probabilities of input features.

        Parameters
        ----------
        columns
            List of input features.

        Returns
        -------
            Log probabilities of input features.

        """
        smooth_fc = np.log(
            base.from_dict(self.feature_counts)[list(self.class_counts.keys())].T.fillna(0)
            + self.alpha
        )[columns]

        smooth_cc = np.log(base.from_dict(self.class_counts) + self.alpha * 2)

        return smooth_fc.subtract(smooth_cc.values)

    def joint_log_likelihood_many(self, X: pd.DataFrame) -> pd.DataFrame:
        """Computes the joint log likelihood of input features.

        Parameters
        ----------
        X
            Term-frequency or TF-IDF pandas dataframe.

        Returns
        -------
        Input samples joint log likelihood.

        """
        unknown = [x for x in X.columns if x not in self.feature_counts]
        missing = [x for x in self.feature_counts if x not in X.columns]

        if unknown:
            X = X.drop(unknown, axis="columns")

        if missing:
            X[missing] = False

        index, columns = X.index, X.columns

        if not self.class_counts or not self.feature_counts:
            return pd.DataFrame(index=index)

        if hasattr(X, "sparse"):
            X = sparse.csr_matrix(X.sparse.to_coo())
            X.data = X.data > self.true_threshold
        else:
            X = X > self.true_threshold

        flp = self._feature_log_prob(columns)
        neg_p = np.log(1 - np.exp(flp))

        return pd.DataFrame(
            X @ (flp - neg_p).T + (np.log(self.p_class_many()) + neg_p.sum(axis=1).T).values,
            index=index,
            columns=self.class_counts.keys(),
        )


# river/naive_bayes/complement.py

from __future__ import annotations

import collections
import math

import numpy as np
import pandas as pd
from scipy import sparse

from river.base import tags

from . import base

__all__ = ["ComplementNB"]


class ComplementNB(base.BaseNB):
    """Naive Bayes classifier for multinomial models.

    Complement Naive Bayes model learns from occurrences between features such as word counts
    and discrete classes. ComplementNB is suitable for imbalance dataset.
    The input vector must contain positive values, such as counts or TF-IDF values.

    Parameters
    ----------
    alpha
        Additive (Laplace/Lidstone) smoothing parameter (use 0 for no smoothing).

    Attributes
    ----------
    class_dist : proba.Multinomial
        Class prior probability distribution.
    feature_counts : collections.defaultdict
        Total frequencies per feature and class.
    class_totals : collections.Counter
        Total frequencies per class.

    Examples
    --------

    >>> import pandas as pd
    >>> from river import compose
    >>> from river import feature_extraction
    >>> from river import naive_bayes

    >>> docs = [
    ...     ("Chinese Beijing Chinese", "yes"),
    ...     ("Chinese Chinese Shanghai", "yes"),
    ...     ("Chinese Macao", "maybe"),
    ...     ("Tokyo Japan Chinese", "no")
    ... ]

    >>> model = compose.Pipeline(
    ...     ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
    ...     ("nb", naive_bayes.ComplementNB(alpha=1))
    ... )

    >>> for sentence, label in docs:
    ...     model = model.learn_one(sentence, label)

    >>> model["nb"].p_class("yes")
    0.5

    >>> model["nb"].p_class("no")
    0.25

    >>> model["nb"].p_class("maybe")
    0.25

    >>> model.predict_proba_one("test")
    {'yes': 0.275, 'maybe': 0.375, 'no': 0.35}

    >>> model.predict_one("test")
    'maybe'

    You can train the model and make predictions in mini-batch mode using the class methods
    `learn_many` and `predict_many`.

    >>> df_docs = pd.DataFrame(docs, columns = ["docs", "y"])

    >>> X = pd.Series([
    ...    "Chinese Beijing Chinese",
    ...    "Chinese Chinese Shanghai",
    ...    "Chinese Macao",
    ...    "Tokyo Japan Chinese"
    ... ])

    >>> y = pd.Series(["yes", "yes", "maybe", "no"])

    >>> model = compose.Pipeline(
    ...     ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
    ...     ("nb", naive_bayes.ComplementNB(alpha=1))
    ... )

    >>> model = model.learn_many(X, y)

    >>> unseen = pd.Series(["Taiwanese Taipei", "Chinese Shanghai"])

    >>> model.predict_proba_many(unseen)
          maybe        no       yes
    0  0.415129  0.361624  0.223247
    1  0.248619  0.216575  0.534807

    >>> model.predict_many(unseen)
    0    maybe
    1      yes
    dtype: object

    References
    ----------
    [^1]: [Rennie, J.D., Shih, L., Teevan, J. and Karger, D.R., 2003. Tackling the poor assumptions of naive bayes text classifiers. In Proceedings of the 20th international conference on machine learning (ICML-03) (pp. 616-623)](https://people.csail.mit.edu/jrennie/papers/icml03-nb.pdf)
    [^2]: [StackExchange discussion](https://stats.stackexchange.com/questions/126009/complement-naive-bayes)

    """

    def __init__(self, alpha=1.0):
        self.alpha = alpha
        self.class_counts = collections.Counter()
        self.feature_counts = collections.defaultdict(collections.Counter)
        self.feature_totals = collections.Counter()
        self.class_totals = collections.Counter()

    def _more_tags(self):
        return {tags.POSITIVE_INPUT}

    def learn_one(self, x, y):
        """Updates the model with a single observation.

        Parameters
        ----------
        x
            Dictionary of term frequencies.
        y
            Target class.

        Returns
        -------
        self

        """
        self.class_counts.update((y,))

        for f, frequency in x.items():
            self.feature_counts[f].update({y: frequency})
            self.feature_totals.update({f: frequency})
            self.class_totals.update({y: frequency})

        return self

    def p_class(self, c):
        return self.class_counts[c] / sum(self.class_counts.values())

    def p_class_many(self) -> pd.DataFrame:
        return base.from_dict(self.class_counts).T[self.class_counts] / sum(
            self.class_counts.values()
        )

    def joint_log_likelihood(self, x):
        """Computes the joint log likelihood of input features.

        Parameters
        ----------
        x
            Dictionary of term frequencies.

        Returns
        -------
        Mapping between classes and joint log likelihood.

        """
        cc = {
            c: {
                f: self.feature_totals[f] + self.alpha - frequency.get(c, 0)
                for f, frequency in self.feature_counts.items()
            }
            for c in self.class_counts
        }

        return {
            c: sum(
                {
                    f: frequency * -math.log(cc[c].get(f, self.alpha) / sum(cc[c].values()))
                    for f, frequency in x.items()
                }.values()
            )
            for c in self.class_counts
        }

    def learn_many(self, X: pd.DataFrame, y: pd.Series):
        """Learn from a batch of count vectors.

        Parameters
        ----------
        X
            Count vectors.
        y
            Target classes.

        Returns
        -------
        self

        """
        y = base.one_hot_encode(y)
        columns, classes = X.columns, y.columns
        y = sparse.csc_matrix(y.sparse.to_coo()).T

        self.class_counts.update({c: count.item() for c, count in zip(classes, y.sum(axis=1))})

        if hasattr(X, "sparse"):
            X = sparse.csr_matrix(X.sparse.to_coo())

        fc = y @ X

        self.class_totals.update({c: count.item() for c, count in zip(classes, fc.sum(axis=1))})

        self.feature_totals.update(
            {c: count.item() for c, count in zip(columns, np.array(fc.sum(axis=0)).flatten())}
        )

        # Update feature counts by slicing the sparse matrix per column.
        # Each column correspond to a class.
        for c, i in zip(classes, range(fc.shape[0])):
            counts = {c: {columns[f]: count for f, count in zip(fc[i].indices, fc[i].data)}}

            # Transform {classe_i: {token_1: f_1, ... token_n: f_n}} into:
            # [{token_1: {classe_i: f_1}},.. {token_n: {class_i: f_n}}]
            for dict_count in [
                {token: {c: f} for token, f in frequencies.items()}
                for c, frequencies in counts.items()
            ]:
                for f, count in dict_count.items():
                    self.feature_counts[f].update(count)

        return self

    def _feature_log_prob(self, unknown: list, columns: list) -> pd.DataFrame:
        """Compute log probabilities of input features.

        Parameters
        ----------
        unknown
            List of features that are not part the vocabulary.
        columns
            List of input features.

        Returns
        -------
        Log probabilities of input features.

        """
        cc = (
            base.from_dict(self.feature_totals).squeeze().T
            + self.alpha
            - base.from_dict(self.feature_counts).fillna(0).T
        )

        sum_cc = cc.sum(axis=1).values

        cc[unknown] = self.alpha

        return -np.log(cc[columns].T / sum_cc)

    def joint_log_likelihood_many(self, X: pd.DataFrame) -> pd.DataFrame:
        """Computes the joint log likelihood of input features.

        Parameters
        ----------
        X
            Term-frequency or TF-IDF pandas dataframe.

        Returns
        -------
        Input samples joint log likelihood.

        """
        index, columns = X.index, X.columns
        unknown = [x for x in columns if x not in self.feature_counts]

        if not self.class_counts or not self.feature_counts:
            return pd.DataFrame(index=index)

        if hasattr(X, "sparse"):
            X = sparse.csr_matrix(X.sparse.to_coo())

        return pd.DataFrame(
            X @ self._feature_log_prob(unknown=unknown, columns=columns),
            index=index,
            columns=self.class_counts.keys(),
        )


# river/naive_bayes/gaussian.py

from __future__ import annotations

import collections
import functools
import math

import pandas as pd
from scipy import special

from river import base, proba

__all__ = ["GaussianNB"]


class GaussianNB(base.Classifier):
    """Gaussian Naive Bayes.

    A Gaussian distribution $G_{cf}$ is maintained for each class $c$ and each feature $f$. Each
    Gaussian is updated using the amount associated with each feature; the details can be be found
    in `proba.Gaussian`. The joint log-likelihood is then obtained by summing the log probabilities
    of each feature associated with each class.

    Examples
    --------

    >>> from river import naive_bayes
    >>> from river import stream
    >>> import numpy as np

    >>> X = np.array([[-1, -1], [-2, -1], [-3, -2], [1, 1], [2, 1], [3, 2]])
    >>> Y = np.array([1, 1, 1, 2, 2, 2])

    >>> model = naive_bayes.GaussianNB()

    >>> for x, y in stream.iter_array(X, Y):
    ...     _ = model.learn_one(x, y)

    >>> model.predict_one({0: -0.8, 1: -1})
    1

    """

    def __init__(self):
        self.class_counts = collections.Counter()
        self.gaussians = collections.defaultdict(
            functools.partial(collections.defaultdict, proba.Gaussian)
        )

    def learn_one(self, x, y):
        self.class_counts.update((y,))

        for i, xi in x.items():
            self.gaussians[y][i].update(xi)

        return self

    def predict_proba_one(self, x):
        """Return probabilities using the log-likelihoods."""
        jll = self.joint_log_likelihood(x)
        if not jll:
            return {}
        lse = special.logsumexp(list(jll.values()))
        return {label: math.exp(ll - lse) for label, ll in jll.items()}

    def p_class(self, c):
        return self.class_counts[c] / sum(self.class_counts.values())

    def joint_log_likelihood(self, x):
        return {
            c: math.log(self.p_class(c))
            + sum(math.log(10e-10 + gaussians[i](xi)) for i, xi in x.items())
            for c, gaussians in self.gaussians.items()
        }

    def joint_log_likelihood_many(self, X: pd.DataFrame):
        pass

    @property
    def _multiclass(self):
        return True


# river/naive_bayes/multinomial.py

from __future__ import annotations

import collections
import math

import numpy as np
import pandas as pd
from scipy import sparse

from river.base import tags

from . import base

__all__ = ["MultinomialNB"]


class MultinomialNB(base.BaseNB):
    """Naive Bayes classifier for multinomial models.

    Multinomial Naive Bayes model learns from occurrences between features such as word counts
    and discrete classes. The input vector must contain positive values, such as
    counts or TF-IDF values.


    Parameters
    ----------
    alpha
        Additive (Laplace/Lidstone) smoothing parameter (use 0 for no smoothing).

    Attributes
    ----------
    class_dist : proba.Multinomial
        Class prior probability distribution.
    feature_counts : collections.defaultdict
        Total frequencies per feature and class.
    class_totals : collections.Counter
        Total frequencies per class.

    Examples
    --------

    >>> import pandas as pd
    >>> from river import compose
    >>> from river import feature_extraction
    >>> from river import naive_bayes

    >>> docs = [
    ...     ("Chinese Beijing Chinese", "yes"),
    ...     ("Chinese Chinese Shanghai", "yes"),
    ...     ("Chinese Macao", "maybe"),
    ...     ("Tokyo Japan Chinese", "no")
    ... ]

    >>> model = compose.Pipeline(
    ...     ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
    ...     ("nb", naive_bayes.MultinomialNB(alpha=1))
    ... )

    >>> for sentence, label in docs:
    ...     model = model.learn_one(sentence, label)

    >>> model["nb"].p_class("yes")
    0.5

    >>> model["nb"].p_class("no")
    0.25

    >>> model["nb"].p_class("maybe")
    0.25

    >>> model.predict_proba_one("test")
    {'yes': 0.413, 'maybe': 0.310, 'no': 0.275}

    >>> model.predict_one("test")
    'yes'

    You can train the model and make predictions in mini-batch mode using the class methods
    `learn_many` and `predict_many`.

    >>> df_docs = pd.DataFrame(docs, columns = ["docs", "y"])

    >>> X = pd.Series([
    ...    "Chinese Beijing Chinese",
    ...    "Chinese Chinese Shanghai",
    ...    "Chinese Macao",
    ...    "Tokyo Japan Chinese"
    ... ])

    >>> y = pd.Series(["yes", "yes", "maybe", "no"])

    >>> model = compose.Pipeline(
    ...     ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
    ...     ("nb", naive_bayes.MultinomialNB(alpha=1))
    ... )

    >>> model = model.learn_many(X, y)

    >>> unseen = pd.Series(["Taiwanese Taipei", "Chinese Shanghai"])

    >>> model.predict_proba_many(unseen)
          maybe        no       yes
    0  0.373272  0.294931  0.331797
    1  0.160396  0.126733  0.712871

    >>> model.predict_many(unseen)
    0    maybe
    1      yes
    dtype: object

    References
    ----------
    [^1]: [Naive Bayes text classification](https://nlp.stanford.edu/IR-book/html/htmledition/naive-bayes-text-classification-1.html)

    """

    def __init__(self, alpha=1.0):
        self.alpha = alpha
        self.class_counts = collections.Counter()
        self.feature_counts = collections.defaultdict(collections.Counter)
        self.class_totals = collections.Counter()

    def _more_tags(self):
        return {tags.POSITIVE_INPUT}

    def learn_one(self, x, y):
        """Updates the model with a single observation.

        Parameters
        ----------
        x
            Dictionary of term frequencies.
        y
            Target class.

        Returns
        -------
        self

        """
        self.class_counts.update((y,))

        for f, frequency in x.items():
            self.feature_counts[f].update({y: frequency})
            self.class_totals.update({y: frequency})

        return self

    @property
    def classes_(self):
        return list(self.class_counts.keys())

    @property
    def n_terms(self):
        return len(self.feature_counts)

    def p_feature_given_class(self, f, c):
        num = self.feature_counts.get(f, {}).get(c, 0.0) + self.alpha
        den = self.class_totals[c] + self.alpha * self.n_terms
        return num / den

    def p_class(self, c) -> float:
        return self.class_counts[c] / sum(self.class_counts.values())

    def p_class_many(self) -> pd.DataFrame:
        return base.from_dict(self.class_counts).T[list(self.class_counts.keys())] / sum(
            self.class_counts.values()
        )

    def joint_log_likelihood(self, x):
        """Computes the joint log likelihood of input features.

        Parameters
        ----------
        x
            Dictionary of term frequencies.

        Returns
        -------
        Mapping between classes and joint log likelihood.

        """
        return {
            c: math.log(self.p_class(c))
            + sum(
                frequency * math.log(self.p_feature_given_class(f, c)) for f, frequency in x.items()
            )
            for c in self.classes_
        }

    def learn_many(self, X: pd.DataFrame, y: pd.Series):
        """Learn from a batch of count vectors.

        Parameters
        ----------
        X
            Count vectors.
        y
            Target classes.

        Returns
        -------
        self

        """
        y = base.one_hot_encode(y)
        columns, classes = X.columns, y.columns
        y = sparse.csc_matrix(y.sparse.to_coo()).T

        self.class_counts.update({c: count.item() for c, count in zip(classes, y.sum(axis=1))})

        if hasattr(X, "sparse"):
            X = sparse.csr_matrix(X.sparse.to_coo())

        fc = y @ X

        self.class_totals.update({c: count.item() for c, count in zip(classes, fc.sum(axis=1))})

        # Update feature counts by slicing the sparse matrix per column.
        # Each column correspond to a class.
        for c, i in zip(classes, range(fc.shape[0])):
            counts = {c: {columns[f]: count for f, count in zip(fc[i].indices, fc[i].data)}}

            # Transform {classe_i: {token_1: f_1, ... token_n: f_n}} into:
            # [{token_1: {classe_i: f_1}},.. {token_n: {class_i: f_n}}]
            for dict_count in [
                {token: {c: f} for token, f in frequencies.items()}
                for c, frequencies in counts.items()
            ]:
                for f, count in dict_count.items():
                    self.feature_counts[f].update(count)

        return self

    def _feature_log_prob(self, columns: list, known: list, unknown: list) -> pd.DataFrame:
        """Compute log probabilities of input features.

        Parameters
        ----------
        columns
            List of input features.
        known
            List of input features that are part of the vocabulary.
        unknown
            List of input features that are not part the vocabulary.

        Returns
        -------
        Log probabilities of input features.

        """
        smooth_fc = np.log(base.from_dict(self.feature_counts).fillna(0).T[known] + self.alpha)
        smooth_fc[unknown] = np.log(self.alpha)

        smooth_cc = np.log(base.from_dict(self.class_totals) + self.alpha * self.n_terms)

        return smooth_fc.subtract(smooth_cc.values, axis="rows")[columns].T

    def joint_log_likelihood_many(self, X: pd.DataFrame) -> pd.DataFrame:
        """Computes the joint log likelihood of input features.

        Parameters
        ----------
        X
            Term-frequency or TF-IDF pandas dataframe.

        Returns
        -------
        Input samples joint log likelihood.

        """
        index, columns = X.index, X.columns
        known, unknown = [], []

        if not self.class_counts or not self.feature_counts:
            return pd.DataFrame(index=index)

        for f in columns:
            if f in self.feature_counts:
                known.append(f)
            else:
                unknown.append(f)

        if hasattr(X, "sparse"):
            X = sparse.csr_matrix(X.sparse.to_coo())

        return pd.DataFrame(
            X @ self._feature_log_prob(columns=columns, known=known, unknown=unknown)
            + np.log(self.p_class_many()).values,
            index=index,
            columns=self.class_totals.keys(),
        )


# river/naive_bayes/test_naive_bayes.py

from __future__ import annotations

import pandas as pd
import pytest
from sklearn import naive_bayes as sk_naive_bayes

from river import compose, feature_extraction, naive_bayes


def river_models():
    """List of Naive Bayes models to test."""
    yield from [
        naive_bayes.MultinomialNB,
        naive_bayes.BernoulliNB,
        naive_bayes.ComplementNB,
    ]


def sklearn_models():
    """Mapping between Naives Bayes river models and sklearn models."""
    yield from [
        (naive_bayes.MultinomialNB, sk_naive_bayes.MultinomialNB),
        (naive_bayes.BernoulliNB, sk_naive_bayes.BernoulliNB),
        (naive_bayes.ComplementNB, sk_naive_bayes.ComplementNB),
    ]


def yield_dataset():
    """Incremental dataset."""
    yield from [
        ("Chinese Beijing Chinese", "yes"),
        ("Chinese Chinese Shanghai", "yes"),
        ("Chinese Macao", "yes"),
        ("Tokyo Japan Chinese", "no"),
    ]


def yield_batch_dataset():
    """Batch dataset."""
    for x, y in yield_dataset():
        yield pd.Series([x]), pd.Series([y])


def yield_unseen_data():
    yield from [
        "Chinese Beijing Chinese",
        "Chinese Chinese Shanghai",
        "Chinese Macao",
        "Tokyo Japan Chinese",
        "new unseen data",
        "Taiwanese Taipei",
        "Chinese Shanghai" "Shanghai",
        "Chinese",
        "Tokyo Macao",
        "Tokyo Tokyo",
        "Macao Macao new",
        "new",
    ]


def yield_batch_unseen_data():
    yield from [pd.Series(x) for x in yield_unseen_data()]


@pytest.mark.parametrize(
    "model",
    [
        pytest.param(
            compose.Pipeline(
                ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
                ("model", model(alpha=alpha)),
            ),
            id=f"{model.__name__} - {alpha}",
        )
        for model in river_models()
        for alpha in [alpha for alpha in range(1, 4)]
    ],
)
def test_learn_one_methods(model):
    """Assert that the methods of the Naives Bayes class behave correctly."""
    assert model.predict_proba_one("not fitted yet") == {}
    assert model.predict_one("not fitted yet") is None

    for x, y in yield_dataset():
        model = model.learn_one(x, y)

    # Check class methods.
    if isinstance(model["model"], naive_bayes.ComplementNB) or isinstance(
        model["model"], naive_bayes.MultinomialNB
    ):
        assert model["model"]._more_tags() == {"positive input"}
    assert model["model"]._multiclass


@pytest.mark.parametrize(
    "model, batch_model",
    [
        pytest.param(
            compose.Pipeline(
                ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
                ("model", model(alpha=alpha)),
            ),
            compose.Pipeline(
                ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
                ("model", model(alpha=alpha)),
            ),
            id=f"{model.__name__} - {alpha}",
        )
        for model in river_models()
        for alpha in [alpha for alpha in range(1, 4)]
    ],
)
def test_learn_many_vs_learn_one(model, batch_model):
    """Assert that the Naive Bayes river models provide the same results when learning in
    incremental and mini-batch modes. The models tested are MultinomialNB, BernoulliNB and
    ComplementNB with differents alpha parameters..
    """
    for x, y in yield_dataset():
        model = model.learn_one(x, y)

    for x, y in yield_batch_dataset():
        batch_model = batch_model.learn_many(x, y)

    assert model["model"].p_class("yes") == batch_model["model"].p_class("yes")
    assert model["model"].p_class("no") == batch_model["model"].p_class("no")

    # Assert batch and incremental models give same results
    for x, x_batch in zip(yield_unseen_data(), yield_batch_unseen_data()):
        assert model.predict_proba_one(x)["yes"] == pytest.approx(
            batch_model.predict_proba_many(x_batch)["yes"][0]
        )

        assert model.predict_proba_one(x)["no"] == pytest.approx(
            batch_model.predict_proba_many(x_batch)["no"][0]
        )

    # Assert class probabilities are the same when trainig Naive Bayes in pure online and in batch.
    assert model["model"].p_class("yes") == batch_model["model"].p_class("yes")
    assert model["model"].p_class("no") == batch_model["model"].p_class("no")

    # Assert conditionnal probabilities are the same when training Naive Bayes in pure online and
    # in batch.
    if isinstance(model["model"], naive_bayes.BernoulliNB) or isinstance(
        model["model"], naive_bayes.MultinomialNB
    ):
        inc_cp = model["model"].p_feature_given_class
        batch_cp = batch_model["model"].p_feature_given_class

        assert inc_cp("Chinese", "yes") == batch_cp("Chinese", "yes")
        assert inc_cp("Tokyo", "yes") == batch_cp("Tokyo", "yes")
        assert inc_cp("Japan", "yes") == batch_cp("Japan", "yes")
        assert inc_cp("Chinese", "no") == batch_cp("Chinese", "no")
        assert inc_cp("Tokyo", "no") == batch_cp("Tokyo", "no")
        assert inc_cp("Japan", "no") == batch_cp("Japan", "no")
        assert inc_cp("unseen", "yes") == batch_cp("unseen", "yes")

    assert model["model"].class_counts == batch_model["model"].class_counts
    assert model["model"].feature_counts == batch_model["model"].feature_counts

    if isinstance(model["model"], naive_bayes.ComplementNB) or isinstance(
        model["model"], naive_bayes.MultinomialNB
    ):
        assert model["model"].class_totals == batch_model["model"].class_totals

    if isinstance(model["model"], sk_naive_bayes.ComplementNB):
        assert model["model"].feature_totals == batch_model["model"].feature_totals


@pytest.mark.parametrize(
    "batch_model",
    [
        pytest.param(
            compose.Pipeline(
                ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
                ("model", model(alpha=alpha)),
            ),
            id=f"{model.__name__} - {alpha}",
        )
        for model in river_models()
        for alpha in [alpha for alpha in range(1, 4)]
    ],
)
def test_learn_many_not_fit(batch_model):
    """Ensure that Naives Bayes models return an empty DataFrame when not yet fitted. Also check
    that the predict_proba_many method keeps the index.
    """
    assert batch_model.predict_proba_many(
        pd.Series(["new", "unseen"], index=["river", "rocks"])
    ).equals(pd.DataFrame(index=["river", "rocks"]))

    assert batch_model.predict_many(pd.Series(["new", "unseen"], index=["river", "rocks"])).equals(
        pd.DataFrame(index=["river", "rocks"])
    )


@pytest.mark.parametrize(
    "model, sk_model, bag",
    [
        pytest.param(
            compose.Pipeline(
                ("tokenize", feature_extraction.BagOfWords(lowercase=False)),
                ("model", model(alpha=alpha)),
            ),
            sk_model(alpha=alpha),
            feature_extraction.BagOfWords(lowercase=False),
            id=f"{model.__name__} - {alpha}",
        )
        for model, sk_model in sklearn_models()
        for alpha in [alpha for alpha in range(1, 4)]
    ],
)
def test_river_vs_sklearn(model, sk_model, bag):
    """Assert that river Naive Bayes models and sklearn Naive Bayes models provide the same results
    when the input data are the same. Also check that the behaviour of Naives Bayes models are the
    same with dense and sparse dataframe. Models tested are MultinomialNB, BernoulliNB and
    ComplementNB with differents alpha parameters.
    """
    for x, y in yield_batch_dataset():
        model = model.learn_many(x, y)

    X = pd.concat([x for x, _ in yield_batch_dataset()])
    y = pd.concat([y for _, y in yield_batch_dataset()])

    sk_model = sk_model.fit(X=bag.transform_many(X), y=y)

    # Assert river produce same results as sklearn using sparse dataframe:
    for sk_preds, river_preds in zip(
        sk_model.predict_proba(bag.transform_many(X)),
        model.predict_proba_many(X).values,
    ):
        for sk_pred, river_pred in zip(sk_preds, river_preds):
            assert river_pred == pytest.approx(1 - sk_pred) or river_pred == pytest.approx(sk_pred)

    # Assert river produce same results as sklearn using dense dataframe:
    for sk_preds, river_preds in zip(
        sk_model.predict_proba(bag.transform_many(X).sparse.to_dense()),
        model["model"].predict_proba_many(bag.transform_many(X).sparse.to_dense()).values,
    ):
        for sk_pred, river_pred in zip(sk_preds, river_preds):
            assert river_pred == pytest.approx(1 - sk_pred) or river_pred == pytest.approx(sk_pred)


# river/neural_net/__init__.py

"""Neural networks."""
from __future__ import annotations

from . import activations
from .mlp import MLPRegressor

__all__ = ["activations", "MLPRegressor"]


# river/neural_net/activations.py

from __future__ import annotations

import abc

import numpy as np

__all__ = ["ReLU", "Sigmoid", "Identity"]


class Activation:
    """An activation function.

    Each activation function is represented by a class and implements two methods. The first method
    is `apply`, which evaluates the activation on an array. The second method is `gradient`, which
    computes the gradient with respect to the input array. Both methods are intended to be pure
    with no side-effects. In other words they do not modify their inputs.

    """

    @abc.abstractstaticmethod
    def apply(self, z):
        """Apply the activation function to a layer output z."""

    @abc.abstractstaticmethod
    def gradient(self, z):
        """Return the gradient with respect to a layer output z."""


class ReLU(Activation):
    """Rectified Linear Unit (ReLU) activation function."""

    @staticmethod
    def apply(z):
        a = np.copy(z)
        a[a < 0] = 0
        return a

    @staticmethod
    def gradient(z):
        a = np.zeros_like(z, dtype=z.dtype)
        a[z > 0] = 1
        return a


class Sigmoid(Activation):
    """Sigmoid activation function."""

    @staticmethod
    def apply(z):
        return 1 / (1 + np.exp(-z))

    @staticmethod
    def gradient(z):
        s = Sigmoid.apply(z)
        return s * (1 - s)


class Identity(Activation):
    """Identity activation function."""

    @staticmethod
    def apply(z):
        return np.copy(z)

    @staticmethod
    def gradient(z):
        return np.ones_like(z, dtype=z.dtype)


# river/neural_net/mlp.py

from __future__ import annotations

import collections
import copy
import functools

import numpy as np
import pandas as pd

from river import base, optim

__all__ = ["MLPRegressor"]


def xavier_init(dims: tuple[int, ...], seed: int | None = None):
    """Xavier weight initialization.

    References
    ----------
    [^1]: Glorot, X. and Bengio, Y., 2010, March. Understanding the difficulty of training deep
        feedforward neural networks. In Proceedings of the thirteenth international conference on
        artificial intelligence and statistics (pp. 249-256).

    """

    rng = np.random.RandomState(seed)

    w = {}  # weights
    b = {}  # biases

    for i in range(len(dims) - 1):
        w[i + 1] = rng.randn(dims[i], dims[i + 1]) / np.sqrt(dims[i])
        b[i + 1] = np.zeros(dims[i + 1], dtype=w[i + 1].dtype)

    return w, b


class MLP:
    """Multi-layer Perceptron.

    The `_forward` method propagates input forward through the network and stores the outputs
    of each layer along the way. The `_backward` method then goes through each layer in reverse
    order and updates the weights within each layer.

    """

    def __init__(
        self,
        hidden_dims: tuple[int],
        activations,
        loss: optim.losses.Loss,
        optimizer: optim.base.Optimizer,
        seed: int | None = None,
    ):
        self.activations = activations
        self.hidden_dims = hidden_dims
        self.loss = loss
        self.optimizer = optimizer
        self.seed = seed
        self._optimizers: collections.defaultdict = collections.defaultdict(
            functools.partial(copy.deepcopy, optimizer)
        )

    @property
    def n_layers(self) -> int:
        """Return the number of layers in the network.

        The number of layers is equal to the number of hidden layers plus 2. The 2 accounts for the
        input layer and the output layer.

        """
        return len(self.hidden_dims) + 2

    def _forward(self, X: pd.DataFrame) -> tuple[dict, dict]:
        """Execute a forward pass through the neural network.

        Parameters
        ----------
        X
            A DataFrame of shape (batch_size, n_features).

        Returns
        -------
        The output of each layer.
        The output of each activation function.

        """

        z = {}  # z = w(x) + b
        a = {}  # a = f(z)

        # As a convention, we sort the columns of the input DataFrame. This allows the user to
        # pass DataFrames with different column orders without having to worry.
        a[1] = X.values[:, np.argsort(X.columns.to_numpy())]

        for i in range(2, self.n_layers + 1):
            z[i] = np.dot(a[i - 1], self.w[i - 1]) + self.b[i - 1]
            a[i] = self.activations[i - 2].apply(z[i])

        return z, a

    def _backward(self, z: dict, a: dict, y: pd.DataFrame):
        """Execute a backward pass through the neural network, aka. backpropagation.

        Parameters
        ----------
        z
            The output of each layer.
        a
            The output of each activation function.
        y
            A DataFrame of shape (batch_size, n_targets).

        """

        y = y.values[:, np.argsort(y.columns.to_numpy())]

        # Determine the partial derivative and delta for the output layer
        y_pred = a[self.n_layers]
        final_activation = self.activations[self.n_layers - 2]
        delta = self.loss.gradient(y, y_pred) * final_activation.gradient(y_pred)
        dw = np.dot(a[self.n_layers - 1].T, delta)

        update_params = {self.n_layers - 1: (dw, delta)}

        # Go through the layers in reverse order
        for i in range(self.n_layers - 2, 0, -1):
            delta = np.dot(delta, self.w[i + 1].T) * self.activations[i - 1].gradient(z[i + 1])
            dw = np.dot(a[i].T, delta)
            update_params[i] = (dw, delta)

        # Update the parameters
        for i, (dw, delta) in update_params.items():
            self._optimizers[i, 0].step(w=self.w[i], g=dw)
            self._optimizers[i, 1].step(w=self.b[i], g=np.mean(delta, axis=0))

    def learn_many(self, X: pd.DataFrame, y: pd.DataFrame):
        """Train the network.

        Parameters
        ----------
        X
            A DataFrame of shape (batch_size, n_features).
        y
            A DataFrame of shape (batch_size, n_targets).

        """

        # We expect y to be 2D, even if there is only one target to predict
        if isinstance(y, pd.Series):
            y = y.to_frame()

        # The weights need initializing during the first call to partial_fit
        if not hasattr(self, "w"):
            self.w, self.b = xavier_init(
                dims=(X.shape[1], *self.hidden_dims, y.shape[1]), seed=self.seed
            )
            self.features = X.columns.to_numpy()
            self.features.sort()
            self.targets = y.columns.to_numpy()

        z, a = self._forward(X)
        self._backward(z, a, y)

        return self

    def __call__(self, X: pd.DataFrame):
        """Make predictions.

        Parameters
        ----------
        X
            A DataFrame of shape (batch_size, n_features).

        """
        _, a = self._forward(X)
        y_pred = a[self.n_layers]
        return pd.DataFrame(y_pred, columns=self.targets, index=X.index)


class MLPRegressor(base.Regressor, MLP):
    """Multi-layer Perceptron for regression.

    This model is still work in progress. Here are some features that still need implementing:

    - `learn_one` and `predict_one` just cast the input `dict` to a single row dataframe and then
        call `learn_many` and `predict_many` respectively. This is very inefficient.
    - Not all of the optimizers in the `optim` module can be used as they are not all vectorised.
    - Emerging and disappearing features are not supported. Each instance/batch has to have the
        same features.
    - The gradient haven't been numerically checked.

    Parameters
    ----------
    hidden_dims
        The dimensions of the hidden layers. For example, specifying `(10, 20)` means that there
        are two hidden layers with 10 and 20 neurons, respectively. Note that the number of layers
        the network contains is equal to the number of hidden layers plus two (to account for the
        input and output layers).
    activations
        The activation functions to use at each layer, including the input and output layers.
        Therefore you need to specify three activation if you specify one hidden layer.
    loss
        Loss function. Defaults to `optim.losses.Squared`.
    optimizer
        Optimizer. Defaults to `optim.SGD` with the learning rate set to `0.01`.
    seed
        Random number generation seed. Set this for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import neural_net as nn
    >>> from river import optim
    >>> from river import preprocessing as pp
    >>> from river import metrics

    >>> model = (
    ...     pp.StandardScaler() |
    ...     nn.MLPRegressor(
    ...         hidden_dims=(5,),
    ...         activations=(
    ...             nn.activations.ReLU,
    ...             nn.activations.ReLU,
    ...             nn.activations.Identity
    ...         ),
    ...         optimizer=optim.SGD(1e-3),
    ...         seed=42
    ...     )
    ... )

    >>> dataset = datasets.TrumpApproval()

    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 1.589827

    You can also use this to process mini-batches of data.

    >>> model = (
    ...     pp.StandardScaler() |
    ...     nn.MLPRegressor(
    ...         hidden_dims=(10,),
    ...         activations=(
    ...             nn.activations.ReLU,
    ...             nn.activations.ReLU,
    ...             nn.activations.ReLU
    ...         ),
    ...         optimizer=optim.SGD(1e-4),
    ...         seed=42
    ...     )
    ... )

    >>> dataset = datasets.TrumpApproval()
    >>> batch_size = 32

    >>> for epoch in range(10):
    ...     for xb in pd.read_csv(dataset.path, chunksize=batch_size):
    ...         yb = xb.pop('five_thirty_eight')
    ...         y_pred = model.predict_many(xb)
    ...         model = model.learn_many(xb, yb)

    >>> model.predict_many(xb)
          five_thirty_eight
    992           39.361609
    993           46.398536
    994           42.094086
    995           40.195802
    996           40.782954
    997           40.839678
    998           40.896403
    999           48.362659
    1000          42.021849

    """

    def __init__(
        self,
        hidden_dims,
        activations,
        loss: optim.losses.Loss | None = None,
        optimizer: optim.base.Optimizer | None = None,
        seed: int | None = None,
    ):
        super().__init__(
            hidden_dims=hidden_dims,
            activations=activations,
            loss=loss or optim.losses.Squared(),
            optimizer=optimizer or optim.SGD(0.01),
            seed=seed,
        )

    @classmethod
    def _default_params(self):
        from . import activations

        return {
            "hidden_dims": (20,),
            "activations": (activations.ReLU, activations.ReLU, activations.Identity),
        }

    def predict_many(self, X):
        if not hasattr(self, "w"):
            return pd.DataFrame({0: 0}, index=X.index)
        return self(X)

    def learn_one(self, x, y):
        # Multi-output
        if isinstance(y, dict):
            return self.learn_many(X=pd.DataFrame([x]), y=pd.DataFrame([y]))

        # Single output
        return self.learn_many(X=pd.DataFrame([x]), y=pd.Series([y]))

    def predict_one(self, x):
        y_pred = self.predict_many(X=pd.DataFrame([x]))

        # Multi-output
        if len(y_pred.columns) > 1:
            return y_pred.iloc[0].to_dict()

        # Single output
        return y_pred.iloc[0, 0]


# river/optim/__init__.py

"""Stochastic optimization."""
from __future__ import annotations

from . import base, initializers, losses, schedulers
from .ada_bound import AdaBound
from .ada_delta import AdaDelta
from .ada_grad import AdaGrad
from .ada_max import AdaMax
from .adam import Adam
from .ams_grad import AMSGrad
from .average import Averager
from .ftrl import FTRLProximal
from .momentum import Momentum
from .nadam import Nadam
from .nesterov import NesterovMomentum
from .rms_prop import RMSProp
from .sgd import SGD

__all__ = [
    "base",
    "AdaBound",
    "AdaDelta",
    "AdaGrad",
    "Adam",
    "AMSGrad",
    "AdaMax",
    "Averager",
    "FTRLProximal",
    "initializers",
    "losses",
    "Momentum",
    "Nadam",
    "NesterovMomentum",
    "Optimizer",
    "RMSProp",
    "schedulers",
    "SGD",
]


# river/optim/ada_bound.py

from __future__ import annotations

import collections
import math

from river import optim, utils

__all__ = ["AdaBound"]


class AdaBound(optim.base.Optimizer):
    """AdaBound optimizer.

    Parameters
    ----------
    lr
        The learning rate.
    beta_1
    beta_2
    eps
    gamma
    final_lr

    Attributes
    ----------
    m : collections.defaultdict
    s : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.AdaBound()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 87.90%

    References
    ----------
    [^1]: [Luo, L., Xiong, Y., Liu, Y. and Sun, X., 2019. Adaptive gradient methods with dynamic bound of learning rate. arXiv preprint arXiv:1902.09843](https://arxiv.org/abs/1902.09843)

    """

    def __init__(self, lr=1e-3, beta_1=0.9, beta_2=0.999, eps=1e-8, gamma=1e-3, final_lr=0.1):
        super().__init__(lr)
        self.base_lr = lr
        self.final_lr = final_lr
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.eps = eps
        self.gamma = gamma
        self.m = collections.defaultdict(float)
        self.v = collections.defaultdict(float)

    def _step_with_dict(self, w, g):
        bias_1 = 1 - self.beta_1 ** (self.n_iterations + 1)
        bias_2 = 1 - self.beta_2 ** (self.n_iterations + 1)

        step_size = self.learning_rate * math.sqrt(bias_2) / bias_1
        self.final_lr *= self.learning_rate / self.base_lr

        lower_bound = self.final_lr * (1 - 1 / (self.gamma * (self.n_iterations + 1) + 1))
        upper_bound = self.final_lr * (1 + 1 / (self.gamma * (self.n_iterations + 1)))

        for i, gi in g.items():
            self.m[i] = self.beta_1 * self.m[i] + (1 - self.beta_1) * gi
            self.v[i] = self.beta_2 * self.v[i] + (1 - self.beta_2) * gi**2

            step_size_bound = step_size / (math.sqrt(self.v[i]) + self.eps)

            w[i] -= utils.math.clamp(step_size_bound, lower_bound, upper_bound) * self.m[i]

        return w


# river/optim/ada_delta.py

from __future__ import annotations

import collections

from river import optim

__all__ = ["AdaDelta"]


class AdaDelta(optim.base.Optimizer):
    """AdaDelta optimizer.

    Parameters
    ----------
    rho
    eps

    Attributes
    ----------
    g2 : collections.defaultdict
    s2 : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.AdaDelta()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 80.56%

    References
    ----------
    [^1]: [Zeiler, M.D., 2012. Adadelta: an adaptive learning rate method. arXiv preprint arXiv:1212.5701.](https://arxiv.org/pdf/1212.5701.pdf)

    """

    def __init__(self, rho=0.95, eps=1e-8):
        super().__init__(lr=None)
        self.rho = rho
        self.eps = eps
        self.g2 = collections.defaultdict(float)
        self.s2 = collections.defaultdict(float)

    def _rms(self, x):
        return (x + self.eps) ** 0.5

    def _step_with_dict(self, w, g):
        for i, gi in g.items():
            # Accumulate the gradient
            self.g2[i] = self.rho * self.g2[i] + (1 - self.rho) * gi**2

            # Compute the update
            step = -self._rms(self.s2[i]) / self._rms(self.g2[i]) * gi

            # Accumulate the update
            self.s2[i] = self.rho * self.s2[i] + (1 - self.rho) * step**2

            # Apply the update
            w[i] += step

        return w


# river/optim/ada_grad.py

from __future__ import annotations

import collections

from river import optim

__all__ = ["AdaGrad"]


class AdaGrad(optim.base.Optimizer):
    """AdaGrad optimizer.

    Parameters
    ----------
    lr
    eps

    Attributes
    ----------
    g2 : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.AdaGrad()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 88.01%

    References
    ----------
    [^1]: [Duchi, J., Hazan, E. and Singer, Y., 2011. Adaptive subgradient methods for online learning and stochastic optimization. Journal of machine learning research, 12(Jul), pp.2121-2159.](http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf)

    """

    def __init__(self, lr=0.1, eps=1e-8):
        super().__init__(lr)
        self.eps = eps
        self.g2 = collections.defaultdict(float)

    def _step_with_dict(self, w, g):
        for i, gi in g.items():
            self.g2[i] += gi**2
            w[i] -= self.learning_rate / (self.g2[i] + self.eps) ** 0.5 * gi

        return w


# river/optim/ada_max.py

from __future__ import annotations

import collections

from river import optim

__all__ = ["AdaMax"]


class AdaMax(optim.base.Optimizer):
    """AdaMax optimizer.

    Parameters
    ----------
    lr
    beta_1
    beta_2
    eps

    Attributes
    ----------
    m : collections.defaultdict
    v : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.AdaMax()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 87.53%

    References
    ----------
    [^1]: [Kingma, D.P. and Ba, J., 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.](https://arxiv.org/pdf/1412.6980.pdf)
    [^2]: [Ruder, S., 2016. An overview of gradient descent optimization algorithms. arXiv preprint arXiv:1609.04747.](http://ruder.io/optimizing-gradient-descent/index.html#adamax)

    """

    def __init__(self, lr=0.1, beta_1=0.9, beta_2=0.999, eps=1e-8):
        super().__init__(lr)
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.eps = eps
        self.m = collections.defaultdict(float)
        self.u = collections.defaultdict(float)

    def _step_with_dict(self, w, g):
        # Correct bias for `m`
        learning_rate = self.learning_rate / (1 - self.beta_1 ** (self.n_iterations + 1))

        for i, gi in g.items():
            self.m[i] = self.beta_1 * self.m[i] + (1 - self.beta_1) * gi
            self.u[i] = max(self.beta_2 * self.u[i], abs(gi))
            w[i] -= learning_rate * self.m[i] / (self.u[i] + self.eps)

        return w


# river/optim/adam.py

from __future__ import annotations

import collections

import numpy as np

from river import optim, utils

__all__ = ["Adam"]


class Adam(optim.base.Optimizer):
    """Adam optimizer.

    Parameters
    ----------
    lr
    beta_1
    beta_2
    eps

    Attributes
    ----------
    m : collections.defaultdict
    v : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.Adam()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 86.50%

    References
    ----------
    [^1]: [Kingma, D.P. and Ba, J., 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980.](https://arxiv.org/pdf/1412.6980.pdf)

    """

    def __init__(self, lr=0.1, beta_1=0.9, beta_2=0.999, eps=1e-8):
        super().__init__(lr)
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.eps = eps
        self.m = None
        self.v = None

    def _step_with_dict(self, w, g):
        if self.m is None:
            self.m = collections.defaultdict(float)
            self.v = collections.defaultdict(float)

        # Correct bias for `v`
        lr = self.learning_rate * (1 - self.beta_2 ** (self.n_iterations + 1)) ** 0.5
        # Correct bias for `m`
        lr /= 1 - self.beta_1 ** (self.n_iterations + 1)

        for i, gi in g.items():
            self.m[i] = self.beta_1 * self.m[i] + (1 - self.beta_1) * gi
            self.v[i] = self.beta_2 * self.v[i] + (1 - self.beta_2) * gi**2
            w[i] -= lr * self.m[i] / (self.v[i] ** 0.5 + self.eps)

        return w

    def _step_with_vector(self, w, g):
        if self.m is None:
            if isinstance(w, np.ndarray):
                self.m = np.zeros_like(w)
                self.v = np.zeros_like(w)
            else:
                self.m = utils.VectorDict()
                self.v = utils.VectorDict()

        lr = self.learning_rate * (1 - self.beta_2 ** (self.n_iterations + 1)) ** 0.5
        lr /= 1 - self.beta_1 ** (self.n_iterations + 1)

        self.m = self.beta_1 * self.m + (1 - self.beta_1) * g
        self.v = self.beta_2 * self.v + (1 - self.beta_2) * g**2
        w -= lr * self.m / (self.v**0.5 + self.eps)

        return w


# river/optim/ams_grad.py

from __future__ import annotations

import collections

from river import optim

__all__ = ["AMSGrad"]


class AMSGrad(optim.base.Optimizer):
    """AMSGrad optimizer.

    Parameters
    ----------
    lr
        The learning rate.
    beta_1
    beta_2
    eps
    correct_bias

    Attributes
    ----------
    m : collections.defaultdict
    v : collections.defaultdict
    v_hat : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.AMSGrad()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 86.57%

    References
    ----------
    [^1]: [Reddi, S.J., Kale, S. and Kumar, S., 2019. On the convergence of adam and beyond. arXiv preprint arXiv:1904.09237](https://arxiv.org/pdf/1904.09237.pdf)

    """

    def __init__(
        self,
        lr: int | float | optim.base.Scheduler = 0.1,
        beta_1=0.9,
        beta_2=0.999,
        eps=1e-8,
        correct_bias=True,
    ):
        super().__init__(lr)
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.eps = eps
        self.correct_bias = correct_bias
        self.m: dict[str, float] = collections.defaultdict(float)
        self.v: dict[str, float] = collections.defaultdict(float)
        self.v_hat: dict[str, float] = collections.defaultdict(float)

    def _step_with_dict(self, w, g):
        lr = self.learning_rate

        if self.correct_bias:
            # Correct bias for `v`
            lr *= (1 - self.beta_2 ** (self.n_iterations + 1)) ** 0.5
            # Correct bias for `m`
            lr /= 1 - self.beta_1 ** (self.n_iterations + 1)

        for i, gi in g.items():
            self.m[i] = self.beta_1 * self.m[i] + (1 - self.beta_1) * gi
            self.v[i] = self.beta_2 * self.v[i] + (1 - self.beta_2) * gi**2
            self.v_hat[i] = max(self.v_hat[i], self.v[i])

            w[i] -= lr * self.m[i] / (self.v_hat[i] ** 0.5 + self.eps)

        return w


# river/optim/average.py

from __future__ import annotations

import collections

from river import optim


class Averager(optim.base.Optimizer):

    """Averaged stochastic gradient descent.

    This is a wrapper that can be applied to any stochastic gradient descent optimiser. Note that
    this implementation differs than what may be found elsewhere. Essentially, the average of the
    weights is usually only used at the end of the optimisation, once all the data has been seen.
    However, in this implementation the optimiser returns the current averaged weights.

    Parameters
    ----------
    optimizer
        An optimizer for which the produced weights will be averaged.
    start
        Indicates the number of iterations to wait before starting the average. Essentially,
        nothing happens differently before the number of iterations reaches this value.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.Averager(optim.SGD(0.01), 100)
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 87.89%

    References
    ----------
    [^1]: [Bottou, L., 2010. Large-scale machine learning with stochastic gradient descent. In Proceedings of COMPSTAT'2010 (pp. 177-186). Physica-Verlag HD.](https://leon.bottou.org/publications/pdf/compstat-2010.pdf)
    [^2]: [Stochastic Algorithms for One-Pass Learning slides by Léon Bottou](https://leon.bottou.org/slides/onepass/onepass.pdf)
    [^3]: [Xu, W., 2011. Towards optimal one pass large scale learning with averaged stochastic gradient descent. arXiv preprint arXiv:1107.2490.](https://arxiv.org/pdf/1107.2490.pdf)

    """

    def __init__(self, optimizer: optim.base.Optimizer, start: int = 0):
        self.optimizer = optimizer
        self.start = start
        self.avg_w: dict[str, float] = collections.defaultdict(float)
        self.n_iterations = 0

    def look_ahead(self, w):
        return self.optimizer.look_ahead(w)

    def _step_with_dict(self, w, g):
        w = self.optimizer.step(w, g)

        # No averaging occurs during the first start iterations
        if self.n_iterations < self.start:
            return w

        for i, wi in w.items():
            self.avg_w[i] += (wi - self.avg_w[i]) / (self.n_iterations - self.start + 1)

        return self.avg_w

    @classmethod
    def _unit_test_params(cls):
        yield {"optimizer": optim.SGD()}


# river/optim/base.py

from __future__ import annotations

import abc
import typing

import numpy as np

from river import base, optim, utils

__all__ = ["Initializer", "Scheduler", "Optimizer", "Loss"]

VectorLike = typing.Union[utils.VectorDict, np.ndarray]  # noqa: UP007


class Initializer(base.Base, abc.ABC):
    """An initializer is used to set initial weights in a model."""

    @abc.abstractmethod
    def __call__(self, shape=1):
        """Returns a fresh set of weights.

        Parameters
        ----------
        shape
            Indicates how many weights to return. If `1`, then a single scalar value will be
            returned.

        """


class Scheduler(base.Base, abc.ABC):
    """Can be used to program the learning rate schedule of an `optim.base.Optimizer`."""

    @abc.abstractmethod
    def get(self, t: int) -> float:
        """Returns the learning rate at a given iteration.

        Parameters
        ----------
        t
            The iteration number.

        """

    def __repr__(self):
        return f"{self.__class__.__name__}({vars(self)})"


class Optimizer(base.Base):
    """Optimizer interface.

    Every optimizer inherits from this base interface.

    Parameters
    ----------
    lr

    Attributes
    ----------
    learning_rate : float
        Returns the current learning rate value.

    """

    def __init__(self, lr: int | float | Scheduler):
        if isinstance(lr, int) or isinstance(lr, float):
            lr = optim.schedulers.Constant(lr)
        self.lr = lr
        self.n_iterations = 0

    @property
    def _mutable_attributes(self):
        return {"lr"}

    @property
    def learning_rate(self) -> float:
        return self.lr.get(self.n_iterations)

    def look_ahead(self, w: dict) -> dict:
        """Updates a weight vector before a prediction is made.

        Parameters:
            w (dict): A dictionary of weight parameters. The weights are modified in-place.

        Returns:
            The updated weights.

        """
        return w

    def _step_with_dict(self, w: dict | VectorLike, g: dict | VectorLike) -> dict:  # type: ignore
        raise NotImplementedError

    def _step_with_vector(self, w: VectorLike, g: VectorLike) -> VectorLike:  # type: ignore
        raise NotImplementedError

    def step(self, w: dict | VectorLike, g: dict | VectorLike) -> dict | VectorLike:  # type: ignore
        """Updates a weight vector given a gradient.

        Parameters
        ----------
        w
            A vector-like object containing weights. The weights are modified in-place.
        g
            A vector-like object of gradients.

        Returns
        -------
        The updated weights.

        """

        if (isinstance(w, utils.VectorDict) or isinstance(w, np.ndarray)) and (
            isinstance(g, utils.VectorDict) or isinstance(g, np.ndarray)
        ):
            try:
                w = self._step_with_vector(w, g)
                self.n_iterations += 1
                return w
            except NotImplementedError:
                pass

        w = self._step_with_dict(w, g)
        self.n_iterations += 1
        return w

    def __repr__(self):
        return f"{self.__class__.__name__}({vars(self)})"


class Loss(base.Base, abc.ABC):
    """Base class for all loss functions."""

    def __repr__(self):
        return f"{self.__class__.__name__}({vars(self)})"

    @abc.abstractmethod
    def __call__(self, y_true, y_pred):
        """Returns the loss.

        Parameters
        ----------
        y_true
            Ground truth(s).
        y_pred
            Prediction(s).

        Returns
        -------
        The loss(es).

        """

    @abc.abstractmethod
    def gradient(self, y_true, y_pred):
        """Return the gradient with respect to y_pred.

        Parameters
        ----------
        y_true
            Ground truth(s).
        y_pred
            Prediction(s).

        Returns
        -------
        The gradient(s).

        """

    @abc.abstractmethod
    def mean_func(self, y_pred):
        """Mean function.

        This is the inverse of the link function. Typically, a loss function takes as input the raw
        output of a model. In the case of classification, the raw output would be logits. The mean
        function can be used to convert the raw output into a value that makes sense to the user,
        such as a probability.

        Parameters
        ----------
        y_pred
            Raw prediction(s).

        Returns
        -------
        The adjusted prediction(s).

        References
        ----------
        [^1]: [Wikipedia section on link and mean function](https://www.wikiwand.com/en/Generalized_linear_model#/Link_function)

        """


# river/optim/ftrl.py

from __future__ import annotations

import collections

import numpy as np

from river import optim

__all__ = ["FTRLProximal"]


class FTRLProximal(optim.base.Optimizer):
    """FTRL-Proximal optimizer.

    Parameters
    ----------
    alpha
    beta
    l1
    l2

    Attributes
    ----------
    z : collections.defaultdict
    n : collections.defaultdict

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.FTRLProximal()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 87.66%

    References
    ----------
    [^1]: [McMahan, H.B., Holt, G., Sculley, D., Young, M., Ebner, D., Grady, J., Nie, L., Phillips, T., Davydov, E., Golovin, D. and Chikkerur, S., 2013, August. Ad click prediction: a view from the trenches. In Proceedings of the 19th ACM SIGKDD international conference on Knowledge discovery and data mining (pp. 1222-1230)](https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf)
    [^2]: [Tensorflow's `FtrlOptimizer`](https://www.tensorflow.org/api_docs/python/tf/train/FtrlOptimizer)

    """

    def __init__(self, alpha=0.05, beta=1.0, l1=0.0, l2=1.0):
        self.alpha = alpha
        self.beta = beta
        self.l1 = l1
        self.l2 = l2
        self.z = collections.defaultdict(float)
        self.n = collections.defaultdict(float)
        self.n_iterations = 0

    def _step_with_dict(self, w, g):
        alpha = self.alpha
        beta = self.beta
        l1 = self.l1
        l2 = self.l2
        z = self.z
        n = self.n

        for i in g:
            if abs(z[i]) > l1:
                w[i] = -(((beta + n[i] ** 0.5) / alpha + l2) ** -1) * (z[i] - np.sign(z[i]) * l1)

        for i, gi in g.items():
            s = ((self.n[i] + gi**2) ** 0.5 - self.n[i] ** 0.5) / self.alpha
            self.z[i] += gi - s * w.get(i, 0)
            self.n[i] += gi**2

        return w


# river/optim/initializers.py

"""Weight initializers."""
from __future__ import annotations

import numpy as np

from river.optim.base import Initializer

__all__ = ["Constant", "Normal", "Zeros"]


class Constant(Initializer):
    """Constant initializer which always returns the same value.

    Parameters
    ----------
    value

    Examples
    --------

    >>> from river import optim

    >>> init = optim.initializers.Constant(value=3.14)

    >>> init(shape=1)
    3.14

    >>> init(shape=2)
    array([3.14, 3.14])

    """

    def __init__(self, value: float):
        self.value = value

    def __call__(self, shape=1):
        return np.full(shape, self.value, dtype=float) if shape != 1 else self.value


class Zeros(Constant):
    """Constant initializer which always returns zeros.

    Examples
    --------

    >>> from river import optim

    >>> init = optim.initializers.Zeros()

    >>> init(shape=1)
    0.0

    >>> init(shape=2)
    array([0., 0.])

    """

    def __init__(self):
        super().__init__(value=0.0)


class Normal(Initializer):
    """Random normal initializer which simulate a normal distribution with specified parameters.

    Parameters
    ----------
    mu
        The mean of the normal distribution
    sigma
        The standard deviation of the normal distribution
    seed
        Random number generation seed that can be set for reproducibility.

    Examples
    --------

    >>> from river import optim

    >>> init = optim.initializers.Normal(mu=0, sigma=1, seed=42)

    >>> init(shape=1)
    0.496714

    >>> init(shape=2)
    array([-0.1382643 ,  0.64768854])

    """

    def __init__(self, mu=0.0, sigma=1.0, seed: int | None = None):
        self.mu = mu
        self.sigma = sigma
        self.seed = seed
        self._rng = np.random.RandomState(seed)

    def __call__(self, shape=1):
        weights = self._rng.normal(loc=self.mu, scale=self.sigma, size=shape)
        if shape == 1:
            return weights[0]
        return weights


# river/optim/losses.py

"""Loss functions.

Each loss function is intended to work with both single values as well as numpy vectors.

"""
from __future__ import annotations

import math

import numpy as np
from scipy import special  # type: ignore

from river import base, utils
from river.optim.base import Loss

__all__ = [
    "Absolute",
    "BinaryLoss",
    "BinaryFocalLoss",
    "Cauchy",
    "CrossEntropy",
    "Hinge",
    "Huber",
    "EpsilonInsensitiveHinge",
    "Log",
    "MultiClassLoss",
    "Poisson",
    "Quantile",
    "RegressionLoss",
    "Squared",
]


def clamp_proba(x: float) -> float:
    return utils.math.clamp(x=x, minimum=1e-15, maximum=1 - 1e-15)


class BinaryLoss(Loss):
    """A loss appropriate for binary classification tasks."""

    def mean_func(self, y_pred):
        if isinstance(y_pred, np.ndarray):
            return 1.0 / (1.0 + np.exp(-y_pred))
        return utils.math.sigmoid(y_pred)


class MultiClassLoss(Loss):
    """A loss appropriate for multi-class classification tasks."""

    def mean_func(self, y_pred):
        if isinstance(y_pred, np.ndarray):
            return special.softmax(y_pred)
        return utils.math.softmax(y_pred)


class RegressionLoss(Loss):
    """A loss appropriate for regression tasks."""

    def mean_func(self, y_pred):
        return y_pred


class Absolute(RegressionLoss):
    """Absolute loss, also known as the mean absolute error or L1 loss.

    Mathematically, it is defined as

    $$L = |p_i - y_i|$$

    It's gradient w.r.t. to $p_i$ is

    $$\\frac{\\partial L}{\\partial p_i} = sgn(p_i - y_i)$$

    Examples
    --------

    >>> from river import optim

    >>> loss = optim.losses.Absolute()
    >>> loss(-42, 42)
    84
    >>> loss.gradient(1, 2)
    1
    >>> loss.gradient(2, 1)
    -1

    """

    def __call__(self, y_true, y_pred):
        if isinstance(y_true, np.ndarray):
            return np.abs(y_pred - y_true)
        return abs(y_pred - y_true)

    def gradient(self, y_true, y_pred):
        if isinstance(y_true, np.ndarray):
            return np.where(y_pred > y_true, 1, -1)

        if y_pred > y_true:
            return 1
        return -1


class Cauchy(RegressionLoss):
    """Cauchy loss function.

    Parameters
    ----------
    C

    References
    ----------
    [^1]: ["Effect of MAE" Kaggle discussion](https://www.kaggle.com/c/allstate-claims-severity/discussion/24520#140163)
    [^2]: [Paris Madness Kaggle kernel](https://www.kaggle.com/raddar/paris-madness)

    """

    def __init__(self, C=80):
        self.C = C

    def __call__(self, y_true, y_pred):
        if isinstance(y_true, np.ndarray):
            return np.abs(y_pred - y_true)
        return abs(y_pred - y_true)

    def gradient(self, y_true, y_pred):
        diff = y_pred - y_true
        return diff / ((diff / self.C) ** 2 + 1)


class CrossEntropy(MultiClassLoss):
    """Cross entropy loss.

    This is a generalization of logistic loss to multiple classes.

    Parameters
    ----------
    class_weight
        A dictionary that indicates what weight to associate with each class.

    Examples
    --------

    >>> from river import optim

    >>> y_true = [0, 1, 2, 2]
    >>> y_pred = [
    ...     {0: 0.29450637, 1: 0.34216758, 2: 0.36332605},
    ...     {0: 0.21290077, 1: 0.32728332, 2: 0.45981591},
    ...     {0: 0.42860913, 1: 0.33380113, 2: 0.23758974},
    ...     {0: 0.44941979, 1: 0.32962558, 2: 0.22095463}
    ... ]

    >>> loss = optim.losses.CrossEntropy()

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(loss(yt, yp))
    1.222454
    1.116929
    1.437209
    1.509797

    >>> for yt, yp in zip(y_true, y_pred):
    ...     print(loss.gradient(yt, yp))
    {0: -0.70549363, 1: 0.34216758, 2: 0.36332605}
    {0: 0.21290077, 1: -0.67271668, 2: 0.45981591}
    {0: 0.42860913, 1: 0.33380113, 2: -0.76241026}
    {0: 0.44941979, 1: 0.32962558, 2: -0.77904537}

    References
    ----------
    [^1]: [What is Softmax regression and how is it related to Logistic regression?](https://github.com/rasbt/python-machine-learning-book/blob/master/faq/softmax_regression.md)

    """

    def __init__(self, class_weight: dict[base.typing.ClfTarget, float] | None = None):
        if class_weight is None:
            class_weight = {}
        self.class_weight = class_weight

    def __call__(self, y_true, y_pred):
        total = 0

        for label, proba in y_pred.items():
            if y_true == label:
                total += self.class_weight.get(label, 1.0) * math.log(clamp_proba(proba))

        return -total

    def gradient(self, y_true, y_pred):
        return {
            label: (
                self.class_weight.get(label, 1.0)
                * (clamp_proba(y_pred.get(label, 0.0)) - (y_true == label))
            )
            for label in {*y_pred.keys(), y_true}
        }


class Hinge(BinaryLoss):
    r"""Computes the hinge loss.

    Mathematically, it is defined as

    $$L = max(0, 1 - p_i * y_i)$$

    It's gradient w.r.t. to $p_i$ is

    $$
    \\frac{\\partial L}{\\partial y_i} = \\left\{
    \\begin{array}{ll}
        \\ 0  &   p_iy_i \geqslant 1  \\\\
        \\ - y_i & p_iy_i < 1
    \\end{array}
    \\right.
    $$

    Parameters
    ----------
    threshold
        Margin threshold. 1 yield the loss used in SVMs, whilst 0 is equivalent to the loss used in
        the Perceptron algorithm.

    Examples
    --------

    >>> from river import optim

    >>> loss = optim.losses.Hinge(threshold=1)
    >>> loss(1, .2)
    0.8

    >>> loss.gradient(1, .2)
    -1

    """

    def __init__(self, threshold=1.0):
        self.threshold = threshold

    def __call__(self, y_true, y_pred):
        y_true = y_true * 2 - 1  # [0, 1] -> [-1, 1]

        if isinstance(y_true, np.ndarray):
            return np.maximum(self.threshold - y_true * y_pred, 0)

        return max(self.threshold - y_true * y_pred, 0)

    def gradient(self, y_true, y_pred):
        y_true = y_true * 2 - 1  # [0, 1] -> [-1, 1]

        if isinstance(y_true, np.ndarray):
            return np.where(y_true * y_pred < self.threshold, -y_true, 0)

        if y_true * y_pred <= self.threshold:
            return -y_true
        return 0


class EpsilonInsensitiveHinge(RegressionLoss):
    """Epsilon-insensitive hinge loss.

    Parameters
    ----------
    eps

    """

    def __init__(self, eps=0.1):
        self.eps = eps

    def __call__(self, y_true, y_pred):
        y_true = y_true * 2 - 1  # [0, 1] -> [-1, 1]

        if isinstance(y_true, np.ndarray):
            return np.maximum(np.abs(y_pred - y_true) - self.eps, 0)

        return max(math.fabs(y_pred - y_true) - self.eps, 0)

    def gradient(self, y_true, y_pred):
        y_true = y_true * 2 - 1  # [0, 1] -> [-1, 1]

        if isinstance(y_true, np.ndarray):
            gradients = np.zeros_like(y_true)
            gradients[y_pred > y_true + self.eps] = 1
            gradients[y_pred + self.eps < y_true] = -1
            return gradients

        if y_pred > y_true + self.eps:
            return 1
        elif y_pred + self.eps < y_true:
            return -1
        return 0


class Log(BinaryLoss):
    """Logarithmic loss.

    This loss function expects each provided `y_pred` to be a logit. In other words if must be
    the raw output of a linear model or a neural network.

    Parameters
    ----------
    weight_pos
    weight_neg

    References
    ----------
    [^1]: [Logit Wikipedia page](https://www.wikiwand.com/en/Logit>)

    """

    def __init__(self, weight_pos=1.0, weight_neg=1.0):
        self.weight_pos = weight_pos
        self.weight_neg = weight_neg

    def __call__(self, y_true, y_pred):
        if isinstance(y_true, np.ndarray):
            weights = np.where(y_true == 0, self.weight_neg, self.weight_pos)
            y_true = 2 * y_true - 1  # map {0, 1} to {-1, 1}
            z = y_pred * y_true
            return weights * np.log(1.0 + np.exp(-z))

        weight = self.weight_pos
        if y_true == 0:
            y_true = -1
            weight = self.weight_neg
        else:
            y_true = int(y_true)

        z = y_pred * y_true
        if z > 18.0:
            return weight * math.exp(-z)
        if z < -18.0:
            return weight * -z
        return weight * math.log(1.0 + math.exp(-z))

    def gradient(self, y_true, y_pred):
        if isinstance(y_true, np.ndarray):
            weights = np.where(y_true == 0, self.weight_neg, self.weight_pos)
            y_true = 2 * y_true - 1  # map {0, 1} to {-1, 1}
            z = y_pred * y_true
            return weights * -y_true / (np.exp(z) + 1.0)

        weight = self.weight_pos
        if y_true == 0:
            y_true = -1
            weight = self.weight_neg
        else:
            y_true = int(y_true)

        z = y_pred * y_true
        if z > 18.0:
            return weight * math.exp(-z) * -y_true
        if z < -18.0:
            return weight * -y_true
        return weight * -y_true / (math.exp(z) + 1.0)


class Quantile(RegressionLoss):
    """Quantile loss.

    Parameters
    ----------
    alpha
        Desired quantile to attain.

    Examples
    --------

    >>> from river import optim

    >>> loss = optim.losses.Quantile(0.5)
    >>> loss(1, 3)
    1.0

    >>> loss.gradient(1, 3)
    0.5

    >>> loss.gradient(3, 1)
    -0.5

    References
    ----------
    [^1]: [Wikipedia article on quantile regression](https://www.wikiwand.com/en/Quantile_regression)
    [^2]: [Derivative from WolframAlpha](https://www.wolframalpha.com/input/?i=derivative+(y+-+p)+*+(alpha+-+Boole(y+-+p))+wrt+p)

    """

    def __init__(self, alpha=0.5):
        self.alpha = alpha

    def __call__(self, y_true, y_pred):
        diff = y_pred - y_true
        return (self.alpha - (diff < 0)) * diff

    def gradient(self, y_true, y_pred):
        return (y_true < y_pred) - self.alpha


class Squared(RegressionLoss):
    """Squared loss, also known as the L2 loss.

    Mathematically, it is defined as

    $$L = (p_i - y_i) ^ 2$$

    It's gradient w.r.t. to $p_i$ is

    $$\\frac{\\partial L}{\\partial p_i} = 2 (p_i - y_i)$$

    One thing to note is that this convention is consistent with Vowpal Wabbit and PyTorch, but
    not with scikit-learn. Indeed, scikit-learn divides the loss by 2, making the 2 disappear in
    the gradient.

    Examples
    --------

    >>> from river import optim

    >>> loss = optim.losses.Squared()
    >>> loss(-4, 5)
    81
    >>> loss.gradient(-4, 5)
    18
    >>> loss.gradient(5, -4)
    -18

    """

    def __call__(self, y_true, y_pred):
        return (y_pred - y_true) * (y_pred - y_true)

    def gradient(self, y_true, y_pred):
        return 2 * (y_pred - y_true)


class Huber(RegressionLoss):
    """Huber loss.

    Variant of the squared loss that is robust to outliers.

    Parameters
    ----------
    epsilon

    References
    ----------
    1. [Huber loss function - Wikipedia](https://en.wikipedia.org/wiki/Huber_Loss_Function)

    """

    def __init__(self, epsilon=0.1):
        self.epsilon = epsilon

    def __call__(self, y_true, y_pred):
        r = y_pred - y_true

        if isinstance(y_true, np.ndarray):
            abs_r = np.abs(r)
            return np.where(
                abs_r <= self.epsilon,
                0.5 * r * r,
                self.epsilon * abs_r - (0.5 * self.epsilon * self.epsilon),
            )

        abs_r = abs(r)
        if abs_r <= self.epsilon:
            return 0.5 * r * r
        return self.epsilon * abs_r - (0.5 * self.epsilon * self.epsilon)

    def gradient(self, y_true, y_pred):
        r = y_pred - y_true

        if isinstance(y_true, np.ndarray):
            abs_r = np.abs(r)
            return np.where(
                abs_r <= self.epsilon, r, np.where(r > 0.0, self.epsilon, -self.epsilon)
            )

        abs_r = abs(r)
        if abs_r <= self.epsilon:
            return r
        elif r > 0.0:
            return self.epsilon
        return -self.epsilon


class BinaryFocalLoss(BinaryLoss):
    """Binary focal loss.

    This implements the "star" algorithm from the appendix of the focal loss paper.

    Parameters
    ----------
    gamma
    beta

    References
    ----------
    1. [Lin, T.Y., Goyal, P., Girshick, R., He, K. and Dollár, P., 2017. Focal loss for dense object detection. In Proceedings of the IEEE international conference on computer vision (pp. 2980-2988)](https://arxiv.org/pdf/1708.02002.pdf)

    """

    def __init__(self, gamma=2, beta=1):
        self.gamma = gamma
        self.beta = beta

    def __call__(self, y_true, y_pred):
        y_true = y_true * 2 - 1  # [0, 1] -> [-1, 1]

        xt = y_true * y_pred

        if isinstance(y_true, np.ndarray):
            pt = 1.0 / (1 + np.exp(-(self.gamma * xt + self.beta)))
            return -np.log(pt) / self.gamma

        pt = utils.math.sigmoid(self.gamma * xt + self.beta)
        return -math.log(pt) / self.gamma

    def gradient(self, y_true, y_pred):
        y_true = y_true * 2 - 1  # [0, 1] -> [-1, 1]

        xt = y_true * y_pred

        if isinstance(y_true, np.ndarray):
            pt = 1.0 / (1 + np.exp(-(self.gamma * xt + self.beta)))
            return y_true * (pt - 1)

        pt = utils.math.sigmoid(self.gamma * xt + self.beta)
        return y_true * (pt - 1)


class Poisson(RegressionLoss):
    """Poisson loss.

    The Poisson loss is usually more suited for regression with count data than the squared loss.

    Mathematically, it is defined as

    $$L = exp(p_i) - y_i \\times p_i$$

    It's gradient w.r.t. to $p_i$ is

    $$\\frac{\\partial L}{\\partial p_i} = exp(p_i) - y_i$$

    """

    def __call__(self, y_true, y_pred):
        if isinstance(y_pred, np.ndarray):
            return np.exp(y_pred) - y_true * y_pred
        return math.exp(y_pred) - y_true * y_pred

    def gradient(self, y_true, y_pred):
        if isinstance(y_pred, np.ndarray):
            return np.exp(y_pred) - y_true
        return math.exp(y_pred) - y_true

    def mean_func(self, y_pred):
        if isinstance(y_pred, np.ndarray):
            return np.exp(y_pred)
        return math.exp(y_pred)


# river/optim/momentum.py

from __future__ import annotations

import collections

from river import optim

__all__ = ["Momentum"]


class Momentum(optim.base.Optimizer):
    """Momentum optimizer.

    Parameters
    ----------
    lr
    rho

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.Momentum()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 84.16%

    """

    def __init__(self, lr=0.1, rho=0.9):
        super().__init__(lr)
        self.rho = rho
        self.s = collections.defaultdict(float)

    def _step_with_dict(self, w, g):
        for i, gi in g.items():
            self.s[i] = self.rho * self.s[i] + self.learning_rate * gi
            w[i] -= self.s[i]

        return w


# river/optim/nadam.py

from __future__ import annotations

import collections
import math

from river import optim

__all__ = ["Nadam"]


class Nadam(optim.base.Optimizer):
    """Nadam optimizer.

    Parameters
    ----------
    lr
    beta_1
    beta_2
    eps

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.Nadam()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 86.60%

    References
    ----------
    [^1]: [Nadam: A combination of adam and nesterov](https://ruder.io/optimizing-gradient-descent/index.html#nadam)

    """

    def __init__(self, lr=0.1, beta_1=0.9, beta_2=0.999, eps=1e-8):
        super().__init__(lr)
        self.beta_1 = beta_1
        self.beta_2 = beta_2
        self.eps = eps
        self.m = collections.defaultdict(float)
        self.v = collections.defaultdict(float)

    def _step_with_dict(self, w, g):
        for i, gi in g.items():
            self.m[i] = self.beta_1 * self.m[i] + (1 - self.beta_1) * gi
            m_hat = self.m[i] / (1 - math.pow(self.beta_1, self.n_iterations + 1))

            self.v[i] = self.beta_2 * self.v[i] + (1 - self.beta_2) * gi**2
            v_hat = self.v[i] / (1 - math.pow(self.beta_2, self.n_iterations + 1))

            w[i] -= (
                self.learning_rate
                * (
                    self.beta_1 * m_hat
                    + (1 - self.beta_1) * gi / (1 - math.pow(self.beta_1, self.n_iterations + 1))
                )
                / (v_hat**0.5 + self.eps)
            )

        return w


# river/optim/nesterov.py

from __future__ import annotations

import collections

from river import optim

__all__ = ["NesterovMomentum"]


class NesterovMomentum(optim.base.Optimizer):
    """Nesterov Momentum optimizer.

    Parameters
    ----------
    lr
    rho

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.NesterovMomentum()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 84.29%

    """

    def __init__(self, lr=0.1, rho=0.9):
        super().__init__(lr)
        self.rho = rho
        self.s = collections.defaultdict(float)

    def look_ahead(self, w):
        for i in w:
            w[i] -= self.rho * self.s[i]

        return w

    def _step_with_dict(self, w, g):
        # Move w back to it's initial position
        for i in w:
            w[i] += self.rho * self.s[i]

        for i, gi in g.items():
            self.s[i] = self.rho * self.s[i] + self.learning_rate * gi
            w[i] -= self.s[i]

        return w


# river/optim/newton.py

from __future__ import annotations

from river import optim, utils

__all__ = ["Newton"]


def sherman_morrison(A_inv: dict, u: dict, v: dict) -> dict:
    """Sherman–Morrison formula.

    This modifies `A_inv` inplace.

    Parameters
    ----------
    A_inv
    u
    v

    Examples
    --------

    >>> import pprint

    >>> A_inv = {
    ...     (0, 0): 0.2,
    ...     (1, 1): 1,
    ...     (2, 2): 1
    ... }
    >>> u = {0: 1, 1: 2, 2: 3}
    >>> v = {0: 4}

    >>> inv = sherman_morrison(A_inv, u, v)
    >>> pprint.pprint(inv)
    {(0, 0): 0.111111,
        (1, 0): -0.888888,
        (1, 1): 1,
        (2, 0): -1.333333,
        (2, 2): 1}

    References
    ----------
    [^1]: [Wikipedia article on the Sherman-Morrison formula](https://www.wikiwand.com/en/Sherman%E2%80%93Morrison_formula)s

    """

    den = 1 + utils.math.dot(utils.math.dotvecmat(u, A_inv), v)

    for k, v in utils.math.matmul2d(
        utils.math.matmul2d(A_inv, utils.math.outer(u, v)), A_inv
    ).items():
        A_inv[k] = A_inv.get(k, 0) - v / den

    return A_inv


class Newton(optim.base.Optimizer):
    """Online Newton Step (ONS) optimizer.

    This optimizer uses second-order information (i.e. the Hessian of the cost function) in
    addition to first-order information (i.e. the gradient of the cost function).

    Parameters
    ----------
    lr
    eps

    References
    ----------
    [^1]: [Hazan, E., Agarwal, A. and Kale, S., 2007. Logarithmic regret algorithms for online convex optimization. Machine Learning, 69(2-3), pp.169-192](https://www.cs.princeton.edu/~ehazan/papers/log-journal.pdf)

    """

    def __init__(self, lr=0.1, eps=1e-5):
        super().__init__(lr)
        self.eps = eps
        self.H_inv = {}

    def _step_with_dict(self, w, g):
        for i in g:
            if (i, i) not in self.H_inv:
                self.H_inv[i, i] = self.eps

        # Update the Hessian
        self.H = sherman_morrison(A_inv=self.H_inv, u=g, v=g)

        # Calculate the update step
        step = utils.math.dotvecmat(x=g, A=self.H_inv)

        # Update the weights
        for i, s in step.items():
            w[i] -= self.learning_rate * s

        return w


# river/optim/rms_prop.py

from __future__ import annotations

import collections

import numpy as np

from river import optim, utils

__all__ = ["RMSProp"]


class RMSProp(optim.base.Optimizer):
    """RMSProp optimizer.

    Parameters
    ----------
    lr
    rho
    eps

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.RMSProp()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 87.24%

    References
    ----------
    [^1]: [Divide the gradient by a running average of itsrecent magnitude](https://www.cs.toronto.edu/~tijmen/csc321/slides/lecture_slides_lec6.pdf)

    """

    def __init__(self, lr=0.1, rho=0.9, eps=1e-8):
        super().__init__(lr)
        self.rho = rho
        self.eps = eps
        self.g2 = None

    def _step_with_dict(self, w, g):
        if self.g2 is None:
            self.g2 = collections.defaultdict(float)

        for i, gi in g.items():
            self.g2[i] = self.rho * self.g2[i] + (1 - self.rho) * gi**2
            w[i] -= self.learning_rate / (self.g2[i] + self.eps) ** 0.5 * gi

        return w

    def _step_with_vector(self, w, g):
        if self.g2 is None:
            if isinstance(w, np.ndarray):
                self.g2 = np.zeros_like(w)
            else:
                self.g2 = utils.VectorDict()

        self.g2 = self.rho * self.g2 + (1 - self.rho) * g**2
        w -= self.learning_rate / (self.g2 + self.eps) ** 0.5 * g

        return w


# river/optim/schedulers.py

"""Learning rate schedulers."""
from __future__ import annotations

import math

from river import optim

__all__ = ["Constant", "InverseScaling", "Optimal"]


class Constant(optim.base.Scheduler):
    """Always uses the same learning rate.

    Parameters
    ----------
    learning_rate

    """

    def __init__(self, learning_rate: int | float):
        self.learning_rate = learning_rate

    @property
    def _mutable_attributes(self):
        return {"learning_rate"}

    def get(self, t):
        return self.learning_rate


class InverseScaling(optim.base.Scheduler):
    r"""Reduces the learning rate using a power schedule.

    Assuming an initial learning rate $\eta$, the learning rate at step $t$ is:

    $$\\frac{eta}{(t + 1) ^ p}$$

    where $p$ is a user-defined parameter.

    Parameters
    ----------
    learning_rate
    power

    """

    def __init__(self, learning_rate: float, power=0.5):
        self.learning_rate = learning_rate
        self.power = power

    def get(self, t):
        return self.learning_rate / pow(t + 1, self.power)


class Optimal(optim.base.Scheduler):
    """Optimal learning schedule as proposed by Léon Bottou.

    Parameters
    ----------
    loss
    alpha

    References
    ----------
    [^1]: [Bottou, L., 2012. Stochastic gradient descent tricks. In Neural networks: Tricks of the trade (pp. 421-436). Springer, Berlin, Heidelberg.](https://cilvr.cs.nyu.edu/diglib/lsml/bottou-sgd-tricks-2012.pdf)

    """

    def __init__(self, loss: optim.losses.Loss, alpha=1e-4):
        self.loss = loss
        self.alpha = alpha

        typw = math.sqrt(1.0 / math.sqrt(self.alpha))
        initial_eta0 = typw / max(1.0, self.loss.gradient(True, -typw))
        self.t0 = 1.0 / (initial_eta0 * self.alpha)

    def get(self, t):
        return 1.0 / (self.alpha * (self.t0 + t))


# river/optim/sgd.py

from __future__ import annotations

from river import optim

__all__ = ["SGD"]


class SGD(optim.base.Optimizer):
    """Plain stochastic gradient descent.

    Parameters
    ----------
    lr

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import optim
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()
    >>> optimizer = optim.SGD(0.1)
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LogisticRegression(optimizer)
    ... )
    >>> metric = metrics.F1()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    F1: 87.85%

    References
    ----------
    [^1]: [Robbins, H. and Monro, S., 1951. A stochastic approximation method. The annals of mathematical statistics, pp.400-407](https://pdfs.semanticscholar.org/34dd/d8865569c2c32dec9bf7ffc817ff42faaa01.pdf)

    """

    def __init__(self, lr=0.01):
        super().__init__(lr)

    def _step_with_dict(self, w, g):
        for i, gi in g.items():
            w[i] -= self.learning_rate * gi
        return w

    def _step_with_vector(self, w, g):
        w -= self.learning_rate * g
        return w


# river/optim/test_.py

from __future__ import annotations

import importlib
import inspect
import math
import random
import typing

import numpy as np
import pytest

from river import optim, utils


def losses() -> typing.Iterable[optim.losses.Loss]:
    for _, loss in inspect.getmembers(
        importlib.import_module("river.optim.losses"),
        lambda x: inspect.isclass(x)
        and not inspect.isabstract(x)
        and not issubclass(x, optim.losses.CrossEntropy),
    ):
        yield loss()


@pytest.mark.parametrize(
    "loss",
    [pytest.param(loss, id=loss.__class__.__name__) for loss in losses()],
)
def test_loss_batch_online_equivalence(loss):
    y_true = np.random.randint(2, size=30)
    y_pred = np.random.uniform(-10, 10, size=30)

    for yt, yp, g in zip(y_true, y_pred, loss.gradient(y_true, y_pred)):
        assert math.isclose(loss.gradient(yt, yp), g, abs_tol=1e-9)


def optimizers() -> typing.Iterable[optim.base.Optimizer]:
    for _, optimizer in inspect.getmembers(
        importlib.import_module("river.optim"),
        lambda x: inspect.isclass(x)
        and issubclass(x, optim.base.Optimizer)
        and x is not optim.base.Optimizer,
    ):
        for params in optimizer._unit_test_params():
            yield optimizer(**params)


@pytest.mark.parametrize(
    "optimizer",
    [pytest.param(optimizer, id=optimizer.__class__.__name__) for optimizer in optimizers()],
)
def test_optimizer_step_with_dict_same_as_step_with_vector_dict(optimizer):
    w_dict = {i: random.uniform(-5, 5) for i in range(10)}
    w_vector = utils.VectorDict(w_dict)

    g_dict = {i: random.uniform(-5, 5) for i in range(10)}
    g_vector = utils.VectorDict(g_dict)

    w_dict = optimizer._step_with_dict(w_dict, g_dict)
    try:
        w_vector = optimizer.clone()._step_with_vector(w_vector, g_vector)
    except NotImplementedError:
        pytest.skip("step_with_vector not implemented")

    for i, w in w_vector.to_dict().items():
        assert math.isclose(w, w_dict[i])


@pytest.mark.parametrize(
    "optimizer",
    [pytest.param(optimizer, id=optimizer.__class__.__name__) for optimizer in optimizers()],
)
def test_optimizer_step_with_dict_same_as_step_with_numpy_array(optimizer):
    w_dict = {i: random.uniform(-5, 5) for i in range(10)}
    w_vector = np.asarray(list(w_dict.values()))

    g_dict = {i: random.uniform(-5, 5) for i in range(10)}
    g_vector = np.asarray(list(g_dict.values()))

    w_dict = optimizer._step_with_dict(w_dict, g_dict)
    try:
        w_vector = optimizer.clone()._step_with_vector(w_vector, g_vector)
    except NotImplementedError:
        pytest.skip("step_with_vector not implemented")

    for i, w in dict(enumerate(w_vector)).items():
        assert math.isclose(w, w_dict[i])


# river/preprocessing/__init__.py

"""Feature preprocessing.

The purpose of this module is to modify an existing set of features so that they can be processed
by a machine learning algorithm. This may be done by scaling numeric parts of the data or by
one-hot encoding categorical features. The difference with the `feature_extraction` module is that
the latter extracts new information from the data

"""
from __future__ import annotations

from .feature_hasher import FeatureHasher
from .impute import PreviousImputer, StatImputer
from .lda import LDA
from .one_hot import OneHotEncoder
from .pred_clipper import PredClipper
from .random_projection import GaussianRandomProjector, SparseRandomProjector
from .scale import (
    AdaptiveStandardScaler,
    Binarizer,
    MaxAbsScaler,
    MinMaxScaler,
    Normalizer,
    RobustScaler,
    StandardScaler,
)
from .scale_target import TargetMinMaxScaler, TargetStandardScaler

__all__ = [
    "AdaptiveStandardScaler",
    "Binarizer",
    "FeatureHasher",
    "GaussianRandomProjector",
    "LDA",
    "MaxAbsScaler",
    "MinMaxScaler",
    "Normalizer",
    "OneHotEncoder",
    "PredClipper",
    "PreviousImputer",
    "RobustScaler",
    "SparseRandomProjector",
    "StandardScaler",
    "StatImputer",
    "TargetMinMaxScaler",
    "TargetStandardScaler",
]


# river/preprocessing/feature_hasher.py

from __future__ import annotations

import collections
import hashlib

import numpy as np

from river import base

__all__ = ["FeatureHasher"]


class FeatureHasher(base.Transformer):
    """Implements the hashing trick.

    Each pair of (name, value) features is hashed into a random integer. A module operator is then
    used to make sure the hash is in a certain range. We use the Murmurhash implementation from
    scikit-learn.

    Parameters
    ----------
    n_features
        The number by which each hash will be moduloed by.
    alternate_sign
        Whether or not half of the hashes will be negated.
    seed
        Set the seed to produce identical results.

    Examples
    --------

    >>> import river

    >>> hasher = river.preprocessing.FeatureHasher(n_features=10, seed=42)

    >>> X = [
    ...     {'dog': 1, 'cat': 2, 'elephant': 4},
    ...     {'dog': 2, 'run': 5}
    ... ]
    >>> for x in X:
    ...     print(hasher.transform_one(x))
    Counter({1: 4, 9: 2, 8: 1})
    Counter({4: 5, 8: 2})

    References
    ----------
    [^1]: [Wikipedia article on feature vectorization using the hashing trick](https://www.wikiwand.com/en/Feature_hashing#/Feature_vectorization_using_hashing_trick)

    """

    def __init__(self, n_features=1048576, seed: int | None = None):
        self.n_features = n_features
        self.seed = seed
        self._salt = np.random.RandomState(seed).bytes(hashlib.blake2s.SALT_SIZE)

    def _hash(self, x):
        hexa = hashlib.blake2s(bytes(x, encoding="utf8"), salt=self._salt).hexdigest()
        return int(hexa, 16)

    def transform_one(self, x):
        x_hashed = collections.Counter()

        for feature, value in x.items():
            if isinstance(value, str):
                feature = f"{feature}={value}"
                value = 1

            i = self._hash(feature) % self.n_features
            x_hashed[i] += value

        return x_hashed


# river/preprocessing/impute.py

from __future__ import annotations

import typing

from river import base, stats

__all__ = ["PreviousImputer", "StatImputer"]


class PreviousImputer(base.Transformer):
    """Imputes missing values by using the most recent value.

    Examples
    --------

    >>> from river import preprocessing

    >>> imputer = preprocessing.PreviousImputer()

    >>> imputer = imputer.learn_one({'x': 1, 'y': 2})
    >>> imputer.transform_one({'y': None})
    {'y': 2}

    >>> imputer.transform_one({'x': None})
    {'x': 1}

    """

    def __init__(self):
        self._latest = {}

    def learn_one(self, x):
        for i, v in x.items():
            if v is not None:
                self._latest[i] = v

        return self

    def transform_one(self, x):
        for i, v in x.items():
            if v is None:
                x[i] = self._latest.get(i)

        return x


class StatImputer(base.Transformer):
    """Replaces missing values with a statistic.

    This transformer allows you to replace missing values with the value of a running statistic.
    During a call to `learn_one`, for each feature, a statistic is updated whenever a numeric feature
    is observed. When `transform_one` is called, each feature with a `None` value is replaced with
    the current value of the corresponding statistic.

    Parameters
    ----------
    imputers
        A list of tuples where each tuple has two elements. The first elements is a
        feature name and the second value is an instance of `stats.base.Univariate`. The second
        value can also be an arbitrary value, such as -1, in which case the missing values will
        be replaced with it.

    Examples
    --------

    >>> from river import preprocessing
    >>> from river import stats

    For numeric data, we can use a `stats.Mean()` to replace missing values by the running
    average of the previously seen values:

    >>> X = [
    ...     {'temperature': 1},
    ...     {'temperature': 8},
    ...     {'temperature': 3},
    ...     {'temperature': None},
    ...     {'temperature': 4}
    ... ]

    >>> imp = preprocessing.StatImputer(('temperature', stats.Mean()))

    >>> for x in X:
    ...     imp = imp.learn_one(x)
    ...     print(imp.transform_one(x))
    {'temperature': 1}
    {'temperature': 8}
    {'temperature': 3}
    {'temperature': 4.0}
    {'temperature': 4}

    For discrete/categorical data, a common practice is to `stats.Mode` to replace missing
    values by the most commonly seen value:

    >>> X = [
    ...     {'weather': 'sunny'},
    ...     {'weather': 'rainy'},
    ...     {'weather': 'sunny'},
    ...     {'weather': None},
    ...     {'weather': 'rainy'},
    ...     {'weather': 'rainy'},
    ...     {'weather': None}
    ... ]

    >>> imp = preprocessing.StatImputer(('weather', stats.Mode()))

    >>> for x in X:
    ...     imp = imp.learn_one(x)
    ...     print(imp.transform_one(x))
    {'weather': 'sunny'}
    {'weather': 'rainy'}
    {'weather': 'sunny'}
    {'weather': 'sunny'}
    {'weather': 'rainy'}
    {'weather': 'rainy'}
    {'weather': 'rainy'}

    You can also choose to replace missing values with a constant value, as so:

    >>> imp = preprocessing.StatImputer(('weather', 'missing'))

    >>> for x in X:
    ...     imp = imp.learn_one(x)
    ...     print(imp.transform_one(x))
    {'weather': 'sunny'}
    {'weather': 'rainy'}
    {'weather': 'sunny'}
    {'weather': 'missing'}
    {'weather': 'rainy'}
    {'weather': 'rainy'}
    {'weather': 'missing'}

    Multiple imputers can be defined by providing a tuple for each feature which you want to
    impute:

    >>> X = [
    ...     {'weather': 'sunny', 'temperature': 8},
    ...     {'weather': 'rainy', 'temperature': 3},
    ...     {'weather': 'sunny', 'temperature': None},
    ...     {'weather': None, 'temperature': 4},
    ...     {'weather': 'snowy', 'temperature': -4},
    ...     {'weather': 'snowy', 'temperature': -3},
    ...     {'weather': 'snowy', 'temperature': -3},
    ...     {'weather': None, 'temperature': None}
    ... ]

    >>> imp = preprocessing.StatImputer(
    ...     ('temperature', stats.Mean()),
    ...     ('weather', stats.Mode())
    ... )

    >>> for x in X:
    ...     imp = imp.learn_one(x)
    ...     print(imp.transform_one(x))
    {'weather': 'sunny', 'temperature': 8}
    {'weather': 'rainy', 'temperature': 3}
    {'weather': 'sunny', 'temperature': 5.5}
    {'weather': 'sunny', 'temperature': 4}
    {'weather': 'snowy', 'temperature': -4}
    {'weather': 'snowy', 'temperature': -3}
    {'weather': 'snowy', 'temperature': -3}
    {'weather': 'snowy', 'temperature': 0.8333}

    A sophisticated way to go about imputation is condition the statistics on a given feature.
    For instance, we might want to replace a missing temperature with the average temperature
    of a particular weather condition. As an example, consider the following dataset where the
    temperature is missing, but not the weather condition:

    >>> X = [
    ...     {'weather': 'sunny', 'temperature': 8},
    ...     {'weather': 'rainy', 'temperature': 3},
    ...     {'weather': 'sunny', 'temperature': None},
    ...     {'weather': 'rainy', 'temperature': 4},
    ...     {'weather': 'sunny', 'temperature': 10},
    ...     {'weather': 'sunny', 'temperature': None},
    ...     {'weather': 'sunny', 'temperature': 12},
    ...     {'weather': 'rainy', 'temperature': None}
    ... ]

    Each missing temperature can be replaced with the average temperature of the corresponding
    weather condition as so:

    >>> from river import compose

    >>> imp = compose.Grouper(
    ...     preprocessing.StatImputer(('temperature', stats.Mean())),
    ...     by='weather'
    ... )

    >>> for x in X:
    ...     imp = imp.learn_one(x)
    ...     print(imp.transform_one(x))
    {'weather': 'sunny', 'temperature': 8}
    {'weather': 'rainy', 'temperature': 3}
    {'weather': 'sunny', 'temperature': 8.0}
    {'weather': 'rainy', 'temperature': 4}
    {'weather': 'sunny', 'temperature': 10}
    {'weather': 'sunny', 'temperature': 9.0}
    {'weather': 'sunny', 'temperature': 12}
    {'weather': 'rainy', 'temperature': 3.5}

    Note that you can also create a `Grouper` with the `*` operator:

    >>> imp = preprocessing.StatImputer(('temperature', stats.Mean())) * 'weather'

    """

    def __init__(self, *imputers):
        self.imputers = imputers
        self.stats = {
            feature: stat if isinstance(stat, stats.base.Univariate) else Constant(stat)
            for feature, stat in imputers
        }

    def learn_one(self, x):
        for i in self.stats:
            if x[i] is not None:
                self.stats[i].update(x[i])

        return self

    def transform_one(self, x):
        # Transformers are supposed to be pure, therefore we make a copy of the features
        x = x.copy()

        for i in self.stats:
            if x[i] is None:
                x[i] = self.stats[i].get()

        return x


class Constant(stats.base.Univariate):
    """Implements the `stats.base.Univariate` interface but always returns the same value.

    Parameters
    ----------
    value

    """

    def __init__(self, value: typing.Any):
        self.value = value

    def update(self, x):
        return self

    def get(self):
        return self.value

    @property
    def name(self):
        return self.value


# river/preprocessing/lda.py

from __future__ import annotations

import functools
import typing
from collections import defaultdict

import numpy as np
from scipy import ndimage, special

from river import base

__all__ = ["LDA"]


class LDA(base.Transformer):
    """Online Latent Dirichlet Allocation with Infinite Vocabulary.

    Latent Dirichlet allocation (LDA) is a probabilistic approach for exploring topics in document
    collections. The key advantage of this variant is that it assumes an infinite vocabulary,
    meaning that the set of tokens does not have to known in advance, as opposed to the
    [implementation from sklearn](https://scikit-learn.org/stable/modules/generated/sklearn.decomposition.LatentDirichletAllocation.html)
    The results produced by this implementation are identical to those from [the original
    implementation](https://github.com/kzhai/PyInfVoc) proposed by the method's authors.

    This class takes as input token counts. Therefore, it requires you to tokenize beforehand. You
    can do so by using a `feature_extraction.BagOfWords` instance, as shown in the example below.

    Parameters
    ----------
    n_components
        Number of topics of the latent Drichlet allocation.
    number_of_documents
        Estimated number of documents.
    alpha_theta
        Hyper-parameter of the Dirichlet distribution of topics.
    alpha_beta
        Hyper-parameter of the Dirichlet process of distribution over words.
    tau
        Learning inertia to prevent premature convergence.
    kappa
        The learning rate kappa controls how quickly new parameters estimates replace the old ones.
        kappa ∈ (0.5, 1] is required for convergence.
    vocab_prune_interval
        Interval at which to refresh the words topics distribution.
    number_of_samples
        Number of iteration to computes documents topics distribution.
    ranking_smooth_factor
    burn_in_sweeps
        Number of iteration necessaries while analyzing a document before updating document topics
        distribution.
    maximum_size_vocabulary
        Maximum size of the stored vocabulary.
    seed
        Random number seed used for reproducibility.

    Attributes
    ----------
    counter : int
        The current number of observed documents.
    truncation_size_prime : int
        Number of distincts words stored in the vocabulary. Updated before processing a document.
    truncation_size : int
        Number of distincts words stored in the vocabulary. Updated after processing a document.
    word_to_index : dict
        Words as keys and indexes as values.
    index_to_word : dict
        Indexes as keys and words as values.
    nu_1 : dict
        Weights of the words. Component of the variational inference.
    nu_2 : dict
        Weights of the words. Component of the variational inference.

    Examples
    --------

    >>> from river import compose
    >>> from river import feature_extraction
    >>> from river import preprocessing

    >>> X = [
    ...    'weather cold',
    ...    'weather hot dry',
    ...    'weather cold rainy',
    ...    'weather hot',
    ...    'weather cold humid',
    ... ]

    >>> lda = compose.Pipeline(
    ...     feature_extraction.BagOfWords(),
    ...     preprocessing.LDA(
    ...         n_components=2,
    ...         number_of_documents=60,
    ...         seed=42
    ...     )
    ... )

    >>> for x in X:
    ...     lda = lda.learn_one(x)
    ...     topics = lda.transform_one(x)
    ...     print(topics)
    {0: 0.5, 1: 2.5}
    {0: 1.5, 1: 2.5}
    {0: 3.5, 1: 0.5}
    {0: 1.5, 1: 1.5}
    {0: 2.5, 1: 1.5}

    References
    ----------
    [^1]: [Zhai, K. and Boyd-Graber, J., 2013, February. Online latent Dirichlet allocation with infinite vocabulary. In International Conference on Machine Learning (pp. 561-569).](http://proceedings.mlr.press/v28/zhai13.pdf)
    [^2]: [PyInfVoc on GitHub](https://github.com/kzhai/PyInfVoc)

    """

    def __init__(
        self,
        n_components=10,
        number_of_documents=1e6,
        alpha_theta=0.5,
        alpha_beta=100.0,
        tau=64.0,
        kappa=0.75,
        vocab_prune_interval=10,
        number_of_samples=10,
        ranking_smooth_factor=1e-12,
        burn_in_sweeps=5,
        maximum_size_vocabulary=4000,
        seed: int | None = None,
    ):
        self.n_components = n_components
        self.number_of_documents = number_of_documents
        self.alpha_theta = alpha_theta
        self.alpha_beta = alpha_beta
        self.tau = tau
        self.kappa = kappa
        self.vocab_prune_interval = vocab_prune_interval
        self.number_of_samples = number_of_samples
        self.ranking_smooth_factor = ranking_smooth_factor
        self.burn_in_sweeps = burn_in_sweeps
        self.maximum_size_vocabulary = maximum_size_vocabulary
        self.seed = seed
        self.rng = np.random.RandomState(seed)

        self.counter = 0
        self.truncation_size_prime = 1
        self.truncation_size = 1

        self.word_to_index: dict[str, int] = {}
        self.index_to_word: dict[int, str] = {}

        self.nu_1: defaultdict = defaultdict(functools.partial(np.ones, 1))
        self.nu_2: defaultdict = defaultdict(functools.partial(np.array, [self.alpha_beta]))

        for topic in range(self.n_components):
            self.nu_1[topic] = np.ones(1)
            self.nu_2[topic] = np.array([self.alpha_beta])

    def learn_transform_one(self, x: dict) -> dict:
        """Equivalent to `lda.learn_one(x).transform_one(x)`s, but faster.

        Parameters
        ----------
        x: A document.

        Returns
        -------
        Component attributions for the input document.

        """

        # Updates number of documents:
        self.counter += 1

        # Extracts words of the document as a list of words:
        word_list: typing.Iterable[str] = x.keys()

        # Update words indexes:
        self._update_indexes(word_list=word_list)

        # Replace the words by their index:
        words_indexes_list = [self.word_to_index[word] for word in word_list]

        # Sample empirical topic assignment:
        (
            statistics,
            batch_document_topic_distribution,
        ) = self._compute_statistics_components(words_indexes_list)

        # Online variational inference
        self._update_weights(statistics=statistics)

        if self.counter % self.vocab_prune_interval == 0:
            self._prune_vocabulary()

        return dict(enumerate(batch_document_topic_distribution))

    def learn_one(self, x):
        self.learn_transform_one(x)
        return self

    def transform_one(self, x):
        # Extracts words of the document as a list of words:
        word_list = x.keys()

        # Update words indexes:
        self._update_indexes(word_list=word_list)

        # Replace the words by their index:
        words_indexes_list = [self.word_to_index[word] for word in word_list]

        # Sample empirical topic assignment:
        _, components = self._compute_statistics_components(words_indexes_list)

        return dict(enumerate(components))

    def _update_indexes(self, word_list: typing.Iterable[str]):
        """
        Adds the words of the document to the index if they are not part of the current vocabulary.
        Updates of the number of distinct words seen.

        Parameters
        ----------
        word_list
            Content of the document as a list of words.

        """
        for word in word_list:
            if word not in self.word_to_index:
                new_index = len(self.word_to_index) + 1
                self.word_to_index[word] = new_index
                self.index_to_word[new_index] = word
                self.truncation_size_prime += 1

    @classmethod
    def _compute_weights(cls, n_components: int, nu_1: dict, nu_2: dict) -> tuple[dict, dict]:
        """Calculates the vocabulary weighting according to the word distribution present in the
        vocabulary.

        The Psi function is the logarithmic derivative of the gamma function.

        Parameters
        ----------
        n_components
            Number of topics.
        nu_1
            Weights of the words of the vocabulary.
        nu_2
            Weights of the words of the vocabulary.

        Returns
        -------
        Weights of the words of the current vocabulary.

        """
        exp_weights = {}
        exp_oov_weights = {}

        for topic in range(n_components):
            psi_nu_1 = special.psi(nu_1[topic])
            psi_nu_2 = special.psi(nu_2[topic])

            psi_nu_1_nu_2 = special.psi(nu_1[topic] + nu_2[topic])

            psi_nu_1_nu_2_minus_psi_nu_2 = np.cumsum([psi_nu_2 - psi_nu_1_nu_2], axis=1)

            exp_oov_weights[topic] = np.exp(psi_nu_1_nu_2_minus_psi_nu_2[0][-1])

            psi_nu_1_nu_2_minus_psi_nu_2 = ndimage.shift(
                input=psi_nu_1_nu_2_minus_psi_nu_2[0], shift=1, cval=0
            )

            exp_weights[topic] = np.exp(psi_nu_1 - psi_nu_1_nu_2 + psi_nu_1_nu_2_minus_psi_nu_2)

        return exp_weights, exp_oov_weights

    def _update_weights(self, statistics):
        """Learn documents and word representations. Calculate the variational approximation.

        Parameters
        ----------
        statistics
            Weights associated to the words.

        """
        reverse_cumulated_phi = {}

        for k in range(self.n_components):
            reverse_cumulated_phi[k] = ndimage.shift(input=statistics[k], shift=-1, cval=0)

            reverse_cumulated_phi[k] = np.flip(reverse_cumulated_phi[k])
            reverse_cumulated_phi[k] = np.cumsum(reverse_cumulated_phi[k])
            reverse_cumulated_phi[k] = np.flip(reverse_cumulated_phi[k])

        # Epsilon will be between 0 and 1.
        # Epsilon value says how much to weight the information we got from this document.
        self.epsilon = (self.tau + self.counter) ** -self.kappa

        for k in range(self.n_components):
            if self.truncation_size < self.truncation_size_prime:
                difference_truncation = self.truncation_size_prime - self.truncation_size

                self.nu_1[k] = np.append(self.nu_1[k], np.ones(difference_truncation))
                self.nu_2[k] = np.append(self.nu_2[k], np.ones(difference_truncation))

            # Variational Approximation
            self.nu_1[k] += self.epsilon * (
                self.number_of_documents * np.array(statistics[k]) + 1 - self.nu_1[k]
            )

            self.nu_2[k] += self.epsilon * (
                self.alpha_beta
                + self.number_of_documents * np.array(reverse_cumulated_phi[k])
                - self.nu_2[k]
            )

        self.truncation_size = self.truncation_size_prime

    def _compute_statistics_components(self, words_indexes_list: list) -> tuple[dict, dict]:
        """Extract latent variables from the document and words.

        Parameters
        ----------
        words_indexes_list
            Ids of the words of the input document.

        Returns
        -------
        Computed statistics over the words. Document reprensetation across topics.

        """
        statistics: defaultdict = defaultdict(lambda: np.zeros(self.truncation_size_prime))

        exp_weights, exp_oov_weights = self._compute_weights(
            n_components=self.n_components, nu_1=self.nu_1, nu_2=self.nu_2
        )

        size_vocab = len(words_indexes_list)

        phi = self.rng.random((self.n_components, size_vocab))

        phi /= np.sum(phi, axis=0)

        phi_sum = np.sum(phi, axis=1)

        for sample_index in range(self.number_of_samples):
            for word_index in range(size_vocab):
                phi_sum -= phi[:, word_index]
                phi_sum = phi_sum.clip(min=0)
                temp_phi = phi_sum + self.alpha_theta

                for k in range(self.n_components):
                    if words_indexes_list[word_index] >= self.truncation_size:
                        temp_phi[k] *= exp_oov_weights[k]
                    else:
                        temp_phi[k] *= exp_weights[k][words_indexes_list[word_index]]

                # Normalize document topic distribution before applying multinomial distribution:
                temp_phi /= temp_phi.sum()

                # Sample a topic based a given probability distribution:
                temp_phi = self.rng.multinomial(1, temp_phi)

                phi[:, word_index] = temp_phi

                phi_sum += temp_phi

                if sample_index >= self.burn_in_sweeps:
                    for k in range(self.n_components):
                        index = words_indexes_list[word_index]

                        statistics[k][index] += temp_phi[k]

        document_topic_distribution = self.alpha_theta + phi_sum

        for k in range(self.n_components):
            statistics[k] /= self.number_of_samples - self.burn_in_sweeps

        return statistics, document_topic_distribution

    def _prune_vocabulary(self):
        """Reduce the size of the index exceeds the maximum size."""
        if self.nu_1[0].shape[0] > self.maximum_size_vocabulary:
            for topic in range(self.n_components):
                # Updates words latent variables
                self.nu_1[topic] = self.nu_1[topic][: self.maximum_size_vocabulary]
                self.nu_2[topic] = self.nu_2[topic][: self.maximum_size_vocabulary]

            new_word_to_index = {}
            new_index_to_word = {}

            for index in range(1, self.maximum_size_vocabulary):
                # Updates words indexes
                word = self.index_to_word[index]
                new_word_to_index[word] = index
                new_index_to_word[index] = word

            self.word_to_index = new_word_to_index
            self.index_to_word = new_index_to_word

            self.truncation_size = self.nu_1[0].shape[0]
            self.truncation_size_prime = self.truncation_size


# river/preprocessing/one_hot.py

from __future__ import annotations

import collections

import pandas as pd

from river import base

__all__ = ["OneHotEncoder"]


class OneHotEncoder(base.MiniBatchTransformer):
    """One-hot encoding.

    This transformer will encode every feature it is provided with.
    If a list or set is provided, this transformer will encode every entry in the list/set.
    You can apply it to a subset of features by composing it
     with `compose.Select` or `compose.SelectType`.

    Parameters
    ----------
    drop_zeros
        Whether or not 0s should be made explicit or not.
    drop_first
        Whether to get `k - 1` dummies out of `k` categorical levels by removing the first key.
        This is useful in some statistical models where perfectly collinear features cause
        problems.

    Examples
    --------

    Let us first create an example dataset.

    >>> from pprint import pprint
    >>> import random
    >>> import string

    >>> random.seed(42)
    >>> alphabet = list(string.ascii_lowercase)
    >>> X = [
    ...     {
    ...         'c1': random.choice(alphabet),
    ...         'c2': random.choice(alphabet),
    ...     }
    ...     for _ in range(4)
    ... ]
    >>> pprint(X)
    [{'c1': 'u', 'c2': 'd'},
        {'c1': 'a', 'c2': 'x'},
        {'c1': 'i', 'c2': 'h'},
        {'c1': 'h', 'c2': 'e'}]

    e can now apply one-hot encoding. All the provided are one-hot encoded, there is therefore
    no need to specify which features to encode.

    >>> from river import preprocessing

    >>> oh = preprocessing.OneHotEncoder()
    >>> for x in X[:2]:
    ...     oh = oh.learn_one(x)
    ...     pprint(oh.transform_one(x))
    {'c1_u': 1, 'c2_d': 1}
    {'c1_a': 1, 'c1_u': 0, 'c2_d': 0, 'c2_x': 1}

    The `drop_zeros` parameter can be set to `True` if you don't want the past features to be included
    in the output. Otherwise, all the past features will be included in the output.

    >>> oh = preprocessing.OneHotEncoder(drop_zeros=True)
    >>> for x in X:
    ...     oh = oh.learn_one(x)
    ...     pprint(oh.transform_one(x))
    {'c1_u': 1, 'c2_d': 1}
    {'c1_a': 1, 'c2_x': 1}
    {'c1_i': 1, 'c2_h': 1}
    {'c1_h': 1, 'c2_e': 1}

    You can encode only `k - 1` features out of `k` by setting `drop_first` to `True`.

    >>> oh = preprocessing.OneHotEncoder(drop_first=True, drop_zeros=True)
    >>> for x in X:
    ...     oh = oh.learn_one(x)
    ...     pprint(oh.transform_one(x))
    {'c2_d': 1}
    {'c2_x': 1}
    {'c2_h': 1}
    {'c2_e': 1}

    A subset of the features can be one-hot encoded by piping a `compose.Select` into the
    `OneHotEncoder`.

    >>> from river import compose

    >>> pp = compose.Select('c1') | preprocessing.OneHotEncoder()

    >>> for x in X:
    ...     pp = pp.learn_one(x)
    ...     pprint(pp.transform_one(x))
    {'c1_u': 1}
    {'c1_a': 1, 'c1_u': 0}
    {'c1_a': 0, 'c1_i': 1, 'c1_u': 0}
    {'c1_a': 0, 'c1_h': 1, 'c1_i': 0, 'c1_u': 0}

    You can preserve the `c2` feature by using a union:

    >>> pp = compose.Select('c1') | preprocessing.OneHotEncoder()
    >>> pp += compose.Select('c2')

    >>> for x in X:
    ...     pp = pp.learn_one(x)
    ...     pprint(pp.transform_one(x))
    {'c1_u': 1, 'c2': 'd'}
    {'c1_a': 1, 'c1_u': 0, 'c2': 'x'}
    {'c1_a': 0, 'c1_i': 1, 'c1_u': 0, 'c2': 'h'}
    {'c1_a': 0, 'c1_h': 1, 'c1_i': 0, 'c1_u': 0, 'c2': 'e'}

    Similar to the above examples, we can also pass values as a list. This will one-hot
    encode all of the entries individually.

    >>> X = [{'c1': ['u', 'a'], 'c2': ['d']},
    ...     {'c1': ['a', 'b'], 'c2': ['x']},
    ...     {'c1': ['i'], 'c2': ['h', 'z']},
    ...     {'c1': ['h', 'b'], 'c2': ['e']}]

    >>> oh = preprocessing.OneHotEncoder(drop_zeros=True)
    >>> for x in X:
    ...     oh = oh.learn_one(x)
    ...     pprint(oh.transform_one(x))
    {'c1_a': 1, 'c1_u': 1, 'c2_d': 1}
    {'c1_a': 1, 'c1_b': 1, 'c2_x': 1}
    {'c1_i': 1, 'c2_h': 1, 'c2_z': 1}
    {'c1_b': 1, 'c1_h': 1, 'c2_e': 1}

    Processing mini-batches is also possible.

    >>> from pprint import pprint
    >>> import random
    >>> import string

    >>> random.seed(42)
    >>> alphabet = list(string.ascii_lowercase)
    >>> X = pd.DataFrame(
    ...     {
    ...         'c1': random.choice(alphabet),
    ...         'c2': random.choice(alphabet),
    ...     }
    ...     for _ in range(3)
    ... )
    >>> X
      c1 c2
    0  u  d
    1  a  x
    2  i  h

    >>> oh = preprocessing.OneHotEncoder(drop_zeros=True)
    >>> df = oh.transform_many(X)
    >>> df.sort_index(axis="columns")
       c1_a  c1_i  c1_u  c2_d  c2_h  c2_x
    0     0     0     1     1     0     0
    1     1     0     0     0     0     1
    2     0     1     0     0     1     0

    >>> oh = preprocessing.OneHotEncoder(drop_zeros=True, drop_first=True)
    >>> df = oh.transform_many(X)
    >>> df.sort_index(axis="columns")
       c1_i  c1_u  c2_d  c2_h  c2_x
    0     0     1     1     0     0
    1     0     0     0     0     1
    2     1     0     0     1     0

    Here's an example where the zeros are kept:

    >>> oh = preprocessing.OneHotEncoder(drop_zeros=False)
    >>> X_init = pd.DataFrame([{"c1": "Oranges", "c2": "Apples"}])
    >>> oh = oh.learn_many(X_init)
    >>> oh = oh.learn_many(X)

    >>> df = oh.transform_many(X)
    >>> df.sort_index(axis="columns")
       c1_Oranges  c1_a  c1_i  c1_u  c2_Apples  c2_d  c2_h  c2_x
    0           0     0     0     1          0     1     0     0
    1           0     1     0     0          0     0     0     1
    2           0     0     1     0          0     0     1     0

    >>> df.dtypes.sort_index()
    c1_Oranges    Sparse[uint8, 0]
    c1_a          Sparse[uint8, 0]
    c1_i          Sparse[uint8, 0]
    c1_u          Sparse[uint8, 0]
    c2_Apples     Sparse[uint8, 0]
    c2_d          Sparse[uint8, 0]
    c2_h          Sparse[uint8, 0]
    c2_x          Sparse[uint8, 0]
    dtype: object

    """

    def __init__(self, drop_zeros=False, drop_first=False):
        self.drop_zeros = drop_zeros
        self.drop_first = drop_first
        self.values = collections.defaultdict(set)

    def learn_one(self, x):
        if self.drop_zeros:
            return self

        for i, xi in x.items():
            if isinstance(xi, list) or isinstance(xi, set):
                for xj in xi:
                    self.values[i].add(xj)
            else:
                self.values[i].add(xi)

        return self

    def transform_one(self, x, y=None):
        oh = {}

        # Add 0s
        if not self.drop_zeros:
            oh = {f"{i}_{v}": 0 for i, values in self.values.items() for v in values}

        # Add 1s
        for i, xi in x.items():
            if isinstance(xi, list) or isinstance(xi, set):
                for xj in xi:
                    oh[f"{i}_{xj}"] = 1
            else:
                oh[f"{i}_{xi}"] = 1

        if self.drop_first:
            oh.pop(min(oh.keys()))

        return oh

    def learn_many(self, X):
        if self.drop_zeros:
            return self

        for col in X.columns:
            self.values[col].update(X[col].unique())

        return self

    def transform_many(self, X):
        oh = pd.get_dummies(X, columns=X.columns, sparse=True, dtype="uint8")

        if not self.drop_zeros:
            seen_in_the_past = {f"{col}_{val}" for col, vals in self.values.items() for val in vals}
            to_add = seen_in_the_past - set(oh.columns)
            for col in to_add:
                oh[col] = pd.arrays.SparseArray([0] * len(oh), dtype="uint8")

        if self.drop_first:
            oh = oh.drop(columns=min(oh.columns))

        return oh


# river/preprocessing/pred_clipper.py

from __future__ import annotations

from river import base, utils

__all__ = ["PredClipper"]


class PredClipper(base.Wrapper, base.Regressor):
    """Clips the target after predicting.

    Parameters
    ----------
    regressor
        Regressor model for which to clip the predictions.
    y_min
        minimum value.
    y_max
        maximum value.

    Examples
    --------
    >>> from river import linear_model
    >>> from river import preprocessing

    >>> dataset = (
    ...     ({'a': 2, 'b': 4}, 80),
    ...     ({'a': 3, 'b': 5}, 100),
    ...     ({'a': 4, 'b': 6}, 120)
    ... )

    >>> model = preprocessing.PredClipper(
    ...     regressor=linear_model.LinearRegression(),
    ...     y_min=0,
    ...     y_max=200
    ... )

    >>> for x, y in dataset:
    ...     _ = model.learn_one(x, y)

    >>> model.predict_one({'a': -100, 'b': -200})
    0

    >>> model.predict_one({'a': 50, 'b': 60})
    200

    """

    def __init__(self, regressor: base.Regressor, y_min: float, y_max: float):
        self.regressor = regressor
        self.y_min = y_min
        self.y_max = y_max

    @property
    def _wrapped_model(self):
        return self.regressor

    def learn_one(self, x, y, **kwargs):
        self.regressor.learn_one(x=x, y=y, **kwargs)
        return self

    def predict_one(self, x, **kwargs):
        y_pred = self.regressor.predict_one(x=x, **kwargs)
        return utils.math.clamp(y_pred, self.y_min, self.y_max)

    @classmethod
    def _unit_test_params(cls):
        import math

        from river import linear_model

        yield {
            "regressor": linear_model.LinearRegression(),
            "y_min": -math.inf,
            "y_max": math.inf,
        }


# river/preprocessing/random_projection.py

from __future__ import annotations

import collections
import random

from river import base


class GaussianRandomProjector(base.Transformer):
    """Gaussian random projector.

    This transformer reduces the dimensionality of inputs through Gaussian random projection.

    The components of the random projections matrix are drawn from `N(0, 1 / n_components)`.

    Parameters
    ----------
    n_components
        Number of components to project the data onto.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()
    >>> model = preprocessing.GaussianRandomProjector(
    ...     n_components=3,
    ...     seed=42
    ... )

    >>> for x, y in dataset:
    ...     x = model.transform_one(x)
    ...     print(x)
    ...     break
    {0: -61289.37139206629, 1: 141312.51039283074, 2: 279165.99370457436}

    >>> model = (
    ...     preprocessing.GaussianRandomProjector(
    ...         n_components=5,
    ...         seed=42
    ...     ) |
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LinearRegression()
    ... )
    >>> evaluate.progressive_val_score(dataset, model, metrics.MAE())
    MAE: 0.860464

    References
    ----------
    [^1]: [Gaussian random projection](https://www.wikiwand.com/en/Gaussian_random_projection#Gaussian_random_projection)
    [^2]: [scikit-learn random projections module](https://scikit-learn.org/stable/modules/random_projection.html)

    """

    def __init__(self, n_components=10, seed: int | None = None):
        self.n_components = n_components
        self.seed = seed
        self._rng = random.Random(seed)
        self._projection_matrix: collections.defaultdict[
            base.typing.FeatureName, float
        ] = collections.defaultdict(self._rand_gauss)

    def _rand_gauss(self):
        return self._rng.gauss(0, 1 / (self.n_components**0.5))

    def transform_one(self, x):
        return {
            i: sum(self._projection_matrix[(i, j)] * x[j] for j in x)
            for i in range(self.n_components)
        }


class SparseRandomProjector(base.Transformer):
    """Sparse random projector.

    This transformer reduces the dimensionality of inputs by projecting them onto a sparse random
    projection matrix.

    Ping Li et al. recommend using a minimum density of `1 / sqrt(n_features)`. The transformer
    is not aware of how many features will be seen, so the user must specify the density manually.

    Parameters
    ----------
    n_components
        Number of components to project the data onto.
    density
        Density of the random projection matrix. The density is defined as the ratio of non-zero
        components in the matrix. It is equal to `1 - sparsity`.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()
    >>> model = preprocessing.SparseRandomProjector(
    ...     n_components=3,
    ...     seed=42
    ... )

    >>> for x, y in dataset:
    ...     x = model.transform_one(x)
    ...     print(x)
    ...     break
    {0: 92.89572746525327, 1: 1344540.5692342375, 2: 0}

    >>> model = (
    ...     preprocessing.SparseRandomProjector(
    ...         n_components=5,
    ...         seed=42
    ...     ) |
    ...     preprocessing.StandardScaler() |
    ...     linear_model.LinearRegression()
    ... )
    >>> evaluate.progressive_val_score(dataset, model, metrics.MAE())
    MAE: 1.296503

    References
    ----------
    [^1]: D. Achlioptas. 2003. Database-friendly random projections: Johnson-Lindenstrauss with binary coins. Journal of Computer and System Sciences 66 (2003) 671-687
    [^2]: Ping Li, Trevor J. Hastie, and Kenneth W. Church. 2006. Very sparse random projections. In Proceedings of the 12th ACM SIGKDD international conference on Knowledge discovery and data mining (KDD'06). ACM, New York, NY, USA, 287-296.

    """

    def __init__(self, n_components=10, density=0.1, seed: int | None = None):
        self.n_components = n_components
        self.density = density
        self.seed = seed
        self._rng = random.Random(seed)
        self._projection_matrix: collections.defaultdict[
            base.typing.FeatureName, dict[int, float]
        ] = collections.defaultdict(self._rand_weights_for_feature)

    def _rand_weights_for_feature(self):
        weights = {}
        for j in range(self.n_components):
            if self._rng.random() < self.density:
                w = (1 / (self.density * self.n_components)) ** 0.5
                # Flip a coin to decide the sign
                weights[j] = w if self._rng.random() < 0.5 else -w
        return weights

    def transform_one(self, x):
        output = {i: 0 for i in range(self.n_components)}
        for j in x:
            for i, w in self._projection_matrix[j].items():
                output[i] += w * x[j]
        return output


# river/preprocessing/scale.py

from __future__ import annotations

import collections
import functools
import itertools
import numbers

import numpy as np
import pandas as pd

from river import base, stats, utils

__all__ = [
    "AdaptiveStandardScaler",
    "Binarizer",
    "MaxAbsScaler",
    "MinMaxScaler",
    "Normalizer",
    "RobustScaler",
    "StandardScaler",
]


def safe_div(a, b):
    """Returns a if b is nil, else divides a by b.

    When scaling, sometimes a denominator might be nil. For instance, during standard scaling
    the denominator can be nil if a feature has no variance.

    """
    return a / b if b else 0.0


class Binarizer(base.Transformer):
    """Binarizes the data to 0 or 1 according to a threshold.

    Parameters
    ----------
    threshold
        Values above this are replaced by 1 and the others by 0.
    dtype
        The desired data type to apply.

    Examples
    --------

    >>> import river
    >>> import numpy as np

    >>> rng = np.random.RandomState(42)
    >>> X = [{'x1': v, 'x2': int(v)} for v in rng.uniform(low=-4, high=4, size=6)]

    >>> binarizer = river.preprocessing.Binarizer()
    >>> for x in X:
    ...     print(binarizer.learn_one(x).transform_one(x))
    {'x1': False, 'x2': False}
    {'x1': True, 'x2': True}
    {'x1': True, 'x2': True}
    {'x1': True, 'x2': False}
    {'x1': False, 'x2': False}
    {'x1': False, 'x2': False}

    """

    def __init__(self, threshold=0.0, dtype=bool):
        self.threshold = threshold
        self.dtype = dtype

    def transform_one(self, x):
        x_tf = x.copy()

        for i, xi in x_tf.items():
            if isinstance(xi, numbers.Number):
                x_tf[i] = self.dtype(xi > self.threshold)

        return x_tf


class StandardScaler(base.MiniBatchTransformer):
    """Scales the data so that it has zero mean and unit variance.

    Under the hood, a running mean and a running variance are maintained. The scaling is slightly
    different than when scaling the data in batch because the exact means and variances are not
    known in advance. However, this doesn't have a detrimental impact on performance in the long
    run.

    This transformer supports mini-batches as well as single instances. In the mini-batch case, the
    number of columns and the ordering of the columns are allowed to change between subsequent
    calls. In other words, this transformer will keep working even if you add and/or remove
    features every time you call `learn_many` and `transform_many`.

    Parameters
    ----------
    with_std
        Whether or not each feature should be divided by its standard deviation.

    Examples
    --------

    >>> import random
    >>> from river import preprocessing

    >>> random.seed(42)
    >>> X = [{'x': random.uniform(8, 12), 'y': random.uniform(8, 12)} for _ in range(6)]
    >>> for x in X:
    ...     print(x)
    {'x': 10.557, 'y': 8.100}
    {'x': 9.100, 'y': 8.892}
    {'x': 10.945, 'y': 10.706}
    {'x': 11.568, 'y': 8.347}
    {'x': 9.687, 'y': 8.119}
    {'x': 8.874, 'y': 10.021}

    >>> scaler = preprocessing.StandardScaler()

    >>> for x in X:
    ...     print(scaler.learn_one(x).transform_one(x))
    {'x': 0.0, 'y': 0.0}
    {'x': -0.999, 'y': 0.999}
    {'x': 0.937, 'y': 1.350}
    {'x': 1.129, 'y': -0.651}
    {'x': -0.776, 'y': -0.729}
    {'x': -1.274, 'y': 0.992}

    This transformer also supports mini-batch updates. You can call `learn_many` and provide a
    `pandas.DataFrame`:

    >>> import pandas as pd
    >>> X = pd.DataFrame.from_dict(X)

    >>> scaler = preprocessing.StandardScaler()
    >>> scaler = scaler.learn_many(X[:3])
    >>> scaler = scaler.learn_many(X[3:])

    You can then call `transform_many` to scale a mini-batch of features:

    >>> scaler.transform_many(X)
        x         y
    0  0.444600 -0.933384
    1 -1.044259 -0.138809
    2  0.841106  1.679208
    3  1.477301 -0.685117
    4 -0.444084 -0.914195
    5 -1.274664  0.992296

    References
    ----------
    [^1]: [Welford's Method (and Friends)](https://www.embeddedrelated.com/showarticle/785.php)
    [^2]: [Batch updates for simple statistics](https://notmatthancock.github.io/2017/03/23/simple-batch-stat-updates.html)

    """

    def __init__(self, with_std=True):
        self.with_std = with_std
        self.counts = collections.Counter()
        self.means = collections.defaultdict(float)
        self.vars = collections.defaultdict(float)

    def learn_one(self, x):
        for i, xi in x.items():
            self.counts[i] += 1
            old_mean = self.means[i]
            self.means[i] += (xi - old_mean) / self.counts[i]
            if self.with_std:
                self.vars[i] += (
                    (xi - old_mean) * (xi - self.means[i]) - self.vars[i]
                ) / self.counts[i]

        return self

    def transform_one(self, x):
        if self.with_std:
            return {i: safe_div(xi - self.means[i], self.vars[i] ** 0.5) for i, xi in x.items()}
        return {i: xi - self.means[i] for i, xi in x.items()}

    def learn_many(self, X: pd.DataFrame):
        """Update with a mini-batch of features.

        Note that the update formulas for mean and variance are slightly different than in the
        single instance case, but they produce exactly the same result.

        Parameters
        ----------
        X
            A dataframe where each column is a feature.

        """

        # Operating on X.values, which is a view to the underlying numpy array, is slightly faster
        # than operating on X
        columns = X.columns
        X = X.values

        # In the rest of this method, old_* refers to the existing statistics, whilst new_* refers
        # to the statistics of the current mini-batch.

        new_means = np.nanmean(X, axis=0)
        # We could call np.var, but we already have the mean so we can be smart
        if self.with_std:
            new_vars = np.einsum("ij,ij->j", X, X) / len(X) - new_means**2
        else:
            new_vars = []
        new_counts = np.sum(~np.isnan(X), axis=0)

        for col, new_mean, new_var, new_count in itertools.zip_longest(
            columns, new_means, new_vars, new_counts
        ):
            old_mean = self.means[col]
            old_var = self.vars[col]
            old_count = self.counts[col]

            a = old_count / (old_count + new_count)
            b = new_count / (old_count + new_count)

            self.means[col] = a * old_mean + b * new_mean
            if self.with_std:
                self.vars[col] = a * old_var + b * new_var + a * b * (old_mean - new_mean) ** 2
            self.counts[col] += new_count

        return self

    def transform_many(self, X: pd.DataFrame):
        """Scale a mini-batch of features.

        Parameters
        ----------
        X
            A dataframe where each column is a feature. An exception will be raised if any of
            the features has not been seen during a previous call to `learn_many`.

        """

        means = np.array([self.means[c] for c in X.columns])
        Xt = X.values - means

        if self.with_std:
            stds = np.array([self.vars[c] ** 0.5 for c in X.columns])
            np.divide(Xt, stds, where=stds > 0, out=Xt)

        return pd.DataFrame(Xt, index=X.index, columns=X.columns, copy=False)


class MinMaxScaler(base.Transformer):
    """Scales the data to a fixed range from 0 to 1.

    Under the hood a running min and a running peak to peak (max - min) are maintained.

    Attributes
    ----------
    min : dict
        Mapping between features and instances of `stats.Min`.
    max : dict
        Mapping between features and instances of `stats.Max`.

    Examples
    --------

    >>> import random
    >>> from river import preprocessing

    >>> random.seed(42)
    >>> X = [{'x': random.uniform(8, 12)} for _ in range(5)]
    >>> for x in X:
    ...     print(x)
    {'x': 10.557707}
    {'x': 8.100043}
    {'x': 9.100117}
    {'x': 8.892842}
    {'x': 10.945884}

    >>> scaler = preprocessing.MinMaxScaler()

    >>> for x in X:
    ...     print(scaler.learn_one(x).transform_one(x))
    {'x': 0.0}
    {'x': 0.0}
    {'x': 0.406920}
    {'x': 0.322582}
    {'x': 1.0}

    """

    def __init__(self):
        self.min = collections.defaultdict(stats.Min)
        self.max = collections.defaultdict(stats.Max)

    def learn_one(self, x):
        for i, xi in x.items():
            self.min[i].update(xi)
            self.max[i].update(xi)

        return self

    def transform_one(self, x):
        return {
            i: safe_div(xi - self.min[i].get(), self.max[i].get() - self.min[i].get())
            for i, xi in x.items()
        }


class MaxAbsScaler(base.Transformer):
    """Scales the data to a [-1, 1] range based on absolute maximum.

    Under the hood a running absolute max is maintained. This scaler is meant for
    data that is already centered at zero or sparse data. It does not shift/center
    the data, and thus does not destroy any sparsity.

    Attributes
    ----------
    abs_max : dict
        Mapping between features and instances of `stats.AbsMax`.

    Examples
    --------

    >>> import random
    >>> from river import preprocessing

    >>> random.seed(42)
    >>> X = [{'x': random.uniform(8, 12)} for _ in range(5)]
    >>> for x in X:
    ...     print(x)
    {'x': 10.557707}
    {'x': 8.100043}
    {'x': 9.100117}
    {'x': 8.892842}
    {'x': 10.945884}

    >>> scaler = preprocessing.MaxAbsScaler()

    >>> for x in X:
    ...     print(scaler.learn_one(x).transform_one(x))
    {'x': 1.0}
    {'x': 0.767216}
    {'x': 0.861940}
    {'x': 0.842308}
    {'x': 1.0}

    """

    def __init__(self):
        self.abs_max = collections.defaultdict(stats.AbsMax)

    def learn_one(self, x):
        for i, xi in x.items():
            self.abs_max[i].update(xi)

        return self

    def transform_one(self, x):
        return {i: safe_div(xi, self.abs_max[i].get()) for i, xi in x.items()}


class RobustScaler(base.Transformer):
    """Scale features using statistics that are robust to outliers.

    This Scaler removes the median and scales the data according to the
    interquantile range.

    Parameters
    ----------
    with_centering
        Whether to centre the data before scaling.
    with_scaling
        Whether to scale data to IQR.
    q_inf
        Desired inferior quantile, must be between 0 and 1.
    q_sup
        Desired superior quantile, must be between 0 and 1.

    Attributes
    ----------
    median : dict
        Mapping between features and instances of `stats.Quantile(0.5)`.
    iqr : dict
        Mapping between features and instances of `stats.IQR`.

    Examples
    --------

    >>> from pprint import pprint
    >>> import random
    >>> from river import preprocessing

    >>> random.seed(42)
    >>> X = [{'x': random.uniform(8, 12)} for _ in range(5)]
    >>> pprint(X)
    [{'x': 10.557707},
        {'x': 8.100043},
        {'x': 9.100117},
        {'x': 8.892842},
        {'x': 10.945884}]

    >>> scaler = preprocessing.RobustScaler()

    >>> for x in X:
    ...     print(scaler.learn_one(x).transform_one(x))
        {'x': 0.0}
        {'x': -1.0}
        {'x': 0.0}
        {'x': -0.12449923287875722}
        {'x': 1.1086595155704708}

    """

    def __init__(self, with_centering=True, with_scaling=True, q_inf=0.25, q_sup=0.75):
        self.with_centering = with_centering
        self.with_scaling = with_scaling
        self.q_inf = q_inf
        self.q_sup = q_sup
        self.median = collections.defaultdict(functools.partial(stats.Quantile, 0.5))
        self.iqr = collections.defaultdict(functools.partial(stats.IQR, self.q_inf, self.q_sup))

    def learn_one(self, x):
        for i, xi in x.items():
            if self.with_centering:
                self.median[i].update(xi)
            if self.with_scaling:
                self.iqr[i].update(xi)

        return self

    def transform_one(self, x):
        x_tf = {}

        for i, xi in x.items():
            x_tf[i] = xi
            if self.with_centering:
                x_tf[i] -= self.median[i].get()
            if self.with_scaling:
                x_tf[i] = safe_div(x_tf[i], self.iqr[i].get())

        return x_tf


class Normalizer(base.Transformer):
    """Scales a set of features so that it has unit norm.

    This is particularly useful when used after a `feature_extraction.TFIDF`.

    Parameters
    ----------
    order
        Order of the norm (e.g. 2 corresponds to the $L^2$ norm).

    Examples
    --------

    >>> from river import preprocessing
    >>> from river import stream

    >>> scaler = preprocessing.Normalizer(order=2)

    >>> X = [[4, 1, 2, 2],
    ...      [1, 3, 9, 3],
    ...      [5, 7, 5, 1]]

    >>> for x, _ in stream.iter_array(X):
    ...     print(scaler.transform_one(x))
    {0: 0.8, 1: 0.2, 2: 0.4, 3: 0.4}
    {0: 0.1, 1: 0.3, 2: 0.9, 3: 0.3}
    {0: 0.5, 1: 0.7, 2: 0.5, 3: 0.1}

    """

    def __init__(self, order=2):
        self.order = order

    def transform_one(self, x):
        norm = utils.math.norm(x, order=self.order)
        return {i: xi / norm for i, xi in x.items()}


class AdaptiveStandardScaler(base.Transformer):
    """Scales data using exponentially weighted moving average and variance.

    Under the hood, a exponentially weighted running mean and variance are maintained for each
    feature. This can potentially provide better results for drifting data in comparison to
    `preprocessing.StandardScaler`. Indeed, the latter computes a global mean and variance for each
    feature, whereas this scaler weights data in proportion to their recency.

    Parameters
    ----------
    fading_factor
        This parameter is passed to `stats.EWVar`. It is expected to be in [0, 1]. More weight is
        assigned to recent samples the closer `fading_factor` is to 1.

    Examples
    --------

    Consider the following series which contains a positive trend.

    >>> import random

    >>> random.seed(42)
    >>> X = [
    ...     {'x': random.uniform(4 + i, 6 + i)}
    ...     for i in range(8)
    ... ]
    >>> for x in X:
    ...     print(x)
    {'x': 5.278}
    {'x': 5.050}
    {'x': 6.550}
    {'x': 7.446}
    {'x': 9.472}
    {'x': 10.353}
    {'x': 11.784}
    {'x': 11.173}

    This scaler works well with this kind of data because it uses statistics that assign higher
    weight to more recent data.

    >>> from river import preprocessing

    >>> scaler = preprocessing.AdaptiveStandardScaler(fading_factor=.6)

    >>> for x in X:
    ...     print(scaler.learn_one(x).transform_one(x))
    {'x': 0.0}
    {'x': -0.816}
    {'x': 0.812}
    {'x': 0.695}
    {'x': 0.754}
    {'x': 0.598}
    {'x': 0.651}
    {'x': 0.124}

    """

    def __init__(self, fading_factor=0.3):
        self.fading_factor = fading_factor
        self.vars = collections.defaultdict(functools.partial(stats.EWVar, self.fading_factor))
        self.means = collections.defaultdict(functools.partial(stats.EWMean, self.fading_factor))

    def learn_one(self, x):
        for i, xi in x.items():
            self.vars[i].update(xi)
            self.means[i].update(xi)
        return self

    def transform_one(self, x):
        return {
            i: safe_div(x[i] - m, s2**0.5 if s2 > 0 else 0)
            for i, m, s2 in ((i, self.means[i].get(), self.vars[i].get()) for i in x)
        }


# river/preprocessing/scale_target.py

from __future__ import annotations

import math

from river import base, compose, stats

__all__ = [
    "TargetMinMaxScaler",
    "TargetStandardScaler",
]


def safe_div(a, b):
    """Returns a if b is nil, else divides a by b.

    When scaling, sometimes a denominator might be nil. For instance, during standard scaling
    the denominator can be nil if a feature has no variance.

    """
    return a / b if b else 0.0


class TargetStandardScaler(compose.TargetTransformRegressor):
    """Applies standard scaling to the target.

    Parameters
    ----------
    regressor
        Regression model to wrap.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     preprocessing.TargetStandardScaler(
    ...         regressor=linear_model.LinearRegression(intercept_lr=0.15)
    ...     )
    ... )
    >>> metric = metrics.MSE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MSE: 2.003724



    """

    def __init__(self, regressor: base.Regressor):
        self.var = stats.Var()
        super().__init__(regressor=regressor, func=self._scale, inverse_func=self._unscale)

    def _update(self, y):
        self.var.update(y)

    def _scale(self, y):
        return safe_div(y - self.var.mean.get(), self.var.get() ** 0.5)

    def _unscale(self, y):
        return y * self.var.get() ** 0.5 + self.var.mean.get()


class TargetMinMaxScaler(compose.TargetTransformRegressor):
    """Applies min-max scaling to the target.

    Parameters
    ----------
    regressor
        Regression model to wrap.

    Examples
    --------

    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     preprocessing.TargetMinMaxScaler(
    ...         regressor=linear_model.LinearRegression(intercept_lr=0.15)
    ...     )
    ... )
    >>> metric = metrics.MSE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MSE: 2.01689

    """

    def __init__(self, regressor: base.Regressor):
        self.min = stats.Min()
        self.max = stats.Max()
        super().__init__(regressor=regressor, func=self._scale, inverse_func=self._unscale)

    def _update(self, y):
        self.min.update(y)
        self.max.update(y)

    def _scale(self, y):
        return safe_div(y - self.min.get(), self.max.get() - self.min.get())

    def _unscale(self, y):
        if self.min.get() == math.inf:
            return y
        return y * (self.max.get() - self.min.get()) + self.min.get()


# river/preprocessing/test_lda.py

"""
The tests performed here confirm that the outputs of the online preprocessing.LDA are exactly the
same as those of the original with a batch_size of size 1. Coverage is 100%.

References
----------
[^1]: Jordan Boyd-Graber, Ke Zhai, Online Latent Dirichlet Allocation with Infinite Vocabulary.
    http://proceedings.mlr.press/v28/zhai13.pdf
[^2]: river's LDA implementation reproduces exactly the same results as the original one from
    PyInfVov (https://github.com/kzhai/PyInfVoc) with a batch size of 1.

"""
from __future__ import annotations

import numpy as np

from river import preprocessing

DOC_SET = [
    "weather cold",
    "weather hot dry",
    "weather cold rainny",
    "weather hot",
    "weather cold humid",
]

REFERENCE_STATISTICS_TWO_COMPONENTS = [
    {0: np.array([0.0, 0.0, 0.0]), 1: np.array([0.0, 1.0, 1.0])},
    {0: np.array([0.0, 0.6, 0.0, 0.8, 0.6]), 1: np.array([0.0, 0.4, 0.0, 0.2, 0.4])},
    {
        0: np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
        1: np.array([0.0, 1.0, 1.0, 0.0, 0.0, 1.0]),
    },
    {
        0: np.array([0.0, 0.2, 0.0, 0.6, 0.0, 0.0]),
        1: np.array([0.0, 0.8, 0.0, 0.4, 0.0, 0.0]),
    },
    {
        0: np.array([0.0, 0.0, 0.2, 0.0, 0.0, 0.0, 0.2]),
        1: np.array([0.0, 1.0, 0.8, 0.0, 0.0, 0.0, 0.8]),
    },
]


REFERENCE_STATISTICS_FIVE_COMPONENTS = [
    {
        0: np.array([0.0, 0.4, 0.2]),
        1: np.array([0.0, 0.2, 0.6]),
        2: np.array([0.0, 0.4, 0.0]),
        3: np.array([0.0, 0.0, 0.0]),
        4: np.array([0.0, 0.0, 0.2]),
    },
    {
        0: np.array([0.0, 0.8, 0.0, 0.4, 0.6]),
        1: np.array([0.0, 0.0, 0.0, 0.2, 0.4]),
        2: np.array([0.0, 0.0, 0.0, 0.0, 0.0]),
        3: np.array([0.0, 0.0, 0.0, 0.2, 0.0]),
        4: np.array([0.0, 0.2, 0.0, 0.2, 0.0]),
    },
    {
        0: np.array([0.0, 0.4, 0.2, 0.0, 0.0, 0.0]),
        1: np.array([0.0, 0.2, 0.6, 0.0, 0.0, 0.6]),
        2: np.array([0.0, 0.4, 0.2, 0.0, 0.0, 0.4]),
        3: np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
        4: np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
    },
    {
        0: np.array([0.0, 0.2, 0.0, 0.4, 0.0, 0.0]),
        1: np.array([0.0, 0.2, 0.0, 0.2, 0.0, 0.0]),
        2: np.array([0.0, 0.4, 0.0, 0.2, 0.0, 0.0]),
        3: np.array([0.0, 0.2, 0.0, 0.2, 0.0, 0.0]),
        4: np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
    },
    {
        0: np.array([0.0, 0.0, 0.2, 0.0, 0.0, 0.0, 0.2]),
        1: np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
        2: np.array([0.0, 0.8, 0.8, 0.0, 0.0, 0.0, 0.6]),
        3: np.array([0.0, 0.2, 0.0, 0.0, 0.0, 0.0, 0.2]),
        4: np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]),
    },
]

REFERENCE_FIVE_COMPONENTS = [
    np.array([1.5, 0.5, 0.5, 0.5, 1.5]),
    np.array([1.5, 1.5, 0.5, 0.5, 1.5]),
    np.array([0.5, 0.5, 3.5, 0.5, 0.5]),
    np.array([0.5, 0.5, 0.5, 2.5, 0.5]),
    np.array([2.5, 0.5, 0.5, 1.5, 0.5]),
]

REFERENCE_COMPONENTS_WITH_PRUNNING = [
    np.array([0.5, 2.5]),
    np.array([3.5, 0.5]),
    np.array([0.5, 3.5]),
    np.array([2.5, 0.5]),
    np.array([2.5, 1.5]),
]

REFERENCE_LEARN_ONE_PREDICT_ONE = [
    np.array([2.5, 0.5]),
    np.array([2.5, 1.5]),
    np.array([1.5, 2.5]),
    np.array([0.5, 2.5]),
    np.array([0.5, 3.5]),
]


def test_extraction_words_ids():
    """
    Assert that input words are split.
    Assert that indexes are updated and extractable.
    """

    lda = preprocessing.LDA(2, number_of_documents=5, seed=42)

    word_indexes_list = []

    for doc in DOC_SET:
        words = doc.split(" ")

        lda._update_indexes(word_list=words)

        word_indexes_list.append([lda.word_to_index[word] for word in words])

    assert word_indexes_list == [[1, 2], [1, 3, 4], [1, 2, 5], [1, 3], [1, 2, 6]]


def test_statistics_two_components():
    """
    Assert that online lda extracts waited statistics on current document.
    """
    n_components = 2

    lda = preprocessing.LDA(n_components, number_of_documents=60, seed=42)

    statistics_list = []

    for doc in DOC_SET:
        word_list = doc.split(" ")

        lda._update_indexes(word_list=word_list)

        word_indexes = [lda.word_to_index[word] for word in word_list]

        statistics, _ = lda._compute_statistics_components(words_indexes_list=word_indexes)

        statistics_list.append(statistics)

        lda._update_weights(statistics=statistics)

    for index, statistics in enumerate(statistics_list):
        for component in range(n_components):
            assert np.array_equal(
                a1=statistics[component],
                a2=REFERENCE_STATISTICS_TWO_COMPONENTS[index][component],
            )


def test_statistics_five_components():
    """
    Assert that online lda extracts waited statistics on current document.
    """

    n_components = 5

    lda = preprocessing.LDA(
        n_components=n_components,
        number_of_documents=60,
        maximum_size_vocabulary=100,
        alpha_beta=100,
        alpha_theta=0.5,
        seed=42,
    )

    statistics_list = []

    for doc in DOC_SET:
        word_list = doc.split(" ")

        lda._update_indexes(word_list=word_list)

        word_indexes = [lda.word_to_index[word] for word in word_list]

        statistics, _ = lda._compute_statistics_components(words_indexes_list=word_indexes)

        statistics_list.append(statistics)

        lda._update_weights(statistics=statistics)

    for index, statistics in enumerate(statistics_list):
        for component in range(n_components):
            assert np.array_equal(
                a1=statistics[component],
                a2=REFERENCE_STATISTICS_FIVE_COMPONENTS[index][component],
            )


def test_five_components():
    """
    Assert that components computed are identical to the original version for n dimensions.
    """

    n_components = 5

    lda = preprocessing.LDA(
        n_components=n_components,
        number_of_documents=60,
        maximum_size_vocabulary=100,
        alpha_beta=100,
        alpha_theta=0.5,
        seed=42,
    )

    components_list = []

    for document in DOC_SET:
        tokens = {token: 1 for token in document.split(" ")}
        components_list.append(lda.learn_transform_one(tokens))

    for index, component in enumerate(components_list):
        assert np.array_equal(a1=list(component.values()), a2=REFERENCE_FIVE_COMPONENTS[index])


def test_prunning_vocabulary():
    """
    Vocabulary prunning is available to improve accuracy and limit memory usage.
    You can perform vocabulary prunning with parameters vocab_prune_interval (int) and
    maximum_size_vocabulary (int).
    """

    lda = preprocessing.LDA(
        n_components=2,
        number_of_documents=60,
        vocab_prune_interval=2,
        maximum_size_vocabulary=3,
        seed=42,
    )

    components_list = []

    for document in DOC_SET:
        tokens = {token: 1 for token in document.split(" ")}
        components_list.append(lda.learn_transform_one(tokens))

    for index, component in enumerate(components_list):
        assert np.array_equal(
            a1=list(component.values()), a2=REFERENCE_COMPONENTS_WITH_PRUNNING[index]
        )


def test_learn_transform():
    """
    Assert that learn_one and transform_one methods returns waited output.
    """

    lda = preprocessing.LDA(
        n_components=2,
        number_of_documents=60,
        vocab_prune_interval=2,
        maximum_size_vocabulary=3,
        seed=42,
    )
    components_list = []

    for document in DOC_SET:
        tokens = {token: 1 for token in document.split(" ")}
        lda = lda.learn_one(x=tokens)

        components_list.append(lda.transform_one(x=tokens))

    for index, component in enumerate(components_list):
        assert np.array_equal(
            a1=list(component.values()), a2=REFERENCE_LEARN_ONE_PREDICT_ONE[index]
        )


# river/preprocessing/test_random_projection.py

from __future__ import annotations

import numpy as np

from river import datasets, preprocessing


def test_gaussian_random_projector_dot_product():
    dataset = datasets.TrumpApproval()
    projector = preprocessing.GaussianRandomProjector(n_components=3)

    for x, y in dataset:
        y = projector.transform_one(x)
        y_arr = np.array(list(y.values()))
        x_arr = np.array(list(x.values()))
        P = np.array(
            [[projector._projection_matrix[i, j] for j in x] for i in range(projector.n_components)]
        )
        np.testing.assert_allclose(x_arr @ P.T, y_arr)


def test_sparse_random_projector_dot_product():
    dataset = datasets.TrumpApproval()
    projector = preprocessing.SparseRandomProjector(n_components=3, density=0.5)

    for x, y in dataset:
        y = projector.transform_one(x)
        y_arr = np.array(list(y.values()))
        x_arr = np.array(list(x.values()))
        P = np.array(
            [
                [projector._projection_matrix[j].get(i, 0) for j in x]
                for i in range(projector.n_components)
            ]
        )
        np.testing.assert_allclose(x_arr @ P.T, y_arr)


def test_sparse_random_projector_size():
    dataset = datasets.TrumpApproval()
    projector = preprocessing.SparseRandomProjector(n_components=3, density=0.5)

    for x, y in dataset:
        projector.transform_one(x)
        break

    n_weights = sum(len(v) for v in projector._projection_matrix.values())
    assert n_weights < len(x) * projector.n_components


# river/preprocessing/test_scale.py

from __future__ import annotations

import math

import numpy as np
import pandas as pd

from river import datasets, preprocessing, stream


def test_standard_scaler_one_many_consistent():
    """Checks that using learn_one or learn_many produces the same result."""

    for with_std in (False, True):
        X = pd.read_csv(datasets.TrumpApproval().path)

        one = preprocessing.StandardScaler(with_std=with_std)
        for x, _ in stream.iter_pandas(X):
            one.learn_one(x)

        many = preprocessing.StandardScaler(with_std=with_std)
        for xb in np.array_split(X, 10):
            many.learn_many(xb)

        for i in X:
            assert math.isclose(one.counts[i], many.counts[i])
            assert math.isclose(one.means[i], many.means[i])
            assert math.isclose(one.vars[i], many.vars[i])


def test_standard_scaler_shuffle_columns():
    """Checks that learn_many works identically whether columns are shuffled or not."""

    X = pd.read_csv(datasets.TrumpApproval().path)

    normal = preprocessing.StandardScaler()
    for xb in np.array_split(X, 10):
        normal.learn_many(xb)

    shuffled = preprocessing.StandardScaler()
    for xb in np.array_split(X, 10):
        cols = np.random.permutation(X.columns)
        shuffled.learn_many(xb[cols])

    for i in X:
        assert math.isclose(shuffled.counts[i], shuffled.counts[i])
        assert math.isclose(shuffled.means[i], shuffled.means[i])
        assert math.isclose(shuffled.vars[i], shuffled.vars[i])


def test_standard_scaler_add_remove_columns():
    """Checks that no exceptions are raised whenever columns are dropped and/or added."""

    X = pd.read_csv(datasets.TrumpApproval().path)

    ss = preprocessing.StandardScaler()
    for xb in np.array_split(X, 10):
        # Pick half of the columns at random
        cols = np.random.choice(X.columns, len(X.columns) // 2, replace=False)
        ss.learn_many(xb[cols])


# river/proba/__init__.py

"""Probability distributions."""
from __future__ import annotations

from . import base
from .beta import Beta
from .gaussian import Gaussian
from .multinomial import Multinomial

__all__ = ["base", "Beta", "Gaussian", "Multinomial"]


# river/proba/base.py

from __future__ import annotations

import abc
import random
import typing

from river import base

__all__ = ["Distribution", "BinaryDistribution", "DiscreteDistribution", "ContinuousDistribution"]


class Distribution(abc.ABC, base.Base):
    """General distribution.

    Parameters
    ----------
    seed
        Random number generator seed for reproducibility.

    """

    def __init__(self, seed: int | None = None):
        self.seed = seed
        self._rng = random.Random(seed)

    @abc.abstractmethod
    def __call__(self, x: typing.Any) -> float:
        """Probability mass/density function."""

    @abc.abstractmethod
    def sample(self):
        """Sample a random value from the distribution."""

    @property
    @abc.abstractmethod
    def n_samples(self):
        """The number of observed samples."""

    @property
    @abc.abstractmethod
    def mode(self):
        """The most likely value in the distribution."""

    def __gt__(self, other):
        return self.mode > other.mode


class DiscreteDistribution(Distribution):
    """A probability distribution for discrete values.

    Parameters
    ----------
    seed
        Random number generator seed for reproducibility.

    """

    @abc.abstractmethod
    def update(self, x: typing.Hashable):
        """Updates the parameters of the distribution given a new observation."""

    @abc.abstractmethod
    def revert(self, x: typing.Hashable):
        """Reverts the parameters of the distribution for a given observation."""


class BinaryDistribution(Distribution):
    """A probability distribution for discrete values.

    Parameters
    ----------
    seed
        Random number generator seed for reproducibility.

    """

    @abc.abstractmethod
    def update(self, x: bool):
        """Updates the parameters of the distribution given a new observation."""

    @abc.abstractmethod
    def revert(self, x: bool):
        """Reverts the parameters of the distribution for a given observation."""


class ContinuousDistribution(Distribution):
    """A probability distribution for continuous values.

    Parameters
    ----------
    seed
        Random number generator seed for reproducibility.

    """

    @abc.abstractmethod
    def update(self, x: float):
        """Updates the parameters of the distribution given a new observation."""

    @abc.abstractmethod
    def revert(self, x: float):
        """Reverts the parameters of the distribution for a given observation."""

    @abc.abstractmethod
    def cdf(self, x: float):
        """Cumulative density function, i.e. P(X <= x)."""


# river/proba/beta.py

from __future__ import annotations

import math

import scipy.special

from river.proba import base

__all__ = ["Beta"]


def _beta_func(a, b):
    """

    A naive implementation with (math.gamma(a) + math.gamma(b)) / math.gamma(a + b) would
    overflow for large values of a and b.

    See https://malishoaib.wordpress.com/2014/04/15/the-beautiful-beta-functions-in-raw-python/
    for more details.

    """
    return math.exp(math.lgamma(a) + math.lgamma(b) - math.lgamma(a + b))


class Beta(base.ContinuousDistribution):
    """Beta distribution for binary data.

    A Beta distribution is very similar to a Bernoulli distribution in that it counts occurrences
    of boolean events. The differences lies in what is being measured. A Binomial distribution
    models the probability of an event occurring, whereas a Beta distribution models the
    probability distribution itself. In other words, it's a probability distribution over
    probability distributions.

    Parameters
    ----------
    alpha
        Initial alpha parameter.
    beta
        Initial beta parameter.
    seed
        Random number generator seed for reproducibility.

    Examples
    --------

    >>> from river import proba

    >>> successes = 81
    >>> failures = 219
    >>> beta = proba.Beta(successes, failures)

    >>> beta(.21), beta(.35)
    (0.867..., 0.165...)

    >>> for success in range(100):
    ...     beta = beta.update(True)
    >>> for failure in range(200):
    ...     beta = beta.update(False)

    >>> beta(.21), beta(.35)
    (2.525...e-05, 0.841...)

    >>> beta.cdf(.35)
    0.994168...

    References
    ----------
    [^1]: [What is the intuition behind beta distribution?](https://stats.stackexchange.com/questions/47771/what-is-the-intuition-behind-beta-distribution)

    """

    def __init__(self, alpha: int = 1, beta: int = 1, seed: int | None = None):
        super().__init__(seed)
        self.alpha = alpha
        self.beta = beta
        self._alpha = alpha
        self._beta = beta

    @property
    def n_samples(self):
        return self._alpha - self.alpha + self._beta - self.beta

    def update(self, x):
        if x:
            self.alpha += 1
        else:
            self.beta += 1
        return self

    def revert(self, x):
        if x:
            self.alpha -= 1
        else:
            self.beta -= 1
        return self

    def __call__(self, p: float):
        return (
            p ** (self.alpha - 1) * (1 - p) ** (self.beta - 1) / _beta_func(self.alpha, self.beta)
        )

    def sample(self):
        return self._rng.betavariate(self.alpha, self.beta)

    @property
    def mode(self):
        try:
            return (self.alpha - 1) / (self.alpha + self.beta - 2)
        except ZeroDivisionError:
            return 0.5

    def cdf(self, x):
        return scipy.special.betainc(self.alpha, self.beta, x)


# river/proba/gaussian.py

from __future__ import annotations

import math

from river import stats
from river.proba import base

__all__ = ["Gaussian"]


class Gaussian(base.ContinuousDistribution):
    """Normal distribution with parameters mu and sigma.

    Parameters
    ----------
    seed
        Random number generator seed for reproducibility.

    Examples
    --------

    >>> from river import proba

    >>> p = proba.Gaussian().update(6).update(7)

    >>> p
    𝒩(μ=6.500, σ=0.707)

    >>> p(6.5)
    0.564189

    >>> p.revert(7)
    𝒩(μ=6.000, σ=0.000)

    """

    def __init__(self, seed=None):
        super().__init__(seed)
        self._var = stats.Var(ddof=1)

    @classmethod
    def _from_state(cls, n, m, sig, ddof):
        new = cls()
        new._var = stats.Var._from_state(n, m, sig, ddof=ddof)
        return new

    @property
    def n_samples(self):
        return self._var.mean.n

    @property
    def mu(self):
        return self._var.mean.get()

    @property
    def sigma(self):
        return self._var.get() ** 0.5

    def __repr__(self):
        return f"𝒩(μ={self.mu:.3f}, σ={self.sigma:.3f})"

    def update(self, x, w=1.0):
        self._var.update(x, w)
        return self

    def revert(self, x, w=1.0):
        self._var.revert(x, w)
        return self

    def __call__(self, x):
        var = self._var.get()
        if var:
            try:
                return math.exp((x - self.mu) ** 2 / (-2 * var)) / math.sqrt(math.tau * var)
            except ValueError:
                return 0.0
            except OverflowError:
                return 0.0
        return 0.0

    def cdf(self, x):
        try:
            return 0.5 * (1.0 + math.erf((x - self.mu) / (self.sigma * math.sqrt(2.0))))
        except ZeroDivisionError:
            return 0.0

    def sample(self):
        return self._rng.gauss(self.mu, self.sigma)

    @property
    def mode(self):
        return self.mu


# river/proba/multinomial.py

from __future__ import annotations

import collections
import typing

from river.proba import base

__all__ = ["Multinomial"]


class Multinomial(base.DiscreteDistribution):
    """Multinomial distribution for categorical data.

    Parameters
    ----------
    events
        An optional list of events that already occurred.
    seed
        Random number generator seed for reproducibility.

    Examples
    --------

    >>> from river import proba

    >>> p = proba.Multinomial(['green'] * 3)
    >>> p = p.update('red')

    >>> p('red')
    0.25

    >>> p = p.update('red').update('red')
    >>> p('green')
    0.5

    >>> p = p.revert('red').revert('red')
    >>> p('red')
    0.25

    You can wrap this with a `utils.Rolling` to measure a distribution over a window:

    >>> from river import utils

    >>> X = ['red', 'green', 'green', 'blue', 'blue']

    >>> dist = utils.Rolling(
    ...     proba.Multinomial(),
    ...     window_size=3
    ... )

    >>> for x in X:
    ...     dist = dist.update(x)
    ...     print(dist)
    ...     print()
    P(red) = 1.000
    <BLANKLINE>
    P(red) = 0.500
    P(green) = 0.500
    <BLANKLINE>
    P(green) = 0.667
    P(red) = 0.333
    <BLANKLINE>
    P(green) = 0.667
    P(blue) = 0.333
    P(red) = 0.000
    <BLANKLINE>
    P(blue) = 0.667
    P(green) = 0.333
    P(red) = 0.000
    <BLANKLINE>

    You can wrap this with a `utils.Rolling` to measure a distribution over a window of time:

    >>> import datetime as dt

    >>> X = ['red', 'green', 'green', 'blue']
    >>> days = [1, 2, 3, 4]

    >>> dist = utils.TimeRolling(
    ...     proba.Multinomial(),
    ...     period=dt.timedelta(days=2)
    ... )

    >>> for x, day in zip(X, days):
    ...     dist = dist.update(x, t=dt.datetime(2019, 1, day))
    ...     print(dist)
    ...     print()
    P(red) = 1.000
    <BLANKLINE>
    P(red) = 0.500
    P(green) = 0.500
    <BLANKLINE>
    P(green) = 1.000
    P(red) = 0.000
    <BLANKLINE>
    P(green) = 0.500
    P(blue) = 0.500
    P(red) = 0.000
    <BLANKLINE>

    """

    def __init__(self, events: dict | list | None = None, seed=None):
        super().__init__(seed)
        self.events = events
        self.counts: typing.Counter[typing.Any] = collections.Counter(events)
        self._n = sum(self.counts.values())

    @property
    def n_samples(self):
        return self._n

    def __iter__(self):
        return iter(self.counts)

    def __len__(self):
        return len(self.counts)

    @property
    def mode(self):
        return self.counts.most_common(1)[0][0]

    def update(self, x):
        self.counts.update([x])
        self._n += 1
        return self

    def revert(self, x):
        self.counts.subtract([x])
        self._n -= 1
        return self

    def sample(self):
        return self._rng.choices(list(self.counts.keys()), weights=list(self.counts.values()))[0]

    def __call__(self, x):
        try:
            return self.counts[x] / self._n
        except ZeroDivisionError:
            return 0.0

    def __repr__(self):
        return "\n".join(f"P({c}) = {self(c):.3f}" for c, _ in self.counts.most_common())


# river/reco/__init__.py

"""Recommender systems module.

Recommender systems (recsys for short) is a large topic. This module is far from comprehensive. It
simply provides models which can contribute towards building a recommender system.

A typical recommender system is made up of a retrieval phase, followed by a ranking phase. The
output of the retrieval phase is a shortlist of the catalogue of items. The items in the shortlist
are then usually ranked according to the expected preference the user will have for each item. This
module focuses on the ranking phase.

Models which inherit from the `Ranker` class have a `rank` method. This allows sorting a set of
items for a given user. Each model also has a `learn_one(user, item, y, context)` which allows
learning user preferences. The `y` parameter is a reward value, the nature of which depends is
specific to each and every recommendation task. Typically the reward is a number or a boolean
value. It is up to the user to determine how to translate a user session into training data.

"""
from __future__ import annotations

from . import base
from .baseline import Baseline
from .biased_mf import BiasedMF
from .funk_mf import FunkMF
from .normal import RandomNormal

__all__ = ["base", "Baseline", "BiasedMF", "FunkMF", "RandomNormal"]


# river/reco/base.py

from __future__ import annotations

import abc
import numbers
import random
import typing

from river import base

ID = typing.Union[str, int]  # noqa: UP007
Reward = typing.Union[numbers.Number, bool]  # noqa: UP007


__all__ = ["Ranker"]


class Ranker(base.Estimator):
    """Base class for ranking models.

    Parameters
    ----------
    seed
        Random number generation seed. Set this for reproducibility.

    """

    def __init__(self, seed: int | None = None):
        self.seed = seed
        self._rng = random.Random(seed)

    @property
    def is_contextual(self):
        return False

    @abc.abstractmethod
    def learn_one(self, user: ID, item: ID, y: Reward, x: dict | None = None):
        """Fits a `user`-`item` pair and a real-valued target `y`.

        Parameters
        ----------
        user
            A user ID.
        item
            An item ID.
        y
            Reward feedback from the user for the item. This may be a boolean or a number.
        x
            Optional context to use.

        """

    @abc.abstractmethod
    def predict_one(self, user: ID, item: ID, x: dict | None = None) -> Reward:
        """Predicts the target value of a set of features `x`.

        Parameters
        ----------
        user
            A user ID.
        item
            An item ID.
        x
            Optional context to use.

        Returns
        -------
        The predicted preference from the user for the item.

        """

    def rank(self, user: ID, items: set[ID], x: dict | None = None) -> list[ID]:
        """Rank models by decreasing order of preference for a given user.

        Parameters
        ----------
        user
            A user ID.
        items
            A set of items to rank.
        x
            Optional context to use.

        """
        preferences = {item: self.predict_one(user, item, x) for item in items}
        return sorted(preferences, key=preferences.__getitem__, reverse=True)  # type: ignore[arg-type]

    def _unit_test_skips(self):
        return {"check_emerging_features", "check_disappearing_features"}


# river/reco/baseline.py

from __future__ import annotations

import collections
import copy

from river import optim, reco, stats, utils

__all__ = ["Baseline"]


class Baseline(reco.base.Ranker):
    """Baseline for recommender systems.

    A first-order approximation of the bias involved in target. The model equation is defined as:

    $$\\hat{y}(x) = \\bar{y} + bu_{u} + bi_{i}$$

    Where $bu_{u}$ and $bi_{i}$ are respectively the user and item biases.

    This model expects a dict input with a `user` and an `item` entries without any type constraint
    on their values (i.e. can be strings or numbers). Other entries are ignored.

    Parameters
    ----------
    optimizer
        The sequential optimizer used for updating the weights.
    loss
        The loss function to optimize for.
    l2
        regularization amount used to push weights towards 0.
    initializer
        Weights initialization scheme.
    clip_gradient
        Clips the absolute value of each gradient value.
    seed
        Random number generation seed. Set this for reproducibility.

    Attributes
    ----------
    global_mean : stats.Mean
        The target arithmetic mean.
    u_biases : collections.defaultdict
        The user bias weights.
    i_biases : collections.defaultdict
        The item bias weights.
    u_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the user bias weights.
    i_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the item bias weights.

    Examples
    --------

    >>> from river import optim
    >>> from river import reco

    >>> dataset = (
    ...     ({'user': 'Alice', 'item': 'Superman'}, 8),
    ...     ({'user': 'Alice', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Alice', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Alice', 'item': 'Notting Hill'}, 2),
    ...     ({'user': 'Alice', 'item': 'Harry Potter'}, 5),
    ...     ({'user': 'Bob', 'item': 'Superman'}, 8),
    ...     ({'user': 'Bob', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Bob', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Bob', 'item': 'Notting Hill'}, 2)
    ... )

    >>> model = reco.Baseline(optimizer=optim.SGD(0.005))

    >>> for x, y in dataset:
    ...     _ = model.learn_one(**x, y=y)

    >>> model.predict_one(user='Bob', item='Harry Potter')
    6.538120

    References
    ----------
    [^1]: [Matrix factorization techniques for recommender systems](https://datajobs.com/data-science-repo/Recommender-Systems-[Netflix].pdf)

    """

    def __init__(
        self,
        optimizer: optim.base.Optimizer | None = None,
        loss: optim.losses.Loss | None = None,
        l2=0.0,
        initializer: optim.initializers.Initializer | None = None,
        clip_gradient=1e12,
        seed=None,
    ):
        super().__init__(seed=seed)

        self.optimizer = optim.SGD() if optimizer is None else copy.deepcopy(optimizer)
        self.u_optimizer = optim.SGD() if optimizer is None else copy.deepcopy(optimizer)
        self.i_optimizer = optim.SGD() if optimizer is None else copy.deepcopy(optimizer)
        self.loss = optim.losses.Squared() if loss is None else loss
        self.l2 = l2

        if initializer is None:
            initializer = optim.initializers.Zeros()
        self.initializer = initializer

        self.clip_gradient = clip_gradient
        self.global_mean = stats.Mean()
        self.u_biases: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(initializer)
        self.i_biases: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(initializer)

    def predict_one(self, user, item, x=None):
        return self.global_mean.get() + self.u_biases[user] + self.i_biases[item]

    def learn_one(self, user, item, y, x=None):
        # Update the global mean
        self.global_mean.update(y)

        # Calculate the gradient of the loss with respect to the prediction
        g_loss = self.loss.gradient(y, self.predict_one(user, item))

        # Clamp the gradient to avoid numerical instability
        g_loss = utils.math.clamp(g_loss, minimum=-self.clip_gradient, maximum=self.clip_gradient)

        # Calculate bias gradients
        u_grad_bias = {user: g_loss + self.l2 * self.u_biases[user]}
        i_grad_bias = {item: g_loss + self.l2 * self.i_biases[item]}

        # Update biases
        self.u_biases = self.u_optimizer.step(self.u_biases, u_grad_bias)
        self.i_biases = self.i_optimizer.step(self.i_biases, i_grad_bias)

        return self


# river/reco/biased_mf.py

from __future__ import annotations

import collections
import copy
import functools

import numpy as np

from river import optim, stats, utils

from .base import Ranker

__all__ = ["BiasedMF"]


class BiasedMF(Ranker):
    """Biased Matrix Factorization for recommender systems.

    The model equation is defined as:

    $$\\hat{y}(x) = \\bar{y} + bu_{u} + bi_{i} + \\langle \\mathbf{v}_u, \\mathbf{v}_i \\rangle$$

    Where $bu_{u}$ and $bi_{i}$ are respectively the user and item biases. The last term being
    simply the dot product between the latent vectors of the given user-item pair:

    $$\\langle \\mathbf{v}_u, \\mathbf{v}_i \\rangle = \\sum_{f=1}^{k} \\mathbf{v}_{u, f} \\cdot \\mathbf{v}_{i, f}$$

    where $k$ is the number of latent factors.

    This model expects a dict input with a `user` and an `item` entries without any type constraint
    on their values (i.e. can be strings or numbers). Other entries are ignored.

    Parameters
    ----------
    n_factors
        Dimensionality of the factorization or number of latent factors.
    bias_optimizer
        The sequential optimizer used for updating the bias weights.
    latent_optimizer
        The sequential optimizer used for updating the latent weights.
    loss
        The loss function to optimize for.
    l2_bias
        Amount of L2 regularization used to push bias weights towards 0.
    l2_latent
        Amount of L2 regularization used to push latent weights towards 0.
    weight_initializer
        Weights initialization scheme.
    latent_initializer
        Latent factors initialization scheme.
    clip_gradient
        Clips the absolute value of each gradient value.
    seed
        Random number generation seed. Set this for reproducibility.

    Attributes
    ----------
    global_mean : stats.Mean
        The target arithmetic mean.
    u_biases : collections.defaultdict
        The user bias weights.
    i_biases : collections.defaultdict
        The item bias weights.
    u_latents : collections.defaultdict
        The user latent vectors randomly initialized.
    i_latents : collections.defaultdict
        The item latent vectors randomly initialized.
    u_bias_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the user bias weights.
    i_bias_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the item bias weights.
    u_latent_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the user latent weights.
    i_latent_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the item latent weights.

    Examples
    --------

    >>> from river import optim
    >>> from river import reco

    >>> dataset = (
    ...     ({'user': 'Alice', 'item': 'Superman'}, 8),
    ...     ({'user': 'Alice', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Alice', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Alice', 'item': 'Notting Hill'}, 2),
    ...     ({'user': 'Alice', 'item': 'Harry Potter'}, 5),
    ...     ({'user': 'Bob', 'item': 'Superman'}, 8),
    ...     ({'user': 'Bob', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Bob', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Bob', 'item': 'Notting Hill'}, 2)
    ... )

    >>> model = reco.BiasedMF(
    ...     n_factors=10,
    ...     bias_optimizer=optim.SGD(0.025),
    ...     latent_optimizer=optim.SGD(0.025),
    ...     latent_initializer=optim.initializers.Normal(mu=0., sigma=0.1, seed=71)
    ... )

    >>> for x, y in dataset:
    ...     _ = model.learn_one(**x, y=y)

    >>> model.predict_one(user='Bob', item='Harry Potter')
    6.489025

    References
    ----------
    [^1]: [Paterek, A., 2007, August. Improving regularized singular value decomposition for collaborative filtering. In Proceedings of KDD cup and workshop (Vol. 2007, pp. 5-8)](https://www.cs.uic.edu/~liub/KDD-cup-2007/proceedings/Regular-Paterek.pdf)
    [^2]: [Matrix factorization techniques for recommender systems](https://datajobs.com/data-science-repo/Recommender-Systems-[Netflix].pdf)

    """

    def __init__(
        self,
        n_factors=10,
        bias_optimizer: optim.base.Optimizer | None = None,
        latent_optimizer: optim.base.Optimizer | None = None,
        loss: optim.losses.Loss | None = None,
        l2_bias=0.0,
        l2_latent=0.0,
        weight_initializer: optim.initializers.Initializer | None = None,
        latent_initializer: optim.initializers.Initializer | None = None,
        clip_gradient=1e12,
        seed=None,
    ):
        super().__init__(seed=seed)

        self.bias_optimizer = bias_optimizer
        self.latent_optimizer = latent_optimizer

        self.n_factors = n_factors
        self.u_bias_optimizer = (
            optim.SGD() if bias_optimizer is None else copy.deepcopy(bias_optimizer)
        )
        self.i_bias_optimizer = (
            optim.SGD() if bias_optimizer is None else copy.deepcopy(bias_optimizer)
        )
        self.u_latent_optimizer = (
            optim.SGD() if latent_optimizer is None else copy.deepcopy(latent_optimizer)
        )
        self.i_latent_optimizer = (
            optim.SGD() if latent_optimizer is None else copy.deepcopy(latent_optimizer)
        )
        self.loss = optim.losses.Squared() if loss is None else loss
        self.l2_bias = l2_bias
        self.l2_latent = l2_latent

        if weight_initializer is None:
            weight_initializer = optim.initializers.Zeros()
        self.weight_initializer = weight_initializer

        if latent_initializer is None:
            latent_initializer = optim.initializers.Normal(sigma=0.1, seed=self.seed)
        self.latent_initializer = latent_initializer

        self.clip_gradient = clip_gradient
        self.global_mean = stats.Mean()

        self.u_biases: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(weight_initializer)
        self.i_biases: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(weight_initializer)

        random_latents = functools.partial(self.latent_initializer, shape=self.n_factors)
        self.u_latents: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(random_latents)
        self.i_latents: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(random_latents)

    @property
    def _mutable_attributes(self):
        return {
            "bias_optimizer",
            "latent_optimizer",
            "loss",
            "l2_bias",
            "l2_latent",
            "weight_initializer",
            "weight_initializer",
            "clip_gradient",
        }

    def predict_one(self, user, item, x=None):
        # Initialize the prediction to the mean
        y_pred = self.global_mean.get()

        # Add the user bias
        y_pred += self.u_biases[user]

        # Add the item bias
        y_pred += self.i_biases[item]

        # Add the dot product of the user and the item latent vectors
        y_pred += np.dot(self.u_latents[user], self.i_latents[item])

        return y_pred

    def learn_one(self, user, item, y, x=None):
        # Update the global mean
        self.global_mean.update(y)

        # Calculate the gradient of the loss with respect to the prediction
        g_loss = self.loss.gradient(y, self.predict_one(user, item))

        # Clamp the gradient to avoid numerical instability
        g_loss = utils.math.clamp(g_loss, minimum=-self.clip_gradient, maximum=self.clip_gradient)

        # Calculate weights gradients
        u_grad_bias = {user: g_loss + self.l2_bias * self.u_biases[user]}
        i_grad_bias = {item: g_loss + self.l2_bias * self.i_biases[item]}
        u_latent_grad = {
            user: g_loss * self.i_latents[item] + self.l2_latent * self.u_latents[user]
        }
        i_latent_grad = {
            item: g_loss * self.u_latents[user] + self.l2_latent * self.i_latents[item]
        }

        # Update weights
        self.u_biases = self.u_bias_optimizer.step(self.u_biases, u_grad_bias)
        self.i_biases = self.i_bias_optimizer.step(self.i_biases, i_grad_bias)
        self.u_latents = self.u_latent_optimizer.step(self.u_latents, u_latent_grad)
        self.i_latents = self.i_latent_optimizer.step(self.i_latents, i_latent_grad)

        return self


# river/reco/funk_mf.py

from __future__ import annotations

import collections
import copy
import functools

import numpy as np

from river import optim, reco, utils

__all__ = ["FunkMF"]


class FunkMF(reco.base.Ranker):
    """Funk Matrix Factorization for recommender systems.

    The model equation is defined as:

    $$\\hat{y}(x) = \\langle \\mathbf{v}_u, \\mathbf{v}_i \\rangle = \\sum_{f=1}^{k} \\mathbf{v}_{u, f} \\cdot \\mathbf{v}_{i, f}$$

    where $k$ is the number of latent factors.

    This model expects a dict input with a `user` and an `item` entries without any type
    constraint on their values (i.e. can be strings or numbers). Other entries are ignored.

    Parameters
    ----------
    n_factors
        Dimensionality of the factorization or number of latent factors.
    optimizer
        The sequential optimizer used for updating the latent factors.
    loss
        The loss function to optimize for.
    l2
        Amount of L2 regularization used to push weights towards 0.
    initializer
        Latent factors initialization scheme.
    clip_gradient
        Clips the absolute value of each gradient value.
    seed
        Random number generation seed. Set this for reproducibility.

    Attributes
    ----------
    u_latents : collections.defaultdict
        The user latent vectors randomly initialized.
    i_latents : collections.defaultdict
        The item latent vectors randomly initialized.
    u_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the user latent weights.
    i_optimizer : optim.base.Optimizer
        The sequential optimizer used for updating the item latent weights.

    Examples
    --------

    >>> from river import optim
    >>> from river import reco

    >>> dataset = (
    ...     ({'user': 'Alice', 'item': 'Superman'}, 8),
    ...     ({'user': 'Alice', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Alice', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Alice', 'item': 'Notting Hill'}, 2),
    ...     ({'user': 'Alice', 'item': 'Harry Potter'}, 5),
    ...     ({'user': 'Bob', 'item': 'Superman'}, 8),
    ...     ({'user': 'Bob', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Bob', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Bob', 'item': 'Notting Hill'}, 2)
    ... )

    >>> model = reco.FunkMF(
    ...     n_factors=10,
    ...     optimizer=optim.SGD(0.1),
    ...     initializer=optim.initializers.Normal(mu=0., sigma=0.1, seed=11),
    ... )

    >>> for x, y in dataset:
    ...     _ = model.learn_one(**x, y=y)

    >>> model.predict_one(user='Bob', item='Harry Potter')
    1.866272

    References
    ----------
    [^1]: [Netflix update: Try this at home](https://sifter.org/simon/journal/20061211.html)
    [^2]: [Matrix factorization techniques for recommender systems](https://datajobs.com/data-science-repo/Recommender-Systems-[Netflix].pdf)

    """

    def __init__(
        self,
        n_factors=10,
        optimizer: optim.base.Optimizer | None = None,
        loss: optim.losses.Loss | None = None,
        l2=0.0,
        initializer: optim.initializers.Initializer | None = None,
        clip_gradient=1e12,
        seed=None,
    ):
        super().__init__(seed=seed)

        self.optimizer = optimizer

        self.n_factors = n_factors
        self.u_optimizer = optim.SGD(0.1) if optimizer is None else copy.deepcopy(optimizer)
        self.i_optimizer = optim.SGD(0.1) if optimizer is None else copy.deepcopy(optimizer)
        self.loss = optim.losses.Squared() if loss is None else loss
        self.l2 = l2

        if initializer is None:
            initializer = optim.initializers.Normal(mu=0.0, sigma=0.1, seed=self.seed)
        self.initializer = initializer

        self.clip_gradient = clip_gradient

        random_latents = functools.partial(self.initializer, shape=self.n_factors)
        self.u_latents: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(random_latents)
        self.i_latents: collections.defaultdict[
            int, optim.initializers.Initializer
        ] = collections.defaultdict(random_latents)

    @property
    def _mutable_attributes(self):
        return {"optimizer", "l2", "loss", "clip_gradient", "initializer"}

    def predict_one(self, user, item, x=None):
        return np.dot(self.u_latents[user], self.i_latents[item])

    def learn_one(self, user, item, y, x=None):
        # Calculate the gradient of the loss with respect to the prediction
        g_loss = self.loss.gradient(y, self.predict_one(user, item))

        # Clamp the gradient to avoid numerical instability
        g_loss = utils.math.clamp(g_loss, minimum=-self.clip_gradient, maximum=self.clip_gradient)

        # Calculate latent gradients
        u_latent_grad = {user: g_loss * self.i_latents[item] + self.l2 * self.u_latents[user]}
        i_latent_grad = {item: g_loss * self.u_latents[user] + self.l2 * self.i_latents[item]}

        # Update latent weights
        self.u_latents = self.u_optimizer.step(self.u_latents, u_latent_grad)
        self.i_latents = self.i_optimizer.step(self.i_latents, i_latent_grad)

        return self


# river/reco/normal.py

from __future__ import annotations

from river import reco, stats

__all__ = ["RandomNormal"]


class RandomNormal(reco.base.Ranker):
    """Predicts random values sampled from a normal distribution.

    The parameters of the normal distribution are fitted with running statistics. They parameters
    are independent of the user, the item, or the context, and are instead fitted globally. This
    recommender therefore acts as a dummy model that any serious model should easily outperform.

    Parameters
    ----------
    seed
        Random number generation seed. Set this for reproducibility.

    Attributes
    ----------
    mean
        stats.Mean
    variance
        stats.Var

    Examples
    --------

    >>> from river import reco

    >>> dataset = (
    ...     ({'user': 'Alice', 'item': 'Superman'}, 8),
    ...     ({'user': 'Alice', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Alice', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Alice', 'item': 'Notting Hill'}, 2),
    ...     ({'user': 'Alice', 'item': 'Harry Potter'}, 5),
    ...     ({'user': 'Bob', 'item': 'Superman'}, 8),
    ...     ({'user': 'Bob', 'item': 'Terminator'}, 9),
    ...     ({'user': 'Bob', 'item': 'Star Wars'}, 8),
    ...     ({'user': 'Bob', 'item': 'Notting Hill'}, 2)
    ... )

    >>> model = reco.RandomNormal(seed=42)

    >>> for x, y in dataset:
    ...     _ = model.learn_one(**x, y=y)

    >>> model.predict_one(user='Bob', item='Harry Potter')
    6.147299621751425

    """

    def __init__(self, seed=None):
        super().__init__(seed=seed)
        self.variance = stats.Var()
        self.mean = stats.Mean()
        self.seed = seed

    def learn_one(self, user, item, y, x=None):
        self.mean.update(y)
        self.variance.update(y)
        return self

    def predict_one(self, user, item, x=None):
        μ = self.mean.get() or 0
        σ = (self.variance.get() or 1) ** 0.5
        return self._rng.gauss(μ, σ)


# river/rules/__init__.py

"""Decision rules-based algorithms."""
from __future__ import annotations

from .amrules import AMRules

__all__ = ["AMRules"]


# river/rules/amrules.py

from __future__ import annotations

import collections
import functools
import io
import math
import typing

from river import anomaly, base, drift, linear_model, stats
from river.tree import split_criterion
from river.tree import splitter as spl

from .base import HoeffdingRule


class MeanRegressor(base.Regressor):
    def __init__(self):
        self.mean = stats.Mean()

    def learn_one(self, x: dict, y: base.typing.RegTarget, w: int = 1):
        self.mean.update(y, w)
        return self

    def predict_one(self, x: dict):
        return self.mean.get()


class AdaptiveRegressor(base.Regressor):
    """This predictor selects between the target mean and the the user-selected regression model
     automatically. Faded error metrics are kept for both the predictors, and the most
    accurate one is selected automatically for each incoming instance.
    """

    def __init__(self, model_predictor: base.Regressor, fading_factor: float):
        self.model_predictor = model_predictor
        self.mean_predictor = MeanRegressor()
        self.fading_factor = fading_factor

        self._mae_mean = 0.0
        self._mae_model = 0.0

    def learn_one(self, x: dict, y: base.typing.RegTarget, w: int = 1):
        abs_error_mean = abs(y - self.mean_predictor.predict_one(x))  # type: ignore
        abs_error_model = abs(y - self.model_predictor.predict_one(x))  # type: ignore

        self._mae_mean = self.fading_factor * self._mae_mean + abs_error_mean
        self._mae_model = self.fading_factor * self._mae_model + abs_error_model

        self.mean_predictor.learn_one(x, y, w)

        for _ in range(int(w)):
            self.model_predictor.learn_one(x, y)

        return self

    def predict_one(self, x: dict):
        if self._mae_mean <= self._mae_model:
            return self.mean_predictor.predict_one(x)
        else:
            return self.model_predictor.predict_one(x)


class RegRule(HoeffdingRule, base.Regressor, anomaly.base.AnomalyDetector):
    def __init__(
        self,
        template_splitter,
        split_criterion,
        pred_model,
        drift_detector,
    ):
        super().__init__(
            template_splitter=template_splitter,
            split_criterion=split_criterion,
        )
        self.pred_model = pred_model
        self.drift_detector = drift_detector

        self._target_stats = stats.Var()
        self._feat_stats = collections.defaultdict(functools.partial(stats.Var))

    def new_nominal_splitter(self):
        return spl.NominalSplitterReg()

    @property
    def statistics(self):
        return self._target_stats

    @statistics.setter
    def statistics(self, target_stats):
        self._target_stats = target_stats

    def _update_target_stats(self, y, w):
        self._target_stats.update(y, w)

    def _update_feature_stats(self, feat_name, feat_val, w):
        if feat_name not in self.nominal_features:
            self._feat_stats[feat_name].update(feat_val, w)

    def drift_test(self, y, y_pred):
        abs_error = abs(y - y_pred)
        self.drift_detector.update(abs_error)

        return self.drift_detector.drift_detected

    def score_one(self, x) -> float:
        """Rule anomaly score.

        The more negative the score, the more anomalous is the instance regarding the subspace
        covered by the rule. Instances whose score is greater than zero are considered normal.
        Different threshold values can be used to decide upon discarding or using the instance
        for training. A small negative threshold is considered to be a conservative choice,
        whereas cutting at zero is a more aggressive choice.

        Parameters
        ----------
        x
            A dictionary of features.

        Returns
        -------
        An anomaly score. The more negative the score, the more anomalous is the instance.


        """
        score = 0.0
        hits = 0

        for feat_name, feat_val in x.items():
            # Skip nominal features
            if feat_name in self.nominal_features:
                continue

            mean = self._feat_stats[feat_name].mean.get()
            var = self._feat_stats[feat_name].get()

            if var > 0:
                # One tailed variant of Chebyshev's inequality: ported from the MOA code
                proba = (2 * var) / (var + (feat_val - mean) ** 2)

                if 0 < proba < 1:
                    score += math.log(proba) - math.log(1 - proba)
                    hits += 1

        return score / hits if hits > 0 else 0.0

    def learn_one(self, x: dict, y: base.typing.RegTarget, w: int = 1):  # type: ignore
        self.update(x, y, w)
        self.pred_model.learn_one(x, y, w)

        return self

    def predict_one(self, x: dict):
        return self.pred_model.predict_one(x)


class AMRules(base.Regressor):
    """Adaptive Model Rules.

    AMRules[^1] is a rule-based algorithm for incremental regression tasks. AMRules relies on the
    Hoeffding bound to build its rule set, similarly to Hoeffding Trees. The Variance-Ratio
    heuristic is used to evaluate rules' splits. Moreover, this rule-based regressor has
    additional capacities not usually found in decision trees.

    Firstly, each created decision rule has a built-in drift detection mechanism. Every time a
    drift is detected, the affected decision rule is removed. In addition, AMRules' rules also
    have anomaly detection capabilities. After a warm-up period, each rule tests whether or not
    the incoming instances are anomalies. Anomalous instances are not used for training.

    Every time no rule is covering an incoming example, a default rule is used to learn
    from it. A rule covers an instance when all of the rule's literals (tests joined by the
    logical operation `and`) match the input case. The default rule is also applied for predicting
    examples not covered by any rules from the rule set.

    Parameters
    ----------
    n_min
        The total weight that must be observed by a rule between expansion attempts.
    delta
        The split test significance. The split confidence is given by `1 - delta`.
    tau
        The tie-breaking threshold.
    pred_type
        The prediction strategy used by the decision rules. Can be either:</br>
        - `"mean"`: outputs the target mean within the partitions defined by the decision
        rules.</br>
        - `"model"`: always use instances of the model passed `pred_model` to make
        predictions.</br>
        - `"adaptive"`: dynamically selects between "mean" and "model" for each incoming example.
        The most accurate option at the moment will be used.
    pred_model
        The regression model that will be replicated for every rule when `pred_type` is either
        `"model"` or `"adaptive"`.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.TEBSTSplitter` is used if `splitter` is `None`.
    drift_detector
        The drift detection model that is used by each rule. Care must be taken to avoid the
        triggering of too many false alarms or delaying too much the concept drift detection.
        By default, `drift.ADWIN` is used if `drift_detector` is `None`.
    fading_factor
        The exponential decaying factor applied to the learning models' absolute errors, that
        are monitored if `pred_type='adaptive'`. Must be between `0` and `1`. The closer
        to `1`, the more importance is going to be given to past observations. On the other hand,
        if its value approaches `0`, the recent observed errors are going to have more influence
        on the final decision.
    anomaly_threshold
        The threshold below which instances will be considered anomalies by the rules.
    m_min
        The minimum total weight a rule must observe before it starts to skip anomalous
        instances during training.
    ordered_rule_set
        If `True`, only the first rule that covers an instance will be used for training or
        prediction. If `False`, all the rules covering an instance will be updated during
        training, and the predictions for an instance will be the average prediction of all
        rules covering that example.
    min_samples_split
        The minimum number of samples each partition of a binary split candidate must have
        to be considered valid.

    Notes
    -----
    AMRules treats all the non-numerical inputs as nominal features. All instances of
    `numbers.Number` will be treated as continuous, even if they represent integer categories.
    When using nominal features, `pred_type` should be set to "mean", otherwise errors will be
    thrown while trying to update the underlying rules' prediction models. Prediction strategies
    other than "mean" can be used, as long as the prediction model passed to `pred_model` supports
    nominal features.

    Raises
    ------
    TypeError
        If one or more input features are non-numeric and the selected `pred_model` is either
        "model" or "adaptive".

    Examples
    --------
    >>> from river import datasets
    >>> from river import drift
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import preprocessing
    >>> from river import rules

    >>> dataset = datasets.TrumpApproval()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     rules.AMRules(
    ...         delta=0.01,
    ...         n_min=50,
    ...         drift_detector=drift.ADWIN()
    ...     )
    ... )

    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 1.117705

    References
    ----------
    [^1]: Duarte, J., Gama, J. and Bifet, A., 2016. Adaptive model rules from high-speed data
    streams. ACM Transactions on Knowledge Discovery from Data (TKDD), 10(3), pp.1-22.

    """

    _PRED_MEAN = "mean"
    _PRED_MODEL = "model"
    _PRED_ADAPTIVE = "adaptive"
    _VALID_PRED = [_PRED_MEAN, _PRED_MODEL, _PRED_ADAPTIVE]

    def __init__(
        self,
        n_min: int = 200,
        delta: float = 1e-7,
        tau: float = 0.05,
        pred_type: str = "adaptive",
        pred_model: base.Regressor | None = None,
        splitter: spl.Splitter | None = None,
        drift_detector: base.DriftDetector | None = None,
        fading_factor: float = 0.99,
        anomaly_threshold: float = -0.75,
        m_min: int = 30,
        ordered_rule_set: bool = True,
        min_samples_split: int = 5,
    ):
        self.n_min = n_min
        self.delta = delta
        self.tau = tau

        if pred_type not in self._VALID_PRED:
            raise ValueError(f"Invalid 'pred_type': {pred_type}")
        self.pred_type = pred_type
        self.pred_model = pred_model if pred_model else linear_model.LinearRegression()

        if splitter is None:
            self.splitter = spl.TEBSTSplitter()
        else:
            self.splitter = splitter  # type: ignore

        self.drift_detector = drift_detector if drift_detector is not None else drift.ADWIN()

        self.fading_factor = fading_factor
        self.anomaly_threshold = anomaly_threshold
        self.m_min = m_min
        self.ordered_rule_set = ordered_rule_set
        self.min_samples_split = min_samples_split

        self._default_rule = self._new_rule()
        self._rules: dict[typing.Hashable, RegRule] = {}

        self._n_drifts_detected: int = 0

    @property
    def _mutable_attributes(self):
        return {
            "n_min",
            "delta",
            "tau",
            "fading_factor",
            "anomaly_threshold",
            "m_min",
            "ordered_rule_set",
        }

    def __len__(self):
        return len(self._rules) + 1

    def __getitem__(self, item):
        return list(self._rules.values())[item]

    @property
    def n_drifts_detected(self) -> int:
        """The number of detected concept drifts."""
        return self._n_drifts_detected

    def _new_rule(self) -> RegRule:
        if self.pred_type == self._PRED_MEAN:
            predictor = MeanRegressor()
        elif self.pred_type == self._PRED_MODEL:
            predictor = self.pred_model.clone()
        else:  # adaptive predictor
            predictor = AdaptiveRegressor(
                model_predictor=self.pred_model.clone(),
                fading_factor=self.fading_factor,
            )  # type: ignore

        return RegRule(
            template_splitter=self.splitter,
            split_criterion=split_criterion.VarianceRatioSplitCriterion(self.min_samples_split),
            pred_model=predictor,
            drift_detector=self.drift_detector.clone(),
        )

    def learn_one(self, x: dict, y: base.typing.RegTarget, w: int = 1) -> AMRules:
        any_covered = False
        to_del = set()

        for rule_id, rule in self._rules.items():
            if not rule.covers(x):
                continue

            # Anomaly detected skip training
            if rule.total_weight > self.m_min and rule.score_one(x) < self.anomaly_threshold:
                continue

            y_pred = rule.predict_one(x)

            in_drift = rule.drift_test(y, y_pred)  # noqa
            if in_drift:
                to_del.add(rule_id)
                self._n_drifts_detected += 1
                continue

            any_covered = True
            rule.learn_one(x, y, w)

            if rule.total_weight - rule.last_expansion_attempt_at >= self.n_min:
                updated_rule, expanded = rule.expand(self.delta, self.tau)

                if expanded:
                    updated_rule.pred_model = rule.pred_model
                    self._rules[rule_id] = updated_rule

            # Only the first matching rule is updated
            if self.ordered_rule_set:
                break

        if not any_covered:
            self._default_rule.learn_one(x, y, w)

            expanded = False
            if (
                self._default_rule.total_weight - self._default_rule.last_expansion_attempt_at
                >= self.n_min
            ):
                updated_rule, expanded = self._default_rule.expand(self.delta, self.tau)

            if expanded:
                updated_rule.pred_model = self._default_rule.pred_model  # noqa
                code = hash(updated_rule)
                self._rules[code] = updated_rule

                self._default_rule = self._new_rule()

        # Drop the rules affected by concept drift
        for rule_id in to_del:
            del self._rules[rule_id]

        return self

    def predict_one(self, x: dict) -> base.typing.RegTarget:
        y_pred = 0
        hits = 0

        for rule in self._rules.values():
            if rule.covers(x):
                y_pred += rule.predict_one(x)
                hits += 1

                if self.ordered_rule_set:
                    break

        if hits > 0:
            return y_pred / hits  # type: ignore
        else:
            return self._default_rule.predict_one(x)

    def anomaly_score(self, x) -> tuple[float, float, float]:
        """Aggregated anomaly score computed using all the rules that cover the input instance.

        Returns the mean anomaly score, the standard deviation of the score, and the proportion
        of rules that cover the instance (support). If the support is zero, it means that the
        default rule was used (not other rule covered `x`).

        Parameters
        ----------
        x
            The input instance.

        Returns
        -------
        mean_anomaly_score, std_anomaly_score, support

        Examples
        --------
        >>> from river import drift
        >>> from river import rules
        >>> from river import tree
        >>> from river.datasets import synth

        >>> dataset = synth.Friedman(seed=42).take(1001)

        >>> model = rules.AMRules(
        ...     n_min=50,
        ...     delta=0.1,
        ...     drift_detector=drift.ADWIN(),
        ...     splitter=tree.splitter.QOSplitter()
        ... )

        >>> for i, (x, y) in enumerate(dataset):
        ...     if i == 1000:
        ...         # Skip the last example
        ...         break
        ...     model = model.learn_one(x, y)

        >>> model.anomaly_score(x)
        (1.0168907243483933, 0.13045786430817402, 1.0)

        """
        var = stats.Var()

        for rule in self._rules.values():
            if rule.covers(x):
                var.update(rule.score_one(x))

        if var.mean.n > 0:
            return var.mean.get(), math.sqrt(var.get()), var.mean.n / len(self._rules)

        # No rule covers the instance. Use the default rule
        return self._default_rule.score_one(x), 0.0, 0

    def debug_one(self, x) -> str:
        """Return an explanation of how `x` is predicted

        Parameters
        ----------
        x
            The input instance.

        Returns
        -------
        A representation of the rules that cover the input and their prediction.

        Examples
        --------
        >>> from river import drift
        >>> from river import rules
        >>> from river import tree
        >>> from river.datasets import synth

        >>> dataset = synth.Friedman(seed=42).take(1001)

        >>> model = rules.AMRules(
        ...     n_min=50,
        ...     delta=0.1,
        ...     drift_detector=drift.ADWIN(),
        ...     splitter=tree.splitter.QOSplitter(),
        ...     ordered_rule_set=False
        ... )

        >>> for i, (x, y) in enumerate(dataset):
        ...     if i == 1000:
        ...         # Skip the last example
        ...         break
        ...     model = model.learn_one(x, y)

        >>> print(model.debug_one(x))
        Rule 0: 3 > 0.5060 and 0 > 0.2538
            Prediction (adaptive): 18.7217
        Rule 1: 1 > 0.2480 and 3 > 0.2573
            Prediction (adaptive): 19.4173
        Final prediction: 19.0695
        <BLANKLINE>

        """
        buffer = io.StringIO()
        _print = functools.partial(print, file=buffer)

        any_covered = False
        for i, rule in enumerate(self._rules.values()):
            if rule.covers(x):
                any_covered = True
                _print(f"Rule {i}: {repr(rule)}")
                _print(f"\tPrediction ({self.pred_type}): {rule.predict_one(x):.4f}")
                if self.ordered_rule_set:
                    break

        if any_covered:
            if not self.ordered_rule_set:
                _print(f"Final prediction: {self.predict_one(x):.4f}")
        else:
            _print("Default rule triggered:")
            _print(f"\tPrediction ({self.pred_type}): {self._default_rule.predict_one(x):.4f}")

        return buffer.getvalue()


# river/rules/base.py

from __future__ import annotations

import abc
import copy
import math
import numbers
import typing

from river import base, tree


class Literal(base.Base):
    __slots__ = "on", "at", "neg"

    def __init__(self, on, at, neg=False):
        self.on = on
        self.at = at
        self.neg = neg

    @abc.abstractmethod
    def __call__(self, x):
        pass

    @abc.abstractmethod
    def describe(self):
        pass


class NumericLiteral(Literal):
    def __init__(self, on, at, neg):
        super().__init__(on, at, neg)

    def __call__(self, x):
        if self.on in x:
            if not self.neg:
                return x[self.on] <= self.at
            else:
                return x[self.on] > self.at

        return False

    def describe(self):
        if not self.neg:
            return f"{self.on} ≤ {self.at:.4f}"
        else:
            return f"{self.on} > {self.at:.4f}"


class NominalLiteral(Literal):
    def __init__(self, on, at, neg):
        super().__init__(on, at, neg)

    def __call__(self, x):
        if self.on in x:
            if not self.neg:
                return x[self.on] == self.at
            else:
                return x[self.on] != self.at

        return False

    def describe(self):
        if not self.neg:
            return f"{self.on} = {self.at}"
        else:
            return f"{self.on} ≠ {self.at}"


class HoeffdingRule(base.Estimator, metaclass=abc.ABCMeta):
    """Base class for the decision rules based on the Hoeffding bound.

    It defines properties and operations shared by all the rule-based systems, for instance, how
    to perform rule expansion, how to monitor inputs, and whether or not a rule is currently
    covering an input datum.

    Parameters
    ----------
    template_splitter
        The attribute observer algorithm used to monitor and perform split attempts in numerical
        features. This splitter is be deep-copied to all the decision rules. The same attribute
        observer algorithms used by the Hoeffding Trees can be used with Hoeffding rules.
    split_criterion
        The criterion used to rank the split candidates. The same split criteria used by the
        Hoeffding Trees can be applied to the Hoeffding rules.
    attributes
        Other parameters passed to the rules via `**kwargs`.
    """

    def __init__(
        self,
        template_splitter: tree.splitter.base.Splitter,
        split_criterion: tree.split_criterion.base.SplitCriterion,
        **attributes,
    ):
        self.template_splitter = template_splitter
        self.split_criterion = split_criterion
        self.literals: list[Literal] = []
        self.splitters: dict[typing.Hashable, tree.splitter.Splitter] = {}

        self._total_weight = 0
        self._last_expansion_attempt_at = 0

        self.nominal_features: set[typing.Hashable] = set()

        self.__dict__.update(attributes)

    def _hoeffding_bound(self, r_heur, delta):
        r"""Compute the Hoeffding bound, used to decide how many samples are necessary to expand
        a decision rule.

        Notes
        -----
        The Hoeffding bound is defined as:

        $\\epsilon = \\sqrt{\\frac{R^2\\ln(1/\\delta))}{2n}}$

        where:

        $\\epsilon$: Hoeffding bound.
        $R$: Range of a random variable.
        $\\delta$: significance level.
        $n$: Number of samples.

        Parameters
        ----------
        r_heur
            Range of the split heuristic.
        delta
            The significance level.
        """
        return math.sqrt((r_heur * r_heur * math.log(1.0 / delta)) / (2.0 * self.total_weight))

    @property  # type: ignore
    @abc.abstractmethod
    def statistics(self):
        pass

    @statistics.setter  # type: ignore
    @abc.abstractmethod
    def statistics(self, target_stats):
        pass

    @property
    def total_weight(self):
        return self._total_weight

    @property
    def last_expansion_attempt_at(self):
        return self._last_expansion_attempt_at

    def expand(self, delta, tau):
        """
        Attempt to expand a decision rule.

        If the expansion succeeds, a new rule with all attributes reset and updated literals is
        returned along a boolean flag set to `True` to indicate the expansion. If the expansion
        fails, `self` is returned along with the boolean flag set to `False`.

        Parameters
        ----------
        delta
            The split test significance.
        tau
            The tie-breaking threshold.

        Returns
        -------
        rule, expanded
            The (potentially expanded) rule and an indicator of whether or not an expansion was
            successfully performed.

        """
        suggestions = []
        for att_id, splitter in self.splitters.items():
            suggestions.append(
                splitter.best_evaluated_split_suggestion(
                    criterion=self.split_criterion,
                    pre_split_dist=self.statistics,
                    att_idx=att_id,
                    binary_only=True,
                )
            )
        suggestions.sort()

        should_expand = False
        if len(suggestions) < 2:
            should_expand = True
        else:
            b_split = suggestions[-1]
            sb_split = suggestions[-2]

            hb = self._hoeffding_bound(self.split_criterion.range_of_merit(self.statistics), delta)

            if b_split.merit > 0 and (b_split.merit - sb_split.merit > hb or hb < tau):
                should_expand = True

        if should_expand:
            b_split = suggestions[-1]
            is_numerical = b_split.numerical_feature
            branch_no = self.split_criterion.select_best_branch(b_split.children_stats)

            if is_numerical:
                lit = NumericLiteral(b_split.feature, b_split.split_info, branch_no != 0)

                literal_updated = False
                for literal in self.literals:
                    if lit.on == literal.on and lit.neg == literal.neg:
                        # Update thresholds rather than adding a new literal
                        if not literal.neg and lit.at < literal.at:
                            literal.at = lit.at
                            literal_updated = True
                            break
                        elif literal.neg and lit.at > literal.at:
                            literal.at = lit.at
                            literal_updated = True
                            break

                # No threshold was updated, thus a new literal is added
                if not literal_updated:
                    self.literals.append(lit)
            else:
                lit = NominalLiteral(b_split.feature, b_split.split_info, branch_no != 0)
                # Add a new literal
                self.literals.append(lit)

            # Reset all the statistics stored in the decision rule
            updated_rule = self.clone()
            # Keep the literals
            updated_rule.literals.extend(self.literals)

            return updated_rule, True

        # If the expansion attempt failed, update the expansion tracker
        self._last_expansion_attempt_at = self.total_weight
        return self, False

    def covers(self, x: dict) -> bool:
        """Check if all the rule's conditions are fulfilled.

        Parameters
        ----------
        x
            Input instance.
        """
        return all(map(lambda lit: lit(x), self.literals))

    @abc.abstractmethod
    def new_nominal_splitter(self):
        """Define the attribute observer used when dealing with nominal features.

        Must be defined by classification and regression decision rule-based algorithms to match
        the target type.
        """
        pass

    def _iter_features(self, x: dict) -> typing.Iterable:
        """Determine how the input instance is looped through when updating the splitters.

        Parameters
        ----------
        x
            The input instance.
        """
        yield from x.items()

    @abc.abstractmethod
    def _update_target_stats(self, y, w):
        pass

    def _update_feature_stats(self, feat_name, feat_val, w):
        # By default, no feature stats are kept
        pass

    def update(self, x, y, w):
        self._total_weight += w
        self._update_target_stats(y, w)

        for feat_name, feat_val in self._iter_features(x):
            try:
                splt = self.splitters[feat_name]
            except KeyError:
                if isinstance(feat_val, numbers.Number):
                    self.splitters[feat_name] = copy.deepcopy(self.template_splitter)
                else:
                    self.nominal_features.add(feat_name)
                    self.splitters[feat_name] = self.new_nominal_splitter()
                splt = self.splitters[feat_name]
            splt.update(feat_val, y, w)
            self._update_feature_stats(feat_name, feat_val, w)

    def __repr__(self):
        return f"{' and '.join([lit.describe() for lit in self.literals])}"


# river/sketch/__init__.py

"""Data containers and collections for sequential data.

This module has summary and sketch structures that operate with constrained amounts
of memory and processing time.

"""
from __future__ import annotations

from .counter import Counter
from .heavy_hitters import HeavyHitters
from .histogram import Histogram
from .set import Set

__all__ = ["Counter", "HeavyHitters", "Histogram", "Set"]


# river/sketch/counter.py

from __future__ import annotations

import math
import random
import typing

import numpy as np

from river import base


class Counter(base.Base):
    """Counting using the Count-Min Sketch (CMS) algorithm.

    Contrary to an exhaustive approach, e.g., using a `collections.Counter`, CMS uses a
    limited and fixed amount of memory. The CMS algorithm uses a sketch structure consisting of a
    matrix $w \\times d$.

    These dimensions are obtained via:

    - $w = \\lceil \\frac{e}{\\epsilon} \\rceil$, where $e$ is the Euler number.

    - $d = \\lceil \\ln\\left(\\frac{1}{\\delta} \\right) \\rceil$.

    Decreasing the values of $\\epsilon$ (`epsilon`) and $\\delta$ (`delta`) increase the accuracy of the
    algorithm, at the cost of increased memory usage. The values of `w` and `d` control the hash tables' capability
    and the amount of hash collisions, respectively.

    CMS works by keeping `d` hash tables with `w` slots each. Elements are mapped to a slot in each hash table.
    These tables store the counting estimates. This implementation assumes the turnstile case described in the paper,
    i.e., count values and updates can be negative.

    The count values obtained by CMS are always overestimates. Suppose $c_i$ and $\\hat{c}_i$ are the ground truth and
    estimated count values, respectively, for a given element $i$. CMS guarantees that $c_i \\le \\hat{c}_i$ and,
    with probability $1 - \\delta$, $\\hat{c}_i \\le c_i + \\epsilon||\\mathbf{c}||_1$. In the expression,
    $||\\mathbf{c}||_1 = \\sum_i |c_i|$.

    Parameters
    ----------
    epsilon
        The approximation error parameter. The error in answering a query is within a factor of `epsilon`
        with probability `delta`.
    delta
        A query estimates have a probability of `1 - delta` of having errors which are a factor of `epsilon`.
        See the CMS description above for more details.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> import collections
    >>> from river import sketch

    >>> cms = sketch.Counter(epsilon=0.005, seed=0)

    >>> # To generate random numbers
    >>> rng = random.Random(7)

    >>> counter = collections.Counter()

    We can check the number of slots per hash table:
    >>> cms.n_slots
    544

    And the number of hash tables:
    >>> cms.n_tables
    3

    Let's compare the sketch against a brute force approach:

    >>> vals = []
    >>> for _ in range(10000):
    ...     v = rng.randint(-1000, 1000)
    ...     cms = cms.update(v)
    ...     counter[v] += 1
    ...     vals.append(v)

    Now, we can compare the estimates of CMS against the exhaustive counting strategy:

    >>> counter[7]
    5
    >>> cms[7]
    12
    >>> counter[532]
    4
    >>> cms[532]
    15

    Keep in mind that CMS is an approximate sketch algorithm. Couting estimates for unseen values
    might not be always reliable:

    >>> cms[1001]
    9

    We can check the number of elements stored by each approach:

    >>> len(counter), len(cms)
    (1982, 1632)

    And also retrieve the total sum of counts:

    >>> cms.total()
    10000

    We can decrease the error by allocating more memory in the CMS:

    >>> cms_a = sketch.Counter(epsilon=0.001, delta=0.01, seed=0)
    >>> for v in vals:
    ...     cms_a = cms_a.update(v)

    >>> cms_a[7]
    5
    >>> cms_a[532]
    4

    We can also obtain estimates of the dot product between two instances of `river.collections.Counter`. This could be useful,
    for instance, to estimate the cosine distance between the data monitored in two different counter sketch instances. Suppose we
    create another CMS instance (the number of slots and hash tables must match) that monitors another sample of the same data
    generating process:

    >>> cms_b = sketch.Counter(epsilon=0.001, delta=0.01, seed=7)

    >>> for _ in range(10000):
    ...     v = rng.randint(-1000, 1000)
    ...     cms_b = cms_b.update(v)

    Now, we can define a cosine distance function:

    >>> def cosine_dist(cms_a, cms_b):
    ...     num = cms_a @ cms_b
    ...     den = math.sqrt(cms_a @ cms_a) * math.sqrt(cms_b @ cms_b)
    ...     return num / den

    And use it to calculate the cosine distance between the elements monitored in `cms_a` and `cms_b`:

    >>> cosine_dist(cms_a, cms_b)
    0.175363...

    References
    ----------
    [^1]: [Cormode, G., & Muthukrishnan, S. (2005). An improved data stream summary: the count-min sketch and its applications. Journal of Algorithms, 55(1), 58-75.](https://www.cse.unsw.edu.au/~cs9314/07s1/lectures/Lin_CS9314_References/cm-latin.pdf)
    [^2]: [Count-Min Sketch](https://florian.github.io//count-min-sketch/)
    [^3]: [Hash functions family generator in Python](https://stackoverflow.com/questions/2255604/hash-functions-family-generator-in-python)

    """

    def __init__(self, epsilon: float = 0.1, delta: float = 0.05, seed: int | None = None):
        self.epsilon = epsilon
        self.delta = delta
        self.seed = seed

        self._w = int(math.ceil(math.e / self.epsilon))
        self._d = int(math.ceil(math.log(1 / self.delta)))

        self._rng = random.Random(self.seed)
        self._masks = [self._rng.getrandbits(64) for _ in range(self._d)]
        self._cms = np.zeros((self._d, self._w), dtype=np.int32)

    def _hash(self, x):
        return tuple(zip(*((i, (hash(x) ^ self._masks[i]) % self._w) for i in range(self._d))))

    def __getitem__(self, x) -> int:
        # Point query
        return min(self._cms[self._hash(x)])

    def __matmul__(self, other: Counter) -> int:
        # Dot product
        if self.n_slots != other.n_slots or self.n_tables != other.n_tables:
            unmatched_dims = (
                f"({self.n_slots}, {self.n_tables}) vs. ({other.n_slots}, {other.n_tables})"
            )
            raise ValueError(f"The number of slots and hash tables do not match: {unmatched_dims}.")

        return min(np.einsum("ij,ij->i", self._cms, other._cms))

    def __len__(self):
        # Return the total number of stored elements in the sketch
        return self._w * self._d

    def update(self, x: typing.Hashable, w: int = 1):
        self._cms[self._hash(x)] += w

        return self

    def total(self) -> int:
        """Return the total count."""
        return sum(self._cms[0, :])

    @property
    def n_slots(self) -> int:
        """The number of slots in each hash table."""
        return self._w

    @property
    def n_tables(self) -> int:
        """The number of stored hash tables."""
        return self._d


# river/sketch/heavy_hitters.py

from __future__ import annotations

import math
import operator
import typing

from river import base


class HeavyHitters(base.Base):
    """Find the Heavy Hitters using the Lossy Count with Forgetting factor algorithm[^1].

    Keep track of the most frequent item(set)s in a data stream and apply a forgetting factor to
    discard previous frequent items that do not often appear anymore. This is an approximation
    algorithm designed to work with a limited amount of memory rather than accounting for every possible
    solution (thus using an unbounded memory footprint). Any hashable type can be passed as input, hence
    tuples or frozensets can also be monitored.

    Considering a data stream where `n` elements were observed so far, the Lossy Count
    algorithm has the following properties:

    - All item(set)s whose true frequency exceeds `support * n` are output. There are no
    false negatives;

    - No item(set) whose true frequency is less than `(support - epsilon) * n` is outputted;

    - Estimated frequencies are less than the true frequencies by at most `epsilon * n`.

    Parameters
    ----------
    support
        The support threshold used to determine if an item is frequent. The value of `support` must
        be in $[0, 1]$. Elements whose frequency is higher than `support` times the number of
        observations seen so far are outputted.
    epsilon
        Error parameter to control the accuracy-memory tradeoff. The value of `epsilon` must be
        in $(0, 1]$ and typically `epsilon` $\\ll$ `support`. The smaller the `epsilon`, the more
        accurate the estimates will be, but the count sketch will have an increased memory
        footprint.
    fading_factor
        Forgetting factor applied to the frequency estimates to reduce the impact of old items.
        The value of `fading_factor` must be in $(0, 1]$.

    Examples
    --------

    >>> import random
    >>> import string
    >>> from river import sketch

    >>> rng = random.Random(42)
    >>> hh = sketch.HeavyHitters()

    We will feed the counter with printable ASCII characters:

    >>> for _ in range(10_000):
    ...     hh = hh.update(rng.choice(string.printable))

    We can retrieve estimates of the `n` top elements and their frequencies. Let's try `n=3`
    >>> hh.most_common(3)
    [(',', 122.099142...), ('[', 116.049510...), ('W', 115.013402...)]

    We can also access estimates of individual elements:

    >>> hh['A']
    99.483575...

    Unobserved elements are handled just fine:
    >>> hh[(1, 2, 3)]
    0.0

    References
    ----------
    [^1]: Veloso, B., Tabassum, S., Martins, C., Espanha, R., Azevedo, R., & Gama, J. (2020).
    Interconnect bypass fraud detection: a case study. Annals of Telecommunications, 75(9), 583-596.

    """

    def __init__(
        self, support: float = 0.001, epsilon: float = 0.005, fading_factor: float = 0.999
    ):
        if support > epsilon:
            raise ValueError("'support' must be smaller than 'epsilon'.")

        self.support = support
        self.epsilon = epsilon
        self.fading_factor = fading_factor

        self._bucket_width = math.ceil(1 / self.epsilon)
        self._n: int = 0
        self._entries: dict[typing.Hashable, tuple[float, float]] = {}
        self._delta: float = self._bucket_width

    def __getitem__(self, index) -> float:
        return self._entries.get(index, (0.0, None))[0]

    def update(self, x: typing.Hashable):
        self._n += 1
        current_bucket = math.ceil(self._n / self._bucket_width)
        freq, delta = 1.0, current_bucket - 1.0
        if x in self._entries:
            freq, delta = self._entries[x]
            freq += 1
        self._entries[x] = (freq, delta)

        # If at bucket boundary then prune low frequency entries.
        if self._n % self._bucket_width == 0:
            prune = []
            for key in self._entries:
                freq, delta = self._entries[key]
                freq *= self.fading_factor
                self._entries[key] = (freq, delta)

                if freq + delta <= current_bucket:
                    prune.append(key)

            for key in prune:
                del self._entries[key]

            self._delta = self._bucket_width + self._delta * self.fading_factor

        return self

    def most_common(self, n: int | None = None) -> list[tuple[typing.Hashable, float]]:
        res = []
        for key in self._entries:
            freq, _ = self._entries[key]
            if freq >= (self.support - self.epsilon) * self._delta:
                res.append((key, freq))

        if n is None:
            n = len(res)

        return sorted(res, key=operator.itemgetter(1), reverse=True)[:n]


# river/sketch/histogram.py

from __future__ import annotations

import collections
import heapq
import itertools
import math

from river import base

__all__ = ["Histogram"]


class Bin:
    """A Bin is an element of a Histogram."""

    __slots__ = ["left", "right", "count"]

    def __init__(self, left, right, count):
        self.left = left
        self.right = right
        self.count = count

    def __iadd__(self, other):
        """Merge with another bin."""
        if other.left < self.left:
            self.left = other.left
        if other.right > self.right:
            self.right = other.right
        self.count += other.count
        return self

    def __lt__(self, other):
        return self.right < other.left

    def __eq__(self, other):
        return self.left == other.left and self.right == other.right

    def __repr__(self):
        return f"[{self.left:.5f}, {self.right:.5f}]: {self.count}"


def coverage_ratio(x: Bin, y: Bin) -> float:
    """Returns the amount of y covered by x.

    Examples
    --------

    >>> coverage_ratio(Bin(1, 2, 0), Bin(1, 2, 0))
    1.0

    >>> coverage_ratio(Bin(1, 3, 0), Bin(2, 4, 0))
    0.5

    >>> coverage_ratio(Bin(1, 3, 0), Bin(3, 5, 0))
    0.0

    >>> coverage_ratio(Bin(1, 3, 0), Bin(0, 4, 0))
    0.5

    >>> coverage_ratio(Bin(0, 4, 0), Bin(1, 3, 0))
    1.0

    >>> coverage_ratio(Bin(1, 3, 0), Bin(0, 1, 0))
    0.0

    >>> coverage_ratio(Bin(1, 1, 0), Bin(1, 1, 0))
    1.0

    """
    if y.left == y.right:
        return float(x.left <= y.left <= x.right)
    return max(0, min(x.right, y.right) - max(x.left, y.left)) / (y.right - y.left)


class Histogram(collections.UserList, base.Base):
    """Streaming histogram.

    Parameters
    ----------
    max_bins
        Maximal number of bins.

    Attributes
    ----------
    n
        Total number of seen values.

    Examples
    --------

    >>> from river import sketch
    >>> import numpy as np

    >>> np.random.seed(42)

    >>> values = np.hstack((
    ...     np.random.normal(-3, 1, 1000),
    ...     np.random.normal(3, 1, 1000),
    ... ))

    >>> hist = sketch.Histogram(max_bins=15)

    >>> for x in values:
    ...     hist = hist.update(x)

    >>> for bin in hist:
    ...     print(bin)
    [-6.24127, -6.24127]: 1
    [-5.69689, -5.19881]: 8
    [-5.12390, -4.43014]: 57
    [-4.42475, -3.72574]: 158
    [-3.71984, -3.01642]: 262
    [-3.01350, -2.50668]: 206
    [-2.50329, -0.81020]: 294
    [-0.80954, 0.29677]: 19
    [0.40896, 0.82733]: 7
    [0.84661, 1.25147]: 24
    [1.26029, 2.30758]: 178
    [2.31081, 3.05701]: 284
    [3.05963, 3.69695]: 242
    [3.69822, 5.64434]: 258
    [6.13775, 6.19311]: 2

    References
    ----------
    [^1]: [Ben-Haim, Y. and Tom-Tov, E., 2010. A streaming parallel decision tree algorithm. Journal of Machine Learning Research, 11(Feb), pp.849-872.](http://jmlr.org/papers/volume11/ben-haim10a/ben-haim10a.pdf)
    [^2]: [Go implementation](https://github.com/VividCortex/gohistogram)

    """

    def __init__(self, max_bins=256):
        super().__init__()
        self.max_bins = max_bins
        self.n = 0

    def update(self, x):
        self.n += 1
        b = Bin(x, x, 1)

        # Insert the bin if the histogram is empty
        if not self:
            self.append(b)
            return self

        # Use bisection to find where to insert
        # We don't use the bisect module in order to save some CPU cycles
        lo = 0
        hi = len(self)
        i = (lo + hi) // 2
        while lo < hi:
            if self[i] < b:
                lo = i + 1
            else:
                hi = i
            i = (lo + hi) // 2

        if i == len(self):
            # x is past the right-most bin
            self.append(b)
        else:
            # Increment the bin counter if x is part of the ith bin
            if x >= self[i].left:
                self[i].count += 1
            # Insert the bin if it is between bin i-1 and bin i
            else:
                self.insert(i, b)

        # Bins have to be merged if there are more than max_bins
        if len(self) > self.max_bins:
            self._shrink(1)

        return self

    def _shrink(self, k):
        """Shrinks the histogram by merging the two closest bins."""

        if k == 1:
            # Find the closest pair of bins
            min_diff = math.inf
            min_idx = None
            for idx, (b1, b2) in enumerate(zip(self.data[:-1], self.data[1:])):
                diff = b2.right - b1.right
                if diff < min_diff:
                    min_diff = diff
                    min_idx = idx

            # Merge the bins
            self[min_idx] += self.pop(min_idx + 1)
            return

        indexes = range(len(self) - 1)

        def bin_distance(i):
            return self[i + 1].right - self[i].right

        for i in sorted(heapq.nsmallest(n=k, iterable=indexes, key=bin_distance), reverse=True):
            self[i] += self.pop(i + 1)  # Calls Bin.__iadd__

    def iter_cdf(self, X, verbose=False):
        """Yields CDF values for a sorted iterable of values.

        This is faster than calling `cdf` with many values.

        Examples
        --------

        >>> from river import sketch

        >>> hist = sketch.Histogram()
        >>> for x in range(4):
        ...     hist = hist.update(x)

        >>> print(hist)
        [0.00000, 0.00000]: 1
        [1.00000, 1.00000]: 1
        [2.00000, 2.00000]: 1
        [3.00000, 3.00000]: 1

        >>> X = [-1, 0, .5, 1, 2.5, 3.5]
        >>> for x, cdf in zip(X, hist.iter_cdf(X)):
        ...     print(x, cdf)
        -1 0.0
        0 0.25
        0.5 0.25
        1 0.5
        2.5 0.75
        3.5 1.0

        """

        bins = iter(self)
        b = next(bins)
        INF = Bin(math.inf, math.inf, 0)

        cdf = 0

        for x in X:
            while x >= b.right:
                cdf += b.count
                b = next(bins, INF)

            if x > b.left:
                yield (cdf + b.count * (x - b.left) / (b.right - b.left)) / self.n
                continue

            yield cdf / self.n

    def cdf(self, x):
        """Cumulative distribution function.

        Examples
        --------

        >>> from river import sketch

        >>> hist = sketch.Histogram()
        >>> for x in range(4):
        ...     hist = hist.update(x)

        >>> print(hist)
        [0.00000, 0.00000]: 1
        [1.00000, 1.00000]: 1
        [2.00000, 2.00000]: 1
        [3.00000, 3.00000]: 1

        >>> hist.cdf(-1)
        0.0

        >>> hist.cdf(0)
        0.25

        >>> hist.cdf(.5)
        0.25

        >>> hist.cdf(1)
        0.5

        >>> hist.cdf(2.5)
        0.75

        >>> hist.cdf(3.5)
        1.0

        """
        return next(self.iter_cdf([x]))

    def __add__(self, other):
        """

        Example:

            >>> h1 = Histogram()
            >>> for b in [Bin(0, 2, 4), Bin(4, 5, 9)]:
            ...     h1.append(b)

            >>> h2 = Histogram()
            >>> for b in [Bin(1, 3, 8), Bin(3, 4, 5)]:
            ...     h2.append(b)

            >>> h1 + h2
            [0.00000, 1.00000]: 2.0
            [1.00000, 2.00000]: 6.0
            [2.00000, 3.00000]: 4.0
            [3.00000, 4.00000]: 5.0
            [4.00000, 5.00000]: 9.0

        """

        xs = iter(self)
        ys = iter(other)
        rights = heapq.merge(
            itertools.chain.from_iterable((b.left, b.right) for b in self),
            itertools.chain.from_iterable((b.left, b.right) for b in other),
        )

        b = Bin(next(rights), next(rights), 0)

        x = next(xs)
        y = next(ys)

        hist = Histogram(max_bins=max(self.max_bins, other.max_bins))

        while True:
            b.count += coverage_ratio(b, x) * x.count
            b.count += coverage_ratio(b, y) * y.count

            if b.count:
                hist.append(b)

            try:
                b = Bin(b.right, next(rights), 0)
            except StopIteration:
                break

            if b.left >= x.right:
                x = next(xs, x)
            if b.left >= y.right:
                y = next(ys, y)

        hist._shrink(k=len(hist) - hist.max_bins)

        return hist

    def __repr__(self):
        return "\n".join(str(b) for b in self)

    def __str__(self):
        return repr(self)


# river/sketch/set.py

from __future__ import annotations

import math
import random
import typing

from river import base


class Set(base.Base):
    """Approximate tracking of observed items using Bloom filters.

    Bloom filters enable using a limited amount of memory to check whether a given item was already
    observed in a stream. They can be used similarly to Python's built-in sets with the difference
    that items are not explicitly stored. For that reason, element removal and set difference are not
    currently supported.

    Bloom filters store a bit array and map incoming items to `k` index positions in the such array.
    The selected positions are set to `True`. Therefore, a binary code representation is created for each item.
    Membership works by projecting the query item and checking if every position of its binary code is
    `True`. If that is not the case, the item was not observed yet. A nice property of Bloom filters is
    that they do not yield false negatives: unobserved items might be signalized as observed, but observed
    items are never signalized as unobserved.

    If more than one item has the same binary code, i.e., hash collisions happen, the accuracy of the Bloom
    filter decreases, and false positives are produced. For instance, a previously unobserved item is signalized
    as observed. Increasing the size of the binary array and the value of `k` increase the filter's accuracy as
    hash collisions are avoided. Nonetheless, even using an increased number of hash functions, hash collisions
    will frequently happen if the array capacity is too small. The length of the bit array and the number of
    hash functions are inferred automatically from the supplied `capacity` and `fp_rate`.


    Parameters
    ----------
    capacity
        The maximum capacity of the Bloom filter, i.e., the maximum number of distinct items to store given
        the selected `fp_rate`.
    fp_rate
        The allowed rate of false positives. The probability of obtaining a true positive is `1 - fp_rate`.
    seed
        Random seed for reproducibility.

    Examples
    --------

    >>> import random
    >>> from river import sketch

    >>> rng = random.Random(42)
    >>> s_set = sketch.Set(capacity=100, seed=0)


    We can retrieve the number of selected hash functions:

    >>> s_set.n_hash
    7

    And the size of the binary array used by the Bloom filter:
    >>> s_set.n_bits
    959

    We can add new items and check for membership using the same calls used by Python's
    standard sets:
    >>> for _ in range(1000):
    ...     s_set.add(rng.randint(0, 200))

    >>> 1 in s_set
    True

    False positives might happen if the capacity is not large enough:
    >>> -10 in s_set
    True

    Iterables can also be supplied to perform multiple updates with a single call to `update`:
    >>> s_set = s_set.update([1, 2, 3, 4, 5, 6, 7])

    We can also combine instances of `sketch.Set` using the intersection and union operations, as long as
    they share the same hash functions and capability. In other words, all they hyperparameters match.
    Let's create two instances that will monitor different portions of a stream of random numbers:

    >>> s1 = sketch.Set(seed=8)
    >>> s2 = sketch.Set(seed=8)

    >>> for _ in range(1000):
    ...     s1.add(rng.randint(0, 5000))

    >>> for _ in range(1000):
    ...     s2.add(rng.randint(0, 5000))

    >>> 43 in s1
    True
    >>> 43 in s2
    False

    We can get the intersection between the two instances by using:

    >>> s_intersection = s1 & s2
    >>> 43 in s_intersection
    False

    We can also obtain the set union:

    >>> s_union = s1 | s2

    >>> 43 in s_union
    True

    The same effect of the non-inplace dunder methods can be achieved via explicit method calls:

    >>> 43 in s1.intersection(s2)
    False

    >>> 43 in s1.union(s2)
    True

    Notes
    -----
    This implementation uses an integer to represent the binary array. Bitwise operations are performed in the
    integer to reflect the Bloom filter updates.

    References
    ----------
    [^1]: [Florian Hartmann's blog article on Bloom Filters](https://florian.github.io/bloom-filters/).
    [^2]: [Wikipedia entry on Bloom filters](https://en.wikipedia.org/wiki/Bloom_filter).

    """

    def __init__(self, capacity: int = 2048, fp_rate: float = 0.01, seed: int | None = None):
        self.capacity = capacity
        self.fp_rate = fp_rate
        self.seed = seed

        # Size of the binary array
        self._asize = int(
            math.ceil(-((self.capacity * math.log(self.fp_rate)) / (math.log(2) ** 2)))
        )
        # Number of hash functions
        self._n_hash = int(math.ceil((self._asize / self.capacity) * math.log(2)))

        self._rng = random.Random(self.seed)

        # Random masks to obtain independent hash functions
        self._masks = [self._rng.getrandbits(64) for _ in range(self._n_hash)]

        # The binary array is represented as an integer that has with one extra bit.
        # The most significant bit is always 1 and it is used to ensure the bit_length
        # does not change during the execution.
        self._bloom = 2 ** (self._asize)

    @property
    def n_hash(self) -> int:
        """Return the number of used hash functions."""
        return self._n_hash

    @property
    def n_bits(self) -> int:
        """Return the size of the binary array used by the Bloom filter."""
        return self._asize

    def _hash(self, x: typing.Hashable):
        return [(hash(x) ^ self._masks[i]) % self._asize for i in range(self._n_hash)]

    def add(self, x: typing.Hashable):
        pos = self._hash(x)

        # Set the corresponding bits to 1
        # https://stackoverflow.com/a/12174125
        for p in pos:
            self._bloom |= 1 << p

    def update(self, values: typing.Iterable):
        for x in values:
            self.add(x)

        return self

    def __contains__(self, x: typing.Hashable):
        proj = []
        pos = self._hash(x)

        # Check if the corresponding bits are 1
        # https://stackoverflow.com/a/45221136
        for p in pos:
            proj.append((self._bloom >> p) & 1 == 1)

        return all(proj)

    def _is_mergeable(self, other: Set) -> bool:
        if len(self._masks) != len(other._masks):
            return False

        # If two instances share the same random masks, they also share the same
        # hash functions
        mask_check = all(m1 == m2 for m1, m2 in zip(self._masks, other._masks))
        return mask_check and self.capacity == other.capacity

    def _check_mergeable(self, other):
        if not self._is_mergeable(other):
            raise ValueError(
                "The supplied 'sketch.Set' instances cannot the combined.",
                "Ensure their 'capacity', 'fp_rate', and 'seed' match.",
            )

    def __iand__(self, other: Set):
        self._check_mergeable(other)

        self._bloom &= other._bloom
        return self

    def __and__(self, other: Set):
        new = self.clone(include_attributes=True)
        new &= other
        return new

    def __ior__(self, other: Set):
        self._check_mergeable(other)

        self._bloom |= other._bloom
        return self

    def __or__(self, other: Set):
        new = self.clone(include_attributes=True)
        new |= other

        return new

    def intersection(self, other: Set):
        """Set intersection.

        Return a new instance that results from the set intersection between the current `Set` object
        and `other`. Dunder operators can be used to replace the method call, i.e., `a &= b` and
        `a & b` for inplace and non-inplace intersections, respectively.

        Parameters
        ----------
        other
            Another instance of `sketch.Set`. All the hyperparameters values of `other` must match
            the values from the instance from which this method is called.

        Raises
        ------
        ValueError
            In case the hyperparameter values of the two involved `sketch.Set` instances do not match.
        """
        return self & other

    def union(self, other: Set):
        """Set union.

        Return a new instance that results from the set union between the current `Set` object and `other`.
        Dunder operators can be used to replace the method call, i.e., `a |= b` and `a | b` for inplace
        and non-inplace unions, respectively.

        Parameters
        ----------
        other
            Another instance of `sketch.Set`. All the hyperparameters values of `other` must match
            the values from the instance from which this method is called.

        Raises
        ------
        ValueError
            In case the hyperparameter values of the two involved `sketch.Set` instances do not match.
        """
        return self | other


# river/stats/__init__.py

"""Running statistics"""
from __future__ import annotations

from . import base
from .auto_corr import AutoCorr
from .count import Count
from .cov import Cov
from .entropy import Entropy
from .ewmean import EWMean
from .ewvar import EWVar
from .iqr import IQR, RollingIQR
from .kurtosis import Kurtosis
from .link import Link
from .mad import MAD
from .maximum import AbsMax, Max, RollingAbsMax, RollingMax
from .mean import BayesianMean, Mean
from .minimum import Min, RollingMin
from .mode import Mode, RollingMode
from .n_unique import NUnique
from .pearson import PearsonCorr
from .ptp import PeakToPeak, RollingPeakToPeak
from .quantile import Quantile, RollingQuantile
from .sem import SEM
from .shift import Shift
from .skew import Skew
from .summing import Sum
from .var import Var

__all__ = [
    "base",
    "AbsMax",
    "AutoCorr",
    "BayesianMean",
    "Count",
    "Cov",
    "Entropy",
    "EWMean",
    "EWVar",
    "IQR",
    "Kurtosis",
    "Link",
    "MAD",
    "Max",
    "Mean",
    "Min",
    "Mode",
    "NUnique",
    "PeakToPeak",
    "PearsonCorr",
    "Quantile",
    "RollingAbsMax",
    "RollingIQR",
    "RollingMax",
    "RollingMin",
    "RollingMode",
    "RollingPeakToPeak",
    "RollingQuantile",
    "SEM",
    "Shift",
    "Skew",
    "Sum",
    "Var",
]


# river/stats/auto_corr.py

from __future__ import annotations

import collections
import numbers

from river import stats


class AutoCorr(stats.base.Univariate):
    """Measures the serial correlation.

    This method computes the Pearson correlation between the current value and the value seen `n`
    steps before.

    Parameters
    ----------
    lag

    Examples
    --------

    The following examples are taken from the [pandas documentation](https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.Series.autocorr.html).

    >>> from river import stats

    >>> auto_corr = stats.AutoCorr(lag=1)
    >>> for x in [0.25, 0.5, 0.2, -0.05]:
    ...     print(auto_corr.update(x).get())
    0
    0
    -1.0
    0.103552

    >>> auto_corr = stats.AutoCorr(lag=2)
    >>> for x in [0.25, 0.5, 0.2, -0.05]:
    ...     print(auto_corr.update(x).get())
    0
    0
    0
    -1.0

    >>> auto_corr = stats.AutoCorr(lag=1)
    >>> for x in [1, 0, 0, 0]:
    ...     print(auto_corr.update(x).get())
    0
    0
    0
    0

    """

    def __init__(self, lag: int):
        self.window: collections.deque[numbers.Number] = collections.deque(maxlen=lag)
        self.lag = lag
        self.pearson = stats.PearsonCorr(ddof=1)

    @property
    def name(self):
        return f"autocorr_{self.lag}"

    def update(self, x):
        # The correlation can be update once enough elements have been seen
        if len(self.window) == self.lag:
            self.pearson.update(x, self.window[0])

        # Add x to the window
        self.window.append(x)

        return self

    def get(self):
        return self.pearson.get()


# river/stats/base.py

from __future__ import annotations

import abc
import numbers

from river import base

__all__ = ["Univariate", "Bivariate"]


class Statistic(abc.ABC, base.Base):
    """A statistic."""

    # Define the format specification used for string representation.
    _fmt = ",.6f"  # Use commas to separate big numbers and show 6 decimals

    @abc.abstractmethod
    def get(self) -> float | None:
        """Return the current value of the statistic."""

    def __repr__(self):
        try:
            value = self.get()
        except NotImplementedError:
            value = None
        fmt_value = None if value is None else f"{value:{self._fmt}}".rstrip("0")
        return f"{self.__class__.__name__}: {fmt_value}"

    def __str__(self):
        return repr(self)

    def __gt__(self, other):
        return self.get() > other.get()


class Univariate(Statistic):
    """A univariate statistic measures a property of a variable."""

    @abc.abstractmethod
    def update(self, x: numbers.Number):
        """Update and return the called instance."""
        raise NotImplementedError

    @property
    def name(self):
        return self.__class__.__name__.lower()

    def __or__(self, other):
        from .link import Link

        return Link(left=self, right=other)


class RollingUnivariate(Univariate):
    """A rolling univariate statistic measures a property of a variable over a window."""

    @property
    @abc.abstractmethod
    def window_size(self):
        pass

    @property
    def name(self):
        return f"{super().name}_{self.window_size}"


class Bivariate(Statistic):
    """A bivariate statistic measures a relationship between two variables."""

    @abc.abstractmethod
    def update(self, x, y):
        """Update and return the called instance."""


# river/stats/count.py

from __future__ import annotations

from river import stats


class Count(stats.base.Univariate):
    """A simple counter.

    Attributes
    ----------
    n : int
        The current number of observations.

    """

    def __init__(self):
        self.n = 0

    def update(self, x=None):
        self.n += 1
        return self

    def get(self):
        return self.n


# river/stats/cov.py

from __future__ import annotations

import copy

import numpy as np

from river import stats


class Cov(stats.base.Bivariate):
    """Covariance.

    Parameters
    ----------
    ddof
        Delta Degrees of Freedom.

    Examples
    --------

    >>> from river import stats

    >>> x = [-2.1,  -1,  4.3]
    >>> y = [   3, 1.1, 0.12]

    >>> cov = stats.Cov()

    >>> for xi, yi in zip(x, y):
    ...     print(cov.update(xi, yi).get())
    0.0
    -1.044999
    -4.286

    This class has a `revert` method, and can thus be wrapped by `utils.Rolling`:

    >>> from river import utils

    >>> x = [-2.1,  -1, 4.3, 1, -2.1,  -1, 4.3]
    >>> y = [   3, 1.1, .12, 1,    3, 1.1, .12]

    >>> rcov = utils.Rolling(stats.Cov(), window_size=3)

    >>> for xi, yi in zip(x, y):
    ...     print(rcov.update(xi, yi).get())
    0.0
    -1.045
    -4.286
    -1.382
    -4.589
    -1.415
    -4.286

    Notes
    -----
    The outcomes of the incremental and parallel updates are consistent with numpy's
    batch processing when $\\text{ddof} \\le 1$.

    References
    ----------
    [^1]: [Wikipedia article on algorithms for calculating variance](https://www.wikiwand.com/en/Algorithms_for_calculating_variance#/Covariance)
    [^2]: Schubert, E. and Gertz, M., 2018, July. Numerically stable parallel computation of
    (co-) variance. In Proceedings of the 30th International Conference on Scientific and
    Statistical Database Management (pp. 1-12).

    """

    def __init__(self, ddof=1):
        self.ddof = ddof
        self.mean_x = stats.Mean()
        self.mean_y = stats.Mean()
        self.cov = 0

    @property
    def n(self):
        return self.mean_x.n

    def update(self, x, y, w=1.0):
        dx = x - self.mean_x.get()
        self.mean_x.update(x, w)
        self.mean_y.update(y, w)
        self.cov += w * (dx * (y - self.mean_y.get()) - self.cov) / max(self.n - self.ddof, 1)
        return self

    def revert(self, x, y, w=1.0):
        dx = x - self.mean_x.get()
        self.mean_x.revert(x, w)
        self.mean_y.revert(y, w)
        self.cov -= w * (dx * (y - self.mean_y.get()) - self.cov) / max(self.n - self.ddof, 1)
        return self

    def update_many(self, X: np.ndarray, Y: np.ndarray):
        dx = X - self.mean_x.get()
        self.mean_x.update_many(X)
        self.mean_y.update_many(Y)
        self.cov += (dx * (Y - self.mean_y.get()) - self.cov).sum() / max(self.n - self.ddof, 1)
        return self

    def get(self):
        return self.cov

    @classmethod
    def _from_state(cls, n, mean_x, mean_y, cov, *, ddof=1):
        new = cls(ddof=ddof)
        new.mean_x = stats.Mean._from_state(n, mean_x)
        new.mean_y = stats.Mean._from_state(n, mean_y)
        new.cov = cov
        return new

    def __iadd__(self, other):
        old_mean_x = self.mean_x.get()
        old_mean_y = self.mean_y.get()
        old_n = self.n

        # Update mean estimates
        self.mean_x += other.mean_x
        self.mean_y += other.mean_y

        if self.n <= self.ddof:
            return self

        # Scale factors
        scale_a = old_n - self.ddof
        scale_b = other.n - other.ddof

        # Scale the covariances
        self.cov = scale_a * self.cov + scale_b * other.cov
        # Apply correction factor
        self.cov += (
            (old_mean_x - other.mean_x.get())
            * (old_mean_y - other.mean_y.get())
            * ((old_n * other.n) / self.n)
        )
        # Reapply scale
        self.cov /= max(self.n - self.ddof, 1)

        return self

    def __add__(self, other):
        result = copy.deepcopy(self)
        result += other
        return result

    def __isub__(self, other):
        if self.n <= self.ddof:
            return self

        old_n = self.n

        # Update mean estimates
        self.mean_x -= other.mean_x
        self.mean_y -= other.mean_y

        if self.n <= self.ddof:
            self.cov = 0
            return self

        # Scale factors
        scale_x = old_n - self.ddof
        scale_b = other.mean_x.n - other.ddof

        # Scale the covariances
        self.cov = scale_x * self.cov - scale_b * other.cov
        # Apply correction
        self.cov -= (
            (self.mean_x.get() - other.mean_x.get())
            * (self.mean_y.get() - other.mean_y.get())
            * ((self.n * other.mean_x.n) / old_n)
        )
        # Re-apply scale factor
        self.cov /= self.n - self.ddof

        return self

    def __sub__(self, other):
        result = copy.deepcopy(self)
        result -= other
        return result


# river/stats/entropy.py

from __future__ import annotations

import collections
import math

from river import stats


class Entropy(stats.base.Univariate):
    """Running entropy.

    Parameters
    ----------
    fading_factor
        Fading factor.
    eps
        Small value that will be added to the denominator to avoid division by zero.

    Attributes
    ----------
    entropy : float
        The running entropy.
    n : int
        The current number of observations.
    counter : collections.Counter
        Count the number of times the values have occurred

    Examples
    --------

    >>> import math
    >>> import random
    >>> import numpy as np
    >>> from scipy.stats import entropy
    >>> from river import stats

    >>> def entropy_list(labels, base=None):
    ...   value,counts = np.unique(labels, return_counts=True)
    ...   return entropy(counts, base=base)

    >>> SEED = 42 * 1337
    >>> random.seed(SEED)

    >>> entro = stats.Entropy(fading_factor=1)

    >>> list_animal = []
    >>> for animal, num_val in zip(['cat', 'dog', 'bird'],[301, 401, 601]):
    ...     list_animal += [animal for i in range(num_val)]
    >>> random.shuffle(list_animal)

    >>> for animal in list_animal:
    ...     _ = entro.update(animal)

    >>> print(f'{entro.get():.6f}')
    1.058093
    >>> print(f'{entropy_list(list_animal):.6f}')
    1.058093

    References
    ----------
    [^1]: [Sovdat, B., 2014. Updating Formulas and Algorithms for Computing Entropy and Gini Index from Time-Changing Data Streams. arXiv preprint arXiv:1403.6348.](https://arxiv.org/pdf/1403.6348.pdf)

    """

    def __init__(self, fading_factor=1, eps=1e-8):
        if 0 < fading_factor <= 1:
            self.fading_factor = fading_factor
        else:
            raise ValueError("fading_factor must be between 0 excluded and 1")
        self.eps = eps
        self.entropy = 0
        self.n = 0
        self.counter = collections.Counter()

    @property
    def name(self):
        return "entropy"

    def update(self, x):
        cx = self.counter.get(x, 0)
        n = self.n
        eps = self.eps
        fading_factor = self.fading_factor

        entropy = self.entropy
        entropy = (n + eps) / (n + 1) * (fading_factor * entropy - math.log((n + eps) / (n + 1)))
        entropy -= (cx + 1) / (n + 1) * math.log((cx + 1) / (n + 1))
        entropy += (cx + eps) / (n + 1) * math.log((cx + eps) / (n + 1))
        self.entropy = entropy

        self.n += 1
        self.counter.update([x])

        return self

    def get(self):
        return self.entropy


# river/stats/ewmean.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class EWMean(stats.base.Univariate):
    """Exponentially weighted mean.

    Parameters
    ----------
    fading_factor
        The closer `fading_factor` is to 1 the more the statistic will adapt to recent values.

    Attributes
    ----------
    mean : float
        The running exponentially weighted mean.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, 3, 5, 4, 6, 8, 7, 9, 11]
    >>> ewm = stats.EWMean(fading_factor=0.5)
    >>> for x in X:
    ...     print(ewm.update(x).get())
    1.0
    2.0
    3.5
    3.75
    4.875
    6.4375
    6.71875
    7.859375
    9.4296875

    References
    ----------
    [^1]: [Finch, T., 2009. Incremental calculation of weighted mean and variance. University of Cambridge, 4(11-5), pp.41-42.](https://fanf2.user.srcf.net/hermes/doc/antiforgery/stats.pdf)
    [^2]: [Exponential Moving Average on Streaming Data](https://dev.to/nestedsoftware/exponential-moving-average-on-streaming-data-4hhl)

    """

    # Note for devs, if you want look the pure python implementation here:
    # https://github.com/online-ml/river/blob/40c3190c9d05671ae4c2dc8b76c163ea53a45fb0/river/stats/ewmean.py
    def __init__(self, fading_factor=0.5):
        if not 0 <= fading_factor <= 1:
            raise ValueError("q is not comprised between 0 and 1")
        self.fading_factor = fading_factor
        self._ewmean = _rust_stats.RsEWMean(fading_factor)
        self.mean = 0

    @property
    def name(self):
        return f"ewm_{self.fading_factor}"

    def update(self, x):
        self._ewmean.update(x)
        return self

    def get(self):
        return self._ewmean.get()


# river/stats/ewvar.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class EWVar(stats.base.Univariate):
    """Exponentially weighted variance.

    To calculate the variance we use the fact that Var(X) = Mean(x^2) - Mean(x)^2 and internally
    we use the exponentially weighted mean of x/x^2 to calculate this.

    Parameters
    ----------
    fading_factor
        The closer `fading_factor` is to 1 the more the statistic will adapt to recent values.

    Attributes
    ----------
    variance : float
        The running exponentially weighted variance.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, 3, 5, 4, 6, 8, 7, 9, 11]
    >>> ewv = stats.EWVar(fading_factor=0.5)
    >>> for x in X:
    ...     print(ewv.update(x).get())
    0.0
    1.0
    2.75
    1.4375
    1.984375
    3.43359375
    1.7958984375
    2.198974609375
    3.56536865234375

    References
    ----------
    [^1]: [Finch, T., 2009. Incremental calculation of weighted mean and variance. University of Cambridge, 4(11-5), pp.41-42.](https://fanf2.user.srcf.net/hermes/doc/antiforgery/stats.pdf)
    [^2]: [Exponential Moving Average on Streaming Data](https://dev.to/nestedsoftware/exponential-moving-average-on-streaming-data-4hhl)

    """

    # Note for devs, if you want look the pure python implementation here:
    # https://github.com/online-ml/river/blob/40c3190c9d05671ae4c2dc8b76c163ea53a45fb0/river/stats/ewvar.py
    def __init__(self, fading_factor=0.5):
        if not 0 <= fading_factor <= 1:
            raise ValueError("q is not comprised between 0 and 1")

        self.fading_factor = fading_factor
        self._ewvar = _rust_stats.RsEWVar(fading_factor)

    @property
    def name(self):
        return f"ewv_{self.fading_factor}"

    def update(self, x):
        self._ewvar.update(x)
        return self

    def get(self):
        return self._ewvar.get()


# river/stats/iqr.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class IQR(stats.base.Univariate):
    """Computes the interquartile range.

    Parameters
    ----------
    q_inf
        Desired inferior quantile, must be between 0 and 1. Defaults to `0.25`.
    q_sup
        Desired superior quantile, must be between 0 and 1. Defaults to `0.75`.

    Examples
    --------

    >>> from river import stats

    >>> iqr = stats.IQR(q_inf=0.25, q_sup=0.75)

    >>> for i in range(0, 1001):
    ...     iqr = iqr.update(i)
    ...     if i % 100 == 0:
    ...         print(iqr.get())
    0.0
    50.0
    100.0
    150.0
    200.0
    250.0
    300.0
    350.0
    400.0
    450.0
    500.0

    """

    # Note for devs, if you want look the pure python implementation here:
    # https://github.com/online-ml/river/blob/40c3190c9d05671ae4c2dc8b76c163ea53a45fb0/river/stats/iqr.py

    def __init__(self, q_inf=0.25, q_sup=0.75):
        super().__init__()
        if q_inf >= q_sup:
            raise ValueError("q_inf must be strictly less than q_sup")
        self.q_inf = q_inf
        self.q_sup = q_sup
        self._iqr = _rust_stats.RsIQR(self.q_inf, self.q_sup)
        self._is_updated = False

    @property
    def name(self):
        return f"{self.__class__.__name__}_{self.q_inf}_{self.q_sup}"

    def update(self, x):
        self._iqr.update(x)
        if not self._is_updated:
            self._is_updated = True
        return self

    def get(self):
        return self._iqr.get()

    def __repr__(self):
        # We surcharge this method to avoid this error on rust side:
        # pyo3_runtime.PanicException: index out of bounds: the len is 0 but the index is 0
        # This error is caused by the `get()` use before the update in the super method.
        value = None
        if self._is_updated:
            value = self.get()
        fmt_value = None if value is None else f"{value:{self._fmt}}".rstrip("0")
        return f"{self.__class__.__name__}: {fmt_value}"


class RollingIQR(stats.base.RollingUnivariate):
    """Computes the rolling interquartile range.

    Parameters
    ----------
    window_size
        Size of the window.
    q_inf
        Desired inferior quantile, must be between 0 and 1. Defaults to `0.25`.
    q_sup
        Desired superior quantile, must be between 0 and 1. Defaults to `0.75`.

    Examples
    --------

    >>> from river import stats
    >>> rolling_iqr = stats.RollingIQR(
    ...     q_inf=0.25,
    ...     q_sup=0.75,
    ...     window_size=101
    ... )

    >>> for i in range(0, 1001):
    ...     rolling_iqr = rolling_iqr.update(i)
    ...     if i % 100 == 0:
    ...         print(rolling_iqr.get())
    0.0
    50.0
    50.0
    50.0
    50.0
    50.0
    50.0
    50.0
    50.0
    50.0
    50.0

    """

    # Note for devs, if you want look the pure python implementation here:
    # https://github.com/online-ml/river/blob/40c3190c9d05671ae4c2dc8b76c163ea53a45fb0/river/stats/iqr.py
    def __init__(self, window_size: int, q_inf=0.25, q_sup=0.75):
        if q_inf >= q_sup:
            raise ValueError("q_inf must be strictly less than q_sup")
        self.q_inf = q_inf
        self.q_sup = q_sup
        self._rolling_iqr = _rust_stats.RsRollingIQR(q_inf, q_sup, window_size)
        self.window_size_value = window_size
        self._is_updated = False

    def update(self, x):
        self._rolling_iqr.update(x)
        if not self._is_updated:
            self._is_updated = True
        return self

    def get(self):
        # HACK: Avoid crash if get is called before update
        # panicked at 'index out of bounds: the len is 0 but the index is 0'
        if not self._is_updated:
            return None
        return self._rolling_iqr.get()

    @property
    def window_size(self):
        return self.window_size_value

    def __repr__(self):
        # We surcharge this method to avoid this error on rust side:
        # pyo3_runtime.PanicException: attempt to subtract with overflow
        # This error is caused by the `get()` use before the update in the super method.
        value = None
        if self._is_updated:
            value = self.get()
        fmt_value = None if value is None else f"{value:{self._fmt}}".rstrip("0")
        return f"{self.__class__.__name__}: {fmt_value}"


# river/stats/kurtosis.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class Kurtosis(stats.base.Univariate):
    """Running kurtosis using Welford's algorithm.

    Parameters
    ----------
    bias
        If `False`, then the calculations are corrected for statistical bias.

    Examples
    --------

    >>> from river import stats
    >>> import scipy.stats
    >>> import numpy as np

    >>> np.random.seed(42)
    >>> X = np.random.normal(loc=0, scale=1, size=10)

    >>> kurtosis = stats.Kurtosis(bias=False)
    >>> for x in X:
    ...     print(kurtosis.update(x).get())
    -3.0
    -2.0
    -1.5
    1.4130027920707047
    0.15367976585756438
    0.46142633246812653
    -1.620647789230658
    -1.3540178492487054
    -1.2310268787102745
    -0.9490372374384453

    >>> for i in range(2, len(X)+1):
    ...     print(scipy.stats.kurtosis(X[:i], bias=False))
    -2.0
    -1.4999999999999998
    1.4130027920707082
    0.15367976585756082
    0.46142633246812403
    -1.620647789230658
    -1.3540178492487063
    -1.2310268787102738
    -0.9490372374384459

    >>> kurtosis = stats.Kurtosis(bias=True)
    >>> for x in X:
    ...     print(kurtosis.update(x).get())
    -3.0
    -2.0
    -1.5
    -1.011599627723906
    -0.9615800585356089
    -0.6989395431537853
    -1.4252699121794408
    -1.311437071070812
    -1.246289111322894
    -1.082283689864171

    >>> for i in range(2, len(X)+1):
    ...     print(scipy.stats.kurtosis(X[:i], bias=True))
    -2.0
    -1.4999999999999998
    -1.0115996277239057
    -0.9615800585356098
    -0.6989395431537861
    -1.425269912179441
    -1.3114370710708125
    -1.2462891113228936
    -1.0822836898641714

    References
    ----------
    [^1]: [Wikipedia article on algorithms for calculating variance](https://www.wikiwand.com/en/Algorithms_for_calculating_variance#/Covariance)

    """

    def __init__(self, bias=False):
        super().__init__()
        self.bias = bias
        self._kurtosis = _rust_stats.RsKurtosis(bias)

    def update(self, x):
        self._kurtosis.update(x)
        return self

    def get(self):
        return self._kurtosis.get()


# river/stats/link.py

from __future__ import annotations

from river import stats


class Link(stats.base.Univariate):
    """A link joins two univariate statistics as a sequence.

    This can be used to pipe the output of one statistic to the input of another. This can be used,
    for instance, to calculate the mean of the variance of a variable. It can also be used to
    compute shifted statistics by piping statistics with an instance of `stats.Shift`.

    Note that a link is not meant to be instantiated via this class definition. Instead, users can
    link statistics together via the `|` operator.

    Parameters
    ----------
    left
    right
        The output from `left`'s `get` method is passed to `right`'s `update` method if `left`'s
        `get` method doesn't produce `None.`

    Examples
    --------

    >>> from river import stats
    >>> stat = stats.Shift(1) | stats.Mean()

    No values have been seen, therefore `get` defaults to the initial value of `stats.Mean`,
    which is 0.

    >>> stat.get()
    0.

    Let us now call `update`.

    >>> stat = stat.update(1)

    The output from `get` will still be 0. The reason is that `stats.Shift` has not enough
    values, and therefore outputs it's default value, which is `None`. The `stats.Mean`
    instance is therefore not updated.

    >>> stat.get()
    0.0

    On the next call to `update`, the `stats.Shift` instance has seen enough values, and
    therefore the mean can be updated. The mean is therefore equal to 1, because that's the
    only value from the past.

    >>> stat = stat.update(3)
    >>> stat.get()
    1.0

    On the subsequent call to update, the mean will be updated with the value 3.

    >>> stat = stat.update(4)
    >>> stat.get()
    2.0

    Note that composing statistics returns a new statistic with it's own name.

    >>> stat.name
    'mean_of_shift_1'

    """

    def __init__(self, left: stats.base.Univariate, right: stats.base.Univariate):
        self.left = left
        self.right = right

    def update(self, x):
        self.left.update(x)
        y = self.left.get()
        if y is not None:
            self.right.update(y)
        return self

    def get(self):
        return self.right.get()

    @property
    def name(self):
        return f"{self.right.name}_of_{self.left.name}"

    def __repr__(self):
        return repr(self.right)


# river/stats/mad.py

from __future__ import annotations

from . import quantile


class MAD(quantile.Quantile):
    """Median Absolute Deviation (MAD).

    The median absolute deviation is the median of the absolute differences between each data point
    and the data's overall median. In an online setting, the median of the data is unknown
    beforehand. Therefore, both the median of the data and the median of the differences of the
    data with respect to the latter are updated online. To be precise, the median of the data is
    updated before the median of the differences. As a consequence, this online version of the MAD
    does not coincide exactly with its batch counterpart.

    Examples
    --------

    >>> from river import stats

    >>> X = [4, 2, 5, 3, 0, 4]

    >>> mad = stats.MAD()
    >>> for x in X:
    ...     print(mad.update(x).get())
    0.0
    2.0
    1.0
    1.0
    1.0
    1.0

    Attributes
    ----------
    median : stats.Median
        The median of the data.

    References
    ----------
    [^1]: [Median absolute deviation article on Wikipedia](https://www.wikiwand.com/en/Median_absolute_deviation)

    """

    #
    def __init__(self):
        super().__init__(q=0.5)
        self.median = quantile.Quantile(q=0.5)

    def update(self, x):
        self.median.update(x)
        super().update(abs(x - self.median.get()))
        return self


# river/stats/maximum.py

from __future__ import annotations

import math

from river import stats, utils


class Max(stats.base.Univariate):
    """Running max.

    Attributes
    ----------
    max : float
        The current max.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 5, -6]
    >>> _max = stats.Max()
    >>> for x in X:
    ...     print(_max.update(x).get())
    1
    1
    3
    3
    5
    5

    """

    def __init__(self):
        self.max = -math.inf

    def update(self, x):
        if x > self.max:
            self.max = x
        return self

    def get(self):
        return self.max


class RollingMax(stats.base.RollingUnivariate):
    """Running max over a window.

    Parameters
    ----------
    window_size
        Size of the rolling window.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 2, 1]
    >>> rolling_max = stats.RollingMax(window_size=2)
    >>> for x in X:
    ...     print(rolling_max.update(x).get())
    1
    1
    3
    3
    2
    2

    """

    def __init__(self, window_size: int):
        self.window = utils.SortedWindow(size=window_size)

    @property
    def window_size(self):
        return self.window.size

    def update(self, x):
        self.window.append(x)
        return self

    def get(self):
        try:
            return self.window[-1]
        except IndexError:
            return None


class AbsMax(stats.base.Univariate):
    """Running absolute max.

    Attributes
    ----------
    abs_max : float
        The current absolute max.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 5, -6]
    >>> abs_max = stats.AbsMax()
    >>> for x in X:
    ...     print(abs_max.update(x).get())
    1
    4
    4
    4
    5
    6

    """

    def __init__(self):
        self.abs_max = 0.0

    def update(self, x):
        if abs(x) > self.abs_max:
            self.abs_max = abs(x)
        return self

    def get(self):
        return self.abs_max


class RollingAbsMax(stats.base.RollingUnivariate):
    """Running absolute max over a window.

    Parameters
    ----------
    window_size
        Size of the rolling window.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 2, 1]
    >>> rolling_absmax = stats.RollingAbsMax(window_size=2)
    >>> for x in X:
    ...     print(rolling_absmax.update(x).get())
    1
    4
    4
    3
    2
    2

    """

    def __init__(self, window_size: int):
        self.window = utils.SortedWindow(size=window_size)

    @property
    def window_size(self):
        return self.window.size

    def update(self, x):
        self.window.append(abs(x))
        return self

    def get(self):
        try:
            return self.window[-1]
        except IndexError:
            return None


# river/stats/mean.py

from __future__ import annotations

import copy

import numpy as np

from river import stats


class Mean(stats.base.Univariate):
    """Running mean.

    Attributes
    ----------
    n : float
        The current sum of weights. If each passed weight was 1, then this is equal to the number
        of seen observations.

    Examples
    --------

    >>> from river import stats

    >>> X = [-5, -3, -1, 1, 3, 5]
    >>> mean = stats.Mean()
    >>> for x in X:
    ...     print(mean.update(x).get())
    -5.0
    -4.0
    -3.0
    -2.0
    -1.0
    0.0

    You can calculate a rolling average by wrapping a `utils.Rolling` around:

    >>> from river import utils

    >>> X = [1, 2, 3, 4, 5, 6]
    >>> rmean = utils.Rolling(stats.Mean(), window_size=2)

    >>> for x in X:
    ...     print(rmean.update(x).get())
    1.0
    1.5
    2.5
    3.5
    4.5
    5.5

    References
    ----------
    [^1]: [West, D. H. D. (1979). Updating mean and variance estimates: An improved method. Communications of the ACM, 22(9), 532-535.](https://dl.acm.org/doi/10.1145/359146.359153)
    [^2]: [Finch, T., 2009. Incremental calculation of weighted mean and variance. University of Cambridge, 4(11-5), pp.41-42.](https://fanf2.user.srcf.net/hermes/doc/antiforgery/stats.pdf)
    [^3]: [Chan, T.F., Golub, G.H. and LeVeque, R.J., 1983. Algorithms for computing the sample variance: Analysis and recommendations. The American Statistician, 37(3), pp.242-247.](https://amstat.tandfonline.com/doi/abs/10.1080/00031305.1983.10483115)

    """

    def __init__(self):
        self.n = 0
        self._mean = 0.0

    def update(self, x, w=1.0):
        self.n += w
        self._mean += (w / self.n) * (x - self._mean)
        return self

    def update_many(self, X: np.ndarray):
        a = self.n / (self.n + len(X))
        b = len(X) / (self.n + len(X))
        self._mean = a * self._mean + b * np.mean(X)
        self.n += len(X)
        return self

    def revert(self, x, w=1.0):
        self.n -= w
        if self.n < 0:
            raise ValueError("Cannot go below 0")
        elif self.n == 0:
            self._mean = 0.0
        else:
            self._mean -= (w / self.n) * (x - self._mean)
        return self

    def get(self):
        return self._mean

    @classmethod
    def _from_state(cls, n, mean):
        new = cls()
        new.n = n
        new._mean = mean

        return new

    def __iadd__(self, other):
        old_n = self.n
        self.n += other.n
        self._mean = (old_n * self._mean + other.n * other.get()) / self.n
        return self

    def __add__(self, other):
        result = copy.deepcopy(self)
        result += other
        return result

    def __isub__(self, other):
        old_n = self.n
        self.n -= other.n

        if self.n > 0:
            self._mean = (old_n * self._mean - other.n * other._mean) / self.n
        else:
            self.n = 0.0
            self._mean = 0.0
        return self

    def __sub__(self, other):
        result = copy.deepcopy(self)
        result -= other
        return result


class BayesianMean(stats.base.Univariate):
    """Estimates a mean using outside information.

    Parameters
    ----------
    prior
    prior_weight

    References
    ----------
    [^1]: [Additive smoothing](https://www.wikiwand.com/en/Additive_smoothing)
    [^2]: [Bayesian average](https://www.wikiwand.com/en/Bayesian_average)
    [^3]: [Practical example of Bayes estimators](https://www.wikiwand.com/en/Bayes_estimator#/Practical_example_of_Bayes_estimators)

    """

    def __init__(self, prior: float, prior_weight: float):
        self.prior = prior
        self.prior_weight = prior_weight
        self._mean = Mean()

    @property
    def name(self):
        return "bayes_mean"

    def update(self, x):
        self._mean.update(x)
        return self

    def revert(self, x):
        self._mean.revert(x)
        return self

    def get(self):
        # Uses the notation from https://www.wikiwand.com/en/Bayes_estimator#/Practical_example_of_Bayes_estimators
        R = self._mean.get()
        v = self._mean.n
        m = self.prior_weight
        C = self.prior

        return (R * v + C * m) / (v + m)


# river/stats/minimum.py

from __future__ import annotations

import math

from river import stats, utils


class Min(stats.base.Univariate):
    """Running min.

    Attributes
    ----------
    min : float
        The current min.

    """

    def __init__(self):
        self.min = math.inf

    def update(self, x):
        if x < self.min:
            self.min = x
        return self

    def get(self):
        return self.min


class RollingMin(stats.base.RollingUnivariate):
    """Running min over a window.

    Parameters
    ----------
    window_size
        Size of the rolling window.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 2, 1]
    >>> rolling_min = stats.RollingMin(2)
    >>> for x in X:
    ...     print(rolling_min.update(x).get())
    1
    -4
    -4
    -2
    -2
    1

    """

    def __init__(self, window_size: int):
        self.window = utils.SortedWindow(size=window_size)

    @property
    def window_size(self):
        return self.window.size

    def update(self, x):
        self.window.append(x)
        return self

    def get(self):
        try:
            return self.window[0]
        except IndexError:
            return None


# river/stats/mode.py

from __future__ import annotations

import collections
import numbers
import typing

from river import stats

__all__ = ["Mode"]


class Mode(stats.base.Univariate):
    """Running mode.

    The mode is simply the most common value. An approximate mode can be computed by setting the
    number of first unique values to count.

    Parameters
    ----------
    k
        Only the first `k` unique values will be included. If `k` equals -1, the exact mode is
        computed.

    Examples
    --------

    >>> from river import stats

    >>> X = ['sunny', 'cloudy', 'cloudy', 'rainy', 'rainy', 'rainy']
    >>> mode = stats.Mode(k=2)
    >>> for x in X:
    ...     print(mode.update(x).get())
    sunny
    sunny
    cloudy
    cloudy
    cloudy
    cloudy

    >>> mode = stats.Mode(k=-1)
    >>> for x in X:
    ...     print(mode.update(x).get())
    sunny
    sunny
    cloudy
    cloudy
    cloudy
    rainy

    """

    def __init__(self, k=25):
        self.k = k
        self.counts = collections.defaultdict(int)

    @property
    def name(self):
        return "mode"

    def update(self, x):
        if self.k == -1 or x in self.counts or len(self.counts) < self.k:
            self.counts[x] += 1
        return self

    def get(self):
        return max(self.counts, key=self.counts.get, default=None)


class RollingMode(stats.base.RollingUnivariate):
    """Running mode over a window.

    The mode is the most common value.

    Parameters
    ----------
    window_size
        Size of the rolling window.

    Attributes
    ----------
    counts : collections.defaultdict
        Value counts.

    Examples
    --------

    >>> from river import stats

    >>> X = ['sunny', 'sunny', 'sunny', 'rainy', 'rainy', 'rainy', 'rainy']
    >>> rolling_mode = stats.RollingMode(window_size=2)
    >>> for x in X:
    ...     print(rolling_mode.update(x).get())
    sunny
    sunny
    sunny
    sunny
    rainy
    rainy
    rainy

    >>> rolling_mode = stats.RollingMode(window_size=5)
    >>> for x in X:
    ...     print(rolling_mode.update(x).get())
    sunny
    sunny
    sunny
    sunny
    sunny
    rainy
    rainy

    """

    def __init__(self, window_size: int):
        self.window: collections.deque[numbers.Number] = collections.deque(maxlen=window_size)
        self.counts: collections.defaultdict[typing.Any, int] = collections.defaultdict(int)

    @property
    def window_size(self):
        return self.window.maxlen

    def update(self, x):
        if len(self.window) >= self.window_size:
            # Subtract the counter of the last element
            first_in = self.window[0]
            self.counts[first_in] -= 1

            # No need to store the value if it's counter is 0
            if self.counts[first_in] == 0:
                self.counts.pop(first_in)

        self.counts[x] += 1
        self.window.append(x)
        return self

    def get(self):
        return max(self.counts, key=self.counts.get, default=None)


# river/stats/moments.py

from __future__ import annotations

from river import stats


class CentralMoments(stats.base.Univariate):
    """Computes central moments using Welford's algorithm.

    Attributes
    ----------
    count : stats.Count
    delta: float
    sum_delta : float
        Mean of sum of differences.
    M1 : float
        Sums of powers of differences from the mean order 1.
    M2 : float)
        Sums of powers of differences from the mean order 2.
    M3 : float
        Sums of powers of differences from the mean order 3.
    M4 : float
        Sums of powers of differences from the mean order 4.

    References
    ----------
    [^1]: [Wikipedia article on algorithms for calculating variance](https://www.wikiwand.com/en/Algorithms_for_calculating_variance#/Covariance)

    """

    def __init__(self):
        self.count = stats.Count()

        self.delta = 0
        self.sum_delta = 0

        self.M1 = 0
        self.M2 = 0
        self.M3 = 0
        self.M4 = 0

    def _update_delta(self, x):
        self.delta = (x - self.sum_delta) / self.count.get()
        return self

    def _update_sum_delta(self):
        self.sum_delta += self.delta
        return self

    def _update_m1(self, x):
        self.M1 = (x - self.sum_delta) * self.delta * (self.count.get() - 1)
        return self

    def _update_m2(self):
        self.M2 += self.M1
        return self

    def _update_m3(self):
        self.M3 += self.M1 * self.delta * (self.count.get() - 2) - 3 * self.delta * self.M2
        return self

    def _update_m4(self):
        delta_square = self.delta**2
        self.M4 += (
            self.M1 * delta_square * (self.count.get() ** 2 - 3 * self.count.get() + 3)
            + 6 * delta_square * self.M2
            - 4 * self.delta * self.M3
        )
        return self


# river/stats/n_unique.py

from __future__ import annotations

import hashlib
import math

import numpy as np

from river import stats


class NUnique(stats.base.Univariate):
    """Approximate number of unique values counter.

    This is basically an implementation of the HyperLogLog algorithm. Adapted from
    [`hypy`](https://github.com/clarkduvall/hypy). The code is a bit too terse but it will do for
    now.

    Parameters
    ----------
    error_rate
        Desired error rate. Memory usage is inversely proportional to this value.
    seed
        Set the seed to produce identical results.

    Attributes
    ----------
    n_bits : int
    n_buckets : int
    buckets : list

    Examples
    --------

    >>> import string
    >>> from river import stats

    >>> alphabet = string.ascii_lowercase
    >>> n_unique = stats.NUnique(error_rate=0.2, seed=42)

    >>> n_unique.update('a').get()
    1

    >>> n_unique.update('b').get()
    2

    >>> for letter in alphabet:
    ...     n_unique = n_unique.update(letter)
    >>> n_unique.get()
    31

    Lowering the `error_rate` parameter will increase the precision.

    >>> n_unique = stats.NUnique(error_rate=0.01, seed=42)
    >>> for letter in alphabet:
    ...     n_unique = n_unique.update(letter)
    >>> n_unique.get()
    26

    References
    ----------
    [^1]: [My favorite algorithm (and data structure): HyperLogLog](https://odino.org/my-favorite-data-structure-hyperloglog/)
    [^2]: [Flajolet, P., Fusy, É., Gandouet, O. and Meunier, F., 2007, June. Hyperloglog: the analysis of a near-optimal cardinality estimation algorithm.](http://algo.inria.fr/flajolet/Publications/FlFuGaMe07.pdf)

    """

    P32 = 2**32

    def __init__(self, error_rate=0.01, seed: int | None = None):
        self.error_rate = error_rate
        self.seed = seed

        self.n_bits = int(math.ceil(math.log((1.04 / error_rate) ** 2, 2)))
        self.n_buckets = 1 << self.n_bits
        self.buckets = [0] * self.n_buckets
        self._salt = np.random.RandomState(seed).bytes(hashlib.blake2s.SALT_SIZE)

    @property
    def name(self):
        return "n_unique"

    def _hash(self, x):
        hexa = hashlib.blake2s(bytes(x, encoding="utf8"), salt=self._salt).hexdigest()
        return int(hexa, 16)

    def update(self, x):
        x = self._hash(x)
        i = x & NUnique.P32 - 1 >> 32 - self.n_bits
        z = 35 - len(bin(NUnique.P32 - 1 & x << self.n_bits | 1 << self.n_bits - 1))
        self.buckets[i] = max(self.buckets[i], z)
        return self

    def get(self):
        a = (
            {16: 0.673, 32: 0.697, 64: 0.709}[self.n_buckets]
            if self.n_buckets <= 64
            else 0.7213 / (1 + 1.079 / self.n_buckets)
        )
        e = a * self.n_buckets * self.n_buckets / sum(1.0 / (1 << x) for x in self.buckets)
        if e <= self.n_buckets * 2.5:
            z = len([r for r in self.buckets if not r])
            return int(self.n_buckets * math.log(float(self.n_buckets) / z) if z else e)
        return int(e if e < NUnique.P32 / 30 else -NUnique.P32 * math.log(1 - e / NUnique.P32))


# river/stats/pearson.py

from __future__ import annotations

from river import stats


class PearsonCorr(stats.base.Bivariate):
    """Online Pearson correlation.

    Parameters
    ----------
    ddof
        Delta Degrees of Freedom.

    Attributes
    ----------
    var_x : stats.Var
        Running variance of `x`.
    var_y : stats.Var
        Running variance of `y`.
    cov_xy : stats.Cov
        Running covariance of `x` and `y`.

    Examples
    --------

    >>> from river import stats

    >>> x = [0, 0, 0, 1, 1, 1, 1]
    >>> y = [0, 1, 2, 3, 4, 5, 6]

    >>> pearson = stats.PearsonCorr()

    >>> for xi, yi in zip(x, y):
    ...     print(pearson.update(xi, yi).get())
    0
    0
    0
    0.774596
    0.866025
    0.878310
    0.866025

    You can also do this in a rolling fashion:

    >>> from river import utils

    >>> x = [0, 0, 0, 1, 1, 1, 1]
    >>> y = [0, 1, 2, 3, 4, 5, 6]

    >>> pearson = utils.Rolling(stats.PearsonCorr(), window_size=4)

    >>> for xi, yi in zip(x, y):
    ...     print(pearson.update(xi, yi).get())
    0
    0
    0
    0.7745966692414834
    0.8944271909999159
    0.7745966692414832
    -4.712160915387242e-09

    """

    def __init__(self, ddof=1):
        self.var_x = stats.Var(ddof=ddof)
        self.var_y = stats.Var(ddof=ddof)
        self.cov_xy = stats.Cov(ddof=ddof)

    @property
    def ddof(self):
        return self.cov_xy.ddof

    def update(self, x, y):
        self.var_x.update(x)
        self.var_y.update(y)
        self.cov_xy.update(x, y)
        return self

    def revert(self, x, y):
        self.var_x.revert(x)
        self.var_y.revert(y)
        self.cov_xy.revert(x, y)
        return self

    def get(self):
        var_x = self.var_x.get()
        var_y = self.var_y.get()
        if var_x and var_y:
            return self.cov_xy.get() / (var_x * var_y) ** 0.5
        return 0


# river/stats/ptp.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class PeakToPeak(stats.base.Univariate):
    """Running peak to peak (max - min).

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 2, 4]
    >>> ptp = stats.PeakToPeak()
    >>> for x in X:
    ...     print(ptp.update(x).get())
    0.
    5.
    7.
    7.
    7.
    8.

    """

    def __init__(self):
        self._ptp = _rust_stats.RsPeakToPeak()
        self._is_updated = False

    @property
    def name(self):
        return "ptp"

    def update(self, x):
        self._ptp.update(x)
        if not self._is_updated:
            self._is_updated = True
        return self

    def get(self):
        if not self._is_updated:
            return 0.0
        return self._ptp.get()


class RollingPeakToPeak(stats.base.RollingUnivariate):
    """Running peak to peak (max - min) over a window.

    Parameters
    ----------
    window_size
        Size of the rolling window.

    Attributes
    ----------
    max : stats.RollingMax
        The running rolling max.
    min : stats.RollingMin
        The running rolling min.

    Examples
    --------

    >>> from river import stats

    >>> X = [1, -4, 3, -2, 2, 1]
    >>> ptp = stats.RollingPeakToPeak(window_size=2)
    >>> for x in X:
    ...     print(ptp.update(x).get())
    0
    5
    7
    5
    4
    1

    """

    def __init__(self, window_size: int):
        self.max = stats.RollingMax(window_size)
        self.min = stats.RollingMin(window_size)

    @property
    def window_size(self):
        return self.max.window_size

    def update(self, x):
        self.max.update(x)
        self.min.update(x)
        return self

    def get(self):
        maximum = self.max.get()
        if maximum is None:
            return None
        minimum = self.min.get()
        if minimum is None:
            return None
        return maximum - minimum


# river/stats/quantile.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class Quantile(stats.base.Univariate):
    """Running quantile.

    Uses the P² algorithm, which is also known as the "Piecewise-Parabolic quantile estimator".
    The code is inspired by LiveStat's implementation [^2].

    Parameters
    ----------
    q
        Determines which quantile to compute, must be comprised between 0 and 1.

    Examples
    --------

    >>> from river import stats
    >>> import numpy as np

    >>> np.random.seed(42 * 1337)
    >>> mu, sigma = 0, 1
    >>> s = np.random.normal(mu, sigma, 500)

    >>> median = stats.Quantile(0.5)
    >>> for x in s:
    ...    _ = median.update(x)
    >>> print(f'The estimated value of the 50th (median) quantile is {median.get():.4f}')
    The estimated value of the 50th (median) quantile is -0.0275

    >>> print(f'The real value of the 50th (median) quantile is {np.median(s):.4f}')
    The real value of the 50th (median) quantile is -0.0135

    >>> percentile_17 = stats.Quantile(0.17)
    >>> for x in s:
    ...    _ = percentile_17.update(x)
    >>> print(f'The estimated value of the 17th quantile is {percentile_17.get():.4f}')
    The estimated value of the 17th quantile is -0.8652

    >>> print(f'The real value of the 17th quantile is {np.percentile(s,17):.4f}')
    The real value of the 17th quantile is -0.9072

    References
    ----------
    [^1]: [The P² Algorithm for Dynamic Univariateal Computing Calculation of Quantiles and Editor Histograms Without Storing Observations](https://www.cse.wustl.edu/~jain/papers/ftp/psqr.pdf)
    [^2]: [LiveStats](https://github.com/cxxr/LiveStats)
    [^3]: [P² quantile estimator: estimating the median without storing values](https://aakinshin.net/posts/p2-quantile-estimator/)

    """

    # Note for devs, if you want look the pure python implementation here:
    # https://github.com/online-ml/river/blob/40c3190c9d05671ae4c2dc8b76c163ea53a45fb0/river/stats/quantile.py
    def __init__(self, q: float = 0.5):
        super().__init__()
        if not 0 < q < 1:
            raise ValueError("q is not comprised between 0 and 1")
        self._quantile = _rust_stats.RsQuantile(q)
        self._is_updated = False
        self.q = q  # Used by anomaly.QuantileFilter

    def update(self, x):
        self._quantile.update(x)
        if not self._is_updated:
            self._is_updated = True
        return self

    def get(self):
        # HACK: Avoid this following error in `QuantileFilter`
        # panicked at 'index out of bounds: the len is 0 but the index is 0'
        if not self._is_updated:
            return None
        return self._quantile.get()

    def __repr__(self):
        # We surcharge this method to avoid this error on rust side:
        # pyo3_runtime.PanicException: index out of bounds: the len is 0 but the index is 0
        # This error is caused by the `get()` use before the update in the super method.
        value = None
        if self._is_updated:
            value = self.get()
        fmt_value = None if value is None else f"{value:{self._fmt}}".rstrip("0")
        return f"{self.__class__.__name__}: {fmt_value}"


class RollingQuantile(stats.base.RollingUnivariate):
    """Running quantile over a window.

    Parameters
    ----------
    q
        Determines which quantile to compute, must be comprised between 0 and 1.
    window_size
        Size of the window.

    Examples
    --------

    >>> from river import stats

    >>> rolling_quantile = stats.RollingQuantile(
    ...     q=.5,
    ...     window_size=101,
    ... )

    >>> for i in range(1001):
    ...     rolling_quantile = rolling_quantile.update(i)
    ...     if i % 100 == 0:
    ...         print(rolling_quantile.get())
    0.0
    50.0
    150.0
    250.0
    350.0
    450.0
    550.0
    650.0
    750.0
    850.0
    950.0

    References
    ----------
    [^1]: [Left sorted](https://stackoverflow.com/questions/8024571/insert-an-item-into-sorted-list-in-python)

    """

    # Note for devs, if you want look the pure python implementation here:
    # https://github.com/online-ml/river/blob/40c3190c9d05671ae4c2dc8b76c163ea53a45fb0/river/stats/quantile.py
    def __init__(self, q: float, window_size: int):
        super().__init__()
        if not 0 <= q <= 1:
            raise ValueError("q is not comprised between 0 and 1")
        self._rolling_quantile = _rust_stats.RsRollingQuantile(q, window_size)
        self.window_size_value = window_size
        self._is_updated = False

    def update(self, x):
        self._rolling_quantile.update(x)
        if not self._is_updated:
            self._is_updated = True
        return self

    def get(self):
        if not self._is_updated:
            return None
        return self._rolling_quantile.get()

    @property
    def window_size(self):
        return self.window_size_value

    def __repr__(self):
        # We surcharge this method to avoid this error on rust side:
        # pyo3_runtime.PanicException: attempt to subtract with overflow
        # This error is caused by the `get()` use before the update in the super method.
        value = None
        if self._is_updated:
            value = self.get()
        fmt_value = None if value is None else f"{value:{self._fmt}}".rstrip("0")
        return f"{self.__class__.__name__}: {fmt_value}"


# river/stats/sem.py

from __future__ import annotations

from . import var


class SEM(var.Var):
    """Running standard error of the mean using Welford's algorithm.

    Parameters
    ----------
    ddof
        Delta Degrees of Freedom. The divisor used in calculations is `n - ddof`, where `n` is the
        number of seen elements.

    Attributes
    ----------
    n : int
        Number of observations.

    Examples
    --------

    >>> from river import stats

    >>> X = [3, 5, 4, 7, 10, 12]

    >>> sem = stats.SEM()
    >>> for x in X:
    ...     print(sem.update(x).get())
    0.0
    1.0
    0.577350
    0.853912
    1.240967
    1.447219

    >>> from river import utils

    >>> X = [1, 4, 2, -4, -8, 0]

    >>> rolling_sem = utils.Rolling(stats.SEM(ddof=1), window_size=3)
    >>> for x in X:
    ...     print(rolling_sem.update(x).get())
    0.0
    1.5
    0.881917
    2.403700
    2.905932
    2.309401

    References
    ----------
    [^1]: [Wikipedia article on algorithms for calculating variance](https://www.wikiwand.com/en/Algorithms_for_calculating_variance#/Covariance)

    """

    def get(self):
        try:
            return (super().get() / self.mean.n) ** 0.5
        except ZeroDivisionError:
            return None


# river/stats/shift.py

from __future__ import annotations

import collections

from river import stats


class Shift(stats.base.Univariate):
    """Shifts a data stream by returning past values.

    This can be used to compute statistics over past data. For instance, if you're computing daily
    averages, then shifting by 7 will be equivalent to computing averages from a week ago.

    Shifting values is useful when you're calculating an average over a target value. Indeed,
    in this case it's important to shift the values in order not to introduce leakage. The
    recommended way to do this is to `feature_extraction.TargetAgg`, which already takes care
    of shifting the target values once.

    Parameters
    ----------
    amount
        Shift amount. The `get` method will return the `t - amount` value, where `t` is the
        current moment.
    fill_value
        This value will be returned by the `get` method if not enough values have been observed.

    Examples
    --------

    It is rare to have to use `Shift` by itself. A more common usage is to compose it with
    other statistics. This can be done via the `|` operator.

    >>> from river import stats

    >>> stat = stats.Shift(1) | stats.Mean()

    >>> for i in range(5):
    ...     stat = stat.update(i)
    ...     print(stat.get())
    0.0
    0.0
    0.5
    1.0
    1.5

    A common usecase for using `Shift` is when computing statistics on shifted data. For
    instance, say you have a dataset which records the amount of sales for a set of shops. You
    might then have a `shop` field and a `sales` field. Let's say you want to look at the
    average amount of sales per shop. You can do this by using a `feature_extraction.Agg`. When
    you call `transform_one`, you're expecting it to return the average amount of sales,
    *without* including today's sales. You can do this by prepending an instance of
    `stats.Mean` with an instance of `stats.Shift`.

    >>> from river import feature_extraction

    >>> agg = feature_extraction.Agg(
    ...     on='sales',
    ...     how=stats.Shift(1) | stats.Mean(),
    ...     by='shop'
    ... )

    Let's define a little example dataset.

    >>> X = iter([
    ...     {'shop': 'Ikea', 'sales': 10},
    ...     {'shop': 'Ikea', 'sales': 15},
    ...     {'shop': 'Ikea', 'sales': 20}
    ... ])

    Now let's call the `learn_one` method to update our feature extractor.

    >>> x = next(X)
    >>> agg = agg.learn_one(x)

    At this point, the average defaults to the initial value of `stats.Mean`, which is 0.

    >>> agg.transform_one(x)
    {'sales_mean_of_shift_1_by_shop': 0.0}

    We can now update our feature extractor with the next data point and check the output.

    >>> agg = agg.learn_one(next(X))
    >>> agg.transform_one(x)
    {'sales_mean_of_shift_1_by_shop': 10.0}

    >>> agg = agg.learn_one(next(X))
    >>> agg.transform_one(x)
    {'sales_mean_of_shift_1_by_shop': 12.5}

    """

    def __init__(self, amount=1, fill_value=None):
        self.amount = amount
        self.fill_value = fill_value
        self.buffer = collections.deque(maxlen=self.amount + 1)

    def update(self, x):
        self.buffer.append(x)
        return self

    def get(self):
        try:
            return self.buffer[-self.amount - 1]
        except IndexError:
            return self.fill_value

    @property
    def name(self):
        return f"shift_{self.amount}"


# river/stats/skew.py

from __future__ import annotations

from river import stats
from river.stats import _rust_stats


class Skew(stats.base.Univariate):
    """Running skew using Welford's algorithm.

    Parameters
    ----------
    bias
        If `False`, then the calculations are corrected for statistical bias.

    Examples
    --------

    >>> from river import stats
    >>> import numpy as np

    >>> np.random.seed(42)
    >>> X = np.random.normal(loc=0, scale=1, size=10)

    >>> skew = stats.Skew(bias=False)
    >>> for x in X:
    ...     print(skew.update(x).get())
    0.0
    0.0
    -1.4802398132849872
    0.5127437186677888
    0.7803466510704751
    1.056115628922055
    0.5057840774320389
    0.3478402420400934
    0.4536710660918704
    0.4123070197493227

    >>> skew = stats.Skew(bias=True)
    >>> for x in X:
    ...     print(skew.update(x).get())
    0.0
    0.0
    -0.6043053732501439
    0.2960327239981376
    0.5234724473423674
    0.7712778043924866
    0.39022088752624845
    0.278892645224261
    0.37425953513864063
    0.3476878073823696

    References
    ----------
    [^1]: [Wikipedia article on algorithms for calculating variance](https://www.wikiwand.com/en/Algorithms_for_calculating_variance#/Covariance)

    """

    def __init__(self, bias=False):
        super().__init__()
        self.bias = bias
        self._skew = _rust_stats.RsSkew(bias)

    @property
    def name(self):
        return "skew"

    def update(self, x):
        self._skew.update(x)
        return self

    def get(self):
        return self._skew.get()


# river/stats/summing.py

from __future__ import annotations

from river import stats


class Sum(stats.base.Univariate):
    """Running sum.

    Attributes
    ----------
    sum : float
        The running sum.

    Examples
    --------

    >>> from river import stats

    >>> X = [-5, -3, -1, 1, 3, 5]
    >>> mean = stats.Sum()
    >>> for x in X:
    ...     print(mean.update(x).get())
    -5.0
    -8.0
    -9.0
    -8.0
    -5.0
    0.0

    >>> from river import utils

    >>> X = [1, -4, 3, -2, 2, 1]
    >>> rolling_sum = utils.Rolling(stats.Sum(), window_size=2)
    >>> for x in X:
    ...     print(rolling_sum.update(x).get())
    1.0
    -3.0
    -1.0
    1.0
    0.0
    3.0

    """

    def __init__(self):
        self.sum = 0.0

    def update(self, x):
        self.sum += x
        return self

    def revert(self, x):
        self.sum -= x
        return self

    def get(self):
        return self.sum


# river/stats/test_parallel.py

from __future__ import annotations

import copy
import math
import random

import numpy as np
import pytest

from river import stats


def test_add_mean():
    A = stats.Mean()
    B = stats.Mean()
    C = stats.Mean()

    X = [random.random() for _ in range(30)]
    Y = [random.random() for _ in range(30)]
    W = [random.random() for _ in range(30)]

    for i, (x, y, w) in enumerate(zip(X, Y, W)):
        A.update(x, w)
        B.update(y, w)
        C.update(x, w).update(y, w)

        D = A + B
        assert math.isclose(C.get(), D.get())

        # TODO: clone should work instead of deepcopy
        E = copy.deepcopy(A)
        E += B
        assert math.isclose(C.get(), E.get())

        assert math.isclose(
            C.get(), np.average(X[: i + 1] + Y[: i + 1], weights=W[: i + 1] + W[: i + 1])
        )


def _weighted_var(X, W, ddof):
    average = np.average(X, weights=W)
    return np.average((X - average) ** 2, weights=W) * len(X) / (len(X) - ddof)


@pytest.mark.parametrize("ddof", [pytest.param(ddof, id=f"{ddof=}") for ddof in [0, 1, 2]])
def test_add_var(ddof):
    A = stats.Var(ddof)
    B = stats.Var(ddof)
    C = stats.Var(ddof)

    X = [random.random() for _ in range(30)]
    Y = [random.random() for _ in range(30)]
    W = [1 for _ in range(30)]

    for i, (x, y, w) in enumerate(zip(X, Y, W)):
        A.update(x, w)
        B.update(y, w)
        C.update(x, w).update(y, w)

        D = A + B
        assert math.isclose(C.get(), D.get())

        E = A.clone(include_attributes=True)
        E += B
        assert math.isclose(C.get(), E.get())

        if i >= ddof:
            assert math.isclose(
                C.get(), _weighted_var(X[: i + 1] + Y[: i + 1], W[: i + 1] + W[: i + 1], ddof=ddof)
            )


@pytest.mark.parametrize(
    "stat",
    [
        pytest.param(stat, id=stat.__class__.__name__)
        for stat in [stats.Mean(), stats.Var(ddof=0), stats.Var(ddof=1), stats.Var(ddof=2)]
    ],
)
def test_sub(stat):
    A = copy.deepcopy(stat)
    B = copy.deepcopy(stat)
    C = copy.deepcopy(stat)

    X = [random.random() for _ in range(30)]
    Y = [random.random() for _ in range(30)]
    W = [random.random() for _ in range(30)]

    for x, y, w in zip(X, Y, W):
        A.update(x, w)
        B.update(y, w)
        C.update(x, w).update(y, w)

    D = C - B
    assert math.isclose(D.get(), A.get())

    C -= B
    assert math.isclose(C.get(), A.get())


@pytest.mark.parametrize(
    "stat",
    [
        pytest.param(stat, id=stat.__class__.__name__)
        for stat in [stats.Mean(), stats.Var(ddof=0), stats.Var(ddof=1), stats.Var(ddof=2)]
    ],
)
def test_sub_back_to_zero(stat):
    A = copy.deepcopy(stat)
    B = copy.deepcopy(stat)
    C = copy.deepcopy(stat)

    x = random.random()
    A.update(x)
    B.update(x)

    D = A - B
    assert math.isclose(D.get(), C.get())


@pytest.mark.parametrize(
    "ddof",
    [pytest.param(ddof, id=f"{ddof=}") for ddof in [0, 1]],
)
def test_add_cov(ddof):
    c1 = stats.Cov(ddof=ddof)
    c2 = stats.Cov(ddof=ddof)

    X = [random.random() for _ in range(30)]
    Y = [random.random() for _ in range(30)]

    W = [random.random() for _ in range(30)]
    Z = [random.random() for _ in range(30)]

    for i, (x, y, w, z) in enumerate(zip(X, Y, W, Z)):
        c1.update(x, y)
        c2.update(w, z)
        if i >= c1.ddof:
            assert math.isclose(
                c1.get(),
                np.cov(X[: i + 1], Y[: i + 1], ddof=ddof)[0, 1],
                abs_tol=1e-10,
            )
            assert math.isclose(
                c2.get(),
                np.cov(W[: i + 1], Z[: i + 1], ddof=ddof)[0, 1],
                abs_tol=1e-10,
            )
            assert math.isclose(
                (c1 + c2).get(),
                np.cov(X[: i + 1] + W[: i + 1], Y[: i + 1] + Z[: i + 1], ddof=ddof)[0, 1],
                abs_tol=1e-10,
            )

    c1 += c2
    assert math.isclose(c1.get(), np.cov(X + W, Y + Z, ddof=ddof)[0, 1], abs_tol=1e-10)


@pytest.mark.parametrize(
    "ddof",
    [pytest.param(ddof, id=f"{ddof=}") for ddof in [0, 1]],
)
def test_sub_cov(ddof):
    c1 = stats.Cov(ddof=ddof)
    c2 = stats.Cov(ddof=ddof)

    X = [random.random() for _ in range(30)]
    Y = [random.random() for _ in range(30)]

    for x, y in zip(X, Y):
        c1.update(x, y)

    for i, (x, y) in enumerate(zip(X, Y)):
        c2.update(x, y)
        if i < len(X) - 2:
            assert math.isclose(
                (c1 - c2).get(), np.cov(X[i + 1 :], Y[i + 1 :], ddof=ddof)[0, 1], abs_tol=1e-10
            )

    # Test inplace subtraction
    X.extend(random.random() for _ in range(3))
    Y.extend(random.random() for _ in range(3))
    for i in range(30, 33):
        c1.update(X[i], Y[i])

    c1 -= c2
    assert math.isclose(c1.get(), np.cov(X[30:33], Y[30:33], ddof=ddof)[0, 1], abs_tol=1e-10)


# river/stats/test_quantile.py

from __future__ import annotations

import random

from river import stats


def test_issue_1178():
    """

    https://github.com/online-ml/river/issues/1178

    >>> from river import stats

    >>> q = stats.Quantile(0.01)
    >>> for x in [5, 0, 0, 0, 0, 0, 0, 0]:
    ...     q = q.update(x)
    ...     print(q)
    Quantile: 5.
    Quantile: 0.
    Quantile: 0.
    Quantile: 0.
    Quantile: 0.
    Quantile: 0.
    Quantile: 0.
    Quantile: 0.

    >>> q = stats.Quantile(0.99)
    >>> for x in [5, 0, 0, 0, 0, 0, 0, 0]:
    ...     q = q.update(x)
    ...     print(q)
    Quantile: 5.
    Quantile: 5.
    Quantile: 5.
    Quantile: 5.
    Quantile: 5.
    Quantile: 0.
    Quantile: 0.277778
    Quantile: 0.827546

    """


def test_ge():
    low = stats.Quantile(0.01)
    high = stats.Quantile(0.99)

    for _ in range(100):
        x = random.random()
        low.update(x)
        high.update(x)
        assert high.get() >= low.get()


# river/stats/test_stats.py

from __future__ import annotations

import collections
import copy
import functools
import importlib
import inspect
import math
import pickle
import random
import statistics

import numpy as np
import pytest
from scipy import stats as sp_stats

from river import stats, utils


def load_stats():
    for _, obj in inspect.getmembers(importlib.import_module("river.stats"), inspect.isclass):
        try:
            if inspect.isabstract(obj):
                continue

            if issubclass(obj, stats.Link):
                yield obj(stats.Shift(1), stats.Mean())
                continue

            sig = inspect.signature(obj)
            yield obj(
                **{
                    param.name: param.default if param.default != param.empty else 1
                    for param in sig.parameters.values()
                }
            )
        except ValueError:
            yield obj()


@pytest.mark.parametrize("stat", load_stats(), ids=lambda stat: stat.__class__.__name__)
def test_pickling(stat):
    assert isinstance(pickle.loads(pickle.dumps(stat)), stat.__class__)
    assert isinstance(copy.deepcopy(stat), stat.__class__)

    # Check the statistic has a working __str__ and name method
    assert isinstance(str(stat), str)

    if isinstance(stat, stats.base.Univariate):
        assert isinstance(stat.name, str)


@pytest.mark.parametrize("stat", load_stats(), ids=lambda stat: stat.__class__.__name__)
def test_pickling_value(stat):
    for i in range(10):
        if isinstance(stat, stats.base.Bivariate):
            stat.update(i, i)
        elif isinstance(stat, stats.NUnique):  # takes string in input
            stat.update(str(i))
        else:
            stat.update(i)

    assert stat.get() == pickle.loads(pickle.dumps(stat)).get()
    assert stat.get() == copy.deepcopy(stat).get()


@pytest.mark.parametrize("stat", load_stats(), ids=lambda stat: stat.__class__.__name__)
def test_repr_with_no_updates(stat):
    assert isinstance(repr(stat), str)
    assert isinstance(str(stat), str)


@pytest.mark.parametrize(
    "stat, func",
    [
        (stats.Kurtosis(bias=True), sp_stats.kurtosis),
        (stats.Kurtosis(bias=False), functools.partial(sp_stats.kurtosis, bias=False)),
        (stats.Mean(), statistics.mean),
        (stats.Skew(bias=True), sp_stats.skew),
        (stats.Skew(bias=False), functools.partial(sp_stats.skew, bias=False)),
        (stats.Var(ddof=0), np.var),
        (stats.Var(), functools.partial(np.var, ddof=1)),
    ],
)
def test_univariate(stat, func):
    X = [random.random() for _ in range(30)]

    for i, x in enumerate(X):
        stat.update(x)
        if i >= 1:
            assert math.isclose(stat.get(), func(X[: i + 1]), abs_tol=1e-10)


# TODO
# def _weighted_variance(X, W):
#     mean = np.average(X, weights=W)
#     return np.average((W - mean) ** 2, weights=W)


@pytest.mark.parametrize(
    "stat, func",
    [
        (stats.Mean(), lambda x, w: np.average(x, weights=w)),
    ],
)
def test_univariate_frequency_weights(stat, func):
    """https://www.wikiwand.com/en/Weighted_arithmetic_mean"""

    X = [random.random() for _ in range(30)]
    W = [random.randint(1, 5) for _ in range(30)]

    for i, (x, w) in enumerate(zip(X, W)):
        stat.update(x, w)
        if i >= 1:
            assert math.isclose(stat.get(), func(X[: i + 1], W[: i + 1]), abs_tol=1e-10)


@pytest.mark.parametrize(
    "stat, func",
    [
        (stats.Mean(), lambda x, w: np.average(x, weights=w)),
    ],
)
def test_univariate_reliability_weights(stat, func):
    """https://www.wikiwand.com/en/Weighted_arithmetic_mean"""

    X = [random.random() for _ in range(30)]
    W = [random.random() for _ in range(30)]

    for i, (x, w) in enumerate(zip(X, W)):
        stat.update(x, w)
        if i >= 1:
            assert math.isclose(stat.get(), func(X[: i + 1], W[: i + 1]), abs_tol=1e-10)


@pytest.mark.parametrize(
    "stat, func",
    [
        # TODO: we shouldn't ignore these types
        (utils.Rolling(stats.Mean(), 3), statistics.mean),  # type: ignore
        (utils.Rolling(stats.Mean(), 10), statistics.mean),  # type: ignore
        (utils.Rolling(stats.Var(ddof=0), 3), np.var),  # type: ignore
        (utils.Rolling(stats.Var(ddof=0), 10), np.var),  # type: ignore
        (
            stats.RollingQuantile(0.0, 10),
            functools.partial(np.quantile, q=0.0, method="linear"),
        ),
        (
            stats.RollingQuantile(0.25, 10),
            functools.partial(np.quantile, q=0.25, method="linear"),
        ),
        (
            stats.RollingQuantile(0.5, 10),
            functools.partial(np.quantile, q=0.5, method="linear"),
        ),
        (
            stats.RollingQuantile(0.75, 10),
            functools.partial(np.quantile, q=0.75, method="linear"),
        ),
        (
            stats.RollingQuantile(1, 10),
            functools.partial(np.quantile, q=1, method="linear"),
        ),
    ],
)
def test_rolling_univariate(stat, func):
    def tail(iterable, n):
        return collections.deque(iterable, maxlen=n)

    n = stat.window_size
    X = [random.random() for _ in range(30)]

    for i, x in enumerate(X):
        stat.update(x)
        if i >= 1:
            assert math.isclose(stat.get(), func(tail(X[: i + 1], n)), abs_tol=1e-10)


@pytest.mark.parametrize(
    "stat, func",
    [
        # TODO: we shouldn't ignore these types
        (utils.Rolling(stats.Mean(), 3), lambda x, w: np.average(x, weights=w)),  # type: ignore
        (utils.Rolling(stats.Mean(), 10), lambda x, w: np.average(x, weights=w)),  # type: ignore
    ],
)
def test_rolling_univariate_sample_weights(stat, func):
    def tail(iterable, n):
        return collections.deque(iterable, maxlen=n)

    n = stat.window_size
    X = [random.random() for _ in range(30)]
    W = [random.randint(1, 5) for _ in range(30)]

    for i, (x, w) in enumerate(zip(X, W)):
        stat.update(x, w)
        if i >= 1:
            assert math.isclose(
                stat.get(), func(tail(X[: i + 1], n), tail(W[: i + 1], n)), abs_tol=1e-10
            )


@pytest.mark.parametrize(
    "stat, func",
    [
        # TODO: we shouldn't ignore these types
        (utils.Rolling(stats.Mean(), 3), lambda x, w: np.average(x, weights=w)),  # type: ignore
        (utils.Rolling(stats.Mean(), 10), lambda x, w: np.average(x, weights=w)),  # type: ignore
    ],
)
def test_rolling_univariate_reliability_weights(stat, func):
    def tail(iterable, n):
        return collections.deque(iterable, maxlen=n)

    n = stat.window_size
    X = [random.random() for _ in range(30)]
    W = [random.random() for _ in range(30)]

    for i, (x, w) in enumerate(zip(X, W)):
        stat.update(x, w)
        if i >= 1:
            assert math.isclose(
                stat.get(), func(tail(X[: i + 1], n), tail(W[: i + 1], n)), abs_tol=1e-10
            )


@pytest.mark.parametrize(
    "stat, func",
    [
        (stats.Cov(), lambda x, y: np.cov(x, y)[0, 1]),
        (stats.PearsonCorr(), lambda x, y: sp_stats.pearsonr(x, y)[0]),
    ],
)
def test_bivariate(stat, func):
    X = [random.random() for _ in range(30)]
    Y = [random.random() * x for x in X]

    for i, (x, y) in enumerate(zip(X, Y)):
        stat.update(x, y)
        if i >= 1:
            assert math.isclose(stat.get(), func(X[: i + 1], Y[: i + 1]), abs_tol=1e-10)


@pytest.mark.parametrize(
    "stat, func",
    [
        (utils.Rolling(stats.PearsonCorr(), 3), lambda x, y: sp_stats.pearsonr(x, y)[0]),  # type: ignore
        (utils.Rolling(stats.PearsonCorr(), 10), lambda x, y: sp_stats.pearsonr(x, y)[0]),  # type: ignore
        (utils.Rolling(stats.Cov(), 3), lambda x, y: np.cov(x, y)[0, 1]),  # type: ignore
        (utils.Rolling(stats.Cov(), 10), lambda x, y: np.cov(x, y)[0, 1]),  # type: ignore
    ],
)
def test_rolling_bivariate(stat, func):
    # Enough alread

    def tail(iterable, n):
        return collections.deque(iterable, maxlen=n)

    n = stat.window_size
    X = [random.random() for _ in range(30)]
    Y = [random.random() * x for x in X]

    for i, (x, y) in enumerate(zip(X, Y)):
        stat.update(x, y)
        if i >= 1:
            x_tail = tail(X[: i + 1], n)
            y_tail = tail(Y[: i + 1], n)
            assert math.isclose(stat.get(), func(x_tail, y_tail), abs_tol=1e-10)


@pytest.mark.parametrize(
    "stat",
    filter(
        lambda stat: hasattr(stat, "update_many")
        and issubclass(stat.__class__, stats.base.Univariate),
        load_stats(),
    ),
    ids=lambda stat: stat.__class__.__name__,
)
def test_update_many_univariate(stat):
    batch_stat = stat.clone()

    for _ in range(5):
        X = np.random.random(10)
        batch_stat.update_many(X)
        for x in X:
            stat.update(x)

    assert math.isclose(batch_stat.get(), stat.get())


@pytest.mark.parametrize(
    "stat",
    filter(
        lambda stat: hasattr(stat, "update_many")
        and issubclass(stat.__class__, stats.base.Bivariate),
        load_stats(),
    ),
    ids=lambda stat: stat.__class__.__name__,
)
def test_update_many_bivariate(stat):
    batch_stat = stat.clone()

    for _ in range(5):
        X = np.random.random(10)
        Y = np.random.random(10)
        batch_stat.update_many(X, Y)
        for x, y in zip(X, Y):
            stat.update(x, y)

    assert math.isclose(batch_stat.get(), stat.get())


# river/stats/test_var.py

from __future__ import annotations

import math
import random

from river import stats


def test_weighted_variance_with_close_numbers():
    """

    Origin of this test: https://github.com/online-ml/river/issues/732

    This test would fail if Var were implemented with a numerically unstable algorithm.

    """

    D = [
        (99.99999978143265, 6),
        (99.99999989071631, 8),
        (99.99999994535816, 6),
        (99.99999997267908, 9),
        (99.99999998633952, 10),
        (99.99999999316977, 3),
        (99.99999999829245, 5),
        (99.99999999957309, 9),
    ]

    var = stats.Var()

    for x, w in D:
        var.update(x, w)

    assert var.get() > 0 and math.isclose(var.get(), 4.648047194845607e-15)


def test_revert():
    for _ in range(5):
        X = [random.random() for _ in range(20)]

        v1 = stats.Var()
        v2 = stats.Var()

        for x in X[:10]:
            v1.update(x)
            v2.update(x)

        for x in X[10:]:
            v2.update(x)
        for x in X[10:]:
            v2.revert(x)

        assert math.isclose(v1.get(), v2.get())


# river/stats/var.py

from __future__ import annotations

import copy

import numpy as np

from river import stats


class Var(stats.base.Univariate):
    """Running variance using Welford's algorithm.

    Parameters
    ----------
    ddof
        Delta Degrees of Freedom. The divisor used in calculations is `n - ddof`, where `n`
        represents the number of seen elements.

    Attributes
    ----------
    mean
        It is necessary to calculate the mean of the data in order to calculate its variance.

    Notes
    -----
    The outcomes of the incremental and parallel updates are consistent with numpy's
    batch processing when $\\text{ddof} \\le 1$.

    Examples
    --------

    >>> from river import stats

    >>> X = [3, 5, 4, 7, 10, 12]

    >>> var = stats.Var()
    >>> for x in X:
    ...     print(var.update(x).get())
    0.0
    2.0
    1.0
    2.916666
    7.7
    12.56666

    You can measure a rolling variance by using a `utils.Rolling` wrapper:

    >>> from river import utils

    >>> X = [1, 4, 2, -4, -8, 0]
    >>> rvar = utils.Rolling(stats.Var(ddof=1), window_size=3)
    >>> for x in X:
    ...     print(rvar.update(x).get())
    0.0
    4.5
    2.333333
    17.333333
    25.333333
    16.0

    References
    ----------
    [^1]: [Wikipedia article on algorithms for calculating variance](https://www.wikiwand.com/en/Algorithms_for_calculating_variance#/Covariance)
    [^2]: [Chan, T.F., Golub, G.H. and LeVeque, R.J., 1983. Algorithms for computing the sample variance: Analysis and recommendations. The American Statistician, 37(3), pp.242-247.](https://amstat.tandfonline.com/doi/abs/10.1080/00031305.1983.10483115)
    [^3]: Schubert, E. and Gertz, M., 2018, July. Numerically stable parallel computation of
    (co-)variance. In Proceedings of the 30th International Conference on Scientific and
    Statistical Database Management (pp. 1-12).

    """

    def __init__(self, ddof=1):
        self.ddof = ddof
        self.mean = stats.Mean()
        self._S = 0

    @property
    def n(self):
        return self.mean.n

    def update(self, x, w=1.0):
        mean_old = self.mean.get()
        self.mean.update(x, w)
        mean_new = self.mean.get()
        self._S += w * (x - mean_old) * (x - mean_new)
        return self

    def revert(self, x, w=1.0):
        mean_old = self.mean.get()
        self.mean.revert(x, w)
        mean_new = self.mean.get()
        self._S -= w * (x - mean_old) * (x - mean_new)
        return self

    def update_many(self, X: np.ndarray):
        mean_old = self.mean.get()
        self.mean.update_many(X)
        mean_new = self.mean.get()
        self._S += np.sum(np.multiply(np.subtract(X, mean_old), np.subtract(X, mean_new)))
        return self

    def get(self):
        if self.n > self.ddof:
            return self._S / (self.n - self.ddof)
        return 0.0

    @classmethod
    def _from_state(cls, n, m, sig, *, ddof=1):
        new = cls(ddof=ddof)
        new.mean = stats.Mean._from_state(n, m)  # noqa
        # scale the second order statistic
        new._S = (n - ddof) * sig

        return new

    def __iadd__(self, other):
        S = (
            self._S
            + other._S
            + (self.mean.get() - other.mean.get()) ** 2 * self.n * other.n / (self.n + other.n)
        )
        self.mean += other.mean
        self._S = S
        return self

    def __add__(self, other):
        result = copy.deepcopy(self)
        result += other
        return result

    def __isub__(self, other):
        self.mean -= other.mean

        S = (
            self._S
            - other._S
            - (self.mean.get() - other.mean.get()) ** 2
            * self.n
            * other.mean.n
            / (self.n + other.mean.n)
        )
        self._S = S

        return self

    def __sub__(self, other):
        result = copy.deepcopy(self)
        result -= other
        return result


# river/stream/__init__.py

"""Streaming utilities.

The module includes tools to iterate over data streams.

"""
from __future__ import annotations

from .cache import Cache
from .iter_arff import iter_arff
from .iter_array import iter_array
from .iter_csv import iter_csv
from .iter_libsvm import iter_libsvm
from .qa import simulate_qa
from .shuffling import shuffle
from .tweet_stream import TwitterLiveStream
from .twitch_chat_stream import TwitchChatStream

__all__ = [
    "Cache",
    "iter_arff",
    "iter_array",
    "iter_csv",
    "iter_libsvm",
    "simulate_qa",
    "shuffle",
    "TwitterLiveStream",
    "TwitchChatStream",
]

try:
    from .iter_pandas import iter_pandas

    __all__ += ["iter_pandas"]
except ImportError:
    pass

try:
    from .iter_sklearn import iter_sklearn_dataset

    __all__ += ["iter_sklearn_dataset"]
except ImportError:
    pass

try:
    from .iter_sql import iter_sql

    __all__ += ["iter_sql"]
except ImportError:
    pass

try:
    from .iter_vaex import iter_vaex

    __all__ += ["iter_vaex"]
except ImportError:
    pass


# river/stream/cache.py

from __future__ import annotations

import glob
import inspect
import os
import pickle
import platform

from river import utils


class Cache:
    """Utility for caching iterables.

    This can be used to save a stream of data to the disk in order to iterate over it faster the
    following time. This can save time depending on the nature of stream. The more processing
    happens in a stream, the more time will be saved. Even in the case where no processing is done
    apart from reading the data, the cache will save some time because it is using the pickle
    binary protocol. It can thus improve the speed in common cases such as reading from a CSV file.

    Parameters
    ----------
    directory
        The path where to store the pickled data streams. If not provided, then it will be
        automatically inferred whenever possible, if not an exception will be raised.

    Attributes
    ----------
    keys : set
        The set of keys that are being cached.

    Examples
    --------

    >>> import time
    >>> from river import datasets
    >>> from river import stream

    >>> dataset = datasets.Phishing()
    >>> cache = stream.Cache()

    The cache can be used by wrapping it around an iterable. Because this is the first time
    are iterating over the data, nothing is cached.

    >>> tic = time.time()
    >>> for x, y in cache(dataset, key='phishing'):
    ...     pass
    >>> toc = time.time()
    >>> print(toc - tic)  # doctest: +SKIP
    0.012813

    If we do the same thing again, we can see the loop is now faster.

    >>> tic = time.time()
    >>> for x, y in cache(dataset, key='phishing'):
    ...     pass
    >>> toc = time.time()
    >>> print(toc - tic)  # doctest: +SKIP
    0.001927

    We can see an overview of the cache. The first line indicates the location of the
    cache.

    >>> cache  # doctest: +SKIP
    /tmp
    phishing - 125.2KiB

    Finally, we can clear the stream from the cache.

    >>> cache.clear('phishing')
    >>> cache  # doctest: +SKIP
    /tmp

    There is also a `clear_all` method to remove all the items in the cache.

    >>> cache.clear_all()

    """

    def __init__(self, directory=None):
        # Guess the directory from the system
        system = platform.system()
        if directory is None:
            directory = {"Linux": "/tmp", "Darwin": "/tmp", "Windows": "C:\\TEMP"}.get(system)

        if directory is None:
            raise ValueError(
                f"There is no default directory defined for {system} systems, "
                "please provide one manually"
            )

        self.directory = directory
        self.keys = set()

        # Check if there is anything already in the cache
        for f in glob.glob(os.path.join(self.directory, "*.river_cache.pkl")):
            key = os.path.basename(f).split(".")[0]
            self.keys.add(key)

    def _get_path(self, key):
        return os.path.join(self.directory, f"{key}.river_cache.pkl")

    def __call__(self, stream, key=None):
        # Try to guess a key from the stream object
        if key is None:
            if inspect.isfunction(stream):
                key = stream.__name__

        if key is None:
            raise ValueError(
                "No default key could be guessed for the given stream, " "please provide one"
            )

        path = self._get_path(key)

        if os.path.exists(path):
            yield from self[key]
            return

        with open(path, "wb") as f:
            pickler = pickle.Pickler(f)
            for el in stream:
                pickler.dump(el)
                yield el
            self.keys.add(key)

    def __getitem__(self, key):
        """Iterates over the stream associated with the given key."""
        with open(self._get_path(key), "rb") as f:
            unpickler = pickle.Unpickler(f)
            while f.peek(1):
                yield unpickler.load()

    def clear(self, key: str):
        """Delete the cached stream associated with the given key.

        Parameters
        ----------
        key

        """
        os.remove(self._get_path(key))
        self.keys.remove(key)

    def clear_all(self):
        """Delete all the cached streams."""
        for key in list(self.keys):
            os.remove(self._get_path(key))
            self.keys.remove(key)

    def __repr__(self):
        return "\n".join(
            [self.directory]
            + [
                f"{key} - {utils.pretty.humanize_bytes(os.path.getsize(self._get_path(key)))}"
                for key in self.keys
            ]
        )


# river/stream/iter_arff.py

from __future__ import annotations

from scipy.io.arff import arffread

from river import base

from . import utils


def iter_arff(
    filepath_or_buffer, target: str | None = None, compression="infer"
) -> base.typing.Stream:
    """Iterates over rows from an ARFF file.

    Parameters
    ----------
    filepath_or_buffer
        Either a string indicating the location of a file, or a buffer object that has a
        `read` method.
    target
        Name of the target field.
    compression
        For on-the-fly decompression of on-disk data. If this is set to 'infer' and
        `filepath_or_buffer` is a path, then the decompression method is inferred for the
        following extensions: '.gz', '.zip'.

    """

    # If a file is not opened, then we open it
    buffer = filepath_or_buffer
    if not hasattr(buffer, "read"):
        buffer = utils.open_filepath(buffer, compression)

    try:
        rel, attrs = arffread.read_header(buffer)
    except ValueError as e:
        msg = f"Error while parsing header, error was: {e}"
        raise arffread.ParseArffError(msg)

    names = [attr.name for attr in attrs]
    types = [float if isinstance(attr, arffread.NumericAttribute) else None for attr in attrs]

    for r in buffer:
        if len(r) == 0:
            continue
        x = {
            name: typ(val) if typ else val
            for name, typ, val in zip(names, types, r.rstrip().split(","))
        }
        try:
            y = x.pop(target) if target else None
        except KeyError as e:
            print(r)
            raise e

        yield x, y

    # Close the file if we opened it
    if buffer is not filepath_or_buffer:
        buffer.close()


# river/stream/iter_array.py

from __future__ import annotations

import itertools
import random

import numpy as np

from river import base


def iter_array(
    X: np.ndarray,
    y: np.ndarray | None = None,
    feature_names: list[base.typing.FeatureName] | None = None,
    target_names: list[base.typing.FeatureName] | None = None,
    shuffle: bool = False,
    seed: int | None = None,
) -> base.typing.Stream:
    """Iterates over the rows from an array of features and an array of targets.

    This method is intended to work with `numpy` arrays, but should also work with Python lists.

    Parameters
    ----------
    X
        A 2D array of features. This can also be a 1D array of strings, which can be the case if
        you're working with text.
    y
        An optional array of targets.
    feature_names
        An optional list of feature names. The features will be labeled with integers if no names
        are provided.
    target_names
        An optional list of output names. The outputs will be labeled with integers if no names are
        provided. Only applies if there are multiple outputs, i.e. if `y` is a 2D array.
    shuffle
        Indicates whether or not to shuffle the input arrays before iterating over them.
    seed
        Random seed used for shuffling the data.

    Examples
    --------

    >>> from river import stream
    >>> import numpy as np

    >>> X = np.array([[1, 2, 3], [11, 12, 13]])
    >>> Y = np.array([True, False])

    >>> dataset = stream.iter_array(
    ...     X, Y,
    ...     feature_names=['x1', 'x2', 'x3']
    ... )
    >>> for x, y in dataset:
    ...     print(x, y)
    {'x1': 1, 'x2': 2, 'x3': 3} True
    {'x1': 11, 'x2': 12, 'x3': 13} False

    This also works with a array of texts:

    >>> X = ["foo", "bar"]
    >>> dataset = stream.iter_array(
    ...     X, Y,
    ...     feature_names=['x1', 'x2', 'x3']
    ... )
    >>> for x, y in dataset:
    ...     print(x, y)
    foo True
    bar False

    """

    # If the first row of X is actually a string, then we assume all the rows are strings and will
    # pass them through
    if isinstance(X[0], str):

        def handle_features(x):
            return x

    # If not we assume each row if a set of features, and will convert them to a dictionary
    else:
        feature_names = list(range(len(X[0]))) if feature_names is None else feature_names

        def handle_features(x):
            return dict(zip(feature_names, xi))

    multioutput = y is not None and not np.isscalar(y[0])
    if multioutput and target_names is None:
        target_names = list(range(len(y[0])))  # type: ignore

    # Shuffle the data
    rng = random.Random(seed)
    if shuffle:
        order = rng.sample(range(len(X)), k=len(X))
        X = X[order]
        y = y if y is None else y[order]

    if multioutput:
        for xi, yi in itertools.zip_longest(X, y if hasattr(y, "__iter__") else []):  # type: ignore
            yield handle_features(xi), dict(zip(target_names, yi))  # type: ignore

    else:
        for xi, yi in itertools.zip_longest(X, y if hasattr(y, "__iter__") else []):  # type: ignore
            yield handle_features(xi), yi


# river/stream/iter_csv.py

from __future__ import annotations

import csv
import datetime as dt
import random

from .. import base
from . import utils

__all__ = ["iter_csv"]


class DictReader(csv.DictReader):
    """Overlay on top of `csv.DictReader` which allows sampling."""

    def __init__(self, fraction, rng, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fraction = fraction
        self.rng = rng

    def __next__(self):
        if self.line_num == 0:
            self.fieldnames

        row = next(self.reader)

        if self.fraction < 1:
            while self.rng.random() > self.fraction:
                row = next(self.reader)

        return dict(zip(self.fieldnames, row))


def iter_csv(
    filepath_or_buffer,
    target: str | list[str] | None = None,
    converters: dict | None = None,
    parse_dates: dict | None = None,
    drop: list[str] | None = None,
    drop_nones=False,
    fraction=1.0,
    compression="infer",
    seed: int | None = None,
    field_size_limit: int | None = None,
    **kwargs,
) -> base.typing.Stream:
    """Iterates over rows from a CSV file.

    Reading CSV files can be quite slow. If, for whatever reason, you're going to loop through
    the same file multiple times, then we recommend that you to use the `stream.Cache` utility.

    Parameters
    ----------
    filepath_or_buffer
        Either a string indicating the location of a file, or a buffer object that has a
        `read` method.
    target
        A single target column is assumed if a string is passed. A multiple output scenario
        is assumed if a list of strings is passed. A `None` value will be assigned to each `y`
        if this parameter is omitted.
    converters
        All values in the CSV are interpreted as strings by default. You can use this parameter to
        cast values to the desired type. This should be a `dict` mapping feature names to callables
        used to parse their associated values. Note that a callable may be a type, such as `float`
        and `int`.
    parse_dates
        A `dict` mapping feature names to a format passed to the `datetime.datetime.strptime`
        method.
    drop
        Fields to ignore.
    drop_nones
        Whether or not to drop fields where the value is a `None`.
    fraction
        Sampling fraction.
    compression
        For on-the-fly decompression of on-disk data. If this is set to 'infer' and
        `filepath_or_buffer` is a path, then the decompression method is inferred for the
        following extensions: '.gz', '.zip'.
    seed
        If specified, the sampling will be deterministic.
    field_size_limit
        If not `None`, this will be passed to the `csv.field_size_limit` function.
    kwargs
        All other keyword arguments are passed to the underlying `csv.DictReader`.

    Returns
    -------

    By default each feature value will be of type `str`. You can use the `converters` and
    `parse_dates` parameters to convert them as you see fit.

    Examples
    --------

    Although this function is designed to handle different kinds of inputs, the most common
    use case is to read a file on the disk. We'll first create a little CSV file to illustrate.

    >>> tv_shows = '''name,year,rating
    ... Planet Earth II,2016,9.5
    ... Planet Earth,2006,9.4
    ... Band of Brothers,2001,9.4
    ... Breaking Bad,2008,9.4
    ... Chernobyl,2019,9.4
    ... '''
    >>> with open('tv_shows.csv', mode='w') as f:
    ...     _ = f.write(tv_shows)

    We can now go through the rows one by one. We can use the `converters` parameter to cast
    the `rating` field value as a `float`. We can also convert the `year` to a `datetime` via
    the `parse_dates` parameter.

    >>> from river import stream

    >>> params = {
    ...     'converters': {'rating': float},
    ...     'parse_dates': {'year': '%Y'}
    ... }
    >>> for x, y in stream.iter_csv('tv_shows.csv', **params):
    ...     print(x, y)
    {'name': 'Planet Earth II', 'year': datetime.datetime(2016, 1, 1, 0, 0), 'rating': 9.5} None
    {'name': 'Planet Earth', 'year': datetime.datetime(2006, 1, 1, 0, 0), 'rating': 9.4} None
    {'name': 'Band of Brothers', 'year': datetime.datetime(2001, 1, 1, 0, 0), 'rating': 9.4} None
    {'name': 'Breaking Bad', 'year': datetime.datetime(2008, 1, 1, 0, 0), 'rating': 9.4} None
    {'name': 'Chernobyl', 'year': datetime.datetime(2019, 1, 1, 0, 0), 'rating': 9.4} None

    The value of `y` is always `None` because we haven't provided a value for the `target`
    parameter. Here is an example where a `target` is provided:

    >>> dataset = stream.iter_csv('tv_shows.csv', target='rating', **params)
    >>> for x, y in dataset:
    ...     print(x, y)
    {'name': 'Planet Earth II', 'year': datetime.datetime(2016, 1, 1, 0, 0)} 9.5
    {'name': 'Planet Earth', 'year': datetime.datetime(2006, 1, 1, 0, 0)} 9.4
    {'name': 'Band of Brothers', 'year': datetime.datetime(2001, 1, 1, 0, 0)} 9.4
    {'name': 'Breaking Bad', 'year': datetime.datetime(2008, 1, 1, 0, 0)} 9.4
    {'name': 'Chernobyl', 'year': datetime.datetime(2019, 1, 1, 0, 0)} 9.4

    Finally, let's delete the example file.

    >>> import os; os.remove('tv_shows.csv')

    """

    # Set the field size limit
    limit = csv.field_size_limit()
    if field_size_limit is not None:
        csv.field_size_limit(field_size_limit)

    # If a file is not opened, then we open it
    buffer = filepath_or_buffer
    if not hasattr(buffer, "read"):
        buffer = utils.open_filepath(buffer, compression)

    for x in DictReader(fraction=fraction, rng=random.Random(seed), f=buffer, **kwargs):
        if drop:
            for i in drop:
                del x[i]

        # Cast the values to the given types
        if converters is not None:
            for i, t in converters.items():
                x[i] = t(x[i])

        # Drop Nones
        if drop_nones:
            for i in list(x):
                if x[i] is None:
                    del x[i]

        # Parse the dates
        if parse_dates is not None:
            for i, fmt in parse_dates.items():
                x[i] = dt.datetime.strptime(x[i], fmt)

        # Separate the target from the features
        y = None
        if isinstance(target, list):
            y = {name: x.pop(name) for name in target}
        elif target is not None:
            y = x.pop(target)

        yield x, y

    # Close the file if we opened it
    if buffer is not filepath_or_buffer:
        buffer.close()

    # Reset the file size limit to it's original value
    csv.field_size_limit(limit)


# river/stream/iter_libsvm.py

from __future__ import annotations

from river import base

from . import utils


def iter_libsvm(
    filepath_or_buffer: str, target_type=float, compression="infer"
) -> base.typing.Stream:
    """Iterates over a dataset in LIBSVM format.

    The LIBSVM format is a popular way in the machine learning community to store sparse datasets.
    Only numerical feature values are supported. The feature names will be considered as strings.

    Parameters
    ----------
    filepath_or_buffer
        Either a string indicating the location of a file, or a buffer object that has a `read`
        method.
    target_type
        The type of the target value.
    compression
        For on-the-fly decompression of on-disk data. If this is set to 'infer' and
        `filepath_or_buffer` is a path, then the decompression method is inferred for the
        following extensions: '.gz', '.zip'.

    Examples
    --------

    >>> import io
    >>> from river import stream

    >>> data = io.StringIO('''+1 x:-134.26 y:0.2563
    ... 1 x:-12 z:0.3
    ... -1 y:.25
    ... ''')

    >>> for x, y in stream.iter_libsvm(data, target_type=int):
    ...     print(y, x)
    1 {'x': -134.26, 'y': 0.2563}
    1 {'x': -12.0, 'z': 0.3}
    -1 {'y': 0.25}

    References
    ----------
    [^1]: [LIBSVM documentation](https://www.csie.ntu.edu.tw/~cjlin/libsvm/)

    """

    # If a file is not opened, then we open it
    buffer = filepath_or_buffer
    should_close = False
    if not hasattr(buffer, "read"):
        should_close = False
        buffer = utils.open_filepath(buffer, compression)

    def split_pair(pair):
        name, value = pair.split(":")
        value = float(value)
        return name, value

    for line in buffer:
        # Remove carriage return and whitespace
        line = line.rstrip()
        # Remove potential end of line comments
        line = line.split("#")[0]

        y, x_str = line.split(" ", maxsplit=1)
        y = target_type(y)
        x = dict([split_pair(pair) for pair in x_str.split(" ")])
        yield x, y

    # Close the file if we opened it
    if should_close:
        buffer.close()  # type: ignore


# river/stream/iter_pandas.py

from __future__ import annotations

import pandas as pd

from river import base, stream


def iter_pandas(
    X: pd.DataFrame, y: pd.Series | pd.DataFrame | None = None, **kwargs
) -> base.typing.Stream:
    """Iterates over the rows of a `pandas.DataFrame`.

    Parameters
    ----------
    X
        A dataframe of features.
    y
        A series or a dataframe with one column per target.
    kwargs
        Extra keyword arguments are passed to the underlying call to `stream.iter_array`.

    Examples
    --------

    >>> import pandas as pd
    >>> from river import stream

    >>> X = pd.DataFrame({
    ...     'x1': [1, 2, 3, 4],
    ...     'x2': ['blue', 'yellow', 'yellow', 'blue'],
    ...     'y': [True, False, False, True]
    ... })
    >>> y = X.pop('y')

    >>> for xi, yi in stream.iter_pandas(X, y):
    ...     print(xi, yi)
    {'x1': 1, 'x2': 'blue'} True
    {'x1': 2, 'x2': 'yellow'} False
    {'x1': 3, 'x2': 'yellow'} False
    {'x1': 4, 'x2': 'blue'} True

    """

    kwargs["feature_names"] = X.columns
    if isinstance(y, pd.DataFrame):
        kwargs["target_names"] = y.columns

    yield from stream.iter_array(X=X.to_numpy(), y=y if y is None else y.to_numpy(), **kwargs)


# river/stream/iter_sklearn.py

from __future__ import annotations

import pandas as pd
import sklearn.utils

from river import base, stream


def iter_sklearn_dataset(dataset: sklearn.utils.Bunch, **kwargs) -> base.typing.Stream:
    """Iterates rows from one of the datasets provided by scikit-learn.

    This allows you to use any dataset from [scikit-learn's `datasets` module](https://scikit-learn.org/stable/modules/classes.html#module-sklearn.datasets). For instance, you can use the `fetch_openml` function to get access to all of the
    datasets from the OpenML website.

    Parameters
    ----------
    dataset
        A scikit-learn dataset.
    kwargs
        Extra keyword arguments are passed to the underlying call to `stream.iter_array`.

    Examples
    --------

    >>> import pprint
    >>> from sklearn import datasets
    >>> from river import stream

    >>> dataset = datasets.load_diabetes()

    >>> for xi, yi in stream.iter_sklearn_dataset(dataset):
    ...     pprint.pprint(xi)
    ...     print(yi)
    ...     break
    {'age': 0.038075906433423026,
     'bmi': 0.061696206518683294,
     'bp': 0.0218723855140367,
     's1': -0.04422349842444599,
     's2': -0.03482076283769895,
     's3': -0.04340084565202491,
     's4': -0.002592261998183278,
     's5': 0.019907486170462722,
     's6': -0.01764612515980379,
     'sex': 0.05068011873981862}
    151.0

    """
    kwargs["X"] = dataset.data
    kwargs["y"] = dataset.target
    try:
        kwargs["feature_names"] = dataset.feature_names
    except AttributeError:
        pass

    if isinstance(kwargs["X"], pd.DataFrame):
        yield from stream.iter_pandas(**kwargs)
    else:
        yield from stream.iter_array(**kwargs)


# river/stream/iter_sql.py

from __future__ import annotations

import sqlalchemy

from river import base

__all__ = ["iter_sql"]


def iter_sql(
    query: str | sqlalchemy.TextClause | sqlalchemy.Select,
    conn: sqlalchemy.Connection,
    target_name: str | None = None,
) -> base.typing.Stream:
    """Iterates over the results from an SQL query.

    By default, SQLAlchemy prefetches results. Therefore, even though you can iterate over the
    resulting rows one by one, the results are in fact loaded in batch. You can modify this
    behavior by configuring the connection you pass to `iter_sql`. For instance, you can set
    the `stream_results` parameter to `True`, as [explained in SQLAlchemy's documentation](https://docs.sqlalchemy.org/en/13/core/connections.html#sqlalchemy.engine.Connection.execution_options). Note, however,
    that this isn't available for all database engines.

    Parameters
    ----------
    query
        SQL query to be executed.
    conn
        An SQLAlchemy construct which has an `execute` method. In other words you can pass an
        engine, a connection, or a session.
    target_name
        The name of the target field. If this is `None`, then `y` will also be `None`.

    Examples
    --------

    As an example we'll create an in-memory database with SQLAlchemy.

    >>> import datetime as dt
    >>> import sqlalchemy

    >>> engine = sqlalchemy.create_engine('sqlite://')

    >>> metadata = sqlalchemy.MetaData()

    >>> t_sales = sqlalchemy.Table('sales', metadata,
    ...     sqlalchemy.Column('shop', sqlalchemy.String, primary_key=True),
    ...     sqlalchemy.Column('date', sqlalchemy.Date, primary_key=True),
    ...     sqlalchemy.Column('amount', sqlalchemy.Integer)
    ... )

    >>> metadata.create_all(engine)

    >>> sales = [
    ...     {'shop': 'Hema', 'date': dt.date(2016, 8, 2), 'amount': 20},
    ...     {'shop': 'Ikea', 'date': dt.date(2016, 8, 2), 'amount': 18},
    ...     {'shop': 'Hema', 'date': dt.date(2016, 8, 3), 'amount': 22},
    ...     {'shop': 'Ikea', 'date': dt.date(2016, 8, 3), 'amount': 14},
    ...     {'shop': 'Hema', 'date': dt.date(2016, 8, 4), 'amount': 12},
    ...     {'shop': 'Ikea', 'date': dt.date(2016, 8, 4), 'amount': 16}
    ... ]

    >>> with engine.connect() as conn:
    ...     _ = conn.execute(t_sales.insert(), sales)
    ...     conn.commit()

    We can now query the database. We will set `amount` to be the target field.

    >>> from river import stream

    >>> with engine.connect() as conn:
    ...     query = sqlalchemy.sql.select(t_sales)
    ...     dataset = stream.iter_sql(query, conn, target_name='amount')
    ...     for x, y in dataset:
    ...         print(x, y)
    {'shop': 'Hema', 'date': datetime.date(2016, 8, 2)} 20
    {'shop': 'Ikea', 'date': datetime.date(2016, 8, 2)} 18
    {'shop': 'Hema', 'date': datetime.date(2016, 8, 3)} 22
    {'shop': 'Ikea', 'date': datetime.date(2016, 8, 3)} 14
    {'shop': 'Hema', 'date': datetime.date(2016, 8, 4)} 12
    {'shop': 'Ikea', 'date': datetime.date(2016, 8, 4)} 16

    This also with raw SQL queries.

    >>> with engine.connect() as conn:
    ...     query = "SELECT * FROM sales WHERE shop = 'Hema'"
    ...     dataset = stream.iter_sql(query, conn, target_name='amount')
    ...     for x, y in dataset:
    ...         print(x, y)
    {'shop': 'Hema', 'date': '2016-08-02'} 20
    {'shop': 'Hema', 'date': '2016-08-03'} 22
    {'shop': 'Hema', 'date': '2016-08-04'} 12

    """

    result_proxy = conn.execute(sqlalchemy.sql.text(query) if isinstance(query, str) else query)

    if target_name is None:
        for row in result_proxy:
            yield dict(row._mapping.items()), None
        return

    for row in result_proxy:
        x = dict(row._mapping.items())
        y = x.pop(target_name)
        yield x, y


# river/stream/iter_vaex.py

from __future__ import annotations

import vaex
from vaex.utils import _ensure_list, _ensure_strings_from_expressions

from river import base


def iter_vaex(
    X: vaex.dataframe.DataFrame,
    y: str | vaex.expression.Expression | None = None,
    features: list[str] | vaex.expression.Expression | None = None,
) -> base.typing.Stream:
    """Yields rows from a ``vaex.DataFrame``.

    Parameters
    ----------
    X
        A vaex DataFrame housing the training featuers.
    y
        The column or expression containing the target variable.
    features
        A list of features used for training. If None, all columns in `X` will be used. Features
        specifying in `y` are ignored.

    """

    features = _ensure_strings_from_expressions(features)
    feature_names = features or X.get_column_names()

    if y:
        y = _ensure_strings_from_expressions(y)
        y = _ensure_list(y)
        feature_names = [feat for feat in feature_names if feat not in y]

    multioutput = len(y) > 1

    if multioutput:
        for i in range(len(X)):
            yield (
                {key: X.evaluate(key, i, i + 1)[0] for key in feature_names},
                {key: X.evaluate(key, i, i + 1)[0] for key in y},
            )

    else:
        for i in range(len(X)):
            yield (
                {key: X.evaluate(key, i, i + 1)[0] for key in feature_names},
                X.evaluate(y[0], i, i + 1)[0],
            )


# river/stream/qa.py

from __future__ import annotations

import bisect
import collections
import datetime as dt
import typing
from copy import deepcopy

from river import base

__all__ = ["simulate_qa"]


class Memento(collections.namedtuple("Memento", "i x y kwargs t_expire")):
    def __lt__(self, other):
        return self.t_expire < other.t_expire


def simulate_qa(
    dataset: base.typing.Dataset,
    moment: str | typing.Callable[[dict], dt.datetime] | None,
    delay: str | int | dt.timedelta | typing.Callable | None,
    copy: bool = True,
):
    """Simulate a time-ordered question and answer session.

    This method allows looping through a dataset in the order in which it arrived. Indeed, it
    usually is the case that labels arrive after features. Being able to go through a dataset in
    arrival order enables assessing a model's performance in a reliable manner. For instance, the
    `evaluate.progressive_val_score` is a high-level method that can be used to score a model on a
    dataset. Under the hood it uses this method to determine the correct arrival order.

    Parameters
    ----------
    dataset
        A stream of (features, target) tuples.
    moment
        The attribute used for measuring time. If a callable is passed, then it is expected
        to take as input a `dict` of features. If `None`, then the observations are implicitly
        timestamped in the order in which they arrive. If a `str` is passed, then it will be
        used to obtain the time from the input features.
    delay
        The amount of time to wait before revealing the target associated with each
        observation to the model. This value is expected to be able to sum with the `moment`
        value. For instance, if `moment` is a `datetime.date`, then `delay` is expected to be a
        `datetime.timedelta`. If a callable is passed, then it is expected to take as input a
        `dict` of features and the target. If a `str` is passed, then it will be used to access
        the relevant field from the features. If `None` is passed, then no delay will be used,
        which leads to doing standard online validation. If a scalar is passed, such an `int`
        or a `datetime.timedelta`, then the delay is constant.
    copy
        If `True`, then a separate copy of the features are yielded the second time
        around. This ensures that inadvertent modifications in downstream code don't have any
        effect.

    Examples
    --------

    The arrival delay isn't usually indicated in a dataset, but it might be able to be inferred
    from the features. As an example, we'll simulate the departure and arrival time of taxi
    trips. Let's first create a time table which records the departure time and the duration of
    seconds of several taxi trips.

    >>> import datetime as dt
    >>> time_table = [
    ...     (dt.datetime(2020, 1, 1, 20,  0, 0),  900),
    ...     (dt.datetime(2020, 1, 1, 20, 10, 0), 1800),
    ...     (dt.datetime(2020, 1, 1, 20, 20, 0),  300),
    ...     (dt.datetime(2020, 1, 1, 20, 45, 0),  400),
    ...     (dt.datetime(2020, 1, 1, 20, 50, 0),  240),
    ...     (dt.datetime(2020, 1, 1, 20, 55, 0),  450)
    ... ]

    We can now create a streaming dataset where the features are the departure dates and the
    targets are the durations.

    >>> dataset = (
    ...     ({'date': date}, duration)
    ...     for date, duration in time_table
    ... )

    Now, we can use `simulate_qa` to iterate over the events in the order in which they are
    meant to occur.

    >>> delay = lambda _, y: dt.timedelta(seconds=y)

    >>> for i, x, y in simulate_qa(dataset, moment='date', delay=delay):
    ...     if y is None:
    ...         print(f'{x["date"]} - trip #{i} departs')
    ...     else:
    ...         arrival_date = x['date'] + dt.timedelta(seconds=y)
    ...         print(f'{arrival_date} - trip #{i} arrives after {y} seconds')
    2020-01-01 20:00:00 - trip #0 departs
    2020-01-01 20:10:00 - trip #1 departs
    2020-01-01 20:15:00 - trip #0 arrives after 900 seconds
    2020-01-01 20:20:00 - trip #2 departs
    2020-01-01 20:25:00 - trip #2 arrives after 300 seconds
    2020-01-01 20:40:00 - trip #1 arrives after 1800 seconds
    2020-01-01 20:45:00 - trip #3 departs
    2020-01-01 20:50:00 - trip #4 departs
    2020-01-01 20:51:40 - trip #3 arrives after 400 seconds
    2020-01-01 20:54:00 - trip #4 arrives after 240 seconds
    2020-01-01 20:55:00 - trip #5 departs
    2020-01-01 21:02:30 - trip #5 arrives after 450 seconds

    This function is extremely practical because it provides a reliable way to evaluate the
    performance of a model in a real scenario. Indeed, it allows to make predictions and
    perform model updates in exactly the same manner that would happen live. For instance, it
    is used in `evaluate.progressive_val_score`, which is a higher level function for
    evaluating models in an online manner.

    """

    # Determine how to insert mementos into the queue
    queue = (
        (lambda q, el: bisect.insort(q, el))
        if callable(delay) or isinstance(delay, str)
        else (lambda q, el: q.append(el))
    )

    # Coerce moment to a function
    get_moment = (
        (lambda _, x: x[moment])
        if isinstance(moment, str)
        else (lambda _, x: moment(x))  # type: ignore
        if callable(moment)
        else (lambda i, _: i)  # type: ignore
    )

    # Coerce delay to a function
    get_delay = (
        (lambda i, _: 0)
        if delay is None
        else (lambda x, _: x[delay])
        if isinstance(delay, str)
        else (lambda _, __: delay)  # type: ignore
        if not callable(delay)
        else delay  # type: ignore
    )

    mementos: list[Memento] = []

    kwargs: list

    for i, (x, y, *kwargs) in enumerate(dataset):
        kwargs = kwargs[0] if kwargs else None

        t = get_moment(i, x)
        d = get_delay(x, y)  # type: ignore

        while mementos:
            # Get the oldest answer
            i_old, x_old, y_old, kwargs_old, t_expire = mementos[0]

            # If the oldest answer isn't old enough then stop
            if t_expire > t:
                break

            # Reveal the ground truth and pop the observation from the queue
            yield (i_old, x_old, y_old, kwargs_old) if kwargs_old else (
                i_old,
                x_old,
                y_old,
            )
            del mementos[0]

        queue(mementos, Memento(i, x, y, kwargs, t + d))
        if copy:
            x = deepcopy(x)
        yield (i, x, None, kwargs) if kwargs else (i, x, None)

    for memento in mementos:
        yield (memento.i, memento.x, memento.y, memento.kwargs) if memento.kwargs else (
            memento.i,
            memento.x,
            memento.y,
        )


# river/stream/shuffling.py

from __future__ import annotations

import itertools
import random
import types
import typing


def shuffle(stream: typing.Iterator, buffer_size: int, seed: int | None = None):
    """Shuffles a stream of data.

    This works by maintaining a buffer of elements. The first `buffer_size` elements are stored in
    memory. Once the buffer is full, a random element inside the buffer is yielded. Every time an
    element is yielded, the next element in the stream replaces it and the buffer is sampled again.
    Increasing `buffer_size` will improve the quality of the shuffling.

    If you really want to stream over your dataset in a "good" random order, the best way is to
    split your dataset into smaller datasets and loop over them in a round-robin fashion. You may
    do this by using the `roundrobin` recipe from the `itertools` module.

    Parameters
    ----------
    stream
        The stream to shuffle.
    buffer_size
        The size of the buffer which contains the elements help in memory. Increasing this will
        increase randomness but will incur more memory usage.
    seed
        Random seed used for sampling.

    Examples
    --------

    >>> from river import stream

    >>> for i in stream.shuffle(range(15), buffer_size=5, seed=42):
    ...     print(i)
    0
    5
    2
    1
    8
    9
    6
    4
    11
    12
    10
    7
    14
    13
    3

    References
    ----------
    [^1]: [Visualizing TensorFlow's streaming shufflers](http://www.moderndescartes.com/essays/shuffle_viz/)

    """

    rng = random.Random(seed)

    # If stream is not a generator, then we coerce it to one
    if not isinstance(stream, types.GeneratorType):
        stream = iter(stream)

    # Initialize the buffer with the first buffer_size elements of the stream
    buffer = list(itertools.islice(stream, buffer_size))

    # Deplete the stream until it is empty
    for element in stream:
        # Pick a random element from the buffer and yield it
        i = rng.randint(0, len(buffer) - 1)
        yield buffer[i]

        # Replace the yielded element from the buffer with the new element from the stream
        buffer[i] = element

    # Shuffle the remaining buffer elements and yield them one by one
    rng.shuffle(buffer)
    yield from buffer


# river/stream/test_iter_csv.py

from __future__ import annotations

import io

from river import stream


def test_iter_csv_custom_converter():
    example = io.StringIO("col1,col2,col3\n,1,2\n5,,4\n3,1,")

    def int_or_none(s):
        try:
            return int(s)
        except ValueError:
            return None

    params = {"converters": {"col1": int_or_none, "col2": int_or_none, "col3": int_or_none}}
    dataset = stream.iter_csv(example, **params)
    assert list(dataset) == [
        ({"col1": None, "col2": 1, "col3": 2}, None),
        ({"col1": 5, "col2": None, "col3": 4}, None),
        ({"col1": 3, "col2": 1, "col3": None}, None),
    ]


def test_iter_csv_drop_nones():
    example = io.StringIO("col1,col2,col3\n,1,2\n5,,4\n3,1,")

    def int_or_none(s):
        try:
            return int(s)
        except ValueError:
            return None

    params = {
        "converters": {"col1": int_or_none, "col2": int_or_none, "col3": int_or_none},
        "drop_nones": True,
    }
    dataset = stream.iter_csv(example, **params)
    assert list(dataset) == [
        ({"col2": 1, "col3": 2}, None),
        ({"col1": 5, "col3": 4}, None),
        ({"col1": 3, "col2": 1}, None),
    ]


# river/stream/test_sql.py

from __future__ import annotations

import json
import os
import zipfile

import pytest
import sqlalchemy as sql

from river import stream


@pytest.fixture
def pokedb():
    engine = sql.create_engine("sqlite://")  # in-memory

    # Load the fixtures
    here = os.path.dirname(os.path.realpath(__file__))
    with zipfile.ZipFile(os.path.join(here, "pokedb.zip")) as z:
        pokemons = json.loads(z.read("pokemons.json"))
        types = json.loads(z.read("types.json"))
        pokemon_types = json.loads(z.read("pokemon_types.json"))

    # Define the tables
    metadata = sql.MetaData()

    t_pokemons = sql.Table(
        "pokemons",
        metadata,
        sql.Column("id", sql.Integer, primary_key=True),
        sql.Column("name", sql.String),
        sql.Column("HP", sql.Integer),
        sql.Column("Attack", sql.Integer),
        sql.Column("Defense", sql.Integer),
        sql.Column("Sp. Attack", sql.Integer),
        sql.Column("Sp. Defense", sql.Integer),
        sql.Column("Speed", sql.Integer),
    )

    t_types = sql.Table(
        "types",
        metadata,
        sql.Column("id", sql.Integer, primary_key=True),
        sql.Column("name", sql.String),
    )

    t_pokemon_types = sql.Table(
        "pokemon_types",
        metadata,
        sql.Column("pokemon_id", sql.Integer, primary_key=True),
        sql.Column("type_id", sql.Integer, primary_key=True),
        sql.Column("no", sql.Integer, primary_key=True),
    )

    # Create the tables
    metadata.create_all(engine)

    # Insert the fixtures
    with engine.connect() as conn:
        conn.execute(t_pokemons.insert(), pokemons)
        conn.execute(t_types.insert(), types)
        conn.execute(t_pokemon_types.insert(), pokemon_types)
        conn.commit()

    return engine


def test_iter_sql(pokedb):
    with pokedb.connect() as conn:
        dataset = stream.iter_sql(query="SELECT * FROM pokemons;", conn=conn)
        x, y = next(dataset)
        assert x["name"] == "Bulbasaur"
        assert y is None

    # This raises an exception because the resource is closed...
    with pytest.raises(sql.exc.ResourceClosedError):
        for x, y in stream.iter_sql(query="SELECT * FROM pokemons;", conn=conn):
            pass

    # ... and yet we can still stream over the results because SQLAlchemy prefetches them
    x, y = next(dataset)
    assert x["name"] == "Ivysaur"

    # The Pokedex from generation 1 contains 151 pokemons, and we've already seen 2 of them
    assert sum(1 for _ in dataset) == 149

    # Check that the stream is depleted
    assert sum(1 for _ in dataset) == 0


def test_iter_sql_join(pokedb):
    query = """
        SELECT
            p.name,
            t1.name AS type_1,
            t2.name AS type_2
        FROM
            pokemons p,
            pokemon_types pt1,
            pokemon_types pt2,
            types t1,
            types t2
        WHERE
            pt1.no = 1 AND
            pt1.pokemon_id = p.id AND
            pt1.type_id = t1.id AND

            pt2.no = 2 AND
            pt2.pokemon_id = p.id AND
            pt2.type_id = t2.id;
    """

    with pokedb.connect() as conn:
        dataset = stream.iter_sql(query=query, conn=conn)
        x, y = next(dataset)
        assert x["name"] == "Bulbasaur"
        assert x["type_1"] == "Grass"
        assert x["type_2"] == "Poison"


# river/stream/tweet_stream.py

from __future__ import annotations

import json
from urllib.parse import urljoin


class TwitterLiveStream:
    """Twitter API v2 live stream client.

        This client gives access to a live stream of Tweets. That is, Tweets that have just been
        published. This is different to `stream.TwitterRecentStream`, which also covers Tweets that
        have been published over recent days, and not necessarily in real-time.

        A list of filtering rules has to be provided. For instance, this allows focusing on a subset of
        topics and/or users.

        !!! note

            Using this requires having the [`requests`](https://requests.readthedocs.io/en/latest/)
            package installed.

        Parameters
        ----------
        rules
            See the documentation[^2] for a comprehensive overview of filtering rules.
        bearer_token
            A bearer token that is available in each account's developer portal.

        Examples
        --------

        The live stream is instantiated by passing a list of filtering rules, as well as a bearer
        token. For instance, we can listen to all the breaking news Tweets from the BBC and CNN.

        >>> from river import stream

        >>> tweets = stream.TwitterLiveStream(
        ...     rules=["from:BBCBreaking", "from:cnnbrk"],
        ...     bearer_token="<insert_bearer_token>"
        ... )

        The stream can then be iterated over, possibly in an infinite loop. This will listen to the
        live feed of Tweets and produce a Tweet right after it's been published.

        ```py
        import logging

        while True:
            try:
                for tweet in tweets:
                    print(tweet)
            except requests.exceptions.RequestException as e:
                logging.warning(str(e))
                time.sleep(10)
        ```

        Here's a Tweet example:

        ```py
        {
            'data': {
                'author_id': '428333',
                'created_at': '2022-08-26T12:59:48.000Z',
                'id': '1563149212774445058',
                'text': "Ukraine's Zaporizhzhia nuclear power plant, which is currently held by
    Russian forces, has been reconnected to Ukraine's electricity grid, according to the
    country's nuclear operator https://t.co/xfylkBs4JR"
            },
            'includes': {
                'users': [
                    {
                        'created_at': '2007-01-02T01:48:14.000Z',
                        'id': '428333',
                        'name': 'CNN Breaking News',
                        'username': 'cnnbrk'
                    }
                ]
            },
            'matching_rules': [{'id': '1563148866333151233', 'tag': 'from:cnnbrk'}]
        }
        ```

        References
        ----------
        [^1]: [Filtered stream introduction](https://developer.twitter.com/en/docs/twitter-api/tweets/filtered-stream/introduction)
        [^2]: [Building rules for filtered stream](https://developer.twitter.com/en/docs/twitter-api/tweets/filtered-stream/integrate/build-a-rule)
        [^3]: [Stream Tweets in real-time](https://developer.twitter.com/en/docs/tutorials/stream-tweets-in-real-time)

    """

    def __init__(self, rules, bearer_token):
        self.rules = rules
        self.bearer_token = bearer_token

    def _request(self, method, endpoint, **kwargs):
        import requests

        url = urljoin("https://api.twitter.com/2/", endpoint)
        r = requests.request(
            method,
            url,
            headers={
                "Authorization": f"Bearer {self.bearer_token}",
                "User-Agent": "v2FilteredStreamPython",
            },
            **kwargs,
        )
        r.raise_for_status()
        return r

    def _get_rules(self):
        return self._request("GET", "tweets/search/stream/rules").json()

    def _delete_rules(self, rules):
        if rule_ids := [rule["id"] for rule in rules.get("data", [])]:
            payload = {"delete": {"ids": rule_ids}}
            return self._request("POST", "tweets/search/stream/rules", json=payload).json()

    def _set_rules(self, rules):
        payload = {"add": rules}
        return self._request("POST", "tweets/search/stream/rules", json=payload).json()

    def __iter__(self):
        existing_rules = self._get_rules()
        self._delete_rules(existing_rules)
        self._set_rules([{"value": rule, "tag": rule} for rule in existing_rules])
        params = {
            "tweet.fields": "created_at",
            "expansions": "author_id",
            "user.fields": "created_at",
        }
        r = self._request("GET", "tweets/search/stream", stream=True, params=params)
        for response_line in r.iter_lines():
            if not response_line:
                continue
            tweet = json.loads(response_line)
            yield tweet


# river/stream/twitch_chat_stream.py

from __future__ import annotations

import dataclasses
import datetime as dt
import enum
import re
import socket
from collections.abc import Iterator

SERVER = "irc.chat.twitch.tv"
PORT = 6667
BUFFER_SIZE = 2048  # bytes
ENCODING = "utf-8"
TIMEOUT = 60  # seconds
CHAT_ITEM_PATTERN = r":(.+)\!.*@.*\.tmi\.twitch\.tv PRIVMSG #([a-zA-Z0-9_]+) :(.+)"
PING_PATTERN = r"\bPING :tmi\.twitch\.tv\b"


class IrcMessage(enum.Enum):
    PASS = enum.auto()
    NICK = enum.auto()
    JOIN = enum.auto()
    PONG = enum.auto()


@dataclasses.dataclass
class ChatMessageItem:
    dt: dt.datetime
    channel: str
    username: str
    msg: str


class TwitchChatStream:
    """Twitch chat stream client.

    This client gives access to a live stream of chat messages in Twitch channels using IRC protocol.
    You need to have a Twitch account and receive an OAuth token from https://twitchapps.com/tmi/.

    Parameters
    ----------
    nickname
        The nickname of your account.
    token
        OAuth token which has been generated.
    channels
        A list of channel names like `["asmongold", "shroud"]` you want to collect messages from.
    buffer_size
        Size of buffer in bytes used for receiving responses from Twitch with IRC (default 2 kB).
    timeout
        A timeout value in seconds for waiting response from Twitch (default 60s). It can be useful if all requested channels are offline or chat is not active enough.

    Returns
    -------
    Stream items are python dictionaries with keys:
    dt (datetime)
        Datetime indicating when response with this message has been received from Twitch.
    channel (str)
        A channel where message has been posted.
    username (str)
        User who posted the message.
    msg (str)
        The message itself.

    Examples
    --------
    The live stream is instantiated by passing your Twitch account nickname, OAuth token and list of channels. Other parameters are optional.

    >>> from river import stream

    >>> twitch_chat = stream.TwitchChatStream(
    ...     nickname="twitch_user1",
    ...     token="oauth:okrip6j6fjio8n5xpy2oum1lph4fbve",
    ...     channels=["asmongold", "shroud"]
    ... )

    The stream can be iterated over like this:

    ```py
    for item in twitch_chat:
        print(item)
    ```

    Here's a single stream item example:
    ```py
    {
        'dt': datetime.datetime(2022, 9, 14, 10, 33, 37, 989560),
        'channel': 'asmongold',
        'username': 'moojiejaa',
        'msg': 'damn this chat mod are wild'
    }
    ```

    References
    ----------
    [^1]: [Twitch IRC doc](https://dev.twitch.tv/docs/irc)

    """

    def __init__(
        self,
        nickname: str,
        token: str,
        channels: list[str],
        buffer_size: int = BUFFER_SIZE,
        timeout: int = TIMEOUT,
    ):
        self.nickname = nickname
        self.token = token
        self.channels = channels
        self.buffer_size = buffer_size
        self.timeout = timeout
        self.chat_item_pattern = re.compile(CHAT_ITEM_PATTERN)
        self.ping_pattern = re.compile(PING_PATTERN)

    def _send(self, s: socket.socket, msg: IrcMessage, payload: str | None = None) -> None:
        text = msg.name
        if payload:
            text += f" {payload}"
        s.send(f"{text}\n".encode(ENCODING))

    def _setup_connection(self, s: socket.socket) -> None:
        s.connect((SERVER, PORT))
        self._send(s, IrcMessage.PASS, self.token)
        self._send(s, IrcMessage.NICK, self.nickname)
        for chan in self.channels:
            self._send(s, IrcMessage.JOIN, f"#{chan}")
        s.settimeout(self.timeout)

    def _is_ping(self, resp: str) -> bool:
        return bool(self.ping_pattern.search(resp))

    def _extract_chat_messages(self, resp: str, dt: dt.datetime) -> Iterator[ChatMessageItem]:
        for m in self.chat_item_pattern.finditer(resp):
            if not m or not m.groups():
                continue
            user, chan, msg = m.groups()
            yield ChatMessageItem(dt=dt, channel=chan, username=user, msg=msg.strip())

    def _gen_items(self, sock: socket.socket) -> Iterator[ChatMessageItem]:
        while True:
            try:
                data = sock.recv(self.buffer_size)
                if not data:
                    continue
                resp = data.decode(ENCODING)
                now = dt.datetime.now()
            except socket.timeout as e:
                raise TimeoutError(f"Twitch did not respond in {self.timeout:,d} seconds") from e
            except UnicodeDecodeError:
                continue

            if self._is_ping(resp):
                self._send(sock, IrcMessage.PONG, ":tmi.twitch.tv")

            yield from self._extract_chat_messages(resp, now)

    def __iter__(self) -> Iterator[dict]:
        with socket.socket() as sock:
            self._setup_connection(sock)
            yield from map(dataclasses.asdict, self._gen_items(sock))


# river/stream/utils.py

from __future__ import annotations

import functools
import gzip
import io
import os
import zipfile


def open_filepath(filepath_or_buffer, compression):
    # Determine the compression from the file extension if "infer" has been specified
    if compression == "infer":
        _, ext = os.path.splitext(filepath_or_buffer)
        compression = {".gz": "gzip", ".zip": "zip"}.get(ext)

    def open_zipfile(path):
        with zipfile.ZipFile(path, "r") as zf:
            f = zf.open(zf.namelist()[0], "r")
            f = io.TextIOWrapper(f)
            return f

    # Determine the file opening method from the compression
    open_func = {
        None: open,
        "gzip": functools.partial(gzip.open, mode="rt"),
        "zip": open_zipfile,
    }[compression]

    # Open the file using the opening method
    return open_func(filepath_or_buffer)


# river/time_series/__init__.py

"""Time series forecasting."""
from __future__ import annotations

from . import base
from .evaluate import evaluate, iter_evaluate
from .holt_winters import HoltWinters
from .metrics import ForecastingMetric, HorizonAggMetric, HorizonMetric
from .snarimax import SNARIMAX

__all__ = [
    "base",
    "evaluate",
    "iter_evaluate",
    "ForecastingMetric",
    "HorizonAggMetric",
    "HorizonMetric",
    "HoltWinters",
    "SNARIMAX",
]


# river/time_series/base.py

from __future__ import annotations

import abc

from river import base

__all__ = ["Forecaster"]


class Forecaster(base.Estimator):
    @property
    def _supervised(self):
        return True

    @abc.abstractmethod
    def learn_one(self, y: float, x: dict | None = None) -> Forecaster:
        """Updates the model.

        Parameters
        ----------
        y
            In the literature this is called the endogenous variable.
        x
            Optional additional features to learn from. In the literature these are called the
            exogenous variables.

        """
        raise NotImplementedError

    @abc.abstractmethod
    def forecast(self, horizon: int, xs: list[dict] | None = None) -> list:
        """Makes forecast at each step of the given horizon.

        Parameters
        ----------
        horizon
            The number of steps ahead to forecast.
        xs
            The set of optional additional features. If given, then it's length should be equal to
            the horizon.

        """


# river/time_series/evaluate.py

from __future__ import annotations

import collections
import numbers
import typing

from river import base, metrics, time_series

TimeSeries = typing.Iterator[
    typing.Tuple[  # noqa: UP006
        typing.Union[dict, None],  # noqa: UP007
        numbers.Number,
        typing.Union[typing.List[dict], None],  # noqa: UP006, UP007
        typing.List[numbers.Number],  # noqa: UP006
    ]
]


def _iter_with_horizon(dataset: base.typing.Dataset, horizon: int) -> TimeSeries:
    """

    Examples
    --------

    >>> from river import datasets
    >>> from river.time_series.evaluate import _iter_with_horizon

    >>> dataset = datasets.AirlinePassengers()

    >>> for x, y, x_horizon, y_horizon in _iter_with_horizon(dataset.take(8), horizon=3):
    ...     print(x['month'].strftime('%Y-%m-%d'), y)
    ...     print([x['month'].strftime('%Y-%m-%d') for x in x_horizon])
    ...     print(list(y_horizon))
    ...     print('---')
    1949-01-01 112
    ['1949-02-01', '1949-03-01', '1949-04-01']
    [118, 132, 129]
    ---
    1949-02-01 118
    ['1949-03-01', '1949-04-01', '1949-05-01']
    [132, 129, 121]
    ---
    1949-03-01 132
    ['1949-04-01', '1949-05-01', '1949-06-01']
    [129, 121, 135]
    ---
    1949-04-01 129
    ['1949-05-01', '1949-06-01', '1949-07-01']
    [121, 135, 148]
    ---
    1949-05-01 121
    ['1949-06-01', '1949-07-01', '1949-08-01']
    [135, 148, 148]
    ---

    """

    x_horizon: collections.deque[dict] = collections.deque(maxlen=horizon)
    y_horizon: collections.deque = collections.deque(maxlen=horizon)

    stream = iter(dataset)

    for _ in range(horizon):
        x, y = next(stream)
        x_horizon.append(x)
        y_horizon.append(y)

    for x, y in stream:
        x_now = x_horizon.popleft()
        y_now = y_horizon.popleft()
        x_horizon.append(x)
        y_horizon.append(y)
        yield x_now, y_now, x_horizon, y_horizon  # type: ignore


def iter_evaluate(
    dataset: base.typing.Dataset,
    model: time_series.base.Forecaster,
    metric: metrics.base.RegressionMetric,
    horizon: int,
    agg_func: typing.Callable[[list[float]], float] | None = None,
    grace_period: int | None = None,
):
    """Evaluates the performance of a forecaster on a time series dataset and yields results.

    This does exactly the same as `evaluate.progressive_val_score`. The only difference is that
    this function returns an iterator, yielding results at every step. This can be useful if you
    want to have control over what you do with the results. For instance, you might want to plot
    the results.

    Parameters
    ----------
    dataset
        A sequential time series.
    model
        A forecaster.
    metric
        A regression metric.
    horizon
    agg_func
    grace_period
        Initial period during which the metric is not updated. This is to fairly evaluate models
        which need a warming up period to start producing meaningful forecasts. The value of this
        parameter is equal to the horizon by default.

    """

    horizon_metric = (
        time_series.HorizonAggMetric(metric, agg_func)
        if agg_func
        else time_series.HorizonMetric(metric)
    )
    steps = _iter_with_horizon(dataset, horizon)

    grace_period = horizon if grace_period is None else grace_period
    for _ in range(grace_period):
        x, y, x_horizon, y_horizon = next(steps)
        model.learn_one(y=y, x=x)  # type: ignore

    for x, y, x_horizon, y_horizon in steps:
        y_pred = model.forecast(horizon, xs=x_horizon)
        horizon_metric.update(y_horizon, y_pred)
        model.learn_one(y=y, x=x)  # type: ignore
        yield x, y, y_pred, horizon_metric


def evaluate(
    dataset: base.typing.Dataset,
    model: time_series.base.Forecaster,
    metric: metrics.base.RegressionMetric,
    horizon: int,
    agg_func: typing.Callable[[list[float]], float] | None = None,
    grace_period: int | None = None,
) -> time_series.HorizonMetric:
    """Evaluates the performance of a forecaster on a time series dataset.

    To understand why this method is useful, it's important to understand the difference between
    nowcasting and forecasting. Nowcasting is about predicting a value at the next time step. This
    can be seen as a special case of regression, where the value to predict is the value at the
    next time step. In this case, the `evaluate.progressive_val_score` function may be used to
    evaluate a model via progressive validation.

    Forecasting models can also be evaluated via progressive validation. This is the purpose of
    this function. At each time step `t`, the forecaster is asked to predict the values at `t + 1`,
    `t + 2`, ..., `t + horizon`. The performance at each time step is measured and returned.

    Parameters
    ----------
    dataset
        A sequential time series.
    model
        A forecaster.
    metric
        A regression metric.
    horizon
    agg_func
    grace_period
        Initial period during which the metric is not updated. This is to fairly evaluate models
        which need a warming up period to start producing meaningful forecasts. The value of this
        parameter is equal to the horizon by default.

    """

    horizon_metric = None
    steps = iter_evaluate(dataset, model, metric, horizon, agg_func, grace_period)
    for *_, horizon_metric in steps:
        pass

    return horizon_metric


# river/time_series/holt_winters.py

from __future__ import annotations

import operator
import statistics
from collections import deque

from river import time_series

__all__ = ["HoltWinters"]


class Component(deque):
    ...


class AdditiveLevel(Component):
    def __init__(self, alpha):
        super().__init__([], maxlen=2)
        self.alpha = alpha

    def update(self, y, trend, season):
        self.append(
            self.alpha * (y - (season[-season.seasonality] if season else 0))
            + (1 - self.alpha) * (self[-1] + (trend[-1] if trend else 0))
        )


class MultiplicativeLevel(Component):
    def __init__(self, alpha):
        super().__init__([], maxlen=2)
        self.alpha = alpha

    def update(self, y, trend, season):
        self.append(
            self.alpha * (y / (season[-season.seasonality] if season else 1))
            + (1 - self.alpha) * (self[-1] + (trend[-1] if trend else 0))
        )


class Trend(Component):
    def __init__(self, beta):
        super().__init__([], maxlen=2)
        self.beta = beta

    def update(self, y, level):
        self.append(self.beta * (level[-1] - level[-2]) + (1 - self.beta) * self[-1])


class AdditiveSeason(Component):
    def __init__(self, gamma, seasonality):
        super().__init__([], maxlen=seasonality + 1)
        self.gamma = gamma
        self.seasonality = seasonality

    def update(self, y, level, trend):
        self.append(
            self.gamma * (y - level[-2] - trend[-2]) + (1 - self.gamma) * self[-self.seasonality]
        )


class MultiplicativeSeason(Component):
    def __init__(self, gamma, seasonality):
        super().__init__([], maxlen=seasonality + 1)
        self.gamma = gamma
        self.seasonality = seasonality

    def update(self, y, level, trend):
        self.append(
            self.gamma * y / (level[-2] + trend[-2]) + (1 - self.gamma) * self[-self.seasonality]
        )


class HoltWinters(time_series.base.Forecaster):
    r"""Holt-Winters forecaster.

    This is a standard implementation of the Holt-Winters forecasting method. Certain
    parametrisations result in special cases, such as simple exponential smoothing.

    Optimal parameters and initialisation values can be determined in a batch setting. However, in
    an online setting, it is necessary to wait and observe enough values. The first
    `k = max(2, seasonality)` values are indeed used to initialize the components.

    **Level initialization**

    $$l = \frac{1}{k} \sum_{i=1}{k} y_i$$

    **Trend initialization**

    $$t = \frac{1}{k - 1} \sum_{i=2}{k} y_i - y_{i-1}$$

    **Trend initialization**

    $$s_i = \frac{y_i}{k}$$

    Parameters
    ----------
    alpha
        Smoothing parameter for the level.
    beta
        Smoothing parameter for the trend.
    gamma
        Smoothing parameter for the seasonality.
    seasonality
        The number of periods in a season. For instance, this should be 4 for quarterly data,
        and 12 for yearly data.
    multiplicative
        Whether or not to use a multiplicative formulation.

    Examples
    --------

    >>> from river import datasets
    >>> from river import metrics
    >>> from river import time_series

    >>> dataset = datasets.AirlinePassengers()

    >>> model = time_series.HoltWinters(
    ...     alpha=0.3,
    ...     beta=0.1,
    ...     gamma=0.6,
    ...     seasonality=12,
    ...     multiplicative=True
    ... )

    >>> metric = metrics.MAE()

    >>> time_series.evaluate(
    ...     dataset,
    ...     model,
    ...     metric,
    ...     horizon=12
    ... )
    +1  MAE: 25.899087
    +2  MAE: 26.26131
    +3  MAE: 25.735903
    +4  MAE: 25.625678
    +5  MAE: 26.093842
    +6  MAE: 26.90249
    +7  MAE: 28.634398
    +8  MAE: 29.284769
    +9  MAE: 31.018351
    +10 MAE: 32.252349
    +11 MAE: 33.518946
    +12 MAE: 33.975057

    References
    ----------
    [^1]: [Exponential smoothing — Wikipedia](https://www.wikiwand.com/en/Exponential_smoothing)
    [^2]: [Exponential smoothing — Forecasting: Principles and Practice](https://otexts.com/fpp2/expsmooth.html)
    [^3]: [What is Exponential Smoothing? — Engineering statistics handbook](https://www.itl.nist.gov/div898/handbook/pmc/section4/pmc43.htm)

    """

    def __init__(
        self,
        alpha,
        beta=None,
        gamma=None,
        seasonality=0,
        multiplicative=False,
    ):
        if seasonality and gamma is None:
            raise ValueError("gamma must be set if seasonality is set")

        if gamma and beta is None:
            raise ValueError("beta must be set if gamma is set")

        self.alpha = alpha
        self.beta = beta
        self.gamma = gamma
        self.seasonality = seasonality
        self.multiplicative = multiplicative
        self.level = MultiplicativeLevel(alpha) if multiplicative else AdditiveLevel(alpha)
        self.trend = Trend(beta) if beta else None
        self.season = (
            (
                MultiplicativeSeason(gamma, seasonality)
                if multiplicative
                else AdditiveSeason(gamma, seasonality)
            )
            if seasonality
            else None
        )
        self._first_values = []
        self._initialized = False

    def learn_one(self, y, x=None):
        if self._initialized:
            self.level.update(y, self.trend, self.season)
            if self.trend is not None:
                self.trend.update(y, self.level)
            if self.season is not None:
                self.season.update(y, self.level, self.trend)
            return self

        self._first_values.append(y)
        if len(self._first_values) < max(2, self.seasonality):
            return self

        # The components can be initialized now that enough values have been observed
        self.level.append(statistics.mean(self._first_values))
        diffs = [b - a for a, b in zip(self._first_values[:-1], self._first_values[1:])]
        if self.trend is not None:
            self.trend.append(statistics.mean(diffs))
        if self.season is not None:
            self.season.extend([y / self.level[-1] for y in self._first_values])

        self._initialized = True

        return self

    def forecast(self, horizon, xs=None):
        op = operator.mul if self.multiplicative else operator.add
        return [
            op(
                self.level[-1] + ((h + 1) * self.trend[-1] if self.trend else 0),
                (
                    self.season[-self.seasonality + h % self.seasonality]
                    if self.season
                    else (1 if self.multiplicative else 0)
                ),
            )
            for h in range(horizon)
        ]


# river/time_series/metrics.py

from __future__ import annotations

import abc
import typing
from numbers import Number

from river import base, metrics


class ForecastingMetric(base.Base, abc.ABC):
    @abc.abstractmethod
    def update(self, y_true: list[Number], y_pred: list[Number]) -> ForecastingMetric:
        """Update the metric at each step along the horizon.

        Parameters
        ----------
        y_true
            Ground truth values at each time step of the horizon.
        y_pred
            Predicted values at each time step of the horizon.

        Returns
        -------
        self

        """

    @abc.abstractmethod
    def get(self) -> list[float]:
        """Return the current performance along the horizon.

        Returns
        -------
        The current performance.

        """


class HorizonMetric(ForecastingMetric):
    """Measures performance at each time step ahead.

    This allows to measure the performance of a model at each time step along the horizon. A copy
    of the provided regression metric is made for each time step. At each time step ahead, the
    metric is thus evaluated on each prediction for said time step, and not for the time steps before
    or after that.

    Parameters
    ----------
    metric
        A regression metric.

    Examples
    --------

    This is used internally by the `time_series.evaluate` function.

    >>> from river import datasets
    >>> from river import metrics
    >>> from river import time_series

    >>> metric = time_series.evaluate(
    ...     dataset=datasets.AirlinePassengers(),
    ...     model=time_series.HoltWinters(alpha=0.1),
    ...     metric=metrics.MAE(),
    ...     horizon=4
    ... )

    >>> metric
    +1 MAE: 40.931286
    +2 MAE: 42.667998
    +3 MAE: 44.158092
    +4 MAE: 43.849617

    """

    def __init__(self, metric: metrics.base.RegressionMetric):
        self.metric = metric
        self.metrics: list[metrics.base.RegressionMetric] = []

    def update(self, y_true, y_pred):
        for t, (yt, yp) in enumerate(zip(y_true, y_pred)):
            try:
                metric = self.metrics[t]
            except IndexError:
                metric = self.metric.clone()
                self.metrics.append(metric)

            metric.update(yt, yp)

        return self

    def get(self):
        return [metric.get() for metric in self.metrics]

    def __repr__(self):
        prefixes = [f"+{t+1}" for t in range(len(self.metrics))]
        prefix_pad = max(map(len, prefixes))
        return "\n".join(
            f"{prefix:<{prefix_pad}} {metric}" for prefix, metric in zip(prefixes, self.metrics)
        )


class HorizonAggMetric(HorizonMetric):
    """Same as `HorizonMetric`, but aggregates the result based on an provided function.

    This allows, for instance, to measure the average performance of a forecasting model along the
    horizon.

    Parameters
    ----------
    metric
        A regression metric.
    agg_func
        A function that takes as input a list of floats and outputs a single float. You may want to
        `min`, `max`, as well as `statistics.mean` and `statistics.median`.

    Examples
    --------

    This is used internally by the `time_series.evaluate` function when you pass an `agg_func`.

    >>> import statistics
    >>> from river import datasets
    >>> from river import metrics
    >>> from river import time_series

    >>> metric = time_series.evaluate(
    ...     dataset=datasets.AirlinePassengers(),
    ...     model=time_series.HoltWinters(alpha=0.1),
    ...     metric=metrics.MAE(),
    ...     agg_func=statistics.mean,
    ...     horizon=4
    ... )

    >>> metric
    mean(MAE): 42.901748

    """

    def __init__(
        self, metric: metrics.base.RegressionMetric, agg_func: typing.Callable[[list[float]], float]
    ):
        super().__init__(metric)
        self.agg_func = agg_func

    def get(self):
        return self.agg_func(super().get())

    def __repr__(self):
        name = f"{self.agg_func.__name__}({self.metric.__class__.__name__})"
        return f"{name}: {self.get():{self.metric._fmt}}".rstrip("0")


# river/time_series/snarimax.py

from __future__ import annotations

import collections
import itertools
import math

from river import base, linear_model, preprocessing, time_series

__all__ = ["SNARIMAX"]


class Differencer:
    """A time series differencer.

    References
    ----------
    [^1]: [Stationarity and differencing](https://otexts.com/fpp2/stationarity.html)
    [^2]: [Backshift notation](https://otexts.com/fpp2/backshift.html)

    """

    def __init__(self, d, m=1):
        self.d = d
        self.m = m

        def n_choose_k(n, k):
            f = math.factorial
            return f(n) // f(k) // f(n - k)

        self.coeffs = {0: 1}
        for k in range(1, d + 1):
            t = k * m
            coeff = (-1 if k % 2 else 1) * n_choose_k(n=d, k=k)
            self.coeffs[t] = coeff

    @property
    def n_required_past_values(self):
        return max(self.coeffs)

    @classmethod
    def from_coeffs(cls, coeffs):
        obj = cls(0, 0)
        obj.coeffs = coeffs
        return obj

    def __mul__(self, other):
        """Compose two differencers together."""
        coeffs = collections.defaultdict(int)

        for (t1, c1), (t2, c2) in itertools.product(self.coeffs.items(), other.coeffs.items()):
            coeffs[t1 + t2] += c1 * c2

        # Remove 0 coefficients
        for t, c in list(coeffs.items()):
            if c == 0:
                del coeffs[t]

        return Differencer.from_coeffs(dict(coeffs))

    def diff(self, p, Y: list):
        """Differentiate by applying each coefficient c at each index t.

        Parameters
        ----------
        Y
            The window of previous values. The first element is assumed to be the most recent
            value.

        """
        total = 0
        for t, c in self.coeffs.items():
            try:
                total += (c * Y[t - 1]) if t else p
            except IndexError:
                break
        return total

    def undiff(self, p, Y: list):
        """Differentiate by applying each coefficient c at each index t.

        Parameters
        ----------
        Y
            The window of previous values. The first element is assumed to be the most recent
            value.

        """
        total = p
        for t, c in self.coeffs.items():
            try:
                if t:
                    total -= c * Y[t - 1]
            except IndexError:
                break
        return total


class SNARIMAX(time_series.base.Forecaster):
    """SNARIMAX model.

    SNARIMAX stands for (S)easonal (N)on-linear (A)uto(R)egressive (I)ntegrated (M)oving-(A)verage
    with e(X)ogenous inputs model.

    This model generalizes many established time series models in a single interface that can be
    trained online. It assumes that the provided training data is ordered in time and is uniformly
    spaced. It is made up of the following components:

    - S (Seasonal)
    - N (Non-linear): Any online regression model can be used, not necessarily a linear regression
        as is done in textbooks.
    - AR (Autoregressive): Lags of the target variable are used as features.
    - I (Integrated): The model can be fitted on a differenced version of a time series. In this
        context, integration is the reverse of differencing.
    - MA (Moving average): Lags of the errors are used as features.
    - X (Exogenous): Users can provide additional features. Care has to be taken to include
        features that will be available both at training and prediction time.

    Each of these components can be switched on and off by specifying the appropriate parameters.
    Classical time series models such as AR, MA, ARMA, and ARIMA can thus be seen as special
    parametrizations of the SNARIMAX model.

    This model is tailored for time series that are homoskedastic. In other words, it might not
    work well if the variance of the time series varies widely along time.

    Parameters
    ----------
    p
        Order of the autoregressive part. This is the number of past target values that will be
        included as features.
    d
        Differencing order.
    q
        Order of the moving average part. This is the number of past error terms that will be
        included as features.
    m
        Season length used for extracting seasonal features. If you believe your data has a
        seasonal pattern, then set this accordingly. For instance, if the data seems to exhibit
        a yearly seasonality, and that your data is spaced by month, then you should set this
        to 12. Note that for this parameter to have any impact you should also set at least one
        of the `p`, `d`, and `q` parameters.
    sp
        Seasonal order of the autoregressive part. This is the number of past target values
        that will be included as features.
    sd
        Seasonal differencing order.
    sq
        Seasonal order of the moving average part. This is the number of past error terms that
        will be included as features.
    regressor
        The online regression model to use. By default, a `preprocessing.StandardScaler`
        piped with a `linear_model.LinearRegression` will be used.

    Attributes
    ----------
    differencer : Differencer
    y_trues : collections.deque
        The `p` past target values.
    errors : collections.deque
        The `q` past error values.

    Examples
    --------

    >>> import datetime as dt
    >>> from river import datasets
    >>> from river import time_series
    >>> from river import utils

    >>> period = 12
    >>> model = time_series.SNARIMAX(
    ...     p=period,
    ...     d=1,
    ...     q=period,
    ...     m=period,
    ...     sd=1
    ... )

    >>> for t, (x, y) in enumerate(datasets.AirlinePassengers()):
    ...     model = model.learn_one(y)

    >>> horizon = 12
    >>> future = [
    ...     {'month': dt.date(year=1961, month=m, day=1)}
    ...     for m in range(1, horizon + 1)
    ... ]
    >>> forecast = model.forecast(horizon=horizon)
    >>> for x, y_pred in zip(future, forecast):
    ...     print(x['month'], f'{y_pred:.3f}')
    1961-01-01 491.988
    1961-02-01 447.593
    1961-03-01 481.405
    1961-04-01 566.278
    1961-05-01 551.561
    1961-06-01 605.414
    1961-07-01 711.140
    1961-08-01 668.204
    1961-09-01 570.517
    1961-10-01 549.589
    1961-11-01 466.344
    1961-12-01 506.945

    Classic ARIMA models learn solely on the time series values. You can also include features
    built at each step.

    >>> import calendar
    >>> import math
    >>> from river import compose
    >>> from river import linear_model
    >>> from river import optim
    >>> from river import preprocessing

    >>> def get_month_distances(x):
    ...     return {
    ...         calendar.month_name[month]: math.exp(-(x['month'].month - month) ** 2)
    ...         for month in range(1, 13)
    ...     }

    >>> def get_ordinal_date(x):
    ...     return {'ordinal_date': x['month'].toordinal()}

    >>> extract_features = compose.TransformerUnion(
    ...     get_ordinal_date,
    ...     get_month_distances
    ... )

    >>> model = (
    ...     extract_features |
    ...     time_series.SNARIMAX(
    ...         p=1,
    ...         d=0,
    ...         q=0,
    ...         m=12,
    ...         sp=3,
    ...         sq=6,
    ...         regressor=(
    ...             preprocessing.StandardScaler() |
    ...             linear_model.LinearRegression(
    ...                 intercept_init=110,
    ...                 optimizer=optim.SGD(0.01),
    ...                 intercept_lr=0.3
    ...             )
    ...         )
    ...     )
    ... )

    >>> for x, y in datasets.AirlinePassengers():
    ...     model = model.learn_one(x, y)

    >>> forecast = model.forecast(horizon=horizon)
    >>> for x, y_pred in zip(future, forecast):
    ...     print(x['month'], f'{y_pred:.3f}')
    1961-01-01 446.874
    1961-02-01 423.998
    1961-03-01 439.957
    1961-04-01 457.958
    1961-05-01 457.303
    1961-06-01 496.554
    1961-07-01 553.798
    1961-08-01 551.388
    1961-09-01 479.620
    1961-10-01 440.613
    1961-11-01 409.914
    1961-12-01 433.774

    References
    ----------
    [^1]: [ARMA - Wikipedia](https://www.wikiwand.com/en/Autoregressive%E2%80%93moving-average_model)
    [^2]: [NARX - Wikipedia](https://www.wikiwand.com/en/Nonlinear_autoregressive_exogenous_model)
    [^3]: [ARIMA - Forecasting: Principles and Practice](https://otexts.com/fpp2/arima.html)
    [^4]: [Anava, O., Hazan, E., Mannor, S. and Shamir, O., 2013, June. Online learning for time series prediction. In Conference on learning theory (pp. 172-184)](https://arxiv.org/pdf/1302.6927.pdf)

    """

    def __init__(
        self,
        p: int,
        d: int,
        q: int,
        m: int = 1,
        sp: int = 0,
        sd: int = 0,
        sq: int = 0,
        regressor: base.Regressor | None = None,
    ):
        self.p = p
        self.d = d
        self.q = q
        self.m = m
        self.sp = sp
        self.sd = sd
        self.sq = sq
        self.regressor = (
            regressor
            if regressor is not None
            else preprocessing.StandardScaler() | linear_model.LinearRegression()
        )
        self.differencer = Differencer(d=d, m=1) * Differencer(d=sd, m=m)
        self.y_hist: collections.deque[float] = collections.deque(maxlen=d + m * sd)
        self.y_diff: collections.deque[float] = collections.deque(maxlen=max(p, m * sp))
        self.errors: collections.deque[float] = collections.deque(maxlen=max(q, m * sq))

    def _add_lag_features(self, x, Y, errors):
        if x is None:
            x = {}

        # AR
        for t in range(self.p):
            try:
                x[f"y-{t+1}"] = Y[t]
            except IndexError:
                break

        # Seasonal AR
        for t in range(self.m - 1, self.m * self.sp, self.m):
            try:
                x[f"sy-{t+1}"] = Y[t]
            except IndexError:
                break

        # MA
        for t in range(self.q):
            try:
                x[f"e-{t+1}"] = errors[t]
            except IndexError:
                break

        # Seasonal MA
        for t in range(self.m - 1, self.m * self.sq, self.m):
            try:
                x[f"se-{t+1}"] = errors[t]
            except IndexError:
                break

        return x

    def learn_one(self, y, x=None):
        # It isn't possible to difference the time series if enough values have not been seen
        if len(self.y_hist) >= self.differencer.n_required_past_values:
            x = self._add_lag_features(x=x, Y=self.y_diff, errors=self.errors)

            # The regressor learns on the differenced version of the time series
            # If d=0 and/or sd=0, then self.differencer.diff is a no-op
            y_diff = self.differencer.diff(y, self.y_hist)
            self.y_diff.appendleft(y_diff)

            y_pred = self.regressor.predict_one(x)
            self.errors.appendleft(y_diff - y_pred)
            self.regressor.learn_one(x, y_diff)

        self.y_hist.appendleft(y)

        return self

    def forecast(self, horizon, xs=None):
        if xs is None:
            xs = [{}] * horizon

        if len(xs) != horizon:
            raise ValueError("the length of xs should be equal to the specified horizon")

        y_hist = collections.deque(self.y_hist)
        y_diff = collections.deque(self.y_diff)
        errors = collections.deque(self.errors)
        forecasts = [None] * horizon

        for t, x in enumerate(xs):
            x = self._add_lag_features(x=x, Y=y_diff, errors=errors)

            y_pred = self.regressor.predict_one(x)
            y_diff.appendleft(y_pred)

            forecasts[t] = self.differencer.undiff(y_pred, y_hist)
            y_hist.appendleft(forecasts[t])

            errors.appendleft(0)

        return forecasts


# river/time_series/test_evaluate.py

from __future__ import annotations

from river import datasets, metrics, stats, time_series


class MeanForecaster(time_series.base.Forecaster):
    def __init__(self):
        self.mean = stats.Mean()

    def learn_one(self, y, x=None):
        self.mean.update(y)
        return self

    def forecast(self, horizon, xs=None):
        return [self.mean.get()] * horizon


def test_forecasts_at_each_step():
    dataset = datasets.AirlinePassengers()
    model = MeanForecaster()
    metric = metrics.MAE()
    horizon = 12
    grace_period = 1

    steps = time_series.iter_evaluate(
        dataset=dataset, model=model, metric=metric, horizon=horizon, grace_period=grace_period
    )

    _, _, y_pred, _ = next(steps)
    assert y_pred == [112] * horizon
    _, _, y_pred, _ = next(steps)
    assert y_pred == [(112 + 118) / 2] * horizon
    _, _, y_pred, _ = next(steps)
    assert y_pred == [(112 + 118 + 132) / 3] * horizon
    _, _, y_pred, _ = next(steps)
    assert y_pred == [(112 + 118 + 132 + 129) / 4] * horizon

    n_steps = sum(
        1
        for _ in time_series.iter_evaluate(
            dataset=dataset, model=model, metric=metric, horizon=horizon, grace_period=grace_period
        )
    )
    assert n_steps == dataset.n_samples - horizon - grace_period


# river/time_series/test_snarimax.py

from __future__ import annotations

import calendar
import math
import random

import pytest
import sympy

from river import compose, datasets, metrics, time_series
from river.time_series.snarimax import Differencer


class Yt(sympy.IndexedBase):
    t = sympy.symbols("t", cls=sympy.Idx)

    def __getitem__(self, idx):
        return super().__getitem__(self.t - idx)


def test_diff_formula():
    """

    >>> import sympy
    >>> from river.time_series.snarimax import Differencer

    >>> Y = Yt('y')
    >>> Y
    y

    >>> p = sympy.symbols('p')
    >>> p
    p

    >>> D = Differencer

    1
    >>> D(0).diff(p, Y)
    p

    (1 - B)
    >>> D(1).diff(p, Y)
    p - y[t]

    (1 - B)^2
    >>> D(2).diff(p, Y)
    p + y[t - 1] - 2*y[t]

    (1 - B^m)
    >>> m = sympy.symbols('m', cls=sympy.Idx)
    >>> D(1, m).diff(p, Y)
    p - y[-m + t + 1]

    (1 - B)(1 - B^m)
    >>> (D(1) * D(1, m)).diff(p, Y)
    p - y[-m + t + 1] + y[-m + t] - y[t]

    (1 - B)(1 - B^12)
    >>> (D(1) * D(1, 12)).diff(p, Y)
    p - y[t - 11] + y[t - 12] - y[t]

    """


def test_diff_example():
    """https://people.duke.edu/~rnau/411sdif.htm

    >>> import pandas as pd
    >>> from river.time_series.snarimax import Differencer

    >>> sales = pd.DataFrame([
    ...     {'date': 'Jan-70', 'autosale': 4.79, 'cpi': 0.297},
    ...     {'date': 'Feb-70', 'autosale': 4.96, 'cpi': 0.298},
    ...     {'date': 'Mar-70', 'autosale': 5.64, 'cpi': 0.300},
    ...     {'date': 'Apr-70', 'autosale': 5.98, 'cpi': 0.302},
    ...     {'date': 'May-70', 'autosale': 6.08, 'cpi': 0.303},
    ...     {'date': 'Jun-70', 'autosale': 6.55, 'cpi': 0.305},
    ...     {'date': 'Jul-70', 'autosale': 6.11, 'cpi': 0.306},
    ...     {'date': 'Aug-70', 'autosale': 5.37, 'cpi': 0.306},
    ...     {'date': 'Sep-70', 'autosale': 5.17, 'cpi': 0.308},
    ...     {'date': 'Oct-70', 'autosale': 5.48, 'cpi': 0.309},
    ...     {'date': 'Nov-70', 'autosale': 4.49, 'cpi': 0.311},
    ...     {'date': 'Dec-70', 'autosale': 4.65, 'cpi': 0.312},
    ...     {'date': 'Jan-71', 'autosale': 5.17, 'cpi': 0.312},
    ...     {'date': 'Feb-71', 'autosale': 5.57, 'cpi': 0.313},
    ...     {'date': 'Mar-71', 'autosale': 6.92, 'cpi': 0.314},
    ...     {'date': 'Apr-71', 'autosale': 7.10, 'cpi': 0.315},
    ...     {'date': 'May-71', 'autosale': 7.02, 'cpi': 0.316},
    ...     {'date': 'Jun-71', 'autosale': 7.58, 'cpi': 0.319},
    ...     {'date': 'Jul-71', 'autosale': 6.93, 'cpi': 0.319},
    ... ])

    >>> sales['autosale/cpi'] = sales.eval('autosale / cpi').round(2)
    >>> Y = sales['autosale/cpi'].to_list()

    >>> diff = Differencer(1)
    >>> sales['(1 - B)'] = [
    ...     diff.diff(p, Y[:i][::-1])
    ...     if i else ''
    ...     for i, p in enumerate(Y)
    ... ]

    >>> sdiff = Differencer(1, 12)
    >>> sales['(1 - B^12)'] = [
    ...     sdiff.diff(p, Y[:i][::-1])
    ...     if i >= 12 else ''
    ...     for i, p in enumerate(Y)
    ... ]

    >>> sales['(1 - B)(1 - B^12)'] = [
    ...     (diff * sdiff).diff(p, Y[:i][::-1])
    ...     if i >= 13 else ''
    ...     for i, p in enumerate(Y)
    ... ]

    >>> sales
          date  autosale    cpi  autosale/cpi (1 - B) (1 - B^12) (1 - B)(1 - B^12)
    0   Jan-70      4.79  0.297         16.13
    1   Feb-70      4.96  0.298         16.64    0.51
    2   Mar-70      5.64  0.300         18.80    2.16
    3   Apr-70      5.98  0.302         19.80     1.0
    4   May-70      6.08  0.303         20.07    0.27
    5   Jun-70      6.55  0.305         21.48    1.41
    6   Jul-70      6.11  0.306         19.97   -1.51
    7   Aug-70      5.37  0.306         17.55   -2.42
    8   Sep-70      5.17  0.308         16.79   -0.76
    9   Oct-70      5.48  0.309         17.73    0.94
    10  Nov-70      4.49  0.311         14.44   -3.29
    11  Dec-70      4.65  0.312         14.90    0.46
    12  Jan-71      5.17  0.312         16.57    1.67       0.44
    13  Feb-71      5.57  0.313         17.80    1.23       1.16              0.72
    14  Mar-71      6.92  0.314         22.04    4.24       3.24              2.08
    15  Apr-71      7.10  0.315         22.54     0.5       2.74              -0.5
    16  May-71      7.02  0.316         22.22   -0.32       2.15             -0.59
    17  Jun-71      7.58  0.319         23.76    1.54       2.28              0.13
    18  Jul-71      6.93  0.319         21.72   -2.04       1.75             -0.53

    """


@pytest.mark.parametrize(
    "differencer",
    [
        Differencer(1),
        Differencer(2),
        Differencer(1, 2),
        Differencer(2, 2),
        Differencer(1, 10),
        Differencer(2, 10),
        Differencer(1) * Differencer(1),
        Differencer(2) * Differencer(1),
        Differencer(1) * Differencer(2),
        Differencer(1) * Differencer(1, 2),
        Differencer(2) * Differencer(1, 2),
        Differencer(1, 2) * Differencer(1, 10),
        Differencer(1, 2) * Differencer(2, 10),
        Differencer(2, 2) * Differencer(1, 10),
        Differencer(2, 2) * Differencer(2, 10),
    ],
)
def test_undiff(differencer):
    Y = [random.random() for _ in range(max(differencer.coeffs))]
    p = random.random()

    diffed = differencer.diff(p, Y)
    undiffed = differencer.undiff(diffed, Y)
    assert math.isclose(undiffed, p)


@pytest.mark.parametrize(
    "snarimax, Y, errors, expected",
    [
        # Non-seasonal parts (p and q)
        (
            time_series.SNARIMAX(p=3, d=0, q=3),
            [1, 2, 3],
            [-4, -5, -6],
            {"e-1": -4, "e-2": -5, "e-3": -6, "y-1": 1, "y-2": 2, "y-3": 3},
        ),
        (
            time_series.SNARIMAX(p=2, d=0, q=3),
            [1, 2, 3],
            [-4, -5, -6],
            {"e-1": -4, "e-2": -5, "e-3": -6, "y-1": 1, "y-2": 2},
        ),
        (
            time_series.SNARIMAX(p=3, d=0, q=2),
            [1, 2, 3],
            [-4, -5, -6],
            {"e-1": -4, "e-2": -5, "y-1": 1, "y-2": 2, "y-3": 3},
        ),
        (
            time_series.SNARIMAX(p=2, d=0, q=2),
            [1, 2, 3],
            [-4, -5, -6],
            {"e-1": -4, "e-2": -5, "y-1": 1, "y-2": 2},
        ),
        # Not enough data, so features too far away are omitted
        (
            time_series.SNARIMAX(p=3, d=0, q=3),
            [1, 2],
            [-4, -5],
            {"e-1": -4, "e-2": -5, "y-1": 1, "y-2": 2},
        ),
        # Seasonal AR
        (
            time_series.SNARIMAX(p=2, d=0, q=2, m=3, sp=2),
            [i for i in range(12)],
            [i for i in range(12)],
            {"e-1": 0, "e-2": 1, "sy-3": 2, "sy-6": 5, "y-1": 0, "y-2": 1},
        ),
        (
            time_series.SNARIMAX(p=2, d=0, q=2, m=2, sp=2),
            [i for i in range(12)],
            [i for i in range(12)],
            {"e-1": 0, "e-2": 1, "sy-2": 1, "sy-4": 3, "y-1": 0, "y-2": 1},
        ),
        (
            time_series.SNARIMAX(p=2, d=0, q=2, m=2, sp=3),
            [i for i in range(12)],
            [i for i in range(12)],
            {"e-1": 0, "e-2": 1, "sy-2": 1, "sy-4": 3, "sy-6": 5, "y-1": 0, "y-2": 1},
        ),
        # Seasonal MA
        (
            time_series.SNARIMAX(p=2, d=0, q=2, m=3, sq=2),
            [i for i in range(12)],
            [i for i in range(12)],
            {"e-1": 0, "e-2": 1, "se-3": 2, "se-6": 5, "y-1": 0, "y-2": 1},
        ),
        (
            time_series.SNARIMAX(p=2, d=0, q=2, m=3, sq=4),
            [i for i in range(12)],
            [i for i in range(12)],
            {"e-1": 0, "e-2": 1, "se-3": 2, "se-6": 5, "se-9": 8, "se-12": 11, "y-1": 0, "y-2": 1},
        ),
        (
            time_series.SNARIMAX(p=1, d=0, q=1, m=2, sq=4),
            [i for i in range(12)],
            [i for i in range(12)],
            {"e-1": 0, "se-2": 1, "se-4": 3, "se-6": 5, "se-8": 7, "y-1": 0},
        ),
    ],
)
def test_add_lag_features(snarimax, Y, errors, expected):
    features = snarimax._add_lag_features(x=None, Y=Y, errors=errors)
    assert features == expected


@pytest.mark.parametrize(
    "snarimax",
    [
        time_series.SNARIMAX(p=1, d=1, q=0, m=12, sp=0, sd=1, sq=0),
        time_series.SNARIMAX(p=0, d=1, q=0, m=12, sp=1, sd=1, sq=0),
        time_series.SNARIMAX(p=1, d=2, q=0, m=12, sp=0, sd=0, sq=0),
        time_series.SNARIMAX(p=1, d=0, q=0, m=12, sp=0, sd=2, sq=0),
    ],
)
def test_no_overflow(snarimax):
    def get_month_distances(x):
        return {
            calendar.month_name[month]: math.exp(-((x["month"].month - month) ** 2))
            for month in range(1, 13)
        }

    def get_ordinal_date(x):
        return {"ordinal_date": x["month"].toordinal()}

    extract_features = compose.TransformerUnion(get_ordinal_date, get_month_distances)

    model = extract_features | snarimax

    time_series.evaluate(
        dataset=datasets.AirlinePassengers(), model=model, metric=metrics.MAE(), horizon=12
    )


# river/tree/__init__.py

"""

This module implements incremental Decision Tree (iDT) algorithms for handling classification
and regression tasks.

Each family of iDT will be presented in a dedicated section.

At any moment, iDT might face situations where an input feature previously used to make
a split decision is missing in an incoming sample. In this case, the most traversed path is
selected to pass down the instance. Moreover, in the case of nominal features, if a new category
arises and the feature is used in a decision node, a new branch is created to accommodate the new
value.

**1. Hoeffding Trees**

This family of iDT algorithms use the Hoeffding Bound to determine whether or not the
incrementally computed best split candidates would be equivalent to the ones obtained in a
batch-processing fashion.

All the available Hoeffding Tree (HT) implementation share some common functionalities:

* Set the maximum tree depth allowed (`max_depth`).

* Handle *Active* and *Inactive* nodes: Active learning nodes update their own
internal state to improve predictions and monitor input features to perform split
attempts. Inactive learning nodes do not update their internal state and only keep the
predictors; they are used to save memory in the tree (`max_size`).

*  Enable/disable memory management.

* Define strategies to sort leaves according to how likely they are going to be split.
This enables deactivating non-promising leaves to save memory.

* Disabling ‘poor’ attributes to save memory and speed up tree construction.
A poor attribute is an input feature whose split merit is much smaller than the current
best candidate. Once a feature is disabled, the tree stops saving statistics necessary
to split such a feature.

* Define properties to access leaf prediction strategies, split criteria, and other
relevant characteristics.

**2. Stochastic Gradient Trees**

Stochastic Gradient Trees (SGT) directly optimize a loss function, rather than relying on split
heuristics to guide the tree growth. F-tests are performed do decide whether a leaf should be
expanded or its prediction value should be updated.

SGTs can deal with binary classification and single-target regression. They also support
dynamic and static feature quantizers to deal with numerical inputs.

"""
from __future__ import annotations

from . import base, splitter
from .extremely_fast_decision_tree import ExtremelyFastDecisionTreeClassifier
from .hoeffding_adaptive_tree_classifier import HoeffdingAdaptiveTreeClassifier
from .hoeffding_adaptive_tree_regressor import HoeffdingAdaptiveTreeRegressor
from .hoeffding_tree_classifier import HoeffdingTreeClassifier
from .hoeffding_tree_regressor import HoeffdingTreeRegressor
from .isoup_tree_regressor import iSOUPTreeRegressor
from .stochastic_gradient_tree import SGTClassifier, SGTRegressor

__all__ = [
    "base",
    "splitter",
    "ExtremelyFastDecisionTreeClassifier",
    "HoeffdingTreeClassifier",
    "HoeffdingAdaptiveTreeClassifier",
    "HoeffdingTreeRegressor",
    "HoeffdingAdaptiveTreeRegressor",
    "iSOUPTreeRegressor",
    "SGTClassifier",
    "SGTRegressor",
]


# river/tree/base.py

"""

This module defines generic branch and leaf implementations. These should be used in River by each
tree-based model. Using these classes makes the code more DRY. The only exception for not doing so
would be for performance, whereby a tree-based model uses a bespoke implementation.

This module defines a bunch of methods to ease the manipulation and diagnostic of trees. Its
intention is to provide utilities for walking over a tree and visualizing it.

"""
from __future__ import annotations

import abc
import typing
from collections import defaultdict
from queue import Queue
from xml.etree import ElementTree as ET

import pandas as pd

from river import base

__all__ = ["Branch", "Leaf"]


class Branch(base.Base, abc.ABC):
    """A generic tree branch.

    Parameters
    ----------
    children
        Child branches and/or leaves.

    """

    def __init__(self, *children):
        self.children = children

    @abc.abstractmethod
    def next(self, x) -> Branch | Leaf:
        """Move to the next node down the tree."""

    @abc.abstractmethod
    def most_common_path(self) -> tuple[int, Leaf | Branch]:
        """Return a tuple with the branch index and the child node related to the most
        traversed path.

        Used in case the split feature is missing from an instance.
        """
        pass

    @property
    @abc.abstractmethod
    def repr_split(self):
        """String representation of the split."""

    def walk(self, x, until_leaf=True) -> typing.Iterable[Branch | Leaf]:
        """Iterate over the nodes of the path induced by x."""
        yield self
        try:
            yield from self.next(x).walk(x, until_leaf)
        except KeyError:
            if until_leaf:
                _, node = self.most_common_path()
                yield node
                yield from node.walk(x, until_leaf)

    def traverse(self, x, until_leaf=True) -> Branch | Leaf:
        """Return the leaf corresponding to the given input."""
        for node in self.walk(x, until_leaf):
            pass
        return node

    @property
    def n_nodes(self):
        """Number of descendants, including thyself."""
        return 1 + sum(child.n_nodes for child in self.children)

    @property
    def n_branches(self):
        """Number of branches, including thyself."""
        return 1 + sum(child.n_branches for child in self.children)

    @property
    def n_leaves(self):
        """Number of leaves."""
        return sum(child.n_leaves for child in self.children)

    @property
    def height(self):
        """Distance to the deepest descendant."""
        return 1 + max(child.height for child in self.children)

    def iter_dfs(self):
        """Iterate over nodes in depth-first order."""
        yield self
        for child in self.children:
            yield from child.iter_dfs()

    def iter_bfs(self):
        """Iterate over nodes in breadth-first order."""

        queue = Queue()

        queue.put(self)

        while not queue.empty():
            node = queue.get()
            yield node
            if isinstance(node, Branch):
                for child in node.children:
                    queue.put(child)

    def iter_leaves(self):
        """Iterate over leaves from the left-most one to the right-most one."""
        for child in self.children:
            yield from child.iter_leaves()

    def iter_branches(self):
        """Iterate over branches in depth-first order."""
        yield self
        for child in self.children:
            yield from child.iter_branches()

    def iter_edges(self):
        """Iterate over edges in depth-first order."""
        for child in self.children:
            yield self, child
            yield from child.iter_edges()

    def to_dataframe(self) -> pd.DataFrame:
        """Build a DataFrame containing one record for each node."""
        node_ids: defaultdict[typing.Hashable, int] = defaultdict(lambda: len(node_ids))  # type: ignore
        nodes = []

        queue: Queue = Queue()
        queue.put((self, None, 0))

        while not queue.empty():
            node, parent, depth = queue.get()
            nodes.append(
                {
                    "node": node_ids[id(node)],
                    "parent": node_ids[id(parent)] if parent else pd.NA,
                    "is_leaf": isinstance(node, Leaf),
                    "depth": depth,
                    **{k: v for k, v in node.__dict__.items() if k != "children"},
                }
            )
            try:
                for child in node.children:
                    queue.put((child, node, depth + 1))
            except AttributeError:
                pass

        return pd.DataFrame.from_records(nodes).set_index("node")

    def _repr_html_(self):
        from river.tree import viz

        div = viz.tree_to_html(self)
        return f"<div>{ET.tostring(div, encoding='unicode')}<style scoped>{viz.CSS}</style></div>"


class Leaf(base.Base):
    """A generic tree node.

    Parameters
    ----------
    kwargs
        Each provided keyword argument is stored in the leaf as an attribute.

    """

    def __init__(self, **kwargs):
        self.__dict__.update(kwargs)

    def walk(self, x, until_leaf=True):  # noqa
        yield self

    @property
    @abc.abstractmethod
    def __repr__(self):
        """String representation for visualization purposes."""

    @property
    def n_nodes(self):
        return 1

    @property
    def n_branches(self):
        return 0

    @property
    def n_leaves(self):
        return 1

    @property
    def height(self):
        return 1

    def iter_dfs(self):
        yield self

    def iter_leaves(self):
        yield self

    def iter_branches(self):  # noqa
        yield from ()

    def iter_edges(self):  # noqa
        yield from ()


# river/tree/extremely_fast_decision_tree.py

from __future__ import annotations

from .hoeffding_tree_classifier import HoeffdingTreeClassifier
from .nodes.branch import DTBranch
from .nodes.efdtc_nodes import (
    BaseEFDTBranch,
    EFDTLeafMajorityClass,
    EFDTLeafNaiveBayes,
    EFDTLeafNaiveBayesAdaptive,
    EFDTNominalBinaryBranch,
    EFDTNominalMultiwayBranch,
    EFDTNumericBinaryBranch,
    EFDTNumericMultiwayBranch,
)
from .nodes.leaf import HTLeaf
from .splitter import Splitter
from .utils import BranchFactory


class ExtremelyFastDecisionTreeClassifier(HoeffdingTreeClassifier):
    """Extremely Fast Decision Tree classifier.

    Also referred to as Hoeffding AnyTime Tree (HATT) classifier.

    Parameters
    ----------
    grace_period
        Number of instances a leaf should observe between split attempts.
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    min_samples_reevaluate
        Number of instances a node should observe before reevaluating the best split.
    split_criterion
        Split criterion to use.</br>
        - 'gini' - Gini</br>
        - 'info_gain' - Information Gain</br>
        - 'hellinger' - Helinger Distance</br>
    delta
        Significance level to calculate the Hoeffding bound. The significance level is given by
        `1 - delta`. Values closer to zero imply longer split decision delays.
    tau
        Threshold below which a split will be forced to break ties.
    leaf_prediction
        Prediction mechanism used at leafs.</br>
        - 'mc' - Majority Class</br>
        - 'nb' - Naive Bayes</br>
        - 'nba' - Naive Bayes Adaptive</br>
    nb_threshold
        Number of instances a leaf should observe before allowing Naive Bayes.
    nominal_attributes
        List of Nominal attributes identifiers. If empty, then assume that all numeric attributes
        should be treated as continuous.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.GaussianSplitter` is used if `splitter` is `None`.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.

    Notes
    -----
    The Extremely Fast Decision Tree (EFDT) [^1] constructs a tree incrementally. The EFDT seeks to
    select and deploy a split as soon as it is confident the split is useful, and then revisits
    that decision, replacing the split if it subsequently becomes evident that a better split is
    available. The EFDT learns rapidly from a stationary distribution and eventually it learns the
    asymptotic batch tree if the distribution from which the data are drawn is stationary.

    References
    ----------
    [^1]:  C. Manapragada, G. Webb, and M. Salehi. Extremely Fast Decision Tree.
    In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data
    Mining (KDD '18). ACM, New York, NY, USA, 1953-1962.
    DOI: https://doi.org/10.1145/3219819.3220005

    Examples
    --------
    >>> from river.datasets import synth
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree

    >>> gen = synth.Agrawal(classification_function=0, seed=42)
    >>> # Take 1000 instances from the infinite data generator
    >>> dataset = iter(gen.take(1000))

    >>> model = tree.ExtremelyFastDecisionTreeClassifier(
    ...     grace_period=100,
    ...     delta=1e-5,
    ...     nominal_attributes=['elevel', 'car', 'zipcode'],
    ...     min_samples_reevaluate=100
    ... )

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 87.29%
    """

    def __init__(
        self,
        grace_period: int = 200,
        max_depth: int | None = None,
        min_samples_reevaluate: int = 20,
        split_criterion: str = "info_gain",
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "nba",
        nb_threshold: int = 0,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
    ):
        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            split_criterion=split_criterion,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            nb_threshold=nb_threshold,
            nominal_attributes=nominal_attributes,
            splitter=splitter,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self.min_samples_reevaluate = min_samples_reevaluate

    @property
    def _mutable_attributes(self):
        return {"grace_period", "delta", "tau", "min_samples_reevaluate"}

    def _new_leaf(self, initial_stats=None, parent=None):
        if initial_stats is None:
            initial_stats = {}
        if parent is None:
            depth = 0
        else:
            depth = parent.depth + 1

        if self._leaf_prediction == self._MAJORITY_CLASS:
            return EFDTLeafMajorityClass(initial_stats, depth, self.splitter)
        elif self._leaf_prediction == self._NAIVE_BAYES:
            return EFDTLeafNaiveBayes(initial_stats, depth, self.splitter)
        else:  # NAIVE BAYES ADAPTIVE (default)
            return EFDTLeafNaiveBayesAdaptive(initial_stats, depth, self.splitter)

    def _branch_selector(self, numerical_feature=True, multiway_split=False) -> type[DTBranch]:
        """Create a new split node."""
        if numerical_feature:
            if not multiway_split:
                return EFDTNumericBinaryBranch
            else:
                return EFDTNumericMultiwayBranch
        else:
            if not multiway_split:
                return EFDTNominalBinaryBranch
            else:
                return EFDTNominalMultiwayBranch

    def learn_one(self, x, y, *, sample_weight=1.0):
        """Incrementally train the model

        Parameters
        ----------
        x
            Instance attributes.
        y
            The label of the instance.
        sample_weight
            The weight of the sample.

        Notes
        -----
        Training tasks:

        * If the tree is empty, create a leaf node as the root.
        * If the tree is already initialized, find the path from root to the corresponding leaf for
        the instance and sort the instance.
        * Reevaluate the best split for each internal node.
        * Attempt to split the leaf.

        Returns
        -------
        self

        """
        # Updates the set of observed classes
        self.classes.add(y)

        self._train_weight_seen_by_model += sample_weight

        if self._root is None:
            self._root = self._new_leaf()
            self._n_active_leaves = 1

        # Sort instance X into a leaf
        self._sort_to_leaf(x, y, sample_weight)
        # Process all nodes, starting from root to the leaf where the instance x belongs.
        self._process_nodes(x, y, sample_weight, self._root, None, None)

        return self

    def _sort_to_leaf(self, x, y, sample_weight):
        """For a given instance, find the corresponding leaf and update it.

        Private function where leaf learn from instance.

        1. Find the node where instance should be.
        2. If no node have been found, create new learning node.
        3.1 Update the node with the provided instance.

        Parameters
        ----------
        x
            Instance attributes.
        y
            The instance label.
        sample_weight
            The weight of the sample.

        """
        node = self._root
        if isinstance(self._root, DTBranch):
            node = self._root.traverse(x, until_leaf=False)
            # Something went wrong in the way: a missing split feature or emerging category
            # Let's deal with these situations
            if isinstance(node, DTBranch):
                while True:
                    if node.max_branches() == -1 and node.feature in x:
                        # Emerging feature in nominal feature: create a new branch
                        leaf = self._new_leaf(parent=node)
                        node.add_child(x[node.feature], leaf)
                        self._n_active_leaves += 1
                        node = leaf
                    else:
                        # Missing split feature: select the most traversed path to continue the
                        # walk
                        _, node = node.most_common_path()
                        if isinstance(node, DTBranch):
                            node = node.traverse(x, until_leaf=False)
                    if isinstance(node, HTLeaf):
                        break
        node.learn_one(x, y, sample_weight=sample_weight, tree=self)

        if self._train_weight_seen_by_model % self.memory_estimate_period == 0:
            self._estimate_model_size()

    def _process_nodes(self, x, y, sample_weight, node, parent, branch_index):
        """Process nodes from the root to the leaf where the instance belongs.

        1. If the node is internal:
            1.1 If the number of samples seen since the last reevaluation are greater than
            `min_samples_reevaluate`, reevaluate the best split for the internal node.
        2. If the node is leaf, attempt to split leaf node.

        Parameters
        ----------
        x
            Instance attributes.
        y
            The label of the instance.
        sample_weight
            The weight of the sample.
        node
            The node to process.
        parent
            The node's parent.
        branch_index
            Parent node's branch index.
        """
        if isinstance(node, BaseEFDTBranch):
            # Update split nodes as the tree is traversed
            node.learn_one(x, y, sample_weight=sample_weight, tree=self)

            old_weight = node.last_split_reevaluation_at
            new_weight = node.total_weight
            stop_flag = False

            if (new_weight - old_weight) >= self.min_samples_reevaluate:
                # Reevaluate the best split
                stop_flag = self._reevaluate_best_split(
                    node,
                    parent,
                    branch_index,
                    # The attribute observer template that will be replicated to new leaves
                    splitter=self.splitter,
                    # Existing attribute observers that will be leveraged
                    splitters=node.splitters,
                )

            if not stop_flag:
                # Move in depth
                try:
                    child_index = node.branch_no(x)
                    child = node.children[child_index]
                except KeyError:
                    child_index, child = node.most_common_path()
                self._process_nodes(x, y, sample_weight, child, node, child_index)
        elif self._growth_allowed and node.is_active():
            if node.depth >= self.max_depth:  # Max depth reached
                node.deactivate()
                self._n_inactive_leaves += 1
                self._n_active_leaves -= 1
            else:
                weight_seen = node.total_weight
                weight_diff = weight_seen - node.last_split_attempt_at
                if weight_diff >= self.grace_period:
                    # Attempt to split
                    self._attempt_to_split(
                        node,
                        parent,
                        branch_index,
                        splitter=self.splitter,
                        splitters=node.splitters,
                    )
                    node.last_split_attempt_at = weight_seen

    def _reevaluate_best_split(self, node, parent, branch_index, **kwargs):
        """Reevaluate the best split for a node.

        If the samples seen so far are not from the same class then:
        1. Find split candidates and select the best one.
        2. Compute the Hoeffding bound.
        3. If the null split candidate is higher than the top split candidate:
            3.1 Kill subtree and replace it with a leaf.
            3.2 Update the tree.
            3.3 Update tree's metrics
        4. If the difference between the top split candidate and the current split is larger than
        the Hoeffding bound:
           4.1 Create a new split node.
           4.2 Update the tree.
           4.3 Update tree's metrics
        5. If the top split candidate is the current split but with different split test:
           5.1 Update the split test of the current split.

        Parameters
        ----------
        node
            The node to reevaluate.
        parent
            The node's parent.
        branch_index
            Parent node's branch index.
        kwargs
            Other parameters passed to the branch node.

        Returns
        -------
            Flag to stop moving in depth.
        """
        stop_flag = False
        if not node.observed_class_distribution_is_pure():
            split_criterion = self._new_split_criterion()
            best_split_suggestions = node.best_split_suggestions(split_criterion, self)
            if len(best_split_suggestions) > 0:
                # Sort the attribute accordingly to their split merit for each attribute
                # (except the null one)
                best_split_suggestions.sort()

                # x_best is the attribute with the highest merit
                x_best = best_split_suggestions[-1]
                id_best = x_best.feature

                # Best split candidate is the null split
                if x_best.feature is None:
                    return True

                # x_current is the current attribute used in this SplitNode
                id_current = node.feature
                x_current = node.find_attribute(id_current, best_split_suggestions)

                # Get x_null
                x_null = BranchFactory(merit=0)

                # Compute Hoeffding bound
                hoeffding_bound = self._hoeffding_bound(
                    split_criterion.range_of_merit(node.stats),
                    self.delta,
                    node.total_weight,
                )

                if x_null.merit - x_best.merit > hoeffding_bound:
                    # Kill subtree & replace the branch by a leaf
                    best_split = self._kill_subtree(node)

                    # update EFDT
                    if parent is None:
                        # Root case : replace the root node by a new split node
                        self._root = best_split
                    else:
                        parent.children[branch_index] = best_split

                    n_active = n_inactive = 0
                    for leaf in node.iter_leaves():
                        if leaf.is_active():
                            n_active += 1
                        else:
                            n_inactive += 1

                    self._n_active_leaves += 1
                    self._n_active_leaves -= n_active
                    self._n_inactive_leaves -= n_inactive
                    stop_flag = True

                    # Manage memory
                    self._enforce_size_limit()

                elif (
                    x_best.merit - x_current.merit > hoeffding_bound or hoeffding_bound < self.tau
                ) and (id_current != id_best):
                    # Create a new branch
                    branch = self._branch_selector(x_best.numerical_feature, x_best.multiway_split)
                    leaves = tuple(
                        self._new_leaf(initial_stats, parent=node)
                        for initial_stats in x_best.children_stats
                    )

                    new_split = x_best.assemble(branch, node.stats, node.depth, *leaves, **kwargs)
                    # Update weights in new_split
                    new_split.last_split_reevaluation_at = node.total_weight

                    n_active = n_inactive = 0
                    for leaf in node.iter_leaves():
                        if leaf.is_active():
                            n_active += 1
                        else:
                            n_inactive += 1

                    self._n_active_leaves -= n_active
                    self._n_inactive_leaves -= n_inactive
                    self._n_active_leaves += len(leaves)

                    if parent is None:
                        # Root case : replace the root node by a new split node
                        self._root = new_split
                    else:
                        parent.children[branch_index] = new_split

                    stop_flag = True

                    # Manage memory
                    self._enforce_size_limit()

                elif (
                    x_best.merit - x_current.merit > hoeffding_bound or hoeffding_bound < self.tau
                ) and (id_current == id_best):
                    branch = self._branch_selector(x_best.numerical_feature, x_best.multiway_split)
                    # Change the branch but keep the existing children nodes
                    new_split = x_best.assemble(
                        branch, node.stats, node.depth, *tuple(node.children), **kwargs
                    )
                    # Update weights in new_split
                    new_split.last_split_reevaluation_at = node.total_weight

                    if parent is None:
                        # Root case : replace the root node by a new split node
                        self._root = new_split
                    else:
                        parent.children[branch_index] = new_split

        return stop_flag

    def _attempt_to_split(self, node, parent, branch_index, **kwargs):
        """Attempt to split a node.

        If the samples seen so far are not from the same class then:

        1. Find split candidates and select the best one.
        2. Compute the Hoeffding bound.
        3. If the difference between the best split candidate and the don't split candidate is
        larger than the Hoeffding bound:
            3.1 Replace the leaf node by a split node.
            3.2 Add a new leaf node on each branch of the new split node.
            3.3 Update tree's metrics

        Parameters
        ----------
        node
            The node to reevaluate.
        parent
            The node's parent.
        branch_index
            Parent node's branch index.
        kwargs
            Other parameters passed to the new branch node.

        """
        if not node.observed_class_distribution_is_pure():  # noqa
            split_criterion = self._new_split_criterion()

            best_split_suggestions = node.best_split_suggestions(split_criterion, self)

            if len(best_split_suggestions) > 0:
                # x_best is the attribute with the highest merit
                best_split_suggestions.sort()
                x_best = best_split_suggestions[-1]

                # Get x_null
                x_null = BranchFactory(merit=0)

                hoeffding_bound = self._hoeffding_bound(
                    split_criterion.range_of_merit(node.stats),
                    self.delta,
                    node.total_weight,
                )

                if x_best.merit - x_null.merit > hoeffding_bound or hoeffding_bound < self.tau:
                    # Create a new branch
                    branch = self._branch_selector(x_best.numerical_feature, x_best.multiway_split)
                    leaves = tuple(
                        self._new_leaf(initial_stats, parent=node)
                        for initial_stats in x_best.children_stats
                    )
                    new_split = x_best.assemble(branch, node.stats, node.depth, *leaves, **kwargs)

                    new_split.last_split_reevaluation_at = node.total_weight

                    self._n_active_leaves -= 1
                    self._n_active_leaves += len(leaves)

                    if parent is None:
                        # root case: replace the root node by a new split node
                        self._root = new_split
                    else:
                        parent.children[branch_index] = new_split

                    # Manage memory
                    self._enforce_size_limit()

    def _kill_subtree(self, node: BaseEFDTBranch):
        """Kill subtree that starts from node.

        Parameters
        ----------
        node
            The node to reevaluate.
        Returns
        -------
            The new leaf.
        """

        leaf = self._new_leaf()
        leaf.depth = node.depth  # type: ignore
        leaf.stats = node.stats
        leaf.splitters = node.splitters

        return leaf


# river/tree/hoeffding_adaptive_tree_classifier.py

from __future__ import annotations

import random

from river import base, drift
from river.utils.norm import normalize_values_in_dict

from .hoeffding_tree_classifier import HoeffdingTreeClassifier
from .nodes.branch import DTBranch
from .nodes.hatc_nodes import (
    AdaBranchClassifier,
    AdaLeafClassifier,
    AdaNomBinaryBranchClass,
    AdaNomMultiwayBranchClass,
    AdaNumBinaryBranchClass,
    AdaNumMultiwayBranchClass,
)
from .splitter import Splitter
from .utils import add_dict_values


class HoeffdingAdaptiveTreeClassifier(HoeffdingTreeClassifier):
    """Hoeffding Adaptive Tree classifier.

    Parameters
    ----------
    grace_period
        Number of instances a leaf should observe between split attempts.
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    split_criterion
        Split criterion to use.</br>
        - 'gini' - Gini</br>
        - 'info_gain' - Information Gain</br>
        - 'hellinger' - Helinger Distance</br>
    delta
        Significance level to calculate the Hoeffding bound. The significance level is given by
        `1 - delta`. Values closer to zero imply longer split decision delays.
    tau
        Threshold below which a split will be forced to break ties.
    leaf_prediction
        Prediction mechanism used at leafs.</br>
        - 'mc' - Majority Class</br>
        - 'nb' - Naive Bayes</br>
        - 'nba' - Naive Bayes Adaptive</br>
    nb_threshold
        Number of instances a leaf should observe before allowing Naive Bayes.
    nominal_attributes
        List of Nominal attributes. If empty, then assume that all numeric attributes should
        be treated as continuous.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.GaussianSplitter` is used if `splitter` is `None`.
    bootstrap_sampling
        If True, perform bootstrap sampling in the leaf nodes.
    drift_window_threshold
        Minimum number of examples an alternate tree must observe before being considered as a
        potential replacement to the current one.
    drift_detector
        The drift detector used to build the tree. If `None` then `drift.ADWIN` is used.
    switch_significance
        The significance level to assess whether alternate subtrees are significantly better
        than their main subtree counterparts.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.
    seed
       Random seed for reproducibility.


    Notes
    -----
    The Hoeffding Adaptive Tree [^1] uses a drift detector to monitor performance of branches in
    the tree and to replace them with new branches when their accuracy decreases.

    The bootstrap sampling strategy is an improvement over the original Hoeffding Adaptive Tree
    algorithm. It is enabled by default since, in general, it results in better performance.

    References
    ----------
    [^1]: Bifet, Albert, and Ricard Gavaldà. "Adaptive learning from evolving data streams."
       In International Symposium on Intelligent Data Analysis, pp. 249-260. Springer, Berlin,
       Heidelberg, 2009.

    Examples
    --------
    >>> from river.datasets import synth
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree

    >>> gen = synth.ConceptDriftStream(stream=synth.SEA(seed=42, variant=0),
    ...                                drift_stream=synth.SEA(seed=42, variant=1),
    ...                                seed=1, position=500, width=50)
    >>> # Take 1000 instances from the infinite data generator
    >>> dataset = iter(gen.take(1000))

    >>> model = tree.HoeffdingAdaptiveTreeClassifier(
    ...     grace_period=100,
    ...     delta=1e-5,
    ...     leaf_prediction='nb',
    ...     nb_threshold=10,
    ...     seed=0
    ... )

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 91.49%

    """

    def __init__(
        self,
        grace_period: int = 200,
        max_depth: int | None = None,
        split_criterion: str = "info_gain",
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "nba",
        nb_threshold: int = 0,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        bootstrap_sampling: bool = True,
        drift_window_threshold: int = 300,
        drift_detector: base.DriftDetector | None = None,
        switch_significance: float = 0.05,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
        seed: int | None = None,
    ):
        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            split_criterion=split_criterion,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            nb_threshold=nb_threshold,
            nominal_attributes=nominal_attributes,
            splitter=splitter,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self.bootstrap_sampling = bootstrap_sampling
        self.drift_window_threshold = drift_window_threshold
        self.drift_detector = drift_detector if drift_detector is not None else drift.ADWIN()
        self.switch_significance = switch_significance
        self.seed = seed

        self._n_alternate_trees = 0
        self._n_pruned_alternate_trees = 0
        self._n_switch_alternate_trees = 0

        self._rng = random.Random(self.seed)

    @property
    def _mutable_attributes(self):
        return {"grace_period", "delta", "tau", "drift_window_threshold", "switch_significance"}

    @property
    def n_alternate_trees(self):
        return self._n_alternate_trees

    @property
    def n_pruned_alternate_trees(self):
        return self._n_pruned_alternate_trees

    @property
    def n_switch_alternate_trees(self):
        return self._n_switch_alternate_trees

    @property
    def summary(self):
        summ = super().summary
        summ.update(
            {
                "n_alternate_trees": self.n_alternate_trees,
                "n_pruned_alternate_trees": self.n_pruned_alternate_trees,
                "n_switch_alternate_trees": self.n_switch_alternate_trees,
            }
        )
        return summ

    def learn_one(self, x, y, *, sample_weight=1.0):
        # Updates the set of observed classes
        self.classes.add(y)

        self._train_weight_seen_by_model += sample_weight

        if self._root is None:
            self._root = self._new_leaf()
            self._n_active_leaves = 1

        self._root.learn_one(x, y, sample_weight=sample_weight, tree=self)

        if self._train_weight_seen_by_model % self.memory_estimate_period == 0:
            self._estimate_model_size()

        return self

    # Override HoeffdingTreeClassifier
    def predict_proba_one(self, x):
        proba = {c: 0.0 for c in self.classes}
        if self._root is not None:
            found_nodes = [self._root]
            if isinstance(self._root, DTBranch):
                found_nodes = self._root.traverse(x, until_leaf=True)
            for leaf in found_nodes:
                dist = leaf.prediction(x, tree=self)
                # Option Tree prediction (of sorts): combine the response of all leaves reached
                # by the instance
                proba = add_dict_values(proba, dist, inplace=True)
            proba = normalize_values_in_dict(proba)

        return proba

    def _new_leaf(self, initial_stats=None, parent=None):
        if initial_stats is None:
            initial_stats = {}

        if parent is not None:
            depth = parent.depth + 1
        else:
            depth = 0

        return AdaLeafClassifier(
            stats=initial_stats,
            depth=depth,
            splitter=self.splitter,
            drift_detector=self.drift_detector.clone(),
            rng=self._rng,
        )

    def _branch_selector(
        self, numerical_feature=True, multiway_split=False
    ) -> type[AdaBranchClassifier]:
        """Create a new split node."""
        if numerical_feature:
            if not multiway_split:
                return AdaNumBinaryBranchClass
            else:
                return AdaNumMultiwayBranchClass
        else:
            if not multiway_split:
                return AdaNomBinaryBranchClass
            else:
                return AdaNomMultiwayBranchClass

    @classmethod
    def _unit_test_params(cls):
        yield {"seed": 1}


# river/tree/hoeffding_adaptive_tree_regressor.py

from __future__ import annotations

import random
import statistics
from copy import deepcopy

from river import base, drift

from .hoeffding_tree_regressor import HoeffdingTreeRegressor
from .nodes.branch import DTBranch
from .nodes.hatr_nodes import (
    AdaBranchRegressor,
    AdaLeafRegAdaptive,
    AdaLeafRegMean,
    AdaLeafRegModel,
    AdaLeafRegressor,
    AdaNomBinaryBranchReg,
    AdaNomMultiwayBranchReg,
    AdaNumBinaryBranchReg,
    AdaNumMultiwayBranchReg,
)
from .splitter import Splitter


class HoeffdingAdaptiveTreeRegressor(HoeffdingTreeRegressor):
    """Hoeffding Adaptive Tree regressor (HATR).

    This class implements a regression version of the Hoeffding Adaptive Tree Classifier. Hence,
    it also uses an ADWIN concept-drift detector instance at each decision node to monitor
    possible changes in the data distribution. If a drift is detected in a node, an alternate
    tree begins to be induced in the background. When enough information is gathered, HATR
    swaps the node where the change was detected by its alternate tree.

    Parameters
    ----------
    grace_period
        Number of instances a leaf should observe between split attempts.
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    delta
        Significance level to calculate the Hoeffding bound. The significance level is given by
        `1 - delta`. Values closer to zero imply longer split decision delays.
    tau
        Threshold below which a split will be forced to break ties.
    leaf_prediction
        Prediction mechanism used at leafs.</br>
        - 'mean' - Target mean</br>
        - 'model' - Uses the model defined in `leaf_model`</br>
        - 'adaptive' - Chooses between 'mean' and 'model' dynamically</br>
    leaf_model
        The regression model used to provide responses if `leaf_prediction='model'`. If not
        provided an instance of `river.linear_model.LinearRegression` with the default
        hyperparameters is used.
    model_selector_decay
        The exponential decaying factor applied to the learning models' squared errors, that
        are monitored if `leaf_prediction='adaptive'`. Must be between `0` and `1`. The closer
        to `1`, the more importance is going to be given to past observations. On the other hand,
        if its value approaches `0`, the recent observed errors are going to have more influence
        on the final decision.
    nominal_attributes
        List of Nominal attributes. If empty, then assume that all numeric attributes should
        be treated as continuous.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.TEBSTSplitter` is used if `splitter` is `None`.
    min_samples_split
        The minimum number of samples every branch resulting from a split candidate must have
        to be considered valid.
    bootstrap_sampling
        If True, perform bootstrap sampling in the leaf nodes.
    drift_window_threshold
        Minimum number of examples an alternate tree must observe before being considered as a
        potential replacement to the current one.
    drift_detector
        The drift detector used to build the tree. If `None` then `drift.ADWIN` is used.
        Only detectors that support arbitrarily valued continuous data can be used for regression.
    switch_significance
        The significance level to assess whether alternate subtrees are significantly better
        than their main subtree counterparts.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.
    seed
       Random seed for reproducibility.

    Notes
    -----
    The Hoeffding Adaptive Tree [^1] uses drift detectors to monitor performance of branches
    in the tree and to replace them with new branches when their accuracy decreases.

    The bootstrap sampling strategy is an improvement over the original Hoeffding Adaptive Tree
    algorithm. It is enabled by default since, in general, it results in better performance.

    To cope with ADWIN's requirements of bounded input data, HATR uses a novel error normalization
    strategy based on the empiral rule of Gaussian distributions. We assume the deviations
    of the predictions from the expected values follow a normal distribution. Hence, we subject
    these errors to a min-max normalization assuming that most of the data lies in the
    $\\left[-3\\sigma, 3\\sigma\\right]$ range. These normalized errors are passed to the ADWIN
    instances. This is the same strategy used by Adaptive Random Forest Regressor.

    References
    ----------
    [^1]: Bifet, Albert, and Ricard Gavaldà. "Adaptive learning from evolving data streams."
    In International Symposium on Intelligent Data Analysis, pp. 249-260. Springer, Berlin,
    Heidelberg, 2009.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     tree.HoeffdingAdaptiveTreeRegressor(
    ...         grace_period=50,
    ...         model_selector_decay=0.3,
    ...         seed=0
    ...     )
    ... )

    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 0.809874
    """

    def __init__(
        self,
        grace_period: int = 200,
        max_depth: int | None = None,
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "adaptive",
        leaf_model: base.Regressor | None = None,
        model_selector_decay: float = 0.95,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        min_samples_split: int = 5,
        bootstrap_sampling: bool = True,
        drift_window_threshold: int = 300,
        drift_detector: base.DriftDetector | None = None,
        switch_significance: float = 0.05,
        binary_split: bool = False,
        max_size: float = 500.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
        seed: int | None = None,
    ):
        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            leaf_model=leaf_model,
            model_selector_decay=model_selector_decay,
            nominal_attributes=nominal_attributes,
            splitter=splitter,
            min_samples_split=min_samples_split,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self.bootstrap_sampling = bootstrap_sampling
        self.drift_window_threshold = drift_window_threshold
        self.drift_detector = drift_detector if drift_detector is not None else drift.ADWIN()
        self.switch_significance = switch_significance
        self.seed = seed

        self._n_alternate_trees = 0
        self._n_pruned_alternate_trees = 0
        self._n_switch_alternate_trees = 0
        self._norm_dist = statistics.NormalDist()

        self._rng = random.Random(self.seed)

    @property
    def _mutable_attributes(self):
        return {"grace_period", "delta", "tau", "drift_window_threshold", "switch_significance"}

    @property
    def n_alternate_trees(self):
        return self._n_alternate_trees

    @property
    def n_pruned_alternate_trees(self):
        return self._n_pruned_alternate_trees

    @property
    def n_switch_alternate_trees(self):
        return self._n_switch_alternate_trees

    @property
    def summary(self):
        summ = super().summary
        summ.update(
            {
                "n_alternate_trees": self.n_alternate_trees,
                "n_pruned_alternate_trees": self.n_pruned_alternate_trees,
                "n_switch_alternate_trees": self.n_switch_alternate_trees,
            }
        )
        return summ

    def learn_one(self, x, y, *, sample_weight=1.0):
        self._train_weight_seen_by_model += sample_weight

        if self._root is None:
            self._root = self._new_leaf()
            self._n_active_leaves = 1
        self._root.learn_one(x, y, sample_weight=sample_weight, tree=self)

        if self._train_weight_seen_by_model % self.memory_estimate_period == 0:
            self._estimate_model_size()

        return self

    def predict_one(self, x):
        pred = 0.0
        if self._root is not None:
            found_nodes = [self._root]
            if isinstance(self._root, DTBranch):
                found_nodes = self._root.traverse(x, until_leaf=True)
            for leaf in found_nodes:
                pred += leaf.prediction(x, tree=self)
            # Mean prediction among the reached leaves
            pred /= len(found_nodes)

        return pred

    def _new_leaf(self, initial_stats=None, parent=None, is_active=True):
        """Create a new learning node.

        The type of learning node depends on the tree configuration.
        """
        if parent is not None:
            depth = parent.depth + 1
        else:
            depth = 0

        leaf_model = None
        if self.leaf_prediction in {self._MODEL, self._ADAPTIVE}:
            if parent is None:
                leaf_model = deepcopy(self.leaf_model)
            else:
                try:
                    leaf_model = deepcopy(parent._leaf_model)  # noqa
                except AttributeError:
                    leaf_model = deepcopy(self.leaf_model)

        if self.leaf_prediction == self._TARGET_MEAN:
            return AdaLeafRegMean(
                initial_stats,
                depth,
                self.splitter,
                drift_detector=self.drift_detector.clone(),
                rng=self._rng,
            )
        elif self.leaf_prediction == self._MODEL:
            return AdaLeafRegModel(
                initial_stats,
                depth,
                self.splitter,
                drift_detector=self.drift_detector.clone(),
                rng=self._rng,
                leaf_model=leaf_model,
            )
        else:  # adaptive learning node
            new_adaptive = AdaLeafRegAdaptive(
                initial_stats,
                depth,
                self.splitter,
                drift_detector=self.drift_detector.clone(),
                rng=self._rng,
                leaf_model=leaf_model,
            )
            if parent is not None and isinstance(parent, AdaLeafRegressor):
                new_adaptive._fmse_mean = parent._fmse_mean  # noqa
                new_adaptive._fmse_model = parent._fmse_model  # noqa

            return new_adaptive

    def _branch_selector(
        self, numerical_feature=True, multiway_split=False
    ) -> type[AdaBranchRegressor]:
        """Create a new split node."""
        if numerical_feature:
            if not multiway_split:
                return AdaNumBinaryBranchReg
            else:
                return AdaNumMultiwayBranchReg
        else:
            if not multiway_split:
                return AdaNomBinaryBranchReg
            else:
                return AdaNomMultiwayBranchReg

    @classmethod
    def _unit_test_params(cls):
        yield {"seed": 1}


# river/tree/hoeffding_tree.py

from __future__ import annotations

import collections
import functools
import io
import math
from abc import ABC, abstractmethod

from river import base
from river.utils.norm import normalize_values_in_dict

from .nodes.branch import (
    DTBranch,
    NominalBinaryBranch,
    NominalMultiwayBranch,
    NumericBinaryBranch,
    NumericMultiwayBranch,
)
from .nodes.leaf import HTLeaf
from .utils import calculate_object_size


class HoeffdingTree(ABC):
    """Base class for Hoeffding Decision Trees.

    This is an **abstract class**, so it cannot be used directly. It defines base operations
    and properties that all the Hoeffding decision trees must inherit or implement according to
    their own design.

    Parameters
    ----------
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.

    """

    def __init__(
        self,
        max_depth: int | None = None,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
    ):
        # Properties common to all the Hoeffding trees
        self._split_criterion: str = ""
        self._leaf_prediction: str = ""

        self.max_depth: float = max_depth if max_depth is not None else math.inf
        self.binary_split: bool = binary_split
        self._max_size: float = max_size
        self._max_byte_size: float = self._max_size * (2**20)  # convert to byte
        self.memory_estimate_period: int = memory_estimate_period
        self.stop_mem_management: bool = stop_mem_management
        self.remove_poor_attrs: bool = remove_poor_attrs
        self.merit_preprune: bool = merit_preprune

        self._root: DTBranch | HTLeaf = None  # type: ignore
        self._n_active_leaves: int = 0
        self._n_inactive_leaves: int = 0
        self._inactive_leaf_size_estimate: float = 0.0
        self._active_leaf_size_estimate: float = 0.0
        self._size_estimate_overhead_fraction: float = 1.0
        self._growth_allowed = True
        self._train_weight_seen_by_model: float = 0.0

    @staticmethod
    def _hoeffding_bound(range_val, confidence, n):
        r"""Compute the Hoeffding bound, used to decide how many samples are necessary at each
        node.

        Notes
        -----
        The Hoeffding bound is defined as:

        $\\epsilon = \\sqrt{\\frac{R^2\\ln(1/\\delta))}{2n}}$

        where:

        $\\epsilon$: Hoeffding bound.
        $R$: Range of a random variable. For a probability the range is 1, and for an
        information gain the range is log *c*, where *c* is the number of classes.
        $\\delta$: Confidence. 1 minus the desired probability of choosing the correct
        attribute at any given node.
        $n$: Number of samples.

        Parameters
        ----------
        range_val
            Range value.
        confidence
            Confidence of choosing the correct attribute.
        n
            Number of processed samples.
        """
        return math.sqrt((range_val * range_val * math.log(1.0 / confidence)) / (2.0 * n))

    @property
    def max_size(self):
        """Max allowed size tree can reach (in MB)."""
        return self._max_size

    @max_size.setter
    def max_size(self, size):
        self._max_size = size
        self._max_byte_size = self._max_size * (2**20)

    @property
    def height(self) -> int:
        if self._root:
            return self._root.height
        return 0

    @property
    def n_nodes(self):
        if self._root:
            return self._root.n_nodes

    @property
    def n_branches(self):
        if self._root:
            return self._root.n_branches

    @property
    def n_leaves(self):
        if self._root:
            return self._root.n_leaves

    @property
    def n_active_leaves(self):
        return self._n_active_leaves

    @property
    def n_inactive_leaves(self):
        return self._n_inactive_leaves

    @property
    def summary(self):
        """Collect metrics corresponding to the current status of the tree
        in a string buffer.
        """
        summary = {
            "n_nodes": self.n_nodes,
            "n_branches": self.n_branches,
            "n_leaves": self.n_leaves,
            "n_active_leaves": self.n_active_leaves,
            "n_inactive_leaves": self.n_inactive_leaves,
            "height": self.height,
            "total_observed_weight": self._train_weight_seen_by_model,
        }
        return summary

    def to_dataframe(self):
        """Return a representation of the current tree structure organized in a
        `pandas.DataFrame` object.

        In case the tree is empty or it only contains a single node (a leaf), `None` is returned.

        Returns
        -------
        df
            A `pandas.DataFrame` depicting the tree structure.
        """
        if self._root is not None and isinstance(self._root, DTBranch):
            return self._root.to_dataframe()

    def _branch_selector(self, numerical_feature=True, multiway_split=False) -> type[DTBranch]:
        """Create a new split node."""
        if numerical_feature:
            if not multiway_split:
                return NumericBinaryBranch
            else:
                return NumericMultiwayBranch
        else:
            if not multiway_split:
                return NominalBinaryBranch
            else:
                return NominalMultiwayBranch

    @abstractmethod
    def _new_leaf(
        self, initial_stats: dict | None = None, parent: HTLeaf | DTBranch | None = None
    ) -> HTLeaf:
        """Create a new learning node.

        The characteristics of the learning node depends on the tree algorithm.

        Parameters
        ----------
        initial_stats
            Target statistics set from the parent node.
        parent
            Parent node to inherit from.

        Returns
        -------
        A new learning node.
        """

    @property
    def split_criterion(self) -> str:
        """Return a string with the name of the split criterion being used by the tree."""
        return self._split_criterion

    @split_criterion.setter  # type: ignore
    @abstractmethod
    def split_criterion(self, split_criterion):
        """Define the split criterion to be used by the tree."""

    @property
    def leaf_prediction(self) -> str:
        """Return the prediction strategy used by the tree at its leaves."""
        return self._leaf_prediction

    @leaf_prediction.setter  # type: ignore
    @abstractmethod
    def leaf_prediction(self, leaf_prediction):
        """Define the prediction strategy used by the tree in its leaves."""

    def _enforce_size_limit(self):
        """Track the size of the tree and disable/enable nodes if required.

        This memory-management routine shared by all the Hoeffding Trees is based on [^1].

        References
        ----------
        [^1]: Kirkby, R.B., 2007. Improving hoeffding trees (Doctoral dissertation,
        The University of Waikato).
        """
        tree_size = self._size_estimate_overhead_fraction * (
            self._active_leaf_size_estimate
            + self._n_inactive_leaves * self._inactive_leaf_size_estimate
        )
        if self._n_inactive_leaves > 0 or tree_size > self._max_byte_size:
            if self.stop_mem_management:
                self._growth_allowed = False
                return
        leaves = self._find_leaves()
        leaves.sort(key=lambda leaf: leaf.calculate_promise())
        max_active = 0
        while max_active < len(leaves):
            max_active += 1
            if (
                (
                    max_active * self._active_leaf_size_estimate
                    + (len(leaves) - max_active) * self._inactive_leaf_size_estimate
                )
                * self._size_estimate_overhead_fraction
            ) > self._max_byte_size:
                max_active -= 1
                break
        cutoff = len(leaves) - max_active
        for i in range(cutoff):
            if leaves[i].is_active():
                leaves[i].deactivate()
                self._n_inactive_leaves += 1
                self._n_active_leaves -= 1
        for i in range(cutoff, len(leaves)):
            if not leaves[i].is_active() and leaves[i].depth < self.max_depth:
                leaves[i].activate()
                self._n_active_leaves += 1
                self._n_inactive_leaves -= 1

    def _estimate_model_size(self):
        """Calculate the size of the model and trigger tracker function
        if the actual model size exceeds the max size in the configuration.

        This memory-management routine shared by all the Hoeffding Trees is based on [^1].

        References
        ----------
        [^1]: Kirkby, R.B., 2007. Improving hoeffding trees (Doctoral dissertation,
        The University of Waikato).
        """
        leaves = self._find_leaves()
        total_active_size = 0
        total_inactive_size = 0
        for leaf in leaves:
            if leaf.is_active():
                total_active_size += calculate_object_size(leaf)
            else:
                total_inactive_size += calculate_object_size(leaf)
        if total_active_size > 0:
            self._active_leaf_size_estimate = total_active_size / self._n_active_leaves
        if total_inactive_size > 0:
            self._inactive_leaf_size_estimate = total_inactive_size / self._n_inactive_leaves
        actual_model_size = calculate_object_size(self)
        estimated_model_size = (
            self._n_active_leaves * self._active_leaf_size_estimate
            + self._n_inactive_leaves * self._inactive_leaf_size_estimate
        )
        self._size_estimate_overhead_fraction = actual_model_size / estimated_model_size
        if actual_model_size > self._max_byte_size:
            self._enforce_size_limit()

    def _deactivate_all_leaves(self):
        """Deactivate all leaves."""
        leaves = self._find_leaves()
        for leaf in leaves:
            leaf.deactivate()
            self._n_inactive_leaves += 1
            self._n_active_leaves -= 1

    def _find_leaves(self) -> list[HTLeaf]:
        """Find learning nodes in the tree.

        Returns
        -------
        List of learning nodes in the tree.
        """
        return [leaf for leaf in self._root.iter_leaves()]

    # Adapted from creme's original implementation
    def debug_one(self, x: dict) -> str | None:
        """Print an explanation of how `x` is predicted.

        Parameters
        ----------
        x
            A dictionary of features.

        Returns
        -------
            A representation of the path followed by the tree to predict `x`; `None` if
            the tree is empty.

        Notes
        -----
        Currently, Label Combination Hoeffding Tree Classifier (for multi-label
        classification) is not supported.
        """
        if self._root is None:
            return

        # We'll redirect all the print statement to a buffer, we'll return the content of the
        # buffer at the end
        buffer = io.StringIO()
        _print = functools.partial(print, file=buffer)

        for node in self._root.walk(x, until_leaf=True):
            if isinstance(node, HTLeaf):
                _print(repr(node))
            elif isinstance(node, DTBranch):
                try:
                    child_index = node.branch_no(x)
                except KeyError:
                    child_index, _ = node.most_common_path()

                _print(node.repr_branch(child_index))

        return buffer.getvalue()

    def draw(self, max_depth: int | None = None):
        """Draw the tree using the `graphviz` library.

        Since the tree is drawn without passing incoming samples, classification trees
        will show the majority class in their leaves, whereas regression trees will
        use the target mean.

        Parameters
        ----------
        max_depth
            Only the root will be drawn when set to `0`. Every node will be drawn when
            set to `None`.

        Notes
        -----
        Currently, Label Combination Hoeffding Tree Classifier (for multi-label
        classification) is not supported.

        Examples
        --------
        >>> from river import datasets
        >>> from river import tree
        >>> model = tree.HoeffdingTreeClassifier(
        ...    grace_period=5,
        ...    delta=1e-5,
        ...    split_criterion='gini',
        ...    max_depth=10,
        ...    tau=0.05,
        ... )
        >>> for x, y in datasets.Phishing():
        ...    model = model.learn_one(x, y)
        >>> dot = model.draw()

        .. image:: ../../docs/img/dtree_draw.svg
            :align: center
        """
        try:
            import graphviz
        except ImportError as e:
            raise ValueError("You have to install graphviz to use the draw method.") from e
        counter = 0

        def iterate(node=None):
            if node is None:
                yield None, None, self._root, 0, None
                yield from iterate(self._root)

            nonlocal counter
            parent_no = counter

            if isinstance(node, DTBranch):
                for branch_index, child in enumerate(node.children):
                    counter += 1
                    yield parent_no, node, child, counter, branch_index
                    if isinstance(child, DTBranch):
                        yield from iterate(child)

        if max_depth is None:
            max_depth = -1

        dot = graphviz.Digraph(
            graph_attr={"splines": "ortho", "forcelabels": "true", "overlap": "false"},
            node_attr={
                "shape": "box",
                "penwidth": "1.2",
                "fontname": "trebuchet",
                "fontsize": "11",
                "margin": "0.1,0.0",
            },
            edge_attr={"penwidth": "0.6", "center": "true", "fontsize": "7  "},
        )

        if isinstance(self, base.Classifier):
            n_colors = len(self.classes)  # type: ignore
        else:
            n_colors = 1

        # Pick a color palette which maps classes to colors
        new_color = functools.partial(next, iter(_color_brew(n_colors)))
        palette: collections.defaultdict = collections.defaultdict(new_color)

        for parent_no, parent, child, child_no, branch_index in iterate():
            if child.depth > max_depth and max_depth != -1:
                continue

            if isinstance(child, DTBranch):
                text = f"{child.feature}"  # type: ignore
            else:
                text = f"{repr(child)}\nsamples: {int(child.total_weight)}"

            # Pick a color, the hue depends on the class and the transparency on the distribution
            if isinstance(self, base.Classifier):
                class_proba = normalize_values_in_dict(child.stats, inplace=False)
                mode = max(class_proba, key=class_proba.get)
                p_mode = class_proba[mode]
                try:
                    alpha = (p_mode - 1 / n_colors) / (1 - 1 / n_colors)
                    fillcolor = str(transparency_hex(color=palette[mode], alpha=alpha))
                except ZeroDivisionError:
                    fillcolor = "#FFFFFF"
            else:
                fillcolor = "#FFFFFF"

            dot.node(f"{child_no}", text, fillcolor=fillcolor, style="filled")

            if parent_no is not None:
                dot.edge(
                    f"{parent_no}",
                    f"{child_no}",
                    xlabel=parent.repr_branch(branch_index, shorten=True),
                )

        return dot


# Utility adapted from the original creme's implementation
def _color_brew(n: int) -> list[tuple[int, int, int]]:
    """Generate n colors with equally spaced hues.

    Parameters
    ----------
    n
        The number of required colors.

    Returns
    -------
        List of n tuples of form (R, G, B) being the components of each color.
    References
    ----------
    https://github.com/scikit-learn/scikit-learn/blob/master/sklearn/tree/_export.py
    """
    colors = []

    # Initialize saturation & value; calculate chroma & value shift
    s, v = 0.75, 0.9
    c = s * v
    m = v - c

    for h in [i for i in range(25, 385, int(360 / n))]:
        # Calculate some intermediate values
        h_bar = h / 60.0
        x = c * (1 - abs((h_bar % 2) - 1))

        # Initialize RGB with same hue & chroma as our color
        rgb = [
            (c, x, 0),
            (x, c, 0),
            (0, c, x),
            (0, x, c),
            (x, 0, c),
            (c, 0, x),
            (c, x, 0),
        ]
        r, g, b = rgb[int(h_bar)]

        # Shift the initial RGB values to match value and store
        colors.append(((int(255 * (r + m))), (int(255 * (g + m))), (int(255 * (b + m)))))

    return colors


# Utility adapted from the original creme's implementation
def transparency_hex(color: tuple[int, int, int], alpha: float) -> str:
    """Apply alpha coefficient on hexadecimal color."""
    return "#{:02x}{:02x}{:02x}".format(
        *tuple([int(round(alpha * c + (1 - alpha) * 255, 0)) for c in color])
    )


# river/tree/hoeffding_tree_classifier.py

from __future__ import annotations

from river import base

from .hoeffding_tree import HoeffdingTree
from .nodes.branch import DTBranch
from .nodes.htc_nodes import LeafMajorityClass, LeafNaiveBayes, LeafNaiveBayesAdaptive
from .nodes.leaf import HTLeaf
from .split_criterion import GiniSplitCriterion, HellingerDistanceCriterion, InfoGainSplitCriterion
from .splitter import GaussianSplitter, Splitter


class HoeffdingTreeClassifier(HoeffdingTree, base.Classifier):
    """Hoeffding Tree or Very Fast Decision Tree classifier.

    Parameters
    ----------
    grace_period
        Number of instances a leaf should observe between split attempts.
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    split_criterion
        Split criterion to use.</br>
        - 'gini' - Gini</br>
        - 'info_gain' - Information Gain</br>
        - 'hellinger' - Helinger Distance</br>
    delta
        Significance level to calculate the Hoeffding bound. The significance level is given by
        `1 - delta`. Values closer to zero imply longer split decision delays.
    tau
        Threshold below which a split will be forced to break ties.
    leaf_prediction
        Prediction mechanism used at leafs.</br>
        - 'mc' - Majority Class</br>
        - 'nb' - Naive Bayes</br>
        - 'nba' - Naive Bayes Adaptive</br>
    nb_threshold
        Number of instances a leaf should observe before allowing Naive Bayes.
    nominal_attributes
        List of Nominal attributes identifiers. If empty, then assume that all numeric
        attributes should be treated as continuous.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.GaussianSplitter` is used if `splitter` is `None`.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.

    Notes
    -----
    A Hoeffding Tree [^1] is an incremental, anytime decision tree induction algorithm that is
    capable of learning from massive data streams, assuming that the distribution generating
    examples does not change over time. Hoeffding trees exploit the fact that a small sample can
    often be enough to choose an optimal splitting attribute. This idea is supported mathematically
    by the Hoeffding bound, which quantifies the number of observations (in our case, examples)
    needed to estimate some statistics within a prescribed precision (in our case, the goodness of
    an attribute).

    A theoretically appealing feature of Hoeffding Trees not shared by other incremental decision
    tree learners is that it has sound guarantees of performance. Using the Hoeffding bound one
    can show that its output is asymptotically nearly identical to that of a non-incremental
    learner using infinitely many examples. Implementation based on MOA [^2].

    References
    ----------

    [^1]: G. Hulten, L. Spencer, and P. Domingos. Mining time-changing data streams.
       In KDD’01, pages 97–106, San Francisco, CA, 2001. ACM Press.

    [^2]: Albert Bifet, Geoff Holmes, Richard Kirkby, Bernhard Pfahringer.
       MOA: Massive Online Analysis; Journal of Machine Learning Research 11: 1601-1604, 2010.

    Examples
    --------

    >>> from river.datasets import synth
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree

    >>> gen = synth.Agrawal(classification_function=0, seed=42)
    >>> # Take 1000 instances from the infinite data generator
    >>> dataset = iter(gen.take(1000))

    >>> model = tree.HoeffdingTreeClassifier(
    ...     grace_period=100,
    ...     delta=1e-5,
    ...     nominal_attributes=['elevel', 'car', 'zipcode']
    ... )

    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 84.58%
    """

    _GINI_SPLIT = "gini"
    _INFO_GAIN_SPLIT = "info_gain"
    _HELLINGER_SPLIT = "hellinger"
    _VALID_SPLIT_CRITERIA = [_GINI_SPLIT, _INFO_GAIN_SPLIT, _HELLINGER_SPLIT]

    _MAJORITY_CLASS = "mc"
    _NAIVE_BAYES = "nb"
    _NAIVE_BAYES_ADAPTIVE = "nba"
    _VALID_LEAF_PREDICTION = [_MAJORITY_CLASS, _NAIVE_BAYES, _NAIVE_BAYES_ADAPTIVE]

    def __init__(
        self,
        grace_period: int = 200,
        max_depth: int | None = None,
        split_criterion: str = "info_gain",
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "nba",
        nb_threshold: int = 0,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        binary_split: bool = False,
        max_size: float = 100.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
    ):
        super().__init__(
            max_depth=max_depth,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )
        self.grace_period = grace_period
        self.split_criterion = split_criterion
        self.delta = delta
        self.tau = tau
        self.leaf_prediction = leaf_prediction
        self.nb_threshold = nb_threshold
        self.nominal_attributes = nominal_attributes

        if splitter is None:
            self.splitter = GaussianSplitter()
        else:
            if not splitter.is_target_class:
                raise ValueError("The chosen splitter cannot be used in classification tasks.")
            self.splitter = splitter  # type: ignore

        # To keep track of the observed classes
        self.classes: set = set()

    @property
    def _mutable_attributes(self):
        return {"grace_period", "delta", "tau"}

    @HoeffdingTree.split_criterion.setter  # type: ignore
    def split_criterion(self, split_criterion):
        if split_criterion not in self._VALID_SPLIT_CRITERIA:
            print(
                "Invalid split_criterion option {}', will use default '{}'".format(
                    split_criterion, self._INFO_GAIN_SPLIT
                )
            )
            self._split_criterion = self._INFO_GAIN_SPLIT
        else:
            self._split_criterion = split_criterion

    @HoeffdingTree.leaf_prediction.setter  # type: ignore
    def leaf_prediction(self, leaf_prediction):
        if leaf_prediction not in self._VALID_LEAF_PREDICTION:
            print(
                "Invalid leaf_prediction option {}', will use default '{}'".format(
                    leaf_prediction, self._NAIVE_BAYES_ADAPTIVE
                )
            )
            self._leaf_prediction = self._NAIVE_BAYES_ADAPTIVE
        else:
            self._leaf_prediction = leaf_prediction

    def _new_leaf(self, initial_stats=None, parent=None):
        if initial_stats is None:
            initial_stats = {}
        if parent is None:
            depth = 0
        else:
            depth = parent.depth + 1

        if self._leaf_prediction == self._MAJORITY_CLASS:
            return LeafMajorityClass(initial_stats, depth, self.splitter)
        elif self._leaf_prediction == self._NAIVE_BAYES:
            return LeafNaiveBayes(initial_stats, depth, self.splitter)
        else:  # Naives Bayes Adaptive (default)
            return LeafNaiveBayesAdaptive(initial_stats, depth, self.splitter)

    def _new_split_criterion(self):
        if self._split_criterion == self._GINI_SPLIT:
            split_criterion = GiniSplitCriterion()
        elif self._split_criterion == self._INFO_GAIN_SPLIT:
            split_criterion = InfoGainSplitCriterion()
        elif self._split_criterion == self._HELLINGER_SPLIT:
            split_criterion = HellingerDistanceCriterion()
        else:
            split_criterion = InfoGainSplitCriterion()

        return split_criterion

    def _attempt_to_split(self, leaf: HTLeaf, parent: DTBranch, parent_branch: int, **kwargs):
        """Attempt to split a leaf.

        If the samples seen so far are not from the same class then:

        1. Find split candidates and select the top 2.
        2. Compute the Hoeffding bound.
        3. If the difference between the top 2 split candidates is larger than the Hoeffding bound:
           3.1 Replace the leaf node by a split node (branch node).
           3.2 Add a new leaf node on each branch of the new split node.
           3.3 Update tree's metrics

        Optional: Disable poor attributes. Depends on the tree's configuration.

        Parameters
        ----------
        leaf
            The leaf to evaluate.
        parent
            The leaf's parent.
        parent_branch
            Parent leaf's branch index.
        kwargs
            Other parameters passed to the new branch.
        """
        if not leaf.observed_class_distribution_is_pure():  # type: ignore
            split_criterion = self._new_split_criterion()

            best_split_suggestions = leaf.best_split_suggestions(split_criterion, self)
            best_split_suggestions.sort()
            should_split = False
            if len(best_split_suggestions) < 2:
                should_split = len(best_split_suggestions) > 0
            else:
                hoeffding_bound = self._hoeffding_bound(
                    split_criterion.range_of_merit(leaf.stats),
                    self.delta,
                    leaf.total_weight,
                )
                best_suggestion = best_split_suggestions[-1]
                second_best_suggestion = best_split_suggestions[-2]
                if (
                    best_suggestion.merit - second_best_suggestion.merit > hoeffding_bound
                    or hoeffding_bound < self.tau
                ):
                    should_split = True
                if self.remove_poor_attrs:
                    poor_atts = set()
                    # Add any poor attribute to set
                    for suggestion in best_split_suggestions:
                        if (
                            suggestion.feature
                            and best_suggestion.merit - suggestion.merit > hoeffding_bound
                        ):
                            poor_atts.add(suggestion.feature)
                    for poor_att in poor_atts:
                        leaf.disable_attribute(poor_att)
            if should_split:
                split_decision = best_split_suggestions[-1]
                if split_decision.feature is None:
                    # Pre-pruning - null wins
                    leaf.deactivate()
                    self._n_inactive_leaves += 1
                    self._n_active_leaves -= 1
                else:
                    branch = self._branch_selector(
                        split_decision.numerical_feature, split_decision.multiway_split
                    )
                    leaves = tuple(
                        self._new_leaf(initial_stats, parent=leaf)
                        for initial_stats in split_decision.children_stats  # type: ignore
                    )

                    new_split = split_decision.assemble(
                        branch, leaf.stats, leaf.depth, *leaves, **kwargs
                    )

                    self._n_active_leaves -= 1
                    self._n_active_leaves += len(leaves)
                    if parent is None:
                        self._root = new_split
                    else:
                        parent.children[parent_branch] = new_split

                # Manage memory
                self._enforce_size_limit()

    def learn_one(self, x, y, *, sample_weight=1.0):
        """Train the model on instance x and corresponding target y.

        Parameters
        ----------
        x
            Instance attributes.
        y
            Class label for sample x.
        sample_weight
            Sample weight.

        Returns
        -------
        self

        Notes
        -----
        Training tasks:

        * If the tree is empty, create a leaf node as the root.
        * If the tree is already initialized, find the corresponding leaf for
          the instance and update the leaf node statistics.
        * If growth is allowed and the number of instances that the leaf has
          observed between split attempts exceed the grace period then attempt
          to split.
        """

        # Updates the set of observed classes
        self.classes.add(y)

        self._train_weight_seen_by_model += sample_weight

        if self._root is None:
            self._root = self._new_leaf()
            self._n_active_leaves = 1

        p_node = None
        node = None
        if isinstance(self._root, DTBranch):
            path = iter(self._root.walk(x, until_leaf=False))
            while True:
                aux = next(path, None)
                if aux is None:
                    break
                p_node = node
                node = aux
        else:
            node = self._root

        if isinstance(node, HTLeaf):
            node.learn_one(x, y, sample_weight=sample_weight, tree=self)
            if self._growth_allowed and node.is_active():
                if node.depth >= self.max_depth:  # Max depth reached
                    node.deactivate()
                    self._n_active_leaves -= 1
                    self._n_inactive_leaves += 1
                else:
                    weight_seen = node.total_weight
                    weight_diff = weight_seen - node.last_split_attempt_at
                    if weight_diff >= self.grace_period:
                        p_branch = p_node.branch_no(x) if isinstance(p_node, DTBranch) else None
                        self._attempt_to_split(node, p_node, p_branch)
                        node.last_split_attempt_at = weight_seen
        else:
            while True:
                # Split node encountered a previously unseen categorical value (in a multi-way
                #  test), so there is no branch to sort the instance to
                if node.max_branches() == -1 and node.feature in x:
                    # Create a new branch to the new categorical value
                    leaf = self._new_leaf(parent=node)
                    node.add_child(x[node.feature], leaf)
                    self._n_active_leaves += 1
                    node = leaf
                # The split feature is missing in the instance. Hence, we pass the new example
                # to the most traversed path in the current subtree
                else:
                    _, node = node.most_common_path()
                    # And we keep trying to reach a leaf
                    if isinstance(node, DTBranch):
                        node = node.traverse(x, until_leaf=False)
                # Once a leaf is reached, the traversal can stop
                if isinstance(node, HTLeaf):
                    break
            # Learn from the sample
            node.learn_one(x, y, sample_weight=sample_weight, tree=self)

        if self._train_weight_seen_by_model % self.memory_estimate_period == 0:
            self._estimate_model_size()

        return self

    def predict_proba_one(self, x):
        proba = {c: 0.0 for c in sorted(self.classes)}
        if self._root is not None:
            if isinstance(self._root, DTBranch):
                leaf = self._root.traverse(x, until_leaf=True)
            else:
                leaf = self._root

            proba.update(leaf.prediction(x, tree=self))
        return proba

    @property
    def _multiclass(self):
        return True


# river/tree/hoeffding_tree_regressor.py

from __future__ import annotations

from copy import deepcopy

from river import base, linear_model

from .hoeffding_tree import HoeffdingTree
from .nodes.branch import DTBranch
from .nodes.htr_nodes import LeafAdaptive, LeafMean, LeafModel
from .nodes.leaf import HTLeaf
from .split_criterion import VarianceReductionSplitCriterion
from .splitter import Splitter, TEBSTSplitter


class HoeffdingTreeRegressor(HoeffdingTree, base.Regressor):
    """Hoeffding Tree regressor.

    Parameters
    ----------
    grace_period
        Number of instances a leaf should observe between split attempts.
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    delta
        Significance level to calculate the Hoeffding bound. The significance level is given by
        `1 - delta`. Values closer to zero imply longer split decision delays.
    tau
        Threshold below which a split will be forced to break ties.
    leaf_prediction
        Prediction mechanism used at leafs.</br>
        - 'mean' - Target mean</br>
        - 'model' - Uses the model defined in `leaf_model`</br>
        - 'adaptive' - Chooses between 'mean' and 'model' dynamically</br>
    leaf_model
        The regression model used to provide responses if `leaf_prediction='model'`. If not
        provided an instance of `river.linear_model.LinearRegression` with the default
        hyperparameters is used.
    model_selector_decay
        The exponential decaying factor applied to the learning models' squared errors, that
        are monitored if `leaf_prediction='adaptive'`. Must be between `0` and `1`. The closer
        to `1`, the more importance is going to be given to past observations. On the other hand,
        if its value approaches `0`, the recent observed errors are going to have more influence
        on the final decision.
    nominal_attributes
        List of Nominal attributes identifiers. If empty, then assume that all numeric attributes
        should be treated as continuous.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.TEBSTSplitter` is used if `splitter` is `None`.
    min_samples_split
        The minimum number of samples every branch resulting from a split candidate must have
        to be considered valid.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.

    Notes
    -----
    The Hoeffding Tree Regressor (HTR) is an adaptation of the incremental tree algorithm of the
    same name for classification. Similarly to its classification counterpart, HTR uses the
    Hoeffding bound to control its split decisions. Differently from the classification algorithm,
    HTR relies on calculating the reduction of variance in the target space to decide among the
    split candidates. The smallest the variance at its leaf nodes, the more homogeneous the
    partitions are. At its leaf nodes, HTR fits either linear models or uses the target
    average as the predictor.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     tree.HoeffdingTreeRegressor(
    ...         grace_period=100,
    ...         model_selector_decay=0.9
    ...     )
    ... )

    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 0.781781
    """

    _TARGET_MEAN = "mean"
    _MODEL = "model"
    _ADAPTIVE = "adaptive"

    _VALID_LEAF_PREDICTION = [_TARGET_MEAN, _MODEL, _ADAPTIVE]

    def __init__(
        self,
        grace_period: int = 200,
        max_depth: int | None = None,
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "adaptive",
        leaf_model: base.Regressor | None = None,
        model_selector_decay: float = 0.95,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        min_samples_split: int = 5,
        binary_split: bool = False,
        max_size: float = 500.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
    ):
        super().__init__(
            max_depth=max_depth,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self._split_criterion: str = "vr"
        self.grace_period = grace_period
        self.delta = delta
        self.tau = tau
        self.leaf_prediction = leaf_prediction
        self.leaf_model = leaf_model if leaf_model else linear_model.LinearRegression()
        self.model_selector_decay = model_selector_decay
        self.nominal_attributes = nominal_attributes
        self.min_samples_split = min_samples_split

        if splitter is None:
            self.splitter = TEBSTSplitter()
        else:
            if splitter.is_target_class:
                raise ValueError("The chosen splitter cannot be used in regression tasks.")
            self.splitter = splitter  # type: ignore

    @property
    def _mutable_attributes(self):
        return {"grace_period", "delta", "tau"}

    @HoeffdingTree.leaf_prediction.setter  # type: ignore
    def leaf_prediction(self, leaf_prediction):
        if leaf_prediction not in self._VALID_LEAF_PREDICTION:
            print(
                'Invalid leaf_prediction option "{}", will use default "{}"'.format(
                    leaf_prediction, self._MODEL
                )
            )
            self._leaf_prediction = self._MODEL
        else:
            self._leaf_prediction = leaf_prediction

    @HoeffdingTree.split_criterion.setter  # type: ignore
    def split_criterion(self, split_criterion):
        if split_criterion != "vr":  # variance reduction
            print(
                "Invalid split_criterion option {}', will use default '{}'".format(
                    split_criterion, "vr"
                )
            )
            self._split_criterion = "vr"
        else:
            self._split_criterion = split_criterion

    def _new_split_criterion(self):
        return VarianceReductionSplitCriterion(min_samples_split=self.min_samples_split)

    def _new_leaf(self, initial_stats=None, parent=None):
        """Create a new learning node.

        The type of learning node depends on the tree configuration.
        """
        if parent is not None:
            depth = parent.depth + 1
        else:
            depth = 0

        leaf_model = None
        if self.leaf_prediction in {self._MODEL, self._ADAPTIVE}:
            if parent is None:
                leaf_model = deepcopy(self.leaf_model)
            else:
                try:
                    leaf_model = deepcopy(parent._leaf_model)  # noqa
                except AttributeError:
                    leaf_model = deepcopy(self.leaf_model)

        if self.leaf_prediction == self._TARGET_MEAN:
            return LeafMean(initial_stats, depth, self.splitter)
        elif self.leaf_prediction == self._MODEL:
            return LeafModel(initial_stats, depth, self.splitter, leaf_model)
        else:  # adaptive learning node
            new_adaptive = LeafAdaptive(initial_stats, depth, self.splitter, leaf_model)
            if parent is not None and isinstance(parent, LeafAdaptive):
                new_adaptive._fmse_mean = parent._fmse_mean  # noqa
                new_adaptive._fmse_model = parent._fmse_model  # noqa

            return new_adaptive

    def learn_one(self, x, y, *, sample_weight=1.0):
        """Train the tree model on sample x and corresponding target y.

        Parameters
        ----------
        x
            Instance attributes.
        y
            Target value for sample x.
        sample_weight
            The weight of the sample.

        Returns
        -------
        self
        """

        self._train_weight_seen_by_model += sample_weight

        if self._root is None:
            self._root = self._new_leaf()
            self._n_active_leaves = 1

        p_node = None
        node = None
        if isinstance(self._root, DTBranch):
            path = iter(self._root.walk(x, until_leaf=False))
            while True:
                aux = next(path, None)
                if aux is None:
                    break
                p_node = node
                node = aux
        else:
            node = self._root

        if isinstance(node, HTLeaf):
            node.learn_one(x, y, sample_weight=sample_weight, tree=self)
            if self._growth_allowed and node.is_active():
                if node.depth >= self.max_depth:  # Max depth reached
                    node.deactivate()
                    self._n_active_leaves -= 1
                    self._n_inactive_leaves += 1
                else:
                    weight_seen = node.total_weight
                    weight_diff = weight_seen - node.last_split_attempt_at
                    if weight_diff >= self.grace_period:
                        p_branch = p_node.branch_no(x) if isinstance(p_node, DTBranch) else None
                        self._attempt_to_split(node, p_node, p_branch)
                        node.last_split_attempt_at = weight_seen
        else:
            while True:
                # Split node encountered a previously unseen categorical value (in a multi-way
                #  test), so there is no branch to sort the instance to
                if node.max_branches() == -1 and node.feature in x:
                    # Create a new branch to the new categorical value
                    leaf = self._new_leaf(parent=node)
                    node.add_child(x[node.feature], leaf)
                    self._n_active_leaves += 1
                    node = leaf
                # The split feature is missing in the instance. Hence, we pass the new example
                # to the most traversed path in the current subtree
                else:
                    _, node = node.most_common_path()
                    # And we keep trying to reach a leaf
                    if isinstance(node, DTBranch):
                        node = node.traverse(x, until_leaf=False)
                # Once a leaf is reached, the traversal can stop
                if isinstance(node, HTLeaf):
                    break
            # Learn from the sample
            node.learn_one(x, y, sample_weight=sample_weight, tree=self)

        if self._train_weight_seen_by_model % self.memory_estimate_period == 0:
            self._estimate_model_size()

        return self

    def predict_one(self, x):
        """Predict the target value using one of the leaf prediction strategies.

        Parameters
        ----------
        x
            Instance for which we want to predict the target.

        Returns
        -------
        Predicted target value.

        """
        pred = 0.0
        if self._root is not None:
            if isinstance(self._root, DTBranch):
                leaf = self._root.traverse(x, until_leaf=True)
            else:
                leaf = self._root

            pred = leaf.prediction(x, tree=self)
        return pred

    def _attempt_to_split(self, leaf: HTLeaf, parent: DTBranch, parent_branch: int, **kwargs):
        """Attempt to split a node.

        If the target's variance is high at the leaf node, then:

        1. Find split candidates and select the top 2.
        2. Compute the Hoeffding bound.
        3. If the ratio between the merit of the second best split candidate and the merit of the
        best one is smaller than 1 minus the Hoeffding bound (or a tie breaking decision
        takes place), then:
           3.1 Replace the leaf node by a split node.
           3.2 Add a new leaf node on each branch of the new split node.
           3.3 Update tree's metrics

        Optional: Disable poor attribute. Depends on the tree's configuration.

        Parameters
        ----------
        leaf
            The node to evaluate.
        parent
            The node's parent in the tree.
        parent_branch
            Parent node's branch index.
        kwargs
            Other parameters passed to the new branch.

        """
        split_criterion = self._new_split_criterion()
        best_split_suggestions = leaf.best_split_suggestions(split_criterion, self)
        best_split_suggestions.sort()
        should_split = False
        if len(best_split_suggestions) < 2:
            should_split = len(best_split_suggestions) > 0
        else:
            hoeffding_bound = self._hoeffding_bound(
                split_criterion.range_of_merit(leaf.stats),
                self.delta,
                leaf.total_weight,
            )
            best_suggestion = best_split_suggestions[-1]
            second_best_suggestion = best_split_suggestions[-2]
            if best_suggestion.merit > 0.0 and (
                second_best_suggestion.merit / best_suggestion.merit < 1 - hoeffding_bound
                or hoeffding_bound < self.tau
            ):
                should_split = True
            if self.remove_poor_attrs:
                poor_attrs = set()
                best_ratio = second_best_suggestion.merit / best_suggestion.merit

                # Add any poor attribute to set
                for suggestion in best_split_suggestions:
                    if (
                        suggestion.feature
                        and suggestion.merit / best_suggestion.merit
                        < best_ratio - 2 * hoeffding_bound
                    ):
                        poor_attrs.add(suggestion.feature)
                for poor_att in poor_attrs:
                    leaf.disable_attribute(poor_att)
        if should_split:
            split_decision = best_split_suggestions[-1]
            if split_decision.feature is None:
                # Pre-pruning - null wins
                leaf.deactivate()
                self._n_inactive_leaves += 1
                self._n_active_leaves -= 1
            else:
                branch = self._branch_selector(
                    split_decision.numerical_feature, split_decision.multiway_split
                )
                leaves = tuple(
                    self._new_leaf(initial_stats, parent=leaf)
                    for initial_stats in split_decision.children_stats  # type: ignore
                )

                new_split = split_decision.assemble(
                    branch, leaf.stats, leaf.depth, *leaves, **kwargs
                )

                self._n_active_leaves -= 1
                self._n_active_leaves += len(leaves)
                if parent is None:
                    self._root = new_split
                else:
                    parent.children[parent_branch] = new_split

            # Manage memory
            self._enforce_size_limit()
        elif (
            len(best_split_suggestions) >= 2
            and best_split_suggestions[-1].merit > 0
            and best_split_suggestions[-2].merit > 0
        ):
            last_check_ratio = best_split_suggestions[-2].merit / best_split_suggestions[-1].merit
            last_check_vr = best_split_suggestions[-1].merit

            leaf.manage_memory(  # type: ignore
                split_criterion, last_check_ratio, last_check_vr, hoeffding_bound
            )


# river/tree/isoup_tree_regressor.py

from __future__ import annotations

from copy import deepcopy

from river import base, tree

from .nodes.branch import DTBranch
from .nodes.isouptr_nodes import LeafAdaptiveMultiTarget, LeafMeanMultiTarget, LeafModelMultiTarget
from .split_criterion import IntraClusterVarianceReductionSplitCriterion
from .splitter import Splitter


class iSOUPTreeRegressor(tree.HoeffdingTreeRegressor, base.MultiTargetRegressor):
    """Incremental Structured Output Prediction Tree (iSOUP-Tree) for multi-target regression.

    This is an implementation of the iSOUP-Tree proposed by A. Osojnik, P. Panov, and
    S. Džeroski [^1].

    Parameters
    ----------
    grace_period
        Number of instances a leaf should observe between split attempts.
    max_depth
        The maximum depth a tree can reach. If `None`, the tree will grow indefinitely.
    delta
        Allowed error in split decision, a value closer to 0 takes longer to
        decide.
    tau
        Threshold below which a split will be forced to break ties.
    leaf_prediction
        Prediction mechanism used at leafs.</br>
        - 'mean' - Target mean</br>
        - 'model' - Uses the model defined in `leaf_model`</br>
        - 'adaptive' - Chooses between 'mean' and 'model' dynamically</br>
    leaf_model
        The regression model(s) used to provide responses if `leaf_prediction='model'`. It can
        be either a regressor (in which case it is going to be replicated to all the targets)
        or a dictionary whose keys are target identifiers, and the values are instances of
        `river.base.Regressor.` If not provided, instances of `river.linear_model.LinearRegression`
        with the default hyperparameters are used for all the targets. If a dictionary is passed
        and not all target models are specified, copies from the first model match in the
        dictionary will be used to the remaining targets.
    model_selector_decay
        The exponential decaying factor applied to the learning models' squared errors, that
        are monitored if `leaf_prediction='adaptive'`. Must be between `0` and `1`. The closer
        to `1`, the more importance is going to be given to past observations. On the other hand,
        if its value approaches `0`, the recent observed errors are going to have more influence
        on the final decision.
    nominal_attributes
        List of Nominal attributes identifiers. If empty, then assume that all numeric attributes
        should be treated as continuous.
    splitter
        The Splitter or Attribute Observer (AO) used to monitor the class statistics of numeric
        features and perform splits. Splitters are available in the `tree.splitter` module.
        Different splitters are available for classification and regression tasks. Classification
        and regression splitters can be distinguished by their property `is_target_class`.
        This is an advanced option. Special care must be taken when choosing different splitters.
        By default, `tree.splitter.TEBSTSplitter` is used if `splitter` is `None`.
    min_samples_split
        The minimum number of samples every branch resulting from a split candidate must have
        to be considered valid.
    binary_split
        If True, only allow binary splits.
    max_size
        The max size of the tree, in Megabytes (MB).
    memory_estimate_period
        Interval (number of processed instances) between memory consumption checks.
    stop_mem_management
        If True, stop growing as soon as memory limit is hit.
    remove_poor_attrs
        If True, disable poor attributes to reduce memory usage.
    merit_preprune
        If True, enable merit-based tree pre-pruning.

    References
    ----------
    [^1]: Aljaž Osojnik, Panče Panov, and Sašo Džeroski. "Tree-based methods for online
        multi-target regression." Journal of Intelligent Information Systems 50.2 (2018): 315-339.

    Examples
    --------

    >>> import numbers
    >>> from river import compose
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import linear_model
    >>> from river import metrics
    >>> from river import preprocessing
    >>> from river import tree

    >>> dataset = datasets.SolarFlare()

    >>> num = compose.SelectType(numbers.Number) | preprocessing.MinMaxScaler()
    >>> cat = compose.SelectType(str) | preprocessing.OneHotEncoder()

    >>> model = tree.iSOUPTreeRegressor(
    ...     grace_period=100,
    ...     leaf_prediction='model',
    ...     leaf_model={
    ...         'c-class-flares': linear_model.LinearRegression(l2=0.02),
    ...         'm-class-flares': linear_model.PARegressor(),
    ...         'x-class-flares': linear_model.LinearRegression(l2=0.1)
    ...     }
    ... )

    >>> pipeline = (num + cat) | model
    >>> metric = metrics.multioutput.MicroAverage(metrics.MAE())

    >>> evaluate.progressive_val_score(dataset, pipeline, metric)
    MicroAverage(MAE): 0.426177

    """

    def __init__(
        self,
        grace_period: int = 200,
        max_depth: int | None = None,
        delta: float = 1e-7,
        tau: float = 0.05,
        leaf_prediction: str = "adaptive",
        leaf_model: base.Regressor | dict | None = None,
        model_selector_decay: float = 0.95,
        nominal_attributes: list | None = None,
        splitter: Splitter | None = None,
        min_samples_split: int = 5,
        binary_split: bool = False,
        max_size: float = 500.0,
        memory_estimate_period: int = 1000000,
        stop_mem_management: bool = False,
        remove_poor_attrs: bool = False,
        merit_preprune: bool = True,
    ):
        super().__init__(
            grace_period=grace_period,
            max_depth=max_depth,
            delta=delta,
            tau=tau,
            leaf_prediction=leaf_prediction,
            leaf_model=leaf_model,  # type: ignore
            model_selector_decay=model_selector_decay,
            nominal_attributes=nominal_attributes,
            splitter=splitter,
            min_samples_split=min_samples_split,
            binary_split=binary_split,
            max_size=max_size,
            memory_estimate_period=memory_estimate_period,
            stop_mem_management=stop_mem_management,
            remove_poor_attrs=remove_poor_attrs,
            merit_preprune=merit_preprune,
        )

        self.split_criterion: str = "icvr"  # intra cluster variance reduction
        self.targets: set = set()

    @property
    def _mutable_attributes(self):
        return {"grace_period", "delta", "tau"}

    @tree.HoeffdingTreeRegressor.split_criterion.setter
    def split_criterion(self, split_criterion):
        if split_criterion == "vr":
            # Corner case due to parent class initialization
            split_criterion = "icvr"
        if split_criterion != "icvr":  # intra cluster variance reduction
            print(
                'Invalid split_criterion option "{}", will use default "{}"'.format(
                    split_criterion, "icvr"
                )
            )
            self._split_criterion = "icvr"
        else:
            self._split_criterion = split_criterion

    def _new_split_criterion(self):
        return IntraClusterVarianceReductionSplitCriterion(min_samples_split=self.min_samples_split)

    def _new_leaf(self, initial_stats=None, parent=None):
        """Create a new learning node. The type of learning node depends on
        the tree configuration.
        """
        if parent is not None:
            depth = parent.depth + 1
        else:
            depth = 0

        leaf_models = None
        if self.leaf_prediction in {self._MODEL, self._ADAPTIVE}:
            if parent is None:
                leaf_models = {}
            else:
                try:
                    leaf_models = deepcopy(parent._leaf_models)  # noqa
                except AttributeError:
                    # Due to an emerging category in a nominal feature, a split node was reached
                    leaf_models = {}

        if self.leaf_prediction == self._TARGET_MEAN:
            return LeafMeanMultiTarget(initial_stats, depth, self.splitter)
        elif self.leaf_prediction == self._MODEL:
            return LeafModelMultiTarget(initial_stats, depth, self.splitter, leaf_models)
        else:  # adaptive learning node
            new_adaptive = LeafAdaptiveMultiTarget(initial_stats, depth, self.splitter, leaf_models)
            if parent is not None and isinstance(parent, LeafAdaptiveMultiTarget):
                new_adaptive._fmse_mean = parent._fmse_mean.copy()  # noqa
                new_adaptive._fmse_model = parent._fmse_model.copy()  # noqa

            return new_adaptive

    def learn_one(self, x, y, *, sample_weight: float = 1.0) -> iSOUPTreeRegressor:  # type: ignore
        """Incrementally train the model with one sample.

        Training tasks:

        * If the tree is empty, create a leaf node as the root.
        * If the tree is already initialized, find the corresponding leaf for
          the instance and update the leaf node statistics.
        * If growth is allowed and the number of instances that the leaf has
          observed between split attempts exceed the grace period then attempt
          to split.

        Parameters
        ----------
        x
            Instance attributes.
        y
            Target values.
        sample_weight
            The weight of the passed sample.
        """
        # Update target set
        self.targets.update(y.keys())

        super().learn_one(x, y, sample_weight=sample_weight)  # type: ignore

        return self

    def predict_one(self, x):
        pred = {}
        if self._root is not None:
            if isinstance(self._root, DTBranch):
                leaf = self._root.traverse(x, until_leaf=True)
            else:
                leaf = self._root

            pred = leaf.prediction(x, tree=self)  # type: ignore
        return pred


# river/tree/losses.py

from __future__ import annotations

import abc
import math

from .utils import GradHess


class Loss(abc.ABC):
    """Base class to implement optimization objectives used in Stochastic Gradient Trees."""

    @abc.abstractmethod
    def compute_derivatives(self, y_true: float, y_pred: float) -> GradHess:
        """Return the gradient and hessian data concerning one instance and its prediction.

        Parameters
        ----------
        y_true
            Target value.
        y_pred
            Predicted target value.
        """
        raise NotImplementedError

    def transfer(self, y: float) -> float:
        """Optionally apply some transformation to the value predicted by the tree before
        returning it.

        For instance, in classification, the softmax operation might be applied.

        Parameters
        ----------
        y
            Value to be transformed.
        """
        return y


class BinaryCrossEntropyLoss(Loss):
    """Loss function used in binary classification tasks."""

    def compute_derivatives(self, y_true, y_pred):
        y_trs = self.transfer(y_pred)

        return GradHess(y_trs - y_true, y_trs * (1.0 - y_trs))

    def transfer(self, y):
        return 1.0 / (1.0 + math.exp(-y))


class SquaredErrorLoss(Loss):
    """Loss function used in regression tasks."""

    def compute_derivatives(self, y_true, y_pred):
        return GradHess(y_pred - y_true, 1.0)


# river/tree/setup.py

from __future__ import annotations

from numpy.distutils.misc_util import Configuration


def configuration(parent_package="", top_path=None):
    config = Configuration("tree", parent_package, top_path)

    # submodules which do not have their own setup.py
    config.add_subpackage("splitter")
    config.add_subpackage("nodes")
    config.add_subpackage("split_criterion")

    return config


if __name__ == "__main__":
    from numpy.distutils.core import setup

    setup(**configuration().todict())


# river/tree/stochastic_gradient_tree.py

from __future__ import annotations

import abc
import math

from scipy.stats import f as f_dist

from river import base, tree

from .losses import BinaryCrossEntropyLoss, SquaredErrorLoss
from .nodes.branch import DTBranch, NominalMultiwayBranch, NumericBinaryBranch
from .nodes.sgt_nodes import SGTLeaf
from .utils import BranchFactory, GradHessMerit


class StochasticGradientTree(base.Estimator, abc.ABC):
    """Base Stochastic Gradient Tree (SGT) class.

    This class defines the main characteristics that are shared by the different SGT
    implementations.

    """

    def __init__(
        self,
        loss_func,
        delta,
        grace_period,
        init_pred,
        max_depth,
        lambda_value,
        gamma,
        nominal_attributes,
        feature_quantizer,
    ):
        # What really defines how a SGT works is its loss function
        self.loss_func = loss_func
        self.delta = delta
        self.grace_period = grace_period
        self.init_pred = init_pred
        self.max_depth = max_depth if max_depth else math.inf

        if lambda_value < 0.0:
            raise ValueError('Invalid value: "lambda_value" must be positive.')

        if gamma < 0.0:
            raise ValueError('Invalid value: "gamma" must be positive.')

        self.lambda_value = lambda_value
        self.gamma = gamma
        self.nominal_attributes = set(nominal_attributes) if nominal_attributes else set()
        self.feature_quantizer = (
            feature_quantizer if feature_quantizer is not None else tree.splitter.StaticQuantizer()
        )

        self._root: SGTLeaf | DTBranch = SGTLeaf(prediction=self.init_pred)

        # set used to check whether categorical feature has been already split
        self._split_features = set()
        self._n_splits = 0
        self._n_node_updates = 0
        self._n_observations = 0

    def _target_transform(self, y):
        """Apply transformation to the raw target input.

        Different strategies are used for classification and regression. By default, use
        an identity function.

        Parameters
        ----------
        y
            The target value, over which the transformation will be applied.
        """
        return y

    def learn_one(self, x, y, *, w=1.0):
        self._n_observations += w

        """ Update Stochastic Gradient Tree with a single instance. """
        y_true_trs = self._target_transform(y)

        p_node = None
        node = None
        if not isinstance(self._root, SGTLeaf):
            path = iter(self._root.walk(x, until_leaf=False))
            while True:
                aux = next(path, None)
                if aux is None:
                    break
                p_node = node
                node = aux
        else:
            node = self._root

        # A leaf could not be reached in a single attempt let's deal with that
        if isinstance(node, NumericBinaryBranch) or isinstance(node, NominalMultiwayBranch):
            while True:
                # Split node encountered a previously unseen categorical value (in a multi-way
                #  test), so there is no branch to sort the instance to
                if node.max_branches() == -1 and node.feature in x:
                    # Create a new branch to the new categorical value
                    leaf = SGTLeaf(depth=node.depth + 1, split_params=node.stats.copy())
                    #
                    node.add_child(x[node.feature], leaf)
                    node = leaf
                # The split feature is missing in the instance. Hence, we pass the new example
                # to the most traversed path in the current subtree
                else:
                    _, node = node.most_common_path()
                    # And we keep trying to reach a leaf
                    if isinstance(node, DTBranch):
                        node = node.traverse(x, until_leaf=False)
                # Once a leaf is reached, the traversal can stop
                if isinstance(node, SGTLeaf):
                    break

            y_pred_raw = self.loss_func.transfer(node.prediction())
            grad_hess = self.loss_func.compute_derivatives(y_true_trs, y_pred_raw)
            node.update(x, grad_hess, self, w)
        else:  # Node is a leaf
            y_pred_raw = self.loss_func.transfer(node.prediction())
            grad_hess = self.loss_func.compute_derivatives(y_true_trs, y_pred_raw)
            node.update(x, grad_hess, self, w)

            if node.total_weight - node.last_split_attempt_at < self.grace_period:
                return self

            # Update split attempt data
            node.last_split_attempt_at = node.total_weight

            # If the maximum depth is reached, attempt to apply a "null split", i.e., update the
            # prediction value
            if node.depth >= self.max_depth:
                # Null split: update the prediction using the new gradient information
                best_split = BranchFactory()
                best_split.merit = GradHessMerit()
                best_split.merit.delta_pred = node.delta_prediction(
                    node.update_stats.mean, self.lambda_value
                )
                dlms = node.update_stats.delta_loss_mean_var(best_split.merit.delta_pred)
                best_split.merit.loss_mean = dlms.mean.get()
                best_split.merit.loss_var = dlms.get()
            else:  # Proceed with the standard split attempt procedure
                best_split = node.find_best_split(self)

            p = self._compute_p_value(best_split.merit, node.total_weight)

            if p < self.delta and best_split.merit.loss_mean < 0:
                p_branch = p_node.branch_no(x) if isinstance(p_node, DTBranch) else None
                node.apply_split(best_split, p_node, p_branch, self)

        return self

    @staticmethod
    def _compute_p_value(merit, n_observations):
        # Null hypothesis: expected loss is zero
        # Alternative hypothesis: expected loss is not zero

        f_value = (
            n_observations * (merit.loss_mean * merit.loss_mean) / merit.loss_var
            if merit.loss_var > 0.0
            else None
        )

        if f_value is None:
            return 1.0

        return 1 - f_dist.cdf(f_value, 1, n_observations - 1)

    @property
    def n_splits(self):
        return self._n_splits

    @property
    def n_node_updates(self):
        return self._n_node_updates

    @property
    def n_observations(self):
        return self._n_observations

    @property
    def height(self) -> int:
        if self._root:
            return self._root.height
        return 0

    @property
    def n_nodes(self):
        if self._root:
            return self._root.n_nodes

    @property
    def n_branches(self):
        if self._root:
            return self._root.n_branches

    @property
    def n_leaves(self):
        if self._root:
            return self._root.n_leaves


class SGTClassifier(StochasticGradientTree, base.Classifier):
    """Stochastic Gradient Tree[^1] for binary classification.

    Binary decision tree classifier that minimizes the binary cross-entropy to guide its growth.

    Stochastic Gradient Trees (SGT) directly minimize a loss function to guide tree growth and
    update their predictions. Thus, they differ from other incrementally tree learners that do
    not directly optimize the loss, but data impurity-related heuristics.

    Parameters
    ----------
    delta
        Define the significance level of the F-tests performed to decide upon creating splits
        or updating predictions.
    grace_period
        Interval between split attempts or prediction updates.
    init_pred
        Initial value predicted by the tree.
    max_depth
        The maximum depth the tree might reach. If set to `None`, the trees will grow
        indefinitely.
    lambda_value
        Positive float value used to impose a penalty over the tree's predictions and force
        them to become smaller. The greater the lambda value, the more constrained are the
        predictions.
    gamma
        Positive float value used to impose a penalty over the tree's splits and force them to
        be avoided when possible. The greater the gamma value, the smaller the chance of a
        split occurring.
    nominal_attributes
        List with identifiers of the nominal attributes. If None, all features containing
        numbers are assumed to be numeric.
    feature_quantizer
        The algorithm used to quantize numeric features. Either a static quantizer (as in the
        original implementation) or a dynamic quantizer can be used. The correct choice and setup
        of the feature quantizer is a crucial step to determine the performance of SGTs.
        Feature quantizers are akin to the attribute observers used in Hoeffding Trees. By
        default, an instance of `tree.splitter.StaticQuantizer` (with default parameters) is
        used if this parameter is not set.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree

    >>> dataset = datasets.Phishing()
    >>> model = tree.SGTClassifier(
    ...     feature_quantizer=tree.splitter.StaticQuantizer(
    ...         n_bins=32, warm_start=10
    ...     )
    ... )
    >>> metric = metrics.Accuracy()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    Accuracy: 82.24%

    References
    ----------
    [^1]: Gouk, H., Pfahringer, B., & Frank, E. (2019, October). Stochastic Gradient Trees.
    In Asian Conference on Machine Learning (pp. 1094-1109).

    """

    def __init__(
        self,
        delta: float = 1e-7,
        grace_period: int = 200,
        init_pred: float = 0.0,
        max_depth: int | None = None,
        lambda_value: float = 0.1,
        gamma: float = 1.0,
        nominal_attributes: list | None = None,
        feature_quantizer: tree.splitter.Quantizer | None = None,
    ):
        super().__init__(
            loss_func=BinaryCrossEntropyLoss(),
            delta=delta,
            grace_period=grace_period,
            init_pred=init_pred,
            max_depth=max_depth,
            lambda_value=lambda_value,
            gamma=gamma,
            nominal_attributes=nominal_attributes,
            feature_quantizer=feature_quantizer,
        )

    def _target_transform(self, y):
        return float(y)

    def predict_proba_one(self, x: dict) -> dict[base.typing.ClfTarget, float]:
        if isinstance(self._root, DTBranch):
            leaf = self._root.traverse(x, until_leaf=True)
        else:
            leaf = self._root

        t_proba = self.loss_func.transfer(leaf.prediction())  # type: ignore

        return {True: t_proba, False: 1 - t_proba}


class SGTRegressor(StochasticGradientTree, base.Regressor):
    """Stochastic Gradient Tree for regression.

    Incremental decision tree regressor that minimizes the mean square error to guide its growth.

    Stochastic Gradient Trees (SGT) directly minimize a loss function to guide tree growth and
    update their predictions. Thus, they differ from other incrementally tree learners that do
    not directly optimize the loss, but a data impurity-related heuristic.

    Parameters
    ----------
    delta
        Define the significance level of the F-tests performed to decide upon creating splits
        or updating predictions.
    grace_period
        Interval between split attempts or prediction updates.
    init_pred
        Initial value predicted by the tree.
    max_depth
        The maximum depth the tree might reach. If set to `None`, the trees will grow
        indefinitely.
    lambda_value
        Positive float value used to impose a penalty over the tree's predictions and force
        them to become smaller. The greater the lambda value, the more constrained are the
        predictions.
    gamma
        Positive float value used to impose a penalty over the tree's splits and force them to
        be avoided when possible. The greater the gamma value, the smaller the chance of a
        split occurring.
    nominal_attributes
        List with identifiers of the nominal attributes. If None, all features containing
        numbers are assumed to be numeric.
    feature_quantizer
        The algorithm used to quantize numeric features. Either a static quantizer (as in the
        original implementation) or a dynamic quantizer can be used. The correct choice and setup
        of the feature quantizer is a crucial step to determine the performance of SGTs.
        Feature quantizers are akin to the attribute observers used in Hoeffding Trees. By
        default, an instance of `tree.splitter.StaticQuantizer` (with default parameters) is
        used if this parameter is not set.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import tree

    >>> dataset = datasets.TrumpApproval()
    >>> model = tree.SGTRegressor(
    ...     delta=0.01,
    ...     lambda_value=0.01,
    ...     grace_period=20,
    ...     feature_quantizer=tree.splitter.DynamicQuantizer(std_prop=0.1)
    ... )
    >>> metric = metrics.MAE()

    >>> evaluate.progressive_val_score(dataset, model, metric)
    MAE: 1.721818

    Notes
    -----
    This implementation enhances the original proposal [^1] by using an incremental strategy to
    discretize numerical features dynamically, rather than relying on a calibration set and
    parameterized number of bins. The strategy used is an adaptation of the Quantization Observer
    (QO) [^2]. Different bin size setting policies are available for selection.
    They directly related to number of split candidates the tree is going to explore, and thus,
    how accurate its split decisions are going to be. Besides, the number of stored bins per
    feature is directly related to the tree's memory usage and runtime.

    References
    ----------
    [^1]: Gouk, H., Pfahringer, B., & Frank, E. (2019, October). Stochastic Gradient Trees.
    In Asian Conference on Machine Learning (pp. 1094-1109).
    [^2]: Mastelini, S.M. and de Leon Ferreira, A.C.P., 2021. Using dynamical quantization
    to perform split attempts in online tree regressors. Pattern Recognition Letters.

    """

    def __init__(
        self,
        delta: float = 1e-7,
        grace_period: int = 200,
        init_pred: float = 0.0,
        max_depth: int | None = None,
        lambda_value: float = 0.1,
        gamma: float = 1.0,
        nominal_attributes: list | None = None,
        feature_quantizer: tree.splitter.Quantizer | None = None,
    ):
        super().__init__(
            loss_func=SquaredErrorLoss(),
            delta=delta,
            grace_period=grace_period,
            init_pred=init_pred,
            max_depth=max_depth,
            lambda_value=lambda_value,
            gamma=gamma,
            nominal_attributes=nominal_attributes,
            feature_quantizer=feature_quantizer,
        )

    def predict_one(self, x: dict) -> base.typing.RegTarget:
        if isinstance(self._root, DTBranch):
            leaf = self._root.traverse(x, until_leaf=True)
        else:
            leaf = self._root
        return self.loss_func.transfer(leaf.prediction())  # type: ignore


# river/tree/test_base.py

from __future__ import annotations

from river.tree.base import Branch, Leaf


class BinaryBranch(Branch):
    def __init__(self, left, right, threshold=None, **kwargs):
        super().__init__(left, right)
        self.threshold = threshold
        self.__dict__.update(kwargs)

    def next(self, x):
        if x < self.threshold:
            return self.children[0]
        return self.children[1]

    def most_common_path(self):
        raise NotImplementedError

    @property
    def repr_split(self):
        return f"x < {self.threshold}"


def test_size():
    tree = BinaryBranch(
        BinaryBranch(
            BinaryBranch(BinaryBranch(Leaf(), Leaf()), Leaf()),
            BinaryBranch(Leaf(), Leaf()),
        ),
        BinaryBranch(Leaf(), Leaf()),
    )

    assert tree.n_nodes == tree.n_branches + tree.n_leaves == 6 + 7
    assert (
        tree.children[0].n_nodes == tree.children[0].n_branches + tree.children[0].n_leaves == 4 + 5
    )
    assert (
        tree.children[1].n_nodes == tree.children[1].n_branches + tree.children[1].n_leaves == 1 + 2
    )
    assert (
        tree.children[1].children[0].n_nodes
        == tree.children[1].children[0].n_branches + tree.children[1].children[0].n_leaves
        == 0 + 1
    )


def test_height():
    tree = BinaryBranch(
        BinaryBranch(
            BinaryBranch(
                BinaryBranch(Leaf(), Leaf()),
                Leaf(),
            ),
            BinaryBranch(Leaf(), Leaf()),
        ),
        BinaryBranch(Leaf(), Leaf()),
    )

    assert tree.height == 5
    assert tree.children[0].height == 4
    assert tree.children[1].height == 2
    assert tree.children[1].children[0].height == 1


def test_iter_dfs():
    tree = BinaryBranch(BinaryBranch(Leaf(no=3), Leaf(no=4), no=2), Leaf(no=5), no=1)

    for i, node in enumerate(tree.iter_dfs(), start=1):
        assert i == node.no


def test_iter_bfs():
    tree = BinaryBranch(BinaryBranch(Leaf(no=4), Leaf(no=5), no=2), Leaf(no=3), no=1)

    for i, node in enumerate(tree.iter_bfs(), start=1):
        assert i == node.no


def test_iter_leaves():
    tree = BinaryBranch(BinaryBranch(Leaf(no=1), Leaf(no=2)), Leaf(no=3))

    for i, leaf in enumerate(tree.iter_leaves(), start=1):
        assert i == leaf.no


def test_iter_branches():
    tree = BinaryBranch(
        BinaryBranch(BinaryBranch(Leaf(), Leaf(), no=3), Leaf(), no=2),
        BinaryBranch(Leaf(), Leaf(), no=4),
        no=1,
    )

    for i, branch in enumerate(tree.iter_branches(), start=1):
        assert i == branch.no


def test_iter_edges():
    tree = BinaryBranch(BinaryBranch(Leaf(no=3), Leaf(no=4), no=2), Leaf(no=5), no=1)

    order = [(1, 2), (2, 3), (2, 4), (1, 5)]

    for i, (parent, child) in enumerate(tree.iter_edges()):
        assert order[i] == (parent.no, child.no)


# river/tree/test_splitter.py

from __future__ import annotations

import pytest

from river import datasets, tree
from river.datasets import synth


def get_regression_data():
    return iter(synth.Friedman(seed=42).take(200))


@pytest.mark.parametrize(
    "dataset, splitter",
    [
        (datasets.Phishing(), tree.splitter.ExhaustiveSplitter()),
        (datasets.Phishing(), tree.splitter.HistogramSplitter()),
        (datasets.Phishing(), tree.splitter.GaussianSplitter()),
    ],
)
def test_class_splitter(dataset, splitter):
    model = tree.HoeffdingTreeClassifier(
        splitter=splitter, grace_period=10, leaf_prediction="mc", delta=0.1
    )

    for x, y in dataset:
        model.learn_one(x, y)

    assert model.height > 0


@pytest.mark.parametrize(
    "dataset, splitter",
    [
        (get_regression_data(), tree.splitter.EBSTSplitter()),
        (get_regression_data(), tree.splitter.TEBSTSplitter()),
        (get_regression_data(), tree.splitter.QOSplitter()),
        (get_regression_data(), tree.splitter.QOSplitter(allow_multiway_splits=True)),
    ],
)
def test_reg_splitter(dataset, splitter):
    model = tree.HoeffdingTreeRegressor(
        splitter=splitter, grace_period=20, delta=0.1, leaf_prediction="mean"
    )

    for x, y in dataset:
        model.learn_one(x, y)

    assert model.height > 0


def test_nominal_reg_splitter():
    dataset = synth.Mv(seed=42).take(200)
    model = tree.HoeffdingTreeRegressor(grace_period=10, leaf_prediction="mean")

    for x, y in dataset:
        model.learn_one(x, y)

    assert model.height > 0

    # Evaluates nominal binary splits
    dataset = synth.Mv(seed=42).take(200)
    model = tree.HoeffdingTreeRegressor(grace_period=10, leaf_prediction="mean", binary_split=True)

    for x, y in dataset:
        model.learn_one(x, y)

    assert model.height > 0


# river/tree/test_trees.py

from __future__ import annotations

import random

import pytest

from river import drift, tree
from river.datasets import synth


def get_classification_data():
    return synth.LED(seed=42).take(500)


def get_regression_data():
    return synth.Friedman(seed=42).take(500)


@pytest.mark.parametrize(
    "dataset, model",
    [
        (
            get_classification_data(),
            tree.HoeffdingTreeClassifier(
                leaf_prediction="mc",
                max_size=0.05,
                grace_period=50,
                memory_estimate_period=50,
                splitter=tree.splitter.ExhaustiveSplitter(),
            ),
        ),
        (
            get_classification_data(),
            tree.HoeffdingAdaptiveTreeClassifier(
                leaf_prediction="mc",
                max_size=0.05,
                grace_period=50,
                memory_estimate_period=50,
                splitter=tree.splitter.ExhaustiveSplitter(),
            ),
        ),
        (
            get_classification_data(),
            tree.ExtremelyFastDecisionTreeClassifier(
                leaf_prediction="mc",
                max_size=0.05,
                grace_period=50,
                memory_estimate_period=50,
                splitter=tree.splitter.ExhaustiveSplitter(),
            ),
        ),
    ],
)
def test_memory_usage_class(dataset, model):
    for x, y in dataset:
        model.learn_one(x, y)

    assert model._raw_memory_usage / (2**20) < 0.065


@pytest.mark.parametrize(
    "dataset, model",
    [
        (
            get_regression_data(),
            tree.HoeffdingTreeRegressor(
                leaf_prediction="mean", max_size=0.5, memory_estimate_period=100
            ),
        ),
        (
            get_regression_data(),
            tree.HoeffdingAdaptiveTreeRegressor(
                leaf_prediction="mean", max_size=0.5, memory_estimate_period=100
            ),
        ),
    ],
)
def test_memory_usage_reg(dataset, model):
    for x, y in dataset:
        model.learn_one(x, y)

    assert model._raw_memory_usage / (2**20) < 0.5


def test_memory_usage_multitarget():
    dataset = get_regression_data()

    model = tree.iSOUPTreeRegressor(
        leaf_prediction="mean",
        max_size=0.5,
        memory_estimate_period=100,
    )

    for x, y in dataset:
        # Handcrafted targets
        y_ = {0: y, 1: 2 * y, 2: 3 * y}
        model.learn_one(x, y_)

    assert model._raw_memory_usage / (2**20) < 0.5


def test_efdt_split_reevaluation():
    dataset = synth.SEA(seed=7, variant=2).take(500)

    model = tree.ExtremelyFastDecisionTreeClassifier(
        leaf_prediction="nb",
        grace_period=50,
        min_samples_reevaluate=10,
        split_criterion="hellinger",
        delta=0.1,
    )

    max_depth = -1
    for x, y in dataset:
        model.learn_one(x, y)

        if model.height > max_depth:
            max_depth = model.height

    assert model.height != max_depth


def test_drift_adaptation_hatc():
    rng = random.Random(42)
    dataset = iter(synth.Sine(seed=8, classification_function=0, has_noise=True))

    model = tree.HoeffdingAdaptiveTreeClassifier(
        leaf_prediction="mc",
        grace_period=10,
        drift_detector=drift.ADWIN(0.1),
        delta=0.1,
        drift_window_threshold=2,
        seed=42,
        max_depth=3,
    )

    for i in range(1000):
        if i % 200 == 0 and i > 0:
            dataset = iter(
                synth.Sine(seed=8, classification_function=rng.randint(0, 3), has_noise=False)
            )

        x, y = next(dataset)
        model.learn_one(x, y)

    assert model._n_switch_alternate_trees > 0


def test_drift_adaptation_hatr():
    dataset = synth.Friedman(seed=7).take(1000)

    model = tree.HoeffdingAdaptiveTreeRegressor(
        leaf_prediction="mean",
        grace_period=10,
        delta=0.1,
        drift_detector=drift.ADWIN(0.1),
        drift_window_threshold=10,
        seed=7,
        max_depth=3,
    )

    for i, (x, y) in enumerate(dataset):
        y_ = y
        if i > 500:
            # Emulate an abrupt drift
            y_ = 1.5 * y
        model.learn_one(x, y_)

    assert model._n_alternate_trees > 0


# river/tree/utils.py

from __future__ import annotations

import collections
import copy
import dataclasses
import functools
import math
import sys
import typing

from river.base.typing import FeatureName
from river.stats import Cov, Var


def do_naive_bayes_prediction(x, observed_class_distribution: dict, splitters: dict):
    """Perform Naive Bayes prediction

    Parameters
    ----------
    x
        The feature values.

    observed_class_distribution
        Observed class distribution.

    splitters
        Attribute (features) observers.

    Returns
    -------
    The probabilities related to each class.

    Notes
    -----
    This method is not intended to be used as a stand-alone method.
    """
    total_weight = sum(observed_class_distribution.values())
    if not observed_class_distribution or total_weight == 0:
        # No observed class distributions, all classes equal
        return {}

    votes = {}
    for class_index, class_weight in observed_class_distribution.items():
        # Prior
        if class_weight > 0:
            votes[class_index] = math.log(class_weight / total_weight)
        else:
            votes[class_index] = 0.0
            continue

        if splitters:
            for att_idx in splitters:
                if att_idx not in x:
                    continue
                obs = splitters[att_idx]
                # Prior plus the log likelihood
                tmp = obs.cond_proba(x[att_idx], class_index)
                votes[class_index] += math.log(tmp) if tmp > 0 else 0.0

    # Max log-likelihood
    max_ll = max(votes.values())
    # Apply the log-sum-exp trick (https://stats.stackexchange.com/a/253319)
    lse = max_ll + math.log(sum(math.exp(log_proba - max_ll) for log_proba in votes.values()))

    for class_index in votes:
        votes[class_index] = math.exp(votes[class_index] - lse)

    return votes


@functools.total_ordering
@dataclasses.dataclass
class BranchFactory:
    """Helper class used to assemble branches designed by the splitters.

    If constructed using the default values, a null-split suggestion is assumed.
    """

    merit: float = -math.inf
    feature: FeatureName | None = None
    split_info: typing.Hashable | list[typing.Hashable] | tuple[
        typing.Hashable, list[typing.Hashable]
    ] | None = None
    children_stats: list | None = None
    numerical_feature: bool = True
    multiway_split: bool = False

    def assemble(
        self,
        branch,  # typing.Type[DTBranch],
        stats: dict | Var,
        depth: int,
        *children,
        **kwargs,
    ):
        return branch(stats, self.feature, self.split_info, depth, *children, **kwargs)

    def __lt__(self, other):
        return self.merit < other.merit

    def __eq__(self, other):
        return self.merit == other.merit


class GradHess:
    """The most basic inner structure of the Stochastic Gradient Trees that carries information
    about the gradient and hessian of a given observation.
    """

    __slots__ = ["gradient", "hessian"]

    def __init__(self, gradient: float = 0.0, hessian: float = 0.0, *, grad_hess=None):
        if grad_hess:
            self.gradient = grad_hess.gradient
            self.hessian = grad_hess.hessian
        else:
            self.gradient = gradient
            self.hessian = hessian

    def __iadd__(self, other):
        self.gradient += other.gradient
        self.hessian += other.hessian

        return self

    def __isub__(self, other):
        self.gradient -= other.gradient
        self.hessian -= other.hessian

        return self

    def __add__(self, other):
        new = copy.deepcopy(self)
        new += other
        return new

    def __sub__(self, other):
        new = copy.deepcopy(self)
        new -= other
        return new


@functools.total_ordering
@dataclasses.dataclass
class GradHessMerit:
    """Class used to keep the split merit of each split candidate, accordingly to its
    gradient and hessian information.

    In Stochastic Gradient Trees, the split merit is given by a combination of the loss mean and
    variance. Additionally, the loss in each resulting tree branch is also accounted.
    """

    loss_mean: float = 0.0
    loss_var: float = 0.0
    delta_pred: float | dict | None = None

    def __lt__(self, other):
        return self.loss_mean < other.loss_mean

    def __eq__(self, other):
        return self.loss_mean == other.loss_mean


class GradHessStats:
    """Class used to monitor and update the gradient/hessian information in Stochastic Gradient
    Trees.

    Represents the aggregated gradient/hessian data in a node (global node statistics), category,
    or numerical feature's discretized bin.
    """

    def __init__(self):
        self.g_var = Var()
        self.h_var = Var()
        self.gh_cov = Cov()

    def __iadd__(self, other):
        self.g_var += other.g_var
        self.h_var += other.h_var
        self.gh_cov += other.gh_cov

        return self

    def __isub__(self, other):
        self.g_var -= other.g_var
        self.h_var -= other.h_var
        self.gh_cov -= other.gh_cov

        return self

    def __add__(self, other):
        new = copy.deepcopy(self)
        new += other

        return new

    def __sub__(self, other):
        new = copy.deepcopy(self)
        new -= other

        return new

    def update(self, gh: GradHess, w: float = 1.0):
        self.g_var.update(gh.gradient, w)
        self.h_var.update(gh.hessian, w)
        self.gh_cov.update(gh.gradient, gh.hessian, w)

    @property
    def mean(self) -> GradHess:
        return GradHess(self.g_var.mean.get(), self.h_var.mean.get())

    @property
    def variance(self) -> GradHess:
        return GradHess(self.g_var.get(), self.h_var.get())

    @property
    def covariance(self) -> float:
        return self.gh_cov.get()

    @property
    def total_weight(self) -> float:
        return self.g_var.mean.n

    # This method ignores correlations between delta_pred and the gradients/hessians! Considering
    # delta_pred is derived from the gradient and hessian sample, this assumption is definitely
    # violated. However, as empirically demonstrated in the original SGT, this fact does not seem
    # to significantly impact on the obtained results.
    def delta_loss_mean_var(self, delta_pred: float) -> Var:
        m = self.mean
        n = self.total_weight
        mean = delta_pred * m.gradient + 0.5 * m.hessian * delta_pred * delta_pred

        variance = self.variance
        covariance = self.covariance

        grad_term_var = delta_pred * delta_pred * variance.gradient
        hess_term_var = 0.25 * variance.hessian * (delta_pred**4.0)
        sigma = max(0.0, grad_term_var + hess_term_var + (delta_pred**3) * covariance)
        return Var._from_state(n, mean, sigma)  # noqa


def calculate_object_size(obj: typing.Any, unit: str = "byte") -> int:
    """Iteratively calculates the `obj` size in bytes.

    Visits all the elements related to obj accounting for their respective
    sizes.

    Parameters
    ----------
    obj
        Object to evaluate.
    unit
        The unit in which the accounted value is going to be returned.
        Values: 'byte', 'kB', 'MB' (Default: 'byte').

    Returns
    -------
    The size of the object and its related properties and objects, in 'unit'.

    """
    seen = set()
    to_visit: collections.deque = collections.deque()
    byte_size = 0

    to_visit.append(obj)

    while True:
        try:
            obj = to_visit.popleft()
        except IndexError:
            break

        # If element was already covered, skip it
        if id(obj) in seen:
            continue

        # Update size accounting
        byte_size += sys.getsizeof(obj)

        # Mark element as seen
        seen.add(id(obj))

        # Add keys and values for size account
        if isinstance(obj, dict):
            for v in obj.values():
                to_visit.append(v)

            for k in obj.keys():
                to_visit.append(k)
        elif hasattr(obj, "__dict__"):
            to_visit.append(obj.__dict__)
        elif hasattr(obj, "__iter__") and not (
            isinstance(obj, str) or isinstance(obj, bytes) or isinstance(obj, bytearray)
        ):
            for i in obj:
                to_visit.append(i)

    if unit == "kB":
        final_size = byte_size / 1024
    elif unit == "MB":
        final_size = byte_size / (2**20)
    else:
        final_size = byte_size

    return int(final_size)


def add_dict_values(dict_a: dict, dict_b: dict, inplace=False) -> dict:
    """Add two dictionaries, summing the values of elements with the same key.

    This function iterates over the keys of dict_b and adds their corresponding
    values to the elements in dict_a. If dict_b has a (key, value) pair that
    does not belong to dict_a, this pair is added to the latter dictionary.

    Parameters
    ----------
    dict_a
        dictionary to update.
    dict_b
        dictionary whose values will be added to `dict_a`.
    inplace
        If `True`, the addition is performed in-place and results are stored in `dict_a`.
        If `False`, `dict_a` is not changed and the results are returned in a new dictionary.

    Returns
    -------
    A dictionary containing the result of the operation. Either a pointer to
    `dict_a` or a new dictionary depending on parameter `inplace`.

    """
    if inplace:
        result = dict_a
    else:
        result = copy.deepcopy(dict_a)

    for k, v in dict_b.items():
        try:
            result[k] += v
        except KeyError:
            result[k] = v
    return result


def round_sig_fig(x, significant_digits=2) -> float:
    """Round considering of significant figures of x, given the select
    `significant_digits` prototype.

    If`significant_digits` match the number of significant figures in `x`, its value
    will be used for rounding; otherwise, decimal places will removed
    accordingly to the significant figures in `x`.

    Parameters
    ----------
    x
        A floating point scalar.
    significant_digits
        The number of intended rounding figures.

    Returns
    -------
        The rounded value of `x`.

    Examples
    --------
    >>> round_sig_fig(1.2345)
    1.2
    >>> round_sig_fig(1.2345, significant_digits=3)
    1.23
    >>> round_sig_fig(0.0)
    0.0
    >>> round_sig_fig(0)
    0
    >>> round_sig_fig(1999, significant_digits=1)
    2000
    >>> round_sig_fig(1999, significant_digits=4)
    1999
    >>> round_sig_fig(0.025, significant_digits=3)
    0.03
    >>> round_sig_fig(0.025, significant_digits=10)
    0.025
    >>> round_sig_fig(0.0250, significant_digits=10)
    0.025
    """
    return round(x, significant_digits - int(math.floor(math.log10(abs(x) + 1))) - 1)


# river/tree/viz.py

from __future__ import annotations

from xml.etree import ElementTree as ET

from river.tree.base import Branch, Leaf


def tree_to_html(tree: Branch) -> ET.Element:
    def add_node(node: Branch | Leaf, parent: ET.Element):
        # We're building this:
        #
        # <li>
        #   <pre>Split information</pre>
        #   <ul>...</ul
        # <li>

        li = ET.Element("li")
        parent.append(li)  # type: ignore

        code = ET.Element("code")

        if isinstance(node, Branch):
            code.text = node.repr_split
            li.append(code)

            ul = ET.Element("ul")

            for child in node.children:
                add_node(node=child, parent=ul)

            li.append(ul)

        else:
            code.text = repr(node)
            li.append(code)

    root = ET.Element("ul", attrib={"class": "tree"})
    add_node(node=tree, parent=root)

    return root


CSS = """
.tree,
.tree ul,
.tree li {
    list-style: none;
    margin: 0;
    padding: 0;
    position: relative;
}

.tree {
    margin: 0 0 1em;
    text-align: center;
}

.tree,
.tree ul {
    display: table;
}

.tree ul {
    width: 100%;
}

.tree li {
    display: table-cell;
    padding: .5em 0;
    vertical-align: top;
}

.tree li:before {
    outline: solid 1px #666;
    content: "";
    left: 0;
    position: absolute;
    right: 0;
    top: 0;
}

.tree li:first-child:before {
    left: 50%;
}

.tree li:last-child:before {
    right: 50%;
}

.tree code,
.tree span {
    border: solid .1em #666;
    display: inline-block;
    margin: 0 .2em .5em;
    padding: .2em .5em;
    position: relative;
}

.tree ul:before,
.tree code:before,
.tree span:before {
    outline: solid 1px #666;
    content: "";
    height: .5em;
    left: 50%;
    position: absolute;
}

.tree ul:before {
    top: -.5em;
}

.tree code:before,
.tree span:before {
    top: -.55em;
}

.tree>li {
    margin-top: 0;
}

.tree>li:before,
.tree>li:after,
.tree>li>code:before,
.tree>li>span:before {
    outline: none;
}
"""


# river/utils/__init__.py

"""Shared utility classes and functions"""
from __future__ import annotations

from . import inspect, math, norm, pretty, random
from .context_managers import log_method_calls
from .data_conversion import dict2numpy, numpy2dict
from .param_grid import expand_param_grid
from .rolling import Rolling, TimeRolling
from .sorted_window import SortedWindow
from .vectordict import VectorDict

__all__ = [
    "dict2numpy",
    "expand_param_grid",
    "inspect",
    "log_method_calls",
    "math",
    "pretty",
    "numpy2dict",
    "random",
    "norm",
    "Rolling",
    "SortedWindow",
    "VectorDict",
    "TimeRolling",
]


# river/utils/context_managers.py

from __future__ import annotations

from river.base.base import log_method_calls

__all__ = ["log_method_calls"]


# river/utils/data_conversion.py

from __future__ import annotations

import numpy as np


def dict2numpy(data) -> np.ndarray:
    """Convert a dictionary containing data to a numpy array.

    There is not restriction to the type of keys in `data`, but values must
    be strictly numeric. To make sure random permutations of the features
    do not impact on the learning algorithms, keys are first converted to
    strings and then sorted prior to the conversion.

    Parameters
    ----------
    data
        A dictionary whose keys represent input attributes and the values
        represent their observed contents.

    Returns
    -------
    An array representation of the values in `data`.

    Examples
    --------
    >>> from river.utils import dict2numpy
    >>> dict2numpy({'a': 1, 'b': 2, 3: 3})
    array([3, 1, 2])

    """
    data_ = {str(k): v for k, v in data.items()}
    return np.asarray(list(x for _, x in sorted(data_.items())))


def numpy2dict(data: np.ndarray) -> dict:
    """Convert a numpy array to a dictionary.

    Parameters
    ----------
    data
        An one-dimensional numpy.array.

    Returns
    -------
    A dictionary where keys are integers $k \\in \\left{0, 1, ..., |\\text{data}| - 1\\right}$,
    and the values are each one of the $k$ entries in `data`.

    Examples
    --------
    >>> import numpy as np
    >>> from river.utils import numpy2dict
    >>> numpy2dict(np.array([1.0, 2.0, 3.0]))
    {0: 1.0, 1: 2.0, 2: 3.0}

    """
    return {k: v for k, v in enumerate(data)}


# river/utils/inspect.py

"""Utilities for inspecting a model's type.

Sometimes we need to check if a model can perform regression, classification, etc. However, for
some models the model's type is only known at runtime. For instance, we can't do
`isinstance(pipeline, base.Regressor)` or `isinstance(wrapper, base.Regressor)`. This submodule
thus provides utilities for determining an arbitrary model's type.

"""
from __future__ import annotations

import inspect

from river import base

# TODO: maybe all of this could be done by monkeypatching isintance for pipelines?


__all__ = [
    "extract_relevant",
    "isactivelearner",
    "isanomalydetector",
    "isanomalyfilter",
    "isclassifier",
    "isclusterer",
    "isdriftdetector",
    "ismoclassifier",
    "ismoregressor",
    "isregressor",
    "istransformer",
]


def extract_relevant(model: base.Estimator):
    """Extracts the relevant part of a model.

    Parameters
    ----------
    model

    """

    if ischildobject(obj=model, class_name="Pipeline"):
        return extract_relevant(model._last_step)  # type: ignore[attr-defined]
    return model


def ischildobject(obj: object, class_name: str) -> bool:
    """Checks weather or not the given object inherits from a class with the given class name.

    Workaround isinstance function to not have to import modules defining classes and become
    dependent on them. class_name is case-sensitive.

    Examples
    --------

    >>> from river import anomaly
    >>> from river import utils

    >>> class_name = "AnomalyDetector"

    >>> utils.inspect.ischildobject(obj=anomaly.HalfSpaceTrees(), class_name=class_name)
    True

    >>> utils.inspect.ischildobject(obj=anomaly.OneClassSVM(), class_name=class_name)
    True

    >>> utils.inspect.ischildobject(obj=anomaly.GaussianScorer(), class_name=class_name)
    False

    """
    parent_classes = inspect.getmro(obj.__class__)
    return any(cls.__name__ == class_name for cls in parent_classes)


def isanomalydetector(model):
    model = extract_relevant(model)
    return ischildobject(obj=model, class_name="AnomalyDetector")


def isanomalyfilter(model):
    model = extract_relevant(model)
    return ischildobject(obj=model, class_name="AnomalyFilter")


def isclassifier(model):
    return isinstance(extract_relevant(model), base.Classifier)


def isclusterer(model):
    return isinstance(extract_relevant(model), base.Clusterer)


def ismoclassifier(model):
    return isinstance(extract_relevant(model), base.MultiLabelClassifier)


def isregressor(model):
    return isinstance(extract_relevant(model), base.Regressor)


def istransformer(model):
    return isinstance(extract_relevant(model), base.Transformer)


def ismoregressor(model):
    return isinstance(extract_relevant(model), base.MultiTargetRegressor)


def isdriftdetector(model):
    return isinstance(extract_relevant(model), base.DriftDetector)


def isactivelearner(model):
    from river import active

    return isinstance(extract_relevant(model), active.base.ActiveLearningClassifier)


# river/utils/math.py

"""Mathematical utility functions (intended for internal purposes).

A lot of this is experimental and has a high probability of changing in the future.

"""
from __future__ import annotations

import functools
import itertools
import math
import operator

import numpy as np
import scipy as sp

__all__ = [
    "argmax",
    "chain_dot",
    "clamp",
    "dot",
    "dotvecmat",
    "matmul2d",
    "minkowski_distance",
    "norm",
    "outer",
    "prod",
    "sigmoid",
    "sign",
    "sherman_morrison",
    "softmax",
    "woodbury_matrix",
    "log_sum_2_exp",
]


def dotvecmat(x, A):
    """Vector times matrix from left side, i.e. transpose(x)A.

    Parameters
    ----------
    x
    A

    Examples
    --------

    >>> from river import utils

    >>> x = {0: 4, 1: 5}

    >>> A = {
    ...     (0, 0): 0, (0, 1): 1,
    ...     (1, 0): 2, (1, 1): 3
    ... }

    >>> C = utils.math.dotvecmat(x, A)
    >>> print(C)
    {0: 10.0, 1: 19.0}

    """

    C = {}

    for (i, xi), ((j, k), ai) in itertools.product(x.items(), A.items()):
        if i != j:
            continue

        C[k] = C.get(k, 0.0) + xi * ai

    return C


def matmul2d(A, B):
    """Multiplication for 2D matrices.

    Parameters
    ----------
    A
    B

    Examples
    --------

    >>> import pprint
    >>> from river import utils

    >>> A = {
    ...     (0, 0): 2, (0, 1): 0, (0, 2): 4,
    ...     (1, 0): 5, (1, 1): 6, (1, 2): 0
    ... }

    >>> B = {
    ...     (0, 0): 1, (0, 1): 1, (0, 2): 0, (0, 3): 0,
    ...     (1, 0): 2, (1, 1): 0, (1, 2): 1, (1, 3): 3,
    ...     (2, 0): 4, (2, 1): 0, (2, 2): 0, (2, 3): 0
    ... }

    >>> C = utils.math.matmul2d(A, B)
    >>> pprint.pprint(C)
    {(0, 0): 18.0,
        (0, 1): 2.0,
        (0, 2): 0.0,
        (0, 3): 0.0,
        (1, 0): 17.0,
        (1, 1): 5.0,
        (1, 2): 6.0,
        (1, 3): 18.0}

    """
    C = {}

    for ((i, k1), x), ((k2, j), y) in itertools.product(A.items(), B.items()):
        if k1 != k2:
            continue
        C[i, j] = C.get((i, j), 0.0) + x * y

    return C


def outer(u: dict, v: dict) -> dict:
    """Outer-product between two vectors.

    Parameters
    ----------
    u
    v

    Examples
    --------

    >>> import pprint
    >>> from river import utils

    >>> u = dict(enumerate((1, 2, 3)))
    >>> v = dict(enumerate((2, 4, 8)))

    >>> uTv = utils.math.outer(u, v)
    >>> pprint.pprint(uTv)
    {(0, 0): 2,
        (0, 1): 4,
        (0, 2): 8,
        (1, 0): 4,
        (1, 1): 8,
        (1, 2): 16,
        (2, 0): 6,
        (2, 1): 12,
        (2, 2): 24}

    """
    return {(ki, kj): vi * vj for (ki, vi), (kj, vj) in itertools.product(u.items(), v.items())}


def minkowski_distance(a: dict, b: dict, p: int):
    """Minkowski distance.

    Parameters
    ----------
    a
    b
    p
        Parameter for the Minkowski distance. When `p=1`, this is equivalent to using the
        Manhattan distance. When `p=2`, this is equivalent to using the Euclidean distance.

    """
    return sum((abs(a.get(k, 0.0) - b.get(k, 0.0))) ** p for k in {*a.keys(), *b.keys()})


def softmax(y_pred: dict):
    """Normalizes a dictionary of predicted probabilities, in-place.

    Parameters
    ----------
    y_pred

    """

    if not y_pred:
        return y_pred

    maximum = max(y_pred.values())
    total = 0.0

    for c, p in y_pred.items():
        y_pred[c] = math.exp(p - maximum)
        total += y_pred[c]

    for c in y_pred:
        y_pred[c] /= total

    return y_pred


def prod(iterable):
    """Product function.

    Parameters
    ----------
    iterable

    """
    return functools.reduce(operator.mul, iterable, 1)


def dot(x: dict, y: dict):
    """Returns the dot product of two vectors represented as dicts.

    Parameters
    ----------
    x
    y

    Examples
    --------

    >>> from river import utils

    >>> x = {'x0': 1, 'x1': 2}
    >>> y = {'x1': 21, 'x2': 3}

    >>> utils.math.dot(x, y)
    42

    """

    if len(x) < len(y):
        return sum(xi * y[i] for i, xi in x.items() if i in y)
    return sum(x[i] * yi for i, yi in y.items() if i in x)


def chain_dot(*xs):
    """Returns the dot product of multiple vectors represented as dicts.

    Parameters
    ----------
    xs

    Examples
    --------

    >>> from river import utils

    >>> x = {'x0': 1, 'x1': 2, 'x2': 1}
    >>> y = {'x1': 21, 'x2': 3}
    >>> z = {'x1': 2, 'x2': 1 / 3}

    >>> utils.math.chain_dot(x, y, z)
    85.0

    """
    keys = min(xs, key=len)
    return sum(prod(x.get(i, 0) for x in xs) for i in keys)


def sigmoid(x: float):
    """Sigmoid function.

    Parameters
    ----------
    x

    """
    if x < -30:
        return 0
    if x > 30:
        return 1
    return 1 / (1 + math.exp(-x))


def clamp(x: float, minimum=0.0, maximum=1.0):
    """Clamp a number.

    This is a synonym of clipping.

    Parameters
    ----------
    x
    minimum
    maximum

    """
    return max(min(x, maximum), minimum)


def norm(x: dict, order=None):
    """Compute the norm of a dictionaries values.

    Parameters
    ----------
    x
    order

    """
    return np.linalg.norm(list(x.values()), ord=order)


def sign(x: float):
    """Sign function.

    Parameters
    ----------
    x

    """
    return -1 if x < 0 else (1 if x > 0 else 0)


def argmax(lst: list):
    """Argmax function.

    Parameters
    ----------
    lst

    """
    return max(range(len(lst)), key=lst.__getitem__)


def sherman_morrison(A: np.ndarray, u: np.ndarray, v: np.ndarray):
    """Sherman-Morrison formula.

    This is an inplace function.

    Parameters
    ----------
    A
    u
    v

    References
    ----------
    [^1]: [Fast rank-one updates to matrix inverse? — Tim Vieira](https://timvieira.github.io/blog/post/2021/03/25/fast-rank-one-updates-to-matrix-inverse/)

    """
    Au = A @ u
    alpha = -1 / (1 + v.T @ Au)
    sp.linalg.blas.dger(alpha, Au, v.T @ A, a=A, overwrite_a=1)


def woodbury_matrix(A: np.ndarray, U: np.ndarray, V: np.ndarray):
    """Woodbury matrix identity.

    This is an inplace function.

    Parameters
    ----------
    A
    U
    V

    References
    ----------
    [^1]: [Matrix inverse mini-batch updates — Max Halford](https://maxhalford.github.io/blog/matrix-inverse-mini-batch/)

    """
    eye = np.eye(len(V))
    Au = A @ U
    A -= Au @ np.linalg.inv(eye + V @ Au) @ V @ A


def log_sum_2_exp(a: float, b: float) -> float:
    """Computation of log( (e^a + e^b) / 2) in an overflow-proof way

    Parameters
    ----------
    a
        First number

    b
        Second number
    """
    # TODO: if |a - b| > 50 skip
    # TODO: try several log and exp implementations
    if a > b:
        return a + math.log((1 + math.exp(b - a)) / 2)
    else:
        return b + math.log((1 + math.exp(a - b)) / 2)


# river/utils/norm.py

from __future__ import annotations

import copy
import math

__all__ = ["normalize_values_in_dict", "scale_values_in_dict"]


def normalize_values_in_dict(dictionary, factor=None, inplace=True, raise_error=False):
    """Normalize the values in a dictionary using the given factor.

    For each element in the dictionary, applies `value/factor`.

    Parameters
    ----------
    dictionary
        Dictionary to normalize.
    factor
        Normalization factor value. If not set, use the sum of values.
    inplace
        If True, perform operation in-place
    raise_error
        In case the normalization factor is either `0` or `None`:</br>
        - `True`: raise an error.
        - `False`: return gracefully (if `inplace=False`, a copy of) `dictionary`.

    Raises
    ------
    ValueError
        In case the normalization factor is either `0` or `None` and `raise_error=True`.

    """
    if factor is None:
        factor = sum(dictionary.values())

    if not inplace:
        dictionary = copy.deepcopy(dictionary)

    if factor == 0 or math.isnan(factor):
        # Can not normalize
        if raise_error:
            raise ValueError(f"Can not normalize, normalization factor is {factor}")
        # return gracefully
        return dictionary

    scale_values_in_dict(dictionary, 1 / factor, inplace=True)

    return dictionary


def scale_values_in_dict(dictionary, multiplier, inplace=True):
    """Scale the values in a dictionary.

    For each element in the dictionary, applies `value * multiplier`.

    Parameters
    ----------
    dictionary
        Dictionary to scale.
    multiplier
        Scaling value.
    inplace
        If True, perform operation in-place

    """

    if not inplace:
        dictionary = copy.deepcopy(dictionary)

    for key, value in dictionary.items():
        dictionary[key] = value * multiplier

    return dictionary


# river/utils/param_grid.py

from __future__ import annotations

import itertools
import typing

import numpy as np

from river import base

__all__ = ["expand_param_grid"]


def expand_param_grid(model: base.Estimator, grid: dict) -> list[base.Estimator]:
    """Expands a grid of parameters.

    This method can be used to generate a list of model parametrizations from a dictionary where
    each parameter is associated with a list of possible parameters. In other words, it expands a
    grid of parameters.

    Typically, this method can be used to create copies of a given model with different parameter
    choices. The models can then be used as part of a model selection process, such as a
    `selection.SuccessiveHalvingClassifier` or a `selection.EWARegressor`.

    The syntax for the parameter grid is quite flexible. It allows nesting parameters and can
    therefore be used to generate parameters for a pipeline.

    Parameters
    ----------
    model
    grid
        The grid of parameters to expand. The provided dictionary can be nested. The only
        requirement is that the values at the leaves need to be lists.

    Examples
    --------

    As an initial example, we can expand a grid of parameters for a single model.

    >>> from river import linear_model
    >>> from river import optim
    >>> from river import utils

    >>> model = linear_model.LinearRegression()

    >>> grid = {'optimizer': [optim.SGD(.1), optim.SGD(.01), optim.SGD(.001)]}
    >>> models = utils.expand_param_grid(model, grid)
    >>> len(models)
    3

    >>> models[0]
    LinearRegression (
      optimizer=SGD (
        lr=Constant (
          learning_rate=0.1
        )
      )
      loss=Squared ()
      l2=0.
      l1=0.
      intercept_init=0.
      intercept_lr=Constant (
        learning_rate=0.01
      )
      clip_gradient=1e+12
      initializer=Zeros ()
    )

    You can expand parameters for multiple choices like so:

    >>> grid = {
    ...     'optimizer': [
    ...         (optim.SGD, {'lr': [.1, .01, .001]}),
    ...         (optim.Adam, {'lr': [.1, .01, .01]})
    ...     ]
    ... }
    >>> models = utils.expand_param_grid(model, grid)
    >>> len(models)
    6

    You may specify a grid of parameters for a pipeline via nesting:

    >>> from river import feature_extraction

    >>> model = (
    ...     feature_extraction.BagOfWords() |
    ...     linear_model.LinearRegression()
    ... )

    >>> grid = {
    ...     'BagOfWords': {
    ...         'strip_accents': [False, True]
    ...     },
    ...     'LinearRegression': {
    ...         'optimizer': [
    ...             (optim.SGD, {'lr': [.1, .01]}),
    ...             (optim.Adam, {'lr': [.1, .01]})
    ...         ]
    ...     }
    ... }

    >>> models = utils.expand_param_grid(model, grid)
    >>> len(models)
    8

    """

    return [model.clone(params) for params in _expand_param_grid(grid)]


def _expand_param_grid(grid: dict) -> typing.Iterator[dict]:
    def expand_tuple(t):
        klass, params = t

        if not isinstance(klass, type):
            raise ValueError(f"Expected first element to be a class, got {klass}")

        if not isinstance(params, dict):
            raise ValueError(f"Expected second element to be a dict, got {params}")

        return (klass(**combo) for combo in _expand_param_grid(params))

    def expand(k, v):
        if isinstance(v, tuple):
            return ((k, el) for el in expand_tuple(v))

        # Example:
        # k = 'lr'
        # v = [0.001, 0.01, 0.1]
        if isinstance(v, list) or isinstance(v, set) or isinstance(v, np.ndarray):
            combos = []

            for el in v:
                if isinstance(el, tuple):
                    for combo in expand_tuple(el):
                        combos.append((k, combo))
                else:
                    combos.append((k, el))

            return combos

        if isinstance(v, dict):
            # Example:
            # k = 'LinearRegression'
            # v = {
            #     'intercept_lr': [0.001],
            #     'l2': [1],
            #     'optimizer': {
            #         optim.Adam: {
            #             'beta_1': [0.1, 0.01, 0.001],
            #             'lr': [0.1, 0.01, 0.001]
            #         },
            #      }
            # }
            return ((k, el) for el in _expand_param_grid(v))

        raise ValueError(f"unsupported type: {type(v)}")

    for key in grid:
        if not isinstance(key, str):
            raise ValueError(f"Expected a key of type str; got {key}")

    # Example:
    # grid = {
    #     'beta_1': [0.1, 0.01, 0.001],
    #     'lr': [0.1, 0.01, 0.001]
    # }
    return (
        dict(el) if isinstance(el[0], tuple) else el[0]
        for el in itertools.product(*(expand(k, v) for k, v in grid.items()))
    )


# river/utils/pretty.py

"""Helper functions for making things readable by humans."""

from __future__ import annotations

import math

__all__ = ["humanize_bytes", "print_table"]


def print_table(
    headers: list[str],
    columns: list[list[str]],
    order: list[int] | None = None,
):
    """Pretty-prints a table.

    Parameters
    ----------
    headers
        The column names.
    columns
        The column values.
    order
        Order in which to print the column the values. Defaults to the order in which the values
        are given.

    """

    # Check inputs
    if len(headers) != len(columns):
        raise ValueError("there must be as many headers as columns")

    if len(set(map(len, columns))) > 1:
        raise ValueError("all the columns must be of the same length")

    # Determine the width of each column based on the maximum length of it's elements
    col_widths = [max(*map(len, col), len(header)) for header, col in zip(headers, columns)]

    # Make a template to print out rows one by one
    row_format = " ".join(["{:" + str(width + 2) + "s}" for width in col_widths])

    # Determine the order in which to print the column values
    if order is None:
        order = list(range(len(columns[0])))

    # Build the table
    table = (
        row_format.format(*headers)
        + "\n"
        + "\n".join(
            row_format.format(*[col[i].rjust(width) for col, width in zip(columns, col_widths)])
            for i in order
        )
    )

    return table


def humanize_bytes(n_bytes: int):
    """Returns a human-friendly byte size.

    Parameters
    ----------
    n_bytes

    """
    suffixes = ["B", "KB", "MB", "GB", "TB", "PB"]
    human = float(n_bytes)
    rank = 0
    if n_bytes != 0:
        rank = int((math.log10(n_bytes)) / 3)
        rank = min(rank, len(suffixes) - 1)
        human = n_bytes / (1024.0**rank)
    f = ("%.2f" % human).rstrip("0").rstrip(".")
    return f"{f} {suffixes[rank]}"


# river/utils/random.py

from __future__ import annotations

import math
import random

__all__ = ["poisson"]


def poisson(rate: float, rng=random) -> int:
    """Sample a random value from a Poisson distribution.

    Parameters
    ----------
    rate
    rng

    References
    ----------
    [^1] [Wikipedia article](https://www.wikiwand.com/en/Poisson_distribution#/Generating_Poisson-distributed_random_variables)

    """

    L = math.exp(-rate)
    k = 0
    p = 1

    while p > L:
        k += 1
        p *= rng.random()

    return k - 1


# river/utils/rolling.py

from __future__ import annotations

import bisect
import collections
import datetime as dt
import typing


@typing.runtime_checkable
class Rollable(typing.Protocol):
    def update(self, *args, **kwargs):
        ...

    def revert(self, *args, **kwargs):
        ...


class BaseRolling:
    def __init__(self, obj: Rollable):
        if not isinstance(obj, Rollable):
            raise ValueError(f"{obj} does not satisfy the necessary protocol")

        self.obj = obj

    def __getattribute__(self, name):
        try:
            return super().__getattribute__(name)
        except AttributeError:
            return super().__getattribute__("obj").__getattribute__(name)

    def __getitem__(self, idx):
        return self.obj[idx]

    def __repr__(self):
        return repr(self.obj)


class Rolling(BaseRolling):
    """A generic wrapper for performing rolling computations.

    This can be wrapped around any object which implements both an `update` and a `revert` method.
    Inputs to `update` are stored in a queue. Elements of the queue are popped when the window is
    full.

    Parameters
    ----------
    obj
        An object that implements both an `update` method and a `rolling `method.
    window_size
        Size of the window.

    Examples
    --------

    For instance, here is how you can compute a rolling average over a window of size 3:

    >>> from river import stats, utils

    >>> X = [1, 3, 5, 7]
    >>> rmean = utils.Rolling(stats.Mean(), window_size=3)

    >>> for x in X:
    ...     print(rmean.update(x).get())
    1.0
    2.0
    3.0
    5.0

    """

    def __init__(self, obj: Rollable, window_size: int):
        super().__init__(obj)
        self.window: collections.deque = collections.deque(maxlen=window_size)

    @property
    def window_size(self):
        return self.window.maxlen

    def update(self, *args, **kwargs):
        if len(self.window) == self.window_size:
            self.obj.revert(*self.window[0][0], **self.window[0][1])
        self.obj.update(*args, **kwargs)
        self.window.append((args, kwargs))
        return self


class TimeRolling(BaseRolling):
    """A generic wrapper for performing time rolling computations.

    This can be wrapped around any object which implements both an `update` and a `revert` method.
    Inputs to `update` are stored in a queue. Elements of the queue are popped when they are too
    old.

    Parameters
    ----------
    obj
        An object that implements both an `update` method and a `rolling `method.
    period
        A duration of time, expressed as a `datetime.timedelta`.

    Examples
    --------

    For instance, here is how you can compute a rolling average over a period of 3 days:

    >>> from river import stats, utils

    >>> X = {
    ...     dt.datetime(2019, 1, 1): 1,
    ...     dt.datetime(2019, 1, 2): 5,
    ...     dt.datetime(2019, 1, 3): 9,
    ...     dt.datetime(2019, 1, 4): 13
    ... }

    >>> rmean = utils.TimeRolling(stats.Mean(), period=dt.timedelta(days=3))
    >>> for t, x in X.items():
    ...     print(rmean.update(x, t=t).get())
    1.0
    3.0
    5.0
    9.0

    """

    def __init__(self, obj: Rollable, period: dt.timedelta):
        super().__init__(obj)
        self.period = period
        self._events: list[tuple[dt.datetime, typing.Any]] = []
        self._latest = dt.datetime(1, 1, 1)

    def update(self, *args, t: dt.datetime, **kwargs):
        self.obj.update(*args, **kwargs)
        bisect.insort_left(self._events, (t, (args, kwargs)))

        # There will only be events to revert if the new event if younger than the previously seen
        # youngest event
        if t > self._latest:
            self._latest = t

            i = 0
            for ti, (argsi, kwargsi) in self._events:
                if ti > t - self.period:
                    break
                self.obj.revert(*argsi, **kwargsi)
                i += 1

            # Remove expired events
            if i > 0:
                self._events = self._events[i:]

        return self


# river/utils/sorted_window.py

from __future__ import annotations

import bisect
import collections


class SortedWindow(collections.UserList):
    """Sorted running window data structure.

    Parameters
    ----------
    size
        Size of the window to compute the rolling quantile.

    Examples
    --------

    >>> from river import utils

    >>> window = utils.SortedWindow(size=3)

    >>> for i in reversed(range(9)):
    ...     print(window.append(i))
    [8]
    [7, 8]
    [6, 7, 8]
    [5, 6, 7]
    [4, 5, 6]
    [3, 4, 5]
    [2, 3, 4]
    [1, 2, 3]
    [0, 1, 2]

    References
    ----------
    [^1]: [Left sorted inserts in Python](https://stackoverflow.com/questions/8024571/insert-an-item-into-sorted-list-in-python)

    """

    def __init__(self, size: int):
        super().__init__()
        self.unsorted_window: collections.deque = collections.deque(maxlen=size)

    @property
    def size(self):
        return self.unsorted_window.maxlen

    def append(self, x):
        if len(self) >= self.size:
            # The window is sorted, and a binary search is more optimized than linear search
            start_deque = bisect.bisect_left(self, self.unsorted_window[0])
            del self[start_deque]

        bisect.insort_left(self, x)
        self.unsorted_window.append(x)

        return self


# river/utils/test_math.py

from __future__ import annotations

import numpy as np

from river import utils


def test_dotvecmat_zero_vector_times_matrix_of_ones():
    A_numpy = np.array([[1, 1], [1, 1]])
    A_river = {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}
    x_numpy = np.array([0, 0])
    x_river = {0: 0, 1: 0}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


def test_dotvecmat_vector_of_ones_times_zero_matrix():
    A_numpy = np.array([[0, 0], [0, 0]])
    A_river = {(0, 0): 0, (0, 1): 0, (1, 0): 0, (1, 1): 0}
    x_numpy = np.array([1, 1])
    x_river = {0: 1, 1: 1}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


def test_dotvecmat_filter_first_matrix_row_with_vector():
    A_numpy = np.array([[1, 1], [1, 1]])
    A_river = {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}
    x_numpy = np.array([1, 0])
    x_river = {0: 1, 1: 0}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


def test_dotvecmat_filter_second_matrix_row_with_vector():
    A_numpy = np.array([[1, 1], [1, 1]])
    A_river = {(0, 0): 1, (0, 1): 1, (1, 0): 1, (1, 1): 1}
    x_numpy = np.array([0, 1])
    x_river = {0: 0, 1: 1}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


def test_dotvecmat_vector_times_identity_matrix():
    A_numpy = np.array([[1, 0], [0, 1]])
    A_river = {(0, 0): 1, (0, 1): 0, (1, 0): 0, (1, 1): 1}
    x_numpy = np.array([1, 1])
    x_river = {0: 1, 1: 1}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


def test_dotvecmat_vector_times_anti_diagonal_identity_matrix():
    A_numpy = np.array([[0, 1], [1, 0]])
    A_river = {(0, 0): 0, (0, 1): 1, (1, 0): 1, (1, 1): 0}
    x_numpy = np.array([1, 1])
    x_river = {0: 1, 1: 1}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


def test_dotvecmat_three_dimensional_vector_times_non_quadratic_matrix():
    A_numpy = np.array([[1, 2, 3, 4, 5], [6, 7, 8, 9, 10], [11, 12, 13, 14, 15]])
    A_river = {
        (0, 0): 1,
        (0, 1): 2,
        (0, 2): 3,
        (0, 3): 4,
        (0, 5): 5,
        (1, 0): 6,
        (1, 1): 7,
        (1, 2): 8,
        (1, 3): 9,
        (1, 5): 10,
        (2, 0): 11,
        (2, 1): 12,
        (2, 2): 13,
        (2, 3): 14,
        (2, 5): 15,
    }
    x_numpy = np.array([1, 2, 3])
    x_river = {0: 1, 1: 2, 2: 3}

    numpy_dotvecmat = x_numpy.dot(A_numpy)
    river_dotvecmat = utils.math.dotvecmat(x_river, A_river)

    assert numpy_dotvecmat[0] == river_dotvecmat[0]
    assert numpy_dotvecmat[1] == river_dotvecmat[1]


# river/utils/test_param_grid.py

from __future__ import annotations

import pytest

from river import compose, linear_model, optim, preprocessing, tree, utils


@pytest.mark.parametrize(
    "model, param_grid, count",
    [
        (
            linear_model.LinearRegression(),
            {
                "optimizer": [
                    (optim.SGD, {"lr": [1, 2]}),
                    (
                        optim.Adam,
                        {
                            "beta_1": [0.1, 0.01, 0.001],
                            "lr": [0.1, 0.01, 0.001, 0.0001],
                        },
                    ),
                ]
            },
            2 + 3 * 4,
        ),
        (
            preprocessing.StandardScaler() | linear_model.LinearRegression(),
            {
                "LinearRegression": {
                    "optimizer": [
                        (optim.SGD, {"lr": [1, 2]}),
                        (
                            optim.Adam,
                            {
                                "beta_1": [0.1, 0.01, 0.001],
                                "lr": [0.1, 0.01, 0.001, 0.0001],
                            },
                        ),
                    ]
                }
            },
            2 + 3 * 4,
        ),
        (
            compose.Pipeline(("Scaler", None), linear_model.LinearRegression()),
            {
                "Scaler": [
                    preprocessing.MinMaxScaler(),
                    preprocessing.MaxAbsScaler(),
                    preprocessing.StandardScaler(),
                ],
                "LinearRegression": {"optimizer": {"lr": [1e-1, 1e-2, 1e-3]}},
            },
            3 * 3,
        ),
    ],
)
def test_expand_param_grid_count(model, param_grid, count):
    assert len(utils.expand_param_grid(model, param_grid)) == count


def test_decision_tree_max_depth():
    model = tree.HoeffdingTreeClassifier()

    max_depths = [1, 2, 3, 4, 5, 6]
    models = utils.expand_param_grid(model, {"max_depth": max_depths})

    for model, max_depth in zip(models, max_depths):
        assert model.max_depth == max_depth


# river/utils/test_rolling.py

from __future__ import annotations

import datetime as dt

import pytest

from river import stats, utils


def test_with_counter():
    """
    >>> from river import utils
    >>> import collections
    >>> collections.Counter.revert = collections.Counter.subtract

    >>> counter = utils.Rolling(collections.Counter(), window_size=3)

    >>> for i in range(5):
    ...     counter = counter.update([i])

    >>> counter
    Counter({2: 1, 3: 1, 4: 1, 0: 0, 1: 0})

    >>> counter.most_common(3)
    [(2, 1), (3, 1), (4, 1)]

    >>> counter[4]
    1

    """


def test_rolling_with_not_rollable():
    with pytest.raises(ValueError):
        utils.Rolling(stats.Quantile(), window_size=10)


def test_time_rolling_with_not_rollable():
    with pytest.raises(ValueError):
        utils.TimeRolling(stats.Quantile(), period=dt.timedelta(seconds=10))


# river/utils/test_vectordict.py

from __future__ import annotations

import numpy as np
import pytest

from river.utils import VectorDict


def test_vectordict():
    # test empty init
    x = dict()
    vx = VectorDict()
    assert vx == x

    # test basics
    x = {"a": 8, "b": -1.2, 4: 2.7}
    vx = VectorDict(x)
    assert vx == x
    assert vx["a"] == 8
    assert vx[4] == 2.7
    vx[9] = 8.9
    assert x[9] == vx[9] == 8.9

    # test copy
    x = {"a": 8, "b": -1.2, 4: 2.7}
    vx = VectorDict(x, copy=True)
    assert vx == x
    vx["a"] = 2
    assert x["a"] == 8
    assert vx["a"] == 2

    # test operations
    x = {"a": 1, "b": -5, "c": -3}
    y = {"a": 2, "b": 0.5, "d": 4}
    vx = VectorDict(x)
    vy = VectorDict(y)
    assert vx == vx == x
    assert +vx == vx == x
    assert -vx == {"a": -1, "b": 5, "c": 3}
    assert vx + 2 == 2 + vx == {"a": 3, "b": -3, "c": -1}
    assert vx * 2 == 2 * vx == {"a": 2, "b": -10, "c": -6}
    assert vx - 2 == {"a": -1, "b": -7, "c": -5}
    assert 2 - vx == -(vx - 2)
    assert vx / 2 == {"a": 0.5, "b": -2.5, "c": -1.5}
    assert 2 / vx == {"a": 2, "b": -0.4, "c": -2 / 3}
    assert vx + vy == vy + vx == {"a": 3, "b": -4.5, "c": -3, "d": 4}
    assert vx - vy == {"a": -1, "b": -5.5, "c": -3, "d": -4}
    assert vx**2 == pow(vx, 2) == {"a": 1, "b": 25, "c": 9}
    assert vx * vy == vy * vx == {"a": 2, "b": -2.5, "d": 0, "c": 0}
    assert vx / vx == {"a": 1.0, "b": 1.0, "c": 1.0}
    with pytest.raises(ZeroDivisionError):
        vx / vy
    assert vx @ vy == vy @ vx == -0.5
    vz = VectorDict(x, copy=True)
    vz += 2
    assert vz == vx + 2
    vz = VectorDict(x, copy=True)
    vz -= 2
    assert vz == vx - 2
    vz = VectorDict(x, copy=True)
    vz *= 2
    assert vz == vx * 2
    vz = VectorDict(x, copy=True)
    vz /= 2
    assert vz == vx / 2
    vz = VectorDict(x, copy=True)
    vz += vy
    assert vz == vx + vy
    vz = VectorDict(x, copy=True)
    vz -= vy
    assert vz == vx - vy
    vz = VectorDict(x, copy=True)
    vz *= vy
    assert vz == vx * vy
    vz = VectorDict(x, copy=True)
    vz /= vz
    assert vz == vx / vx
    vz = VectorDict(x, copy=True)
    vz **= 2
    assert vz == vx**2

    # test default_factory
    x = {"a": 1, "b": -5}
    y = {"b": 0.5, "d": 4, "e": 3, "f": 8}
    counter = iter(range(100))
    vx = VectorDict(x, default_factory=counter.__next__)
    vy = VectorDict(y)
    assert vx @ vy == 16.5
    assert counter.__next__() == 3
    assert x["f"] == 2

    # test mask
    x = {"a": 1, "b": -5, "e": 2}
    y = {"b": 0.5, "d": 4, "e": 3, "f": 8}
    z = {"b": 4, "d": 2, "g": -1}
    vx = VectorDict(x)
    vy = VectorDict(y)
    assert vx + vy == vy + vx == {"a": 1, "b": -4.5, "d": 4, "e": 5, "f": 8}
    vy = VectorDict(y, mask=z)
    assert vx + vy == vy + vx == {"a": 1, "b": -4.5, "d": 4, "e": 2}
    vy = VectorDict(y).with_mask(z.keys())
    assert vx + vy == vy + vx == {"a": 1, "b": -4.5, "d": 4, "e": 2}
    vy = VectorDict(y).with_mask(x)
    assert vy / vx == {"b": -0.1, "a": 0.0, "e": 1.5}

    # test export
    x = {"a": 1, "b": -5}
    vx = VectorDict(x)
    nx = vx.to_numpy(["b", "c"])
    assert isinstance(nx, np.ndarray)
    assert (vx.to_numpy(["b", "c"]) == np.array([-5, 0])).all()

    # other methods
    x = {"a": 1, "b": -5}
    vx = VectorDict(x)
    assert vx.abs() == abs(vx) == {"a": 1, "b": 5}
    assert vx.min() == -5
    assert vx.max() == 1
    assert vx.with_mask(["a"]).min() == 1
    assert vx.with_mask(["b"]).max() == -5
    assert vx.minimum(-2) == {"a": -2, "b": -5}
    assert vx.maximum(-2) == {"a": 1, "b": -2}
    y = {"b": 0.5, "c": 4}
    vy = VectorDict(y)
    assert vx.minimum(vy) == vy.minimum(vx) == {"a": 0, "b": -5, "c": 0}
    assert vx.maximum(vy) == vy.maximum(vx) == {"a": 1, "b": 0.5, "c": 4}


# setup.py

from __future__ import annotations

import io
import os
import platform
import subprocess
import sys

import setuptools  # type: ignore
from setuptools_rust import Binding, RustExtension  # type: ignore

try:
    from numpy import __version__ as numpy_version
    from numpy import get_include
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy"])
    from numpy import __version__ as numpy_version
    from numpy import get_include

try:
    from Cython.Build import cythonize
except ImportError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "Cython"])
    from Cython.Build import cythonize  # type: ignore

NAME = "river"
DESCRIPTION = "Online machine learning in Python"
LONG_DESCRIPTION_CONTENT_TYPE = "text/markdown"
URL = "https://github.com/online-ml/river"
EMAIL = "maxhalford25@gmail.com"
AUTHOR = "Max Halford"
REQUIRES_PYTHON = ">=3.8.0"

here = os.path.abspath(os.path.dirname(__file__))
with open(os.path.join(here, "README.md"), encoding="utf-8") as f:
    long_description = "\n" + f.read()

about: dict = {}
with open(os.path.join(here, NAME, "__version__.py")) as f:
    exec(f.read(), about)

setuptools.setup(
    name=NAME,
    version=about["__version__"],
    description=DESCRIPTION,
    long_description=long_description,
    long_description_content_type=LONG_DESCRIPTION_CONTENT_TYPE,
    author=AUTHOR,
    author_email=EMAIL,
    python_requires=REQUIRES_PYTHON,
    url=URL,
    packages=setuptools.find_packages(exclude=("tests",)),
    install_requires=(base_packages := [f"numpy>={numpy_version}", "scipy>=1.5", "pandas>=1.3"]),
    extras_require={
        "dev": base_packages
        + [
            "black>=22.1.0",
            "graphviz>=0.10.1",
            "gym>=0.26.1",
            "matplotlib>=3.0.2",
            "mypy>=0.980",
            "pre-commit>=2.9.2",
            "pytest>=4.5.0",
            "ruff>=0.0.213",
            "scikit-learn>=1.0.1",
            "sqlalchemy>=2.0.0",
            "sympy>=1.10.1",
        ],
        "compat": base_packages
        + [
            "scikit-learn",
            "sqlalchemy>=2.0.0",
            "vaex",
        ],
        "docs": base_packages
        + [
            "dominate",
            "flask",
            "ipykernel",
            "jupyter-client",
            "mike",
            "mkdocs",
            "mkdocs-awesome-pages-plugin",
            "mkdocs-charts-plugin",
            "mkdocs-material",
            "nbconvert",
            "numpydoc",
            "python-slugify",
            "spacy",
            "tabulate",
            "watermark",
        ],
        "extra": ["river_extra"],
        "deep": ["deep-river"],
        "torch": ["river_torch"],
        "benchmark": base_packages
        + [
            "scikit-learn==1.1.2",
            "scipy==1.9.3",
            "torch==1.13.0",
            "vowpalwabbit==9.6.0",
            "torch==1.13.0",
            "dominate==2.7.0",
            "slugify==0.0.1",
            "watermark==2.3.1",
            "river-torch==0.1.2",
        ],
        ":python_version == '3.6'": ["dataclasses"],
    },
    include_package_data=True,
    license="BSD-3",
    classifiers=[
        # Trove classifiers
        # Full list: https://pypi.python.org/pypi?%3Aaction=list_classifiers
        "License :: OSI Approved :: BSD License",
        "Programming Language :: Python",
        "Programming Language :: Python :: 3",
        "Programming Language :: Python :: 3.8",
        "Programming Language :: Python :: 3.9",
        "Programming Language :: Python :: 3.10",
        "Programming Language :: Python :: Implementation :: CPython",
        "Programming Language :: Python :: Implementation :: PyPy",
    ],
    ext_modules=cythonize(
        module_list=[
            setuptools.Extension(
                "*",
                sources=["**/*.pyx"],
                include_dirs=[get_include()],
                libraries=[] if platform.system() == "Windows" else ["m"],
                define_macros=[("NPY_NO_DEPRECATED_API", "NPY_1_7_API_VERSION")],
            )
        ],
        compiler_directives={
            "language_level": 3,
            "binding": True,
            "embedsignature": True,
        },
    ),
    rust_extensions=[RustExtension("river.stats._rust_stats", binding=Binding.PyO3)],
    # rust extensions are not zip safe, just like C-extensions.
    zip_safe=False,
)


# river/__init__.py

"""River is a library for incremental learning. Incremental learning is a machine learning regime
where the observations are made available one by one. It is also known as online learning,
iterative learning, or sequential learning. This is in contrast to batch learning where all the
data is processed at once. Incremental learning is desirable when the data is too big to fit in
memory, or simply when it isn't available all at once. river's API is heavily inspired from that of
scikit-learn, enough so that users who are familiar with scikit-learn should feel right at home.
"""
from __future__ import annotations

from .__version__ import __version__  # noqa: F401


# river/__version__.py

from __future__ import annotations

VERSION = (0, 17, 0)

__version__ = ".".join(map(str, VERSION))  # noqa: F401


# river/api.py

"""River API module."""
from __future__ import annotations

from . import (
    active,
    anomaly,
    bandit,
    base,
    cluster,
    compat,
    compose,
    conf,
    covariance,
    datasets,
    drift,
    dummy,
    ensemble,
    evaluate,
    facto,
    feature_extraction,
    feature_selection,
    forest,
    imblearn,
    linear_model,
    metrics,
    misc,
    model_selection,
    multiclass,
    multioutput,
    naive_bayes,
    neighbors,
    neural_net,
    optim,
    preprocessing,
    proba,
    reco,
    rules,
    sketch,
    stats,
    stream,
    time_series,
    tree,
    utils,
)

__all__ = [
    "active",
    "anomaly",
    "base",
    "bandit",
    "cluster",
    "compat",
    "compose",
    "conf",
    "covariance",
    "datasets",
    "dummy",
    "drift",
    "ensemble",
    "evaluate",
    "facto",
    "feature_extraction",
    "feature_selection",
    "forest",
    "imblearn",
    "linear_model",
    "metrics",
    "misc",
    "model_selection",
    "multiclass",
    "multioutput",
    "naive_bayes",
    "neighbors",
    "neural_net",
    "optim",
    "preprocessing",
    "proba",
    "reco",
    "rules",
    "sketch",
    "stats",
    "stream",
    "time_series",
    "tree",
    "utils",
]


# river/conftest.py

from __future__ import annotations

collect_ignore = []

try:
    import sklearn  # noqa: F401
except ImportError:
    collect_ignore.append("compat/sklearn.py")
    collect_ignore.append("compat/test_sklearn.py")

try:
    import sqlalchemy  # noqa: F401
except ImportError:
    collect_ignore.append("stream/iter_sql.py")
    collect_ignore.append("stream/test_sql.py")

try:
    import torch  # noqa: F401
except ImportError:
    collect_ignore.append("compat/pytorch.py")

try:
    import vaex  # noqa: F401
except ImportError:
    collect_ignore.append("stream/iter_vaex.py")


# river/dummy.py

"""Dummy estimators.

This module is here for testing purposes, as well as providing baseline performances.

"""
from __future__ import annotations

import collections

from river import base, stats

__all__ = ["NoChangeClassifier", "PriorClassifier", "StatisticRegressor"]


class NoChangeClassifier(base.Classifier):
    """Dummy classifier which returns the last class seen.

    The predict_one method will output the last class seen whilst predict_proba_one will
    return 1 for the last class seen and 0 for the others.

    Attributes
    ----------
    last_class
        The last class seen.
    classes
        The set of classes seen.

    Examples
    --------

    Taken from example 2.1 from
    [this page](https://www.cms.waikato.ac.nz/~abifet/book/chapter_2.html).

    >>> import pprint
    >>> from river import dummy

    >>> sentences = [
    ...     ('glad happy glad', '+'),
    ...     ('glad glad joyful', '+'),
    ...     ('glad pleasant', '+'),
    ...     ('miserable sad glad', '−')
    ... ]

    >>> model = dummy.NoChangeClassifier()

    >>> for sentence, label in sentences:
    ...     model = model.learn_one(sentence, label)

    >>> new_sentence = 'glad sad miserable pleasant glad'
    >>> model.predict_one(new_sentence)
    '−'

    >>> pprint.pprint(model.predict_proba_one(new_sentence))
    {'+': 0, '−': 1}

    """

    def __init__(self):
        self.last_class = None
        self.classes = set()

    @property
    def _multiclass(self):
        return True

    def learn_one(self, x, y):
        self.last_class = y
        self.classes.add(y)
        return self

    def predict_one(self, x):
        return self.last_class

    def predict_proba_one(self, x):
        probas = {c: 0 for c in self.classes}
        probas[self.last_class] = 1
        return probas


class PriorClassifier(base.Classifier):
    """Dummy classifier which uses the prior distribution.

    The `predict_one` method will output the most common class whilst `predict_proba_one` will
    return the normalized class counts.

    Attributes
    ----------
    counts : collections.Counter
        Class counts.
    n : int
        Total number of seen instances.

    Examples
    --------

    Taken from example 2.1 from
    [this page](https://www.cms.waikato.ac.nz/~abifet/book/chapter_2.html)

    >>> from river import dummy

    >>> sentences = [
    ...     ('glad happy glad', '+'),
    ...     ('glad glad joyful', '+'),
    ...     ('glad pleasant', '+'),
    ...     ('miserable sad glad', '−')
    ... ]

    >>> model = dummy.PriorClassifier()

    >>> for sentence, label in sentences:
    ...     model = model.learn_one(sentence, label)

    >>> new_sentence = 'glad sad miserable pleasant glad'
    >>> model.predict_one(new_sentence)
    '+'
    >>> model.predict_proba_one(new_sentence)
    {'+': 0.75, '−': 0.25}

    References
    ----------
    [^1]: [Krichevsky–Trofimov estimator](https://www.wikiwand.com/en/Krichevsky%E2%80%93Trofimov_estimator)

    """

    def __init__(self):
        self.counts = collections.Counter()
        self.n = 0

    @property
    def _multiclass(self):
        return True

    def learn_one(self, x, y):
        self.counts.update([y])
        self.n += 1
        return self

    def predict_proba_one(self, x):
        return {label: count / self.n for label, count in self.counts.items()}


class StatisticRegressor(base.Regressor):
    """Dummy regressor that uses a univariate statistic to make predictions.

    Parameters
    ----------
    statistic

    Examples
    --------

    >>> from pprint import pprint
    >>> from river import dummy
    >>> from river import stats

    >>> sentences = [
    ...     ('glad happy glad', 3),
    ...     ('glad glad joyful', 3),
    ...     ('glad pleasant', 2),
    ...     ('miserable sad glad', -3)
    ... ]

    >>> model = dummy.StatisticRegressor(stats.Mean())

    >>> for sentence, score in sentences:
    ...     model = model.learn_one(sentence, score)

    >>> new_sentence = 'glad sad miserable pleasant glad'
    >>> model.predict_one(new_sentence)
    1.25

    """

    def __init__(self, statistic: stats.base.Univariate):
        self.statistic = statistic

    @classmethod
    def _unit_test_params(cls):
        yield {"statistic": stats.Mean()}

    def learn_one(self, x, y):
        self.statistic.update(y)
        return self

    def predict_one(self, x):
        return self.statistic.get()


# river/test_estimators.py

"""General tests that all estimators need to pass."""
from __future__ import annotations

import importlib
import inspect

import pytest

from river import (
    anomaly,
    base,
    checks,
    compose,
    facto,
    feature_extraction,
    feature_selection,
    imblearn,
    linear_model,
    model_selection,
    multiclass,
    neighbors,
    neural_net,
    preprocessing,
    time_series,
)

try:
    from river.compat.pytorch import PyTorch2RiverBase

    PYTORCH_INSTALLED = True
except ImportError:
    PYTORCH_INSTALLED = False
from river.compat.river_to_sklearn import River2SKLBase
from river.compat.sklearn_to_river import SKL2RiverBase


def iter_estimators():
    for submodule in importlib.import_module("river.api").__all__:

        def is_estimator(obj):
            return inspect.isclass(obj) and issubclass(obj, base.Estimator)

        for _, obj in inspect.getmembers(
            importlib.import_module(f"river.{submodule}"), is_estimator
        ):
            yield obj


def iter_estimators_which_can_be_tested():
    ignored = (
        River2SKLBase,
        SKL2RiverBase,
        compose.FuncTransformer,
        compose.Grouper,
        compose.Pipeline,
        compose.Prefixer,
        compose.Renamer,
        compose.Suffixer,
        compose.TargetTransformRegressor,
        facto.FFMClassifier,
        facto.FFMRegressor,
        facto.FMClassifier,
        facto.FMRegressor,
        facto.FwFMClassifier,
        facto.FwFMRegressor,
        facto.HOFMClassifier,
        facto.HOFMRegressor,
        feature_extraction.Agg,
        feature_extraction.TargetAgg,
        feature_selection.PoissonInclusion,
        imblearn.RandomOverSampler,
        imblearn.RandomUnderSampler,
        imblearn.RandomSampler,
        model_selection.SuccessiveHalvingClassifier,
        neighbors.NearestNeighbors,
        neural_net.MLPRegressor,
        preprocessing.PreviousImputer,
        preprocessing.OneHotEncoder,
        preprocessing.StatImputer,
        time_series.base.Forecaster,
    )

    if PYTORCH_INSTALLED:
        ignored = (*ignored, PyTorch2RiverBase)

    def can_be_tested(estimator):
        return not inspect.isabstract(estimator) and not issubclass(estimator, ignored)

    for estimator in filter(can_be_tested, iter_estimators()):
        for params in estimator._unit_test_params():
            yield estimator(**params)


@pytest.mark.parametrize(
    "estimator, check",
    [
        pytest.param(
            estimator,
            check,
            id=f"{estimator}:{check.__name__}",
        )
        for estimator in list(iter_estimators_which_can_be_tested())
        + [
            preprocessing.StandardScaler() | linear_model.LinearRegression(),
            preprocessing.StandardScaler() | linear_model.PAClassifier(),
            (
                preprocessing.StandardScaler()
                | preprocessing.TargetStandardScaler(
                    regressor=linear_model.LinearRegression(),
                )
            ),
            (
                preprocessing.StandardScaler()
                | multiclass.OneVsRestClassifier(linear_model.LogisticRegression())
            ),
            (
                preprocessing.StandardScaler()
                | multiclass.OneVsRestClassifier(linear_model.PAClassifier())
            ),
            preprocessing.MinMaxScaler() + preprocessing.StandardScaler(),
            (
                feature_extraction.PolynomialExtender()
                | preprocessing.StandardScaler()
                | linear_model.LinearRegression()
            ),
            preprocessing.MinMaxScaler() | anomaly.HalfSpaceTrees(),
        ]
        for check in checks.yield_checks(estimator)
        if check.__name__ not in estimator._unit_test_skips()
    ],
)
def test_check_estimator(estimator, check):
    check(estimator.clone())


# river/neighbors/__init__.py

"""Neighbors-based learning.

Also known as *lazy* methods. In these methods, generalisation of the training data is delayed
until a query is received.

"""
from __future__ import annotations

from .base import NearestNeighbors
from .knn_classifier import KNNClassifier
from .knn_regressor import KNNRegressor

__all__ = [
    "NearestNeighbors",
    "KNNClassifier",
    "KNNRegressor",
]


# river/neighbors/base.py

from __future__ import annotations

import collections
import operator
import typing


class DistanceFunc(typing.Protocol):
    def __call__(self, a: typing.Any, b: typing.Any, **kwargs) -> float:
        ...


class FunctionWrapper:
    """Wrapper used to make distance function work with KNNClassifier and
    KNNRegressor.

    The k-NN-based classifier and regressor store tuples with `(x, y)`, but only
    `x` is used for distance calculations. This wrapper makes sure `x` is accessed
    when calculating the distances.

    Parameters
    ----------
    distance_function
        The custom distance function to be wrapped.

    """

    def __init__(self, distance_function: DistanceFunc):
        self.distance_function = distance_function

    def __call__(self, a, b):
        # Access x, which is stored in a tuple (x, y)
        return self.distance_function(a[0], b[0])


class NearestNeighbors:
    """Exact nearest neighbors search data structure.

    Parameters
    ----------
    window_size
        Size of the sliding window use to search neighbors with.
    min_distance_keep
        The minimum distance (similarity) to consider adding a point to the window.
        E.g., a value of 0.0 will add even exact duplicates.
    distance_func
        A distance function which accepts two input items to compare.

    Notes
    -----
    Updates are by default stored by the FIFO (first in first out) method,
    which means that when the size limit is reached, old samples are dumped to
    give room for new samples. This is circular, meaning that older points
    are dumped first. This also gives the implementation a temporal aspect,
    because older samples are replaced with newer ones.

    The parameter `min_dinstance_keep` controls the addition of new items to the
    window - items that are far enough away (> min_distance_keep) are added to
    the window. Thus a value of 0 indicates that we add all points, and
    increasing from 0 makes it less likely we will keep a new item.

    """

    def __init__(
        self,
        window_size: int,
        min_distance_keep: float,
        distance_func: DistanceFunc | FunctionWrapper,
    ):
        self.window_size = window_size

        # A minimum distance (similarity) to determine adding to window
        # The model will perform better with a more diverse window
        # Since the distance function can be anything, it could be < 0
        self.min_distance_keep = min_distance_keep

        self.distance_func = distance_func
        self.window: collections.deque = collections.deque(maxlen=self.window_size)

    def append(self, item: typing.Any, extra: typing.Any | None = None):
        """Add a point to the window, optionally with extra metadata.

        Parameters
        ----------
        item
            The data intended to be provided to the distance function. It is always
            the first item in the window, and typically this will be a tuple
            (x,y) with features `x` and class or value `y`.
        extra:
            An extra set of metadata to add to the window that is not passed to
            the distance function, and allows easy customization without needing
            to always write a custom distance function.

        """
        self.window.append((item, *(extra or [])))

    def update(
        self,
        item: typing.Any,
        n_neighbors: int = 1,
        extra: typing.Any | None = None,
    ):
        """Update the window with a new point, only added if > min distance.

        If min distance is 0, we do not need to do the calculation. The item
        (and extra metadata) will not be added to the window if it is too close
        to an existing point.

        Parameters
        ----------
        item
            The data intended to be provided to the distance function.
        extra
            Metadata that is separate from the item that should also be added
            to the window, but is not included to be passed to the distance
            function.

        Returns
        -------
        A boolean (true/false) to indicate if the point was added.

        """
        # If min distance is 0, we add all points
        if self.min_distance_keep == 0:
            self.append(item, extra=extra)
            return True

        # Don't add VERY similar points to window
        nearest = self.find_nearest(item, n_neighbors)

        # Distance always in the last index, (item <extra> distance)
        if not nearest or nearest[0][-1] < self.min_distance_keep:
            self.append(item, extra=extra)
            return True
        return False

    def find_nearest(self, item: typing.Any, n_neighbors: int = 1):
        """Find the `n_neighbors` closest points to `item`, along with their distances."""
        # Compute the distances to each point in the window
        # The window is (item, <extra>, distance)
        points = ((*p, self.distance_func(item, p[0])) for p in self.window)

        # Return the k closest points (last index is distance)
        return sorted(points, key=operator.itemgetter(-1))[:n_neighbors]


# river/neighbors/knn_classifier.py

from __future__ import annotations

import functools

from river import base, utils
from river.neighbors import NearestNeighbors
from river.neighbors.base import DistanceFunc, FunctionWrapper


class KNNClassifier(base.Classifier):
    """K-Nearest Neighbors (KNN) for classification.

    This works by storing a buffer with the `window_size` most recent observations.
    A brute-force search is used to find the `n_neighbors` nearest observations
    in the buffer to make a prediction. See the NearestNeighbors parent class for more
    details.

    Parameters
    ----------
    n_neighbors
        The number of nearest neighbors to search for.
    window_size
        The maximum size of the window storing the last observed samples.
    min_distance_keep
        The minimum distance (similarity) to consider adding a point to the window.
        E.g., a value of 0.0 will add even exact duplicates. Default is 0.05 to add
        similar but not exactly the same points.
    weighted
        Weight the contribution of each neighbor by it's inverse distance.
    cleanup_every
        This determines at which rate old classes are cleaned up. Classes that
        have been seen in the past but that are not present in the current
        window are dropped. Classes are never dropped when this is set to 0.
    distance_func
        An optional distance function that should accept an a=, b=, and any
        custom set of kwargs. If not defined, the Minkowski distance is used with
        p=2 (Euclidean distance). See the example section for more details.
    softmax
        Whether or not to use softmax normalization to normalize the neighbors contributions.
        Votes are divided by the total number of votes if this is `False`.

    Notes
    -----
    Note that since the window is moving and we keep track of all classes that
    are added at some point, a class might be returned in a result (with a
    value of 0) if it is no longer in the window. You can call
    model.clean_up_classes(), or set `cleanup_every` to a non-zero value.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import neighbors
    >>> from river import preprocessing

    >>> dataset = datasets.Phishing()

    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     neighbors.KNNClassifier(window_size=50)
    ... )

    >>> evaluate.progressive_val_score(dataset, model, metrics.Accuracy())
    Accuracy: 84.55%

    When defining a custom distance function you can rely on `functools.partial` to set default
    parameter values. For instance, let's use the Manhattan function instead of the default Euclidean distance:

    >>> import functools
    >>> from river import utils
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     neighbors.KNNClassifier(
    ...         window_size=50,
    ...         distance_func=functools.partial(utils.math.minkowski_distance, p=1)
    ...     )
    ... )
    >>> evaluate.progressive_val_score(dataset, model, metrics.Accuracy())
    Accuracy: 86.87%

    """

    def __init__(
        self,
        n_neighbors: int = 5,
        window_size: int = 1000,
        min_distance_keep: float = 0.0,
        weighted: bool = True,
        cleanup_every: int = 0,
        distance_func: DistanceFunc | None = None,
        softmax: bool = False,
    ):
        self.n_neighbors = n_neighbors
        self.window_size = window_size
        self.min_distance_keep = min_distance_keep
        self.distance_func = (
            distance_func
            if distance_func is not None
            else functools.partial(utils.math.minkowski_distance, p=2)
        )

        self.weighted = weighted
        self.cleanup_every = cleanup_every
        self.classes: set[base.typing.ClfTarget] = set()
        self.softmax = softmax
        self._cleanup_counter = cleanup_every

        self._nn = NearestNeighbors(
            window_size=self.window_size,
            min_distance_keep=min_distance_keep,
            distance_func=FunctionWrapper(self.distance_func),
        )

    @property
    def _multiclass(self):
        return True

    @classmethod
    def _unit_test_params(cls):
        yield {"n_neighbors": 3, "window_size": 30}

    def clean_up_classes(self):
        """Clean up classes added to the window.

        Classes that are added (and removed) from the window may no longer be valid.
        This method cleans up the window and and ensures only known classes
        are added, and we do not consider "None" a class. It is called every
        `cleanup_every` step, or can be called manually.

        """
        self.classes = {x for x in self.window if x[0][1] is not None}

    def learn_one(self, x, y):
        # Only add the class y to known classes if we actually add the point!
        if self._nn.update((x, y), n_neighbors=self.n_neighbors):
            self.classes.add(y)

        # Ensure classes known to instance reflect window
        self._run_class_cleanup()
        return self

    def _run_class_cleanup(self):
        """Helper function to run class cleanup, accounting for _cleanup_counter."""
        # clean up classes every cleanup_every steps
        if self.cleanup_every:
            self._cleanup_counter -= 1
            if self._cleanup_counter == 0:
                self.clean_up_classes()
                self._cleanup_counter = self.cleanup_every

        return self

    def predict_proba_one(self, x):
        nearest = self._nn.find_nearest((x, None), n_neighbors=self.n_neighbors)

        # Default prediction for every class we know is 0.
        # If class_cleanup is false this can include classes not in window
        y_pred = {c: 0.0 for c in self.classes}

        # No nearest points? Return the default (normalized)
        # Note that normalization otherwise happens at the end
        if not nearest:
            default_pred = 1 / len(self.classes) if self.classes else 0.0
            return {c: default_pred for c in self.classes}

        # If the closest is an exact match AND has a class, return it
        if nearest[0][-1] == 0 and nearest[0][0][1] is not None:
            # Update the class in our prediction from 0 to 1, 100% certain!
            y_pred[nearest[0][0][1]] = 1.0
            return y_pred

        for neighbor in nearest:
            (x, y), distance = neighbor

            # Weighted votes by inverse distance
            if self.weighted:
                y_pred[y] += 1.0 / distance

            # Uniform votes
            else:
                y_pred[y] += 1.0

        # Normalize votes into real [0, 1] probabilities
        if self.softmax:
            return utils.math.softmax(y_pred)

        # Otherwise normalize by the total sum
        total = sum(y_pred.values())
        for y in y_pred:
            y_pred[y] /= total
        return y_pred


# river/neighbors/knn_regressor.py

from __future__ import annotations

import functools
import statistics

from river import base, utils
from river.neighbors import NearestNeighbors
from river.neighbors.base import DistanceFunc, FunctionWrapper


class KNNRegressor(base.Regressor):
    """K-Nearest Neighbors regressor.

    This non-parametric regression method keeps track of the last `window_size`
    training samples. Predictions are obtained by aggregating the values of the
    closest n_neighbors stored samples with respect to a query sample.

    Parameters
    ----------
    n_neighbors
        The number of nearest neighbors to search for.
    window_size
        The maximum size of the window storing the last observed samples.
    aggregation_method
        The method to aggregate the target values of neighbors.
            | 'mean'
            | 'median'
            | 'weighted_mean'
    min_distance_keep
        The minimum distance (similarity) to consider adding a point to the window.
        E.g., a value of 0.0 will add even exact duplicates.
    distance_func
        An optional distance function that should accept an a=, b=, and any
        custom set of kwargs. If not defined, the Minkowski distance is used with
        p=2 (Euclidean distance). See the example section for more details.

    Examples
    --------
    >>> from river import datasets
    >>> from river import evaluate
    >>> from river import metrics
    >>> from river import neighbors
    >>> from river import preprocessing

    >>> dataset = datasets.TrumpApproval()

    >>> model = neighbors.KNNRegressor(window_size=50)
    >>> evaluate.progressive_val_score(dataset, model, metrics.RMSE())
    RMSE: 1.427746

    When defining a custom distance function you can rely on `functools.partial` to set default
    parameter values. For instance, let's use the Manhattan function instead of the default Euclidean distance:

    >>> import functools
    >>> from river.utils.math import minkowski_distance
    >>> model = (
    ...     preprocessing.StandardScaler() |
    ...     neighbors.KNNRegressor(
    ...         window_size=50,
    ...         distance_func=functools.partial(minkowski_distance, p=1)
    ...     )
    ... )
    >>> evaluate.progressive_val_score(dataset, model, metrics.RMSE())
    RMSE: 1.460385

    """

    _MEAN = "mean"
    _MEDIAN = "median"
    _WEIGHTED_MEAN = "weighted_mean"

    def __init__(
        self,
        n_neighbors: int = 5,
        window_size: int = 1000,
        aggregation_method: str = "mean",
        min_distance_keep: float = 0.0,
        distance_func: DistanceFunc | None = None,
    ):
        self.n_neighbors = n_neighbors
        self.window_size = window_size
        self.min_distance_keep = min_distance_keep
        self.distance_func = (
            distance_func
            if distance_func is not None
            else functools.partial(utils.math.minkowski_distance, p=2)
        )

        self._nn = NearestNeighbors(
            window_size=self.window_size,
            min_distance_keep=min_distance_keep,
            distance_func=FunctionWrapper(self.distance_func),
        )

        self._check_aggregation_method(aggregation_method)
        self.aggregation_method = aggregation_method

    def _check_aggregation_method(self, method):
        """Ensure validation method is known to the model.

        Raises a ValueError if not.

        Parameters
        ----------
        method
            The suplied aggregration method.
        """
        if method not in {self._MEAN, self._MEDIAN, self._WEIGHTED_MEAN}:
            raise ValueError(
                "Invalid aggregation_method: {}.\n"
                "Valid options are: {}".format(
                    method, {self._MEAN, self._MEDIAN, self._WEIGHTED_MEAN}
                )
            )

    def learn_one(self, x, y):
        self._nn.update((x, y), n_neighbors=self.n_neighbors)
        return self

    def predict_one(self, x):
        # Find the nearest neighbors!
        nearest = self._nn.find_nearest((x, None), n_neighbors=self.n_neighbors)

        if not nearest:
            return 0.0

        # For each in nearest, call it 'item"
        # item[0] is the original item (x, y)
        # item[-1] is the distance

        # If the closest distance is 0 (it's the same) return it's output (y)
        # BUT only if the output (y) is not None.
        if nearest[0][-1] == 0 and nearest[0][0][1] is not None:
            return nearest[0][0][1]

        # Only include neighbors in the sum that are not None
        neighbor_vals = [n[0][1] for n in nearest if n[0][1] is not None]

        if self.aggregation_method == self._MEDIAN:
            return statistics.median(neighbor_vals)

        dists = [n[-1] for n in nearest if n[0][1] is not None]
        sum_ = sum(1 / d for d in dists)

        if self.aggregation_method == self._MEAN or sum_ == 0.0:
            return statistics.mean(neighbor_vals)

        # weighted mean based on distance
        return sum(y / d for y, d in zip(neighbor_vals, dists)) / sum_


# river/neighbors/ann/swinn.py
from __future__ import annotations

import collections
import functools
import heapq
import itertools
import math
import operator
import random
import typing

from river import utils
from river.neighbors.base import BaseNN, DistanceFunc, FunctionWrapper

from .nn_vertex import Vertex


class SWINN(BaseNN):
    """Sliding WIndow-based Nearest Neighbor (SWINN) search using Graphs.

    Extends the NNDescent algorithm[^1] to handle vertex addition and removal in a FIFO data
    ingestion policy. SWINN builds and keeps a directed graph where edges connect the nearest
    neighbors. Any distance metric can be used to build the graph. By using a directed graph,
    the user must set the desired number of neighbors. More neighbors imply more accurate
    search queries at the cost of increased running time and memory usage. Note that although
    the number of directed neighbors is limited by the user, there is no direct control on the
    number of reverse neighbors, i.e., the number of vertices that have an edge to a given vertex.

    The basic idea of SWINN and NNDescent is that "the neighbor of my neighbors might as well be
    my neighbor". Hence, the connections are constantly revisited to improve the graph structure.
    The algorithm for creating and maintaining the search graph can be described
    in general lines as follows:

    * Start with a random neighborhood graph;

    * For each node in the search graph: refine the current neighborhood by checking if there
    are better neighborhood options among the neighbors of the current neighbors;

    * If the total number of neighborhood changes is smaller than a given stopping criterion, then stop.

    SWINN adds strategies to remove vertices from the search graph and pruning redundant edges. SWINN is
    more efficient when the selected `maxlen` is greater than 500. For small sized data windows, using
    the lazy/exhaustive search, i.e., `neighbors.LazySearch` might be a better idea.

    Parameters
    ----------
    graph_k
        The maximum number of direct nearest neighbors each node has.
    maxlen
        The maximum size of the data buffer.
    warm_up
        How many data instances to observe before starting the search graph.
    dist_func
        The distance function used to compare two items. If not set, use the Minkowski distance
        with `p=2`.
    max_candidates
        The maximum number of vertices to consider when performing local neighborhood joins. If not set
        SWINN will use `min(50, max(50, self.graph_k))`.
    delta
        Early stop parameter for the neighborhood refinement procedure. NNDescent will stop running
        if the maximum number of iterations is reached or the number of edge changes after an iteration
        is smaller than or equal to `delta * graph_k * n_nodes`. In the last expression, `n_nodes`
        refers to the number of graph nodes involved in the (local) neighborhood refinement.
    prune_prob
        The probability of removing redundant edges. Must be between `0` and `1`. If set to zero,
        no edge will be pruned. When set to one, every potentially redundant edge will be dropped.
    n_iters
        The maximum number of NNDescent iterations to perform to refine the search index.
    seed
        Random seed for reproducibility.

    Notes
    -----
    There is an accuracy/speed trade-off between `graph_k` and `sample_rate`. To ensure a single
    connected component, and thus an effective search index, one can increase `graph_k`. The
    `connectivity` method is a helper to determine whether the search index has a single connected component.
    However, search accuracy might come at the cost of increased memory usage and slow processing. To alleviate
    that, one can rely on decreasing the `sample_rate` to avoid exploring all the undirected edges of a node
    during search queries and local graph refinements. Moreover, the edge pruning procedures also help
    decreasing the computational costs. Note that, anything that limits the number of explored neighbors or
    prunes edges might have a negative impact on search accuracy.

    References
    ----------
    [^1]: Dong, W., Moses, C., & Li, K. (2011, March). Efficient k-nearest neighbor graph construction for
    generic similarity measures. In Proceedings of the 20th international conference on World wide web (pp. 577-586).

    """

    def __init__(
        self,
        graph_k: int = 20,
        dist_func: DistanceFunc | FunctionWrapper | None = None,
        maxlen: int = 1000,
        warm_up: int = 500,
        max_candidates: int = None,
        delta: float = 0.0001,
        prune_prob: float = 0.0,
        n_iters: int = 10,
        seed: int = None,
    ):
        self.graph_k = graph_k
        if dist_func is None:
            dist_func = functools.partial(utils.math.minkowski_distance, p=2)
        self.dist_func = dist_func

        self.maxlen = maxlen
        self.warm_up = warm_up
        if max_candidates is None:
            self.max_candidates = min(50, max(50, self.graph_k))
        else:
            self.max_candidates = max_candidates

        self.delta = delta
        self.prune_prob = prune_prob

        self.n_iters = n_iters
        self.seed = seed

        self._data: collections.deque[Vertex] = collections.deque(maxlen=self.maxlen)
        self._uuid = itertools.cycle(range(self.maxlen))
        self._rng = random.Random(self.seed)
        self._index = False

    def __len__(self):
        return len(self._data)

    def __getitem__(self, i):
        return self._data[i]

    def __iter__(self):
        yield from self._data

    def _init_graph(self):
        """Create a random nearest neighbor graph."""
        n_nodes = len(self)

        nodes = set([i for i in range(n_nodes)])
        for nid in range(n_nodes):
            nodes.remove(nid)
            ns = self._rng.sample(tuple(nodes), self.graph_k)
            dists = [math.inf for _ in range(self.graph_k)]
            self[nid].fill([self[n] for n in ns], dists)
            nodes.add(nid)

    def _fix_graph(self):
        """Connect every isolated node in the graph to their nearest neighbors."""

        for node in list(Vertex._isolated):
            if not node.is_isolated():
                continue
            neighbors, dists = self._search(node.item, self.graph_k)
            node.fill(neighbors, dists)

        # Update class property
        Vertex._isolated.clear()

    def _safe_node_removal(self):
        """Remove the oldest data point from the search graph.

        Make sure nodes are accessible from any given starting point after removing the oldest
        node in the search graph. New traversal paths will be added in case the removed node was
        the only bridge between its neighbors.

        """
        node = self._data.popleft()
        # Get previous neighborhood info
        rns = node.r_neighbors()[0]
        ns = node.neighbors()[0]
        node.farewell()

        # Nodes whose only direct neighbor was the removed node
        rns = {rn for rn in rns if not rn.has_neighbors()}
        # Nodes whose only reverse neighbor was the removed node
        ns = {n for n in ns if not n.has_rneighbors()}

        affected = list(rns | ns)
        isolated = rns.intersection(ns)

        # First we handle the unreachable nodes
        for al in isolated:
            neighbors, dists = self._search(al.item, self.graph_k)
            al.fill(neighbors, dists)

        rns -= isolated
        ns -= isolated
        ns = tuple(ns)

        # Nodes with no direct neighbors
        for rn in rns:
            seed = None
            # Check the group of nodes without reverse neighborhood for seeds
            # Thus we can join two separate groups
            if len(ns) > 0:
                seed = self._rng.choice(ns)

            # Use the search index to create new connections
            neighbors, dists = self._search(rn.item, self.graph_k, seed=seed, exclude=rn)
            rn.fill(neighbors, dists)

        self._refine(affected)

    def _refine(self, nodes: list[Vertex] = None):
        """Update the nearest neighbor graph to improve the edge distances.

        Parameters
        ----------
        nodes
            The list of nodes for which the neighborhood refinement will be applied.
            If `None`, all nodes will have their neighborhood enhanced.
        """

        if nodes is None:
            nodes = [n for n in self]

        min_changes = self.delta * self.graph_k * len(nodes)

        tried = set()
        for _ in range(self.n_iters):
            total_changes = 0

            new = collections.defaultdict(set)
            old = collections.defaultdict(set)

            # Expand undirected neighborhood
            for node in nodes:
                neighbors = node.neighbors()[0]
                flags = node.sample_flags

                for neigh, flag in zip(neighbors, flags):
                    # To avoid evaluating previous neighbors again
                    tried.add((node.uuid, neigh.uuid))
                    if flag:
                        new[node].add(neigh)
                        new[neigh].add(node)
                    else:
                        old[node].add(neigh)
                        old[neigh].add(node)

            # Limits the maximum number of edges to explore and update sample flags
            for node in nodes:
                if len(new[node]) > self.max_candidates:
                    new[node] = self._rng.sample(tuple(new[node]), self.max_candidates)  # type: ignore
                else:
                    new[node] = new[node]

                if len(old[node]) > self.max_candidates:
                    old[node] = self._rng.sample(tuple(old[node]), self.max_candidates)  # type: ignore
                else:
                    old[node] = old[node]

                node.sample_flags = new[node]

            # Perform local joins an attempt to improve the neighborhood
            for node in nodes:
                # The origin of the join must have a boolean flag set to true
                for n1 in new[node]:
                    # Consider connections between vertices whose boolean flags are both true
                    for n2 in new[node]:
                        if n1.uuid == n2.uuid or n1.is_neighbor(n2):
                            continue

                        if (n1.uuid, n2.uuid) in tried or (n2.uuid, n1.uuid) in tried:
                            continue

                        dist = self.dist_func(n1.item, n2.item)
                        total_changes += n1.push_edge(n2, dist, self.graph_k)
                        total_changes += n2.push_edge(n1, dist, self.graph_k)

                        tried.add((n1.uuid, n2.uuid))

                    # Or one of the connections has a boolean flag set to false
                    for n2 in old[node]:
                        if n1.uuid == n2.uuid or n1.is_neighbor(n2):
                            continue

                        if (n1.uuid, n2.uuid) in tried or (n2.uuid, n1.uuid) in tried:
                            continue

                        dist = self.dist_func(n1.item, n2.item)
                        total_changes += n1.push_edge(n2, dist, self.graph_k)
                        total_changes += n2.push_edge(n1, dist, self.graph_k)

                        tried.add((n1.uuid, n2.uuid))

            # Stopping criterion
            if total_changes <= min_changes:
                break

        # Reduce the number of edges, if needed
        for n in nodes:
            n.prune(self.prune_prob, self.max_candidates, self._rng)

        # Ensure that no node is isolated in the graph
        self._fix_graph()

    def append(self, item: typing.Any, **kwargs):
        """Add a new item to the search index.

        Data is stored using the FIFO strategy. Both the data buffer and the search graph are updated. The
        addition of a new item will trigger the removal of the oldest item, if the maximum size was
        reached. All edges of the removed node are also dropped and safety procedures are applied to ensure
        its neighbors keep accessible. The addition of a new item also trigger local neighborhood refinement
        procedures, to ensure the search index is effective and the node degree constraints are met.

        Parameters
        ----------
        item
            Item to be added.
        kwargs
            Not used in this implementation.

        """
        node = Vertex(item, next(self._uuid))
        if not self._index:
            self._data.append(node)
            if len(self) >= self.warm_up:
                self._init_graph()
                self._refine()
                self._index = True
            return

        # A slot will be replaced, so let's update the search graph first
        if len(self) == self.maxlen:
            self._safe_node_removal()

        # Assign the closest neighbors to the new item
        neighbors, dists = self._search(node.item, self.graph_k)

        # Add the new element to the buffer
        self._data.append(node)
        node.fill(neighbors, dists)

    def _linear_scan(self, item, k):
        # Lazy search while the warm-up period is not finished
        points = [(p.item, self.dist_func(item, p.item)) for p in self]

        if points:
            return tuple(map(list, zip(*sorted(points, key=operator.itemgetter(-1))[:k])))

        return None

    def _search(self, item, k, epsilon: float = 0.1, seed=None, exclude=None) -> tuple[list, list]:
        # Limiter for the distance bound
        distance_scale = 1 + epsilon
        # Distance threshold for early stops
        distance_bound = math.inf

        if exclude is None:
            exclude = set()
        else:
            exclude = {exclude.uuid}

        if seed is None:
            # Make sure the starting point for the search is valid
            while True:
                # Random seed point to start the search
                seed = self[self._rng.randint(0, len(self) - 1)]
                if not seed.is_isolated() and seed.uuid not in exclude:
                    break

        dist = self.dist_func(item, seed.item)

        # To avoid computing distances more than once for a given node
        visited = {seed.uuid}
        visited |= exclude

        # Search pool is a minimum heap
        pool = [(dist, seed)]

        # Results are stored in a maximum heap
        result = [(-dist, seed)]

        c_dist, c_n = heapq.heappop(pool)
        while c_dist < distance_bound:
            tns = [n for n in c_n.all_neighbors() if n.uuid not in visited]

            for n in tns:
                dist = self.dist_func(item, n.item)

                if len(result) < k:
                    heapq.heappush(result, (-dist, n))
                    heapq.heappush(pool, (dist, n))
                    distance_bound = distance_scale * -result[0][0]
                elif dist < -result[0][0]:
                    heapq.heapreplace(result, (-dist, n))
                    heapq.heappush(pool, (dist, n))
                    distance_bound = distance_scale * -result[0][0]
                visited.add(n.uuid)
            if len(pool) == 0:
                break
            c_dist, c_n = heapq.heappop(pool)

        result.sort(reverse=True)
        neighbors, dists = map(list, zip(*((r[1], -r[0]) for r in result)))

        return neighbors, dists

    def search(
        self, item: typing.Any, n_neighbors: int, epsilon: float = 0.1, **kwargs
    ) -> tuple[list, list]:
        """Search the underlying nearest neighbor graph given a query item.

        In case not enough samples were observed, i.e., the number of stored samples is smaller than
        `warm_up`, then the search switches to a brute force strategy.

        Parameters
        ----------
        item
            The query item to search for nearest neighbors.
        n_neighbors
            The number of nearest neighbors to return.
        epsilon
            Distance bound to aid in avoiding local minima while traversing the search graph. Let $d_k$
            be the distance of the query item to current $k$-th nearest neighbor. At any given
            moment, any point whose distance to the query item is smaller than or equal to
            $(1 + \\epsilon) * d_k$ is kept as a potential path. After every addition to the heap of
            candidate nodes, $d_k$ is updated.
        kwargs
            Not used in this implementation.

        Returns
        -------
        neighbors, dists
            A tuple containing the id of the neighbors in the buffer and the respective distances to them.

        """

        if len(self) <= self.warm_up:
            return self._linear_scan(item, n_neighbors)

        neighbors, dists = self._search(item, n_neighbors, epsilon)
        return [n.item for n in neighbors], dists

    def connectivity(self) -> list[int]:
        """Get a list with the size of each connected component in the search graph.

        This metric provides an overview of reachability in the search index by using Kruskal's
        algorithm to build a forest of connected components.

        We want our search index to have a single connected component, i.e., the case where we get
        a list containing a single number which is equal to `maxlen`. If that is not the case, not
        every node in the search graph can be reached from any given starting point. You may want to try
        increasing `graph_k` to improve connectivity. However, keep in mind the following aspects:
        1) computing this metric is a costly operation ($O(E\\log V)$), where $E$ and $V$ are, respectively,
        the number of edges and vertices in the search graph; 2) often, connectivity comes at the price of
        increased computational costs. Tweaking the `sample_rate` might help in such situations. The best
        possible scenario is to decrease the value of `graph_k` while keeping a single connected
        component.

        Returns
        -------
        A list of the number of elements in each connected component of the graph.

        """
        forest = set()
        trees = {n: {n} for n in self}

        edges = [((n1, n2), w) for n1 in self for n2, w in n1.edges.items()]
        edges.sort(key=operator.itemgetter(1))

        for (n1, n2), _ in edges:
            if trees[n1].isdisjoint(trees[n2]):
                forest.discard(frozenset(trees[n1]))
                forest.discard(frozenset(trees[n2]))

                u = trees[n1] | trees[n2]
                # Update the trees
                for v in u:
                    trees[v] = u

                forest.add(frozenset(u))

        return [len(tree) for tree in forest]
