from flask import Flask, request, jsonify
import joblib
import pandas as pd
import numpy as np

app = Flask(__name__)

# Load the preprocessor and model
try:
    # Adjust these paths if your model and preprocessor are in a different location relative to app.py
    model = joblib.load('best_decision_tree_model.joblib')
    preprocessor = joblib.load('preprocessor.joblib')
    print("Model and Preprocessor loaded successfully!")
except Exception as e:
    print(f"Error loading model or preprocessor: {e}")
    model = None
    preprocessor = None

@app.route('/predict', methods=['POST'])
def predict():
    if model is None or preprocessor is None:
        return jsonify({'error': 'Model or preprocessor not loaded'}), 500

    try:
        data = request.get_json(force=True)

        # Log received data for debugging
        print(f"Received data: {data}")

        # Ensure the input data has the expected format (e.g., list of dictionaries)
        if not isinstance(data, list):
            data = [data]

        # Convert input data to a Pandas DataFrame
        input_df = pd.DataFrame(data)

        # Define expected columns based on your preprocessing setup
        # These should match the columns BEFORE one-hot encoding for categorical features
        # and direct numerical features.
        expected_features = ['Age', 'Gender', 'CityTier'] # Adjust if you have more features

        # Reorder columns to match the training data order if necessary
        # This assumes all expected_features are present in input_df
        input_df = input_df[expected_features]

        # Preprocess the input data
        # The preprocessor expects a DataFrame with raw features
        X_processed = preprocessor.transform(input_df)

        # Convert the processed data back to DataFrame for consistency if needed,
        # but sklearn models generally accept numpy arrays.
        # Optionally, if you need feature names after preprocessing:
        # categorical_features = ['Gender', 'CityTier'] # These should be defined here or passed
        # numerical_features = ['Age'] # These should be defined here or passed
        # categorical_feature_names = preprocessor.named_transformers_['cat'].get_feature_names_out(categorical_features)
        # all_feature_names = list(categorical_feature_names) + numerical_features
        # X_processed_df = pd.DataFrame(X_processed, columns=all_feature_names)

        predictions = model.predict(X_processed)
        prediction_proba = model.predict_proba(X_processed)[:, 1] # Probability of the positive class

        results = []
        for i in range(len(predictions)):
            results.append({
                'prediction': int(predictions[i]),
                'probability': float(prediction_proba[i])
            })

        return jsonify(results)

    except Exception as e:
        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    # In a production environment, you might use a production-ready WSGI server like Gunicorn
    # For local development, this is fine.
    app.run(host='0.0.0.0', port=5000, debug=True)