from flask import Flask, request, jsonify
import joblib
import numpy as np
import os


from huggingface_hub import hf_hub_download
import joblib

# Load the trained model
model = joblib.load("model_compressed.joblib")


# Initialize app
app = Flask(__name__)

@app.route("/")
def home():
    return jsonify({"message": "Supermarket Revenue Prediction API is running!"})

@app.route("/predict", methods=["POST"])
def predict():
    try:
        data = request.get_json(force=True)
        features = np.array(data["features"])

        # Case 1: single row
        if features.ndim == 1:
            features = features.reshape(1, -1)
            prediction = model.predict(features)[0]
            return jsonify({"predicted_revenue": float(prediction)})

        # Case 2: multiple rows
        else:
            predictions = model.predict(features).tolist()
            return jsonify({"predictions": predictions})

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


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860, debug=True)