import json
from nltk_processing import tokenization,temming,bag_of_word
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader
from model import NeuralNet

with open("data.json","r") as f:
    data=json.load(f)
all_words=[]
tags=[]
xy=[]
for intent in data["intents"]:
    tag=intent["tag"]
    tags.append(tag)
    for pattern in intent["patterns"]:
        w=tokenization(pattern)
        all_words.extend(w)# this use extend because w is array
        xy.append((w,tag))
ignore_words=['?',',','.','!']
all_words=[temming(w) for w in all_words if w not in ignore_words]
all_words=sorted(set(all_words))
tags=sorted(set(tags))

X_train=[]
Y_train=[]
for (pattern_sentence,tag) in xy:
    bag=bag_of_word(pattern_sentence,all_words)
    X_train.append(bag)
    lable=tags.index(tag)#get index of tag in tags array
    Y_train.append(lable)
    
X_train=np.array(X_train)
Y_train=np.array(Y_train)
'''print(X_train)
print(Y_train)'''
class ChatDataset(Dataset):
    def __init__(self):
        self.n_samples=len(X_train)
        self.x_data=X_train
        self.y_data=Y_train
    def __getitem__(self,index):
        return self.x_data[index],self.y_data[index]
    def __len__(self):
        return self.n_samples
batch_size=8
hidden_size=8
output_size=len(tags)
input_size=len(X_train[0])
learning_rate=0.001
num_epochs=1000



dataset=ChatDataset()
train_loader=DataLoader(dataset=dataset,batch_size=batch_size,shuffle=True,num_workers=0)

device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model= NeuralNet(input_size, hidden_size, output_size).to(device)

# loss and optimizer
criterion=nn.CrossEntropyLoss()
optimizer=torch.optim.Adam(model.parameters(), lr=learning_rate)

for epoch in range(num_epochs):
    for words,labels in train_loader:
        words=words.to(device)
        labels=labels.to(dtype=torch.long).to(device)
        #forward
        outputs=model(words)
        loss = criterion(outputs,labels)
        #backward and optimizer step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if (epoch+1)%100==0:
        print (f'epoch {epoch+1}/num_epochs, loss={loss.item():.4f}')
print (f'epoch final, loss={loss.item():.4f}')
data={
"model_state": model.state_dict(),
"input_size": input_size,
"hidden_size": hidden_size,
"output_size": output_size,
"all_words": all_words,
"tags": tags
}
FILE = "data.pth"
torch.save(data,FILE)
print(f'Training complete, file save to {FILE}')
