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

# ---------------------
# 1. Download NLTK Data
# ---------------------
nltk.download('punkt')
nltk.download('stopwords')
nltk.download('wordnet')

# ---------------------
# 2. 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.json, "r") as f:
    data = json.load(f)

# Convert dataset into a dictionary: disease name -> details
disease_data = {d["name"]: d for d in data["diseases"]}
print("Sample Disease (Acne):")
print(json.dumps(disease_data.get("Acne", {}), indent=2))

# ---------------------
# 3. Advanced Preprocessing Functions
# ---------------------
lemmatizer = WordNetLemmatizer()
stop_words = set(stopwords.words('english'))

# Mapping for common medical synonyms to standardize terminology
medical_synonyms = {
    "itchy": "itch",
    "pruritic": "itch",
    "rash": "eruption",
    "redness": "red",
    "swollen": "swelling",
    "bumps": "lesion"
}

def advanced_preprocess_text(text):
    # Lowercase and remove unwanted characters (retain only letters and spaces)
    text = text.lower()
    text = re.sub(r'[^a-zA-Z\s]', ' ', text)
    tokens = word_tokenize(text)
    # Replace tokens using the medical synonym mapping
    tokens = [medical_synonyms.get(token, token) for token in tokens]
    # Remove stopwords and lemmatize
    tokens = [lemmatizer.lemmatize(token) for token in tokens if token not in stop_words]
    return " ".join(tokens)

# ---------------------
# 4. Prepare Disease Descriptions & Embeddings Using BERT
# ---------------------
# For each disease, combine symptoms, causes, cure, and home remedies into a detailed description.
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)

# Initialize SentenceTransformer (BERT-based model)
bert_model = SentenceTransformer('all-MiniLM-L6-v2')
# Compute embeddings for each disease description (as torch tensors)
disease_embeddings = bert_model.encode(disease_descriptions, convert_to_tensor=True)
print("Computed disease embeddings for", len(disease_embeddings), "diseases.")

# ---------------------
# 5. Detailed User Input Collection
# ---------------------
def get_detailed_user_input(primary_symptom, location, associated_symptoms, duration, severity, additional_info):
    # Combine all responses into one detailed description.
    combined_input = " ".join([primary_symptom, location, associated_symptoms, duration, severity, additional_info])
    return combined_input

# ---------------------
# 6. Chatbot Prediction Function Using BERT Similarity
# ---------------------
def chatbot_response(user_input):
    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]
    return predicted_disease

# ---------------------
# 7. Gradio Chatbot Function
# ---------------------
def gradio_chatbot(primary_symptom, location, associated_symptoms, duration, severity, additional_info):
    detailed_input = get_detailed_user_input(primary_symptom, location, associated_symptoms, duration, severity, additional_info)
    predicted = chatbot_response(detailed_input)
    details = disease_data.get(predicted, {})
    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}\n\n"
                f"**Causes:** {causes}\n\n"
                f"**Cure:** {cure}\n\n"
                f"**Home Remedies:** {home_remedies}")
    return response

# ---------------------
# 8. Gradio Interface Setup
# ---------------------
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."
)

# ---------------------
# 9. Launch the Gradio App on Hugging Face Spaces
# ---------------------
iface.launch()