"""
Cleaner — Production-grade data cleaning, normalization, and encoding utilities.
"""

from __future__ import annotations

import re
from typing import Any, Literal

import numpy as np
import pandas as pd
from sklearn.preprocessing import (
    LabelEncoder, MinMaxScaler, OneHotEncoder, RobustScaler, StandardScaler,
)

from bizanalyticsai.core.dataframe import BizFrame


def clean(
    df: pd.DataFrame,
    *,
    drop_duplicates: bool = True,
    strip_strings: bool = True,
    fix_column_names: bool = True,
    drop_empty_cols_threshold: float = 0.95,
    infer_dtypes: bool = True,
) -> BizFrame:
    """
    One-shot comprehensive data cleaning.

    Steps performed:
    1. Fix column names (snake_case, remove special chars)
    2. Strip leading/trailing whitespace from strings
    3. Drop fully duplicated rows
    4. Drop columns with > threshold missing values
    5. Infer better dtypes (numeric, datetime, categorical)

    Args:
        df:                        Input DataFrame.
        drop_duplicates:           Remove exact duplicate rows.
        strip_strings:             Strip whitespace from object columns.
        fix_column_names:          Normalize column names to snake_case.
        drop_empty_cols_threshold: Drop columns with this fraction of nulls.
        infer_dtypes:              Auto-convert dtypes where possible.

    Returns:
        Cleaned BizFrame.
    """
    out = df.copy()

    if fix_column_names:
        out.columns = [_snake(c) for c in out.columns]

    if strip_strings:
        str_cols = out.select_dtypes(include="object").columns
        out[str_cols] = out[str_cols].apply(lambda s: s.str.strip())

    if drop_duplicates:
        before = len(out)
        out = out.drop_duplicates()
        if len(out) < before:
            pass  # logger can be added here

    # Drop mostly-empty columns
    null_frac = out.isna().mean()
    drop_cols = null_frac[null_frac > drop_empty_cols_threshold].index.tolist()
    if drop_cols:
        out = out.drop(columns=drop_cols)

    if infer_dtypes:
        out = _infer_dtypes(out)

    return BizFrame(out)


def fill_missing(
    df: pd.DataFrame,
    strategy: Literal["mean", "median", "mode", "zero", "ffill", "bfill", "drop"] = "median",
    cols: list[str] | None = None,
    fill_value: Any = None,
) -> BizFrame:
    """
    Fill missing values using a chosen strategy.

    Args:
        df:         Input DataFrame.
        strategy:   One of mean/median/mode/zero/ffill/bfill/drop.
        cols:       Columns to apply to. Defaults to all columns.
        fill_value: Explicit fill value (overrides strategy).

    Returns:
        BizFrame with missing values handled.
    """
    out  = df.copy()
    cols = cols or list(out.columns)

    if fill_value is not None:
        out[cols] = out[cols].fillna(fill_value)
        return BizFrame(out)

    for col in cols:
        s = out[col]
        if strategy == "drop":
            out = out.dropna(subset=[col])
        elif strategy == "ffill":
            out[col] = s.ffill()
        elif strategy == "bfill":
            out[col] = s.bfill()
        elif strategy == "zero":
            out[col] = s.fillna(0)
        elif strategy == "mode":
            mode = s.mode()
            out[col] = s.fillna(mode.iloc[0] if len(mode) else np.nan)
        elif pd.api.types.is_numeric_dtype(s):
            if strategy == "mean":
                out[col] = s.fillna(s.mean())
            elif strategy == "median":
                out[col] = s.fillna(s.median())
        else:
            # For non-numeric, default to mode
            mode = s.mode()
            out[col] = s.fillna(mode.iloc[0] if len(mode) else "")

    return BizFrame(out)


def remove_outliers(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    method: Literal["iqr", "zscore", "percentile"] = "iqr",
    threshold: float = 3.0,
    action: Literal["drop", "clip", "flag"] = "clip",
) -> BizFrame:
    """
    Detect and handle outliers in numeric columns.

    Methods:
    - iqr:        Interquartile range (robust, recommended)
    - zscore:     Standard deviations from mean (assumes normality)
    - percentile: Explicit percentile bounds

    Actions:
    - drop:  Remove rows containing outliers
    - clip:  Clip values to the bounds (Winsorization)
    - flag:  Add a boolean column <col>_is_outlier without removing

    Returns:
        BizFrame with outliers handled.
    """
    out  = df.copy()
    num_cols = cols or list(out.select_dtypes(include="number").columns)
    mask = pd.Series(False, index=out.index)

    for col in num_cols:
        s = out[col].dropna()
        if method == "iqr":
            q1, q3 = s.quantile(0.25), s.quantile(0.75)
            iqr = q3 - q1
            lo, hi = q1 - threshold * iqr, q3 + threshold * iqr
        elif method == "zscore":
            mean, std = s.mean(), s.std()
            lo, hi = mean - threshold * std, mean + threshold * std
        else:  # percentile
            lo = s.quantile(0.01)
            hi = s.quantile(0.99)

        col_mask = (out[col] < lo) | (out[col] > hi)
        mask |= col_mask.fillna(False)

        if action == "clip":
            out[col] = out[col].clip(lo, hi)
        elif action == "flag":
            out[f"{col}_is_outlier"] = col_mask

    if action == "drop":
        out = out[~mask]

    return BizFrame(out)


