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

# Initialize the Flask application
super_kart_api = Flask("Super Kart Price Predictor")

# Load the trained machine learning model
model_path = "super_kart_model_v1_0.joblib"

try:
    model = joblib.load(model_path)
    print(f"Model loaded successfully from {model_path}")
except FileNotFoundError:
    raise FileNotFoundError(f"Model file not found at {model_path}. Ensure it's uploaded to the repo root.")

# Expected feature names from the model (adjust if your training columns differ)
EXPECTED_COLUMNS = [
    'Product_Type_Baking Goods', 'Product_Type_Breads', 'Product_Type_Breakfast', 'Product_Type_Canned',
    'Product_Type_Dairy', 'Product_Type_Frozen Foods', 'Product_Type_Fruits and Vegetables',
    'Product_Type_Hard Drinks', 'Product_Type_Health and Hygiene', 'Product_Type_Household',
    'Product_Type_Meat', 'Product_Type_Others', 'Product_Type_Seafood', 'Product_Type_Snack Foods',
    'Product_Type_Soft Drinks', 'Product_Type_Starchy Foods', 'Store_Type_Departmental Store',
    'Store_Type_Food Mart', 'Store_Type_Supermarket Type1', 'Store_Type_Supermarket Type2',
    'Product_Sugar_Content', 'Store_Size', 'Store_Location_City_Type', 'Product_Weight',
    'Product_Allocated_Area', 'Product_MRP', 'Store_Establishment_Year'
]

# Define a route for the home page (GET request)
@super_kart_api.get('/')
def home():
    return "Welcome to the Super Kart Price Prediction API!"

# Define an endpoint for single product sales prediction (POST request)
@super_kart_api.post('/v1/sales')
def predict_sales():
    input_data = request.get_json()
    
    sample = {
        'Product_Weight': input_data['Product_Weight'],
        'Product_Sugar_Content': input_data['Product_Sugar_Content'],
        'Product_Allocated_Area': input_data['Product_Allocated_Area'],
        'Product_Type': input_data['Product_Type'],
        'Product_MRP': input_data['Product_MRP'],
        'Store_Establishment_Year': input_data['Store_Establishment_Year'],
        'Store_Size': input_data['Store_Size'],
        'Store_Location_City_Type': input_data['Store_Location_City_Type'],
        'Store_Type': input_data['Store_Type']
    }
    
    features_df = pd.DataFrame([sample])
    
    # Apply one-hot encoding
    features_df = pd.get_dummies(features_df, columns=['Product_Type', 'Store_Type'], drop_first=True)
    
    # Apply ordinal encoding
    sugar_mapping = {'No Sugar': 0, 'Low Sugar': 1, 'Regular': 2}
    size_mapping = {'Small': 0, 'Medium': 1, 'High': 2}
    city_mapping = {'Tier 3': 0, 'Tier 2': 1, 'Tier 1': 2}
    
    features_df['Product_Sugar_Content'] = features_df['Product_Sugar_Content'].map(sugar_mapping)
    features_df['Store_Size'] = features_df['Store_Size'].map(size_mapping)
    features_df['Store_Location_City_Type'] = features_df['Store_Location_City_Type'].map(city_mapping)
    
    # Align with expected columns (add missing as 0, drop extras)
    features_df = features_df.reindex(columns=EXPECTED_COLUMNS, fill_value=0)
    
    # Make prediction
    predicted_sales = model.predict(features_df)[0]
    predicted_sales = round(float(predicted_sales), 2)
    
    return jsonify({'Predicted Sales Total (in dollars)': predicted_sales})


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