import json
import random

import spacy
from spacy.training.example import Example

# Lade das deutsche Basismodell
nlp = spacy.load("de_core_news_sm")

with open('../annotations.json', encoding='utf-8') as f:
    TRAINING_DATA = json.load(f)

# Wenn das Modell keine benutzerdefinierten Entitäten hat, füge die Entitätenerkennung hinzu
if "ner" not in nlp.pipe_names:
    ner = nlp.create_pipe("ner")
    nlp.add_pipe("ner", last=True)
else:
    ner = nlp.get_pipe("ner")

# Füge die Entitäten hinzu
ner.add_label("START_DATE")
ner.add_label("END_DATE")
ner.add_label("DATE")
ner.add_label("OTHER")

# Trainingsdaten in Beispiele umwandeln
examples = []
for text, annotations in TRAINING_DATA:
    doc = nlp.make_doc(text)
    example = Example.from_dict(doc, annotations)
    examples.append(example)

# Beginne mit dem Training
optimizer = nlp.begin_training()
for epoch in range(30):  # Anzahl der Epochen
    print(f"Epoch {epoch + 1}")
    losses = {}
    # Shuffle und trainiere das Modell mit den Beispielen
    random.shuffle(examples)
    # Trainiere mit den Beispielen
    for example in examples:
        nlp.update([example], drop=0.5, losses=losses)
    print(losses)

# Speichere das trainierte Modell
nlp.to_disk("models/date_model")