import numpy as np


class Dataset:
    def __init__(self, data_file: str = "pytorch_codebase.txt", max_length: int = 128):
        # Load the data in binary mode
        with open(data_file, "r") as f:
            data = list(f.read())
        self.vocab = np.array(sorted(set(data)))
        print(f"Data has {len(data)} characters, {len(self.vocab)} unique.")
        vocab_map = {ch: i for i, ch in enumerate(self.vocab)}
        self.index_data = np.array([vocab_map[ch] for ch in data], dtype=np.uint8)
        self.max_length = max_length

    def split(self, train_ratio: float = 0.7, val_ratio: float = 0.15):
        total_len = len(self.index_data)
        train_end = int(total_len * train_ratio)
        val_end = int(total_len * (train_ratio + val_ratio))

        train_data = self.index_data[:train_end]
        val_data = self.index_data[train_end:val_end]
        test_data = self.index_data[val_end:]

        return (
            Dataset.from_array(train_data, self.vocab, self.max_length),
            Dataset.from_array(val_data, self.vocab, self.max_length),
            Dataset.from_array(test_data, self.vocab, self.max_length),
        )

    @classmethod
    def from_array(cls, array: np.ndarray, vocab: np.ndarray, max_length: int):
        instance = cls.__new__(cls)
        instance.index_data = array
        instance.vocab = vocab
        instance.max_length = max_length
        return instance

    def __len__(self):
        return len(self.index_data) // self.max_length

    def __getitem__(self, idx):
        start = idx * self.max_length
        end = start + self.max_length
        return self.index_data[start:end]
