from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression
from sklearn.decomposition import PCA
import numpy as np
from sklearn.base import BaseEstimator, TransformerMixin
import sklearn

from sklearn.utils._metadata_requests import (
    MetadataRouter,
    MethodMapping,
    process_routing,
)

sklearn.set_config(enable_metadata_routing=True)

class MyTransformer(TransformerMixin, BaseEstimator):
    def fit(self, X, y=None, sample_weight=None):
        return self
    def transform(self, X):
        return X

class MyMetaTransformer(TransformerMixin, BaseEstimator):
    def __init__(self, estimator):
        self.estimator = estimator

    def fit(self, X, y=None, **params):
        routed_params = process_routing(self, "fit", **params)
        self.estimator.fit(X, y, **routed_params.estimator.fit)

    def transform(self, X):
        return self.estimator.transform(X)

    def get_metadata_routing(self):
        router = MetadataRouter(owner=self.__class__.__name__).add(
            estimator=self.estimator,
            method_mapping=MethodMapping().add(caller="fit", callee="fit"),
        )
        return router


class MyPipeline(BaseEstimator):
    def __init__(self, steps):
        self.steps = steps

    def fit(self, X, y=None, **params):
        routed_params = process_routing(self, "fit", **params)
        for name, step in self.steps[:-1]:
            if hasattr(step, "fit_transform"):
                X = step.fit_transform(X, y, **routed_params[name].fit_transform)
            else:
                X = step.fit(X, y, **routed_params[name].fit).transform(X)
        self.steps[-1][1].fit(X, y, **routed_params[self.steps[-1][0]].fit)

    def get_metadata_routing(self):
        router = MetadataRouter(owner=self.__class__.__name__)
        for name, step in self.steps:
            if hasattr(step, "fit_transform"):
                router.add(
                    **{name: step},
                    method_mapping=MethodMapping().add(caller="fit", callee="fit_transform"),
                )
            else:
                router.add(
                    **{name: step},
                    method_mapping=MethodMapping().add(caller="fit", callee="fit"),
                )
        return router


X, y = load_iris(return_X_y=True)
sample_weights = np.ones(len(y))
estimator = MyMetaTransformer(MyTransformer().set_fit_request(sample_weight=True))
pipe = MyPipeline([("trs", estimator), ("estimator", LogisticRegression().set_fit_request(sample_weight=True))])
pipe.fit(X, y, sample_weight=sample_weights)