from collections import defaultdict

import torch
import os
import json
import numpy as np
import time
from mock import patch

from django.http import JsonResponse
from django.shortcuts import render
from django.utils.datastructures import MultiValueDict
from django.views import View
from django.views.generic import TemplateView

from capreolus.collection import Collection
from capreolus.demo_app.constants import WEIGHTS_DIR_NAME, WEIGHTS_FILE_NAME, NUM_RESULTS_TO_SHOW
from capreolus.demo_app.forms import QueryForm, BM25Form
from capreolus.demo_app.utils import search_files_or_folders_in_directory, copy_keys_to_each_other
from capreolus.extractor.embedding import EmbeddingHolder
from capreolus.index import Index
from capreolus.reranker.reranker import Reranker
from capreolus.pipeline import Pipeline
from capreolus.tokenizer import Tokenizer
from capreolus.utils.common import padlist, get_default_cache_dir, get_default_results_dir
from capreolus.utils.loginit import get_logger


logger = get_logger(__name__)


class ConfigsView(TemplateView):
    """
    Returns a dict of configurations based on available experiment results
    """

    template_name = "demo.html"

    def get(self, request, *args, **kwargs):
        available_configs = self.get_config_from_results()
        available_indices = self.get_available_indices()
        context = {"configs": available_configs, "available_indices": available_indices}

        return render(request, self.template_name, context)

    @staticmethod
    def get_available_indices():
        cache_path = os.environ.get("CAPREOLUS_CACHE", get_default_cache_dir())
        index_dirs = search_files_or_folders_in_directory(cache_path, "index")
        index_dirs_with_done = [
            index_dir for index_dir in index_dirs if len(search_files_or_folders_in_directory(index_dir, "done"))
        ]

        return index_dirs_with_done

    @staticmethod
    def get_config_from_results():
        # TODO: Using 'results' if there's no env varialbe - this logic is duplicated. Move to a single place?
        results_folder = os.environ.get("CAPREOLUS_RESULTS", get_default_results_dir())
        config_files = search_files_or_folders_in_directory(results_folder, "config.json")
        configs = []

        for file in config_files:
            with open(file) as json_file:
                data = json.load(json_file)
                configs.append(data)

        return configs


class ExperimentView(TemplateView):
    """
        Displays a query box, along with the chosen experiment configs
    """

    template_name = "experiment.html"

    def get(self, request, *args, **kwargs):
        context = {"config": sorted(request.GET.items())}
        return render(request, self.template_name, context)


class QuerySuggestionView(View):
    """
    Used in auto-completing query suggestions in the demo page
    This works only if the index that the user chose was created from a collection
    """

    def get(self, request, *args, **kwargs):
        query = request.GET.dict()["q"]
        target_index = request.GET.dict()["target_index"]
        collection = Collection.get_collection_from_index_path(target_index)

        if collection is None:
            return JsonResponse([], safe=False)
        else:
            suggestions = collection.get_query_suggestions(query)
            return JsonResponse(suggestions, safe=False)


class QueryView(TemplateView):
    def get(self, request, *args, **kwargs):
        model_name = request.GET.dict()["reranker"]
        if model_name == "BM25":
            return BM25View().get(request, *args, **kwargs)
        else:
            return NeuralQueryView().get(request, *args, **kwargs)


