# -*- coding: utf-8 -*-
from pyloader import DataReader
from pyloader import Dataset
from pyloader import DataCollator
from pyloader import DataLoader
import json
import torch

# Step 1: Inherit the class `DataReader` and implement the method `read_file()`
class JsonLineDataReader(DataReader):
    """Data reader for reading files with the format of json line."""
    def __init__(self, batch_size=4, shuffle=False, random_seed=42):
        super(JsonLineDataReader, self).__init__(
            batch_size,
            shuffle=shuffle,
            random_seed=random_seed
        )
    
    def read_file(self, path):
        """
        Read files from disk.
        Args:
            path (string): the file path.
        Returns:
            dict_samples (Dict[str, List]): input samples, type: a dict of lists.
            n_sample (int): the number of samples.
        """
        print("Reading data from [{}]".format(path))
        dict_samples = {}
        n_sample = 0
        with open(path, 'r', encoding='utf-8') as f:
            for line in f:
                json_sample = json.loads(line)
                n_sample += 1
                for k, v in json_sample.items():
                    if k == "post" or k == "response":
                        if not k in dict_samples:
                            dict_samples[k] = [v]
                        else:
                            dict_samples[k].append(v)
        print("Total data: {}".format(n_sample))
        return dict_samples, n_sample


# Step 2: Inherit the class `Dataset` and implement the method `batch_encode()`
class TextDataset(Dataset):
    """Dataset wrapping for text."""
    def __init__(self, vocab_file=None, max_seq_len=100, max_vocab_size=30000):
        self.max_seq_len = max_seq_len
        self.max_vocab_size = max_vocab_size
        self.stoi = {}
        self.itos = {}
        if vocab_file is not None:
            self.load_vocab(vocab_file)
    
    def load_vocab(self, vocab_file):
        with open(vocab_file, 'r', encoding='utf-8') as fr:
            for idx, line in enumerate(fr):
                if idx < self.max_vocab_size:
                    w = line.strip().split('\t')[0]
                    self.stoi[w] = idx
                    self.itos[idx] = w
        assert len(self.stoi) == len(self.itos)
        assert len(self.stoi) <= self.max_vocab_size

    def batch_encode(self, batch_data):
        """Encode batch data from raw format to numeric format."""
        post_ids = []
        post_mask = []
        resp_ids = []
        resp_maks = []
        for k, data_list in batch_data.items():
            if k == "post" or k == "response":
                for text in data_list:
                    if self.stoi:
                        text_ids = [self.stoi.get(w, 0) for w in text.split()]
                    else:
                        text_ids = [int(w) for w in text.split()]
                    text_masks = [1] * len(text_ids)
                    if len(text_ids) > self.max_seq_len:
                        input_id = text_ids[:self.max_seq_len]
                        input_mask = text_masks[:self.max_seq_len]
                    else:
                        input_id = text_ids + [0] * (self.max_seq_len-len(text_ids))
                        input_mask = text_masks + [0] * (self.max_seq_len-len(text_ids))
                    if k == "post":
                        post_ids.append(input_id)
                        post_mask.append(input_mask)
                    else:
                        resp_ids.append(input_id)
                        resp_maks.append(input_mask)
        batch = {
            "post_ids": post_ids, 
            "post_mask": post_mask,
            "resp_ids": resp_ids,
            "resp_maks": resp_maks
        }
        return batch


# Step 3: Inherit the class `DataCollator` and implement the method `collate_batch()`
class PyDataCollator(DataCollator):
    """Simple PyTorch-based data collator."""
    
    def collate_batch(self, features):
        # Special handling for labels.
        # Ensure that tensor is created with the correct type
        if "label" in features and features["label"] is not None:
            first = features["label"][0]
            if type(first) is int:
                labels = torch.tensor([f for f in features["label"]], dtype=torch.long)
            else:
                labels = torch.tensor([f for f in features["label"]], dtype=torch.float)
            batch = {"label": labels}
        else:
            batch = {}

        # Handling of all other possible attributes
        for k, v in features.items():
            if k not in ("label") and v is not None:
                batch[k] = torch.tensor([f for f in features[k]], dtype=torch.long)
        return batch


def main_test():
    data_dir = "data/chunk_files/"
    vocab_file = "data/vocab.txt"
    datareader = JsonLineDataReader(batch_size=16, shuffle=True)
    dataset = TextDataset(vocab_file, max_seq_len=60, max_vocab_size=30000)
    datacollator = PyDataCollator()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # define a DataLoader
    train_loader = DataLoader(
        datareader, dataset, datacollator, data_dir,
        num_epoch=5, max_queue_size=10, shuffle=True
    )

    # iteratively load data for training
    for batch_step, inputs in enumerate(train_loader):
        for k, v in inputs.items():
            inputs[k] = v.to(device)
            # show batch data
            print("step: {}\t{}:\n{}".format(batch_step, k, inputs[k]))
            
            # define your code
            # ...


if __name__ == "__main__":
    main_test()
