    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)

    if model.value == 'multiplicative':
        new_ts = ts / other
    elif model.value == 'additive':
        new_ts = ts - other
    else:
        raise_log(ValueError('Invalid parameter; must be either ADDITIVE or MULTIPLICATIVE. Was: {}'.format(model)))
    return new_ts

"""
Utils for time series statistics
--------------------------------
"""

import math
from typing import Tuple, Optional, List, Union

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
from scipy.signal import argrelmax
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import acf, pacf

from warnings import warn
from ..logging import raise_log, get_logger, raise_if_not, raise_if
from ..timeseries import TimeSeries
from .missing_values import fill_missing_values
from .utils import SeasonalityMode, ModelMode

logger = get_logger(__name__)

def check_seasonality(ts: TimeSeries,
                      m: Optional[int] = None,
                      max_lag: int = 24,
                      alpha: float = 0.05):
    """
    Checks whether the TimeSeries `ts` is seasonal with period `m` or not.

    If `m` is None, we work under the assumption that there is a unique seasonality period, which is inferred
    from the Auto-correlation Function (ACF).

    Parameters
    ----------
    ts
        The time series to check for seasonality.
    m
        The seasonality period to check.
    max_lag
        The maximal lag allowed in the ACF.
    alpha
        The desired confidence level (default 5%).

    Returns
    -------
    Tuple[bool, int]
        A tuple `(season, m)`, where season is a boolean indicating whether the series has seasonality or not
        and `m` is the seasonality period.
    """

    ts._assert_univariate()

    if m is not None and (m < 2 or not isinstance(m, int)):
        raise_log(ValueError('m must be an integer greater than 1.'), logger)

    if m is not None and m > max_lag:
        raise_log(ValueError('max_lag must be greater than or equal to m.'), logger)

    n_unique = np.unique(ts.values()).shape[0]

    if n_unique == 1:  # Check for non-constant TimeSeries
        return False, 0

    r = acf(ts.values(), nlags=max_lag, fft=False)  # In case user wants to check for seasonality higher than 24 steps.

    # Finds local maxima of Auto-Correlation Function
    candidates = argrelmax(r)[0]

    if len(candidates) == 0:
        logger.info('The ACF has no local maximum for m < max_lag = {}.'.format(max_lag))
        return False, 0

    if m is not None:
        # Check for local maximum when m is user defined.
        test = m not in candidates

        if test:
            return False, m

        candidates = [m]

    # Remove r[0], the auto-correlation at lag order 0, that introduces bias.
    r = r[1:]

    # The non-adjusted upper limit of the significance interval.
    band_upper = r.mean() + norm.ppf(1 - alpha / 2) * r.var()

    # Significance test, stops at first admissible value. The two '-1' below
    # compensate for the index change due to the restriction of the original r to r[1:].
    for candidate in candidates:
        stat = _bartlett_formula(r, candidate - 1, len(ts))
        if r[candidate - 1] > stat * band_upper:
            return True, candidate
    return False, 0

def _bartlett_formula(r: np.ndarray,
                      m: int,
                      length: int) -> float:
    """
    Computes the standard error of `r` at order `m` with respect to `length` according to Bartlett's formula.

    Parameters
    ----------
    r
        The array whose standard error is to be computed.
    m
        The order of the standard error.
    length
        The size of the underlying sample to be used.

    Returns
    -------
    float
        The standard error of `r` with order `m`.
    """

    if m == 1:
        return math.sqrt(1 / length)
    else:
        return math.sqrt((1 + 2 * sum(map(lambda x: x ** 2, r[:m - 1]))) / length)

def extract_trend_and_seasonality(ts: TimeSeries,
                                  freq: int = None,
                                  model: Union[SeasonalityMode, ModelMode] = ModelMode.MULTIPLICATIVE) -> \
        Tuple[TimeSeries, TimeSeries]:
    """
    Extracts trend and seasonality from a TimeSeries instance using `statsmodels.seasonal_decompose`.

    Parameters
    ----------
    ts
        The series to decompose
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode, SeasonalityMode` Enum member.
        Either MULTIPLICATIVE or ADDITIVE.
        Defaults ModelMode.MULTIPLICATIVE.

    Returns
    -------
        A tuple of (trend, seasonal) time series.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    decomp = seasonal_decompose(ts.pd_series(), period=freq, model=model.value, extrapolate_trend='freq')

    season = TimeSeries.from_times_and_values(ts.time_index, decomp.seasonal)
    trend = TimeSeries.from_times_and_values(ts.time_index, decomp.trend)

    return trend, season

def remove_from_series(ts: TimeSeries,
                       other: TimeSeries,
                       model: Union[SeasonalityMode, ModelMode]) -> TimeSeries:
    """
    Removes the TimeSeries `other` from the TimeSeries `ts` as specified by `model`.
    Use e.g. to remove an additive or multiplicative trend from a series.

    Parameters
    ----------
    ts
        The TimeSeries to be modified.
    other
        The TimeSeries to remove.
    model
        The type of model considered.
        Must be `from darts import ModelMode, SeasonalityMode` Enums member.
        Either MULTIPLICATIVE or ADDITIVE.

    Returns
    -------
    TimeSeries
        A TimeSeries defined by removing `other` from `ts`.
    """

def remove_seasonality(ts: TimeSeries,
                       freq: int = None,
                       model: SeasonalityMode = SeasonalityMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a seasonality of order `frequency` using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be a `from darts import SeasonalityMode` Enum member.
        Either SeasonalityMode.MULTIPLICATIVE or SeasonalityMode.ADDITIVE.
        Defaults SeasonalityMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the seasonality-adjusted 'ts'.
    """

    ts._assert_univariate()
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    _, seasonality = extract_trend_and_seasonality(ts, freq, model)
    new_ts = remove_from_series(ts, seasonality, model)
    return new_ts

