# ✅ Install dependencies first
# pip install transformers accelerate gradio translate

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import gradio as gr
from translate import Translator

# Load IBM Granite model
model_id = "ibm-granite/granite-3.3-2b-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
generator = pipeline("text-generation", model=model, tokenizer=tokenizer)

# Translate helpers
def translate_to_english(text, lang_code):
    if lang_code != "en":
        translator = Translator(from_lang=lang_code, to_lang="en")
        return translator.translate(text)
    return text

def translate_to_user_lang(text, lang_code):
    if lang_code != "en":
        translator = Translator(from_lang="en", to_lang=lang_code)
        return translator.translate(text)
    return text

# Core functions
def identify_disease(symptoms, lang_code="en"):
    symptoms_en = translate_to_english(symptoms, lang_code)
    prompt = f"You are a medical assistant. A user reports these symptoms: {symptoms_en}. What possible disease or condition could this indicate?"
    output = generator(prompt, max_new_tokens=150, do_sample=True)[0]["generated_text"]
    result = output[len(prompt):].strip()
    return translate_to_user_lang(result, lang_code)

def suggest_remedy(disease, lang_code="en"):
    disease_en = translate_to_english(disease, lang_code)
    prompt = f"Suggest effective natural home remedies for treating {disease_en}."
    output = generator(prompt, max_new_tokens=150, do_sample=True)[0]["generated_text"]
    result = output[len(prompt):].strip()
    return translate_to_user_lang(result, lang_code)

def preventive_measures(disease, lang_code="en"):
    disease_en = translate_to_english(disease, lang_code)
    prompt = f"What are the best preventive measures to avoid {disease_en}?"
    output = generator(prompt, max_new_tokens=150, do_sample=True)[0]["generated_text"]
    result = output[len(prompt):].strip()
    return translate_to_user_lang(result, lang_code)

def diet_recommendations(disease, lang_code="en"):
    disease_en = translate_to_english(disease, lang_code)
    prompt = f"Suggest a healthy diet plan for someone suffering from {disease_en}."
    output = generator(prompt, max_new_tokens=150, do_sample=True)[0]["generated_text"]
    result = output[len(prompt):].strip()
    return translate_to_user_lang(result, lang_code)

def first_aid_advice(condition, lang_code="en"):
    condition_en = translate_to_english(condition, lang_code)
    prompt = f"What first aid steps should be taken immediately for {condition_en}?"
    output = generator(prompt, max_new_tokens=150, do_sample=True)[0]["generated_text"]
    result = output[len(prompt):].strip()
    return translate_to_user_lang(result, lang_code)

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("## 🤖 HealthAI - Your Intelligent Healthcare Assistant")

    with gr.Tab("🩺 Symptoms Identifier"):
        symptoms_input = gr.Textbox(label="Enter your symptoms")
        disease_output = gr.Textbox(label="Predicted Disease")
        btn1 = gr.Button("Identify Disease")
        btn1.click(fn=identify_disease, inputs=symptoms_input, outputs=disease_output)

    with gr.Tab("🌿 Home Remedies"):
        disease_input = gr.Textbox(label="Enter disease name")
        remedy_output = gr.Textbox(label="Suggested Home Remedy")
        btn2 = gr.Button("Get Remedy")
        btn2.click(fn=suggest_remedy, inputs=disease_input, outputs=remedy_output)

    with gr.Tab("🛡️ Preventive Measures"):
        prevent_input = gr.Textbox(label="Enter disease name")
        prevent_output = gr.Textbox(label="Preventive Measures")
        btn3 = gr.Button("Get Advice")
        btn3.click(fn=preventive_measures, inputs=prevent_input, outputs=prevent_output)

    with gr.Tab("🥗 Diet Recommendations"):
        diet_input = gr.Textbox(label="Enter disease name")
        diet_output = gr.Textbox(label="Recommended Diet")
        btn4 = gr.Button("Get Diet Plan")
        btn4.click(fn=diet_recommendations, inputs=diet_input, outputs=diet_output)

    with gr.Tab("🆘 First Aid Help"):
        first_aid_input = gr.Textbox(label="Enter medical condition or emergency")
        first_aid_output = gr.Textbox(label="First Aid Advice")
        btn5 = gr.Button("Get First Aid")
        btn5.click(fn=first_aid_advice, inputs=first_aid_input, outputs=first_aid_output)

demo.launch(ssr_mode=False)