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

app = Flask(__name__)

# Load the pre-trained model and scaler
try:
    kmeans_model = joblib.load("kmeans_model.joblib")
    scaler_model = joblib.load("scaler_model.joblib")
    print("Models loaded successfully!")
except Exception as e:
    print(f"Error loading models: {e}")
    kmeans_model = None
    scaler_model = None

@app.route("/predict", methods=["POST"])
def predict():
    if kmeans_model is None or scaler_model is None:
        return jsonify({"error": "Model not loaded. Please check deployment logs."}), 500

    data = request.get_json(force=True)
    
    try:
        # Extract features and ensure order
        age = data.get("age")
        annual_income = data.get("annual_income")
        spending_score = data.get("spending_score")

        if age is None or annual_income is None or spending_score is None:
            return jsonify({"error": "Missing required input features (age, annual_income, spending_score)."}), 400

        # Create a DataFrame for scaling
        features_df = pd.DataFrame([[age, annual_income, spending_score]],
                                   columns=['Age', 'Annual Income (k$)', 'Spending Score (1-100)'])

        # Scale the input features using the loaded scaler
        scaled_features = scaler_model.transform(features_df)

        # Predict the cluster
        prediction = kmeans_model.predict(scaled_features)
        cluster_id = int(prediction[0]) # Convert numpy int to Python int

        return jsonify({"cluster_id": cluster_id})

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(debug=True) # debug=True for local testing, set to False for deployment