import streamlit as st
import pandas as pd
import joblib
from PIL import Image

# Custom styling for the app
st.markdown(
    """
    <style>
    .main {
        background-color: #f4f4f9;
    }
    .stButton > button {
        background-color: #4CAF50;
        color: white;
        border-radius: 8px;
        padding: 10px 20px;
        font-size: 16px;
    }
    .stButton > button:hover {
        background-color: #45a049;
        color: white;
    }
    .stSidebar {
        background-color: #F7F7F7;
    }
    .header {
        font-size: 36px;
        font-weight: bold;
        color: #4CAF50;
        text-align: center;
    }
    .subheader {
        font-size: 20px;
        color: #333;
        margin-top: 0;
    }
    .footer {
        font-size: 14px;
        color: #666;
        text-align: center;
        margin-top: 30px;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# Load the model
@st.cache_resource
def load_model():
    return joblib.load("best_model.pkl")


model = load_model()

# Sidebar customization
with st.sidebar:
    st.image("MAT_primary-logo@2x.jpg", use_container_width=True)
    st.header("Retention Predictor")
    st.write("Enter the required details below to predict retention probability.")

# Main app title
st.markdown('<div class="header">Medication Retention Probability Predictor</div>', unsafe_allow_html=True)
#st.markdown('<p class="subheader">Analyze patient retention using advanced ML models.</p>', unsafe_allow_html=True)

# Input fields for features
st.sidebar.subheader("Patient Episode Details")
previous_bup_episodes = st.sidebar.number_input(
    "Number of Previous Buprenorphine Episodes", min_value=0, value=20, step=1
)
previous_failed_bup_episodes = st.sidebar.number_input(
    "Number of Previous Failed Buprenorphine Episodes", min_value=0, value=20, step=1
)

st.sidebar.subheader("Duration Details")
days_covered = st.sidebar.number_input(
    "Days Covered in Current Episode", min_value=0, value=30, step=1
)
days_since_start = st.sidebar.number_input(
    "Days Since Start of the Episode", min_value=1, value=90, step=1
)

st.sidebar.subheader("Other Features")
prior_BUP_days_supply = st.sidebar.number_input(
    "Prior BUP Days Supply", min_value=0, value=30, step=1
)
YearOfInitiation = st.sidebar.number_input(
    "Year of Initiation", min_value=1900, max_value=2100, value=2023, step=1
)
Previous_Meth_Episodes = st.sidebar.radio(
    "Previous Methadone Episodes (Yes/No)",
    options=[0, 1],
    format_func=lambda x: "Yes" if x == 1 else "No",
)

# Calculate derived features and predict
if st.sidebar.button("🔍 Predict Retention Probability"):
    # Calculate derived features
    if previous_bup_episodes > 0:
        Failed_Bup_Rate = previous_failed_bup_episodes / previous_bup_episodes
    else:
        Failed_Bup_Rate = 0.0

    if days_since_start > 0:
        proportionofDaysCovered = days_covered / days_since_start
    else:
        proportionofDaysCovered = 0.0

    # Prepare input data
    input_data = pd.DataFrame(
        {
            "proportionofDaysCovered": [proportionofDaysCovered],
            "Failed_Bup_Rate": [Failed_Bup_Rate],
            "prior_BUP_days_supply": [prior_BUP_days_supply],
            "YearOfInitiation": [YearOfInitiation],
            "Previous_Meth_Episodes": [Previous_Meth_Episodes],
        }
    )

    # Make prediction
    proba_output = model.predict_proba(input_data)
    print("Predict_proba output shape:", proba_output.shape)

    # Handle both 1D and 2D outputs
    retention_probability = (
        proba_output[:, 1][0] if proba_output.ndim == 2 else proba_output[0]
    )

    # Display the prediction result in a styled card
    st.markdown(
        f"""
        <div style="padding: 20px; border: 1px solid #4CAF50; border-radius: 8px; background-color: #FFFFFF;">
            <h2 style="color: #4CAF50; text-align: center;">Prediction Result</h2>
            <p style="font-size: 18px; color: #333; text-align: center;">
                The predicted retention probability is:
            </p>
            <p style="font-size: 28px; font-weight: bold; color: #4CAF50; text-align: center;">
                {retention_probability:.2f}
            </p>
        </div>
        """,
        unsafe_allow_html=True,
    )

    # Display calculated features in a visually appealing manner
    st.markdown("### **Calculated Features**")
    st.markdown(
        f"""
        <ul style="font-size: 18px; color: #333;">
            <li><strong>Failed Buprenorphine Rate:</strong> {Failed_Bup_Rate:.2%}</li>
            <li><strong>Proportion of Days Covered:</strong> {proportionofDaysCovered:.2%}</li>
        </ul>
        """,
        unsafe_allow_html=True,
    )

# Footer section
st.markdown(
    """
    <div style="text-align: center; font-size: 14px; color: #666; margin-top: 30px;">
        Developed by <strong>Nahiyan Noor</strong>
    </div>
    """,
    unsafe_allow_html=True,
)