import gradio as gr
import torch
import json
import random
import re
import numpy as np
from huggingface_hub import hf_hub_download
from sklearn.preprocessing import LabelEncoder

print("🤖 Loading Asad AI...")

# Download model from your Hub repo
model_path = hf_hub_download("Asad-ullah008/asad-ai", "asad_ai_best.pth")
info_path = hf_hub_download("Asad-ullah008/asad-ai", "model_info.json")
data_path = hf_hub_download("Asad-ullah008/asad-ai", "training_data.json")

with open(info_path, 'r') as f:
    info = json.load(f)
with open(data_path, 'r') as f:
    training_data = json.load(f)

class AsadAIModel(torch.nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super().__init__()
        self.network = torch.nn.Sequential(
            torch.nn.Linear(input_size, hidden_size),
            torch.nn.BatchNorm1d(hidden_size),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(hidden_size, hidden_size // 2),
            torch.nn.BatchNorm1d(hidden_size // 2),
            torch.nn.ReLU(),
            torch.nn.Dropout(0.2),
            torch.nn.Linear(hidden_size // 2, output_size)
        )
    def forward(self, x):
        return self.network(x)

model = AsadAIModel(info['input_size'], info['hidden_size'], info['output_size'])
model.load_state_dict(torch.load(model_path, map_location='cpu'))
model.eval()

le = LabelEncoder()
le.classes_ = np.array(info['tags'])
vocab = info['vocab']

def clean_text(text):
    text = text.lower().strip()
    text = re.sub(r'[^\w\s]', '', text)
    return text

def text_to_bow(text, vocab):
    words = clean_text(text).split()
    bow = np.zeros(len(vocab), dtype=np.float32)
    for word in words:
        if word in vocab:
            bow[vocab.index(word)] = 1.0
    return bow

def respond(message, history):
    bow = text_to_bow(message, vocab)
    input_tensor = torch.FloatTensor(bow).unsqueeze(0)
    with torch.no_grad():
        output = model(input_tensor)
        probs = torch.softmax(output, dim=1)
        conf, pred = torch.max(probs, 1)
    tag = le.inverse_transform(pred.numpy())[0]
    for intent in training_data['intents']:
        if intent['tag'] == tag:
            return random.choice(intent['responses'])
    return "Samajh nahi aaya"

# Fixed: No 'theme' parameter
gr.ChatInterface(
    fn=respond,
    title="🤖 Asad AI Chatbot",
    description="Urdu/Hindi mein baat karein! Main aapka personal assistant hoon."
).launch()