import streamlit as st
import joblib
import pandas as pd
import numpy as np
from tensorflow.keras.models import load_model 

MODEL_PATH = 'src/heart.h5' 
SCALER_PATH = 'src/scaler_heart.joblib' 
FEATURES_PATH = 'src/final_features.joblib'

ORIGINAL_CATEGORICAL_COLS = ['sex', 'cp', 'fbs', 'restecg', 'exang', 'slope', 'ca', 'thal']
CONTINUOUS_COLS = ['age', 'trestbps', 'chol', 'thalach', 'oldpeak']
ORIGINAL_ALL_COLS = CONTINUOUS_COLS + ORIGINAL_CATEGORICAL_COLS

@st.cache_resource
def load_assets():
    try:
        model = load_model(MODEL_PATH)
        scaler = joblib.load(SCALER_PATH)
        final_features_list = joblib.load(FEATURES_PATH) # Load the correct feature order
        return model, scaler, final_features_list
    except Exception as e:
        st.error(f"Error loading assets. Ensure all three files are uploaded. Error: {e}")
        return None, None, None

def preprocess_and_predict(model, scaler, final_features_list, input_data):
    
    # 1. Create DataFrame from inputs
    df_input = pd.DataFrame([input_data])

    # 2. One-Hot Encoding (OHE) for Categorical Features (drop_first=True)
    df_processed = pd.get_dummies(df_input, columns=ORIGINAL_CATEGORICAL_COLS, drop_first=True)
    
    # 3. Align Columns: Add missing dummy columns and reorder (CRITICAL STEP)
    for feature in final_features_list:
        if feature not in df_processed.columns:
            # Add missing dummy variables (e.g., cp_1, cp_2) with value 0
            df_processed[feature] = 0
            
    # Select and order the final feature array using the loaded list
    final_features_df = df_processed[final_features_list].copy()
    
    # 4. Scaling Numerical Data (CRITICAL STEP)
    numerical_part = final_features_df[CONTINUOUS_COLS]
    final_features_df[CONTINUOUS_COLS] = scaler.transform(numerical_part)
    
    # 5. Predict (The model expects a numpy array)
    prediction_proba = model.predict(final_features_df.values)
    
    return float(prediction_proba[0])

# --- Streamlit Interface ---
st.set_page_config(page_title="Heart Disease Predictor", layout="centered")
st.title("❤️ Heart Disease Prediction (Neural Network)")
st.markdown("Enter patient data (all 13 features) to predict the probability of heart disease.")

model, scaler, final_features_list = load_assets()

if model is not None and scaler is not None and final_features_list is not None:
    st.sidebar.header("Patient Data Input (13 Features)")
    
    # Continuous Features
    age = st.sidebar.slider("Age:", min_value=18, max_value=100, value=50)
    trestbps = st.sidebar.number_input("Resting BP (trestbps):", min_value=90, max_value=200, value=120)
    chol = st.sidebar.number_input("Cholesterol (chol):", min_value=100, max_value=600, value=250)
    thalach = st.sidebar.number_input("Max Heart Rate (thalach):", min_value=60, max_value=220, value=150)
    oldpeak = st.sidebar.number_input("ST Depression (oldpeak):", min_value=0.0, max_value=6.2, value=1.0, step=0.1)

    # Categorical Features (Use indices matching the original dataset)
    sex = st.sidebar.selectbox("Sex (1=Male, 0=Female):", options=[1, 0])
    cp = st.sidebar.selectbox("Chest Pain Type (cp):", options=[0, 1, 2, 3], index=0)
    fbs = st.sidebar.selectbox("Fasting Blood Sugar > 120 (fbs):", options=[0, 1])
    restecg = st.sidebar.selectbox("Resting ECG (restecg):", options=[0, 1, 2], index=1)
    exang = st.sidebar.selectbox("Exercise Induced Angina (exang):", options=[0, 1])
    slope = st.sidebar.selectbox("Slope of Peak ST (slope):", options=[0, 1, 2], index=1)
    ca = st.sidebar.selectbox("Major Vessels (ca):", options=[0, 1, 2, 3, 4], index=0)
    thal = st.sidebar.selectbox("Thal (thal):", options=[1, 2, 3], index=1)
    
    # Collect inputs
    input_data = {
        'age': age, 'sex': sex, 'cp': cp, 'trestbps': trestbps, 'chol': chol,
        'fbs': fbs, 'restecg': restecg, 'thalach': thalach, 'exang': exang, 
        'oldpeak': oldpeak, 'slope': slope, 'ca': ca, 'thal': thal
    }

    if st.button("Predict Probability"):
        with st.spinner('Calculating probability...'):
            prediction_proba = preprocess_and_predict(model, scaler, final_features_list, input_data)
            
            st.success("Prediction Successful!")
            st.markdown("### Predicted Heart Disease Probability:")
            st.markdown(f"**{prediction_proba * 100:.1f}%**")
            
            risk = "HIGH RISK" if prediction_proba > 0.5 else "LOW RISK"
            st.markdown(f"Outcome: **{risk}**")