"""
Model selection utilities
-------------------------
Utilities that help in model selection e.g. by splitting a dataset.
"""

from typing import Sequence, Optional, Union, Tuple
from darts import TimeSeries

MODEL_AWARE = 'model-aware'
SIMPLE = 'simple'


class SplitTimeSeriesSequence(Sequence):
    """
    This class is primarily meant to be instantiated from ``train_test_split()`` function.
    """
    def __init__(self,
                 type: str,
                 data: Sequence[TimeSeries],
                 test_size: Optional[Union[float, int]] = 0.25,
                 axis: Optional[int] = 0,
                 input_size: Optional[int] = None,
                 horizon: Optional[int] = None,
                 vertical_split_type: Optional[str] = SIMPLE):

        if type not in ['train', 'test']:
            raise AttributeError('Value for type parameter should be either `train` or `test`')

        if not data:
            raise AttributeError('The `data` parameter cannot be empty.')

        if axis not in [0, 1]:
            raise AttributeError('The `axis` parameter should be either 0 or 1.')

        if vertical_split_type not in [SIMPLE, MODEL_AWARE]:
            raise AttributeError('`vertical_split_type` can be eiter `simple` or `model-aware`.')

        if axis == 1 and vertical_split_type == MODEL_AWARE and (horizon == 0 or input_size == 0):
                raise AttributeError("You need to provide non-zero `horizon` and `input_size` parameters when axis=1")

        self.type = type
        self.data = data
        self.test_size = test_size
        self.axis = axis
        self.input_size = input_size
        self.horizon = horizon
        self.vertical_split_type = vertical_split_type
        self._horizontal_split_index = None

    def _get_horizontal_split_index(self):
        if not self._horizontal_split_index:
            if 0 < self.test_size < 1:
                self._horizontal_split_index = int(len(self.data) * (1 - self.test_size))
            else:
                self._horizontal_split_index = len(self.data) - self.test_size

        return self._horizontal_split_index

    def _get_vertical_split_indices(self, ts_length):

        if self.vertical_split_type == SIMPLE:
            if 0 < self.test_size < 1:
                test_size = int(ts_length * self.test_size)
            else:
                test_size = self.test_size

            test_start_index = ts_length - test_size
            train_end_index = test_start_index

            if test_start_index < 0:
                raise AttributeError("`test_size` is bigger then timeseries length")

        else: # model-aware split

            # an example on how these are calculated:
            # test_size = 4
            # input_size = 5 (-)
            # horizon = 3 (+)
            # len(data) = 16 (*)

            # 0         5         10          15 <- index
            # * * * * * * * * * * * * * * * *
            #                 - - - - - + + +   < test sample 1
            #               - - - - - + + +     < test sample 2
            #             - - - - - + + +       < test sample 3
            #           - - - - - + + +         < test sample 4
            #           ^         ^
            #           |         train_end_index = 10 (note the [:train_end_index] slice stops at 9)
            #           test_start_index = 5

            if 0 < self.test_size < 1:
                test_size = int((ts_length - self.horizon) * self.test_size)
            else:
                test_size = self.test_size

            train_end_index = ts_length - self.horizon - test_size + 1
            test_start_index = ts_length - self.horizon - self.input_size - test_size + 1

            if train_end_index < 0:
                train_end_index = 0

            if train_end_index < self.input_size + self.horizon:
                # Note: note that we don't have to check for test_start_index < 0 because train_end_index will always
                # raise exception first.
                raise AttributeError("Not enough data to create training and test sets")

        return train_end_index, test_start_index

    def __getitem__(self, i: int) -> TimeSeries:
        if self.axis == 0:
            split_index = self._get_horizontal_split_index()
            if self.type == 'train':
                if i >= split_index:
                    raise IndexError('Exceeded the size of the training sequence.')
                return self.data[i]
            else:
                if i + split_index >= len(self.data):
                    raise IndexError('Exceeded the size of the test sequence.')
                return self.data[split_index + i]
        else: # axis == 1
            train_end_index, test_start_index = self._get_vertical_split_indices(len(self.data[i]))
            if self.type == 'train':
                return self.data[i][:train_end_index]
            else:
                return self.data[i][test_start_index:]

    def __len__(self):
        if self.axis == 0:
            split_index = self._get_horizontal_split_index()
            if self.type == 'train':
                return split_index
            else:
                return len(self.data) - split_index
        else:
            return len(self.data)

    @classmethod
    def make_splitter(cls,
                      data: Union[TimeSeries, Sequence[TimeSeries]],
                      test_size: Optional[Union[float, int]] = 0.25,
                      axis: Optional[int] = 0,
                      input_size: Optional[int] = 0,
                      horizon: Optional[int] = 0,
                      vertical_split_type: Optional[str] = SIMPLE,
                      lazy: bool = False
                      ) -> Union[Tuple[TimeSeries, TimeSeries], Tuple[Sequence[TimeSeries], Sequence[TimeSeries]]]:

        if not isinstance(data, Sequence):
            axis = 1
            data = [data]  # convert to sequence for unified processing later
            single_timeseries = True
        else:
            single_timeseries = False

        train_set = cls(type='train', data=data, test_size=test_size, axis=axis, input_size=input_size, horizon=horizon,
                        vertical_split_type=vertical_split_type)

        test_set = cls(type='test', data=data, test_size=test_size, axis=axis, input_size=input_size, horizon=horizon,
                       vertical_split_type=vertical_split_type)

        if single_timeseries:
            return train_set[0], test_set[0]
        else:
            if lazy:
                return train_set, test_set
            else:
                return list(train_set), list(test_set)


