"""This dummy dataset is intended for testing and debugging purposes, not for
benchmarking. If your approach can fit this dataset it means that it is able to learn
how to estimate RUL. It does not mean it is good at it. """

from typing import Tuple, List, Optional, Union

import numpy as np
from sklearn import preprocessing  # type: ignore

from rul_datasets import utils
from rul_datasets.reader import AbstractReader, scaling


class DummyReader(AbstractReader):
    """
    This reader represents a simple, small dummy dataset that can be uses to test or
    debug RUL estimation approaches. It contains ten runs for each split with a
    single feature which makes it easy to hold in memory even on low-end computers.
    The dataset is so simple that it can be sufficiently fit by a three-layer
    perceptron in less than 50 epochs.

    Each run is randomly generated by sampling a run length between 90 and 110 time
    steps and creating a piece-wise linear RUL function `y(t)` with a maximum value of
    `max_rul`. The feature `x(t)` is then calculated as:

    ```python
    x(t) = exp(-0.05 * y(t) + N(offset, 0.01)) + N(0, noise_factor)
    ```

    where `N(loc, scale)` is a function drawing a sample from a normal distribution
    with a mean of `loc` and a standard deviation of `scale`. The `dev`, `val` and
    `test` splits are all generated the same way with a different fixed random seed.
    This makes generating the dataset deterministic.

    The dummy dataset contains two sub-datasets. The first has uses an `offset` of
    0.5 and a `noise_factor` of 0.01. The second uses an `offset` of 0.75 and a
    `noise_factor` of 0.02. Both use a default window size of 10 and are min-max
    scaled between -1 and 1 with a scaler fitted on the `dev` split.

    Examples:
        >>> import rul_datasets
        >>> fd1 = rul_datasets.reader.DummyReader(fd=1)
        >>> features, labels = fd1.load_split("dev")
        >>> features[0].shape
        (81, 10, 1)
    """

    _FDS = [1, 2]

    _DEFAULT_WINDOW_SIZE = 10
    _NOISE_FACTOR = {1: 0.01, 2: 0.02}
    _OFFSET = {1: 0.5, 2: 0.75}
    _SPLIT_SEED = {"dev": 42, "val": 1337, "test": 101}

    def __init__(
        self,
        fd: int,
        window_size: Optional[int] = None,
        max_rul: Optional[int] = 50,
        percent_broken: Optional[float] = None,
        percent_fail_runs: Optional[Union[float, List[int]]] = None,
        truncate_val: bool = False,
        truncate_degraded_only: bool = False,
    ):
        """
        Create a new dummy reader for one of the two sub-datasets. The maximum RUL
        value is set to 50 by default. Please be aware that changing this value will
        lead to different features, too, as they are calculated based on the RUL
        values.

        For more information about using readers, refer to the [reader]
        [rul_datasets.reader] module page.

        Args:
            fd: Index of the selected sub-dataset
            window_size: Size of the sliding window. Default defined per sub-dataset.
            max_rul: Maximum RUL value of targets.
            percent_broken: The maximum relative degradation per time series.
            percent_fail_runs: The percentage or index list of available time series.
            truncate_val: Truncate the validation data with `percent_broken`, too.
            truncate_degraded_only: Only truncate the degraded part of the data
                                    (< max RUL).
        """
        super(DummyReader, self).__init__(
            fd,
            window_size,
            max_rul,
            percent_broken,
            percent_fail_runs,
            truncate_val,
            truncate_degraded_only,
        )

        features, _ = self._generate_split("dev")
        scaler = preprocessing.MinMaxScaler(feature_range=(-1, 1))
        self.scaler = scaling.fit_scaler(features, scaler)

    @property
    def dataset_name(self) -> str:
        return "xjtu-sy"

    @property
    def fds(self) -> List[int]:
        """Indices of available sub-datasets."""
        return self._FDS

    def default_window_size(self, fd: int) -> int:
        return self._DEFAULT_WINDOW_SIZE

    def prepare_data(self) -> None:
        """This function has no effect as there is nothing to prepare."""
        pass

    def load_complete_split(
        self, split: str, alias: str
    ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        features, targets = self._generate_split(split, alias)
        features = scaling.scale_features(features, self.scaler)

        return features, targets

    def _generate_split(
        self, split: str, alias: Optional[str] = None
    ) -> Tuple[List[np.ndarray], List[np.ndarray]]:
        alias = split if alias is None else alias
        features = []
        targets = []
        rng = np.random.default_rng(self._SPLIT_SEED[split])
        for i in range(10):
            t = self._generate_targets(rng)
            f = self._generate_features(rng, t)
            f = utils.extract_windows(f, self.window_size)
            t = t[: -(self.window_size - 1)]
            features.append(f)
            targets.append(t)
        if alias == "test":
            features, targets = self._truncate_test_split(rng, features, targets)

        return features, targets

    def _generate_targets(self, rng):
        length = rng.integers(90, 110)
        t = np.clip(np.arange(length, 1, -1), a_min=0, a_max=self.max_rul)
        t = t.astype(float)

        return t

    def _generate_features(self, rng, targets):
        steady = -0.05 * targets + self._OFFSET[self.fd] + rng.normal() * 0.01
        noise = rng.normal(size=targets.shape) * self._NOISE_FACTOR[self.fd]
        f = np.exp(steady) + noise

        return f[:, None]

    def _truncate_test_split(self, rng, features, targets):
        """Extract a single window from a random position of the time series."""
        for i in range(len(features)):
            run_len = len(features[i])
            cutoff = rng.integers(run_len // 2, run_len - 1)
            features[i] = features[i][None, cutoff]
            targets[i] = targets[i][None, cutoff]

        return features, targets
