import streamlit as st

# Page configuration - MUST be first Streamlit command
st.set_page_config(
    page_title="Risk Level Prediction App - Tugas Akhir MBD",
    page_icon="🩺",
    layout="wide",
    initial_sidebar_state="collapsed"
)

import pandas as pd
import numpy as np
import joblib
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime
import json

# Load model components
@st.cache_resource
def load_model():
    """Load the trained model components"""
    try:
        components = joblib.load('./src/model.joblib')
        return components
    except FileNotFoundError:
        st.error("Model file 'RiskLevel_prediction_components.joblib' not found!")
        st.stop()
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        st.stop()

def predict_RiskLevel(data, model_components):
    """
    Predict maternal health risk category using the trained decision tree model.

    Parameters:
    -----------
    data : dict
        Dictionary with input features. Must include:
        - Age: float or int
        - SystolicBP: float
        - DiastolicBP: float
        - BS: float (Blood Sugar)
        - BodyTemp: float
        - HeartRate: float

    Returns:
    --------
    dict
        Dictionary containing:
        - prediction: int (class index)
        - prediction_label: str (Low Risk / Mid Risk / High Risk)
        - probability: float
    """
    # Load model components
    components = joblib.load('./src/model.joblib')

    # Get needed components
    model = components['model']
    encoding_maps = components['encoding_maps']
    feature_names = components['feature_names']

    # Convert to DataFrame for easier processing
    if isinstance(data, dict):
        df = pd.DataFrame([data])
    else:
        df = data.copy()

    # Apply encodings to each categorical column
    for column in df.columns:
        if column in encoding_maps and column != 'RiskLevel':
            df[column] = df[column].map(encoding_maps[column])

    # Ensure we're only using the features that the model was trained on
    df_for_pred = df[feature_names].copy()

    # Make prediction
    prediction = model.predict(df_for_pred)[0]
    probabilities = model.predict_proba(df_for_pred)[0]

    # Pastikan prediction bertipe int native
    prediction = int(prediction.item() if hasattr(prediction, 'item') else prediction)

    # Get income label (mapping from numeric prediction back to label)
    RiskLevel_map_inverse = {v: k for k, v in encoding_maps['RiskLevel'].items()}
    prediction_label = RiskLevel_map_inverse[prediction]

    return {
        'prediction': int(prediction),
        'prediction_label': prediction_label,
        'probability': float(probabilities[prediction]),
        'probabilities': probabilities.tolist()
    }

def validate_inputs(data):
    errors = []

    # Berdasarkan data Anda:
    # Age: min = 10, max = 50
    if not (10 <= data['Age'] <= 50):
        errors.append("Age should be between 10 and 50")

    # SystolicBP: min = 70, max = 140
    if not (70 <= data['SystolicBP'] <= 140):
        errors.append("Systolic BP should be between 70 and 140")

    # DiastolicBP: min = 50, max = 100
    if not (50 <= data['DiastolicBP'] <= 100):
        errors.append("Diastolic BP should be between 50 and 100")

    # Blood Sugar (BS): min = 2, max = 18 (semua bilangan bulat, tidak memakai koma)
    if not (2 <= data['BS'] <= 18):
        errors.append("Blood Sugar should be between 2 and 18")

    # Body Temp: min = 96, max = 102 (integer, tidak memakai koma/desimal)
    if not (96 <= data['BodyTemp'] <= 102):
        errors.append("Body Temperature should be between 96 and 102 °F")

    # Heart Rate: min = 66, max = 101
    if not (66 <= data['HeartRate'] <= 101):
        errors.append("Heart Rate should be between 66 and 101 bpm")

    return errors

def export_prediction(data, result):
    """Export prediction result to JSON"""
    export_data = {
        'timestamp': datetime.now().isoformat(),
        'input_data': data,
        'prediction': {
            'class': result['prediction_label'],
            'confidence': result['probability'],
            'raw_prediction': result['prediction']
        }
    }
    return json.dumps(export_data, indent=2)

def reset_session_state():
    """Reset all input values to default"""
    keys_to_reset = [
        'Age', 	'SystolicBP',	'DiastolicBP',	'BS',	'BodyTemp', 'HeartRate',	'RiskLevel'
    ]
    for key in keys_to_reset:
        if key in st.session_state:
            del st.session_state[key]

# Load model
model_components = load_model()

# Main app
st.title("🩺 Maternal Health Risk App - Tugas Akhir MBD")
st.markdown("Predict the **Risk Level** for maternal health using basic physiological data.")

# Create two columns for layout
col1, col2 = st.columns([2, 1])

