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

app = Flask(__name__)

# Load the trained model and scaler
model = joblib.load('xgboost_model.joblib')
scaler = joblib.load('scaler.joblib')

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


#@app.route('/predict', methods=['POST'])
#def predict():
#    try:
#        json_ = request.json
#        query_df = pd.DataFrame(json_)

        # Ensure the order of columns matches the training data
        # This assumes the order of columns in X_train is known or can be derived
        # For this example, we'll assume X_train is available globally or its columns can be recreated
        # In a real deployment, you'd save/load the column order or use a more robust preprocessing pipeline
        
        # --- Recreate df_encoded column order and apply transformations ---
        # This is a placeholder. In a production system, you'd save the columns of X_train
        # and ensure incoming data matches that structure after one-hot encoding.

        # Example of creating dummy columns for categorical features if not present in the query_df
        # and then reordering. This needs to be robust for all possible categories.

        # Ensure all columns expected by the model are present, fill missing with 0 for one-hot encoded features
        # and apply scaling to numerical features.
        
        # Numerical features that were scaled
#        numerical_features_to_scale = ['Product_Weight', 'Product_Allocated_Area', 'Product_MRP', 'Store_Age', 'Product_Visibility']

        # Categorical features that were one-hot encoded
#        ohe_columns = ['Product_Sugar_Content', 'Product_Type', 'Store_Id', 'Store_Size', 'Store_Location_City_Type', 'Store_Type']

        # Need to re-create the full list of columns that were used for training
        # This is crucial for consistent prediction. Ideally, you save X.columns during training.
        # For demonstration, let's create a simplified version based on `df_encoded` structure from notebook state.
        # NOTE: This part needs to be precise based on your exact `df_encoded` columns. 
        # A more robust solution involves saving the `X_train.columns` list during model training.
        
        # Assuming X_train's columns are known from previous steps (from `df_encoded` before dropping target)
        # For a more robust solution, save X.columns to a file and load it here.
        # In this specific context, the notebook state shows X as a global variable, so we can use its columns.
#        expected_columns = model.feature_names_in_ if hasattr(model, 'feature_names_in_') else X.columns
        
        # Apply scaler to numerical features in query_df
#        query_df[numerical_features_to_scale] = scaler.transform(query_df[numerical_features_to_scale])
        
        # Handle categorical features using one-hot encoding, similar to training
#       query_df_processed = pd.get_dummies(query_df, columns=ohe_columns, drop_first=True)

        # Align columns - crucial for production deployment
        # Add missing columns with 0 and remove extra ones
#        missing_cols = set(expected_columns) - set(query_df_processed.columns)
#        for c in missing_cols:
#            query_df_processed[c] = 0
#        query_df_processed = query_df_processed[expected_columns]

#        prediction = model.predict(query_df_processed)
#        return jsonify({'prediction': list(prediction)}) 

#    except Exception as e:
#        return jsonify({'error': str(e)}), 400

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)