# FeatureEngineering.py
from sklearn.base import BaseEstimator, TransformerMixin
import numpy as np
import pandas as pd

class FeatureEngineering(BaseEstimator, TransformerMixin):
    def __init__(self):
        pass

    def fit(self, X, y=None):
        return self

    def transform(self, X):
        X_ = X.copy()
        X_['StaffRatio'] = X_['StaffOnline'] / X_['StaffEmployed']
        X_['TotalArea'] = X_['StoreArea'] + X_['PickingArea']
        X_['Year'] = X_['Date'].dt.year
        X_['Month'] = X_['Date'].dt.month
        X_['Weekday'] = X_['Date'].dt.weekday
        X_['SlotHour'] = X_['Slot'].str.split(":").str[0].astype(int)

        # Keep categorical SpecialEvent as-is, and also add an indicator if needed
        X_['IsSpecialEvent'] = X_['SpecialEvent'].apply(lambda x: 0 if pd.isna(x) or x == "" else 1)

        return X_