#imports
import random
import json
import pickle
import numpy as np

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Activation, Flatten
from tensorflow.keras.optimizers import SGD

import nltk
# nltk.download('punkt')
# nltk.download('wordnet')
from nltk.stem import WordNetLemmatizer

lemmatizer = WordNetLemmatizer()

intents = json.loads(open('intents.json').read())

words = []
classes = []
documents = []
ignore_chars = ['?', '!', '.', ',']

for intent in intents['intents']:
    for pattern in intent['patterns']:
        w = nltk.word_tokenize(pattern)
        words.extend(w)
        documents.append((w, intent['tag']))
        if intent['tag'] not in classes:
            classes.append(intent['tag'])

words = [lemmatizer.lemmatize(w.lower()) for w in words if w not in ignore_chars]
words = sorted(list(set(words)))

classes = sorted(list(set(classes)))

#saving the 2 words files
pickle.dump(words, open('words.pkl', 'wb'))
pickle.dump(classes, open('classes.pkl', 'wb'))

training = []
output = [0] * len(classes)

# ["hi", "hello", "see", "my", "you"]

for document in documents:
    # (["See", "You"], "greeting")

    bag = [] # [0, 0, 1, 0, 1]
    pattern_words = document[0]
    pattern_words = [lemmatizer.lemmatize(w.lower()) for w in pattern_words]

    for w in words:
        bag.append(1) if w in pattern_words else bag.append(0)
    
    output_row = list(output) # [0, 0` 0, 0, 0, 0, 0, 0]
    output_row[classes.index(document[1])] = 1 # [0, 0` 0, 0, 0, 0, 0, 0]
    training.append([bag, output_row])
'''
[
    [0, 1, 0, 0, 1], [0, 1, 0, 0, 0]
    [1, 0, 0, 0, 0], []
]

'''


random.shuffle(training)
training = np.array(training)

train_X = list(training[:,0])
train_Y = list(training[:,1])

model = Sequential()
model.add(Dense(128, input_shape=(len(train_X[0]),), activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(len(train_Y[0]), activation='softmax'))

sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)

model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])

hist = model.fit(np.array(train_X), np.array(train_Y), epochs=200, batch_size=5, verbose=1)
model.save('chatbot_model.h5', hist)
print("Training completed!")