import gradio as gr
import pandas as pd
from joblib import load

# Load the model at startup
rf_model = load('churn_prediction_model.joblib')

def predict_churn(file):
    try:
        # Load the uploaded file
        if file.name.endswith('.csv'):
            test_data = pd.read_csv(file.name)
        else:
            test_data = pd.read_excel(file.name)

        # Ensure 'major_issue' is one-hot encoded if needed
        if 'major_issue_Technical Issue' not in test_data.columns:
            test_data = pd.get_dummies(test_data, columns=['major_issue'], drop_first=True)

        # Ensure all training-time columns are present
        required_columns = [
            'late_payments_last_year', 'missed_payments_last_year', 'plan_tenure',
            'num_employees', 'avg_monthly_contribution', 'annual_revenue',
            'support_calls_last_year', 'support_engagement_per_year',
            'major_issue_Technical Issue'
        ]
        for col in required_columns:
            if col not in test_data.columns:
                test_data[col] = 0

        # Extract features
        test_data_features = test_data[required_columns]

        # Make predictions
        test_data['predicted_churn'] = rf_model.predict(test_data_features)

        # Select output columns
        output_columns = [
            'customer_id', 'late_payments_last_year', 'missed_payments_last_year',
            'plan_tenure', 'num_employees', 'avg_monthly_contribution',
            'annual_revenue', 'support_calls_last_year', 'support_engagement_per_year',
            'major_issue_Technical Issue', 'predicted_churn'
        ]
        output_data = test_data[output_columns]

        # Save to temporary file
        temp_output = "output.csv"
        output_data.to_csv(temp_output, index=False)
        
        return temp_output

    except Exception as e:
        raise gr.Error(f"Error processing file: {str(e)}")

# Create Gradio interface
iface = gr.Interface(
    fn=predict_churn,
    inputs=gr.File(
        type="file",
        label="Upload Customer Data File",
        file_types=[".csv", ".xlsx", ".xls"]
    ),
    outputs=gr.File(label="Download Predictions"),
    title="Customer Churn Prediction",
    description="Upload your customer data file to predict churn probability",
    cache_examples=False
)

# Launch the app
if __name__ == "__main__":
    iface.launch()