# NER: https://spacy.io/usage/linguistic-features#named-entities
# Compatible with: spaCy v2.0.0+

# Standard libs
from __future__ import unicode_literals, print_function
import random
from pathlib import Path

# spacy als Standard für Textverarbeitung
import spacy
from spacy.util import minibatch, compounding
from spacy import displacy

# Check, welche Version genutzt wird. Ändert sich häufig mit starken Anpassungen.
print(spacy.__version__)

# Laden des deutschen Sprachmodells
# https://spacy.io/models/de#de_core_news_sm
nlp = spacy.load('de_core_news_sm')

# Einlesen eines Satzes als Unicode-String
doc = nlp(u'Die Bundesregierung wechselte ihren Sitz von Bonn nach Berlin.')

# Ausgabe der erkannten Objekte (Start, Ende) und Art
for ent in doc.ents:
    print(ent.text, ent.start_char, ent.end_char, ent.label_)

# Visualisierung der erkannten Textbestandteile
displacy.render(doc, style='ent', jupyter=True)

# Visualisierung der Satzstruktur (Dependency Model)
displacy.render(doc, style='dep', jupyter=True)

doc = nlp(u'Theresa Mays Plan B könnte erneut Plan A sein, mit dem die Premierministerin in der vergangenen Woche so krachend im Unterhaus gescheitert ist: das mit Brüssel ausgehandelte Austrittsabkommen. Sie könnte jedoch eine wichtige Änderung vorschlagen: May will, nach Berichten britischer Medien, den Backstopp, die Auffanglösung für die Grenze auf der irischen Insel, noch einmal mit der EU nachverhandeln.')

displacy.render(doc, style='ent', jupyter=True)

# Hinzufügen der Objekte Bundesregierung (ORG) und London/München (LOC)
TRAIN_DATA = [
    ("Die Bundesregierung wechselte ihren Sitz von Bonn nach Berlin.", {"entities": [(4, 19, "ORG"), 
                                                                                     (45, 49, "LOC"),
                                                                                     (55, 61, "LOC")]}),
    ("London, Muenchen", {"entities": [(1, 7, "LOC"), (9, 17, "LOC")]}),
]

# Typische Anzahl für Lerniterationen n = 100, hier "nur" 10 fürs schnelle Durchlaufen
def main(model=None, output_dir="./spacy_modell/", n_iter=10):
    # NER Pipeline laden    
    ner = nlp.get_pipe("ner")
    # Datenvorbereitung
    for _, annotations in TRAIN_DATA:
        for ent in annotations.get("entities"):
            ner.add_label(ent[2])
    # Andere Pipelines fürs Lernen abschalten
    other_pipes = [pipe for pipe in nlp.pipe_names if pipe != "ner"]
    with nlp.disable_pipes(*other_pipes):  # only train NER
        # Gewichte des neuen Modells setzen
        if model is None:
            nlp.begin_training()
        for itn in range(n_iter):
            random.shuffle(TRAIN_DATA)
            losses = {}
            batches = minibatch(TRAIN_DATA, size=compounding(4.0, 32.0, 1.001))
            for batch in batches:
                texts, annotations = zip(*batch)
                nlp.update(
                    texts,  
                    annotations,  
                    drop=0.5,  
                    losses=losses,
                )
            print("Losses", losses)

    # Test der trainierten Daten
    for text, _ in TRAIN_DATA:
        doc = nlp(text)
        print("Entities", [(ent.text, ent.label_) for ent in doc.ents])
        print("Tokens", [(t.text, t.ent_type_, t.ent_iob) for t in doc])

    # Modell speichern
    nlp.to_disk(Path(output_dir))

# Methode aufrufen
main()

nlp = spacy.load("./spacy_modell/")
doc = nlp(u'Hannover ist die schönste Stadt in Deutschland.')

for ent in doc.ents:
    print(ent.text, ent.start_char, ent.end_char, ent.label_)

displacy.render(doc, style='ent', jupyter=True)

doc = nlp(u'Die Bundesregierung wechselte ihren Sitz von Bonn nach Berlin.')
displacy.render(doc, style='ent', jupyter=True)

options = {'compact': True, 'color': 'blue'}
displacy.render(doc, style='dep', jupyter=True, options=options)