"""
Data cleaning transformations.

Includes operations for handling missing data, duplicates, and outliers.
"""

from typing import Any

import pandas as pd
import polars as pl


# Pandas implementations
def drop_duplicates_pandas(df: pd.DataFrame, subset: list[str] | None = None, keep: str = "first") -> pd.DataFrame:
    """Drop duplicate rows."""
    return df.drop_duplicates(subset=subset, keep=keep)


def drop_na_pandas(df: pd.DataFrame, subset: list[str] | None = None, how: str = "any") -> pd.DataFrame:
    """Drop rows with missing values."""
    return df.dropna(subset=subset, how=how)


def fill_missing_pandas(
    df: pd.DataFrame,
    column: str,
    strategy: str = "mean",
    fill_value: Any = None,
) -> pd.DataFrame:
    """
    Fill missing values in a column.

    Args:
        strategy: 'mean', 'median', 'mode', 'constant', 'ffill', 'bfill'
        fill_value: Value to use if strategy is 'constant'
    """
    df = df.copy()

    if strategy == "mean":
        df[column] = df[column].fillna(df[column].mean())
    elif strategy == "median":
        df[column] = df[column].fillna(df[column].median())
    elif strategy == "mode":
        mode_val = df[column].mode()
        if len(mode_val) > 0:
            df[column] = df[column].fillna(mode_val[0])
    elif strategy == "constant":
        df[column] = df[column].fillna(fill_value)
    elif strategy == "ffill":
        df[column] = df[column].fillna(method="ffill")
    elif strategy == "bfill":
        df[column] = df[column].fillna(method="bfill")
    else:
        raise ValueError(f"Unknown strategy: {strategy}")

    return df


def remove_outliers_pandas(
    df: pd.DataFrame,
    column: str,
    method: str = "iqr",
    threshold: float = 1.5,
) -> pd.DataFrame:
    """
    Remove outliers from a column.

    Args:
        method: 'iqr' (Interquartile Range) or 'zscore'
        threshold: Threshold for outlier detection
    """
    df = df.copy()

    if method == "iqr":
        Q1 = df[column].quantile(0.25)
        Q3 = df[column].quantile(0.75)
        IQR = Q3 - Q1
        lower = Q1 - threshold * IQR
        upper = Q3 + threshold * IQR
        df = df[(df[column] >= lower) & (df[column] <= upper)]
    elif method == "zscore":
        from scipy import stats
        z_scores = stats.zscore(df[column].dropna())
        abs_z_scores = abs(z_scores)
        df = df[abs_z_scores < threshold]
    else:
        raise ValueError(f"Unknown method: {method}")

    return df


# Polars implementations
def drop_duplicates_polars(df: pl.DataFrame, subset: list[str] | None = None, keep: str = "first") -> pl.DataFrame:
    """Drop duplicate rows."""
    return df.unique(subset=subset, keep=keep)


def drop_na_polars(df: pl.DataFrame, subset: list[str] | None = None, how: str = "any") -> pl.DataFrame:
    """Drop rows with missing values."""
    if subset:
        if how == "any":
            # Drop if any of the columns have null
            for col in subset:
                df = df.filter(pl.col(col).is_not_null())
        else:  # 'all'
            # Drop only if all of the columns have null
            mask = pl.lit(True)
            for col in subset:
                mask = mask & pl.col(col).is_null()
            df = df.filter(~mask)
    else:
        df = df.drop_nulls()
    return df


def fill_missing_polars(
    df: pl.DataFrame,
    column: str,
    strategy: str = "mean",
    fill_value: Any = None,
) -> pl.DataFrame:
    """Fill missing values in a column."""
    if strategy == "mean":
        fill_val = df[column].mean()
        df = df.with_columns(pl.col(column).fill_null(fill_val))
    elif strategy == "median":
        fill_val = df[column].median()
        df = df.with_columns(pl.col(column).fill_null(fill_val))
    elif strategy == "mode":
        mode_val = df[column].mode().first()
        df = df.with_columns(pl.col(column).fill_null(mode_val))
    elif strategy == "constant":
        df = df.with_columns(pl.col(column).fill_null(fill_value))
    elif strategy == "ffill":
        df = df.with_columns(pl.col(column).fill_null(strategy="forward"))
    elif strategy == "bfill":
        df = df.with_columns(pl.col(column).fill_null(strategy="backward"))
    else:
        raise ValueError(f"Unknown strategy: {strategy}")

    return df


def remove_outliers_polars(
    df: pl.DataFrame,
    column: str,
    method: str = "iqr",
    threshold: float = 1.5,
) -> pl.DataFrame:
    """Remove outliers from a column."""
    if method == "iqr":
        Q1 = df[column].quantile(0.25)
        Q3 = df[column].quantile(0.75)
        IQR = Q3 - Q1
        lower = Q1 - threshold * IQR
        upper = Q3 + threshold * IQR
        df = df.filter((pl.col(column) >= lower) & (pl.col(column) <= upper))
    elif method == "zscore":
        # Polars doesn't have built-in zscore, so we calculate manually
        mean = df[column].mean()
        std = df[column].std()
        df = df.filter((pl.col(column) - mean).abs() < threshold * std)
    else:
        raise ValueError(f"Unknown method: {method}")

    return df
