
import csv
import os
import pickle
import string

from nltk import word_tokenize
from collections import defaultdict
from nltk import FreqDist
import visualkeras

import nltk
from nltk.corpus import stopwords
from sklearn.model_selection import train_test_split
from sklearn.naive_bayes import MultinomialNB
#from sklarn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn import metrics
import random
from sklearn.metrics import *
import numpy as np
from sklearn.metrics import accuracy_score
from sklearn.pipeline import make_pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
import seaborn as sns
import matplotlib.pyplot as plt
import nlpaug.augmenter.sentence as nas
import torch
from keras.preprocessing.sequence import *
from keras.preprocessing.text import Tokenizer
from keras.models import Sequential
from keras.layers import *
import tensorflow as tf
from keras.layers import Dense
from sklearn.preprocessing import LabelEncoder
from keras.models import load_model
from transformers import pipeline


# Variable que determina les 'stopwords' en castellà
stop_words = set(stopwords.words('spanish'))

# Directori root per les mostres de les que disposo (Local)
BASE_DIR = 'C:/Users/Guillem/OneDrive - Universitat de Barcelona/4rt/TFG/Files/Muestras_ALL'
TEST_DIR = 'C:/Users/Guillem/OneDrive - Universitat de Barcelona/4rt/TFG/Files/Muestras_ALL_TestFiles'
LABELS = ['DR','F','FD','IT','NS','TG']
test_files = []

# Variable per determinar el nombre de fitxers per categoria
counts = {'DR':0,'F':0,'FD':0,'IT':0,'NS':0,'TG':0}


# Funció usada per crear el dataset en format csv a partir dels fitxers TXT
def create_data_set(docs):
    with open('data.csv', 'w', encoding='utf8', newline='') as outfile:
        for label in LABELS:
            dir = '%s/%s' % (BASE_DIR, label)
            for filename in os.listdir(dir):
                fullfilename = '%s/%s' % (dir, filename)
                #print(fullfilename)
                with open(fullfilename, 'r') as file:
                    text = file.read()
                    text = text.splitlines()
                    new_text = ""
                    for item in text:
                        new_text = new_text + item
                    #if random.randint(0,10) == 1:
                    #    test_files.append((label,new_text))
                    writer = csv.writer(outfile, delimiter=',')
                    writer.writerow([label, filename, new_text])
                    docs.append((label, new_text))
                    if label == 'DR':
                        counts['DR'] = counts['DR'] + 1
                    elif label == 'F':
                        counts['F'] = counts['F'] + 1
                    elif label == 'FD':
                        counts['FD'] = counts['FD'] + 1
                    elif label == 'IT':
                        counts['IT'] = counts['IT'] + 1
                    elif label == 'NS':
                        counts['NS'] = counts['NS'] + 1
                    elif label == 'TG':
                        counts['TG'] = counts['TG'] + 1
                    else:
                        raise TypeError("We are missing some label!")
                    #outfile.write('%s\t%s\t%s\n' % (label, filename, text))
    print(len(docs))
    return docs


