def load_text_data(file_path):
    """
    加载文本数据文件
    """
    text_data = ""
    with open(file_path, "r", encoding="utf-8") as file:
        text_data = file.read()

    return text_data


def load_text_data_and_split(filepath, train_ratio=0.9):
    """
    加载文本数据文件并进行训练/验证集划分
    """
    text_data = load_text_data(filepath)

    # Train/Validation 划分比例
    split_idx = int(train_ratio * len(text_data))
    train_data = text_data[:split_idx]
    val_data = text_data[split_idx:]

    return train_data, val_data


import torch
from torch.utils.data import Dataset


class GPTDataset(Dataset):
    """
    使用滑动窗口构造GPT预训练数据集

    参数:
        txt (str): 输入文本字符串
        tokenizer: 分词器对象，用于将文本转换为token ID
        max_length (int): 每个样本的最大长度（输入序列长度）
        stride (int): 滑动窗口的步长，控制样本之间的重叠程度
    """

    def __init__(self, txt, tokenizer, max_length, stride):
        # 初始化一个列表来保存所有样本(input_ids, target_ids)对
        self.samples = []

        # 将整个文本编码为token ID序列
        token_ids = tokenizer.encode(txt, allowed_special={"<|endoftext|>"})

        # 使用滑动窗口将文本分割为长度为最大长度的重叠序列
        for i in range(0, len(token_ids) - max_length, stride):
            # 当前窗口的token作为输入数据
            input_ids = token_ids[i : i + max_length]

            # 向右偏移一位的token作为监督数据
            target_ids = token_ids[i + 1 : i + max_length + 1]

            # 转换为张量并添加到样本列表中
            self.samples.append((torch.tensor(input_ids), torch.tensor(target_ids)))

    def __len__(self):
        # 返回数据集中样本的总数
        return len(self.samples)

    def __getitem__(self, idx):
        # 获取单个样本，形式为(input_ids, target_ids)
        return self.samples[idx]


import lightning as L
from torch.utils.data import DataLoader


class GPTDataModule(L.LightningDataModule):

    def __init__(
        self,
        batch_size,
        tokenizer,
        train_data_file,
        max_length,
        stride,
        train_ratio=0.9,
    ):
        super().__init__()

        # 初始化数据模块的属性
        self.batch_size = batch_size  # 设置批次大小
        self.tokenizer = tokenizer  # 文本转换器
        self.train_data_file = train_data_file  # 训练数据文件路径
        self.max_length = max_length  # 最大序列长度
        self.stride = stride  # 滑动窗口步长
        self.train_ratio = train_ratio  # 训练集比例

        # 初始化数据集
        self.train_dataset: GPTDataset  # 训练数据集
        self.val_dataset: GPTDataset  # 评估数据集
        self.test_dataset: GPTDataset  # 测试数据集

    def prepare_data(self):
        """
        准备数据的方法
        用于下载数据集或进行一次性数据预处理操作
        """
        pass

    def setup(self, stage=None):
        """
        设置数据集的方法
        """

        # 1️⃣ 加载文本数据并进行训练集和验证集的划分
        train_data, val_data = load_text_data_and_split(
            self.train_data_file, self.train_ratio
        )

        # 2️⃣ 初始化训练数据集、评估集和测试集
        self.train_dataset = GPTDataset(
            train_data, self.tokenizer, max_length=self.max_length, stride=self.stride
        )
        self.val_dataset = GPTDataset(
            val_data, self.tokenizer, max_length=self.max_length, stride=self.stride
        )
        self.test_dataset = self.val_dataset

    def train_dataloader(self):
        """
        返回训练数据加载器

        Returns:
            DataLoader: 训练数据的DataLoader对象
        """
        return DataLoader(self.train_dataset, batch_size=self.batch_size, shuffle=True)

    def val_dataloader(self):
        """
        返回验证数据加载器

        Returns:
            DataLoader: 验证数据的DataLoader对象
        """
        return DataLoader(self.val_dataset, batch_size=self.batch_size)

    def test_dataloader(self):
        """
        返回测试数据加载器

        Returns:
            DataLoader: 测试数据的DataLoader对象
        """
        return DataLoader(self.test_dataset, batch_size=self.batch_size)
