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

# Initialize Flask app
app = Flask("Sales Forecaster")

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

# Root endpoint
@app.get('/')
def home():
    return "Welcome to the Forecasting Model API!"

# Single prediction endpoint
@app.post('/v1/params')
#@app.route('/v1/params', methods=['POST'])
def predict_sales():
    # Get JSON data from the request
    input_dict = request.get_json()

    # Extract features
    sample = {
        'Product_Family': input_dict['Product_Family'],
        'Product_Weight': input_dict['Product_Weight'],
        'Product_Sugar_Content': input_dict['Product_Sugar_Content'],
        'Product_Allocated_Area': input_dict['Product_Allocated_Area'],
        'Product_Type': input_dict['Product_Type'],
        'Product_MRP': input_dict['Product_MRP'],
        'years_of_operation': input_dict['years_of_operation'],
        'Store_Size': input_dict['Store_Size'],
        'Store_Location_City_Type': input_dict['Store_Location_City_Type'],
        'Store_Type': input_dict['Store_Type']
    }

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

    # Make prediction
    prediction = model.predict(input_data).tolist()[0]

    return jsonify({'Prediction': prediction})

# Batch prediction endpoint
@app.post('/v1/paramsbatch')
#@app.route('/v1/paramsbatch', methods=['POST'])
def predict_sales_batch():
    # Get uploaded CSV file
    file = request.files['file']
    input_data = pd.read_csv(file)

    # Make predictions
    predictions = model.predict(input_data).tolist()

    # If CustomerId exists, return mapping
    if 'CustomerId' in input_data.columns:
        cust_id_list = input_data['CustomerId'].values.tolist()
        output_dict = dict(zip(cust_id_list, predictions))
        return jsonify(output_dict)

    # Otherwise just return predictions as list
    return jsonify({'Predictions': predictions})

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