class NeuralQueryView(TemplateView):
    """
        Perform a query against the chosen pipeline and display the results
    """

    template_name = "experiment.html"

    def get(self, request, *args, should_render=True, **kwargs):
        # The expected query params differ based on the NIR model. We use django forms to help us do the validation
        formcls = QueryForm.FORM_CLASSES[request.GET.dict()["reranker"]]
        form = formcls(request.GET.dict())

        if form.is_valid():
            # TODO: Keep a list of allowed/known experiment config and use that to extract config from request.GET
            # request.GET might contain key value pairs that won't be recognized by sacred
            query_string = form.cleaned_data.pop("query")
            target_index = form.cleaned_data.pop("target_index")
            config = form.cleaned_data

            results = self.get_most_relevant_doc_based_on_config(config, query_string, target_index)
            # Hack. Re-adding target_index because the html page would need it for subsequent queries
            config["target_index"] = target_index
            context = {
                "config": sorted(config.items()),
                "query": query_string,
                "has_result": True,
                "results": results,
                "target_index": target_index,
            }
            if should_render:
                return render(request, self.template_name, context)
            else:
                return context
        else:
            raise Exception("Malformed query: {0}".format(form.errors))

    # TODO: Move this to the Pipeline class maybe?
    @staticmethod
    def get_weight_file(pipeline):
        """
        Given a pipeline, get the path to a weight file.
        If there are multiple weight files, returns the path to the first one
        Here weight file = a file created from torch.model.save()
        """
        run_path = os.path.join(pipeline.model_path, pipeline.cfg["fold"])
        weight_dir = os.path.join(run_path, WEIGHTS_DIR_NAME)
        weight_file = search_files_or_folders_in_directory(weight_dir, WEIGHTS_FILE_NAME)
        return weight_file[0]

    @staticmethod
    def get_tokenizer(pipeline, config, tokenizer_name):
        tokenizer_class = Tokenizer.ALL[tokenizer_name]
        tokenizer = tokenizer_class(pipeline, keepstops=config.get("keepstops", False), use_cache=False)
        tokenizer.create()
        return tokenizer

    @staticmethod
    def get_tokens_from_docs_and_query(tokenizer, docs, query_string):
        """
        Returns one big list of tokens, based on the terms that appear in supplied documents and the query string
        """
        all_tokens = [tokenizer.tokenize(doc) for doc in docs]
        # Just flattening
        all_tokens = [token for sublist in all_tokens for token in sublist]
        all_tokens.extend(tokenizer.tokenize(query_string))
        return all_tokens

    @staticmethod
    def create_tensor_from_docs(docs, tokenizer, embedding_holder, maxdoclen):
        """
        Returns a stacked tensor, with each element being a document's tensor representation
        """
        doc_features = []
        for doc in docs:
            doc_tokens = tokenizer.tokenize(doc)
            # Doc features should be indices into the embedding layer. This is an optimization enforced by pytorch
            doc_token_indices = embedding_holder.get_index_array_from_tokens(doc_tokens, maxdoclen)
            doc_features.append(torch.from_numpy(doc_token_indices))

        return torch.stack(doc_features)

    @staticmethod
    def create_tensor_from_query_string(query_string, index, tokenizer, embedding_holder, doc_features_len, maxqlen):
        """
        Creates a tensor from the query string, and stacks it repeatedly (i.e duplicate) so as to match the dimensions
        of the document features tensor
        :param query_string
        :param tokenizer - A tokenizer instance
        :param embedding_holder - An instance of EmbeddingHolder
        :param doc_features_len - The length of the document features tensor. This dictates the size of the resulting
        query features tensor created by this method
        :param maxqlen - Maximum query length
        """
        query_toks = tokenizer.tokenize(query_string)
        query_feature = torch.from_numpy(embedding_holder.get_index_array_from_tokens(query_toks, maxqlen))
        query_features = torch.stack([query_feature for i in range(0, doc_features_len)])

        query_idfs = []
        for i in range(0, doc_features_len):
            query_idfs.append(torch.from_numpy(np.array(padlist([index.getidf(tok) for tok in query_toks], maxqlen))))
        query_idfs = torch.stack(query_idfs)

        return query_features, query_idfs

    @staticmethod
    def add_model_required_params_to_config(config, embedding_holder):
        # So that MODEL_BASE.validate_params() would pass
        # We are adding params required by all the NIR model classes. The alternative is to do if condition checks
        # for the model class and add only the relevant params. But this is cleaner, and we are not saving anything
        # with the if checks
        config["pad_token"] = 0
        config["stoi"] = embedding_holder.get_stoi()
        config["itos"] = embedding_holder.get_itos()
        config["nvocab"] = embedding_holder.get_nvocab()

        return config

    @staticmethod
    def construct_result_dicts(doc_ids, docs, relevances):
        result_dicts = []

        for i in range(len(doc_ids)):
            result_dicts.append({"doc_id": doc_ids[i], "doc": docs[i], "relevance": relevances[i]})
        return result_dicts

    @staticmethod
    def get_most_relevant_doc_based_on_config(config, query_string, target_index):
        """
        1. Instantiate various classes based on config
        2. Get the most relevant doc
        """
        # We still need to init a pipeline because it pre-processes some config params, and we rely on that to
        # construct paths e.t.c.
        config = config.copy()  # because we end up modifying config
        pipeline = Pipeline(config)
        pipeline.from_demo_app = True
        pipeline.initialize(config)
        path_dict = pipeline.get_paths(config)
        index_path = target_index
        index_class = Index.get_index_from_index_path(index_path)
        index = index_class(pipeline.collection, index_path, None)  # TODO: Pass a proper index_key
        model_class = Reranker.ALL[config["reranker"]]
        tokenizer = NeuralQueryView.get_tokenizer(pipeline, config, index_class.name)
        embedding_holder = EmbeddingHolder.get_instance(config.get("embeddings", "glove6b"))
        trained_weight_path = path_dict["trained_weight_path"]
        config = NeuralQueryView.add_model_required_params_to_config(config, embedding_holder)

        return NeuralQueryView.do_query(
            config,
            query_string,
            pipeline,
            index,
            tokenizer,
            embedding_holder,
            model_class,
            trained_weight_path=trained_weight_path,
        )

    @staticmethod
    def do_query(config, query_string, pipeline, index, tokenizer, embedding_holder, model_class, trained_weight_path=None):
        """
        1. Do a bm25 search to get top 100 results. This is an optimization technique - we don't want to feed the entire
        dataset to the NIR model's forward pass
        2. Based on the documents retrieved above, create an embedding layer to be used in the NIR model
        3. Instantiate an NIR model, load the trained weights, and do a forward pass with the 1000 docs
        4. Return the document with the highest score
        """
        # 1. Do bm25 search and tokenize the results
        # TODO: Handle the case where bm25 search returns no results
        api_start = time.time()
        doc_ids, docs = BM25View.do_query(query_string, index)

        all_tokens = NeuralQueryView.get_tokens_from_docs_and_query(tokenizer, docs, query_string)

        # 2. Form an embedding layer and doc and query features
        embeddings = torch.from_numpy(embedding_holder.create_indexed_embedding_layer_from_tokens(all_tokens)).to(pipeline.device)
        doc_features = NeuralQueryView.create_tensor_from_docs(docs, tokenizer, embedding_holder, config["maxdoclen"]).to(
            pipeline.device
        )
        query_features, query_idf_features = NeuralQueryView.create_tensor_from_query_string(
            query_string, index, tokenizer, embedding_holder, len(doc_features), config["maxqlen"]
        )
        query_features = query_features.to(pipeline.device)
        query_idf_features = query_idf_features.to(pipeline.device)

        # 3. Do a forward pass of the NIR model class, and get the max scoring document
        # TODO: Remove the dependence of NIR reranker on pipeline. Pass the configs explicitly
        model_instance = model_class(embeddings, None, config)
        model_instance.build()
        model_instance.to(pipeline.device)
        # model_instance = model_class.alternate_init(pipeline, embeddings, config).to(pipeline.device)
        if trained_weight_path is not None:
            model_instance.load(trained_weight_path)
        scores = model_instance.test(query_features, query_idf_features, doc_features)
        max_scoring_doc_ids = [doc_ids[i] for i in reversed(torch.argsort(scores))]
        max_scoring_doc_ids = max_scoring_doc_ids[:NUM_RESULTS_TO_SHOW]
        # TODO: Get this in a single go, instead of calling `getdoc()` repeatedly
        docs = [index.getdoc(doc_id) for doc_id in max_scoring_doc_ids]
        # Trimming to first 250 chars. TODO: Make this configurable
        docs = [doc[:250] for doc in docs]

        relevance_fetch_start = time.time()
        collection = Collection.get_collection_from_index_path(index.index_path)
        relevances = collection.get_relevance(query_string, max_scoring_doc_ids) if collection is not None else [0] * len(doc_ids)
        result_dicts = NeuralQueryView.construct_result_dicts(max_scoring_doc_ids, docs, relevances)
        relevance_fetch_stop = time.time()

        api_stop = time.time()
        logger.debug("Took {0} seconds to get the most relevant doc".format(api_stop - api_start))
        logger.debug(
            "Determining the relevance of the fetched document took {0} seconds".format(
                relevance_fetch_stop - relevance_fetch_start
            )
        )
        return result_dicts


