import json
import os
import random

import nltk
import numpy as np
import pandas as pd
import torch
import torchvision as tv
from PIL import Image
from torch.utils.data import DataLoader, Dataset, sampler
from tqdm import tqdm

from .utils import Config, get_training_data
from .vocabulary import Vocabulary


class FlickrDataset(Dataset):
    def __init__(
        self,
        transform: tv.transforms,
        mode: str,
        batch_size: int,
        vocab_threshold: int,
        vocab_file: str,
        start_word: str,
        end_word: int,
        unk_word: str,
        caption_file: str,
        image_id_file: str,
        vocab_from_file: bool,
        image_folder: str,
    ):
        self.transform = transform
        self.mode = mode
        self.caption_file = caption_file
        self.image_id_file = image_id_file
        self.batch_size = batch_size
        self.vocab = Vocabulary(
            vocab_threshold,
            vocab_file,
            caption_file,
            image_id_file,
            vocab_from_file,
            start_word,
            end_word,
            unk_word,
        )
        self.image_folder = image_folder
        self.df_data = get_training_data(self.image_id_file, self.caption_file)

        if self.mode in ["train", "validation"]:
            all_captions = self.df_data.CAPTION.values.tolist()
            all_tokens = [
                nltk.tokenize.word_tokenize(str(caption).lower())
                for caption in tqdm(all_captions)
            ]
            self.caption_lengths = [len(token) for token in all_tokens]

    def __getitem__(self, index):

        if self.mode in ["train", "validation"]:
            item = self.df_data.iloc[index]
            image_id = item["IMAGE_ID"]
            caption = item["CAPTION"]

            image_file = os.path.join(self.image_folder, image_id)
            image = Image.open(image_file).convert("RGB")
            image = self.transform(image)

            tokens = nltk.tokenize.word_tokenize(str(caption).lower())
            caption = []
            caption.append(self.vocab(self.vocab.start_word))
            caption.extend([self.vocab(token) for token in tokens])
            caption.append(self.vocab(self.vocab.end_word))

            caption = torch.Tensor(caption).long()

            # return pre-processed image and caption
            return image, caption
        else:
            item = self.df_data.iloc[index]
            image_id = item["IMAGE_ID"]
            image_file = os.path.join(self.image_folder, image_id)
            PIL_image = Image.open(image_file).convert("RGB")

            original_image = np.array(PIL_image)
            transformed_image = self.transform(PIL_image)

            return original_image, transformed_image

    def get_train_indices(self):
        sel_length = np.random.choice(self.caption_lengths)
        all_indices = np.where(
            [
                self.caption_lengths[i] == sel_length
                for i in np.arange(len(self.caption_lengths))
            ]
        )[0]
        indices = list(np.random.choice(all_indices, size=self.batch_size))
        return indices

    def __len__(self):
        return len(self.df_data)


def get_data_loader(
    transform: tv.transforms,
    caption_file: str,
    image_id_file: str,
    image_folder: str,
    config: Config,
    vocab_file: str,
    mode: str = "train",
    batch_size: int = 1,
    vocab_threshold=None,
    start_word: str = "<start>",
    end_word: str = "<end>",
    unk_word: str = "<unk>",
    vocab_from_file: bool = True,
    num_workers: int = 0,
):
    """Returns the data loader

    :param transform: [description]
    :type transform: tv.transforms
    :param mode: [description], defaults to "train"
    :type mode: str, optional
    :param batch_size: [description], defaults to 1
    :type batch_size: int, optional
    :param vocab_threshold: [description], defaults to None
    :type vocab_threshold: [type], optional
    :param vocab_file: [description], defaults to "output/vocab.pkl"
    :type vocab_file: str, optional
    :param start_word: [description], defaults to "<start>"
    :type start_word: str, optional
    :param end_word: [description], defaults to "<end>"
    :type end_word: str, optional
    :param unk_word: [description], defaults to "<unk>"
    :type unk_word: str, optional
    :param vocab_from_file: [description], defaults to True
    :type vocab_from_file: bool, optional
    :param num_workers: [description], defaults to 0
    :type num_workers: int, optional
    
    """

    assert mode in [
        "train",
        "validation",
        "test",
    ], f"mode: '{mode}' must be one of ['train','validation','test']"
    if vocab_from_file == False:
        assert (
            mode == "train"
        ), f"mode: '{mode}', but to generate vocab from caption file, mode must be 'train' "

    if mode == "train":
        if vocab_from_file == True:
            assert os.path.exists(
                vocab_file
            ), "vocab_file does not exist.  Change vocab_from_file to False to create vocab_file."
        assert image_id_file.find(
            "train"
        ), f"double check image_id_file: {image_id_file}. File name should have the substring 'train'"
        assert os.path.exists(
            image_id_file
        ), f"image id file: {image_id_file} doesn't not exist."
        assert os.path.exists(
            caption_file
        ), f"caption file: {caption_file} doesn't not exist."
        assert os.path.isdir(
            config.IMAGE_DATA_DIR
        ), f"{config.IMAGE_DATA_DIR} not a directory"
        assert (
            len(os.listdir(config.IMAGE_DATA_DIR)) != 0
        ), f"{config.IMAGE_DATA_DIR} is empty."

    if mode == "validation":
        assert image_id_file.find(
            "dev"
        ), f"double check image_id_file: {image_id_file}. File name should have the substring 'dev' "
        assert os.path.exists(
            image_id_file
        ), f"image id file: {image_id_file} doesn't not exist."
        assert os.path.exists(
            caption_file
        ), f"caption file: {caption_file} doesn't not exist."
        assert os.path.isdir(
            config.IMAGE_DATA_DIR
        ), f"{config.IMAGE_DATA_DIR} not a directory"
        assert (
            len(os.listdir(config.IMAGE_DATA_DIR)) != 0
        ), f"{config.IMAGE_DATA_DIR} is empty."
        assert os.path.exists(
            vocab_file
        ), f"Must first generate {vocab_file} from training data."
        assert vocab_from_file == True, "Change vocab_from_file to True."

    if mode == "test":
        assert (
            batch_size == 1
        ), "Please change batch_size to 1 if testing your model."
        assert image_id_file.find(
            "test"
        ), f"double check image_id_file: {image_id_file}. File name should have the substring 'test'"
        assert os.path.exists(
            vocab_file
        ), f"Must first generate {vocab_file} from training data."
        assert vocab_from_file == True, "Change vocab_from_file to True."

    img_folder = config.IMAGE_DATA_DIR
    annotations_file = caption_file

    # image caption dataset
    dataset = FlickrDataset(
        transform,
        mode,
        batch_size,
        vocab_threshold,
        vocab_file,
        start_word,
        end_word,
        unk_word,
        caption_file,
        image_id_file,
        vocab_from_file,
        image_folder,
    )

    if mode in ["train", "validation"]:
        # Randomly sample a caption length, and sample indices with that length.
        indices = dataset.get_train_indices()
        # Create and assign a batch sampler to retrieve a batch with the sampled indices.
        initial_sampler = sampler.SubsetRandomSampler(indices=indices)
        # data loader for COCO dataset.
        data_loader = DataLoader(
            dataset=dataset,
            num_workers=num_workers,
            batch_sampler=sampler.BatchSampler(
                sampler=initial_sampler,
                batch_size=dataset.batch_size,
                drop_last=False,
            ),
        )
    else:
        data_loader = DataLoader(
            dataset=dataset,
            batch_size=dataset.batch_size,
            shuffle=False,
            num_workers=num_workers,
        )

    return data_loader
