import torch as tr
import numpy as np

from Model.model import WordModel

import string
import json
import random
import re

import nltk
from nltk.stem import PorterStemmer
nltk.download('punkt')

with open('./Model/intents.json', 'r', encoding='UTF-8') as fp:
    intents = json.load(fp)

CONFIG = tr.load('./Model/save.pth')


def load_chat_model():
    model = WordModel(IN_DIMS=CONFIG['input_size'], HIDDEN_DIMS=CONFIG['hidden_size'], VOCAB_SIZE=CONFIG['output_size'])
    model.load_state_dict(CONFIG['model_state'])
    model.eval()
    return model


class Agent:

    MODEL_NAME = 'Garmento Agent'
    stemmer = PorterStemmer()
    model = load_chat_model()


    def preprocess_text(self, text:str)->list[str]:
        text = text.lower()
        text = re.sub(pattern=f'[{string.punctuation}]', repl='', string=text)
        text = nltk.tokenize.word_tokenize(text)
        return text

    def stemming(self, world_list:list[str])->list[str]:
        return list(map(self.stemmer.stem, world_list))

    def bag_of_words(self, tokens:list[str], vocab:list[str])->list[str]:
        tokens = self.stemming(tokens)
        bow = np.zeros(shape=(len(vocab)), dtype=np.float32)
        for idx, token in enumerate(vocab):
            if token in tokens:
                bow[idx] = 1
        return bow
    
    def agent_response(self, user_input:str)->str:
        tokens = self.preprocess_text(user_input)
        vector = self.bag_of_words(tokens=tokens, vocab=CONFIG['VOCAB'])
        vector = tr.from_numpy(vector[None, :])
        logits = self.model(vector)
        tag_id = tr.argmax(logits, dim=-1)
        tag = CONFIG['tags'][tag_id.item()]
        confidence = tr.softmax(logits, dim=-1)

        for intent in intents['intents']:
            if intent['tag'] == tag:
                response = f"{random.choice(intent['responses'])}"
                if confidence.max() < 0.7:
                    response += ". <br> I am an AI-agent trained to represent my organisation. If my response is not relevant to your query. Kindly connect to a our customer support team. ☎️ Contact information : +91-999-9999-999  "
        
        return response, confidence.max()