# Funció per entrenar el model de NB
def train_classifier(docs):
    """X = []
    y = []
    for x_,y_ in docs:
        X.append(y_)
        y.append(x_)

    #for i in range(0, len(X) - 1):
    #    try:
    #        new_text = text_aug_rand(X[i],i)
    #        X.append(new_text[0])
    #        y.append(y[i])
    #    except ValueError:
    #        pass

    level_ds(X,y)
    print(len(X))
    X_train, X_test, y_train, y_test = train_test_split(X,y, test_size=0.20, random_state=42)"""
    X_train, X_test, y_train, y_test = split_train_test(docs)

    vectorizer = CountVectorizer(stop_words='english',
                                 ngram_range=(1,3),
                                 min_df=3, analyzer='word')
    """
    dtm = vectorizer.fit_transform(X_train)
    naive_bayes_classifier = MultinomialNB().fit(dtm, y_train)
    evaluate_classifier("Naive Bayes TRAIN", naive_bayes_classifier, vectorizer, X_train, y_train)
    evaluate_classifier("Naive Bayes TEST", naive_bayes_classifier, vectorizer, X_test, y_test)
    clf_filename = 'naive_bayes_classifier.pkl'
    pickle.dump(naive_bayes_classifier, open(clf_filename,'wb'))
    vec_filename = 'count_vectorizer.pk1'
    pickle.dump(vectorizer, open(vec_filename, 'wb'))
    """
    # Model building
    model = make_pipeline(TfidfVectorizer(), MultinomialNB())
    # Training the model with the training data
    model.fit(X_train, y_train)
    # Predicting the test data categories
    predicted_categories = model.predict(X_test)

    # plotting the confusion matrix
    mat = confusion_matrix(y_test, predicted_categories)
    sns.set()
    sns.heatmap(mat.T, square=True, annot=True, fmt="d", xticklabels=LABELS,yticklabels=LABELS)
    plt.xlabel("true labels")
    plt.ylabel("predicted label")
    plt.show()

    mlcm = multilabel_confusion_matrix(y_test, predicted_categories)
    index = 0
    for item in mlcm:
        sns.set()
        sns.heatmap(item.T, square=True, annot=True, fmt="d", xticklabels=LABELS[index], yticklabels=LABELS[index])
        plt.xlabel("true labels")
        plt.ylabel("predicted label")
        plt.show()
        index += 1

    print("Accuracy: {}".format(accuracy_score(y_test, predicted_categories)))

    """evaluate_classifier("Naive Bayes TEST", model, vectorizer, X_test, y_test)
    clf_filename = 'naive_bayes_classifier.pkl'
    pickle.dump(model, open(clf_filename, 'wb'))
    vec_filename = 'count_vectorizer.pk1'
    pickle.dump(vectorizer, open(vec_filename, 'wb'))"""


# Funció per a la creació manual d'un test set
def create_test_set():
    for label in LABELS:
        dir = '%s/%s' % (TEST_DIR, label)
        for filename in os.listdir(dir):
            fullfilename = '%s/%s' % (dir, filename)
            with open(fullfilename, 'r') as file:
                text = file.read()
                text = text.splitlines()
                new_text = ""
                for item in text:
                    new_text = new_text + item
                test_files.append((label,new_text))


# Funció per a l'obtenció manual de tokens sense puntuacións
def get_tokens(text):
    tokens = word_tokenize(text)
    tokens = [t for t in tokens if not t in stop_words]
    return tokens


# Funció per imprimir un resum del dataset (quantitat per label)
def print_frequency_dist(docs):
    tokens = defaultdict(list)
    for doc in docs:
        doc_label = doc[0]
        doc_text = clean_text(doc[1])
        doc_tokens = get_tokens(doc_text)
        tokens[doc_label].extend(doc_tokens)

    for category_label, category_tokens in tokens.items():
        print(category_label)
        fd = FreqDist(category_tokens)
        print(fd.most_common(50))


# Neteja manual del text
def clean_text(text):
    text = text.translate(str.maketrans('','',string.punctuation))
    text = text.lower()
    return text


# Funció per avaluar el classificador amb totes les mètriques
def evaluate_classifier(title, classifier, vectorizer, X_test, y_test):
    X_test_tfidf = vectorizer.transform(X_test)
    y_pred = classifier.predict(X_test_tfidf)
    precision = metrics.precision_score(y_test,y_pred, average='micro')
    recall = metrics.recall_score(y_test,y_pred,average='micro')
    f1 = metrics.f1_score(y_test,y_pred, average='micro')
    print((title, precision, recall, f1))


# Funció per a la generació de text (mostres sintètiques)
def text_aug_gpt2(text,i):
    aug = nas.ContextualWordEmbsForSentenceAug(model_path='gpt2')
    augmented_text = aug.augment(text, num_thread=4)
    return augmented_text


# Funció per a la generació de text (mostres sintètiques)
def text_aug_rand(text,i):
    aug = nas.RandomSentAug(mode="random")
    augmented_text = aug.augment(text, num_thread=4)
    return augmented_text


# Funció per a la generació de text (mostres sintètiques)
def text_aug_pipe(text,i):
    generator = pipeline("text-generation", model = 'gpt2')
    augmented_text = generator(text, max_length = len(text), num_return_sequences=1)
    return augmented_text


