import os
import re
import json
import torch

from functional import pseq

from torchtext.data.dataset import Dataset
from torchtext.data.example import Example
from torchtext.data import Field, RawField, BucketIterator
from torchtext.vocab import Vocab, pretrained_aliases, Vectors
from torchtext.utils import download_from_url

from qanta.wikipedia.cached_wikipedia import extract_wiki_sentences

DS_VERSION = "2018.04.18"


ftp_patterns = {
    "\n",
    ", for 10 points,",
    ", for ten points,",
    "--for 10 points--",
    "for 10 points, ",
    "for 10 points--",
    "for ten points, ",
    "for 10 points ",
    "for ten points ",
    ", ftp," "ftp,",
    "ftp",
}

regex_pattern = "|".join([re.escape(p) for p in ftp_patterns])
regex_pattern += r"|\[.*?\]|\(.*?\)"


def str_split(text):
    return text.split()


def create_qb_tokenizer(
    unigrams=True,
    bigrams=False,
    trigrams=False,
    zero_length_token="zerolengthunk",
    strip_qb_patterns=True,
):
    def tokenizer(text):
        if strip_qb_patterns:
            text = (
                re.sub(
                    "\s+", " ", re.sub(regex_pattern, " ", text, flags=re.IGNORECASE)
                )
                .strip()
                .capitalize()
            )
        import nltk

        tokens = nltk.word_tokenize(text)
        if len(tokens) == 0:
            return [zero_length_token]
        else:
            ngrams = []
            if unigrams:
                ngrams.extend(tokens)
            if bigrams:
                ngrams.extend([f"{w0}++{w1}" for w0, w1 in nltk.bigrams(tokens)])
            if trigrams:
                ngrams.extend(
                    [f"{w0}++{w1}++{w2}" for w0, w1, w2 in nltk.trigrams(tokens)]
                )

            if len(ngrams) == 0:
                ngrams.append(zero_length_token)
            return ngrams

    return tokenizer


class LongField(RawField):
    def __init__(self):
        super().__init__()

    def preprocess(self, x):
        return int(x)

    def process(self, batch, **kwargs):
        return torch.LongTensor(batch)


class QBVocab(Vocab):
    def load_vectors(self, vectors):
        if not isinstance(vectors, list):
            vectors = [vectors]
        for idx, vector in enumerate(vectors):
            if isinstance(vector, str):
                # Convert the string pretrained vector identifier
                # to a Vectors object
                if vector not in pretrained_aliases:
                    raise ValueError(
                        "Got string input vector {}, but allowed pretrained "
                        "vectors are {}".format(vector, list(pretrained_aliases.keys()))
                    )
                vectors[idx] = pretrained_aliases[vector]()
            elif not isinstance(vector, Vectors):
                raise ValueError(
                    "Got input vectors of type {}, expected str or "
                    "Vectors object".format(type(vector))
                )

        tot_dim = sum(v.dim for v in vectors)
        self.vectors = torch.rand(len(self), tot_dim) * 0.08 * 2 - 0.08
        for i, token in enumerate(self.itos):
            start_dim = 0
            for v in vectors:
                end_dim = start_dim + v.dim
                self.vectors[i][start_dim:end_dim] = v[token.strip()]
                start_dim = end_dim

            assert start_dim == tot_dim


class QBTextField(Field):
    vocab_cls = QBVocab


HTTP_PREFIX = "http://s3-us-west-2.amazonaws.com/pinafore-us-west-2/qanta-jmlr-datasets"
S3_URL_PATTERN = os.path.join(HTTP_PREFIX, "qanta.torchtext.{fold}.2018.04.18.json")
s3_wiki = "https://s3-us-west-2.amazonaws.com/pinafore-us-west-2/datasets/wikipedia/wiki_lookup.json"


