from sklearn.base import BaseEstimator, TransformerMixin

class DatetimeDummyExtractor(BaseEstimator, TransformerMixin):
    """
    Extracts comprehensive date/time features from a datetime column.
    'month' and 'weekday' can optionally be returned as categorical strings.

    Parameters
    ----------
    datetime_column : str
        Name of the datetime column to extract features from.
    features : list of str, optional
        List of features to extract. Possible values:
        'year', 'month', 'day', 'weekday', 'hour', 'minute', 'second',
        'quarter', 'dayofyear', 'is_weekend'.
        Default is all features.
    month_as_category : bool, default=True
        If True, month numbers are converted to string names (e.g., 'Jan', 'Feb').
    weekday_as_category : bool, default=True
        If True, weekdays are converted to string names (e.g., 'Mon', 'Tue').
    """

    ALL_FEATURES = ['year', 'month', 'day', 'weekday', 'hour', 'minute', 'second',
                    'quarter', 'dayofyear', 'is_weekend']

    MONTH_NAMES = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun', 'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
    WEEKDAY_NAMES = ['Mon', 'Tue', 'Wed', 'Thu', 'Fri', 'Sat', 'Sun']

    def __init__(self, datetime_column, features=None, month_as_category=True, weekday_as_category=True):
        self.datetime_column = datetime_column
        self.features = features or self.ALL_FEATURES
        self.month_as_category = month_as_category
        self.weekday_as_category = weekday_as_category

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

    def transform(self, X):
        X = X.copy()
        dt = X[self.datetime_column]

        if 'year' in self.features:
            X[f'{self.datetime_column}_year'] = dt.dt.year
        if 'month' in self.features:
            if self.month_as_category:
                X[f'{self.datetime_column}_month'] = dt.dt.month.apply(lambda x: self.MONTH_NAMES[x - 1])
            else:
                X[f'{self.datetime_column}_month'] = dt.dt.month
        if 'day' in self.features:
            X[f'{self.datetime_column}_day'] = dt.dt.day
        if 'weekday' in self.features:
            if self.weekday_as_category:
                X[f'{self.datetime_column}_weekday'] = dt.dt.weekday.apply(lambda x: self.WEEKDAY_NAMES[x])
            else:
                X[f'{self.datetime_column}_weekday'] = dt.dt.weekday
        if 'hour' in self.features:
            X[f'{self.datetime_column}_hour'] = dt.dt.hour
        if 'minute' in self.features:
            X[f'{self.datetime_column}_minute'] = dt.dt.minute
        if 'second' in self.features:
            X[f'{self.datetime_column}_second'] = dt.dt.second
        if 'quarter' in self.features:
            X[f'{self.datetime_column}_quarter'] = dt.dt.quarter
        if 'dayofyear' in self.features:
            X[f'{self.datetime_column}_dayofyear'] = dt.dt.dayofyear
        if 'is_weekend' in self.features:
            X[f'{self.datetime_column}_is_weekend'] = dt.dt.weekday >= 5

        return X
