import numpy as np
import pandas as pd
from .base import ModuleBase
from itertools import combinations


class CustomLogRatio(ModuleBase):

    def __init__(self, pairs=None):

        self.pairs = pairs
        self.fitted_pairs = None

    def fit(self, X: pd.DataFrame, pairs=None):

        self.fitted_pairs = pairs if pairs is not None else self.pairs
        if self.fitted_pairs is None:
            raise ValueError("Pairs must be provided either at init or in fit.")

        self.fitted_pairs = [
            (c1, c2) for c1, c2 in self.fitted_pairs if c1 in X.columns and c2 in X.columns
        ]
        return self

    def transform(self, X: pd.DataFrame) -> pd.DataFrame:
        if self.fitted_pairs is None:
            raise ValueError("Transformer is not fitted yet. Call `fit` first.")

        if len(self.fitted_pairs) == 0:
            return pd.DataFrame(index=X.index)

        col_indices = np.array([[X.columns.get_loc(c1), X.columns.get_loc(c2)]
                                for c1, c2 in self.fitted_pairs])
        selected = X.values[:, col_indices]  # shape: (n_samples, n_pairs, 2)

        numerator = selected[:, :, 0]
        denominator = selected[:, :, 1]

        with np.errstate(divide='ignore', invalid='ignore'):
            ratio = np.divide(numerator, denominator)
            ratio[denominator == 0] = np.nan
            ratio[numerator == 0] = np.nan
            feature_array = np.log(ratio)

        feature_names = [f"{c1}__logratio__{c2}" for c1, c2 in self.fitted_pairs]

        return pd.DataFrame(feature_array, columns=feature_names, index=X.index)

    def fit_transform(self, X: pd.DataFrame, pairs=None) -> pd.DataFrame:
        return self.fit(X, pairs).transform(X)
