import gradio as gr
import pandas as pd
import pickle
import numpy as np

def load_model():
    """Load the trained sklearn model from disk"""
    with open('model.pkl', 'rb') as f:
        model = pickle.load(f)
    return model

def process_input(df):
    """Process input dataframe to match model requirements"""
    # Add any necessary preprocessing steps here
    # For example: handling missing values, scaling, encoding
    return df

def predict_mortality(csv_file):
    try:
        # Read the CSV file
        df = pd.read_csv(csv_file.name)
        
        # Load the model
        model = load_model()
        
        # Process the input data
        processed_df = process_input(df)
        
        # Make predictions
        predictions = model.predict(processed_df)
        probabilities = model.predict_proba(processed_df)
        
        # Add predictions to the original dataframe
        df['Mortality_Risk'] = predictions
        df['Death_Probability'] = probabilities[:, 1]  # Assuming 1 is the positive class
        
        # Format the results
        results_df = df.copy()
        results_df['Mortality_Risk'] = results_df['Mortality_Risk'].map({1: 'High Risk', 0: 'Low Risk'})
        results_df['Death_Probability'] = results_df['Death_Probability'].round(3)
        
        return results_df
    
    except Exception as e:
        return f"Error processing file: {str(e)}"

# Create the Gradio interface
iface = gr.Interface(
    fn=predict_mortality,
    inputs=gr.File(label="Upload Patient Data (CSV)"),
    outputs=gr.Dataframe(label="Prediction Results"),
    title="Patient Mortality Risk Prediction",
    description="""
    Upload a CSV file containing patient data to predict mortality risk.
    The model will return the original data with two additional columns:
    - Mortality_Risk: High Risk or Low Risk classification
    - Death_Probability: Probability of death (0-1)
    """
)

if __name__ == "__main__":
    iface.launch(share=True)