#

# Licensed to the Apache Software Foundation (ASF) under one or more

# contributor license agreements. See the NOTICE file distributed with

# this work for additional information regarding copyright ownership.

# The ASF licenses this file to You under the Apache License, Version 2.0

# (the "License"); you may not use this file except in compliance with

# the License. You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

#

"""Trackers for calculating median in windowed fashion."""

from typing import Optional

from apache_beam.ml.anomaly.specifiable import specifiable

from apache_beam.ml.anomaly.univariate.base import BaseTracker

from apache_beam.ml.anomaly.univariate.quantile import BufferedSlidingQuantileTracker # pylint: disable=line-too-long

from apache_beam.ml.anomaly.univariate.quantile import QuantileTracker

DEFAULT_WINDOW_SIZE = 1000

@specifiable

class MedianTracker(BaseTracker):

"""Tracks the median of a stream of values using a quantile tracker.

This wrapper class encapsulates a `QuantileTracker` configured specifically

for the 0.5 quantile (median).

Args:

quantile_tracker: An optional `QuantileTracker` instance. If not provided,

a `BufferedSlidingQuantileTracker` with a default window size 1000 and

q=0.5 is created.

Raises:

AssertionError: If the provided quantile_tracker is not initialized with

q=0.5.

"""

def __init__(self, quantile_tracker: Optional[QuantileTracker] = None):

self._quantile_tracker = quantile_tracker or BufferedSlidingQuantileTracker(

DEFAULT_WINDOW_SIZE, 0.5)

assert self._quantile_tracker._q == 0.5, \

"quantile_tracker must be initialized with q = 0.5"

def push(self, x):

"""Pushes a new value and updates the internal quantile tracker.

Args:

x: The new value to be pushed.

"""

self._quantile_tracker.push(x)

def get(self):

"""Calculates and returns the median (q = 0.5).

Returns:

float: The median of the values in the window setting specified in the

internal quantile tracker. Returns NaN if the window is empty.

"""

return self._quantile_tracker.get()