from sklearn.base import RegressorMixin, TransformerMixin
from sklearn.utils.validation import check_is_fitted, check_array, check_X_y

class HybridBoostingRegressor(RegressorMixin, TransformerMixin):
    def __init__(self, kwargs):
        self.names = []
        self.models = []
        self.features_list = []
        self.targets = []

        for name, model, features, target in kwargs:
            self.names.append(name)
            self.models.append(model)
            self.features_list.append(features)

            if len(target) == 1:
              self.targets.append(target[0])
            else:
              self.targets.append(target)

    def filter_X(self, X, features):
        if len(features) > 0:
          X = X[features]
        return X

    def get_y(self, y, y_intermediate, target):
        if len(target) == 0:
          y = y
        else:
          y = y_intermediate[target]

        return y

    def fit(self, X, y, y_intermediate = None, **fit_params):
        if y_intermediate is None or y_intermediate.shape[0] == 0:
          for target in self.targets:
            if len(target) > 0:
              return Exception(f"y_intermediate not found for {target}")

        y_res = y.copy()
        y_intermediate_res = y_intermediate.copy()

        self.fitted_models_ = []
        for model, features, target in zip(self.models, self.features_list, self.targets):

            X = self.filter_X(X.copy(), features)
            y = self.get_y(y_res, y_intermediate_res, target)

            # Check that X and y have correct shape
            X, y = check_X_y(X, y)

            model.fit(X, y, **fit_params)

            self.fitted_models_.append(model)

            if len(target) == 0:
              pred = model.predict(X)
            else:
              pred = model.predict(X)

            # residual
            if len(target) != 0:
              y_intermediate_res[target] = self.accumulate(y_intermediate_res[target], pred)
            y_res = self.get_residual(y_res, pred)

        return self

    def predict(self, X):
        check_is_fitted(self, "fitted_models_")

        pred = self.pred_init
        for fitted_model, features, target in zip(self.fitted_models_, self.features_list, self.targets):
            X = self.filter_X(X.copy(), features)

            # Input validation
            X = check_array(X)

            pred = self.accumulate(pred, fitted_model.predict(X))

        return pred

class AdditiveHybridBoostingRegressor(HybridBoostingRegressor):
    """
    y_hat = f_1(x_1) + f_1(x_2) + ...
    """
    def __init__(self, kwargs):
      super().__init__(kwargs)
      self.pred_init = 0
    def get_residual(self, y, pred):
      return y-pred
    def accumulate(self, y, pred):
      return y+pred

class MultiplicativeHybridBoostingRegressor(HybridBoostingRegressor):
    """
    y_hat = f_1(x_1) * f_1(x_2) * ...
    """
    def __init__(self, kwargs):
      super().__init__(kwargs)
      self.pred_init = 1
    def get_residual(self, y, pred):
      return y/pred
    def accumulate(self, y, pred):
      return y*pred

x = np.arange(1, 1000, 1)
X = pd.DataFrame().assign(
  x1 = x,
  x2 = x**2
)
y_intermediate = pd.DataFrame().assign(
  trend = 10000*X["x1"] + X["x2"],
  seasonality = np.sin(X["x1"]) + np.sin(X["x1"])
)

y = y_intermediate["trend"] + y_intermediate["seasonality"]

from sklearn.linear_model import LinearRegression
from sklearn.ensemble import RandomForestRegressor

model = AdditiveHybridBoostingRegressor(
	[
		('lr', LinearRegression(), [], ["trend"]),
		('rf', RandomForestRegressor(n_estimators=100, random_state=42), [], ["seasonality"])
	]
)
model.fit(X, y, y_intermediate)
model.predict(X)