import os
import json
import re
import numpy as np
import nltk
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
import torch
from torch.nn.functional import cosine_similarity
from sentence_transformers import SentenceTransformer
import gradio as gr

# Download necessary NLTK data
nltk.download('punkt_tab')
nltk.download('stopwords')
nltk.download('wordnet')

# Load dataset
dataset_path = "dataset.json"
if not os.path.exists(dataset_path):
    raise FileNotFoundError(f"{dataset_path} not found. Please ensure the file is in the current directory.")

with open(dataset_path, "r") as f:
    data = json.load(f)

disease_data = {d["name"]: d for d in data["diseases"]}

# Preprocessing
lemmatizer = WordNetLemmatizer()
stop_words = set(stopwords.words('english'))

medical_synonyms = {
    "itchy": "itch", "pruritic": "itch", "rash": "eruption",
    "redness": "red", "swollen": "swelling", "bumps": "lesion"
}


def advanced_preprocess_text(text):
    text = text.lower()
    text = re.sub(r'[^a-zA-Z\s]', ' ', text)
    tokens = word_tokenize(text)
    tokens = [medical_synonyms.get(token, token) for token in tokens]
    tokens = [lemmatizer.lemmatize(token) for token in tokens if token not in stop_words]
    return " ".join(tokens)


# Prepare disease descriptions
disease_names = []
disease_descriptions = []

for disease_name, details in disease_data.items():
    disease_names.append(disease_name)
    description = ""
    if "symptoms" in details:
        description += "Symptoms: " + " ".join(details["symptoms"]) + ". "
    if "causes" in details:
        description += "Causes: " + " ".join(details["causes"]) + ". "
    if "cure" in details:
        description += "Cure: " + " ".join(details["cure"]) + ". "
    if "home_remedies" in details:
        description += "Home Remedies: " + " ".join(details["home_remedies"]) + ". "
    processed_desc = advanced_preprocess_text(description)
    disease_descriptions.append(processed_desc)

# Load BERT model
bert_model = SentenceTransformer('all-MiniLM-L6-v2')
disease_embeddings = bert_model.encode(disease_descriptions, convert_to_tensor=True)


def gradio_chatbot(primary_symptom, location, associated_symptoms, duration, severity, additional_info):
    user_input = " ".join([
        primary_symptom, location, associated_symptoms, duration, severity, additional_info
    ])
    processed_input = advanced_preprocess_text(user_input)
    user_embedding = bert_model.encode(processed_input, convert_to_tensor=True)
    similarities = cosine_similarity(user_embedding, disease_embeddings)
    best_match_idx = torch.argmax(similarities).item()
    predicted_disease = disease_names[best_match_idx]
    details = disease_data.get(predicted_disease, {})

    causes = ", ".join(details.get("causes", ["Not available"]))
    cure = ", ".join(details.get("cure", ["Not available"]))
    home_remedies = ", ".join(details.get("home_remedies", ["Not available"]))

    response = (f"🩺 **Predicted Disease:** {predicted_disease}\n\n"
                f"**Causes:** {causes}\n\n"
                f"**Cure:** {cure}\n\n"
                f"**Home Remedies:** {home_remedies}")
    return response


iface = gr.Interface(
    fn=gradio_chatbot,
    inputs=[
        gr.Textbox(label="1. Primary Symptom", placeholder="e.g., itching, rash"),
        gr.Textbox(label="2. Location on Body", placeholder="e.g., face, arms"),
        gr.Textbox(label="3. Associated Symptoms", placeholder="e.g., swelling, redness"),
        gr.Textbox(label="4. Duration", placeholder="e.g., 2 days, 1 week"),
        gr.Textbox(label="5. Severity (scale 1-10)", placeholder="e.g., 5"),
        gr.Textbox(label="6. Additional Observations", placeholder="e.g., recent exposure, diet changes")
    ],
    outputs="markdown",
    title="Advanced Skin Disease AI Chatbot",
    description="Answer a few questions about your symptoms and get a predicted skin disease along with recommended solutions."
)

iface.launch()