"""
This is a modification of codes from two great existing kernels (see acknowledgement below), with early stopping added in the training part. While the public score is closer to the baseline model, I hope this will be helpful to anyone wishing to apply early stopping in the training of spaCy NER in this or other projects.
"""
"""
Acknowledgement: 

* https://www.kaggle.com/rohitsingh9990/ner-training-using-spacy-ensemble

* https://www.kaggle.com/tanulsingh077/twitter-sentiment-extaction-analysis-eda-and-model

Reference: 

* https://spacy.io/usage/training
"""
import numpy as np # linear algebra
import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)
import os
import time
from datetime import datetime
import nltk
import spacy
from spacy.util import minibatch, compounding
import random

import warnings
warnings.filterwarnings("ignore")

pd.set_option('display.max_colwidth', -1)
spacy.prefer_gpu()
BASE_PATH = '../input/tweet-sentiment-extraction/'
MODELS_BASE_PATH = '../working/models/'
train_df = pd.read_csv(BASE_PATH + 'train.csv')
test_df = pd.read_csv( BASE_PATH + 'test.csv')
submission_df = pd.read_csv( BASE_PATH + 'sample_submission.csv')

train_df = train_df.dropna()
train_df = train_df.applymap(lambda x : x.strip())
test_df = test_df.applymap(lambda x : x.strip())

print(f"train_df size: {len(train_df) :,}")
print(f"test_df size: {len(test_df) :,}")
def jaccard(str1, str2): 
    a = set(str1.lower().split()) 
    b = set(str2.lower().split())
    c = a.intersection(b)
    return float(len(c)) / (len(a) + len(b) - len(c))


def get_training_data(sentiment):
    tmp_df = train_df[train_df.sentiment == sentiment]    
    train_data = []
    for textID, text, selected_text, _ in tmp_df.values:
        start = text.find(selected_text)
        end = start + len(selected_text)
        train_data.append((text, {"entities": [[start, end, 'selected_text']]}))
    return train_data


def save_model(output_dir, nlp, new_model_name):
    output_dir = f'{MODELS_BASE_PATH}{output_dir}'
    if output_dir is not None:        
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        nlp.meta["name"] = new_model_name
        nlp.to_disk(output_dir)
        print("Saved model to", output_dir)

        
def predict_entities(text, model):
    doc = model(text)
    start, end = -1, -1
    for ent in doc.ents:
        start = text.find(ent.text)
        end = start + len(ent.text)  
        break #We're only interested in the first one (and there should only be one!)
        
    if start > -1:
        return text[start: end]
    else:        
        return text
    
    
def get_jaccard(data, model, to_print=False, N=-1):
    jaccard_sum = 0
    for text, annotations in data[:N]:
        start = annotations['entities'][0][0]
        end = annotations['entities'][0][1]
        selected_text = text[start:end]
        
        pred = predict_entities(text, model)
        jaccard_sum += jaccard(selected_text, pred)
        
        if to_print:
            print(f"{text}; 'predict': {pred}; 'actual': {selected_text}")        
        
    jaccard_mean = jaccard_sum / len(data)
    
    return jaccard_mean

# Code modified from https://spacy.io/usage/training

def train(train_data, sentiment, val_size=0.2, patience=2, model=None):
    """Load the model, set up the pipeline and train the entity recognizer."""
    if model is not None:
        nlp = spacy.load(model)  # load existing spaCy model
        print("Loaded model '%s'" % model)
    else:
        nlp = spacy.blank("en")  # create blank Language class
        print("Created blank 'en' model")

    # create the built-in pipeline components and add them to the pipeline
    # nlp.create_pipe works for built-ins that are registered with spaCy
    if "ner" not in nlp.pipe_names:
        ner = nlp.create_pipe("ner")
        nlp.add_pipe(ner, last=True)
    # otherwise, get it so we can add labels
    else:
        ner = nlp.get_pipe("ner")

    # add labels
    for _, annotations in train_data:
        for ent in annotations.get("entities"):
            ner.add_label(ent[2])

    # get names of other pipes to disable them during training
    pipe_exceptions = ["ner"]    
    other_pipes = [pipe for pipe in nlp.pipe_names if pipe not in pipe_exceptions]
    
    # only train NER
    with nlp.disable_pipes(*other_pipes) and warnings.catch_warnings():
        # show warnings for misaligned entity spans once
        warnings.filterwarnings("once", category=UserWarning, module='spacy')

        # reset and initialize the weights randomly – but only if we're
        # training a new model
        if model is None:
            nlp.begin_training()
            
        last_val_jaccard = 0
        patience_ = 0
        itn = 0
            
        while True:
            random.shuffle(train_data)
            losses = {}
            
            # Split into train and validation set
            N = int(np.floor(len(train_data) * (1-val_size)))
            train = train_data[:N]
            val   = train_data[N:]
            
            # batch up the examples using spaCy's minibatch
            batches = minibatch(train, size=compounding(4.0, 30.0, 1.001))
            for batch in batches:
                texts, annotations = zip(*batch)
                nlp.update(
                    texts,  # batch of texts
                    annotations,  # batch of annotations
                    drop=0.5,  # dropout - make it harder to memorise data
                    losses=losses,
                )      
                    
            #Calculate Jaccard for train and val dataset
            train_jaccard = get_jaccard(train, nlp)
            val_jaccard = get_jaccard(val, nlp)
                                    
            print(f"iter {itn}. Losses: {losses['ner'] :.2f}; Train Jaccard: {train_jaccard :.4f}; Val Jaccard: {val_jaccard :.4f};")
                                                                
            if val_jaccard > last_val_jaccard:
                save_model(f"{sentiment}", nlp, f"nlp_{sentiment}")
                last_val_jaccard = val_jaccard
                patience_ = 0
            else:
                patience_ += 1   
                print(f"patience count increased to {patience_}")                
                                        
            if patience_ >= patience:
                break
                
            itn += 1
                
    return nlp
""" Training the NER models, one for each sentiment
"""

to_train_model  = True
# to_train_model  = False

# model_tags = ['positive', 'negative', 'neutral']
model_tags = ['positive', 'negative']

val_size = 0.2 
patience = 2

if to_train_model:
    t0 = time.time()
   
    for sentiment in model_tags:
        print(f"Training '{sentiment}' model:")
        train_data = get_training_data(sentiment)
        nlp = train(train_data, sentiment, val_size=val_size, patience=patience)

        print(f"Training NER for '{sentiment}' takes {(time.time()-t0)/60 :.2f} minutes")
        t0 = time.time()
        print('')
""" Loading the models
"""

models = dict()
for sentiment in model_tags:
    models[sentiment] = spacy.load(f"{MODELS_BASE_PATH}{sentiment}")    
%%time
""" Apply models to test data
"""

selected_texts = []

for textID, text, sentiment in test_df.values:
    if sentiment == 'neutral' or len(text.split()) < 3:
        selected_texts.append(text)    
    else:
        pred = predict_entities(text, models[sentiment])
        selected_texts.append(pred)
             
# for textID, text, sentiment in test_df.values:    
#     selected_texts.append(predict_entities(text, models[sentiment]))

""" Submission 
"""
submission_df['selected_text'] = selected_texts
submission_df.to_csv("submission.csv", index=False)
display(submission_df.head(10))

dt_string = datetime.now().strftime("%d/%m/%Y %H:%M:%S")
print(f"Successfully submitted at {dt_string}")