import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np

# Modelo
class SimpleChat(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(SimpleChat, self).__init__()

        # Camada de codificador RNN
        self.encoder = nn.GRU(input_size, hidden_size)

        # Camada de decodificador RNN
        self.decoder = nn.GRU(hidden_size, hidden_size)

        # Mapeamento linear para a saída
        self.linear = nn.Linear(hidden_size, output_size)

    def forward(self, input_seq, hidden=None):
        # Codificador RNN
        encoder_output, encoder_hidden = self.encoder(input_seq, hidden)

        # Decodificador RNN
        decoder_output, _ = self.decoder(encoder_output)

        # Mapeamento linear da saída
        output = self.linear(decoder_output)

        return output

# Treinamento do modelo
def train(model, input_data, target_data, criterion, optimizer, epochs=1000):
    for epoch in range(epochs):
        optimizer.zero_grad()
        output = model(input_data)
        loss = criterion(output, target_data)
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 100 == 0:
            print(f'Epoch [{epoch + 1}/{epochs}], Loss: {loss.item():.4f}')

# Chat em tempo real
def chatbot(model, input_size):
    model.eval()
    hidden = None

    # Mapeamento de respostas em texto
    text_responses = {
        0: "Me chamo Theo",
        1: "Como posso ajudar?",
        2: "Desculpe, não consigo responder a isso.",
        3: "Isso parece interessante!",
        4: "Não entendi o que você disse.",
        5: "Qual é o seu endereço?",
        6: "Qual o seu nome completo?",
        # Adicione mais respostas 
    }

    while True:
        user_input = input("\n=> ")
        if user_input.lower().strip() == "sair":
            print("__Encerrando Chat__")
            break

        # Convertendo entrada do usuário em um tensor
        input_tensor = torch.rand((len(user_input), 1, input_size))
        output = model(input_tensor, hidden)

        # Arredondando as saídas para inteiros
        rounded_output = np.round(output.squeeze().detach().numpy()).astype(int)

        # Obtendo resposta em texto
        text_response = text_responses.get(rounded_output[0], "Desculpe, não sei o que dizer.")

        print(f'Chatbot: {text_response}')

# Uso
input_size = 10
hidden_size = 20
output_size = 1

# Instâncias
model = SimpleChat(input_size, hidden_size, output_size)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Dados de entrada e de saída de exemplo
input_data = torch.rand(5, 1, input_size)
target_data = torch.rand(5, 1, output_size)

train(model, input_data, target_data, criterion, optimizer)

print('Fim do treinamento!')

# Iniciar interação em tempo real
chatbot(model, input_size)