class TimeSeriesDataset:
    def __init__(self, look_back=1, train_size=0.67):
        self.look_back = look_back
        self.train_size = train_size

    def load_data(self):
        df = sns.load_dataset('flights')
        df['Date']  = pd.to_datetime(df['year'].astype(str) + '-' + df['month'].astype(str) + '-01')
        df = df.set_index('Date')
        df = df.drop(['year', 'month'], axis=1)
        df = df.values.astype(float)
        df = self.MinMaxScaler(df)
        train_size = int(len(df) * self.train_size)
        train, test = df[0:train_size,:], df[train_size:len(df),:]
        return train, test

    def MinMaxScaler(self, data):
        numerator = data - np.min(data, 0)
        denominator = np.max(data, 0) - np.min(data, 0)
        return numerator / (denominator + 1e-7)

    def create_dataset(self, dataset):
        dataX, dataY = []
        for i in range(len(dataset)-self.look_back):
            a = dataset[i:(i + self.look_back), 0]
            dataX.append(a)
            dataY.append(dataset[i + self.look_back, 0])
        return np.array(dataX), np.array(dataY)

    def get_train_test(self):
        train, test = self.load_data()
        trainX, trainY = self.create_dataset(train)
        testX, testY = self.create_dataset(test)
        return trainX, trainY, testX, testY