"""
Encoding transformations for categorical variables.
"""

import pandas as pd
import polars as pl
from sklearn.preprocessing import LabelEncoder


# Pandas implementations
def onehot_encode_pandas(
    df: pd.DataFrame,
    column: str,
    drop_first: bool = False,
) -> pd.DataFrame:
    """One-hot encode a categorical column."""
    df = df.copy()
    dummies = pd.get_dummies(df[column], prefix=column, drop_first=drop_first)
    df = pd.concat([df.drop(columns=[column]), dummies], axis=1)
    return df


def label_encode_pandas(df: pd.DataFrame, column: str) -> pd.DataFrame:
    """Label encode a categorical column."""
    df = df.copy()
    le = LabelEncoder()
    df[column] = le.fit_transform(df[column].astype(str))
    return df


def encode_categorical_pandas(
    df: pd.DataFrame,
    column: str,
    method: str = "onehot",
    **kwargs,
) -> pd.DataFrame:
    """
    Encode categorical column.

    Args:
        method: 'onehot' or 'label'
    """
    if method == "onehot":
        return onehot_encode_pandas(df, column, **kwargs)
    elif method == "label":
        return label_encode_pandas(df, column)
    else:
        raise ValueError(f"Unknown encoding method: {method}")


# Polars implementations
def onehot_encode_polars(
    df: pl.DataFrame,
    column: str,
    drop_first: bool = False,
) -> pl.DataFrame:
    """One-hot encode a categorical column."""
    # Get unique values
    unique_values = df[column].unique().to_list()
    
    if drop_first and len(unique_values) > 0:
        unique_values = unique_values[1:]
    
    # Create binary columns
    for value in unique_values:
        col_name = f"{column}_{value}"
        df = df.with_columns(
            (pl.col(column) == value).cast(pl.Int64).alias(col_name)
        )
    
    # Drop original column
    df = df.drop(column)
    return df


def label_encode_polars(df: pl.DataFrame, column: str) -> pl.DataFrame:
    """Label encode a categorical column."""
    # Create a mapping from categories to integers
    unique_values = df[column].unique().sort()
    mapping = {val: idx for idx, val in enumerate(unique_values.to_list())}
    
    # Apply mapping
    df = df.with_columns(
        pl.col(column).map_dict(mapping).alias(column)
    )
    return df


def encode_categorical_polars(
    df: pl.DataFrame,
    column: str,
    method: str = "onehot",
    **kwargs,
) -> pl.DataFrame:
    """Encode categorical column."""
    if method == "onehot":
        return onehot_encode_polars(df, column, **kwargs)
    elif method == "label":
        return label_encode_polars(df, column)
    else:
        raise ValueError(f"Unknown encoding method: {method}")
