import gradio as gr
import torch
import torch.nn as nn
import torch.optim as optim

class TinyTextAI(nn.Module):
    def __init__(self, vocab_size, num_classes):
        super(TinyTextAI, self).__init__()
        self.embedding = nn.EmbeddingBag(vocab_size, 16, sparse=False)
        self.fc1 = nn.Linear(16, 12)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(12, num_classes)

    def forward(self, text, offsets):
        embedded = self.embedding(text, offsets)
        x = self.relu(self.fc1(embedded))
        return self.fc2(x)

data = {
    "hello": 0, "hi": 0, "hey": 0,
    "how are you": 1, "status": 1, "up": 1,
    "name": 2, "who": 2, "identity": 2,
    "bye": 3, "goodbye": 3, "exit": 3
}

responses = {
    0: "Hello! I am feeling smarter now.",
    1: "Systems are nominal. My parameters are tuned!",
    2: "I am a Version 2 Tiny AI.",
    3: "Goodbye! Come back to train me more soon."
}

vocab = {word: i for i, word in enumerate(set(" ".join(data.keys()).split()))}
vocab["<UNK>"] = len(vocab)
vocab_size = len(vocab)
num_classes = len(responses)

model = TinyTextAI(vocab_size, num_classes)
optimizer = optim.Adam(model.parameters(), lr=0.05)
criterion = nn.CrossEntropyLoss()

def prepare_text(text):
    tokens = [vocab.get(w, vocab["<UNK>"]) for w in text.lower().split()]
    if not tokens: return torch.tensor([vocab["<UNK>"]]), torch.tensor([0])
    return torch.tensor(tokens, dtype=torch.int64), torch.tensor([0], dtype=torch.int64)

def train_model(epochs):
    model.train()
    log = []
    for epoch in range(int(epochs)):
        total_loss = 0
        for text, label in data.items():
            input_tensor, offsets = prepare_text(text)
            target = torch.tensor([label], dtype=torch.int64)
            optimizer.zero_grad()
            output = model(input_tensor, offsets)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
        if epoch % 20 == 0:
            log.append(f"Epoch {epoch} - Error: {total_loss:.4f}")
    return "\n".join(log)

def chat(user_input):
    model.eval()
    input_tensor, offsets = prepare_text(user_input)
    with torch.no_grad():
        output = model(input_tensor, offsets)
        prediction = torch.argmax(output, dim=1).item()
    return responses[prediction]

with gr.Blocks() as demo:
    gr.Markdown("# 🚀 Tiny AI v2: The Hidden Layer")
    with gr.Row():
        epochs = gr.Number(label="Training Rounds", value=200)
        btn_train = gr.Button("Re-Train AI")
    status = gr.Textbox(label="Neural Progress")
    chat_input = gr.Textbox(label="Talk to the AI")
    btn_chat = gr.Button("Send")
    chat_output = gr.Textbox(label="Response")
    
    btn_train.click(train_model, inputs=epochs, outputs=status)
    btn_chat.click(chat, inputs=chat_input, outputs=chat_output)

demo.launch()