class QuizBowl(Dataset):
    name = "quizbowl"
    dirname = ""
    urls = [S3_URL_PATTERN.format(fold=fold) for fold in ["train", "val", "dev"]]

    @staticmethod
    def sort_key(example):
        if hasattr(example, "text"):
            return len(example.text)
        elif hasattr(example, "unigram"):
            return len(example.unigram)
        elif hasattr(example, "bigram"):
            return len(example.bigram)
        elif hasattr(example, "trigram"):
            return len(example.trigram)
        else:
            raise ValueError("Not valid length fields")

    def __init__(
        self,
        path,
        qanta_id_field,
        sent_field,
        page_field,
        text_field,
        unigram_field,
        bigram_field,
        trigram_field,
        example_mode="sentence",
        use_wiki=False,
        n_wiki_sentences=3,
        replace_title_mentions="",
        **kwargs,
    ):
        from unidecode import unidecode

        if use_wiki and "train" in path:
            base_path = os.path.dirname(path)
            filename = os.path.basename(s3_wiki)
            output_file = os.path.join(base_path, filename)
            if not os.path.exists(output_file):
                download_from_url(s3_wiki, output_file)
            with open(output_file) as f:
                self.wiki_lookup = json.load(f)
        else:
            self.wiki_lookup = {}
        self.path = path
        self.example_mode = example_mode

        text_dependent_fields = []
        if text_field is not None:
            text_dependent_fields.append(("text", text_field))
        if unigram_field is not None:
            text_dependent_fields.append(("unigram", unigram_field))
        if bigram_field is not None:
            text_dependent_fields.append(("bigram", bigram_field))
        if trigram_field is not None:
            text_dependent_fields.append(("trigram", trigram_field))

        example_fields = {
            "qanta_id": [("qanta_id", qanta_id_field)],
            "sent": [("sent", sent_field)],
            "page": [("page", page_field)],
            "text": text_dependent_fields,
        }

        examples = []
        answer_set = set()
        with open(path) as f:
            for ex in json.load(f)["questions"]:
                if example_mode == "sentence":
                    sentences = [
                        ex["text"][start:end] for start, end in ex["tokenizations"]
                    ]
                    for i, s in enumerate(sentences):
                        examples.append(
                            Example.fromdict(
                                {
                                    "qanta_id": ex["qanta_id"],
                                    "sent": i,
                                    "text": unidecode(s),
                                    "page": ex["page"],
                                },
                                example_fields,
                            )
                        )
                        answer_set.add(ex["page"])
                elif example_mode == "question":
                    examples.append(
                        Example.fromdict(
                            {
                                "qanta_id": ex["qanta_id"],
                                "sent": -1,
                                "text": unidecode(ex["text"]),
                                "page": ex["page"],
                            },
                            example_fields,
                        )
                    )
                    answer_set.add(ex["page"])
                else:
                    raise ValueError(
                        f"Valid modes are 'sentence' and 'question', but '{example_mode}' was given"
                    )

        if use_wiki and n_wiki_sentences > 0 and "train" in path:
            print("Loading wikipedia")
            pages = [
                (p, self.wiki_lookup[p]["text"])
                for p in answer_set
                if p in self.wiki_lookup
            ]

            def extract(args):
                title, text = args
                sentences = extract_wiki_sentences(
                    title,
                    text,
                    n_wiki_sentences,
                    replace_title_mentions=replace_title_mentions,
                )
                return title, sentences

            for page, sentences in pseq(pages).map(extract).list():
                for i, s in enumerate(sentences):
                    examples.append(
                        Example.fromdict(
                            {"qanta_id": -1, "sent": i, "text": s, "page": page},
                            example_fields,
                        )
                    )

        dataset_fields = {
            "qanta_id": qanta_id_field,
            "sent": sent_field,
            "page": page_field,
        }
        if text_field is not None:
            dataset_fields["text"] = text_field
        if unigram_field is not None:
            dataset_fields["unigram"] = unigram_field
        if bigram_field is not None:
            dataset_fields["bigram"] = bigram_field
        if trigram_field is not None:
            dataset_fields["trigram"] = trigram_field

        super(QuizBowl, self).__init__(examples, dataset_fields, **kwargs)

    @classmethod
    def splits(
        cls,
        example_mode="sentence",
        use_wiki=False,
        n_wiki_sentences=5,
        replace_title_mentions="",
        root=".data",
        train=f"qanta.torchtext.train.{DS_VERSION}.json",
        validation=f"qanta.torchtext.val.{DS_VERSION}.json",
        test=f"qanta.torchtext.dev.{DS_VERSION}.json",
        **kwargs,
    ):
        remaining_kwargs = kwargs.copy()
        del remaining_kwargs["qanta_id_field"]
        del remaining_kwargs["sent_field"]
        del remaining_kwargs["page_field"]

        remaining_kwargs.pop("text_field", None)
        remaining_kwargs.pop("unigram_field", None)
        remaining_kwargs.pop("bigram_field", None)
        remaining_kwargs.pop("trigram_field", None)
        return super(QuizBowl, cls).splits(
            root=root,
            train=train,
            validation=validation,
            test=test,
            example_mode=example_mode,
            qanta_id_field=kwargs["qanta_id_field"],
            sent_field=kwargs["sent_field"],
            page_field=kwargs["page_field"],
            text_field=kwargs.get("text_field", None),
            unigram_field=kwargs.get("unigram_field", None),
            bigram_field=kwargs.get("bigram_field", None),
            trigram_field=kwargs.get("trigram_field", None),
            use_wiki=use_wiki,
            n_wiki_sentences=n_wiki_sentences,
            replace_title_mentions=replace_title_mentions,
            **remaining_kwargs,
        )

    @classmethod
    def iters(
        cls,
        lower=True,
        example_mode="sentence",
        use_wiki=False,
        n_wiki_sentences=5,
        replace_title_mentions="",
        batch_size=128,
        device=0,
        root=".data",
        vectors="glove.6B.300d",
        unigrams=True,
        bigrams=False,
        trigrams=False,
        combined_ngrams=True,
        combined_max_vocab_size=None,
        unigram_max_vocab_size=None,
        bigram_max_vocab_size=None,
        trigram_max_vocab_size=None,
        sort_within_batch=None,
        **kwargs,
    ):
        QANTA_ID = LongField()
        SENT = LongField()
        PAGE = Field(sequential=False, tokenize=str_split)
        if combined_ngrams:
            tokenizer = create_qb_tokenizer(
                unigrams=unigrams, bigrams=bigrams, trigrams=trigrams
            )
            TEXT = QBTextField(
                batch_first=True, tokenize=tokenizer, include_lengths=True, lower=lower
            )
            train, val, dev = cls.splits(
                qanta_id_field=QANTA_ID,
                sent_field=SENT,
                text_field=TEXT,
                page_field=PAGE,
                root=root,
                example_mode=example_mode,
                use_wiki=use_wiki,
                n_wiki_sentences=n_wiki_sentences,
                replace_title_mentions=replace_title_mentions,
                **kwargs,
            )
            TEXT.build_vocab(train, vectors=vectors, max_size=combined_max_vocab_size)
            PAGE.build_vocab(train)
        else:
            if unigrams:
                unigram_tokenizer = create_qb_tokenizer(
                    unigrams=True, bigrams=False, trigrams=False
                )
                UNIGRAM_TEXT = QBTextField(
                    batch_first=True,
                    tokenize=unigram_tokenizer,
                    include_lengths=True,
                    lower=lower,
                )
            else:
                UNIGRAM_TEXT = None

            if bigrams:
                bigram_tokenizer = create_qb_tokenizer(
                    unigrams=False, bigrams=True, trigrams=False
                )
                BIGRAM_TEXT = QBTextField(
                    batch_first=True,
                    tokenize=bigram_tokenizer,
                    include_lengths=True,
                    lower=lower,
                )
            else:
                BIGRAM_TEXT = None

            if trigrams:
                trigram_tokenizer = create_qb_tokenizer(
                    unigrams=False, bigrams=False, trigrams=True
                )
                TRIGRAM_TEXT = QBTextField(
                    batch_first=True,
                    tokenize=trigram_tokenizer,
                    include_lengths=True,
                    lower=lower,
                )
            else:
                TRIGRAM_TEXT = None

            train, val, dev = cls.splits(
                qanta_id_field=QANTA_ID,
                sent_field=SENT,
                page_field=PAGE,
                unigram_field=UNIGRAM_TEXT,
                bigram_field=BIGRAM_TEXT,
                trigram_field=TRIGRAM_TEXT,
                root=root,
                example_mode=example_mode,
                use_wiki=use_wiki,
                n_wiki_sentences=n_wiki_sentences,
                replace_title_mentions=replace_title_mentions,
                **kwargs,
            )
            if UNIGRAM_TEXT is not None:
                UNIGRAM_TEXT.build_vocab(
                    train, vectors=vectors, max_size=unigram_max_vocab_size
                )
            if BIGRAM_TEXT is not None:
                BIGRAM_TEXT.build_vocab(train, max_size=bigram_max_vocab_size)
            if TRIGRAM_TEXT is not None:
                TRIGRAM_TEXT.build_vocab(train, max_size=trigram_max_vocab_size)
            PAGE.build_vocab(train)

        return BucketIterator.splits(
            (train, val, dev),
            batch_size=batch_size,
            device=device,
            repeat=False,
            sort_within_batch=sort_within_batch,
        )
