from sklearn.compose import ColumnTransformer
from sklearn.base import TransformerMixin, BaseEstimator, clone
from sklearn.utils.metaestimators import _BaseComposition

from sklearn.utils.validation import check_array, check_is_fitted, _check_feature_names_in
from itertools import chain
import numpy as np

class StreamColumnTransformer(ColumnTransformer, TransformerMixin, _BaseComposition):

    _required_parameters = ["transformers"]

    def __init__(
        self,
        transformers,
        *,
        remainder="drop",
        sparse_threshold=0.3,
        n_jobs=None,
        transformer_weights=None,
        verbose=False,
        verbose_feature_names_out=True,
    ):
        self.transformers = transformers
        self.remainder = remainder
        self.sparse_threshold = sparse_threshold
        self.n_jobs = n_jobs
        self.transformer_weights = transformer_weights
        self.verbose = verbose
        self.verbose_feature_names_out = verbose_feature_names_out
        
    def partial_fit(self, X, y=None):
        self._check_n_features(X, reset=True)

        for _, transformer, col in self.transformers:
            transformer.partial_fit(X[col], y)
        
        self.transformers_ = self.transformers

    def get_feature_names_out(self, input_features=None):
        """Get output feature names for transformation.

        Parameters
        ----------
        input_features : array-like of str or None, default=None
            Input features.

            - If `input_features` is `None`, then `feature_names_in_` is
              used as feature names in. If `feature_names_in_` is not defined,
              then names are generated: `[x0, x1, ..., x(n_features_in_)]`.
            - If `input_features` is an array-like, then `input_features` must
              match `feature_names_in_` if `feature_names_in_` is defined.

        Returns
        -------
        feature_names_out : ndarray of str objects
            Transformed feature names.
        """
        input_features = _check_feature_names_in(self, input_features)

        # List of tuples (name, feature_names_out)
        transformer_with_feature_names_out = []
        for name, trans, column, _ in self._iter(fitted=True):
            feature_names_out = trans.get_feature_names_out()
            if feature_names_out is None:
                continue
            transformer_with_feature_names_out.append((name, feature_names_out))

        if not transformer_with_feature_names_out:
            # No feature names
            return np.array([], dtype=object)

        if self.verbose_feature_names_out:
            # Prefix the feature names out with the transformers name
            names = list(
                chain.from_iterable(
                    (f"{name}__{i}" for i in feature_names_out)
                    for name, feature_names_out in transformer_with_feature_names_out
                )
            )
            return np.asarray(names, dtype=object)

        # verbose_feature_names_out is False
        # Check that names are all unique without a prefix
        feature_names_count = Counter(
            chain.from_iterable(s for _, s in transformer_with_feature_names_out)
        )
        top_6_overlap = [
            name for name, count in feature_names_count.most_common(6) if count > 1
        ]
        top_6_overlap.sort()
        if top_6_overlap:
            if len(top_6_overlap) == 6:
                # There are more than 5 overlapping names, we only show the 5
                # of the feature names
                names_repr = str(top_6_overlap[:5])[:-1] + ", ...]"
            else:
                names_repr = str(top_6_overlap)
            raise ValueError(
                f"Output feature names: {names_repr} are not unique. Please set "
                "verbose_feature_names_out=True to add prefixes to feature names"
            )

        return np.concatenate(
            [name for _, name in transformer_with_feature_names_out],
        )