class DocumentView(TemplateView):
    """
        Displays a single document
    """

    template_name = "document.html"

    def get(self, request, *args, **kwargs):
        index_class = Index.get_index_from_index_path(request.GET["target_index"])
        # Hack - set the collection and index key as None since we won't need them for the purpose of WSDM demo
        index = index_class(None, request.GET["target_index"], None)
        doc_id = request.GET["doc_id"]
        context = {"doc": index.getdoc(doc_id)}

        return render(request, self.template_name, context)


class QueryDictParserMixin:
    @staticmethod
    def get_two_configs_from_query_dict(query_dict):
        """
        Parses the django query dict and separates it into two different experiment configs.
        eg: "/api/?model=KNRM&model=DRMM&blah..." should be parsed so that "model=KNRM" goes to one config while
        "model=DRMM" goes to the other
        """
        assert isinstance(query_dict, MultiValueDict)

        config_1, config_2 = dict(), dict()

        for key, values in query_dict.lists():
            if len(values) == 2:
                config_1[key] = values[0]
                config_2[key] = values[1]
            elif len(values) == 1:
                config_1[key] = values[0]
                config_2[key] = values[0]
            else:
                # TODO: This is a 500 error. Replace this with a 400 bad request error
                raise Exception("Bad request")

        return config_1, config_2


