from config import *
from utils import *
import torch.utils.data as data
from torch.nn.utils.rnn import pad_sequence
import json
import torch
from model import get_padding_mask,get_subsequent_mask

class Dataset(data.Dataset):
    def __init__(self,type="train"):
        super().__init__()
        if type == 'train':
            file_path = TRAIN_SAMPLE_PATH
        elif type == 'val':
            file_path = VAL_SAMPLE_PATH
        # 读取文件
        with open(file_path, encoding='utf-8') as file:
            self.lines = json.loads(file.read())
        # 词表引入
        _, self.en_vocab2id = get_vocab('en')
        _, self.zh_vocab2id = get_vocab('zh')
    def __len__(self):
        return len(self.lines)
    def __getitem__(self, index):
        en_text, zh_text = self.lines[index]
        source = [self.en_vocab2id.get(v.lower(), UNK_ID) for v in divided_en(en_text)] # 英文转换成编号
        # 原代码没有过滤Jieba返回的纯空白token，空格不在中文词表中，
        # 因此会在训练和验证时被错误地转换为<unk>。
        # target = [self.zh_vocab2id.get(v.lower(), UNK_ID) for v in divided_zh(zh_text)]

        # 数据映射改进：与data_processor.py生成词表时的规则保持一致，
        # 只过滤空格、制表符等纯空白token，正常的中文、英文名称和标点均保留。
        target = [
            self.zh_vocab2id.get(v.lower(), UNK_ID)
            for v in divided_zh(zh_text)
            if v.strip()
        ] # 中文转换成编号
        return source, target, zh_text
    
    #数据对齐和整理
    #因为每个批次中不同的句子长度不一致，需要把一个批次中的所有句子的长度处理成一样的长度
    def collate_fn(self, batch):
        #利用 * 号操作符，可以将元组解压为列表
        batch_src, batch_tgt, tgt_text = zip(*batch)
        #处理encoder的输入
        src_x = pad_sequence([torch.LongTensor(src) for src in batch_src], True, PAD_ID)
        src_mask = get_padding_mask(src_x, PAD_ID)
        #处理decoder的输入 
        tgt_x = [torch.LongTensor([SOS_ID] + tgt) for tgt in batch_tgt]
        tgt_x = pad_sequence(tgt_x, True, PAD_ID)
        tgt_pad_mask = get_padding_mask(tgt_x, PAD_ID)
        tgt_subsqueent_mask = get_subsequent_mask(tgt_x.size(1))
        tgt_mask = tgt_pad_mask | tgt_subsqueent_mask
        tgt_mask = tgt_mask != 0
        #处理标签
        tgt_y = [torch.LongTensor(tgt+[EOS_ID]) for tgt in batch_tgt]
        tgt_y = pad_sequence(tgt_y, True, PAD_ID)
        return src_x, src_mask, tgt_x, tgt_mask, tgt_y, tgt_text
    

if __name__ == '__main__':
    val = Dataset('train')
    loader = data.DataLoader(val, batch_size=2, collate_fn=val.collate_fn)
    print(next(iter(loader)))