# Funció per anivellar el dataset (mostres sintètiques)
def level_ds(X,y):
    max_val = max(counts.values())
    threshold = 0.01
    op = max_val - threshold * max_val
    print(op)
    counter = 0
    print("Creating new samples... ")
    for i in range(0, len(X) - 1):
        if counts[y[i]] < op:
            try:
                new_text = text_aug_rand(X[i], i)
                counts[y[i]] = counts[y[i]] + 1
                X.append(new_text[0])
                y.append(y[i])
            except ValueError or IndexError or OSError:
                try:
                    new_text = text_aug_gpt2(X[i], i)
                    counts[y[i]] = counts[y[i]] + 1
                    X.append(new_text[0])
                    y.append(y[i])
                except ValueError or IndexError or OSError:
                    pass
    print("New DS:")
    print(counts)
    return X,y


# Funció per combinar la generació de mostres sintètiques i generació del train/test set
def split_train_test(docs):
    X = []
    y = []
    for x_, y_ in docs:
        X.append(y_)
        y.append(x_)
    X,y = level_ds(X, y)
    print(len(X))
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.20, random_state=42)
    return  X_train, X_test, y_train, y_test


# Funció per a generar i entrenar el classificador de NN
def train_deep_classifier(docs):
    X = []
    y = []
    for x_, y_ in docs:
        X.append(y_)
        y.append(x_)
    X, y = level_ds(X, y)

    # Preprocess the data
    tokenizer = Tokenizer(num_words=20000)
    tokenizer.fit_on_texts(X)
    X = tokenizer.texts_to_sequences(X)
    X = tf.keras.preprocessing.sequence.pad_sequences(X, maxlen=1000)

    lb = LabelEncoder()
    y = lb.fit_transform(y)
    y = np.asarray(y)

    # Split the data into training and testing sets
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.30, random_state=42)

    # Create the model
    model = Sequential()

    # Define the layers
    model.add(Embedding(input_dim=20000, output_dim=128, input_length=1000))
    model.add(Bidirectional(LSTM(64, dropout=0.1, recurrent_dropout=0.2)))
    model.add(Dropout(0.2))
    model.add(Dense(128, activation='relu'))
    model.add(Dense(64, activation='relu'))
    model.add(Dense(6, activation='softmax'))

    # Compile the model
    optimizer = tf.keras.optimizers.Adam(learning_rate=0.01)
    model.compile(loss='sparse_categorical_crossentropy', optimizer=optimizer, metrics=['accuracy'])

    # Use callbacks
    mycall = [tf.keras.callbacks.EarlyStopping(min_delta=0.00001,
                                     patience=30,
                                     verbose=1,
                                     monitor="val_loss",
                                     restore_best_weights=True),
                                     tf.keras.callbacks.ReduceLROnPlateau(monitor="val_loss",factor=0.25,patience=30,
                                                                          verbose=1)]

    # Train the model
    model.fit(X_train, y_train, batch_size=128, epochs=60, validation_data=(X_test, y_test),callbacks=mycall)

    y_pred = np.argmax(model.predict(X_test), axis=-1)
    cm = confusion_matrix(y_test, y_pred)
    sns.set()
    sns.heatmap(cm.T, square=True, annot=True, fmt="d", xticklabels=LABELS, yticklabels=LABELS)
    plt.xlabel("true labels")
    plt.ylabel("predicted label")
    plt.show()

    mlcm = multilabel_confusion_matrix(y_test, y_pred)
    index = 0

    """
    for item in mlcm:
        sns.set()
        sns.heatmap(item.T, square=True, annot=True, fmt="d", xticklabels=LABELS[index], yticklabels=LABELS[index])
        plt.xlabel("true labels")
        plt.ylabel("predicted label")
        plt.show()
        index += 1"""

    # Evaluate the model
    loss, accuracy = model.evaluate(X_test, y_test)
    print('Test loss:', loss)
    print('Test accuracy:', accuracy)

    # Save classifier
    """print("Saving the model")
    clf_filename = 'deep_classifier.pkl'
    pickle.dump(model, open(clf_filename, 'wb'))
    """
    toc_filename = 'count_tokenizer.pk1'
    pickle.dump(tokenizer, open(toc_filename, 'wb'))
    model.save("model.h5")
    print("Saved model to disk")


