import streamlit as st
from transformers import pipeline

# Load your trained model and tokenizer
model_path = "quadranttechnologies/Clinical_Decision_Support"  # Update with your model path
classifier = pipeline("text-classification", model=model_path)

# Define the Streamlit app
st.title("Clinical Decision Support System")
st.write("Provide patient details to get a medical recommendation.")

# Input fields for user
age = st.number_input("Age", min_value=0, max_value=120, step=1, value=50)
gender = st.selectbox("Gender", options=["Male", "Female"])
weight = st.number_input("Weight (kg)", min_value=0, max_value=300, step=1, value=70)
smoking_status = st.selectbox("Smoking Status", options=["Never", "Former", "Current"])
diabetes = st.selectbox("Diabetes", options=["No", "Yes"])
hypertension = st.selectbox("Hypertension", options=["No", "Yes"])
cholesterol = st.number_input("Cholesterol (mg/dL)", min_value=0, max_value=500, step=1, value=200)
heart_disease_history = st.selectbox("Heart Disease History", options=["No", "Yes"])
symptoms = st.text_input("Symptoms", value="Chest pain")
risk_score = st.number_input("Risk Score", min_value=0.0, max_value=10.0, step=0.1, value=5.0)

# Button to get recommendation
if st.button("Get Recommendation"):
    # Convert inputs to model format
    input_text = (
        f"Age: {age}, Gender: {gender}, Weight: {weight}, Smoking Status: {smoking_status}, "
        f"Diabetes: {1 if diabetes == 'Yes' else 0}, Hypertension: {1 if hypertension == 'Yes' else 0}, "
        f"Cholesterol: {cholesterol}, Heart Disease History: {1 if heart_disease_history == 'Yes' else 0}, "
        f"Symptoms: {symptoms}, Risk Score: {risk_score}"
    )

    # Get prediction from the model
    prediction = classifier(input_text)
    recommendation_label = prediction[0]['label']
    
    # Map label to recommendation title (ensure you have reverse mapping loaded)
    reverse_label_mapping = {
        "LABEL_0": "Maintain healthy lifestyle",
        "LABEL_1": "Immediate cardiologist consultation",
        "LABEL_2": "Start statins, monitor regularly",
        "LABEL_3": "Lifestyle changes, monitor",
        "LABEL_4": "No immediate action",
        "LABEL_5": "Increase statins, lifestyle changes",
        "LABEL_6": "Start ACE inhibitors, monitor"
    }
    recommendation = reverse_label_mapping.get(recommendation_label, "Unknown Recommendation")
    
    # Display the recommendation
    st.subheader("Recommendation")
    st.write(recommendation)