import nltk
nltk.download('punkt')
nltk.download('wordnet')
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 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, training, train_x, train_y, sgd, hist
    
    lemmatizer = WordNetLemmatizer()
    words = []
    classes = []
    documents = []
    ignore_words = ['?', '!']

    with open('intents.json') as file:
        intents = json.load(file)

    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)

        # Save data and model locally
        with open('words.pkl', 'wb') as f:
            pickle.dump(words, f)
        with open('classes.pkl', 'wb') as f:
            pickle.dump(classes, f)
        with open('training.pkl', 'wb') as f:  # Save training data
            pickle.dump(training, f)
        model.save('chatbot_model.h5') 

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

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