# Funció per a testejar el model classificador de NN
def classify_deep(text):
    # Load model from disk
    model = load_model('model.h5')
    """clf_filename = 'deep_classifier.pkl'
    nb_clf = pickle.load(open(clf_filename, 'rb'))
    """
    toc_filename = 'count_tokenizer.pk1'
    tokenizer = pickle.load(open(toc_filename, 'rb'))

    tokenizer.fit_on_texts(text)
    text = tokenizer.texts_to_sequences(text)
    text = tf.keras.preprocessing.sequence.pad_sequences(text, maxlen=1000)

    y_pred = np.argmax(model.predict(text), axis=-1)

    return y_pred


# Funció per a testejar el model classificador de NB
def classify(text):
    clf_filename = 'naive_bayes_classifier.pkl'
    nb_clf = pickle.load(open(clf_filename, 'rb'))
    vec_filename = 'count_vectorizer.pk1'
    vectorizer = pickle.load(open(vec_filename, 'rb'))
    pred = nb_clf.predict(vectorizer.transform([text]))
    return(pred[0])


# Funció per testejar manualment els models
def test_classifier():
    create_test_set()
    print("Test set created")
    # Testing on a couple documents
    testing = 0

    y_original = []
    x = []

    for i in range(0, len(test_files)):
        y_original.append(test_files[i][0])
        x.append(test_files[i][1])

    y_predicted = classify_deep(x)
    lb = LabelEncoder()
    y_original = lb.fit_transform(y_original)
    y_original = np.asarray(y_original)
    """print(len(x))
    print(y_original)
    print(y_predicted)
    print("AAAAS")"""
    cm = confusion_matrix(y_original, y_predicted)
    sns.set()
    sns.heatmap(cm.T, square=True, annot=True, fmt="d", xticklabels=LABELS, yticklabels=LABELS)
    plt.xlabel("true labels")
    plt.ylabel("predicted label")
    plt.show()

    for i in range(0,len(x)):
        if y_original[i]==y_predicted[i]:
            testing +=1
    # Accuracy of the test files
    acc = testing / len(test_files)
    print("Manual Accuracy: ", acc)


# Funció per al càlcul manual i aleatori de la accuraccy del model
def manual_accuracy():
    r = []
    for root, dirs, files in os.walk("Files/Muestras_ALL_TestFiles/"):
        for name in files:
            r.append(os.path.join(root, name))

    original = []
    for item in r:
        if "DR" in item:
            original.append(0)
        elif "F" in item:
            original.append(1)
        elif "FD" in item:
            original.append(2)
        elif "IT" in item:
            original.append(3)
        elif "NS" in item:
            original.append(4)
        elif "TG" in item:
            original.append(5)

    predicted = []
    for item in r:
        f = open(item, "r")
        f = f.read()
        predicted.append(classify_deep([f])[0])

    print(predicted)
    print(original)
    print(len(original))
    print(len(predicted))

    manual = {0:0,1:0,2:0,3:0,4:0,5:0}
    plus = 0
    for i in range(0, len(predicted)):
        if predicted[i] == original[i]:
            plus += 1
            manual[original[i]] += 1

    """div = []
    for i,j in counts:
        div.append(j)
    for i in range(0,len(manual)):
        manual[i] = manual[i]/div[i]
    """
    print("Plus: ", plus)
    print("Accuracy: ", plus / len(original))

    print("_______________________")
    print(counts)
    print(manual)



# Funció main per desenvolupar i testejar totes les funcions explicades anteriorment
if __name__ == '__main__':

    """
    docs = []
    docs = create_data_set(docs)
    #print_frequency_dist(docs)
    print(counts)
    #train_classifier(docs)
    #train_deep_classifier(docs)
    test_classifier()
    manual_accuracy()
    """
    """
    f = open("Files/Muestras_ALL/DR/DR (1).txt", "r")
    f = f.read()
    print(classify_deep([f]))

    f = open("sample_imgs/DNI_E_example.txt", "r")
    f = f.read()
    print(classify_deep([f]))"""



