# aivora/generators/ratio.py

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


class LogRatioGenerator(AivoraComponent):
    def __init__(
        self,
        columns=None,
        symmetric=False,
        max_pairs=None,
        min_denominator=1e-8,
        max_abs_log=20,
    ):
        """
        LogRatio generator safe against zero, extreme values, and symmetric pairs.

        Parameters
        ----------
        columns : list or None
            Columns to include. If None, use all columns.
        symmetric : bool
            If True, generate both f1/f2 and f2/f1.
        max_pairs : int or None
            Optional cap on number of pairs to avoid explosion.
        min_denominator : float
            Threshold below which numerator or denominator is considered invalid → NaN.
        max_abs_log : float
            Threshold for |log(ratio)| above which value is set to NaN.
        """
        self.columns = columns
        self.symmetric = symmetric
        self.max_pairs = max_pairs
        self.min_denominator = min_denominator
        self.max_abs_log = max_abs_log

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

        if self.columns is not None:
            self.use_cols_ = [c for c in self.columns if c in self.columns_]
        else:
            self.use_cols_ = self.columns_

        pairs = list(itertools.combinations(self.use_cols_, 2))

        if self.symmetric:
            pairs += [(b, a) for a, b in pairs]

        if self.max_pairs is not None:
            pairs = pairs[: self.max_pairs]

        self.pairs_ = pairs
        self.col_index_ = {c: i for i, c in enumerate(self.columns_)}

        return self

    def transform(self, X):
        X_values = X.values.astype(float)
        n_samples = X_values.shape[0]

        ratio_features = []
        ratio_names = []

        for f1, f2 in self.pairs_:
            i1 = self.col_index_[f1]
            i2 = self.col_index_[f2]

            numerator = X_values[:, i1]
            denominator = X_values[:, i2]

            valid_mask = np.abs(denominator) > self.min_denominator

            ratio = np.full_like(numerator, np.nan, dtype=float)

            safe_ratio = numerator[valid_mask] / denominator[valid_mask]


            positive_mask = safe_ratio > 0

            log_ratio = np.full_like(safe_ratio, np.nan)
            log_ratio[positive_mask] = np.log(safe_ratio[positive_mask])

            log_ratio[np.abs(log_ratio) > self.max_abs_log] = np.nan

            ratio_indices = np.where(valid_mask)[0]
            ratio[ratio_indices] = log_ratio

            ratio_features.append(ratio)
            ratio_names.append(f"{f1}_over_{f2}")

        if ratio_features:
            ratio_matrix = np.column_stack(ratio_features)
            full_matrix = np.hstack([X_values, ratio_matrix])
            full_columns = self.columns_ + ratio_names
        else:
            full_matrix = X_values
            full_columns = self.columns_

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