"""Quality assurance based on statistical methods."""

from __future__ import annotations

from datetime import timedelta
from typing import TYPE_CHECKING

import numpy as np
import pandas as pd

from avoca.qa_class.zscore import ExtremeValues

if TYPE_CHECKING:

    from avoca.utils.torch_models import MultipleRegressionModel


class RollingWindow(ExtremeValues):
    """Detect in rolling windows.

    The method is based on outliers in a rolling window using the median and standard deviation.
    The training is done directly on the fitted data.

    :param variable: The variable to check for extreme values.
    :param threshold: The threshold for the z-score. To flag values.
    :param use_log_normal: If True, the log of the values will be used to calculate the z-score.
        This can be useful if the values are log-normal distributed.
    :param only_greater: If True, only values greater than the threshold will be flagged.
        The values lower than the negative threshold will not be flagged.
        By default, this is True if use_log_normal is True, and False otherwise.
    """

    require_datetime_index = True

    rolling_window: timedelta

    def __init__(
        self,
        *args,
        rolling_window: timedelta = timedelta(days=7),
        threshold: float = 1.5,
        **kwargs,
    ):
        super().__init__(*args, threshold=threshold, **kwargs)
        self.rolling_window = rolling_window

    def fit(self, df: pd.DataFrame):

        self.check_columns_or_raise(df, columns=self._stats_columns)

        self.df_train = df[self._stats_columns]

    def assign(self, df: pd.DataFrame) -> dict[str, pd.Index]:
        df = df[self._stats_columns]
        df = self._clean_data(df)
        if self.use_log_normal:
            # Replace <=0 with NaN
            df = df.where(df > 0, np.nan)
            df = df.map(lambda x: np.log(x))

        rolling = df.rolling(window=self.rolling_window)
        means = rolling.median()
        stds = rolling.std()

        self.rolling_median = means
        self.rolling_std = stds

        thresholds = means + stds * self.threshold

        df_fail = df > thresholds
        if not self.only_greater:
            df_fail = df_fail | (df < (means - stds * self.threshold))

        out_dict = {}
        for compound in self.compounds:
            col = (compound, self.variable)
            this_c_fail = df_fail[col]
            out_dict[compound] = this_c_fail.loc[this_c_fail].index

        return out_dict

    def plot(self):

        import matplotlib.pyplot as plt

        fig, axes = plt.subplots(
            len(self.compounds), 1, figsize=(6, 3 * len(self.compounds)), sharex=True
        )

        x = self.dt if hasattr(self, "dt") else self.df_train.index
        x = pd.Series(x, index=self.df_train.index)

        outliers = self.assign(self.df_train)

        for i, compound in enumerate(self.compounds):
            ax = axes[i]
            col = (compound, self.variable)
            ax.scatter(
                x,
                self.df_train[col],
                s=1,
                label="darkblue",
            )
            median = self.rolling_median[col]
            std = self.rolling_std[col]
            top, bottom = median + std * self.threshold, median - std * self.threshold

            ax.fill_between(
                x,
                top,
                bottom,
                color="lightgray",
                label="Rolling threshold",
                alpha=0.5,
            )

            outlier_indices = outliers[compound]
            ax.scatter(
                x.loc[outlier_indices],
                self.df_train.loc[outlier_indices, col],
                s=10,
                marker="x",
                color="red",
                label="Extreme values",
            )
            ax.set_title(
                f"{compound} +- {self.threshold} std",
                # Under teh top line
                y=0.8,
            )
            ax.tick_params(axis="x", rotation=25)

        return fig, axes
