from __future__ import annotations

import itertools
import numpy as np
import pandas as pd


def interaction_terms(df: pd.DataFrame, cols: list[str]) -> pd.DataFrame:
    """Add pairwise interaction terms for selected numeric columns."""
    out = df.copy()
    for a, b in itertools.combinations(cols, 2):
        out[f"{a}__x__{b}"] = out[a] * out[b]
    return out


def polynomial_features(
    df: pd.DataFrame, cols: list[str], degree: int = 2
) -> pd.DataFrame:
    """Add simple power terms up to degree."""
    out = df.copy()
    for c in cols:
        for d in range(2, degree + 1):
            out[f"{c}__pow{d}"] = out[c] ** d
    return out


def drop_low_variance(df: pd.DataFrame, threshold: float = 1e-8) -> pd.DataFrame:
    """Drop numeric columns with variance below threshold."""
    keep = [
        c
        for c in df.columns
        if not pd.api.types.is_numeric_dtype(df[c]) or df[c].var(ddof=0) > threshold
    ]
    return df[keep].copy()


def drop_high_correlation(df: pd.DataFrame, threshold: float = 0.95) -> pd.DataFrame:
    """Drop one of highly correlated pairs."""
    num = df.select_dtypes(include="number")
    corr = num.corr().abs()
    upper = corr.where(~np.tril(np.ones(corr.shape)).astype(bool))
    to_drop = [col for col in upper.columns if any(upper[col] > threshold)]
    return df.drop(columns=to_drop)
