from modules import DataCleaner, NaNImputerWrapper, FeatureEng, AdjustedProbClassifier
import gradio as gr
import joblib
import numpy as np
import pandas as pd
from xgboost import XGBClassifier

# load the ML pipeline
pipeline = joblib.load(r'full_xgb_pipeline_adjusted_prob.joblib')

feedback_options = [
    'Products always in Stock', 'Quality Customer Care', 'Reasonable Price',
    'User Friendly Website', 'No Reason Specified', 'Poor Website',
    'Poor Customer Service', 'Poor Product Quality', 'Too many ads'
]

membership_options = [
    'No Membership', 'Basic Membership', 'Silver Membership',
    'Gold Membership', 'Platinum Membership', 'Premium Membership'
]

def predict(points_in_wallet, feedback, membership_category):
    try:
        # convert points_in_wallet to float and validate
        points_in_wallet = float(points_in_wallet)
        if points_in_wallet < 0:
            return "Error: points_in_wallet cannot be negative."
        
        X_df = pd.DataFrame({
            'points_in_wallet': [points_in_wallet],
            'feedback': [feedback],
            'membership_category': [membership_category]
        })
        
        prediction = pipeline.predict(X_df)
        return f"Churn Risk Score: {prediction[0]}"
    except Exception as e:
        return f"Prediction error: {str(e)}"

# Gradio Interface
demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Number(label="Points in Wallet", value=0, precision=2),
        gr.Dropdown(choices=feedback_options, label="Feedback Category"),
        gr.Dropdown(choices=membership_options, label="Membership Category")
    ],
    outputs=gr.Textbox(label="Prediction"),
    title="Customer Churn Risk Score Predictor",
    description="Enter wallet points, select feedback and membership to get a prediction."
)

if __name__ == "__main__":
    demo.launch()