import pandas as pd
import pytz
from sklearn.base import BaseEstimator, TransformerMixin


class ValueTypeCastDatetimeTransformer(BaseEstimator, TransformerMixin):
    """
    Transformer to cast a column to datetime, with optional custom format and optional per-row timezone conversion.

    Parameters
    ----------
    column : str
        Name of the column to cast to datetime.
    datetime_format : str, optional
        Optional datetime format for parsing.
    tz_column : str, optional
        Name of the column containing the timezone strings (e.g., 'Europe/Paris').
        If provided, the datetime will be converted to the corresponding local timezone per row.
    """

    def __init__(self, column, datetime_format=None, tz_column=None):
        self.column = column
        self.datetime_format = datetime_format
        self.tz_column = tz_column

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

    def transform(self, X):
        X = X.copy()

        X[self.column] = pd.to_datetime(X[self.column], format=self.datetime_format, errors="raise")

        if self.tz_column is not None:
            X[self.column] = X[self.column].dt.tz_localize('UTC')

            for tz in X[self.tz_column].unique():
                mask = X[self.tz_column] == tz
                X.loc[mask, self.column] = X.loc[mask, self.column].dt.tz_convert(pytz.timezone(tz))

        return X
