import itertools
import pandas as pd
import numpy as np
from ..core.component import AivoraComponent


class InteractionGenerator(AivoraComponent):
    def __init__(self, pairs=None, auto=True):
        self.pairs = pairs
        self.auto = auto
        self.pairs_ = None

    def fit(self, X, y=None):
        self.columns_ = list(X.columns)

        if self.pairs is None and self.auto:
            self.pairs_ = list(itertools.combinations(X.columns, 2))
        else:
            self.pairs_ = self.pairs

        if self.pairs_ is None:
            self.pairs_ = [] 

        return self

    def transform(self, X):
        X_values = X.values
        col_index = {c: i for i, c in enumerate(self.columns_)}

        new_features = []
        new_names = []

        for f1, f2 in self.pairs_:
            i1 = col_index[f1]
            i2 = col_index[f2]
            interaction = X_values[:, i1] * X_values[:, i2]

            new_features.append(interaction)
            new_names.append(f"{f1}_{f2}")

        if new_features:
            new_matrix = np.column_stack(new_features)
            full_matrix = np.hstack([X_values, new_matrix])
            full_columns = self.columns_ + new_names
        else:
            full_matrix = X_values
            full_columns = self.columns_

        return pd.DataFrame(full_matrix, columns=full_columns, index=X.index)

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