def normalize(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    method: Literal["minmax", "zscore", "robust"] = "minmax",
    suffix: str = "_norm",
    inplace: bool = False,
) -> BizFrame:
    """
    Normalize numeric columns using a scaling method.

    Args:
        df:      Input DataFrame.
        cols:    Columns to normalize. Defaults to all numeric.
        method:  minmax (0-1), zscore (μ=0, σ=1), robust (median/IQR-based).
        suffix:  Suffix for new columns. Use "" to overwrite originals.
        inplace: Overwrite original columns (equivalent to suffix="").

    Returns:
        BizFrame with normalized columns added (or overwritten).
    """
    out  = df.copy()
    cols = cols or list(out.select_dtypes(include="number").columns)

    scaler_map = {
        "minmax": MinMaxScaler(),
        "zscore": StandardScaler(),
        "robust": RobustScaler(),
    }
    scaler = scaler_map[method]
    sfx = "" if inplace else suffix

    data = out[cols].copy()
    data_filled = data.fillna(data.median())
    scaled = scaler.fit_transform(data_filled)

    for i, col in enumerate(cols):
        out[f"{col}{sfx}"] = scaled[:, i]

    return BizFrame(out)


def encode(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    method: Literal["label", "onehot", "ordinal"] = "label",
    drop_original: bool = False,
) -> BizFrame:
    """
    Encode categorical columns.

    Methods:
    - label:   Integer encoding (0, 1, 2, …)
    - onehot:  One-hot encoding (creates new binary columns)
    - ordinal: Same as label but preserves order if specified

    Returns:
        BizFrame with encoded columns.
    """
    out  = df.copy()
    cols = cols or list(out.select_dtypes(include=["object", "category"]).columns)

    if method in ("label", "ordinal"):
        le = LabelEncoder()
        for col in cols:
            out[f"{col}_enc"] = le.fit_transform(out[col].astype(str).fillna("__null__"))
    elif method == "onehot":
        enc = OneHotEncoder(sparse_output=False, handle_unknown="ignore")
        encoded = enc.fit_transform(out[cols].astype(str).fillna("__null__"))
        new_cols = enc.get_feature_names_out(cols)
        ohe_df   = pd.DataFrame(encoded, columns=new_cols, index=out.index)
        out      = pd.concat([out, ohe_df], axis=1)

    if drop_original:
        out = out.drop(columns=cols)

    return BizFrame(out)


def deduplicate(
    df: pd.DataFrame,
    subset: list[str] | None = None,
    keep: Literal["first", "last", "none"] = "first",
) -> BizFrame:
    """Remove duplicate rows, optionally based on a subset of columns."""
    _keep: str | bool = keep if keep in ("first", "last") else False
    return BizFrame(df.drop_duplicates(subset=subset, keep=_keep))


def cast_dates(
    df: pd.DataFrame,
    cols: list[str] | None = None,
    fmt: str | None = None,
    errors: Literal["raise", "coerce", "ignore"] = "coerce",
) -> BizFrame:
    """Convert string columns to datetime."""
    out  = df.copy()
    cols = cols or list(out.select_dtypes(include="object").columns)
    for col in cols:
        try:
            out[col] = pd.to_datetime(out[col], format=fmt, errors=errors)
        except Exception:
            pass
    return BizFrame(out)


# ── Internal helpers ──────────────────────────────────────────────────────────
def _snake(name: str) -> str:
    """Convert column name to snake_case."""
    s = str(name).strip()
    s = re.sub(r"[\s\-]+", "_", s)
    s = re.sub(r"[^\w]", "", s)
    s = re.sub(r"([A-Z]+)([A-Z][a-z])", r"\1_\2", s)
    s = re.sub(r"([a-z\d])([A-Z])", r"\1_\2", s)
    return s.lower().strip("_")


def _infer_dtypes(df: pd.DataFrame) -> pd.DataFrame:
    """Try to infer better dtypes for object columns."""
    out = df.copy()
    for col in out.select_dtypes(include="object").columns:
        # Try numeric
        converted = pd.to_numeric(out[col], errors="coerce")
        if converted.notna().sum() / max(len(converted), 1) > 0.9:
            out[col] = converted
            continue
        # Try datetime
        try:
            converted_dt = pd.to_datetime(out[col], errors="coerce", infer_datetime_format=True)
            if converted_dt.notna().sum() / max(len(converted_dt), 1) > 0.9:
                out[col] = converted_dt
                continue
        except Exception:
            pass
        # Low cardinality → category
        if out[col].nunique() / max(len(out), 1) < 0.05:
            out[col] = out[col].astype("category")
    return out