with col1:
    st.subheader("📝 Input Features")

    # Create form for inputs
    with st.form("prediction_form"):
        # Medical Measurements Information
        st.markdown("**Medical Measurements**")
        col_demo1, col_demo2 = st.columns(2)

        with col_demo1:
            Age = st.number_input("Age", min_value=10, max_value=50, value=25, key="Age")
            SystolicBP = st.number_input("Systolic Blood Pressure (mmHg)", min_value=70, max_value=140, value=120, key="SystolicBP")
            DiastolicBP = st.number_input("Diastolic Blood Pressure (mmHg)", min_value=50, max_value=100, value=80, key="DiastolicBP")

        with col_demo2:
            BS = st.number_input("Blood Sugar Level", min_value=2, max_value=18, value=5, key="BS")
            BodyTemp = st.number_input("Body Temperature (°F)", min_value=96, max_value=102, value=98, key="BodyTemp")
            HeartRate = st.number_input("Heart Rate (bpm)", min_value=66, max_value=101, value=85, key="HeartRate")

        st.divider()

        # Buttons
        col_btn1, col_btn2, col_btn3 = st.columns(3)
        with col_btn1:
            predict_button = st.form_submit_button("🔮 Predict", type="primary")
        with col_btn2:
            reset_button = st.form_submit_button("🔄 Reset")
        with col_btn3:
            export_button = st.form_submit_button("📤 Export Last Result")

# Handle reset button
if reset_button:
    reset_session_state()
    st.rerun()

# Handle prediction
if predict_button:
    # Collect input data
    input_data = {
        'Age': Age,
        'SystolicBP': SystolicBP,
        'DiastolicBP': DiastolicBP,
        'BS': BS,
        'BodyTemp': BodyTemp,
        'HeartRate': HeartRate
    }

    # Validate inputs
    validation_errors = validate_inputs(input_data)

    if validation_errors:
        with col2:
            st.error("❌ Validation Errors:")
            for error in validation_errors:
                st.error(f"• {error}")
    else:
        # Make prediction
        try:
            result = predict_RiskLevel(input_data, model_components)

            # Store result in session state for export
            st.session_state['last_prediction'] = {
                'input_data': input_data,
                'result': result
            }

            with col2:
                st.subheader("🎯 Prediction Results")

                # Display prediction
                prediction_color = {
                    'low risk': 'green',
                    'mid risk': 'orange',
                    'high risk': 'red'
                }.get(result['prediction_label'], 'blue')

                st.markdown(f"**Predicted Risk:** :{prediction_color}[{result['prediction_label']}]")

                # Confidence level with gauge
                confidence = result['probability'] * 100

                fig_gauge = go.Figure(go.Indicator(
                    mode = "gauge+number+delta",
                    value = confidence,
                    domain = {'x': [0, 1], 'y': [0, 1]},
                    title = {'text': "Confidence Level (%)"},
                    gauge = {
                        'axis': {'range': [None, 100]},
                        'bar': {'color': prediction_color},
                        'steps': [
                            {'range': [0, 50], 'color': "lightgray"},
                            {'range': [50, 80], 'color': "yellow"},
                            {'range': [80, 100], 'color': "lightgreen"}
                        ],
                        'threshold': {
                            'line': {'color': "red", 'width': 4},
                            'thickness': 0.75,
                            'value': 90
                        }
                    }
                ))
                fig_gauge.update_layout(height=300, margin=dict(l=20, r=20, t=40, b=20))
                st.plotly_chart(fig_gauge, use_container_width=True)

                # Probability breakdown
                prob_df = pd.DataFrame({
                    'Class': ['low risk', 'mid risk', 'high risk'],
                    'Probability': result['probabilities']
                })

                fig_bar = px.bar(
                    prob_df,
                    x='Class',
                    y='Probability',
                    title='Risk Class Distribution',
                    color='Probability',
                    color_continuous_scale=['orange', 'green']
                )
                fig_bar.update_layout(height=300, margin=dict(l=20, r=20, t=40, b=20))
                st.plotly_chart(fig_bar, use_container_width=True)

        except Exception as e:
            with col2:
                st.error(f"❌ Prediction Error: {str(e)}")

# Feature Importance section
st.subheader("📊 Feature Importance")

if 'model' in model_components:
    try:
        feature_names = model_components['feature_names']
        feature_importance = model_components['model'].feature_importances_

        importance_df = pd.DataFrame({
            'Feature': feature_names,
            'Importance': feature_importance
        }).sort_values('Importance', ascending=True)

        fig_importance = px.bar(
            importance_df,
            x='Importance',
            y='Feature',
            orientation='h',
            title='Feature Importance in Decision Tree Model',
            color='Importance',
            color_continuous_scale='viridis'
        )
        fig_importance.update_layout(height=400, margin=dict(l=20, r=20, t=40, b=20))
        st.plotly_chart(fig_importance, use_container_width=True)

    except Exception as e:
        st.error(f"Error displaying feature importance: {str(e)}")

# Handle export
if export_button:
    if 'last_prediction' in st.session_state:
        export_data = export_prediction(
            st.session_state['last_prediction']['input_data'],
            st.session_state['last_prediction']['result']
        )

        st.download_button(
            label="📥 Download Prediction Results",
            data=export_data,
            file_name=f"RiskLevel_prediction_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
            mime="application/json"
        )
    else:
        st.warning("⚠️ No prediction results to export. Please make a prediction first.")

# Footer
st.markdown("---")
st.markdown("*Built with Streamlit • Dr. Eng. Farrikh Alzami, M.Kom*")