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

# Initialize Flask app
sales_prediction_api = Flask(__name__)

# 👇 REQUIRED for Hugging Face Gunicorn
application = sales_prediction_api

# Load models
dt_model = joblib.load("decision_tree_model.pkl")
xgb_model = joblib.load("xgboost_model.pkl")

# Home route
@sales_prediction_api.route("/")
def home():
    return "✅ SuperKart Sales Prediction API is running"

# Prediction route
@sales_prediction_api.route("/predict", methods=["POST"])
def predict():
    data = request.get_json()

    # Model choice
    model_choice = data.get("model", "dt")

    # Extract features (MATCHES STREAMLIT KEYS)
    sample = {
        "Product_Weight": data["Product_Weight"],
        "Product_Sugar_Content": data["Product_Sugar_Content"],
        "Product_Allocated_Area": data["Product_Allocated_Area"],
        "Product_Type": data["Product_Type"],
        "Product_MRP": data["Product_MRP"],
        "Store_Size": data["Store_Size"],
        "Store_Location_City_Type": data["Store_Location_City_Type"],
        "Store_Type": data["Store_Type"],
        "Store_Age": data["Store_Age"]
    }

    sample_df = pd.DataFrame([sample])

    # Select model
    if model_choice == "dt":
        prediction = dt_model.predict(sample_df)[0]
    elif model_choice == "xgb":
        prediction = xgb_model.predict(sample_df)[0]
    else:
        return jsonify({"error": "Invalid model choice"}), 400

    return jsonify({
        "Prediction": float(prediction),
        "ModelUsed": model_choice
    })

if __name__ == '__main__':
    sales_prediction_api.run(debug=True)