from __future__ import annotations

from collections import defaultdict
from copy import deepcopy
from typing import DefaultDict, Optional

import numpy as np
import pandas as pd
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn.preprocessing import LabelEncoder


class CustomLabelEncoder(BaseEstimator, TransformerMixin):
    def __init__(self, col_names: list[str]) -> None:
        self.col_names = col_names
        self.col2encoder: DefaultDict[str, LabelEncoder] = defaultdict(
            lambda: LabelEncoder()
        )

    def fit(
        self, X: pd.DataFrame, y: Optional[np.ndarray] = None
    ) -> CustomLabelEncoder:
        for col in self.col_names:
            encoder = self.col2encoder[col]
            encoder.fit(X[col])
        return self

    def transform(
        self, X: pd.DataFrame, y: Optional[np.ndarray] = None
    ) -> pd.DataFrame:
        X = deepcopy(X)

        for col in self.col_names:
            encoder = self.col2encoder[col]
            X[col] = encoder.transform(X[col]) + 1

        return X

    def inverse_transform(
        self, X: pd.DataFrame, y: Optional[np.ndarray] = None
    ) -> pd.DataFrame:
        X = deepcopy(X)

        for col in self.col_names:
            encoder = self.col2encoder[col]
            X[col] = encoder.inverse_transform(X[col] - 1)

        return X
