app.py

"""
Streamlit Application for Predictive Maintenance Project
Interactive web app for EDA, model visualization, and runtime predictions
"""

import streamlit as st
import pandas as pd
import numpy as np
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import pickle
import warnings
warnings.filterwarnings('ignore')

# Import custom modules
from preprocessing import DataPreprocessor
from model import PredictiveMaintenanceModel

# Page configuration
st.set_page_config(
    page_title="Predictive Maintenance System",
    page_icon="🔧",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
    <style>
    .main-header {
        font-size: 3rem;
        font-weight: bold;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 2rem;
    }
    .sub-header {
        font-size: 1.5rem;
        color: #ff7f0e;
        margin-top: 2rem;
        margin-bottom: 1rem;
    }
    .metric-card {
        background-color: #f0f2f6;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 0.5rem 0;
    }
    </style>
""", unsafe_allow_html=True)

# Initialize session state
if 'model' not in st.session_state:
    st.session_state.model = None
if 'preprocessor' not in st.session_state:
    st.session_state.preprocessor = None
if 'data' not in st.session_state:
    st.session_state.data = None

@st.cache_data
def load_data():
    """Load and cache the dataset"""
    df = pd.read_csv('ai4i2020.csv')
    # Create additional features
    df['Temperature difference [K]'] = (
        df['Process temperature [K]'] - df['Air temperature [K]']
    )
    df['Power [W]'] = (
        df['Rotational speed [rpm]'] * df['Torque [Nm]'] / 9.5488
    )
    return df

def train_model():
    """Train the model and preprocessor"""
    with st.spinner("Training model... This may take a moment."):
        preprocessor = DataPreprocessor('ai4i2020.csv')
        X_train, X_test, y_train, y_test, feature_columns = preprocessor.prepare_data()
        
        model = PredictiveMaintenanceModel()
        model.train(X_train, y_train)
        
        # Evaluate model
        results = model.evaluate(X_test, y_test)
        
        st.session_state.model = model
        st.session_state.preprocessor = preprocessor
        st.session_state.feature_columns = feature_columns
        
        return results

# Main App
def main():
    # Header
    st.markdown('<h1 class="main-header">🔧 Predictive Maintenance System</h1>', unsafe_allow_html=True)
    st.markdown("---")
    
    # Sidebar Navigation
    st.sidebar.title("Navigation")
    page = st.sidebar.radio(
        "Select Page",
        ["Introduction", "Exploratory Data Analysis", "Model & Predictions", "Conclusion"]
    )
    
    # Load data
    df = load_data()
    st.session_state.data = df
    
    if page == "Introduction":
        show_introduction(df)
    elif page == "Exploratory Data Analysis":
        show_eda(df)
    elif page == "Model & Predictions":
        show_model_predictions(df)
    elif page == "Conclusion":
        show_conclusion()

def show_introduction(df):
    """Introduction page"""
    st.markdown('<h2 class="sub-header">📋 Project Introduction</h2>', unsafe_allow_html=True)
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.markdown("""
        ### About the Dataset
        
        This project uses the **AI4I 2020 Predictive Maintenance Dataset**, which contains 
        synthetic data simulating predictive maintenance scenarios for industrial machinery.
        
        #### Dataset Overview:
        - **Total Records**: 10,000 machines
        - **Features**: 14 attributes including temperature, rotational speed, torque, and tool wear
        - **Target**: Machine failure prediction (binary classification)
        - **Failure Types**: Tool Wear Failure (TWF), Heat Dissipation Failure (HDF), 
          Power Failure (PWF), Overstrain Failure (OSF), and Random Failure (RNF)
        
        #### Project Goals:
        1. **Exploratory Data Analysis**: Understand patterns and relationships in the data
        2. **Predictive Modeling**: Build a machine learning model to predict machine failures
        3. **Maintenance Scheduling**: Estimate when maintenance is needed and how urgent it is
        4. **Interactive Visualization**: Present findings through an interactive web application
        
        #### Key Features:
        - Comprehensive EDA with 15+ different analyses
        - Random Forest Classifier for failure prediction
        - Real-time predictions based on user input
        - Maintenance urgency assessment
        - Time-to-failure estimation
        """)
    
    with col2:
        st.markdown("### Dataset Statistics")
        st.metric("Total Machines", f"{len(df):,}")
        st.metric("Features", len(df.columns))
        st.metric("Machine Failures", f"{df['Machine failure'].sum():,}")
        st.metric("Failure Rate", f"{(df['Machine failure'].mean()*100):.2f}%")
        
        st.markdown("### Machine Types")
        type_counts = df['Type'].value_counts()
        type_meanings = {'L': 'Low Quality/Load', 'M': 'Medium Quality/Load', 'H': 'High Quality/Load'}
        for machine_type, count in type_counts.items():
            meaning = type_meanings.get(machine_type, '')
            st.metric(f"Type {machine_type} ({meaning})", f"{count:,}")
    
    st.markdown("---")
    st.markdown("### Dataset Preview")
    preview_mode = st.radio(
        "Preview mode",
        options=["First 20 rows", "Show all (10,000 rows)"],
        index=0,
        horizontal=True,
    )
    if preview_mode == "First 20 rows":
        st.dataframe(df.head(20), use_container_width=True)
    else:
        st.dataframe(df, use_container_width=True)
    
    st.markdown("### Dataset Information")
    with st.expander("View Column Descriptions"):
        st.markdown("""
        - **UDI**: Unique identifier for each machine
        - **Product ID**: Product identifier
        - **Type**: Machine type (L = Low Quality/Load, M = Medium Quality/Load, H = High Quality/Load)
        - **Air temperature [K]**: Air temperature in Kelvin
        - **Process temperature [K]**: Process temperature in Kelvin
        - **Rotational speed [rpm]**: Rotational speed in revolutions per minute
        - **Torque [Nm]**: Torque in Newton meters
        - **Tool wear [min]**: Tool wear in minutes
        - **Machine failure**: Binary target (0 = no failure, 1 = failure)
        - **TWF, HDF, PWF, OSF, RNF**: Different failure type indicators
        """)

def show_eda(df):
    """EDA page"""
    st.markdown('<h2 class="sub-header">📊 Exploratory Data Analysis</h2>', unsafe_allow_html=True)
    
    # Analysis selection
    analysis_type = st.selectbox(
        "Select Analysis Type",
        [
            "Summary Statistics",
            "Data Types & Unique Values",
            "Target Distribution",
            "Feature Distributions",
            "Correlation Analysis",
            "Failure Analysis by Type",
            "Tool Wear Analysis",
            "Temperature Analysis",
            "Power & Rotational Speed Analysis",
            "Outlier Detection",
            "Pairwise Relationships",
            "Failure Type Breakdown",
            "Time to Failure Estimation",
            "Grouped Aggregations"
        ]
    )
    
    st.markdown("---")
    
    if analysis_type == "Summary Statistics":
        st.subheader("Summary Statistics")
        st.dataframe(df.describe(), use_container_width=True)
        
        # Key metrics
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.metric("Mean Air Temp", f"{df['Air temperature [K]'].mean():.2f} K")
        with col2:
            st.metric("Mean Process Temp", f"{df['Process temperature [K]'].mean():.2f} K")
        with col3:
            st.metric("Mean Rotational Speed", f"{df['Rotational speed [rpm]'].mean():.0f} rpm")
        with col4:
            st.metric("Mean Torque", f"{df['Torque [Nm]'].mean():.2f} Nm")
    
    elif analysis_type == "Data Types & Unique Values":
        st.subheader("Data Types and Unique Values")
        
        info_df = pd.DataFrame({
            'Column': df.columns,
            'Data Type': df.dtypes.astype(str),
            'Unique Values': [df[col].nunique() for col in df.columns],
            'Non-Null Count': df.count().values
        })
        
        st.dataframe(info_df, use_container_width=True)
        
        st.subheader("Machine Type Distribution")
        type_counts = df['Type'].value_counts()
        type_meanings = {'L': 'Low Quality/Load', 'M': 'Medium Quality/Load', 'H': 'High Quality/Load'}
        
        col1, col2 = st.columns(2)
        with col1:
            fig = px.bar(
                x=type_counts.index,
                y=type_counts.values,
                title="Machine Type Distribution",
                labels={'x': 'Machine Type', 'y': 'Count'},
                color=type_counts.index,
                color_discrete_sequence=['#e74c3c', '#3498db', '#2ecc71']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            for machine_type, count in type_counts.items():
                meaning = type_meanings.get(machine_type, '')
                st.metric(f"Type {machine_type} ({meaning})", f"{count:,}")
    
    elif analysis_type == "Target Distribution":
        st.subheader("Machine Failure Distribution")
        
        col1, col2 = st.columns(2)
        
        with col1:
            failure_counts = df['Machine failure'].value_counts()
            fig = px.pie(
                values=failure_counts.values,
                names=['No Failure', 'Failure'],
                title="Failure Distribution",
                color_discrete_sequence=['#2ecc71', '#e74c3c']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            st.metric("No Failure", f"{failure_counts[0]:,} ({(failure_counts[0]/len(df)*100):.2f}%)")
            st.metric("Failure", f"{failure_counts[1]:,} ({(failure_counts[1]/len(df)*100):.2f}%)")
            
            # Failure types
            st.subheader("Failure Types Breakdown")
            failure_types = {
                'TWF': 'Tool Wear Failure',
                'HDF': 'Heat Dissipation Failure',
                'PWF': 'Power Failure',
                'OSF': 'Overstrain Failure',
                'RNF': 'Random Failure'
            }
            for ft_code, ft_name in failure_types.items():
                count = df[ft_code].sum()
                st.metric(ft_name, f"{count} ({(count/len(df)*100):.2f}%)")
    
    elif analysis_type == "Feature Distributions":
        st.subheader("Feature Distributions")
        
        feature = st.selectbox(
            "Select Feature",
            ['Air temperature [K]', 'Process temperature [K]', 'Rotational speed [rpm]', 
             'Torque [Nm]', 'Tool wear [min]', 'Temperature difference [K]', 'Power [W]']
        )
        
        col1, col2 = st.columns(2)
        
        with col1:
            fig = px.histogram(
                df, x=feature, nbins=50,
                title=f"Distribution of {feature}",
                color_discrete_sequence=['#3498db']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            fig = px.box(
                df, y=feature,
                title=f"Box Plot of {feature}",
                color_discrete_sequence=['#9b59b6']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        # Statistics
        st.subheader("Statistics")
        col1, col2, col3, col4 = st.columns(4)
        with col1:
            st.metric("Mean", f"{df[feature].mean():.2f}")
        with col2:
            st.metric("Median", f"{df[feature].median():.2f}")
        with col3:
            st.metric("Std Dev", f"{df[feature].std():.2f}")
        with col4:
            st.metric("Skewness", f"{df[feature].skew():.2f}")
    
    elif analysis_type == "Correlation Analysis":
        st.subheader("Correlation Analysis")
        
        numerical_cols = [
            'Air temperature [K]', 'Process temperature [K]',
            'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]',
            'Temperature difference [K]', 'Power [W]', 'Machine failure'
        ]
        
        corr_matrix = df[numerical_cols].corr()
        
        fig = px.imshow(
            corr_matrix,
            text_auto=True,
            aspect="auto",
            title="Correlation Heatmap",
            color_continuous_scale="RdBu"
        )
        st.plotly_chart(fig, use_container_width=True)
        
        st.subheader("Correlation with Machine Failure")
        failure_corr = corr_matrix['Machine failure'].sort_values(ascending=False)
        fig = px.bar(
            x=failure_corr.index,
            y=failure_corr.values,
            title="Feature Correlation with Machine Failure",
            labels={'x': 'Feature', 'y': 'Correlation'},
            color=failure_corr.values,
            color_continuous_scale="RdYlGn"
        )
        st.plotly_chart(fig, use_container_width=True)
    
    elif analysis_type == "Failure Analysis by Type":
        st.subheader("Failure Analysis by Machine Type")
        st.info("**Machine Type Meanings**: L = Low Quality/Load, M = Medium Quality/Load, H = High Quality/Load")
        
        failure_by_type = df.groupby('Type')['Machine failure'].agg(['count', 'sum', 'mean']).reset_index()
        failure_by_type.columns = ['Type', 'Total Machines', 'Failures', 'Failure Rate']
        failure_by_type['Failure Rate'] = failure_by_type['Failure Rate'] * 100
        failure_by_type['Type_Label'] = failure_by_type['Type'].map({
            'L': 'L (Low Quality/Load)',
            'M': 'M (Medium Quality/Load)',
            'H': 'H (High Quality/Load)'
        })
        
        st.dataframe(failure_by_type[['Type', 'Total Machines', 'Failures', 'Failure Rate']], use_container_width=True)
        
        fig = px.bar(
            failure_by_type,
            x='Type_Label',
            y='Failure Rate',
            title="Failure Rate by Machine Type",
            color='Type',
            color_discrete_sequence=['#e74c3c', '#3498db', '#2ecc71'],
            labels={'Type_Label': 'Machine Type'}
        )
        st.plotly_chart(fig, use_container_width=True)
    
    elif analysis_type == "Tool Wear Analysis":
        st.subheader("Tool Wear Analysis")
        
        col1, col2 = st.columns(2)
        
        with col1:
            fig = px.scatter(
                df,
                x='Tool wear [min]',
                y='Machine failure',
                color='Machine failure',
                title="Tool Wear vs Machine Failure",
                color_discrete_sequence=['#2ecc71', '#e74c3c']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            tool_wear_by_failure = df.groupby('Machine failure')['Tool wear [min]'].agg(['mean', 'median', 'std'])
            st.dataframe(tool_wear_by_failure, use_container_width=True)
        
        # Tool wear distribution by failure status
        fig = px.histogram(
            df,
            x='Tool wear [min]',
            color='Machine failure',
            nbins=50,
            title="Tool Wear Distribution by Failure Status",
            barmode='overlay',
            color_discrete_sequence=['#2ecc71', '#e74c3c']
        )
        st.plotly_chart(fig, use_container_width=True)
    
    elif analysis_type == "Temperature Analysis":
        st.subheader("Temperature Analysis")
        
        col1, col2 = st.columns(2)
        
        with col1:
            fig = px.scatter(
                df,
                x='Air temperature [K]',
                y='Process temperature [K]',
                color='Machine failure',
                title="Temperature Relationship",
                color_discrete_sequence=['#2ecc71', '#e74c3c']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            temp_by_failure = df.groupby('Machine failure')[
                ['Air temperature [K]', 'Process temperature [K]', 'Temperature difference [K]']
            ].mean()
            st.dataframe(temp_by_failure, use_container_width=True)
    
    elif analysis_type == "Power & Rotational Speed Analysis":
        st.subheader("Power and Rotational Speed Analysis")
        
        power_stats = df.groupby('Machine failure')[
            ['Rotational speed [rpm]', 'Torque [Nm]', 'Power [W]']
        ].agg(['mean', 'std', 'min', 'max'])
        
        st.dataframe(power_stats, use_container_width=True)
        
        col1, col2 = st.columns(2)
        
        with col1:
            fig = px.box(
                df,
                x='Machine failure',
                y='Rotational speed [rpm]',
                title="Rotational Speed by Failure Status",
                color='Machine failure',
                color_discrete_sequence=['#2ecc71', '#e74c3c']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            fig = px.box(
                df,
                x='Machine failure',
                y='Power [W]',
                title="Power by Failure Status",
                color='Machine failure',
                color_discrete_sequence=['#2ecc71', '#e74c3c']
            )
            st.plotly_chart(fig, use_container_width=True)
        
        # Scatter plot: Power vs Rotational Speed
        fig = px.scatter(
            df,
            x='Rotational speed [rpm]',
            y='Power [W]',
            color='Machine failure',
            title="Power vs Rotational Speed",
            color_discrete_sequence=['#2ecc71', '#e74c3c']
        )
        st.plotly_chart(fig, use_container_width=True)
    
    elif analysis_type == "Outlier Detection":
        st.subheader("Outlier Detection")
        
        feature = st.selectbox(
            "Select Feature for Outlier Detection",
            ['Air temperature [K]', 'Process temperature [K]', 'Rotational speed [rpm]', 
             'Torque [Nm]', 'Tool wear [min]']
        )
        
        method = st.radio(
            "Detection method",
            options=["IQR (robust, default)", "Z-score"],
            index=0,
            horizontal=True,
            help="IQR is robust to skew; Z-score highlights extreme standardized values."
        )
        
        if method == "IQR (robust, default)":
            iqr_mult = st.slider("IQR multiplier", 0.5, 3.0, 1.5, 0.1,
                                 help="Lower the multiplier to surface milder outliers.")
            Q1 = df[feature].quantile(0.25)
            Q3 = df[feature].quantile(0.75)
            IQR = Q3 - Q1
            lower_bound = Q1 - iqr_mult * IQR
            upper_bound = Q3 + iqr_mult * IQR
            
            outliers = df[(df[feature] < lower_bound) | (df[feature] > upper_bound)]
            
            col1, col2 = st.columns(2)
            with col1:
                st.metric("Lower Bound", f"{lower_bound:.2f}")
                st.metric("Upper Bound", f"{upper_bound:.2f}")
            with col2:
                st.metric("Outlier Count", len(outliers))
                st.metric("Outlier Percentage", f"{(len(outliers)/len(df)*100):.2f}%")
            
            fig = px.box(df, y=feature, title=f"Box Plot with Outliers - {feature}")
            st.plotly_chart(fig, use_container_width=True)
        
        else:
            z_thresh = st.slider("Z-score threshold", 2.0, 5.0, 3.0, 0.1,
                                 help="Lower threshold to surface more anomalies.")
            mean = df[feature].mean()
            std = df[feature].std()
            if std == 0:
                outliers = df.iloc[0:0]
                zscores = pd.Series([0]*len(df), index=df.index)
            else:
                zscores = (df[feature] - mean) / std
                outliers = df[zscores.abs() > z_thresh]
            
            col1, col2 = st.columns(2)
            with col1:
                st.metric("Mean", f"{mean:.2f}")
                st.metric("Std Dev", f"{std:.2f}")
            with col2:
                st.metric("Outlier Count", len(outliers))
                st.metric("Outlier Percentage", f"{(len(outliers)/len(df)*100):.2f}%")
            
            fig = px.histogram(df, x=feature, nbins=60, opacity=0.7,
                               title=f"{feature} with Z-score Threshold (>|{z_thresh}|)")
            # Overlay threshold lines
            fig.add_vline(x=mean + z_thresh*std, line_dash="dash", line_color="red")
            fig.add_vline(x=mean - z_thresh*std, line_dash="dash", line_color="red")
            st.plotly_chart(fig, use_container_width=True)
    
    elif analysis_type == "Pairwise Relationships":
        st.subheader("Pairwise Feature Relationships")
        
        feature1 = st.selectbox("Select First Feature", 
                               ['Tool wear [min]', 'Temperature difference [K]', 
                                'Rotational speed [rpm]', 'Torque [Nm]'])
        feature2 = st.selectbox("Select Second Feature",
                               ['Tool wear [min]', 'Temperature difference [K]',
                                'Rotational speed [rpm]', 'Torque [Nm]'])
        
        if feature1 != feature2:
            fig = px.scatter(
                df,
                x=feature1,
                y=feature2,
                color='Machine failure',
                title=f"{feature1} vs {feature2}",
                color_discrete_sequence=['#2ecc71', '#e74c3c']
            )
            st.plotly_chart(fig, use_container_width=True)
            
            correlation = df[feature1].corr(df[feature2])
            st.metric("Correlation", f"{correlation:.4f}")
    
    elif analysis_type == "Failure Type Breakdown":
        st.subheader("Detailed Failure Type Analysis")
        
        failure_types = {
            'TWF': 'Tool Wear Failure',
            'HDF': 'Heat Dissipation Failure',
            'PWF': 'Power Failure',
            'OSF': 'Overstrain Failure',
            'RNF': 'Random Failure'
        }
        
        for ft_code, ft_name in failure_types.items():
            with st.expander(f"{ft_name} ({ft_code})"):
                failed_machines = df[df[ft_code] == 1]
                if len(failed_machines) > 0:
                    st.metric("Count", len(failed_machines))
                    col1, col2, col3 = st.columns(3)
                    with col1:
                        st.metric("Avg Tool Wear", f"{failed_machines['Tool wear [min]'].mean():.2f} min")
                    with col2:
                        st.metric("Avg Temp Diff", f"{failed_machines['Temperature difference [K]'].mean():.2f} K")
                    with col3:
                        st.metric("Avg Rotational Speed", f"{failed_machines['Rotational speed [rpm]'].mean():.0f} rpm")
    
    elif analysis_type == "Time to Failure Estimation":
        st.subheader("Time to Failure Estimation")
        
        # Analyze tool wear progression for machines that failed
        failed_machines = df[df['Machine failure'] == 1]
        
        if len(failed_machines) > 0:
            avg_tool_wear_at_failure = failed_machines['Tool wear [min]'].mean()
            median_tool_wear_at_failure = failed_machines['Tool wear [min]'].median()
            
            col1, col2 = st.columns(2)
            with col1:
                st.metric("Average Tool Wear at Failure", f"{avg_tool_wear_at_failure:.2f} minutes")
            with col2:
                st.metric("Median Tool Wear at Failure", f"{median_tool_wear_at_failure:.2f} minutes")
            
            # Estimate time remaining for machines not yet failed
            non_failed = df[df['Machine failure'] == 0].copy()
            if len(non_failed) > 0:
                non_failed['Estimated Time to Failure'] = (
                    avg_tool_wear_at_failure - non_failed['Tool wear [min]']
                )
                non_failed['Estimated Time to Failure'] = non_failed['Estimated Time to Failure'].clip(lower=0)
                
                st.subheader("Time to Failure Estimates (for non-failed machines)")
                
                immediate = (non_failed['Estimated Time to Failure'] < 10).sum()
                soon = ((non_failed['Estimated Time to Failure'] >= 10) & 
                       (non_failed['Estimated Time to Failure'] < 50)).sum()
                remaining = (non_failed['Estimated Time to Failure'] >= 50).sum()
                
                col1, col2, col3 = st.columns(3)
                with col1:
                    st.metric("Immediate Maintenance (< 10 min)", f"{immediate:,}")
                with col2:
                    st.metric("Maintenance Soon (10-50 min)", f"{soon:,}")
                with col3:
                    st.metric("Time Remaining (> 50 min)", f"{remaining:,}")
                
                # Distribution chart
                fig = px.histogram(
                    non_failed,
                    x='Estimated Time to Failure',
                    nbins=50,
                    title="Distribution of Estimated Time to Failure",
                    color_discrete_sequence=['#3498db']
                )
                st.plotly_chart(fig, use_container_width=True)
    
    elif analysis_type == "Grouped Aggregations":
        st.subheader("Grouped Aggregations by Type and Failure Status")
        
        grouped = df.groupby(['Type', 'Machine failure']).agg({
            'Tool wear [min]': ['mean', 'std', 'max'],
            'Temperature difference [K]': ['mean', 'std'],
            'Rotational speed [rpm]': ['mean', 'std'],
            'Torque [Nm]': ['mean', 'std']
        })
        
        st.dataframe(grouped, use_container_width=True)
        
        # Visualizations
        st.subheader("Tool Wear by Type and Failure Status")
        fig = px.box(
            df,
            x='Type',
            y='Tool wear [min]',
            color='Machine failure',
            title="Tool Wear Distribution by Type and Failure Status",
            color_discrete_sequence=['#2ecc71', '#e74c3c']
        )
        st.plotly_chart(fig, use_container_width=True)
        
        st.subheader("Temperature Difference by Type and Failure Status")
        fig = px.box(
            df,
            x='Type',
            y='Temperature difference [K]',
            color='Machine failure',
            title="Temperature Difference by Type and Failure Status",
            color_discrete_sequence=['#2ecc71', '#e74c3c']
        )
        st.plotly_chart(fig, use_container_width=True)

def show_model_predictions(df):
    """Model and predictions page"""
    st.markdown('<h2 class="sub-header">🤖 Machine Learning Model & Predictions</h2>', unsafe_allow_html=True)
    
    # Train model section
    if st.session_state.model is None:
        st.info("⚠️ Model not trained yet. Click the button below to train the model.")
        if st.button("Train Model", type="primary"):
            results = train_model()
            st.success("✅ Model trained successfully!")
            st.session_state.model_results = results
    
    if st.session_state.model is not None:
        model = st.session_state.model
        preprocessor = st.session_state.preprocessor
        
        # Model performance section
        st.subheader("Model Performance")
        
        if 'model_results' in st.session_state:
            results = st.session_state.model_results
            col1, col2, col3, col4 = st.columns(4)
            with col1:
                st.metric("Accuracy", f"{results['accuracy']:.4f}")
            with col2:
                st.metric("Precision", f"{results['precision']:.4f}")
            with col3:
                st.metric("Recall", f"{results['recall']:.4f}")
            with col4:
                st.metric("F1-Score", f"{results['f1_score']:.4f}")
        
        # Feature importance
        st.subheader("Feature Importance")
        feature_importance = model.get_feature_importance()
        fig = px.bar(
            feature_importance,
            x='importance',
            y='feature',
            orientation='h',
            title="Feature Importance",
            color='importance',
            color_continuous_scale="Viridis"
        )
        st.plotly_chart(fig, use_container_width=True)
        
        st.markdown("---")
        
        # Runtime prediction section
        st.subheader("🔮 Runtime Prediction - Predict Maintenance Needs")
        st.markdown("Enter machine parameters below to predict if maintenance is needed:")
        
        col1, col2 = st.columns(2)
        
        with col1:
            machine_type = st.selectbox(
                "Machine Type", 
                ['L', 'M', 'H'],
                format_func=lambda x: f"{x} ({'Low Quality/Load' if x == 'L' else 'Medium Quality/Load' if x == 'M' else 'High Quality/Load'})"
            )
            air_temp = st.slider("Air Temperature (K)", 
                                min_value=295.0, max_value=305.0, value=298.0, step=0.1)
            process_temp = st.slider("Process Temperature (K)",
                                    min_value=305.0, max_value=315.0, value=309.0, step=0.1)
        
        with col2:
            rotational_speed = st.slider("Rotational Speed (rpm)",
                                        min_value=1000, max_value=3000, value=1500, step=10)
            torque = st.slider("Torque (Nm)",
                             min_value=10.0, max_value=80.0, value=40.0, step=0.1)
            tool_wear = st.slider("Tool Wear (minutes)",
                                min_value=0, max_value=300, value=50, step=1)
        
        if st.button("Predict Maintenance Status", type="primary"):
            # Create input data
            input_data = pd.DataFrame({
                'Type': [machine_type],
                'Air temperature [K]': [air_temp],
                'Process temperature [K]': [process_temp],
                'Rotational speed [rpm]': [rotational_speed],
                'Torque [Nm]': [torque],
                'Tool wear [min]': [tool_wear]
            })
            
            # Preprocess
            X_new = preprocessor.preprocess_new_data(input_data)
            
            # Predict
            maintenance_pred = model.predict_maintenance(X_new, tool_wear_values=[tool_wear])
            
            # Display results
            st.markdown("### Prediction Results")
            
            failure_prob = maintenance_pred['Failure_Probability'].iloc[0]
            time_to_failure = maintenance_pred['Time_to_Failure_Minutes'].iloc[0]
            status = maintenance_pred['Maintenance_Status'].iloc[0]
            urgency = maintenance_pred['Maintenance_Urgency'].iloc[0]
            
            # Color coding
            if urgency == "CRITICAL":
                st.error(f"🚨 **{status}**")
            elif urgency == "HIGH":
                st.warning(f"⚠️ **{status}**")
            elif urgency == "MEDIUM":
                st.info(f"ℹ️ **{status}**")
            else:
                st.success(f"✅ **{status}**")
            
            col1, col2, col3 = st.columns(3)
            with col1:
                st.metric("Failure Probability", f"{failure_prob:.2%}")
            with col2:
                st.metric("Estimated Time to Maintenance", f"{time_to_failure:.1f} minutes")
                if time_to_failure > 0:
                    st.caption(f"Based on historical data: avg failure at 120 min tool wear")
                else:
                    st.caption("Tool wear exceeded average failure threshold (120 min)")
            with col3:
                st.metric("Maintenance Urgency", urgency)
            
            # Detailed information
            with st.expander("View Detailed Information"):
                st.markdown(f"""
                **Machine Parameters:**
                - Type: {machine_type} ({'Low Quality/Load' if machine_type == 'L' else 'Medium Quality/Load' if machine_type == 'M' else 'High Quality/Load'})
                - Air Temperature: {air_temp} K
                - Process Temperature: {process_temp} K
                - Rotational Speed: {rotational_speed} rpm
                - Torque: {torque} Nm
                - Tool Wear: {tool_wear} minutes
                
                **Prediction Details:**
                - Failure Predicted: {'Yes' if maintenance_pred['Failure_Predicted'].iloc[0] == 1 else 'No'}
                - Failure Probability: {failure_prob:.2%}
                - Estimated Time to Maintenance: {time_to_failure:.1f} minutes
                  *Based on historical analysis: Machines in this dataset typically require maintenance when tool wear reaches approximately 120 minutes. 
                  This estimate projects when your machine will reach that threshold based on current tool wear level.*
                - Maintenance Status: {status}
                - Urgency Level: {urgency}
                
                **Recommendation:**
                {get_maintenance_recommendation(urgency, time_to_failure, failure_prob)}
                """)
        
        st.markdown("---")
        
        # Batch prediction section
        st.subheader("Batch Prediction")
        st.markdown("Upload a CSV file with machine data for batch predictions:")
        
        uploaded_file = st.file_uploader("Choose a CSV file", type="csv")
        
        if uploaded_file is not None:
            try:
                batch_data = pd.read_csv(uploaded_file)
                st.dataframe(batch_data.head(), use_container_width=True)
                
                if st.button("Predict for Batch", type="primary"):
                    # Check required columns
                    required_cols = ['Type', 'Air temperature [K]', 'Process temperature [K]',
                                    'Rotational speed [rpm]', 'Torque [Nm]', 'Tool wear [min]']
                    
                    if all(col in batch_data.columns for col in required_cols):
                        X_batch = preprocessor.preprocess_new_data(batch_data)
                        tool_wear_batch = batch_data['Tool wear [min]'].values
                        batch_predictions = model.predict_maintenance(X_batch, tool_wear_batch)
                        
                        # Combine with original data
                        results_df = pd.concat([batch_data, batch_predictions], axis=1)
                        
                        st.success("✅ Batch prediction complete!")
                        st.dataframe(results_df, use_container_width=True)
                        
                        # Summary statistics
                        st.subheader("Batch Prediction Summary")
                        col1, col2, col3, col4 = st.columns(4)
                        with col1:
                            st.metric("Total Machines", len(batch_data))
                        with col2:
                            st.metric("Critical Maintenance", 
                                    (batch_predictions['Maintenance_Urgency'] == 'CRITICAL').sum())
                        with col3:
                            st.metric("High Priority", 
                                    (batch_predictions['Maintenance_Urgency'] == 'HIGH').sum())
                        with col4:
                            st.metric("Average Time to Failure", 
                                    f"{batch_predictions['Time_to_Failure_Minutes'].mean():.1f} min")
                    else:
                        st.error(f"CSV must contain these columns: {', '.join(required_cols)}")
            except Exception as e:
                st.error(f"Error processing file: {str(e)}")

def get_maintenance_recommendation(urgency, time_to_failure, failure_prob):
    """Get maintenance recommendation based on prediction"""
    if urgency == "CRITICAL":
        return "**IMMEDIATE ACTION REQUIRED**: Stop the machine immediately and perform maintenance. Failure is imminent."
    elif urgency == "HIGH":
        if time_to_failure < 60:
            return f"**URGENT**: Schedule maintenance within {int(time_to_failure/60)} hour(s) or immediately if possible. The machine shows high risk of failure."
        else:
            return f"**Schedule maintenance within {int(time_to_failure/60)} hours**. The machine shows high risk of failure."
    elif urgency == "MEDIUM":
        if time_to_failure < 20:
            return f"**Schedule maintenance within the next 20-30 minutes**. Monitor the machine very closely. Estimated time to maintenance: {time_to_failure:.0f} minutes."
        elif time_to_failure < 60:
            return f"**Schedule maintenance within the next hour**. Monitor the machine closely. Estimated time to maintenance: {time_to_failure:.0f} minutes."
        elif time_to_failure < 120:
            return f"**Plan maintenance within the next 2 hours**. Monitor the machine closely. Estimated time to maintenance: {int(time_to_failure/60)} hours."
        else:
            return f"**Plan maintenance within the next few days**. Monitor the machine regularly. Estimated time to maintenance: {int(time_to_failure/60)} hours."
    else:
        if time_to_failure < 120:
            return f"**Monitor regularly**. No immediate action needed, but plan maintenance soon. Estimated time to maintenance: {int(time_to_failure/60)} hours."
        else:
            return f"**No immediate action needed**. Continue regular monitoring. Estimated time to maintenance: {int(time_to_failure/60)} hours."

def show_conclusion():
    """Conclusion page"""
    st.markdown('<h2 class="sub-header">📝 Conclusion & Key Takeaways</h2>', unsafe_allow_html=True)
    
    st.markdown("""
    ### Project Summary
    
    This predictive maintenance project successfully analyzed the AI4I 2020 dataset and built 
    a machine learning model to predict machine failures and estimate maintenance needs.
    
    ### Key Findings
    
    1. **Dataset Characteristics**:
       - The dataset contains 10,000 machine records with 14 features
       - Machine failure rate is approximately 3.39% (imbalanced dataset)
       - Five different failure types were identified: TWF, HDF, PWF, OSF, and RNF
    
    2. **Important Features**:
       - Tool wear is a critical indicator of machine health
       - Temperature difference between process and air temperature correlates with failures
       - Machine type (L = Low Quality/Load, M = Medium Quality/Load, H = High Quality/Load) affects failure rates differently
       - Rotational speed and torque relationships are important predictors
    
    3. **Model Performance**:
       - Random Forest Classifier achieved good performance on the imbalanced dataset
       - The model can effectively predict machine failures
       - Feature importance analysis revealed tool wear and temperature as key predictors
    
    4. **Maintenance Insights**:
       - Machines with tool wear > 100 minutes are at higher risk
       - Temperature differences > 10K indicate potential heat dissipation issues
       - Early detection can prevent costly downtime
    
    ### Applications
    
    This system can be used in real-world industrial settings to:
    - **Prevent unexpected failures** by predicting maintenance needs
    - **Optimize maintenance schedules** based on actual machine conditions
    - **Reduce downtime** through proactive maintenance
    - **Save costs** by avoiding catastrophic failures
    
    ### Future Improvements
    
    1. **Model Enhancement**:
       - Try ensemble methods (XGBoost, LightGBM)
       - Implement time-series analysis for sequential data
       - Add anomaly detection algorithms
    
    2. **Feature Engineering**:
       - Create more domain-specific features
       - Include historical maintenance records
       - Add environmental factors
    
    3. **System Integration**:
       - Real-time data streaming
       - Integration with IoT sensors
       - Automated alert system
    
    ### Technical Stack
    
    - **Data Analysis**: Pandas, NumPy
    - **Visualization**: Matplotlib, Seaborn, Plotly
    - **Machine Learning**: Scikit-learn (Random Forest)
    - **Web Application**: Streamlit
    - **Preprocessing**: StandardScaler, LabelEncoder
    
    ### Acknowledgments
    
    Dataset: AI4I 2020 Predictive Maintenance Dataset
    
    ---
    
    **Project completed for Introduction to Data Science Course (IDS F24)**
    """)
    
    st.markdown("---")
    st.markdown("### Thank you for using the Predictive Maintenance System! 🔧")

if __name__ == "__main__":
    main()