def train_test_split(
        data: Union[TimeSeries, Sequence[TimeSeries]],
        test_size: Optional[Union[float, int]] = 0.25,
        axis: Optional[int] = 0,
        input_size: Optional[int] = 0,
        horizon: Optional[int] = 0,
        vertical_split_type: Optional[str] = SIMPLE,
        lazy: bool = False
        ) -> Union[Tuple[TimeSeries, TimeSeries], Tuple[Sequence[TimeSeries], Sequence[TimeSeries]]]:
    """
    Splits the dataset into training and test dataset. Supports splitting along the sample axis and time axis.

    If the input type is single TimeSeries, then only splitting over time axis is available, thus ``input_size`` and
    ``horizon`` have to be provided.

    When splitting over the time axis, splitter tries to greedy satisfy the requested test set size, i.e. when one of
    the timeseries in sequence is too small, all samples will go to the test set and the exception will be raised.


    Time axis split with ``model-aware`` split type enabled tries to reclaim as much datapoints for training as
    possible by partially overlapping test set with training set. This is possible because only forecasted part of
    test set cannot be used for training. The formula to calculate last available timestep for training set is
    following:

        train end index = ts_length - self.horizon - test_size

    And the formula to calculate the first timestep of test dataset is following:

        test start index = timeseries length - horizon - input_size - test_size + 1

    Parameters
    ----------
    data
        original dataset to split into training and test

    test_size
        size of the test set. If the value is between 0 and 1, parameter is treated as a split proportion. Otherwise
        it is treated as a absolute number of samples from each timeseries that will be in the test set. [default = 0.25]

    axis
        Axis to split the dataset on. When 0 (default) it is split on samples. Otherwise, if ``axis = 1``,
        timeseries are split along time axis. Note that for single timeseries the default option is 1 (0 makes
        no sense). [default: 0 for sequence of timeseries, 1 for timeseries]

    input_size
        size of the input. Only valid with ``vertical_split_type == 'model-aware'``. [default: None]

    horizon
        forecast horizon. Only valid with ``vertical_split_type == 'model-aware'``. [default: None]

    vertical_split_type
        can be either ``simple``, where the exact number from test size will be deducted from timeseries for test set and
        remaining will go to training set; or ``model-aware``, where you have to provide ``input_size`` and ``horizon``
        as well. Note, that second option is more efficient timestep-wise, since training and test sets will be
        partially overlapping. [default: ``simple``]

    lazy
        by default, train and test datasets are returned as a list of timeseries. This may be memory
        inefficient if dataset is large, so setting this flag allows instead to return a ``Sequence`` object
        loading the data lazily. Warning: turning ``lazy`` on disables some sanity checks for the datasets
        that may result in exceptions during sample generation. [default: False]

    Returns
    -------
    tuple of two Sequence[TimeSeries], or tuple of two Timeseries
        Training and test datasets tuple.
    """

    return SplitTimeSeriesSequence.make_splitter(data, test_size, axis, input_size, horizon, vertical_split_type, lazy)
