import os
import json
import omegaconf
import itertools
import copy

import numpy as np
from PIL import Image

import torch
from torch.utils.data.dataset import Dataset

from utils import get_class_names
from lightning_modules.utils import sliding_windows_by_words

class KVUDataset(Dataset):
    def __init__(
        self,
        cfg, 
        tokenizer_layoutxlm,
        feature_extractor,
        mode=None,
    ):
        super(KVUDataset, self).__init__()
        
        self.dataset_root_path = cfg.dataset_root_path
        if not isinstance(self.dataset_root_path, omegaconf.listconfig.ListConfig):
            self.dataset_root_path = [self.dataset_root_path]
            
        self.backbone_type = cfg.model.backbone
        self.max_seq_length = cfg.train.max_seq_length
        self.window_size = cfg.train.max_num_words
        self.slice_interval = cfg.train.slice_interval
        
        self.tokenizer_layoutxlm = tokenizer_layoutxlm
        self.feature_extractor = feature_extractor
        
        self.stage = cfg.stage
        self.mode = mode

        self.pad_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(self.tokenizer_layoutxlm._pad_token)
        self.cls_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(self.tokenizer_layoutxlm._cls_token)
        self.sep_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(self.tokenizer_layoutxlm._sep_token)
        self.unk_token_id_layoutxlm = tokenizer_layoutxlm.convert_tokens_to_ids(self.tokenizer_layoutxlm._unk_token)
        
        self.examples = self._load_examples()

        self.class_names = get_class_names(self.dataset_root_path)
        self.class_idx_dic = dict(
            [(class_name, idx) for idx, class_name in enumerate(self.class_names)]
        )
        
    def _load_examples(self):
        examples = []
        for dataset_dir in self.dataset_root_path:
            with open(
                os.path.join(dataset_dir, f"preprocessed_files_{self.mode}.txt"),
                "r",
                encoding="utf-8",
            ) as fp:
                for line in fp.readlines():
                    preprocessed_file = os.path.join(dataset_dir, line.strip())
                    examples.append(
                        json.load(open(preprocessed_file, "r", encoding="utf-8"))
                    )

        return examples

    def __len__(self):
        return len(self.examples)
    
    
    def __getitem__(self, index):
        json_obj = self.examples[index]

        width = json_obj["meta"]["imageSize"]["width"]
        height = json_obj["meta"]["imageSize"]["height"]
        img_path = json_obj["meta"]["image_path"]

        images = [Image.open(json_obj["meta"]["image_path"]).convert("RGB")]
        image_features = torch.from_numpy(self.feature_extractor(images)['pixel_values'][0].copy())
        
        word_windows, parse_class_windows, parse_relation_windows = sliding_windows_by_words(
                                                                    json_obj["words"], 
                                                                    json_obj['parse']['class'], 
                                                                    json_obj['parse']['relations'], 
                                                                    self.window_size, self.slice_interval)
        outputs = {}
        if self.stage == 1:
            if self.mode == 'train':
                i = np.random.randint(0, len(word_windows), 1)[0]
                outputs['windows'] = self.preprocess(word_windows[i], parse_class_windows[i], parse_relation_windows[i], 
                                        {'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
                                        max_seq_length=self.max_seq_length)
            else:
                outputs['windows'] = []
                for i in range(len(word_windows)):
                    single_window = self.preprocess(word_windows[i], parse_class_windows[i], parse_relation_windows[i], 
                                        {'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
                                        max_seq_length=self.max_seq_length)
                    outputs['windows'].append(single_window)

        elif self.stage == 2:
            outputs['documents'] = self.preprocess(json_obj["words"], json_obj['parse']['class'], json_obj['parse']['relations'], 
                                    {'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
                                    max_seq_length=2048)

            windows = []
            for i in range(len(word_windows)):
                _words = word_windows[i]
                _parse_class = parse_class_windows[i]
                _parse_relation = parse_relation_windows[i]
                windows.append(
                    self.preprocess(_words, _parse_class, _parse_relation, 
                                    {'image': image_features, 'width': width, 'height': height, 'img_path': img_path}, 
                                    self.max_seq_length)
                    )
            outputs['windows'] = windows
        return outputs
            
    def preprocess(self, words, parse_class, parse_relation, feature_maps, max_seq_length):        
        input_ids_layoutxlm = np.ones(max_seq_length, dtype=int) * self.pad_token_id_layoutxlm

        attention_mask_layoutxlm = np.zeros(max_seq_length, dtype=int)

        bbox = np.zeros((max_seq_length, 8), dtype=np.float32)
        
        are_box_first_tokens = np.zeros(max_seq_length, dtype=np.bool_)

        itc_labels = np.zeros(max_seq_length, dtype=int)
        stc_labels = np.ones(max_seq_length, dtype=np.int64) * max_seq_length
        el_labels = np.ones((max_seq_length,), dtype=int) * max_seq_length
        el_labels_from_key = np.ones((max_seq_length,), dtype=int) * max_seq_length
        list_layoutxlm_tokens = []
        
        list_bbs = []
        box2token_span_map = []


        box_to_token_indices = []
        cum_token_idx = 0

        cls_bbs = [0.0] * 8
        len_overlap_tokens = 0
        len_non_overlap_tokens = 0 
        len_valid_tokens = 0
        
        for word_idx, word in enumerate(words):
            this_box_token_indices = []

            layoutxlm_tokens = word["layoutxlm_tokens"]
            bb = word["boundingBox"]
            len_valid_tokens += len(layoutxlm_tokens)
            if word_idx < self.slice_interval:
                len_non_overlap_tokens += len(layoutxlm_tokens)
                # print(word_idx, layoutxlm_tokens, non_overlap_tokens)
            
            if len(layoutxlm_tokens) == 0:
                layoutxlm_tokens.append(self.unk_token_id_layoutxlm)

            if len(list_layoutxlm_tokens) + len(layoutxlm_tokens) > max_seq_length - 2:
                break
            
            box2token_span_map.append(
                [len(list_layoutxlm_tokens) + 1, len(list_layoutxlm_tokens) + len(layoutxlm_tokens) + 1]
            )  # including st_idx
            list_layoutxlm_tokens += layoutxlm_tokens

            # min, max clipping
            for coord_idx in range(4):
                bb[coord_idx][0] = max(0.0, min(bb[coord_idx][0], feature_maps['width']))
                bb[coord_idx][1] = max(0.0, min(bb[coord_idx][1], feature_maps['height']))

            bb = list(itertools.chain(*bb))
            bbs = [bb for _ in range(len(layoutxlm_tokens))]
            
            for _ in layoutxlm_tokens:
                cum_token_idx += 1
                this_box_token_indices.append(cum_token_idx)

            list_bbs.extend(bbs)
            box_to_token_indices.append(this_box_token_indices)
            
        sep_bbs = [feature_maps['width'], feature_maps['height']] * 4

        # For [CLS] and [SEP]
        list_layoutxlm_tokens = (
            [self.cls_token_id_layoutxlm]
            + list_layoutxlm_tokens[: max_seq_length - 2]
            + [self.sep_token_id_layoutxlm]
        )
        
        if len(list_bbs) == 0:
            # When len(json_obj["words"]) == 0 (no OCR result)
            list_bbs = [cls_bbs] + [sep_bbs]
        else:  # len(list_bbs) > 0
            list_bbs = [cls_bbs] + list_bbs[: max_seq_length - 2] + [sep_bbs]

        len_list_layoutxlm_tokens = len(list_layoutxlm_tokens)
        input_ids_layoutxlm[:len_list_layoutxlm_tokens] = list_layoutxlm_tokens
        attention_mask_layoutxlm[:len_list_layoutxlm_tokens] = 1

        bbox[:len_list_layoutxlm_tokens, :] = list_bbs

        # Normalize bbox -> 0 ~ 1
        bbox[:, [0, 2, 4, 6]] = bbox[:, [0, 2, 4, 6]] / feature_maps['width']
        bbox[:, [1, 3, 5, 7]] = bbox[:, [1, 3, 5, 7]] / feature_maps['height']

        if self.backbone_type in ("layoutlm", "layoutxlm", "xlm-roberta"):
            bbox = bbox[:, [0, 1, 4, 5]]
            bbox = bbox * 1000
            bbox = bbox.astype(int)
        else:
            assert False

        st_indices = [
            indices[0]
            for indices in box_to_token_indices
            if indices[0] < max_seq_length
        ]
        are_box_first_tokens[st_indices] = True

        # Label for entity extraction
        classes_dic = parse_class
        for class_name in self.class_names:
            if class_name == "others":
                continue
            if class_name not in classes_dic:
                continue

            for word_list in classes_dic[class_name]:
                is_first, last_word_idx = True, -1
                for word_idx in word_list:
                    if word_idx >= len(box_to_token_indices):
                        break
                    box2token_list = box_to_token_indices[word_idx]
                    for converted_word_idx in box2token_list:
                        if converted_word_idx >= max_seq_length:
                            break  # out of idx

                        if is_first:
                            itc_labels[converted_word_idx] = self.class_idx_dic[
                                class_name
                            ]
                            is_first, last_word_idx = False, converted_word_idx
                        else:
                            stc_labels[converted_word_idx] = last_word_idx
                            last_word_idx = converted_word_idx
                            
            
        # Label for entity linking
        relations = parse_relation
        for relation in relations:
            if relation[0] >= len(box2token_span_map) or relation[1] >= len(
                box2token_span_map
            ):
                continue
            if (
                box2token_span_map[relation[0]][0] >= max_seq_length
                or box2token_span_map[relation[1]][0] >= max_seq_length
            ):
                continue

            word_from = box2token_span_map[relation[0]][0]
            word_to = box2token_span_map[relation[1]][0]
            # el_labels[word_to] = word_from
            
             
            #### 1st relation => ['key, 'value']
            #### 2st relation => ['header', 'key'or'value']
            if itc_labels[word_from] == 2 and itc_labels[word_to] == 3:
                el_labels_from_key[word_to] = word_from # pair of (key-value)
            if itc_labels[word_from] == 4 and (itc_labels[word_to] in (2, 3)):
                el_labels[word_to] = word_from # pair of (header, key) or (header-value)
                
        assert len_list_layoutxlm_tokens == len_valid_tokens + 2
        len_overlap_tokens = len_valid_tokens - len_non_overlap_tokens
        # overlap_tokens = max_seq_length - non_overlap_tokens  
        ntokens = max_seq_length if max_seq_length == 512 else len_valid_tokens + 2
        
        # ntokens = max_seq_length
        
        input_ids_layoutxlm = torch.from_numpy(input_ids_layoutxlm[:ntokens])
        
        attention_mask_layoutxlm = torch.from_numpy(attention_mask_layoutxlm[:ntokens])
        
        bbox = torch.from_numpy(bbox[:ntokens])
        are_box_first_tokens = torch.from_numpy(are_box_first_tokens[:ntokens])

        itc_labels = itc_labels[:ntokens]
        stc_labels = stc_labels[:ntokens]
        el_labels = el_labels[:ntokens]
        el_labels_from_key = el_labels_from_key[:ntokens]
        
        itc_labels = np.where(itc_labels != max_seq_length, itc_labels, ntokens)
        stc_labels = np.where(stc_labels != max_seq_length, stc_labels, ntokens)
        el_labels = np.where(el_labels != max_seq_length, el_labels, ntokens)
        el_labels_from_key = np.where(el_labels_from_key != max_seq_length, el_labels_from_key, ntokens)

        itc_labels = torch.from_numpy(itc_labels)
        stc_labels = torch.from_numpy(stc_labels)
        el_labels = torch.from_numpy(el_labels)
        el_labels_from_key = torch.from_numpy(el_labels_from_key)
        

        return_dict = {
            "img_path": feature_maps['img_path'],
            "len_overlap_tokens": len_overlap_tokens, 
            'len_valid_tokens': len_valid_tokens,
            "image": feature_maps['image'],
            "input_ids_layoutxlm": input_ids_layoutxlm,
            "attention_mask_layoutxlm": attention_mask_layoutxlm,
            "bbox": bbox,
            "itc_labels": itc_labels,
            "are_box_first_tokens": are_box_first_tokens,
            "stc_labels": stc_labels,
            "el_labels": el_labels,
            "el_labels_from_key": el_labels_from_key,
        }
        return return_dict     
    


class KVUPredefinedDataset(KVUDataset):
    def __init__(self, cfg, tokenizer_layoutxlm, feature_extractor, mode=None):
        super().__init__(cfg, tokenizer_layoutxlm, feature_extractor, mode)
        self.max_windows = cfg.train.max_windows
    
    def __getitem__(self, index):
        json_obj = self.examples[index]

        width = json_obj["meta"]["imageSize"]["width"]
        height = json_obj["meta"]["imageSize"]["height"]
        img_path = json_obj["meta"]["image_path"]

        images = [Image.open(json_obj["meta"]["image_path"]).convert("RGB")]
        image_features = torch.from_numpy(self.feature_extractor(images)['pixel_values'][0].copy())
        
        
        word_windows, parse_class_windows, parse_relation_windows = sliding_windows_by_words(
                                                                    json_obj["words"], 
                                                                    json_obj['parse']['class'], 
                                                                    json_obj['parse']['relations'], 
                                                                    self.window_size, self.slice_interval)
        
        word_windows = word_windows[: self.max_windows] if len(word_windows) >= self.max_windows else word_windows + [[]] * (self.max_windows - len(word_windows)) 
        parse_class_windows = parse_class_windows[: self.max_windows] if len(parse_class_windows) >= self.max_windows else parse_class_windows + [[]] * (self.max_windows - len(parse_class_windows)) 
        parse_relation_windows = parse_relation_windows[: self.max_windows] if len(parse_relation_windows) >= self.max_windows else parse_relation_windows + [[]] * (self.max_windows - len(parse_relation_windows)) 

        
        outputs = {}
        # outputs['labels'] = self.preprocess(json_obj["words"], json_obj['parse']['class'], json_obj['parse']['relations'], 
        #                             {'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
        #                             max_seq_length=self.max_seq_length*self.max_windows)
        
        outputs['windows'] = []
        for i in range(len(self.max_windows)):
            single_window = self.preprocess(word_windows[i], parse_class_windows[i], parse_relation_windows[i], 
                                {'image': image_features, 'width': width, 'height': height, 'img_path': img_path},
                                max_seq_length=self.max_seq_length)
            outputs['windows'].append(single_window)
        
        outputs['windows'] = torch.cat(outputs["windows"], dim=0) 

        return outputs


class KVUEmbeddingDataset(Dataset):
    def __init__(
        self,
        cfg, 
        mode=None,
    ):
        super(KVUEmbeddingDataset, self).__init__()
        
        self.dataset_root_path = cfg.dataset_root_path
        if not isinstance(self.dataset_root_path, omegaconf.listconfig.ListConfig):
            self.dataset_root_path = [self.dataset_root_path]
        
        self.stage = cfg.stage
        self.mode = mode
        
        self.examples = self._load_examples()

    def _load_examples(self):
        examples = []
        for dataset_dir in self.dataset_root_path:
            with open(
                os.path.join(dataset_dir, f"preprocessed_files_{self.mode}.txt"),
                "r",
                encoding="utf-8",
            ) as fp:
                for line in fp.readlines():
                    preprocessed_file = os.path.join(dataset_dir, line.strip())
                    examples.append(
                        preprocessed_file.replace('preprocessed', 'embedding_matrix').replace('.json', '.npz') 
                    )
        return examples

    def __len__(self):
        return len(self.examples)
    
    def __getitem__(self, index):
        input_embeddings = np.load(self.examples[index])
        
        return_dict = {
            "embeddings": torch.from_numpy(input_embeddings["embeddings"]).type(torch.HalfTensor),
            "attention_mask_layoutxlm": torch.from_numpy(input_embeddings["attention_mask_layoutxlm"]),
            "are_box_first_tokens": torch.from_numpy(input_embeddings["are_box_first_tokens"]),
            "bbox": torch.from_numpy(input_embeddings["bbox"]),
            "itc_labels": torch.from_numpy(input_embeddings["itc_labels"]),
            "stc_labels": torch.from_numpy(input_embeddings["stc_labels"]),
            "el_labels": torch.from_numpy(input_embeddings["el_labels"]),
            "el_labels_from_key": torch.from_numpy(input_embeddings["el_labels_from_key"]),
        }
        return return_dict     
    
    
class DocumentKVUDataset(KVUDataset):
    def __init__(self, cfg, tokenizer_layoutxlm, feature_extractor, mode=None):
        super().__init__(cfg, tokenizer_layoutxlm, feature_extractor, mode)
        self.self.max_window_count = cfg.train.max_window_count
    
    def __getitem__(self, idx):
        json_obj = self.examples[idx]

        width = json_obj["meta"]["imageSize"]["width"]
        height = json_obj["meta"]["imageSize"]["height"]

        images = [Image.open(json_obj["meta"]["image_path"]).convert("RGB")]
        feature_maps = torch.from_numpy(self.feature_extractor(images)['pixel_values'][0].copy())

        n_words = len(json_obj['words'])
        output_dicts = {'windows': [], 'documents': []}
        box_to_token_indices_document = []
        box2token_span_map_document = []
        n_empty_windows = 0

        for i in range(self.max_window_count):
            input_ids = np.ones(self.max_seq_length, dtype=int) * 1
            bbox = np.zeros((self.max_seq_length, 8), dtype=np.float32)
            attention_mask = np.zeros(self.max_seq_length, dtype=int)

            itc_labels = np.zeros(self.max_seq_length, dtype=int)
            are_box_first_tokens = np.zeros(self.max_seq_length, dtype=np.bool_)

            # stc_labels stores the index of the previous token.
            # A stored index of max_seq_length (512) indicates that
            # this token is the initial token of a word box.
            stc_labels = np.ones(self.max_seq_length, dtype=np.int64) * self.max_seq_length
            el_labels = np.ones((self.max_seq_length,), dtype=int) * self.max_seq_length
            el_labels_from_key = np.ones((self.max_seq_length,), dtype=int) * self.max_seq_length

            start_word_idx = i * self.window_size
            stop_word_idx = min(n_words, (i+1)*self.window_size)

            if start_word_idx >= stop_word_idx:
                n_empty_windows += 1
                output_dicts.append(output_dicts[-1])
                
                box_to_token_indices_to_mod = copy.deepcopy(box_to_token_indices)
                for i_box in range(len(box_to_token_indices_to_mod)):
                    for j in range(len(box_to_token_indices_to_mod[i_box])):
                        box_to_token_indices_to_mod[i_box][j] += i * self.max_seq_length    
                for element in box_to_token_indices_to_mod:
                    box_to_token_indices_document.append(element)

                box2token_span_map_to_mod = copy.deepcopy(box2token_span_map)
                for i_box in range(len(box2token_span_map_to_mod)):
                    for j in range(len(box2token_span_map_to_mod[i_box])):
                        box2token_span_map_to_mod[i_box][j] += i * self.max_seq_length
                for element in box2token_span_map_to_mod:
                    box2token_span_map_document.append(element)

                continue

            list_tokens = []
            list_bbs = []
            box2token_span_map = []

            box_to_token_indices = []
            cum_token_idx = 0

            cls_bbs = [0.0] * 8

            # Parse words
            for word_idx, word in enumerate(json_obj["words"][start_word_idx:stop_word_idx]):
                this_box_token_indices = []
                tokens = word["tokens"]
                bb = word["boundingBox"]
                if len(tokens) == 0:
                    tokens.append(self.unk_token_id)

                if len(list_tokens) + len(tokens) > self.max_seq_length - 2:
                    break ### be able to apply sliding window here
                
                box2token_span_map.append(
                    [len(list_tokens) + 1, len(list_tokens) + len(tokens) + 1]
                )  # including st_idx
                list_tokens += tokens

                # min, max clipping
                for coord_idx in range(4):
                    bb[coord_idx][0] = max(0.0, min(bb[coord_idx][0], width))
                    bb[coord_idx][1] = max(0.0, min(bb[coord_idx][1], height))

                bb = list(itertools.chain(*bb))
                bbs = [bb for _ in range(len(tokens))]

                for _ in tokens:
                    cum_token_idx += 1
                    this_box_token_indices.append(cum_token_idx)

                list_bbs.extend(bbs)
                box_to_token_indices.append(this_box_token_indices)
            
            sep_bbs = [width, height] * 4

            # For [CLS] and [SEP]
            list_tokens = (
                [self.cls_token_id]
                + list_tokens[: self.max_seq_length - 2]
                + [self.sep_token_id]
            )
            if len(list_bbs) == 0:
                # When len(json_obj["words"]) == 0 (no OCR result)
                list_bbs = [cls_bbs] + [sep_bbs]
            else:  # len(list_bbs) > 0
                list_bbs = [cls_bbs] + list_bbs[: self.max_seq_length - 2] + [sep_bbs]

            len_list_tokens = len(list_tokens)
            input_ids[:len_list_tokens] = list_tokens
            attention_mask[:len_list_tokens] = 1

            bbox[:len_list_tokens, :] = list_bbs

            # Normalize bbox -> 0 ~ 1
            bbox[:, [0, 2, 4, 6]] = bbox[:, [0, 2, 4, 6]] / width
            bbox[:, [1, 3, 5, 7]] = bbox[:, [1, 3, 5, 7]] / height

            if self.backbone_type in ('layoutlm', 'layoutxlm', 'xlm-roberta'):
                bbox = bbox[:, [0, 1, 4, 5]]
                bbox = bbox * 1000
                bbox = bbox.astype(int)
            else:
                assert False

            st_indices = [
                indices[0]
                for indices in box_to_token_indices
                if indices[0] < self.max_seq_length
            ]
            are_box_first_tokens[st_indices] = True

            # Parse word groups
            classes_dic = json_obj["parse"]["class"]
            for class_name in self.class_names:
                if class_name == "others":
                    continue
                if class_name not in classes_dic:
                    continue

                for word_list in classes_dic[class_name]:
                    word_list = [w for w in word_list if w >= start_word_idx and w < stop_word_idx]
                    if len(word_list) == 0:
                        continue # no more word left
                    word_list = [w - start_word_idx for w in word_list]
                    
                    is_first, last_word_idx = True, -1
                    for word_idx in word_list:
                        if word_idx >= len(box_to_token_indices):
                            break
                        box2token_list = box_to_token_indices[word_idx]
                        for converted_word_idx in box2token_list:
                            if converted_word_idx >= self.max_seq_length:
                                break  # out of idx

                            if is_first:
                                itc_labels[converted_word_idx] = self.class_idx_dic[
                                    class_name
                                ]
                                is_first, last_word_idx = False, converted_word_idx
                            else:
                                stc_labels[converted_word_idx] = last_word_idx
                                last_word_idx = converted_word_idx
            
            # Parse relation
            relations = json_obj["parse"]["relations"]
            for relation in relations:
                relation = [r for r in relation if r >= start_word_idx and r < stop_word_idx]
                if len(relation) != 2:
                    continue # relation popped due to window inconsistent
                relation[0] -= start_word_idx
                relation[1] -= start_word_idx
                
                word_from = box2token_span_map[relation[0]][0]
                word_to = box2token_span_map[relation[1]][0]
                
                #### 1st relation => ['key, 'value']
                #### 2st relation => ['header', 'key'or'value']
                if itc_labels[word_from] == 2 and itc_labels[word_to] == 3:
                    el_labels_from_key[word_to] = word_from # pair of (key-value)
                if itc_labels[word_from] == 4 and (itc_labels[word_to] in (2, 3)):
                    el_labels[word_to] = word_from # pair of (header, key) or (header-value)
            
            input_ids = torch.from_numpy(input_ids)
            bbox = torch.from_numpy(bbox)
            attention_mask = torch.from_numpy(attention_mask)

            itc_labels = torch.from_numpy(itc_labels)
            are_box_first_tokens = torch.from_numpy(are_box_first_tokens)
            stc_labels = torch.from_numpy(stc_labels)
            el_labels = torch.from_numpy(el_labels)
            el_labels_from_key = torch.from_numpy(el_labels_from_key)

            box_to_token_indices_to_mod = copy.deepcopy(box_to_token_indices)
            for i_box in range(len(box_to_token_indices_to_mod)):
                for j in range(len(box_to_token_indices_to_mod[i_box])):
                    box_to_token_indices_to_mod[i_box][j] += i * self.max_seq_length    
            for element in box_to_token_indices_to_mod:
                box_to_token_indices_document.append(element)

            box2token_span_map_to_mod = copy.deepcopy(box2token_span_map)
            for i_box in range(len(box2token_span_map_to_mod)):
                for j in range(len(box2token_span_map_to_mod[i_box])):
                    box2token_span_map_to_mod[i_box][j] += i * self.max_seq_length
            for element in box2token_span_map_to_mod:
                box2token_span_map_document.append(element)

            return_dict = {
                "image": feature_maps,    
                "input_ids": input_ids,
                "bbox": bbox,
                "attention_mask": attention_mask,
                "itc_labels": itc_labels,
                "are_box_first_tokens": are_box_first_tokens,
                "stc_labels": stc_labels,
                "el_labels": el_labels,
                "el_labels_from_key": el_labels_from_key,
            }

            output_dicts["windows"].append(return_dict)
        
        
        # Parse whole document labels
        attention_mask = torch.cat([o['attention_mask'] for o in output_dicts])
        are_box_first_tokens = torch.cat([o['are_box_first_tokens'] for o in output_dicts])
        if n_empty_windows > 0:
            attention_mask[self.max_seq_length * (self.max_window_count - n_empty_windows):] = torch.from_numpy(np.zeros(self.max_seq_length * n_empty_windows, dtype=int))
            are_box_first_tokens[self.max_seq_length * (self.max_window_count - n_empty_windows):] = torch.from_numpy(np.zeros(self.max_seq_length * n_empty_windows, dtype=np.bool_))
        bbox = torch.cat([o['bbox'] for o in output_dicts])

        self.max_seq_length_document = self.max_seq_length * self.max_window_count
        itc_labels = np.zeros(self.max_seq_length_document, dtype=int)
        stc_labels = np.ones(self.max_seq_length_document, dtype=np.int64) * self.max_seq_length_document
        el_labels = np.ones((self.max_seq_length_document,), dtype=int) * self.max_seq_length_document
        el_labels_from_key = np.ones((self.max_seq_length_document,), dtype=int) * self.max_seq_length_document

        # Parse word groups
        classes_dic = json_obj["parse"]["class"]
        for class_name in self.class_names:
            if class_name == "others":
                continue
            if class_name not in classes_dic:
                continue
            
            word_lists = classes_dic[class_name]

            for word_list in word_lists:
                is_first, last_word_idx = True, -1
                for word_idx in word_list:
                    if word_idx >= len(box_to_token_indices_document):
                        break
                    box2token_list = box_to_token_indices_document[word_idx]
                    for converted_word_idx in box2token_list:
                        if converted_word_idx >= self.max_seq_length_document:
                            break  # out of idx

                        if is_first:
                            itc_labels[converted_word_idx] = self.class_idx_dic[
                                class_name
                            ]
                            is_first, last_word_idx = False, converted_word_idx
                        else:
                            stc_labels[converted_word_idx] = last_word_idx
                            last_word_idx = converted_word_idx

        # Parse relation
        relations = json_obj["parse"]["relations"]
        
        for relation in relations:
            if relation[0] >= len(box2token_span_map_document) or relation[1] >= len(
                box2token_span_map_document
            ):
                continue
            if (
                box2token_span_map_document[relation[0]][0] >= self.max_seq_length_document
                or box2token_span_map_document[relation[1]][0] >= self.max_seq_length_document
            ):
                continue

            word_from = box2token_span_map_document[relation[0]][0]
            word_to = box2token_span_map_document[relation[1]][0]
            
            if itc_labels[word_from] == 2 and itc_labels[word_to] == 3:
                el_labels_from_key[word_to] = word_from # pair of (key-value)
            if itc_labels[word_from] == 4 and (itc_labels[word_to] in (2, 3)):
                el_labels[word_to] = word_from # pair of (header, key) or (header-value)

        itc_labels = torch.from_numpy(itc_labels)
        stc_labels = torch.from_numpy(stc_labels)
        el_labels = torch.from_numpy(el_labels)
        el_labels_from_key = torch.from_numpy(el_labels_from_key)

        return_dict = {
            "img_path": json_obj["meta"]["image_path"],
            "attention_mask": attention_mask,
            "bbox": bbox,
            "itc_labels": itc_labels,
            "are_box_first_tokens": are_box_first_tokens,
            "stc_labels": stc_labels,
            "el_labels": el_labels,
            "el_labels_from_key": el_labels_from_key,
            "n_empty_windows": n_empty_windows
        }

        output_dicts['documents'] = return_dict
        return output_dicts 