def remove_trend(ts: TimeSeries,
                 model: ModelMode = ModelMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a trend using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode` Enum member.
        Either ModelMode.MULTIPLICATIVE or ModelMode.ADDITIVE.
        Defaults to modelMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the trend-adjusted 'ts'.
    """

    ts._assert_univariate()

    trend, _ = extract_trend_and_seasonality(ts, model=model)
    new_ts = remove_from_series(ts, trend, model)
    return new_ts

def plot_acf(ts: TimeSeries,
             m: Optional[int] = None,
             max_lag: int = 24,
             alpha: float = 0.05,
             bartlett_confint: bool = True,
             fig_size: Tuple[int, int] = (10, 5),
             axis: Optional[plt.axis] = None) -> None:
    """
    Plots the ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <http://example.com>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    alpha : float, default: 0.05
        The confidence interval to display.
    bartlett_confint : bool, default: True
        The boolean value indicating whether the confidence interval should be
        calculated using Bartlett's formula. If set to True, the confidence interval
        can be used in the model identification stage for fitting ARIMA models.
        If set to False, the confidence interval can be used to test for randomness
        (i.e. there is no time dependence in the data) of the data.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)),
             'max_lag must be greater than or equal to 1 and less than len(ts).')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = acf(ts.values(), nlags=max_lag, fft=False, alpha=alpha, bartlett_confint=bartlett_confint)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_pacf(ts: TimeSeries,
              m: Optional[int] = None,
              max_lag: int = 24,
              method: str = "ywadjusted",
              alpha: float = 0.05,
              fig_size: Tuple[int, int] = (10, 5),
              axis: Optional[plt.axis] = None) -> None:
    """
    Plots the Partial ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <https://github.com/statsmodels/statsmodels>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    method : str, default: "ywadjusted"
        The method to be used for the PACF calculation.
        - "yw" or "ywadjusted" : Yule-Walker with sample-size adjustment in
          denominator for acovf. Default.
        - "ywm" or "ywmle" : Yule-Walker without adjustment.
        - "ols" : regression of time series on lags of it and on constant.
        - "ols-inefficient" : regression of time series on lags using a single
          common sample to estimate all pacf coefficients.
        - "ols-adjusted" : regression of time series on lags with a bias
          adjustment.
        - "ld" or "ldadjusted" : Levinson-Durbin recursion with bias
          correction.
        - "ldb" or "ldbiased" : Levinson-Durbin recursion without bias
          correction.
    alpha : float, default: 0.05
        The confidence interval to display.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)//2),
             'max_lag must be greater than or equal to 1 and less than len(ts)//2.')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = pacf(ts.values(), nlags=max_lag, method=method, alpha=alpha)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_hist(data: Union[TimeSeries, List[float], np.ndarray],
              bins: Optional[Union[int, np.ndarray, List[float]]] = None,
              density: bool = False,
              title: Optional[str] = None,
              fig_size: Optional[Tuple[int, int]] = None,
              ax: Optional[plt.axis] = None) -> None:
    """ This function plots the histogram of values in a TimeSeries instance or an array-like.

    All types of TimeSeries are supported (uni-, multivariate, deterministic, stochastic).
    Depending on the number of components in `data`, up to four histograms can be plotted on one figure.
    All stochastic samples will be displayed with the corresponding component.

    If `data` is an array-like, ALL values will be displayed in the same histogram.

    Parameters
    ----------
    data
        TimeSeries instance or an array-like from which to plot the histogram.
    bins
        Optionally, either an integer value for the number of bins to be displayed
        or an array-like of floats determining the position of bins.
    density
        bool, if `density` is set to True, the bin counts will be converted to probability density
    title
        The title of the figure to be displayed
    fig_size
        The size of the figure to be displayed.
    ax
        Optionally, an axis object to plot the histogram on.
    """

    n_plots_max = 4

    if isinstance(data, TimeSeries):
        if len(data.components) > n_plots_max:
            logger.warning("TimeSeries contains more than 4 components. Only the first 4 components will be displayed.")

        components = list(data.components[:n_plots_max])
        values = data[components].all_values(copy=False).flatten(order='F')
    else:
        values = data if isinstance(data, np.ndarray) else np.array(data)
        if len(values.shape) > 1:
            logger.warning("Input array-like data with `dim>1d` will be flattened and displayed in one histogram.")

        components = ['Data']
        values = values.flatten(order='F')

    # compute the number of columns and rows for subplots depending on shape of input data
    n_components = len(components)
    n_cols = 1 if n_components == 1 else 2
    n_rows = math.ceil(n_components / n_cols)

    title = 'Histogram' if title is None else title
    if ax is None:
        fig = plt.figure(constrained_layout=True, figsize=fig_size)
        gs = fig.add_gridspec(n_rows, n_cols)
        fig.suptitle(title)
        ax_all = [fig.add_subplot(gs[i]) for i in range(n_components)]
    else:
        if n_components > 1:
            logger.warning("Only the first component is plotted when calling plot_hist() with a given `ax`")
        ax.set_title(title)
        ax_all = [ax]

    n_entries = len(values) // n_components
    for i, label, ax_i in zip(range(n_components), components, ax_all):
        ax_i.hist(values[i * n_entries:(i + 1) * n_entries], bins=bins, density=density, label=label)
        ax_i.set_xlabel('value')
        ax_i.set_ylabel('count' if not density else 'probability density')
        ax_i.legend()

def plot_residuals_analysis(residuals: TimeSeries,
                            num_bins: int = 20,
                            fill_nan: bool = True) -> None:
    """ Plots data relevant to residuals.

    This function takes a univariate TimeSeries instance of residuals and plots their values,
    their distribution and their ACF.
    Please note that if the residual TimeSeries instance contains NaN values, the plots
    might be displayed incorrectly. If `fill_nan` is set to True, the missing values will
    be interpolated.

    Parameters
    ----------
    residuals
        Univariate TimeSeries instance representing residuals.
    num_bins
        Optionally, an integer value determining the number of bins in the histogram.
    fill_nan:
        A boolean value indicating whether NaN values should be filled in the residuals.
    """

    residuals._assert_univariate()

    fig = plt.figure(constrained_layout=True, figsize=(8, 6))
    gs = fig.add_gridspec(2, 2)

    if fill_nan:
        residuals = fill_missing_values(residuals)

    # plot values
    ax1 = fig.add_subplot(gs[:1, :])
    residuals.plot(ax=ax1)
    ax1.set_ylabel('value')
    ax1.set_title('Residual values')

    # plot histogram and distribution
    res_mean, res_std = np.mean(residuals.univariate_values()), np.std(residuals.univariate_values())
    res_min, res_max = min(residuals.univariate_values()), max(residuals.univariate_values())
    x = np.linspace(res_min, res_max, 100)
    ax2 = fig.add_subplot(gs[1:, 1:])
    plot_hist(residuals, bins=num_bins, ax=ax2)
    ax2.plot(x, norm(res_mean, res_std).pdf(x) * len(residuals) * (res_max - res_min) / num_bins)
    ax2.yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    ax2.set_title('Distribution')
    ax2.set_ylabel('count')
    ax2.set_xlabel('value')

    # plot ACF
    ax3 = fig.add_subplot(gs[1:, :1])
    plot_acf(residuals, axis=ax3)
    ax3.set_ylabel('ACF value')
    ax3.set_xlabel('lag')
    ax3.set_title('ACF')


"""
Utils for time series statistics
--------------------------------
"""

import math
from typing import Tuple, Optional, List, Union

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
from scipy.signal import argrelmax
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import acf, pacf

from warnings import warn
from ..logging import raise_log, get_logger, raise_if_not, raise_if
from ..timeseries import TimeSeries
from .missing_values import fill_missing_values
from .utils import SeasonalityMode, ModelMode

logger = get_logger(__name__)

def check_seasonality(ts: TimeSeries,
                      m: Optional[int] = None,
                      max_lag: int = 24,
                      alpha: float = 0.05):
    """
    Checks whether the TimeSeries `ts` is seasonal with period `m` or not.

    If `m` is None, we work under the assumption that there is a unique seasonality period, which is inferred
    from the Auto-correlation Function (ACF).

    Parameters
    ----------
    ts
        The time series to check for seasonality.
    m
        The seasonality period to check.
    max_lag
        The maximal lag allowed in the ACF.
    alpha
        The desired confidence level (default 5%).

    Returns
    -------
    Tuple[bool, int]
        A tuple `(season, m)`, where season is a boolean indicating whether the series has seasonality or not
        and `m` is the seasonality period.
    """

    ts._assert_univariate()

    if m is not None and (m < 2 or not isinstance(m, int)):
        raise_log(ValueError('m must be an integer greater than 1.'), logger)

    if m is not None and m > max_lag:
        raise_log(ValueError('max_lag must be greater than or equal to m.'), logger)

    n_unique = np.unique(ts.values()).shape[0]

    if n_unique == 1:  # Check for non-constant TimeSeries
        return False, 0

    r = acf(ts.values(), nlags=max_lag, fft=False)  # In case user wants to check for seasonality higher than 24 steps.

    # Finds local maxima of Auto-Correlation Function
    candidates = argrelmax(r)[0]

    if len(candidates) == 0:
        logger.info('The ACF has no local maximum for m < max_lag = {}.'.format(max_lag))
        return False, 0

    if m is not None:
        # Check for local maximum when m is user defined.
        test = m not in candidates

        if test:
            return False, m

        candidates = [m]

    # Remove r[0], the auto-correlation at lag order 0, that introduces bias.
    r = r[1:]

    # The non-adjusted upper limit of the significance interval.
    band_upper = r.mean() + norm.ppf(1 - alpha / 2) * r.var()

    # Significance test, stops at first admissible value. The two '-1' below
    # compensate for the index change due to the restriction of the original r to r[1:].
    for candidate in candidates:
        stat = _bartlett_formula(r, candidate - 1, len(ts))
        if r[candidate - 1] > stat * band_upper:
            return True, candidate
    return False, 0

def _bartlett_formula(r: np.ndarray,
                      m: int,
                      length: int) -> float:
    """
    Computes the standard error of `r` at order `m` with respect to `length` according to Bartlett's formula.

    Parameters
    ----------
    r
        The array whose standard error is to be computed.
    m
        The order of the standard error.
    length
        The size of the underlying sample to be used.

    Returns
    -------
    float
        The standard error of `r` with order `m`.
    """

    if m == 1:
        return math.sqrt(1 / length)
    else:
        return math.sqrt((1 + 2 * sum(map(lambda x: x ** 2, r[:m - 1]))) / length)

def extract_trend_and_seasonality(ts: TimeSeries,
                                  freq: int = None,
                                  model: Union[SeasonalityMode, ModelMode] = ModelMode.MULTIPLICATIVE) -> \
        Tuple[TimeSeries, TimeSeries]:
    """
    Extracts trend and seasonality from a TimeSeries instance using `statsmodels.seasonal_decompose`.

    Parameters
    ----------
    ts
        The series to decompose
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode, SeasonalityMode` Enum member.
        Either MULTIPLICATIVE or ADDITIVE.
        Defaults ModelMode.MULTIPLICATIVE.

    Returns
    -------
        A tuple of (trend, seasonal) time series.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    decomp = seasonal_decompose(ts.pd_series(), period=freq, model=model.value, extrapolate_trend='freq')

    season = TimeSeries.from_times_and_values(ts.time_index, decomp.seasonal)
    trend = TimeSeries.from_times_and_values(ts.time_index, decomp.trend)

    return trend, season

def remove_from_series(ts: TimeSeries,
                       other: TimeSeries,
                       model: Union[SeasonalityMode, ModelMode]) -> TimeSeries:
    """
    Removes the TimeSeries `other` from the TimeSeries `ts` as specified by `model`.
    Use e.g. to remove an additive or multiplicative trend from a series.

    Parameters
    ----------
    ts
        The TimeSeries to be modified.
    other
        The TimeSeries to remove.
    model
        The type of model considered.
        Must be `from darts import ModelMode, SeasonalityMode` Enums member.
        Either MULTIPLICATIVE or ADDITIVE.

    Returns
    -------
    TimeSeries
        A TimeSeries defined by removing `other` from `ts`.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)

    if model.value == 'multiplicative':
        new_ts = ts / other
    elif model.value == 'additive':
        new_ts = ts - other
    else:
        raise_log(ValueError('Invalid parameter; must be either ADDITIVE or MULTIPLICATIVE. Was: {}'.format(model)))
    return new_ts

def remove_seasonality(ts: TimeSeries,
                       freq: int = None,
                       model: SeasonalityMode = SeasonalityMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a seasonality of order `frequency` using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be a `from darts import SeasonalityMode` Enum member.
        Either SeasonalityMode.MULTIPLICATIVE or SeasonalityMode.ADDITIVE.
        Defaults SeasonalityMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the seasonality-adjusted 'ts'.
    """

    ts._assert_univariate()
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    _, seasonality = extract_trend_and_seasonality(ts, freq, model)
    new_ts = remove_from_series(ts, seasonality, model)
    return new_ts

def remove_trend(ts: TimeSeries,
                 model: ModelMode = ModelMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a trend using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode` Enum member.
        Either ModelMode.MULTIPLICATIVE or ModelMode.ADDITIVE.
        Defaults to modelMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the trend-adjusted 'ts'.
    """

    ts._assert_univariate()

    trend, _ = extract_trend_and_seasonality(ts, model=model)
    new_ts = remove_from_series(ts, trend, model)
    return new_ts

def plot_acf(ts: TimeSeries,
             m: Optional[int] = None,
             max_lag: int = 24,
             alpha: float = 0.05,
             bartlett_confint: bool = True,
             fig_size: Tuple[int, int] = (10, 5),
             axis: Optional[plt.axis] = None) -> None:
    """
    Plots the ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <http://example.com>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    alpha : float, default: 0.05
        The confidence interval to display.
    bartlett_confint : bool, default: True
        The boolean value indicating whether the confidence interval should be
        calculated using Bartlett's formula. If set to True, the confidence interval
        can be used in the model identification stage for fitting ARIMA models.
        If set to False, the confidence interval can be used to test for randomness
        (i.e. there is no time dependence in the data) of the data.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)),
             'max_lag must be greater than or equal to 1 and less than len(ts).')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = acf(ts.values(), nlags=max_lag, fft=False, alpha=alpha, bartlett_confint=bartlett_confint)

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_pacf(ts: TimeSeries,
              m: Optional[int] = None,
              max_lag: int = 24,
              method: str = "ywadjusted",
              alpha: float = 0.05,
              fig_size: Tuple[int, int] = (10, 5),
              axis: Optional[plt.axis] = None) -> None:
    """
    Plots the Partial ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <https://github.com/statsmodels/statsmodels>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    method : str, default: "ywadjusted"
        The method to be used for the PACF calculation.
        - "yw" or "ywadjusted" : Yule-Walker with sample-size adjustment in
          denominator for acovf. Default.
        - "ywm" or "ywmle" : Yule-Walker without adjustment.
        - "ols" : regression of time series on lags of it and on constant.
        - "ols-inefficient" : regression of time series on lags using a single
          common sample to estimate all pacf coefficients.
        - "ols-adjusted" : regression of time series on lags with a bias
          adjustment.
        - "ld" or "ldadjusted" : Levinson-Durbin recursion with bias
          correction.
        - "ldb" or "ldbiased" : Levinson-Durbin recursion without bias
          correction.
    alpha : float, default: 0.05
        The confidence interval to display.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)//2),
             'max_lag must be greater than or equal to 1 and less than len(ts)//2.')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = pacf(ts.values(), nlags=max_lag, method=method, alpha=alpha)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_hist(data: Union[TimeSeries, List[float], np.ndarray],
              bins: Optional[Union[int, np.ndarray, List[float]]] = None,
              density: bool = False,
              title: Optional[str] = None,
              fig_size: Optional[Tuple[int, int]] = None,
              ax: Optional[plt.axis] = None) -> None:
    """ This function plots the histogram of values in a TimeSeries instance or an array-like.

    All types of TimeSeries are supported (uni-, multivariate, deterministic, stochastic).
    Depending on the number of components in `data`, up to four histograms can be plotted on one figure.
    All stochastic samples will be displayed with the corresponding component.

    If `data` is an array-like, ALL values will be displayed in the same histogram.

    Parameters
    ----------
    data
        TimeSeries instance or an array-like from which to plot the histogram.
    bins
        Optionally, either an integer value for the number of bins to be displayed
        or an array-like of floats determining the position of bins.
    density
        bool, if `density` is set to True, the bin counts will be converted to probability density
    title
        The title of the figure to be displayed
    fig_size
        The size of the figure to be displayed.
    ax
        Optionally, an axis object to plot the histogram on.
    """

    n_plots_max = 4

    if isinstance(data, TimeSeries):
        if len(data.components) > n_plots_max:
            logger.warning("TimeSeries contains more than 4 components. Only the first 4 components will be displayed.")

        components = list(data.components[:n_plots_max])
        values = data[components].all_values(copy=False).flatten(order='F')
    else:
        values = data if isinstance(data, np.ndarray) else np.array(data)
        if len(values.shape) > 1:
            logger.warning("Input array-like data with `dim>1d` will be flattened and displayed in one histogram.")

        components = ['Data']
        values = values.flatten(order='F')

    # compute the number of columns and rows for subplots depending on shape of input data
    n_components = len(components)
    n_cols = 1 if n_components == 1 else 2
    n_rows = math.ceil(n_components / n_cols)

    title = 'Histogram' if title is None else title
    if ax is None:
        fig = plt.figure(constrained_layout=True, figsize=fig_size)
        gs = fig.add_gridspec(n_rows, n_cols)
        fig.suptitle(title)
        ax_all = [fig.add_subplot(gs[i]) for i in range(n_components)]
    else:
        if n_components > 1:
            logger.warning("Only the first component is plotted when calling plot_hist() with a given `ax`")
        ax.set_title(title)
        ax_all = [ax]

    n_entries = len(values) // n_components
    for i, label, ax_i in zip(range(n_components), components, ax_all):
        ax_i.hist(values[i * n_entries:(i + 1) * n_entries], bins=bins, density=density, label=label)
        ax_i.set_xlabel('value')
        ax_i.set_ylabel('count' if not density else 'probability density')
        ax_i.legend()

def plot_residuals_analysis(residuals: TimeSeries,
                            num_bins: int = 20,
                            fill_nan: bool = True) -> None:
    """ Plots data relevant to residuals.

    This function takes a univariate TimeSeries instance of residuals and plots their values,
    their distribution and their ACF.
    Please note that if the residual TimeSeries instance contains NaN values, the plots
    might be displayed incorrectly. If `fill_nan` is set to True, the missing values will
    be interpolated.

    Parameters
    ----------
    residuals
        Univariate TimeSeries instance representing residuals.
    num_bins
        Optionally, an integer value determining the number of bins in the histogram.
    fill_nan:
        A boolean value indicating whether NaN values should be filled in the residuals.
    """

    residuals._assert_univariate()

    fig = plt.figure(constrained_layout=True, figsize=(8, 6))
    gs = fig.add_gridspec(2, 2)

    if fill_nan:
        residuals = fill_missing_values(residuals)

    # plot values
    ax1 = fig.add_subplot(gs[:1, :])
    residuals.plot(ax=ax1)
    ax1.set_ylabel('value')
    ax1.set_title('Residual values')

    # plot histogram and distribution
    res_mean, res_std = np.mean(residuals.univariate_values()), np.std(residuals.univariate_values())
    res_min, res_max = min(residuals.univariate_values()), max(residuals.univariate_values())
    x = np.linspace(res_min, res_max, 100)
    ax2 = fig.add_subplot(gs[1:, 1:])
    plot_hist(residuals, bins=num_bins, ax=ax2)
    ax2.plot(x, norm(res_mean, res_std).pdf(x) * len(residuals) * (res_max - res_min) / num_bins)
    ax2.yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    ax2.set_title('Distribution')
    ax2.set_ylabel('count')
    ax2.set_xlabel('value')

    # plot ACF
    ax3 = fig.add_subplot(gs[1:, :1])
    plot_acf(residuals, axis=ax3)
    ax3.set_ylabel('ACF value')
    ax3.set_xlabel('lag')
    ax3.set_title('ACF')


"""
Utils for time series statistics
--------------------------------
"""

import math
from typing import Tuple, Optional, List, Union

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
from scipy.signal import argrelmax
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import acf, pacf

from warnings import warn
from ..logging import raise_log, get_logger, raise_if_not, raise_if
from ..timeseries import TimeSeries
from .missing_values import fill_missing_values
from .utils import SeasonalityMode, ModelMode

logger = get_logger(__name__)

def check_seasonality(ts: TimeSeries,
                      m: Optional[int] = None,
                      max_lag: int = 24,
                      alpha: float = 0.05):
    """
    Checks whether the TimeSeries `ts` is seasonal with period `m` or not.

    If `m` is None, we work under the assumption that there is a unique seasonality period, which is inferred
    from the Auto-correlation Function (ACF).

    Parameters
    ----------
    ts
        The time series to check for seasonality.
    m
        The seasonality period to check.
    max_lag
        The maximal lag allowed in the ACF.
    alpha
        The desired confidence level (default 5%).

    Returns
    -------
    Tuple[bool, int]
        A tuple `(season, m)`, where season is a boolean indicating whether the series has seasonality or not
        and `m` is the seasonality period.
    """

    ts._assert_univariate()

    if m is not None and (m < 2 or not isinstance(m, int)):
        raise_log(ValueError('m must be an integer greater than 1.'), logger)

    if m is not None and m > max_lag:
        raise_log(ValueError('max_lag must be greater than or equal to m.'), logger)

    n_unique = np.unique(ts.values()).shape[0]

    if n_unique == 1:  # Check for non-constant TimeSeries
        return False, 0

    r = acf(ts.values(), nlags=max_lag, fft=False)  # In case user wants to check for seasonality higher than 24 steps.

    # Finds local maxima of Auto-Correlation Function
    candidates = argrelmax(r)[0]

    if len(candidates) == 0:
        logger.info('The ACF has no local maximum for m < max_lag = {}.'.format(max_lag))
        return False, 0

    if m is not None:
        # Check for local maximum when m is user defined.
        test = m not in candidates

        if test:
            return False, m

        candidates = [m]

    # Remove r[0], the auto-correlation at lag order 0, that introduces bias.
    r = r[1:]

    # The non-adjusted upper limit of the significance interval.
    band_upper = r.mean() + norm.ppf(1 - alpha / 2) * r.var()

    # Significance test, stops at first admissible value. The two '-1' below
    # compensate for the index change due to the restriction of the original r to r[1:].
    for candidate in candidates:
        stat = _bartlett_formula(r, candidate - 1, len(ts))
        if r[candidate - 1] > stat * band_upper:
            return True, candidate
    return False, 0

def _bartlett_formula(r: np.ndarray,
                      m: int,
                      length: int) -> float:
    """
    Computes the standard error of `r` at order `m` with respect to `length` according to Bartlett's formula.

    Parameters
    ----------
    r
        The array whose standard error is to be computed.
    m
        The order of the standard error.
    length
        The size of the underlying sample to be used.

    Returns
    -------
    float
        The standard error of `r` with order `m`.
    """

    if m == 1:
        return math.sqrt(1 / length)
    else:
        return math.sqrt((1 + 2 * sum(map(lambda x: x ** 2, r[:m - 1]))) / length)

def extract_trend_and_seasonality(ts: TimeSeries,
                                  freq: int = None,
                                  model: Union[SeasonalityMode, ModelMode] = ModelMode.MULTIPLICATIVE) -> \
        Tuple[TimeSeries, TimeSeries]:
    """
    Extracts trend and seasonality from a TimeSeries instance using `statsmodels.seasonal_decompose`.

    Parameters
    ----------
    ts
        The series to decompose
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode, SeasonalityMode` Enum member.
        Either MULTIPLICATIVE or ADDITIVE.
        Defaults ModelMode.MULTIPLICATIVE.

    Returns
    -------
        A tuple of (trend, seasonal) time series.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    decomp = seasonal_decompose(ts.pd_series(), period=freq, model=model.value, extrapolate_trend='freq')

    season = TimeSeries.from_times_and_values(ts.time_index, decomp.seasonal)
    trend = TimeSeries.from_times_and_values(ts.time_index, decomp.trend)

    return trend, season

def remove_from_series(ts: TimeSeries,
                       other: TimeSeries,
                       model: Union[SeasonalityMode, ModelMode]) -> TimeSeries:
    """
    Removes the TimeSeries `other` from the TimeSeries `ts` as specified by `model`.
    Use e.g. to remove an additive or multiplicative trend from a series.

    Parameters
    ----------
    ts
        The TimeSeries to be modified.
    other
        The TimeSeries to remove.
    model
        The type of model considered.
        Must be `from darts import ModelMode, SeasonalityMode` Enums member.
        Either MULTIPLICATIVE or ADDITIVE.

    Returns
    -------
    TimeSeries
        A TimeSeries defined by removing `other` from `ts`.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)

    if model.value == 'multiplicative':
        new_ts = ts / other
    elif model.value == 'additive':
        new_ts = ts - other
    else:
        raise_log(ValueError('Invalid parameter; must be either ADDITIVE or MULTIPLICATIVE. Was: {}'.format(model)))
    return new_ts

def remove_seasonality(ts: TimeSeries,
                       freq: int = None,
                       model: SeasonalityMode = SeasonalityMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a seasonality of order `frequency` using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be a `from darts import SeasonalityMode` Enum member.
        Either SeasonalityMode.MULTIPLICATIVE or SeasonalityMode.ADDITIVE.
        Defaults SeasonalityMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the seasonality-adjusted 'ts'.
    """

    ts._assert_univariate()
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    _, seasonality = extract_trend_and_seasonality(ts, freq, model)
    new_ts = remove_from_series(ts, seasonality, model)
    return new_ts

def remove_trend(ts: TimeSeries,
                 model: ModelMode = ModelMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a trend using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode` Enum member.
        Either ModelMode.MULTIPLICATIVE or ModelMode.ADDITIVE.
        Defaults to modelMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the trend-adjusted 'ts'.
    """

    ts._assert_univariate()

    trend, _ = extract_trend_and_seasonality(ts, model=model)
    new_ts = remove_from_series(ts, trend, model)
    return new_ts

def plot_acf(ts: TimeSeries,
             m: Optional[int] = None,
             max_lag: int = 24,
             alpha: float = 0.05,
             bartlett_confint: bool = True,
             fig_size: Tuple[int, int] = (10, 5),
             axis: Optional[plt.axis] = None) -> None:
    """
    Plots the ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <http://example.com>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    alpha : float, default: 0.05
        The confidence interval to display.
    bartlett_confint : bool, default: True
        The boolean value indicating whether the confidence interval should be
        calculated using Bartlett's formula. If set to True, the confidence interval
        can be used in the model identification stage for fitting ARIMA models.
        If set to False, the confidence interval can be used to test for randomness
        (i.e. there is no time dependence in the data) of the data.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)),
             'max_lag must be greater than or equal to 1 and less than len(ts).')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = acf(ts.values(), nlags=max_lag, fft=False, alpha=alpha, bartlett_confint=bartlett_confint)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_pacf(ts: TimeSeries,
              m: Optional[int] = None,
              max_lag: int = 24,
              method: str = "ywadjusted",
              alpha: float = 0.05,
              fig_size: Tuple[int, int] = (10, 5),
              axis: Optional[plt.axis] = None) -> None:
    """
    Plots the Partial ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <https://github.com/statsmodels/statsmodels>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    method : str, default: "ywadjusted"
        The method to be used for the PACF calculation.
        - "yw" or "ywadjusted" : Yule-Walker with sample-size adjustment in
          denominator for acovf. Default.
        - "ywm" or "ywmle" : Yule-Walker without adjustment.
        - "ols" : regression of time series on lags of it and on constant.
        - "ols-inefficient" : regression of time series on lags using a single
          common sample to estimate all pacf coefficients.
        - "ols-adjusted" : regression of time series on lags with a bias
          adjustment.
        - "ld" or "ldadjusted" : Levinson-Durbin recursion with bias
          correction.
        - "ldb" or "ldbiased" : Levinson-Durbin recursion without bias
          correction.
    alpha : float, default: 0.05
        The confidence interval to display.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)//2),
             'max_lag must be greater than or equal to 1 and less than len(ts)//2.')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = pacf(ts.values(), nlags=max_lag, method=method, alpha=alpha)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_hist(data: Union[TimeSeries, List[float], np.ndarray],
              bins: Optional[Union[int, np.ndarray, List[float]]] = None,
              density: bool = False,
              title: Optional[str] = None,
              fig_size: Optional[Tuple[int, int]] = None,
              ax: Optional[plt.axis] = None) -> None:
    """ This function plots the histogram of values in a TimeSeries instance or an array-like.

    All types of TimeSeries are supported (uni-, multivariate, deterministic, stochastic).
    Depending on the number of components in `data`, up to four histograms can be plotted on one figure.
    All stochastic samples will be displayed with the corresponding component.

    If `data` is an array-like, ALL values will be displayed in the same histogram.

    Parameters
    ----------
    data
        TimeSeries instance or an array-like from which to plot the histogram.
    bins
        Optionally, either an integer value for the number of bins to be displayed
        or an array-like of floats determining the position of bins.
    density
        bool, if `density` is set to True, the bin counts will be converted to probability density
    title
        The title of the figure to be displayed
    fig_size
        The size of the figure to be displayed.
    ax
        Optionally, an axis object to plot the histogram on.
    """

    n_plots_max = 4

    if isinstance(data, TimeSeries):
        if len(data.components) > n_plots_max:
            logger.warning("TimeSeries contains more than 4 components. Only the first 4 components will be displayed.")

        components = list(data.components[:n_plots_max])
        values = data[components].all_values(copy=False).flatten(order='F')
    else:
        values = data if isinstance(data, np.ndarray) else np.array(data)
        if len(values.shape) > 1:
            logger.warning("Input array-like data with `dim>1d` will be flattened and displayed in one histogram.")

        components = ['Data']
        values = values.flatten(order='F')

    # compute the number of columns and rows for subplots depending on shape of input data
    n_components = len(components)
    n_cols = 1 if n_components == 1 else 2
    n_rows = math.ceil(n_components / n_cols)

    title = 'Histogram' if title is None else title
    if ax is None:
        fig = plt.figure(constrained_layout=True, figsize=fig_size)
        gs = fig.add_gridspec(n_rows, n_cols)
        fig.suptitle(title)
        ax_all = [fig.add_subplot(gs[i]) for i in range(n_components)]
    else:
        if n_components > 1:
            logger.warning("Only the first component is plotted when calling plot_hist() with a given `ax`")
        ax.set_title(title)
        ax_all = [ax]

    n_entries = len(values) // n_components
    for i, label, ax_i in zip(range(n_components), components, ax_all):
        ax_i.hist(values[i * n_entries:(i + 1) * n_entries], bins=bins, density=density, label=label)
        ax_i.set_xlabel('value')
        ax_i.set_ylabel('count' if not density else 'probability density')
        ax_i.legend()

def plot_residuals_analysis(residuals: TimeSeries,
                            num_bins: int = 20,
                            fill_nan: bool = True) -> None:
    """ Plots data relevant to residuals.

    This function takes a univariate TimeSeries instance of residuals and plots their values,
    their distribution and their ACF.
    Please note that if the residual TimeSeries instance contains NaN values, the plots
    might be displayed incorrectly. If `fill_nan` is set to True, the missing values will
    be interpolated.

    Parameters
    ----------
    residuals
        Univariate TimeSeries instance representing residuals.
    num_bins
        Optionally, an integer value determining the number of bins in the histogram.
    fill_nan:
        A boolean value indicating whether NaN values should be filled in the residuals.
    """

    residuals._assert_univariate()

    fig = plt.figure(constrained_layout=True, figsize=(8, 6))
    gs = fig.add_gridspec(2, 2)

    if fill_nan:
        residuals = fill_missing_values(residuals)

    # plot values
    ax1 = fig.add_subplot(gs[:1, :])
    residuals.plot(ax=ax1)
    ax1.set_ylabel('value')
    ax1.set_title('Residual values')

    # plot histogram and distribution
    res_mean, res_std = np.mean(residuals.univariate_values()), np.std(residuals.univariate_values())
    res_min, res_max = min(residuals.univariate_values()), max(residuals.univariate_values())
    x = np.linspace(res_min, res_max, 100)
    ax2 = fig.add_subplot(gs[1:, 1:])
    plot_hist(residuals, bins=num_bins, ax=ax2)
    ax2.plot(x, norm(res_mean, res_std).pdf(x) * len(residuals) * (res_max - res_min) / num_bins)
    ax2.yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    ax2.set_title('Distribution')
    ax2.set_ylabel('count')
    ax2.set_xlabel('value')

    # plot ACF
    ax3 = fig.add_subplot(gs[1:, :1])
    plot_acf(residuals, axis=ax3)
    ax3.set_ylabel('ACF value')
    ax3.set_xlabel('lag')
    ax3.set_title('ACF')


"""
Utils for time series statistics
--------------------------------
"""

import math
from typing import Tuple, Optional, List, Union

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import norm
from scipy.signal import argrelmax
from statsmodels.tsa.seasonal import seasonal_decompose
from statsmodels.tsa.stattools import acf, pacf

from warnings import warn
from ..logging import raise_log, get_logger, raise_if_not, raise_if
from ..timeseries import TimeSeries
from .missing_values import fill_missing_values
from .utils import SeasonalityMode, ModelMode

logger = get_logger(__name__)

def check_seasonality(ts: TimeSeries,
                      m: Optional[int] = None,
                      max_lag: int = 24,
                      alpha: float = 0.05):
    """
    Checks whether the TimeSeries `ts` is seasonal with period `m` or not.

    If `m` is None, we work under the assumption that there is a unique seasonality period, which is inferred
    from the Auto-correlation Function (ACF).

    Parameters
    ----------
    ts
        The time series to check for seasonality.
    m
        The seasonality period to check.
    max_lag
        The maximal lag allowed in the ACF.
    alpha
        The desired confidence level (default 5%).

    Returns
    -------
    Tuple[bool, int]
        A tuple `(season, m)`, where season is a boolean indicating whether the series has seasonality or not
        and `m` is the seasonality period.
    """

    ts._assert_univariate()

    if m is not None and (m < 2 or not isinstance(m, int)):
        raise_log(ValueError('m must be an integer greater than 1.'), logger)

    if m is not None and m > max_lag:
        raise_log(ValueError('max_lag must be greater than or equal to m.'), logger)

    n_unique = np.unique(ts.values()).shape[0]

    if n_unique == 1:  # Check for non-constant TimeSeries
        return False, 0

    r = acf(ts.values(), nlags=max_lag, fft=False)  # In case user wants to check for seasonality higher than 24 steps.

    # Finds local maxima of Auto-Correlation Function
    candidates = argrelmax(r)[0]

    if len(candidates) == 0:
        logger.info('The ACF has no local maximum for m < max_lag = {}.'.format(max_lag))
        return False, 0

    if m is not None:
        # Check for local maximum when m is user defined.
        test = m not in candidates

        if test:
            return False, m

        candidates = [m]

    # Remove r[0], the auto-correlation at lag order 0, that introduces bias.
    r = r[1:]

    # The non-adjusted upper limit of the significance interval.
    band_upper = r.mean() + norm.ppf(1 - alpha / 2) * r.var()

    # Significance test, stops at first admissible value. The two '-1' below
    # compensate for the index change due to the restriction of the original r to r[1:].
    for candidate in candidates:
        stat = _bartlett_formula(r, candidate - 1, len(ts))
        if r[candidate - 1] > stat * band_upper:
            return True, candidate
    return False, 0

def _bartlett_formula(r: np.ndarray,
                      m: int,
                      length: int) -> float:
    """
    Computes the standard error of `r` at order `m` with respect to `length` according to Bartlett's formula.

    Parameters
    ----------
    r
        The array whose standard error is to be computed.
    m
        The order of the standard error.
    length
        The size of the underlying sample to be used.

    Returns
    -------
    float
        The standard error of `r` with order `m`.
    """

    if m == 1:
        return math.sqrt(1 / length)
    else:
        return math.sqrt((1 + 2 * sum(map(lambda x: x ** 2, r[:m - 1]))) / length)

def extract_trend_and_seasonality(ts: TimeSeries,
                                  freq: int = None,
                                  model: Union[SeasonalityMode, ModelMode] = ModelMode.MULTIPLICATIVE) -> \
        Tuple[TimeSeries, TimeSeries]:
    """
    Extracts trend and seasonality from a TimeSeries instance using `statsmodels.seasonal_decompose`.

    Parameters
    ----------
    ts
        The series to decompose
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode, SeasonalityMode` Enum member.
        Either MULTIPLICATIVE or ADDITIVE.
        Defaults ModelMode.MULTIPLICATIVE.

    Returns
    -------
        A tuple of (trend, seasonal) time series.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    decomp = seasonal_decompose(ts.pd_series(), period=freq, model=model.value, extrapolate_trend='freq')

    season = TimeSeries.from_times_and_values(ts.time_index, decomp.seasonal)
    trend = TimeSeries.from_times_and_values(ts.time_index, decomp.trend)

    return trend, season

def remove_from_series(ts: TimeSeries,
                       other: TimeSeries,
                       model: Union[SeasonalityMode, ModelMode]) -> TimeSeries:
    """
    Removes the TimeSeries `other` from the TimeSeries `ts` as specified by `model`.
    Use e.g. to remove an additive or multiplicative trend from a series.

    Parameters
    ----------
    ts
        The TimeSeries to be modified.
    other
        The TimeSeries to remove.
    model
        The type of model considered.
        Must be `from darts import ModelMode, SeasonalityMode` Enums member.
        Either MULTIPLICATIVE or ADDITIVE.

    Returns
    -------
    TimeSeries
        A TimeSeries defined by removing `other` from `ts`.
    """

    ts._assert_univariate()
    raise_if_not(model in ModelMode or model in SeasonalityMode,
                 "Unknown value for model_mode: {}.".format(model), logger)

    if model.value == 'multiplicative':
        new_ts = ts / other
    elif model.value == 'additive':
        new_ts = ts - other
    else:
        raise_log(ValueError('Invalid parameter; must be either ADDITIVE or MULTIPLICATIVE. Was: {}'.format(model)))
    return new_ts

def remove_seasonality(ts: TimeSeries,
                       freq: int = None,
                       model: SeasonalityMode = SeasonalityMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a seasonality of order `frequency` using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    freq
        The seasonality period to use.
    model
        The type of decomposition to use.
        Must be a `from darts import SeasonalityMode` Enum member.
        Either SeasonalityMode.MULTIPLICATIVE or SeasonalityMode.ADDITIVE.
        Defaults SeasonalityMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the seasonality-adjusted 'ts'.
    """

    ts._assert_univariate()
    raise_if_not(model is not SeasonalityMode.NONE, "The model must be either MULTIPLICATIVE or ADDITIVE.")

    _, seasonality = extract_trend_and_seasonality(ts, freq, model)
    new_ts = remove_from_series(ts, seasonality, model)
    return new_ts

def remove_trend(ts: TimeSeries,
                 model: ModelMode = ModelMode.MULTIPLICATIVE) -> TimeSeries:
    """
    Adjusts the TimeSeries `ts` for a trend using the `model` decomposition.

    Parameters
    ----------
    ts
        The TimeSeries to adjust.
    model
        The type of decomposition to use.
        Must be `from darts import ModelMode` Enum member.
        Either ModelMode.MULTIPLICATIVE or ModelMode.ADDITIVE.
        Defaults to modelMode.MULTIPLICATIVE.
    Returns
    -------
    TimeSeries
        A new TimeSeries instance that corresponds to the trend-adjusted 'ts'.
    """

    ts._assert_univariate()

    trend, _ = extract_trend_and_seasonality(ts, model=model)
    new_ts = remove_from_series(ts, trend, model)
    return new_ts

def plot_acf(ts: TimeSeries,
             m: Optional[int] = None,
             max_lag: int = 24,
             alpha: float = 0.05,
             bartlett_confint: bool = True,
             fig_size: Tuple[int, int] = (10, 5),
             axis: Optional[plt.axis] = None) -> None:
    """
    Plots the ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <http://example.com>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    alpha : float, default: 0.05
        The confidence interval to display.
    bartlett_confint : bool, default: True
        The boolean value indicating whether the confidence interval should be
        calculated using Bartlett's formula. If set to True, the confidence interval
        can be used in the model identification stage for fitting ARIMA models.
        If set to False, the confidence interval can be used to test for randomness
        (i.e. there is no time dependence in the data) of the data.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)),
             'max_lag must be greater than or equal to 1 and less than len(ts).')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = acf(ts.values(), nlags=max_lag, fft=False, alpha=alpha, bartlett_confint=bartlett_confint)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_pacf(ts: TimeSeries,
              m: Optional[int] = None,
              max_lag: int = 24,
              method: str = "ywadjusted",
              alpha: float = 0.05,
              fig_size: Tuple[int, int] = (10, 5),
              axis: Optional[plt.axis] = None) -> None:
    """
    Plots the Partial ACF of `ts`, highlighting it at lag `m`, with corresponding significance interval.
    This function uses the `Statsmodels module <https://github.com/statsmodels/statsmodels>`_.

    Parameters
    ----------
    ts : TimeSeries
        The TimeSeries whose ACF should be plotted.
    m : int, optional
        Optionally, a time lag to highlight on the plot.
    max_lag : int, default: 24
        The maximal lag order to consider.
    method : str, default: "ywadjusted"
        The method to be used for the PACF calculation.
        - "yw" or "ywadjusted" : Yule-Walker with sample-size adjustment in
          denominator for acovf. Default.
        - "ywm" or "ywmle" : Yule-Walker without adjustment.
        - "ols" : regression of time series on lags of it and on constant.
        - "ols-inefficient" : regression of time series on lags using a single
          common sample to estimate all pacf coefficients.
        - "ols-adjusted" : regression of time series on lags with a bias
          adjustment.
        - "ld" or "ldadjusted" : Levinson-Durbin recursion with bias
          correction.
        - "ldb" or "ldbiased" : Levinson-Durbin recursion without bias
          correction.
    alpha : float, default: 0.05
        The confidence interval to display.
    fig_size : tuple of int, default: (10, 5)
        The size of the figure to be displayed.
    axis : plt.axis, optional
        Optionally, an axis object to plot the ACF on.
    """

    ts._assert_univariate()
    raise_if(max_lag is None or not (1 <= max_lag < len(ts)//2),
             'max_lag must be greater than or equal to 1 and less than len(ts)//2.')
    raise_if(m is not None and not (0 <= m <= max_lag),
             'm must be greater than or equal to 0 and less than or equal to max_lag.')
    raise_if(alpha is None or not (0 < alpha < 1), 'alpha must be greater than 0 and less than 1.')

    r, confint = pacf(ts.values(), nlags=max_lag, method=method, alpha=alpha)

    if axis is None:
        plt.figure(figsize=fig_size)
        axis = plt

    for i in range(len(r)):
        axis.plot((i, i),
                  (0, r[i]),
                  color=('#b512b8' if m is not None and i == m else 'black'),
                  lw=(1 if m is not None and i == m else .5))

    # Adjusts the upper band of the confidence interval to center it on the x axis.
    upp_band = [confint[lag][1] - r[lag] for lag in range(1, max_lag + 1)]

    axis.fill_between(np.arange(1, max_lag + 1), upp_band, [-x for x in upp_band], color='#003DFD', alpha=.25)
    axis.plot((0, max_lag + 1), (0, 0), color='black')

def plot_hist(data: Union[TimeSeries, List[float], np.ndarray],
              bins: Optional[Union[int, np.ndarray, List[float]]] = None,
              density: bool = False,
              title: Optional[str] = None,
              fig_size: Optional[Tuple[int, int]] = None,
              ax: Optional[plt.axis] = None) -> None:
    """ This function plots the histogram of values in a TimeSeries instance or an array-like.

    All types of TimeSeries are supported (uni-, multivariate, deterministic, stochastic).
    Depending on the number of components in `data`, up to four histograms can be plotted on one figure.
    All stochastic samples will be displayed with the corresponding component.

    If `data` is an array-like, ALL values will be displayed in the same histogram.

    Parameters
    ----------
    data
        TimeSeries instance or an array-like from which to plot the histogram.
    bins
        Optionally, either an integer value for the number of bins to be displayed
        or an array-like of floats determining the position of bins.
    density
        bool, if `density` is set to True, the bin counts will be converted to probability density
    title
        The title of the figure to be displayed
    fig_size
        The size of the figure to be displayed.
    ax
        Optionally, an axis object to plot the histogram on.
    """

    n_plots_max = 4

    if isinstance(data, TimeSeries):
        if len(data.components) > n_plots_max:
            logger.warning("TimeSeries contains more than 4 components. Only the first 4 components will be displayed.")

        components = list(data.components[:n_plots_max])
        values = data[components].all_values(copy=False).flatten(order='F')
    else:
        values = data if isinstance(data, np.ndarray) else np.array(data)
        if len(values.shape) > 1:
            logger.warning("Input array-like data with `dim>1d` will be flattened and displayed in one histogram.")

        components = ['Data']
        values = values.flatten(order='F')

    # compute the number of columns and rows for subplots depending on shape of input data
    n_components = len(components)
    n_cols = 1 if n_components == 1 else 2
    n_rows = math.ceil(n_components / n_cols)

    title = 'Histogram' if title is None else title
    if ax is None:
        fig = plt.figure(constrained_layout=True, figsize=fig_size)
        gs = fig.add_gridspec(n_rows, n_cols)
        fig.suptitle(title)
        ax_all = [fig.add_subplot(gs[i]) for i in range(n_components)]
    else:
        if n_components > 1:
            logger.warning("Only the first component is plotted when calling plot_hist() with a given `ax`")
        ax.set_title(title)
        ax_all = [ax]

    n_entries = len(values) // n_components
    for i, label, ax_i in zip(range(n_components), co

mponents, ax_all):
        ax_i.hist(values[i * n_entries:(i + 1) * n_entries], bins=bins, density=density, label=label)
        ax_i.set_xlabel('value')
        ax_i.set_ylabel('count' if not density else 'probability density')
        ax_i.legend()

def plot_residuals_analysis(residuals: TimeSeries,
                            num_bins: int = 20,
                            fill_nan: bool = True) -> None:
    """ Plots data relevant to residuals.

    This function takes a univariate TimeSeries instance of residuals and plots their values,
    their distribution and their ACF.
    Please note that if the residual TimeSeries instance contains NaN values, the plots
    might be displayed incorrectly. If `fill_nan` is set to True, the missing values will
    be interpolated.

    Parameters
    ----------
    residuals
        Univariate TimeSeries instance representing residuals.
    num_bins
        Optionally, an integer value determining the number of bins in the histogram.
    fill_nan:
        A boolean value indicating whether NaN values should be filled in the residuals.
    """

    residuals._assert_univariate()

    fig = plt.figure(constrained_layout=True, figsize=(8, 6))
    gs = fig.add_gridspec(2, 2)

    if fill_nan:
        residuals = fill_missing_values(residuals)

    # plot values
    ax1 = fig.add_subplot(gs[:1, :])
    residuals.plot(ax=ax1)
    ax1.set_ylabel('value')
    ax1.set_title('Residual values')

    # plot histogram and distribution
    res_mean, res_std = np.mean(residuals.univariate_values()), np.std(residuals.univariate_values())
    res_min, res_max = min(residuals.univariate_values()), max(residuals.univariate_values())
    x = np.linspace(res_min, res_max, 100)
    ax2 = fig.add_subplot(gs[1:, 1:])
    plot_hist(residuals, bins=num_bins, ax=ax2)
    ax2.plot(x, norm(res_mean, res_std).pdf(x) * len(residuals) * (res_max - res_min) / num_bins)
    ax2.yaxis.set_major_locator(plt.MaxNLocator(integer=True))
    ax2.set_title('Distribution')
    ax2.set_ylabel('count')
    ax2.set_xlabel('value')

    # plot ACF
    ax3 = fig.add_subplot(gs[1:, :1])
    plot_acf(residuals, axis=ax3)
    ax3.set_ylabel('ACF value')
    ax3.set_xlabel('lag')
    ax3.set_title('ACF')


def predict(self, features, num_predictions_per_location=1):

"""Mask Head.

Contains Mask prediction head classes for different meta architectures.
All the mask prediction heads have a predict function that receives the
`features` as the first argument and returns `mask_predictions`.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
from six.moves import range
import tensorflow.compat.v1 as tf
import tf_slim as slim

from object_detection.predictors.heads import head
from object_detection.utils import ops

class MaskRCNNMaskHead(head.Head):
  """Mask RCNN mask prediction head.

  Please refer to Mask RCNN paper:
  http://example.com
  """

  def __init__(self,
               num_classes,
               conv_hyperparams_fn=None,
               mask_height=14,
               mask_width=14,
               mask_prediction_num_conv_layers=2,
               mask_prediction_conv_depth=256,
               masks_are_class_agnostic=False,
               convolve_then_upsample=False):
    """Constructor.

    Args:
      num_classes: number of classes.  Note that num_classes *does not*
        include the background category, so if groundtruth labels take values
        in {0, 1, .., K-1}, num_classes=K (and not K+1, even though the
        assigned classification targets can range from {0,... K}).
      conv_hyperparams_fn: A function to generate tf-slim arg_scope with
        hyperparameters for convolution ops.
      mask_height: Desired output mask height. The default value is 14.
      mask_width: Desired output mask width. The default value is 14.
      mask_prediction_num_conv_layers: Number of convolution layers applied to
        the image_features in mask prediction branch.
      mask_prediction_conv_depth: The depth for the first conv2d_transpose op
        applied to the image_features in the mask prediction branch. If set
        to 0, the depth of the convolution layers will be automatically chosen
        based on the number of object classes and the number of channels in the
        image features.
      masks_are_class_agnostic: Boolean determining if the mask-head is
        class-agnostic or not.
      convolve_then_upsample: Whether to apply convolutions on mask features
        before upsampling using nearest neighbor resizing. Otherwise, mask
        features are resized to [`mask_height`, `mask_width`] using bilinear
        resizing before applying convolutions.

    Raises:
      ValueError: conv_hyperparams_fn is None.
    """
    super(MaskRCNNMaskHead, self).__init__()
    self._num_classes = num_classes
    self._conv_hyperparams_fn = conv_hyperparams_fn
    self._mask_height = mask_height
    self._mask_width = mask_width
    self._mask_prediction_num_conv_layers = mask_prediction_num_conv_layers
    self._mask_prediction_conv_depth = mask_prediction_conv_depth
    self._masks_are_class_agnostic = masks_are_class_agnostic
    self._convolve_then_upsample = convolve_then_upsample
    if conv_hyperparams_fn is None:
      raise ValueError('conv_hyperparams_fn is None.')

  def _get_mask_predictor_conv_depth(self,
                                     num_feature_channels,
                                     num_classes,
                                     class_weight=3.0,
                                     feature_weight=2.0):
    """Computes the depth of the mask predictor convolutions.

    Computes the depth of the mask predictor convolutions given feature channels
    and number of classes by performing a weighted average of the two in
    log space to compute the number of convolution channels. The weights that
    are used for computing the weighted average do not need to sum to 1.

    Args:
      num_feature_channels: An integer containing the number of feature
        channels.
      num_classes: An integer containing the number of classes.
      class_weight: Class weight used in computing the weighted average.
      feature_weight: Feature weight used in computing the weighted average.

    Returns:
      An integer containing the number of convolution channels used by mask
        predictor.
    """
    num_feature_channels_log = math.log(float(num_feature_channels), 2.0)
    num_classes_log = math.log(float(num_classes), 2.0)
    weighted_num_feature_channels_log = (
        num_feature_channels_log * feature_weight)
    weighted_num_classes_log = num_classes_log * class_weight
    total_weight = feature_weight + class_weight
    num_conv_channels_log = round(
        (weighted_num_feature_channels_log + weighted_num_classes_log) /
        total_weight)
    return int(math.pow(2.0, num_conv_channels_log))

"""Performs mask prediction.

    Args:
      features: A float tensor of shape [batch_size, height, width, channels]
        containing features for a batch of images.
      num_predictions_per_location: Int containing number of predictions per
        location.

    Returns:
      instance_masks: A float tensor of shape
          [batch_size, 1, num_classes, mask_height, mask_width].

    Raises:
      ValueError: If num_predictions_per_location is not 1.
    """
    if num_predictions_per_location != 1:
      raise ValueError('Only num_predictions_per_location=1 is supported')
    num_conv_channels = self._mask_prediction_conv_depth
    if num_conv_channels == 0:
      num_feature_channels = features.get_shape().as_list()[3]
      num_conv_channels = self._get_mask_predictor_conv_depth(
          num_feature_channels, self._num_classes)
    with slim.arg_scope(self._conv_hyperparams_fn()):
      if not self._convolve_then_upsample:
        features = tf.image.resize_bilinear(
            features, [self._mask_height, self._mask_width],
            align_corners=True)
      for _ in range(self._mask_prediction_num_conv_layers - 1):
        features = slim.conv2d(
            features,
            num_outputs=num_conv_channels,
            kernel_size=[3, 3])
      if self._convolve_then_upsample:
        # Replace Transposed Convolution with a Nearest Neighbor upsampling step
        # followed by 3x3 convolution.
        height_scale = self._mask_height // features.shape[1].value
        width_scale = self._mask_width // features.shape[2].value
        features = ops.nearest_neighbor_upsampling(
            features, height_scale=height_scale, width_scale=width_scale)
        features = slim.conv2d(
            features,
            num_outputs=num_conv_channels,
            kernel_size=[3, 3])

      num_masks = 1 if self._masks_are_class_agnostic else self._num_classes
      mask_predictions = slim.conv2d(
          features,
          num_outputs=num_masks,
          activation_fn=None,
          normalizer_fn=None,
          kernel_size=[3, 3])
      return tf.expand_dims(
          tf.transpose(mask_predictions, perm=[0, 3, 1, 2]),
          axis=1,
          name='MaskPredictor')

class ConvolutionalMaskHead(head.Head):
  """Convolutional class prediction head."""

  def __init__(self,
               is_training,
               num_classes,
               use_dropout,
               dropout_keep_prob,
               kernel_size,
               use_depthwise=False,
               mask_height=7,
               mask_width=7,
               masks_are_class_agnostic=False):
    """Constructor.

    Args:
      is_training: Indicates whether the BoxPredictor is in training mode.
      num_classes: Number of classes.
      use_dropout: Option to use dropout or not.  Note that a single dropout
        op is applied here prior to both box and class predictions, which stands
        in contrast to the ConvolutionalBoxPredictor below.
      dropout_keep_prob: Keep probability for dropout.
        This is only used if use_dropout is True.
      kernel_size: Size of final convolution kernel.  If the
        spatial resolution of the feature map is smaller than the kernel size,
        then the kernel size is automatically set to be
        min(feature_width, feature_height).
      use_depthwise: Whether to use depthwise convolutions for prediction
        steps. Default is False.
      mask_height: Desired output mask height. The default value is 7.
      mask_width: Desired output mask width. The default value is 7.
      masks_are_class_agnostic: Boolean determining if the mask-head is
        class-agnostic or not.

    Raises:
      ValueError: if min_depth > max_depth.
    """
    super(ConvolutionalMaskHead, self).__init__()
    self._is_training = is_training
    self._num_classes = num_classes
    self._use_dropout = use_dropout
    self._dropout_keep_prob = dropout_keep_prob
    self._kernel_size = kernel_size
    self._use_depthwise = use_depthwise
    self._mask_height = mask_height
    self._mask_width = mask_width
    self._masks_are_class_agnostic = masks_are_class_agnostic

  def predict(self, features, num_predictions_per_location):
    """Predicts boxes.

    Args:
      features: A float tensor of shape [batch_size, height, width, channels]
        containing image features.
      num_predictions_per_location: Number of box predictions to be made per
        spatial location.

    Returns:
      mask_predictions: A float tensors of shape
        [batch_size, num_anchors, num_masks, mask_height, mask_width]
        representing the mask predictions for the proposals.
    """
    image_feature = features
    # Add a slot for the background class.
    if self._masks_are_class_agnostic:
      num_masks = 1
    else:
      num_masks = self._num_classes
    num_mask_channels = num_masks * self._mask_height * self._mask_width
    net = image_feature
    if self._use_dropout:
      net = slim.dropout(net, keep_prob=self._dropout_keep_prob)
    if self._use_depthwise:
      mask_predictions = slim.separable_conv2d(
          net, None, [self._kernel_size, self._kernel_size],
          padding='SAME', depth_multiplier=1, stride=1,
          rate=1, scope='MaskPredictor_depthwise')
      mask_predictions = slim.conv2d(
          mask_predictions,
          num_predictions_per_location * num_mask_channels,
          [1, 1],
          activation_fn=None,
          normalizer_fn=None,
          normalizer_params=None,
          scope='MaskPredictor')
    else:
      mask_predictions = slim.conv2d(
          net,
          num_predictions_per_location * num_mask_channels,
          [self._kernel_size, self._kernel_size],
          activation_fn=None,
          normalizer_fn=None,
          normalizer_params=None,
          scope='MaskPredictor')
    batch_size = features.get_shape().as_list()[0]
    if batch_size is None:
      batch_size = tf.shape(features)[0]
    mask_predictions = tf.reshape(
        mask_predictions,
        [batch_size, -1, num_masks, self._mask_height, self._mask_width])
    return mask_predictions

# TODO(alirezafathi): See if possible to unify Weight Shared with regular
# convolutional mask head.
class WeightSharedConvolutionalMaskHead(head.Head):
  """Weight shared convolutional mask prediction head."""

  def __init__(self,
               num_classes,
               kernel_size=3,
               use_dropout=False,
               dropout_keep_prob=0.8,
               mask_height=7,
               mask_width=7,
               masks_are_class_agnostic=False):
    """Constructor.

    Args:
      num_classes: number of classes.  Note that num_classes *does not*
        include the background category, so if groundtruth labels take values
        in {0, 1, .., K-1}, num_classes=K (and not K+1, even though the
        assigned classification targets can range from {0,... K}).
      kernel_size: Size of final convolution kernel.
      use_dropout: Whether to apply dropout to class prediction head.
      dropout_keep_prob: Probability of keeping activiations.
      mask_height: Desired output mask height. The default value is 7.
      mask_width: Desired output mask width. The default value is 7.
      masks_are_class_agnostic: Boolean determining if the mask-head is
        class-agnostic or not.
    """
    super(WeightSharedConvolutionalMaskHead, self).__init__()
    self._num_classes = num_classes
    self._kernel_size = kernel_size
    self._use_dropout = use_dropout
    self._dropout_keep_prob = dropout_keep_prob
    self._mask_height = mask_height
    self._mask_width = mask_width
    self._masks_are_class_agnostic = masks_are_class_agnostic

  def predict(self, features, num_predictions_per_location):
    """Predicts boxes.

    Args:
      features: A float tensor of shape [batch_size, height, width, channels]
        containing image features.
      num_predictions_per_location: Number of box predictions to be made per
        spatial location.

    Returns:
      mask_predictions: A tensor of shape
        [batch_size, num_anchors, num_classes, mask_height, mask_width]
        representing the mask predictions for the proposals.
    """
    mask_predictions_net = features
    if self._masks_are_class_agnostic:
      num_masks = 1
    else:
      num_masks = self._num_classes
    num_mask_channels = num_masks * self._mask_height * self._mask_width
    if self._use_dropout:
      mask_predictions_net = slim.dropout(
          mask_predictions_net, keep_prob=self._dropout_keep_prob)
    mask_predictions = slim.conv2d(
        mask_predictions_net,
        num_predictions_per_location * num_mask_channels,
        [self._kernel_size, self._kernel_size],
        activation_fn=None, stride=1, padding='SAME',
        normalizer_fn=None,
        scope='MaskPredictor')
    batch_size = features.get_shape().as_list()[0]
    if batch_size is None:
      batch_size = tf.shape(features)[0]
    mask_predictions = tf.reshape(
        mask_predictions,
        [batch_size, -1, num_masks, self._mask_height, self._mask_width])
    return mask_predictions