import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# Load the IBM Granite instruct model from Hugging Face
MODEL_NAME = "ibm-granite/granite-3.3-2b-instruct"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
    device_map="auto"
)

def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        do_sample=True,
        top_p=0.9,
        temperature=0.7,
        pad_token_id=tokenizer.eos_token_id,
    )
    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return text[len(prompt):].strip()

def symptoms_identifier(symptoms):
    if not symptoms.strip():
        return "Please enter symptoms."
    prompt = f"Based on the following symptoms, identify the most likely disease:\nSymptoms: {symptoms}\nDisease:"
    return generate_response(prompt)

def home_remedies(disease):
    if not disease.strip():
        return "Please enter a disease."
    prompt = f"Suggest a natural, easy home remedy for the disease:\nDisease: {disease}\nHome Remedy:"
    return generate_response(prompt)

# Custom CSS styling
custom_css = """
body {
    background-color: #ffffff;
    color: #6b7280;
    font-family: 'Poppins', sans-serif;
    margin: 0; padding: 0;
}
h1 {
    font-weight: 700;
    font-size: 3rem;
    color: #111827;
    margin-bottom: 0.5rem;
}
h2 {
    font-weight: 600;
    font-size: 1.5rem;
    color: #111827;
    margin-bottom: 1rem;
}
.card {
    background: white;
    border-radius: 0.75rem;
    box-shadow: rgba(203, 213, 224, 0.5) 0px 4px 6px -1px;
    padding: 2rem;
    margin: 1rem;
    box-sizing: border-box;
}
.gr-button {
    background-color: #111827;
    color: white;
    border-radius: 0.5rem;
    padding: 0.75rem 1.5rem;
    font-weight: 600;
    font-size: 1rem;
    transition: all 0.3s ease;
    border: none;
    cursor: pointer;
}
.gr-button:hover {
    background-color: #374151;
    transform: scale(1.05);
}
.gr-textbox {
    border-radius: 0.5rem;
    border: 1px solid #d1d5db;
    padding: 0.75rem 1rem;
    font-size: 1rem;
    color: #111827;
    font-family: 'Poppins', sans-serif;
    transition: border-color 0.3s ease;
}
.gr-textbox:focus {
    border-color: #111827;
    outline: none;
    box-shadow: 0 0 6px rgb(17 24 39 / 0.3);
}
@media (min-width: 768px) {
    .flex-row {
        display: flex;
        justify-content: center;
        max-width: 1200px;
        margin: 0 auto;
        gap: 2rem;
    }
    .flex-column {
        flex: 1;
        min-width: 0; /* For Firefox */
    }
}
@media (max-width: 767px) {
    .flex-column {
        margin: 1rem 0;
    }
}
"""

with gr.Blocks(css=custom_css) as demo:
    gr.Markdown("<h1 style='text-align:center;'>HealthAI</h1><p style='text-align:center; font-size:1.25rem; color:#4b5563; max-width:700px; margin:auto;'>Generative AI-powered health assistant for symptom identification and natural home remedies.</p>")
    with gr.Row(variant="panel", elem_classes="flex-row"):
        with gr.Column(elem_classes="flex-column"):
            gr.Markdown("## Symptoms Identifier")
            symptoms_input = gr.Textbox(label="Enter Symptoms", placeholder="e.g. fever, headache, fatigue", lines=4)
            symptoms_output = gr.Textbox(label="Predicted Disease", interactive=False, lines=2)
            symptoms_button = gr.Button("Identify Disease")
            symptoms_button.click(symptoms_identifier, inputs=symptoms_input, outputs=symptoms_output)
        with gr.Column(elem_classes="flex-column"):
            gr.Markdown("## Home Remedies")
            disease_input = gr.Textbox(label="Enter Disease", placeholder="e.g. common cold", lines=2)
            remedy_output = gr.Textbox(label="Recommended Home Remedy", interactive=False, lines=4)
            remedy_button = gr.Button("Get Home Remedy")
            remedy_button.click(home_remedies, inputs=disease_input, outputs=remedy_output)
    gr.Markdown("<p style='text-align:center; margin-top:3rem; color:#9ca3af;'>Powered by <a href='https://huggingface.co/ibm-granite/granite-3.3-2b-instruct' target='_blank' rel='noopener noreferrer'>ibm-granite/granite-3.3-2b-instruct</a></p>")

if __name__ == "__main__":
    demo.launch()