import nltk
nltk.download('punkt')
nltk.download('wordnet')
nltk.download('omw-1.4') # This is important for WordNetLemmatizer
from nltk.stem import WordNetLemmatizer

import json
import pickle

import numpy as np
from tensorflow.keras import Sequential
from tensorflow.keras.layers import Dense, Dropout
from tensorflow.keras.optimizers import SGD

import redis
import os
from dotenv import load_dotenv
import uvicorn
from fastapi import FastAPI

load_dotenv()

app = FastAPI()

@app.on_event("startup")
async def startup_event():
    global lemmatizer, words, classes, documents, ignore_words, intents, model, r, training, train_x, train_y, sgd, hist
    
    lemmatizer = WordNetLemmatizer()
    words = []
    classes = []
    documents = []
    ignore_words = ['?', '!']

    with open('intents.json') as file:
        intents = json.load(file)
    
    redis_password = os.getenv("REDIS_PASSWORD")
    r = redis.Redis(host=os.getenv("REDIS_HOST"), port=int(os.getenv("REDIS_PORT")), db=int(os.getenv("REDIS_DB")), password=redis_password)

    while True:
        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_words]

        training = []
        output_empty = [0] * len(classes)
        for doc in documents:
            bag = []
            pattern_words = doc[0]
            pattern_words = [lemmatizer.lemmatize(word.lower()) for word in pattern_words]
            for w in words:
                bag.append(1) if w in pattern_words else bag.append(0)

            output_row = list(output_empty)
            output_row[classes.index(doc[1])] = 1

            training.append([bag, output_row])

        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=300, batch_size=5, verbose=1)

        r.set('words', pickle.dumps(words))
        r.set('classes', pickle.dumps(classes))
        r.set('train_x', pickle.dumps(train_x))
        r.set('train_y', pickle.dumps(train_y))
        r.set('model', pickle.dumps(model))

        print("Data and model saved to Redis. Re-training...")



if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)