import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import streamlit.components.v1 as components


# Set Streamlit page config for a clinical look
st.set_page_config(
    page_title="Symptom2Disease Clinical AI",
    page_icon="🩺",
    layout="centered"
)

# Clinical-themed header
st.title("🩺 Symptom2Disease Clinical AI")
st.markdown("""
Welcome to the **Symptom2Disease** clinical assistant.  
Enter your symptoms below and our AI will suggest possible diseases.
""")

# Load model and tokenizer
@st.cache_resource
def load_model():
    tokenizer = AutoTokenizer.from_pretrained("alibidaran/Symptom2disease")
    model = AutoModelForSequenceClassification.from_pretrained("alibidaran/Symptom2disease")
    return tokenizer, model

tokenizer, model = load_model()

label_2id = {
    'Psoriasis': 0, 'Varicose Veins': 1, 'Typhoid': 2, 'Chicken pox': 3, 'Impetigo': 4,
    'Dengue': 5, 'Fungal infection': 6, 'Common Cold': 7, 'Pneumonia': 8, 'Dimorphic Hemorrhoids': 9,
    'Arthritis': 10, 'Acne': 11, 'Bronchial Asthma': 12, 'Hypertension': 13, 'Migraine': 14,
    'Cervical spondylosis': 15, 'Jaundice': 16, 'Malaria': 17, 'urinary tract infection': 18,
    'allergy': 19, 'gastroesophageal reflux disease': 20, 'drug reaction': 21, 'peptic ulcer disease': 22,
    'diabetes': 23
}
id2label = {i: v for v, i in label_2id.items()}

def detect_symptom(symptoms):
    inputs_id = tokenizer(symptoms, padding=True, truncation=True, return_tensors="pt")
    output = model(inputs_id['input_ids'])
    preds = torch.nn.functional.softmax(output.logits, -1).topk(5)
    results = {id2label[preds.indices[0][i].item()]: float(preds.values[0][i].item()) for i in range(5)}
    return results

# Example symptoms for quick testing
examples = [
    "I can't stop sneezing and I feel really tired and crummy. My throat is really sore",
    "I have been experiencing a severe headache that is accompanied by pain behind my eyes.",
    "There are small red spots all over my body that I can't explain. It's worrying me.",
    "I've been having a really hard time going to the bathroom lately. It's really painful"
]

with st.expander("📝 Example symptom descriptions"):
    for ex in examples:
        st.code(ex)

# User input
symptoms = st.text_area(
    "Describe your symptoms:",
    placeholder="e.g. I have a persistent cough and mild fever for the last 3 days.",
    height=100
)

if st.button("Analyze Symptoms", type="primary"):
    if symptoms.strip():
        with st.spinner("Analyzing symptoms..."):
            results = detect_symptom(symptoms)
        st.subheader("Top 5 Predicted Diseases")
        for disease, prob in results.items():
            st.markdown(f"**{disease}**")
            st.progress(prob)
            st.caption(f"Probability: {prob:.2%}")
    else:
        st.warning("Please enter your symptoms above.")

st.markdown("---")
st.info("**Disclaimer:** This tool is for informational purposes only and does not provide medical advice. Please consult a healthcare professional for diagnosis and treatment.")
components.iframe(
    src="http://127.0.0.1:8080",
    height=200)