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

app = Flask('SuperKart Sales Predictor')

# Load the saved model
saved_model = joblib.load("final_xgb_pipeline_model.joblib")

@app.get('/')
def home():
  return 'Welcome to the SuperKart Sales Predictor'

@app.post('/predict')
def predict_sales():
  store_data = request.get_json()

  sample = {
      'Product_Weight' :	store_data['Product_Weight'],
      'Product_Sugar_Content' : store_data['Product_Sugar_Content'],	
      'Product_Allocated_Area'	: store_data['Product_Allocated_Area'],
      'Product_Type'	: store_data['Product_Type'],
      'Product_MRP'	: store_data['Product_MRP'],
      'Store_Id'	: store_data['Store_Id'],
      'Store_Size' : store_data['Store_Size'],	
      'Store_Location_City_Type' : store_data['Store_Location_City_Type'],	
      'Store_Type' : 	store_data['Store_Type'],
      'Store_Age' : store_data['Store_Age']
  }

  input_data = pd.DataFrame([sample])

  predictions = saved_model.predict(input_data).tolist()[0]

  return jsonify({'prediction': predictions})

@app.post('/predict_batch')
def predict_batch():
  file = request.files['file']
  input_data = pd.read_csv(file)

  predictions = saved_model.predict(input_data).tolist()

  store_id_list = input_data.storeID.values.tolist()
  output_dict = dict(zip(store_id_list, predictions))

  return output_dict


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