"""
split the data into train and test
"""
from abc import abstractmethod
from typing import Tuple, Iterator
import pandas as pd

from commons.data_provision import LoadConfig
from commons.data_processing import split_train_val


class Split:
    """Represents a single split of the data from a set generated by a Splitter,
    Uses the indices to return the split data

    :param train_index: indices of the training set
    :param test_index: indices of the testing set
    """
    def __init__(self, train_index, test_index):
        self.train_index = train_index
        self.test_index = test_index

    @abstractmethod
    def data(self, data: pd.DataFrame) -> Tuple[pd.DataFrame, pd.DataFrame]:
        return data.loc[self.train_index].reset_index(drop=True), data.loc[self.test_index].reset_index(drop=True)


class Splitter:
    """Abstract base class for data Split generators"""
    @abstractmethod
    def split(self, data: pd.DataFrame, load_config: LoadConfig) -> Iterator[Split]:
        pass


class RandomSplitter(Splitter):
    """
    Splits the data into train and test on a single date. The split_date is included in the test set.
    Split_date defaults to end_date
    """
    def __init__(self, split_size=None):
        self.split_size = split_size

    def split(self, data: pd.DataFrame, load_config: LoadConfig) -> Iterator[Split]:
        splitting_size = self.split_size if self.split_size else load_config.split_size
        train_index, test_index = split_train_val(raw_data=data,
                                                  val_size=splitting_size,
                                                  random_state=42,
                                                  return_index=True)
        yield Split(train_index, test_index)
