"""
TODO: Fill this
"""
from pathlib import Path
import random

from spacy.symbols import PERSON
from spacy.tokens import Span
from spacy.util import minibatch, compounding
import spacy

from wellcomeml.ml.spacy_knowledge_base import SpacyKnowledgeBase


class SpacyEntityLinker(object):
    def __init__(self, kb_path, n_iter=50, print_output=True):
        self.kb_path = kb_path
        self.n_iter = n_iter
        self.print_output = print_output

    def _remove_examples_not_in_kb(self, kb, data):
        # Remove examples with unknown identifiers to the knowlege base
        kb_ids = kb.get_entity_strings()
        for text, annotation in data:
            for offset, kb_id_dict in annotation["links"].items():
                new_dict = {}
                for kb_id, value in kb_id_dict.items():
                    if kb_id in kb_ids:
                        new_dict[kb_id] = value
                    else:
                        print(
                            "Removed",
                            kb_id,
                            "from training because it is not in the KB.",
                        )
                annotation["links"][offset] = new_dict
        return data

    def train(self, data):
        """
        Input:
            data: list of training data in the form
                [('A sentence about Farrar',
                {'links': {(17, 22): {'Q1': 1.0, 'Q2': 0.0}}})]

        See https://spacy.io/usage/linguistic-features#entity-linking
        for where I got this code from
        """
        # TODO: Replace n_iter with self.n_iter
        n_iter = self.n_iter

        kb = SpacyKnowledgeBase()
        kb = kb.load(self.kb_path)
        print("Loaded Knowledge Base from '%s'" % self.kb_path)

        nlp = spacy.blank("en", vocab=kb.vocab)
        nlp.vocab.vectors.name = "spacy_pretrained_vectors"

        entity_linker = nlp.create_pipe("entity_linker")
        entity_linker.set_kb(kb)
        nlp.add_pipe(entity_linker, last=True)

        data = self._remove_examples_not_in_kb(kb, data)

        other_pipes = [pipe for pipe in nlp.pipe_names if pipe != "entity_linker"]
        with nlp.disable_pipes(*other_pipes):
            optimizer = nlp.begin_training()
            for itn in range(n_iter):
                random.shuffle(data)
                losses = {}
                batches = minibatch(data, size=compounding(4.0, 32.0, 1.001))
                for batch in batches:
                    texts, annotations = zip(*batch)
                    nlp.update(
                        texts, annotations, drop=0.2, losses=losses, sgd=optimizer,
                    )
                if self.print_output:
                    print(itn, "Losses", losses)

        self.nlp = nlp
        return self.nlp

    def _get_token_nums(self, doc, char_idx):
        """
        Convert a character index to a token index
        i.e. what number token is character number char_idx in ?
        """
        for i, token in enumerate(doc):
            if char_idx > token.idx:
                continue
            if char_idx == token.idx:
                return i
            if char_idx < token.idx:
                return i

    def predict(self, data):
        """
        See how well the model predicts which entity you are referring to in your data
        Input:
            data: list of test data in the form
                [('A sentence about Farrar',
                {'links': {(17, 22): {'Q1': 1.0, 'Q2': 0.0}}})]
        Output:
           pred_entities_ids: ['Q1', 'Q1', 'Q2']
    """
        # TODO: Replace nlp_el with self.nlp
        nlp_el = self.nlp

        # IMPORTANT
        #
        # This code assumes each training example has one entity
        # e.g. (start, end) = list(values['links'].keys())[0]
        #
        # If the code has more than one entities we predict the
        # first only  e.g. pred_entity_id = doc.ents[0].kb_id_

        pred_entities_ids = []
        for text, values in data:
            (start, end) = list(values["links"].keys())[0]

            doc = nlp_el.tokenizer(text)

            # Set entity span to PERSON
            doc.ents = [
                Span(
                    doc,
                    self._get_token_nums(doc, start),
                    self._get_token_nums(doc, end),
                    label=PERSON,
                )
            ]

            doc = nlp_el.get_pipe("entity_linker")(doc)

            pred_entity_id = doc.ents[0].kb_id_
            pred_entities_ids.append(pred_entity_id)

        return pred_entities_ids

    def save(self, output_dir):
        output_dir = Path(output_dir)
        if not output_dir.exists():
            output_dir.mkdir()
        self.nlp.to_disk(output_dir)
        print("Saved model to", output_dir)

    def load(self, output_dir):
        print("Loading from", output_dir)
        self.nlp = spacy.load(output_dir)
        return self.nlp
