#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File         :datasets.py
@Description  :DataSets  for NLP_query
@Time         :2023/01/17 15:21:45
@Author       :KangQing
@Version      :1.0
'''


import sys
sys.path.append('./')
import json
import os
import jieba
from torch.utils.data import Dataset
import torch
import numpy as np
import logging
from transformers import InputFeatures
jieba.setLogLevel(logging.INFO)


class InputExample():
    def __init__(self,id:str,text_a:str,text_b:str=None,label:str=None) -> None:  
        self.id = id
        self.text_a = text_a
        self.text_b = text_b
        self.label = label
    def __str__(self) -> str:
        return json.dumps({'id':self.id,'text_a':self.text_a,'text_b':self.text_b,'label':self.label},indent=2,ensure_ascii=False)+'\n'


class QQR_data():
    def __init__(self,data_path='data') -> None:
        self.data_path = data_path
    
    def get_data(self,json_data_path):
        with open(json_data_path,'r',encoding='utf-8') as f:
            data = json.load(f,encoding='utf-8')
        
        examples = []
        for example in data:
            examples.append(InputExample(example['id'],example['query1'],example['query2'],example['label'] if example['label']!="" else None))
        return examples

    def get_labels(self):
        return ['0','1','2']
    
    def get_train_data(self):
        path = os.path.join(self.data_path,'KUAKE-QQR_{}.json'.format('train'))
        return self.get_data(path)
    
    def get_dev_data(self):
        path = os.path.join(self.data_path,'KUAKE-QQR_{}.json'.format('dev'))
        return self.get_data(path)
    
    def get_test_data(self):
        path = os.path.join(self.data_path,'KUAKE-QQR_{}.json'.format('test'))
        return self.get_data(path)
    

    
class QQRDataset(Dataset):
    def __init__(self,examples_list,labels_list,w2v_map,max_length):
        self.examples_list = examples_list
        self.label2id = {label:idx for idx,label in enumerate(labels_list)}
        self.id2label = {idx:label for idx,label in enumerate(labels_list)}
        self.w2v_map = w2v_map
        self.max_length = max_length
    
    def __len__(self):
        return len(self.examples_list)
    
    def _tokenize(self,text):
        token_list = list(jieba.cut(text))
        token_ids = []
        for token in token_list:
            if(token in self.w2v_map):
                token_ids.append(self.w2v_map.get(token))
            else:
                if(len(token)>1):
                    for character in token:
                        token_ids.append(self.w2v_map.get(token) if self.w2v_map.get(token)!=None else np.random.choice(len(self.w2v_map),1).item())
                else:
                    token_ids.append(np.random.choice(len(self.w2v_map),1).item())
        
        token_ids,attention_mask = self._pad_and_cut(token_ids)
        return token_ids,attention_mask
    
    def _pad_and_cut(self,token_ids):
        
        #Generate attention mask
        attention_mask = None
        
        if(len(token_ids)>self.max_length):
            token_ids = token_ids[:self.max_length]
            attention_mask = [1]*self.max_length
        else:
            attention_mask = [1]*len(token_ids)
            diff = self.max_length - len(token_ids)
            token_ids.extend([0]*diff)
            attention_mask.extend([0]*diff)
        
        return torch.tensor(token_ids,dtype=torch.long),torch.tensor(attention_mask,dtype=torch.long)
    
    
    def __getitem__(self,index):
        example = self.examples_list[index]
        idx = example.id
        text_a = example.text_a
        text_b = example.text_b
        if(example.label in self.label2id):
            label = self.label2id[example.label]
        else:
            label = 3
        
        text_a_inputs_id,text_a_attention_mask = self._tokenize(text_a)
        text_b_inputs_id,text_b_attention_mask = self._tokenize(text_b)
        
        label = torch.tensor(label,dtype=torch.long)

        
        return {
            'text_a_inputs_id':text_a_inputs_id,
            'text_b_inputs_id':text_b_inputs_id,
            'text_a_attention_mask':text_a_attention_mask,
            'text_b_attention_mask':text_b_attention_mask,
            'labels':label,
            'text_a':text_a,
            'text_b':text_b,
            'idx':idx
        }
        
class BertClassificationDataset(Dataset):
    
    def __init__(
        self,
        examples,
        tokenizer,
        label_list,
        max_length,
        processer=None
        ):
        super().__init__()
        
        self.examples = examples
        self.max_length = max_length
        self.tokenizer =tokenizer
        self.processor = processer
        
        self.label2id = {label:idx for idx,label in enumerate(label_list)}
        self.id2label = {idx:label for idx,label in enumerate(label_list)}
        
    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, index):
        example = self.examples[index]
        if(example.label in self.label2id):
            label = self.label2id[example.label]
        else:
            label = 3
        
        inputs = self.tokenizer(
            text = example.text_a,
            text_pair = example.text_b,
            padding = 'max_length',
            truncation = True,
            max_length = self.max_length
        )
        
        input_ids = torch.tensor(inputs.get('input_ids'),dtype=torch.long)
        attention_mask = torch.tensor(inputs.get('attention_mask'),dtype=torch.long)
        token_type_ids = torch.tensor(inputs.get('token_type_ids'),dtype=torch.long)
        label = torch.tensor(label,dtype=torch.long)
        
        return {
            'labels':label,
            'text_a':example.text_a,
            'text_b':example.text_b,
            'idx':example.id,
            'input_ids':input_ids,
            'token_type_ids':token_type_ids,
            'attention_mask':attention_mask
        }