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

# -------------------------------
# Initialize Flask App
# -------------------------------
sales_prediction_api = Flask('Superkart Sales Prediction')

# -------------------------------
# Load the trained Gradient Boosting model
# -------------------------------
model = joblib.load("gradient_tuned.joblib")

# -------------------------------
# Root endpoint (Health check)
# -------------------------------
@sales_prediction_api.get("/")
def home():
    """Simple health check endpoint"""
    return 'Welcome to the SuperKart Sales Prediction API'

# -------------------------------
# Single Prediction Endpoint
# -------------------------------
@sales_prediction_api.post("/v1/salesdata")
def predict_sales():
    """Predict sales for a single product/store record"""
    try:
        # Parse incoming JSON
        sales_data = request.get_json()

        # Prepare input dictionary with only the features used by the model
        sample = {
            'Product_Weight': sales_data['Product_Weight'],
            'Product_Allocated_Area': sales_data['Product_Allocated_Area'],
            'Product_MRP': sales_data['Product_MRP'],
            'Product_Sugar_Content': sales_data['Product_Sugar_Content'],
            'Product_Type': sales_data['Product_Type'],
            'Store_Id': sales_data['Store_Id'],
            'Store_Size': sales_data['Store_Size'],
            'Store_Location_City_Type': sales_data['Store_Location_City_Type'],
        }

        # Convert to DataFrame
        input_df = pd.DataFrame([sample])

        # Predict sales
        predicted_sales = model.predict(input_df)[0]
        predicted_sales = round(float(predicted_sales), 2)

        # Return JSON response
        return jsonify({"predicted_sales": predicted_sales})

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

# -------------------------------
# Batch Prediction Endpoint
# -------------------------------
@sales_prediction_api.post("/v1/salesdatabatch")
def predict_sales_batch():
    """Predict sales for multiple product/store records from a CSV file"""
    try:
        # Get uploaded CSV file
        file = request.files['file']
        input_df = pd.read_csv(file)

        # Keep only the features used by the model
        features = [
            'Product_Weight', 'Product_Allocated_Area', 'Product_MRP',
            'Product_Sugar_Content', 'Product_Type', 'Store_Id',
            'Store_Size', 'Store_Location_City_Type'
        ]
        input_features = input_df[features]

        # Predict sales
        predicted_sales = model.predict(input_features)
        predicted_sales = [round(float(sale), 2) for sale in predicted_sales]

        # Attach predictions to Product IDs if available
        if 'Product_Id' in input_df.columns:
            product_ids = input_df['Product_Id'].tolist()
            output_dict = dict(zip(product_ids, predicted_sales))
        else:
            output_dict = {"predicted_sales": predicted_sales}

        return jsonify(output_dict)

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

# -------------------------------
# Run Flask App
# -------------------------------
if __name__ == '__main__':
    sales_prediction_api.run(debug=True)