class CompareExperimentsView(TemplateView, QueryDictParserMixin):
    """
        Displays a query box, along with configs of 2 different experiments
    """

    def get(self, request, *args, **kwargs):
        config_1, config_2 = self.get_two_configs_from_query_dict(request.GET)
        zipped_configs = list(zip(sorted(config_1.items()), sorted(config_2.items())))
        context = {"configs": zipped_configs, "config_1": sorted(config_1.items()), "config_2": sorted(config_2.items())}

        return render(request, "experiment_comparison.html", context)


class CompareQueryView(TemplateView, QueryDictParserMixin):
    def get(self, request, *args, **kwargs):
        config_1, config_2 = self.get_two_configs_from_query_dict(request.GET)
        if "BM25" in [config_1["reranker"], config_2["reranker"]]:
            return CompareBM25AndNeuralQueryView().get(request, *args, **kwargs)
        else:
            return CompareNeuralQueryView().get(request, *args, **kwargs)


class CompareNeuralQueryView(TemplateView, QueryDictParserMixin):
    """
    Perform a query against two different NIR reranker and display the results for comparison
    """

    template_name = "experiment_comparison.html"

    def get(self, request, *args, **kwargs):
        config_1, config_2 = self.get_two_configs_from_query_dict(request.GET)
        formcls_1 = QueryForm.FORM_CLASSES[config_1["reranker"]]
        form_1 = formcls_1(config_1)
        formcls_2 = QueryForm.FORM_CLASSES[config_2["reranker"]]
        form_2 = formcls_2(config_2)

        if form_1.is_valid() and form_2.is_valid():
            query_string = form_1.cleaned_data.pop("query")
            form_2.cleaned_data.pop("query")
            target_index = form_1.cleaned_data.pop("target_index")
            form_2.cleaned_data.pop("target_index")

            config_1 = form_1.cleaned_data
            config_2 = form_2.cleaned_data
            results_1 = NeuralQueryView.get_most_relevant_doc_based_on_config(config_1, query_string, target_index)
            results_2 = NeuralQueryView.get_most_relevant_doc_based_on_config(config_2, query_string, target_index)

            combined_results = self.combine_results(results_1, results_2)

            # Hack. Re-adding target_index because the html page would need it for subsequent queries
            config_1["target_index"], config_2["target_index"] = target_index, target_index

            # Need to do this because we are highlighting differences between the configs in the UI
            config_1, config_2 = copy_keys_to_each_other(config_1, config_2)
            zipped_configs = list(zip(sorted(config_1.items()), sorted(config_2.items())))
            context = {
                "configs": zipped_configs,
                "model_1": config_1["reranker"],
                "model_2": config_2["reranker"],
                "query": query_string,
                "has_result": True,
                "combined_results": sorted(combined_results.items()),
                "target_index": target_index,
            }

            return render(request, self.template_name, context)
        else:
            raise Exception("Malformed query: {0}, {1}".format(form_1.errors, form_2.errors))

    @staticmethod
    def combine_results(result_1, result_2):
        combined_results = defaultdict(lambda: {"config_1_rank": None, "config_2_rank": None, "relevance": 0})

        for i, result in enumerate(result_1):
            doc_id = result["doc_id"]
            combined_results[doc_id]["config_1_rank"] = i + 1
            combined_results[doc_id]["relevance"] = result["relevance"]

        for i, result in enumerate(result_2):
            doc_id = result["doc_id"]
            combined_results[doc_id]["config_2_rank"] = i + 1
            combined_results[doc_id]["relevance"] = result["relevance"]

        return combined_results


