import numpy as np
import tensorflow as tf
from tensorflow.keras.utils import Sequence

class TimeSeriesDataset(Sequence):
    def __init__(self, sequences, next_points, batch_size=32, shuffle=True):
        self.sequences = sequences
        self.next_points = next_points
        self.batch_size = batch_size
        self.shuffle = shuffle
        self.indices = np.arange(len(self.sequences))
        self.on_epoch_end()

    def __len__(self):
        return int(np.ceil(len(self.sequences) / self.batch_size))

    def __getitem__(self, index):
        batch_indices = self.indices[index*self.batch_size:(index+1)*self.batch_size]
        X = self.sequences[batch_indices]
        y = self.next_points[batch_indices]
        return X, y

    def on_epoch_end(self):
        if self.shuffle:
            np.random.shuffle(self.indices)