# -----------------------------------------------------------
# Stacked Cross Attention Network implementation based on 
# https://arxiv.org/abs/1803.08024.
# "Stacked Cross Attention for Image-Text Matching"
# Kuang-Huei Lee, Xi Chen, Gang Hua, Houdong Hu, Xiaodong He
#
# Writen by Kuang-Huei Lee, 2018
# ---------------------------------------------------------------
"""Data provider"""

import torch
import torch.utils.data as data
# import torchvision.transforms as transforms
import os
import nltk
from PIL import Image
import numpy as np
import h5py
# import json as jsonmod
import json
import pdb


class PrecompDataset(data.Dataset):
    """
    Load precomputed captions and image features
    Possible options: f30k_precomp, coco_precomp
    """

    def __init__(self, data_path, data_split, vocab):
        self.vocab = vocab
        self.data_name = data_path.split('/')[-1]
        loc = data_path + '/'

        # Captions
        self.captions = []
        with open(loc+'%s_caps.txt' % data_split, 'rb') as f:
            for line in f:
                self.captions.append(line.strip())

        # Tags
        if self.data_name == 'f30k_precomp':
            self.tags = []
            with open(loc+'%s_tags.txt' % data_split, 'rb') as f:
                for line in f:
                    self.tags.append(line.strip().decode('utf-8'))

        # Image features
        # self.images = np.load(loc+'%s_ims.npy' % data_split)
        self.images = h5py.File(loc+'%s_ims.h5' % data_split,'r')['feats']
        # self.images = h5py.File(loc+'%s_ims_fc.h5' % data_split,'r')['feats']
        self.boxes  = h5py.File(loc+'%s_ims_box.h5' % data_split,'r')['boxes']
        self.files  = json.load(open(loc+'%s_files.json' % data_split,'r'))
        assert self.images.shape[0]==self.boxes.shape[0]==len(self.files)

        self.length = len(self.captions)
        # rkiros data has redundancy in images, we divide by 5, 10crop doesn't
        if self.images.shape[0] != self.length:
            self.im_div = 5
        else:
            self.im_div = 1
        # the development set for coco is large and so validation would be slow
        if data_split == 'dev':
            self.length = 5000

    def __getitem__(self, index):
        # handle the image redundancy
        # img_id = index/self.im_div
        img_id = index//self.im_div
        # pdb.set_trace()
        image = torch.Tensor(self.images[img_id])
        box   = self.boxes[img_id]
        file = self.files[img_id]
        caption = self.captions[index]
        #=========================================
        if self.data_name == 'f30k_precomp':
            tag = self.tags[index]
        #=========================================
        vocab = self.vocab

        # Convert caption (string) to word ids.
        # pdb.set_trace()
        # tokens = nltk.tokenize.word_tokenize(
        #     str(caption).lower().decode('utf-8'))
        # tokens = nltk.tokenize.word_tokenize(caption.lower().decode('utf-8'))
        tokens = caption.lower().decode('utf-8').split()
        caption = []
        caption.append(vocab('<start>'))
        caption.extend([vocab(token) for token in tokens])
        caption.append(vocab('<end>'))
        target = torch.Tensor(caption)
        #=========================================
        if self.data_name == 'f30k_precomp':
            tag_mask = []
            tag = tag.split()
            tag.append('PAD')
            tag.insert(0,'PAD')
            for t in tag:
                # tag_mask.append(1)
                if t in ['NN', 'NNS','NNP','NNPS']:
                # if t in ['NN', 'NNS','NNP','NNPS','JJ','JJR','JJS']:
                # if t in ['NN', 'NNS','NNP','NNPS','JJ','JJR','JJS','VB','VBD','VBG','VBN','VBP']:
                    tag_mask.append(1)
                else:
                    tag_mask.append(0)
            tag_mask = torch.ByteTensor(tag_mask)
            # if not tag_mask.any():
            #     pdb.set_trace()
            assert  tag_mask.size(0) == target.size(0)
        else:
            tag_mask = torch.ones_like(target).byte()

        #=========================================
        return image, target, index, img_id, tag_mask

    def __len__(self):
        return self.length


def collate_fn(data):
    """Build mini-batch tensors from a list of (image, caption) tuples.
    Args:
        data: list of (image, caption) tuple.
            - image: torch tensor of shape (3, 256, 256).
            - caption: torch tensor of shape (?); variable length.

    Returns:
        images: torch tensor of shape (batch_size, 3, 256, 256).
        targets: torch tensor of shape (batch_size, padded_length).
        lengths: list; valid length for each padded caption.
    """
    # Sort a data list by caption length
    data.sort(key=lambda x: len(x[1]), reverse=True)
    images, captions, ids, img_ids, tag_masks = zip(*data)

    # Merge images (convert tuple of 3D tensor to 4D tensor)
    images = torch.stack(images, 0)

    # Merget captions (convert tuple of 1D tensor to 2D tensor)
    lengths = [len(cap) for cap in captions]
    targets = torch.zeros(len(captions), max(lengths)).long()
    for i, cap in enumerate(captions):
        end = lengths[i]
        targets[i, :end] = cap[:end]

    return images, targets, lengths, ids, img_ids, tag_masks


def get_precomp_loader(data_path, data_split, vocab, opt, batch_size=100,
                       shuffle=True, num_workers=2):
    """Returns torch.utils.data.DataLoader for custom coco dataset."""
    dset = PrecompDataset(data_path, data_split, vocab)

    data_loader = torch.utils.data.DataLoader(dataset=dset,
                                              batch_size=batch_size,
                                              shuffle=shuffle,
                                              pin_memory=True,
                                              collate_fn=collate_fn)
    return data_loader


def get_loaders(data_name, vocab, batch_size, workers, opt):
    dpath = os.path.join(opt.data_path, data_name)
    train_loader = get_precomp_loader(dpath, 'train', vocab, opt,
                                      batch_size, True, workers)
    val_loader = get_precomp_loader(dpath, 'dev', vocab, opt,
                                    batch_size, False, workers)
    return train_loader, val_loader


def get_test_loader(split_name, data_name, vocab, batch_size,
                    workers, opt):
    dpath = os.path.join(opt.data_path, data_name)
    test_loader = get_precomp_loader(dpath, split_name, vocab, opt,
                                     batch_size, False, workers)
    return test_loader
