import io
import json
import os
import pickle
from collections import defaultdict, Counter, OrderedDict

import numpy as np
import torch
import torchvision
from nltk.tokenize import sent_tokenize, word_tokenize
from torch.utils.data.dataset import Dataset
from torchvision import datasets
from torchvision.datasets.utils import download_and_extract_archive
from tqdm import tqdm


class Cub(Dataset):
    def __init__(
        self,
        root,
        download=False,
        train=True,
        image_tx=None,
    ):
        self.resources = [
            "http://www.robots.ox.ac.uk/~yshi/mmdgm/datasets/cub.zip",
        ]
        self.root = root
        if download:
            self.download()
        if not self._check_exists():
            raise RuntimeError(
                "Dataset not found." + " You can use download=True to download it"
            )
        if image_tx is None:
            image_tx = torchvision.transforms.Compose(
                [torchvision.transforms.ToTensor()]
            )
        if train:
            self.images = datasets.ImageFolder(
                os.path.join(self.raw_folder, "cub/train"), transform=image_tx
            )
        else:
            self.images = datasets.ImageFolder(
                os.path.join(self.raw_folder, "cub/test"), transform=image_tx
            )
        tx = lambda data: torch.Tensor(data)
        self.sentences = CUBSentences(
            self.raw_folder, train=train, transform=tx, max_sequence_length=32
        )

    @property
    def raw_folder(self) -> str:
        return os.path.join(self.root, self.__class__.__name__, "raw")

    def __getitem__(self, index):
        batch = {"index": index}
        image, label = self.images[index // 10]
        batch["image"] = image
        batch["label"] = label
        batch["sentence"] = self.sentences[index][0]
        return batch

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

    def _check_raw_exists(self) -> bool:
        return os.path.exists(
            os.path.join(
                self.raw_folder,
                "cub.zip",
            )
        )

    def _check_exists(self) -> bool:
        return os.path.exists(
            os.path.join(
                self.raw_folder,
                "cub",
            )
        )

    def download(self) -> None:
        """Download the data if it doesn't exist in processed_folder already."""

        if not self._check_raw_exists():
            os.makedirs(self.raw_folder, exist_ok=True)
            import ssl

            ssl._create_default_https_context = ssl._create_unverified_context
            # download files
            for url in self.resources:
                filename = url.rpartition("/")[2]
                download_and_extract_archive(
                    url, download_root=self.raw_folder, filename=filename
                )


class OrderedCounter(Counter, OrderedDict):
    """Counter that remembers the order elements are first encountered."""

    def __repr__(self):
        return "%s(%r)" % (self.__class__.__name__, OrderedDict(self))

    def __reduce__(self):
        return self.__class__, (OrderedDict(self),)


class CUBSentences(Dataset):
    def __init__(self, root_data_dir, train=True, transform=None, **kwargs):
        super().__init__()
        self.data_dir = os.path.join(root_data_dir, "cub")
        self.train = train
        self.max_sequence_length = kwargs.get("max_sequence_length", 32)
        self.min_occ = kwargs.get("min_occ", 3)
        self.transform = transform
        os.makedirs(os.path.join(root_data_dir, "lang_emb"), exist_ok=True)

        self.gen_dir = os.path.join(
            self.data_dir, "oc{}_msl{}".format(self.min_occ, self.max_sequence_length)
        )
        if train:
            self.raw_data_path = os.path.join(self.data_dir, "text_trainvalclasses.txt")
        else:
            self.raw_data_path = os.path.join(self.data_dir, "text_testclasses.txt")

        os.makedirs(self.gen_dir, exist_ok=True)
        self.data_file = "cub.{}.s{}".format(
            "train" if self.train else "test", self.max_sequence_length
        )
        self.vocab_file = "cub.vocab"

        if not os.path.exists(os.path.join(self.gen_dir, self.data_file)):
            print(
                "Data file not found for {} split at {}. Creating new... (this may take a while)".format(
                    "train" if self.train else "test",
                    os.path.join(self.gen_dir, self.data_file),
                )
            )
            self._create_data()

        else:
            self._load_data()

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

    def __getitem__(self, idx):
        sent = self.data[str(idx)]["idx"]
        if self.transform is not None:
            sent = self.transform(sent)
        return sent, self.data[str(idx)]["length"]

    @property
    def vocab_size(self):
        return len(self.w2i)

    @property
    def pad_idx(self):
        return self.w2i["<pad>"]

    @property
    def eos_idx(self):
        return self.w2i["<eos>"]

    @property
    def unk_idx(self):
        return self.w2i["<unk>"]

    def get_w2i(self):
        return self.w2i

    def get_i2w(self):
        return self.i2w

    def _load_data(self, vocab=True):
        with open(os.path.join(self.gen_dir, self.data_file), "rb") as file:
            self.data = json.load(file)

        if vocab:
            self._load_vocab()

    def _load_vocab(self):
        if not os.path.exists(os.path.join(self.gen_dir, self.vocab_file)):
            self._create_vocab()
        with open(os.path.join(self.gen_dir, self.vocab_file), "r") as vocab_file:
            vocab = json.load(vocab_file)
        self.w2i, self.i2w = vocab["w2i"], vocab["i2w"]

    def _create_data(self):
        if self.train and not os.path.exists(
            os.path.join(self.gen_dir, self.vocab_file)
        ):
            self._create_vocab()
        else:
            self._load_vocab()

        with open(self.raw_data_path, "r") as file:
            text = file.read()
            sentences = sent_tokenize(text)

        data = defaultdict(dict)
        pad_count = 0

        for i, line in enumerate(tqdm(sentences)):
            words = word_tokenize(line)

            tok = words[: self.max_sequence_length - 1]
            tok = tok + ["<eos>"]
            length = len(tok)
            if self.max_sequence_length > length:
                tok.extend(["<pad>"] * (self.max_sequence_length - length))
                pad_count += 1
            idx = [self.w2i.get(w, self.w2i["<exc>"]) for w in tok]

            id = len(data)
            data[id]["tok"] = tok
            data[id]["idx"] = idx
            data[id]["length"] = length

        print(
            "{} out of {} sentences are truncated with max sentence length {}.".format(
                len(sentences) - pad_count, len(sentences), self.max_sequence_length
            )
        )
        with io.open(os.path.join(self.gen_dir, self.data_file), "wb") as data_file:
            data = json.dumps(data, ensure_ascii=False)
            data_file.write(data.encode("utf8", "replace"))

        self._load_data(vocab=False)

    def _create_vocab(self):

        assert self.train, "Vocabulary can only be created for training file."

        with open(self.raw_data_path, "r") as file:
            text = file.read()
            sentences = sent_tokenize(text)

        occ_register = OrderedCounter()
        w2i = dict()
        i2w = dict()

        special_tokens = ["<exc>", "<pad>", "<eos>"]
        for st in special_tokens:
            i2w[len(w2i)] = st
            w2i[st] = len(w2i)

        texts = []
        unq_words = []

        for i, line in enumerate(sentences):
            words = word_tokenize(line)
            occ_register.update(words)
            texts.append(words)

        for w, occ in occ_register.items():
            if occ > self.min_occ and w not in special_tokens:
                i2w[len(w2i)] = w
                w2i[w] = len(w2i)
            else:
                unq_words.append(w)

        assert len(w2i) == len(i2w)

        print(
            "Vocabulary of {} keys created, {} words are excluded (occurrence <= {}).".format(
                len(w2i), len(unq_words), self.min_occ
            )
        )

        vocab = dict(w2i=w2i, i2w=i2w)
        with io.open(os.path.join(self.gen_dir, self.vocab_file), "wb") as vocab_file:
            data = json.dumps(vocab, ensure_ascii=False)
            vocab_file.write(data.encode("utf8", "replace"))

        with open(os.path.join(self.gen_dir, "cub.unique"), "wb") as unq_file:
            pickle.dump(np.array(unq_words), unq_file)

        with open(os.path.join(self.gen_dir, "cub.all"), "wb") as a_file:
            pickle.dump(occ_register, a_file)

        self._load_vocab()
