"""
Feature engineering transformations.
"""

from typing import Callable

import pandas as pd
import polars as pl
import numpy as np


# Pandas implementations
def create_feature_pandas(
    df: pd.DataFrame,
    new_column: str,
    expression: str,
) -> pd.DataFrame:
    """
    Create a new feature using a Python expression.

    Example: expression = "df['col1'] + df['col2']"
    """
    df = df.copy()
    try:
        df[new_column] = eval(expression, {"df": df, "np": np})
    except Exception as e:
        raise ValueError(f"Failed to evaluate expression: {e}")
    return df


def binning_pandas(
    df: pd.DataFrame,
    column: str,
    bins: int | list[float],
    labels: list[str] | None = None,
    new_column: str | None = None,
) -> pd.DataFrame:
    """
    Bin a numeric column into categories.

    Args:
        bins: Number of bins or list of bin edges
        labels: Labels for bins
        new_column: Name of new column (default: {column}_binned)
    """
    df = df.copy()
    target_column = new_column or f"{column}_binned"
    
    df[target_column] = pd.cut(df[column], bins=bins, labels=labels)
    return df


def polynomial_features_pandas(
    df: pd.DataFrame,
    columns: list[str],
    degree: int = 2,
) -> pd.DataFrame:
    """Create polynomial features from columns."""
    from sklearn.preprocessing import PolynomialFeatures

    df = df.copy()
    poly = PolynomialFeatures(degree=degree, include_bias=False)
    
    # Generate polynomial features
    poly_features = poly.fit_transform(df[columns])
    feature_names = poly.get_feature_names_out(columns)
    
    # Add new features
    for i, name in enumerate(feature_names):
        if name not in columns:  # Skip original features
            df[name] = poly_features[:, i]
    
    return df


# Polars implementations (simpler versions)
def create_feature_polars(
    df: pl.DataFrame,
    new_column: str,
    source_columns: list[str],
    operation: str = "sum",
) -> pl.DataFrame:
    """
    Create a new feature from existing columns.

    Args:
        operation: 'sum', 'product', 'mean', 'difference'
    """
    if operation == "sum":
        df = df.with_columns(
            pl.sum_horizontal(source_columns).alias(new_column)
        )
    elif operation == "product":
        # Multiply columns
        expr = pl.lit(1)
        for col in source_columns:
            expr = expr * pl.col(col)
        df = df.with_columns(expr.alias(new_column))
    elif operation == "mean":
        df = df.with_columns(
            pl.mean_horizontal(source_columns).alias(new_column)
        )
    elif operation == "difference":
        if len(source_columns) == 2:
            df = df.with_columns(
                (pl.col(source_columns[0]) - pl.col(source_columns[1])).alias(new_column)
            )
    else:
        raise ValueError(f"Unknown operation: {operation}")
    
    return df


def binning_polars(
    df: pl.DataFrame,
    column: str,
    bins: list[float],
    labels: list[str] | None = None,
    new_column: str | None = None,
) -> pl.DataFrame:
    """Bin a numeric column into categories."""
    target_column = new_column or f"{column}_binned"
    
    # Use cut function
    df = df.with_columns(
        pl.col(column).cut(bins, labels=labels).alias(target_column)
    )
    return df
