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

sales_revenue_predictor_api = Flask("Superkart sales Revenue Predictor")

# Load trained model
model = joblib.load("Sales_revenue_prediction_model_v1_0.joblib")

# Define the expected columns (based on your dataset)
EXPECTED_COLUMNS = [
    'Product_Id', 'Product_Weight', 'Product_Sugar_Content',
    'Product_Allocated_Area', 'Product_Type', 'Product_MRP',
    'Store_Id', 'Store_Establishment_Year', 'Store_Size',
    'Store_Location_City_Type', 'Store_Type'
]

@sales_revenue_predictor_api.get('/')
def home():
    return "Welcome to the SuperKart Sales Prediction API!"

@sales_revenue_predictor_api.post('/v1/sales')
def predict_sales():
    sales_data = request.get_json()

    # Manually build DataFrame with missing/default values
    input_data = pd.DataFrame([sales_data])

    # Add missing expected columns if any
    missing_cols = set(EXPECTED_COLUMNS) - set(input_data.columns)
    if missing_cols:
        return jsonify({"error": f"columns are missing: {missing_cols}"}), 400

    # Predict
    prediction = model.predict(input_data)[0]
    return jsonify({"Sales": round(float(prediction), 2)})

@sales_revenue_predictor_api.post('/v1/salesbatch')
def predict_sales_batch():
    file = request.files['file']
    input_data = pd.read_csv(file)

    # Check for missing columns
    missing_cols = set(EXPECTED_COLUMNS) - set(input_data.columns)
    if missing_cols:
        return jsonify({"error": f"columns are missing: {missing_cols}"}), 400

    # Predict
    predictions = model.predict(input_data).tolist()
    predictions = [round(float(p), 2) for p in predictions]

    # Use ID column or row index
    ids = input_data['Product_Id'].tolist() if 'Product_Id' in input_data.columns else list(range(1, len(predictions) + 1))

    return jsonify(dict(zip(ids, predictions)))

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