import streamlit as st
import numpy as np
import pandas as pd
import logging
from data.data_loader import load_data, preprocess_data
from visualization.plots import generate_shap_plots, plot_feature_importance
from models.model import train_model, predict
from models.evaluation import calculate_metrics

def setup_logging():
    """Configure logging settings"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(levelname)s - %(message)s'
    )

def is_valid_number(value):
    """Validate if input string is a valid number"""
    try:
        float(value)
        return True
    except ValueError:
        return False

def main():
    st.set_page_config(
        page_title="30-Day Readmission Risk Predictor",
        layout="wide"
    )
    
    setup_logging()
    
    st.title("30-Day Mental Health Hospital Readmission Risk Predictor")
    
    st.write("""
    This application predicts the risk of patient readmission within 30 days
    of discharge from a mental health hospital using machine learning.
    """)
    
    try:
        # Load and preprocess data
        data = load_data()
        if data is None:
            st.error("Error loading data")
            return
            
        X_train, X_test, y_train, y_test, feature_names, scaler = preprocess_data(data)
        if X_train is None:
            st.error("Error preprocessing data")
            return
        
        # Ensure X_train and X_test are DataFrames with feature names
        X_train_df = pd.DataFrame(X_train, columns=feature_names)
        X_test_df = pd.DataFrame(X_test, columns=feature_names)
        
        # Train model
        model = train_model(X_train_df, y_train)
        if model is None:
            st.error("Error training model")
            return
            
        # Make predictions
        y_pred, y_pred_proba = predict(model, X_test_df)
        if y_pred is None:
            st.error("Error making predictions")
            return
            
        # Calculate and display metrics
        metrics = calculate_metrics(y_test, y_pred, y_pred_proba)
        if metrics:
            st.header("Model Performance Metrics")
            col1, col2, col3, col4, col5 = st.columns(5)
            with col1:
                st.metric("Accuracy", f"{metrics['accuracy']:.2f}")
            with col2:
                st.metric("Precision", f"{metrics['precision']:.2f}")
            with col3:
                st.metric("Recall", f"{metrics['recall']:.2f}")
            with col4:
                st.metric("F1 Score", f"{metrics['f1']:.2f}")
            with col5:
                st.metric("AUC-ROC", f"{metrics['auc_roc']:.2f}")
        
        # Define categorical features and their options
        categorical_features = {
            'age': ['18-25', '26-35', '36-45', '46-55', '56-65', '65+'],
            'length_of_stay': ['1-3 days', '4-7 days', '8-14 days', '15+ days']
        }

        # Define numerical features and their ranges
        numerical_features = {
            'previous_admissions': {'min': 0, 'max': 100},
            'num_procedures': {'min': 0, 'max': 50},
            'num_medications': {'min': 0, 'max': 100},
            'num_diagnoses': {'min': 0, 'max': 50},
            'num_lab_procedures': {'min': 0, 'max': 200},
            'num_outpatient': {'min': 0, 'max': 100},
            'num_emergency': {'min': 0, 'max': 100},
            'num_inpatient': {'min': 0, 'max': 100}
        }

        # Interactive Prediction
        st.header("Predict Readmission Risk")
        st.write("Enter patient information to predict readmission risk")
        
        # Create columns for input fields
        col1, col2 = st.columns(2)
        
        # Store all input values
        input_values = {}
        
        # First column of inputs
        with col1:
            # Handle categorical features
            for feature in categorical_features:
                input_values[feature] = st.selectbox(
                    f"{feature.replace('_', ' ').title()}",
                    options=categorical_features[feature],
                    key=f"select_{feature}"
                )
            
            # Handle first half of numerical features
            numerical_features_list = list(numerical_features.keys())
            for feature in numerical_features_list[:len(numerical_features_list)//2]:
                input_text = st.text_input(
                    f"{feature.replace('_', ' ').title()}",
                    value="0",
                    key=f"input_{feature}"
                )
                
                if not is_valid_number(input_text):
                    st.error(f"Please enter a valid number for {feature}")
                    input_values[feature] = 0
                else:
                    value = float(input_text)
                    if value < numerical_features[feature]['min'] or value > numerical_features[feature]['max']:
                        st.warning(f"Value should be between {numerical_features[feature]['min']} and {numerical_features[feature]['max']}")
                    input_values[feature] = value

        # Second column of inputs
        with col2:
            # Handle second half of numerical features
            for feature in numerical_features_list[len(numerical_features_list)//2:]:
                input_text = st.text_input(
                    f"{feature.replace('_', ' ').title()}",
                    value="0",
                    key=f"input_{feature}"
                )
                
                if not is_valid_number(input_text):
                    st.error(f"Please enter a valid number for {feature}")
                    input_values[feature] = 0
                else:
                    value = float(input_text)
                    if value < numerical_features[feature]['min'] or value > numerical_features[feature]['max']:
                        st.warning(f"Value should be between {numerical_features[feature]['min']} and {numerical_features[feature]['max']}")
                    input_values[feature] = value

        # Prediction button
        if st.button("Predict"):
            # Convert categorical inputs to numerical
            input_data = input_values.copy()
            
            # Process age categories
            age_mapping = {'18-25': 21.5, '26-35': 30.5, '36-45': 40.5, 
                          '46-55': 50.5, '56-65': 60.5, '65+': 70}
            input_data['age'] = age_mapping[input_data['age']]
            
            # Process length of stay categories
            los_mapping = {'1-3 days': 2, '4-7 days': 5.5, '8-14 days': 11, 
                          '15+ days': 15}
            input_data['length_of_stay'] = los_mapping[input_data['length_of_stay']]
            
            # Create DataFrame and make prediction
            input_df = pd.DataFrame([input_data], columns=feature_names)
            input_scaled = pd.DataFrame(
                scaler.transform(input_df),
                columns=feature_names
            )
            
            _, probabilities = predict(model, input_scaled)
            
            if probabilities is not None:
                prediction = probabilities[0]
                st.write(f"Predicted probability of readmission: {prediction:.2%}")
                
                risk_category = (
                    "High" if prediction > 0.7
                    else "Medium" if prediction > 0.3
                    else "Low"
                )
                st.write(f"Risk Category: {risk_category}")
        
        # Model Insights Section
        st.header("Model Insights")
        
        # Feature Importance Plot
        st.subheader("Feature Importance")
        st.write("""
        This plot shows the relative importance of each feature in making predictions.
        Higher values indicate more important features.
        """)
        fig_importance = plot_feature_importance(model, X_train_df)
        if fig_importance:
            st.plotly_chart(fig_importance)
        
        # SHAP Analysis
        st.subheader("SHAP Analysis")
        st.write("""
        SHAP (SHapley Additive exPlanations) values show how each feature
        contributes to predictions. Features in red increase the prediction,
        while features in blue decrease it.
        """)
        
        # Create and display SHAP plot with custom width
        fig_shap = generate_shap_plots(model, X_test_df)
        if fig_shap:
            # Use a container with custom width
            with st.container():
                st.pyplot(fig_shap, use_container_width=False)
        else:
            st.error("Error generating SHAP plot")
        
    except Exception as e:
        st.error(f"An error occurred: {str(e)}")
        logging.error(f"Error in main: {str(e)}")

if __name__ == "__main__":
    main()