class BM25View(TemplateView):
    template_name = "experiment.html"

    def get(self, request, *args, **kwargs):
        request_dict = request.GET.dict()
        form = BM25Form(request_dict)

        if form.is_valid():
            cleaned_data = form.cleaned_data
            results = self.get_bm25_results(
                cleaned_data["query"], cleaned_data["target_index"], cleaned_data.get("b"), cleaned_data.get("k1")
            )
            context = {
                "config": sorted(cleaned_data.items()),
                "query": request_dict["query"],
                "has_result": True,
                "results": results,
                "target_index": request_dict["target_index"],
            }
            return render(request, self.template_name, context)
        else:
            raise ValueError("Query has errors: {0}".format(form.errors))

    @staticmethod
    def get_bm25_results(query_string, target_index, b, k1):
        index_class = Index.get_index_from_index_path(target_index)
        index = index_class(None, target_index, None)

        bm25_kwargs = {"n": NUM_RESULTS_TO_SHOW}
        if b is not None:
            bm25_kwargs["b"] = b
        if k1 is not None:
            bm25_kwargs["k1"] = k1

        doc_ids, docs = BM25View.do_query(query_string, index, **bm25_kwargs)
        docs = [doc[:250] for doc in docs]
        collection = Collection.get_collection_from_index_path(index.index_path)
        relevances = collection.get_relevance(query_string, doc_ids) if collection is not None else [0] * len(doc_ids)
        result_dicts = NeuralQueryView.construct_result_dicts(doc_ids, docs, relevances)
        return result_dicts

    @staticmethod
    def do_query(query_string, index, k1=1.2, b=0.75, n=100):
        """
        Does a bm25 search and returns the most relevant 1000 docs and their ids
        """
        with patch("pyserini.setup.configure_classpath") as mock_setup:
            mock_setup.return_value = None
            from pyserini.search import pysearch

            searcher = pysearch.SimpleSearcher(index.index_path)
            searcher.set_bm25_similarity(k1, b)
            hits = searcher.search(query_string, n)
            doc_ids = [hit.docid for hit in hits]
            docs = [index.getdoc(doc_id) for doc_id in doc_ids]

            return doc_ids, docs


class CompareBM25AndNeuralQueryView(TemplateView, QueryDictParserMixin):
    template_name = "experiment_comparison.html"

    def get(self, request, *args, **kwargs):
        config_1, config_2 = self.get_two_configs_from_query_dict(request.GET)
        nir_config = config_2 if config_2["reranker"] != "BM25" else config_1
        bm25_config = config_2 if config_2["reranker"] == "BM25" else config_1

        bm25_form = BM25Form(bm25_config)
        query_string = nir_config["query"]
        target_index = nir_config["target_index"]
        nir_form_class = QueryForm.FORM_CLASSES[nir_config["reranker"]]
        nir_form = nir_form_class(nir_config)

        if bm25_form.is_valid() and nir_form.is_valid():
            bm25_data = bm25_form.cleaned_data
            bm25_results = BM25View.get_bm25_results(query_string, target_index, bm25_data.get("b"), bm25_data.get("k1"))
            nir_data = nir_form.cleaned_data

            # TODO: Make pipeline initialization robust to extra parameters. We have to remove query and target index
            # because pipeline init errors out otherwise
            nir_data.pop("target_index")
            nir_data.pop("query")
            nir_results = NeuralQueryView.get_most_relevant_doc_based_on_config(nir_data, query_string, target_index)
            combined_results = CompareNeuralQueryView.combine_results(bm25_results, nir_results)
            nir_data["target_index"] = target_index
            bm25_config = {"reranker": "BM25", "b": bm25_data["b"], "k1": bm25_data["k1"], "target_index": target_index}

            # Need to do this because we are highlighting differences between the configs in the UI
            bm25_config, nir_data = copy_keys_to_each_other(bm25_config, nir_data)
            zipped_configs = list(zip(sorted(bm25_config.items()), sorted(nir_data.items())))
            context = {
                "configs": zipped_configs,
                "model_1": "BM25",
                "model_2": nir_data["reranker"],
                "query": nir_config["query"],
                "has_result": True,
                "combined_results": sorted(combined_results.items()),
                "target_index": nir_config["target_index"],
            }

            return render(request, self.template_name, context)
        else:
            raise ValueError("Query is malformed: BM25: {0}, NIR: {1}".format(bm25_form.errors, nir_form.errors))
