data_processor.py

"""Data loading, profiling, and preprocessing utilities."""

from __future__ import annotations

import os
import re
import tempfile

import numpy as np
import pandas as pd

# Temp workspace that persists for the app session
TEMP_DIR = os.path.join(tempfile.gettempdir(), "ml_analysis_space")
os.makedirs(TEMP_DIR, exist_ok=True)

CLEANED_PATH = os.path.join(TEMP_DIR, "cleaned_data.csv")

SUPPORTED_EXTENSIONS = [".csv", ".tsv", ".xlsx", ".xls", ".json", ".parquet", ".txt"]

# ------------------------------------------------------------------ load ----

ENCODINGS = ("utf-8", "utf-8-sig", "cp1252", "latin-1")


def _read_text_table(file_path: str, sep: str | None = None) -> pd.DataFrame:
    """Read a delimited text file, trying several encodings.

    With sep=None pandas sniffs the delimiter (handles ; | tab exports).
    Malformed rows are skipped instead of failing the whole file.
    """
    last_err: Exception | None = None
    for encoding in ENCODINGS:
        try:
            return pd.read_csv(
                file_path,
                sep=sep,
                engine="python",
                encoding=encoding,
                on_bad_lines="skip",
            )
        except UnicodeDecodeError as e:
            last_err = e
    raise ValueError(f"Could not decode file with any of {ENCODINGS}: {last_err}")


def _tidy_columns(df: pd.DataFrame) -> pd.DataFrame:
    """Fix messy headers: strip whitespace, name blank/auto columns, dedupe."""
    df = df.copy()
    names = []
    for i, col in enumerate(df.columns):
        name = re.sub(r"\s+", " ", str(col)).strip()
        if not name or name.lower().startswith("unnamed:"):
            name = f"column_{i}"
        names.append(name)
    seen: dict[str, int] = {}
    unique = []
    for name in names:
        if name in seen:
            seen[name] += 1
            name = f"{name}_{seen[name]}"
        else:
            seen[name] = 0
        unique.append(name)
    df.columns = unique
    return df


def load_data(file_path: str) -> pd.DataFrame:
    """Load a dataset from CSV, TSV, Excel, JSON, or Parquet."""
    ext = os.path.splitext(file_path)[1].lower()
    if ext in (".csv", ".txt"):
        df = _read_text_table(file_path)
    elif ext == ".tsv":
        df = _read_text_table(file_path, sep="\t")
    elif ext in (".xlsx", ".xls"):
        df = pd.read_excel(file_path)
    elif ext == ".json":
        df = pd.read_json(file_path)
    elif ext == ".parquet":
        df = pd.read_parquet(file_path)
    else:
        raise ValueError(
            f"Unsupported file type '{ext}'. Supported: {', '.join(SUPPORTED_EXTENSIONS)}"
        )
    if df.shape[0] == 0 or df.shape[1] == 0:
        raise ValueError("The file was read but contains no data.")
    return _tidy_columns(df)


# --------------------------------------------------------------- profile ----


def profile_data(df: pd.DataFrame) -> dict:
    """Return a summary profile of the dataset."""
    numeric_cols = df.select_dtypes(include=np.number).columns.tolist()
    categorical_cols = df.select_dtypes(exclude=np.number).columns.tolist()
    profile = {
        "n_rows": int(df.shape[0]),
        "n_cols": int(df.shape[1]),
        "columns": df.columns.tolist(),
        "dtypes": {c: str(t) for c, t in df.dtypes.items()},
        "numeric_columns": numeric_cols,
        "categorical_columns": categorical_cols,
        "missing_counts": df.isna().sum().to_dict(),
        "missing_total": int(df.isna().sum().sum()),
        "duplicate_rows": int(df.duplicated().sum()),
        "memory_kb": round(df.memory_usage(deep=True).sum() / 1024, 1),
    }
    return profile


def profile_text(profile: dict) -> str:
    """Human-readable summary of a profile dict."""
    lines = [
        f"**Rows:** {profile['n_rows']:,} | **Columns:** {profile['n_cols']} "
        f"| **Memory:** {profile['memory_kb']} KB",
        f"**Numeric columns ({len(profile['numeric_columns'])}):** "
        + (", ".join(profile["numeric_columns"]) or "none"),
        f"**Categorical columns ({len(profile['categorical_columns'])}):** "
        + (", ".join(profile["categorical_columns"]) or "none"),
        f"**Missing values:** {profile['missing_total']:,} | "
        f"**Duplicate rows:** {profile['duplicate_rows']:,}",
    ]
    missing = {k: v for k, v in profile["missing_counts"].items() if v > 0}
    if missing:
        lines.append(
            "**Columns with missing values:** "
            + ", ".join(f"{k} ({v})" for k, v in missing.items())
        )
    return "\n\n".join(lines)


# ------------------------------------------------------------ preprocess ----

_NUMERIC_JUNK = re.compile(r"[\s$€£,%]")
_MISSING_TOKENS = {"", "nan", "none", "null", "na", "n/a", "-", "?", "missing"}


def _normalize_missing_tokens(df: pd.DataFrame) -> tuple[pd.DataFrame, int]:
    """Turn placeholder strings like 'N/A', '?', '-' into real NaN."""
    before = int(df.isna().sum().sum())
    for col in df.select_dtypes(exclude=np.number).columns:
        mask = df[col].astype(str).str.strip().str.lower().isin(_MISSING_TOKENS)
        if mask.any():
            df.loc[mask, col] = np.nan
    return df, int(df.isna().sum().sum()) - before


def _coerce_numeric_strings(df: pd.DataFrame) -> tuple[pd.DataFrame, list[str]]:
    """Convert text columns that are really numbers ('$1,234', '45%') to numeric."""
    converted = []
    for col in df.select_dtypes(include="object").columns:
        stripped = df[col].astype(str).str.replace(_NUMERIC_JUNK, "", regex=True)
        num = pd.to_numeric(stripped, errors="coerce")
        notna = df[col].notna()
        if notna.any() and num[notna].notna().mean() >= 0.8:
            df[col] = num
            converted.append(col)
    return df, converted


def _clip_outliers(df: pd.DataFrame, target_column: str | None) -> tuple[pd.DataFrame, list[str]]:
    """Clip numeric values outside 1.5*IQR to the whisker bounds."""
    clipped = []
    for col in df.select_dtypes(include=np.number).columns:
        if col == target_column:
            continue
        q1, q3 = df[col].quantile([0.25, 0.75])
        iqr = q3 - q1
        if not iqr:
            continue
        lo, hi = q1 - 1.5 * iqr, q3 + 1.5 * iqr
        n = int(((df[col] < lo) | (df[col] > hi)).sum())
        if n:
            df[col] = df[col].clip(lo, hi)
            clipped.append(f"{col} ({n})")
    return df, clipped


def preprocess_data(
    df: pd.DataFrame,
    missing_strategy: str = "Impute (mean/mode)",
    drop_duplicates: bool = True,
    encode_categoricals: bool = True,
    scaling: str = "None",
    clip_outliers: bool = False,
    target_column: str | None = None,
) -> tuple[pd.DataFrame, list[str]]:
    """Clean the dataset and return (cleaned_df, list of steps applied)."""
    steps = []
    df = df.copy()

    # Text tidy-up: strip whitespace in string cells, normalize missing tokens
    for col in df.select_dtypes(include="object").columns:
        df[col] = df[col].str.strip()
    df, n_tokens = _normalize_missing_tokens(df)
    if n_tokens:
        steps.append(
            f"Converted {n_tokens} placeholder values ('N/A', '?', '-', 'null'...) to missing"
        )

    # Text columns that are actually numeric ('$1,234', '45%', '1 000')
    df, converted = _coerce_numeric_strings(df)
    if converted:
        steps.append(f"Converted numeric-looking text columns to numbers: {', '.join(converted)}")

    # Drop columns that are entirely empty
    empty_cols = [c for c in df.columns if df[c].isna().all()]
    if empty_cols:
        df = df.drop(columns=empty_cols)
        steps.append(f"Dropped fully-empty columns: {', '.join(empty_cols)}")

    # Drop constant columns — they carry no signal
    const_cols = [
        c for c in df.columns
        if c != target_column and df[c].nunique(dropna=False) <= 1
    ]
    if const_cols:
        df = df.drop(columns=const_cols)
        steps.append(f"Dropped constant columns: {', '.join(const_cols)}")

    # Duplicates
    if drop_duplicates:
        n = int(df.duplicated().sum())
        if n:
            df = df.drop_duplicates().reset_index(drop=True)
            steps.append(f"Removed {n} duplicate rows")

    # Missing values
    n_missing = int(df.isna().sum().sum())
    if n_missing:
        if missing_strategy.startswith("Drop"):
            before = len(df)
            df = df.dropna().reset_index(drop=True)
            steps.append(f"Dropped {before - len(df)} rows with missing values")
        else:
            use_median = "median" in missing_strategy.lower()
            for col in df.columns:
                if df[col].isna().any():
                    if pd.api.types.is_numeric_dtype(df[col]):
                        fill = df[col].median() if use_median else df[col].mean()
                        df[col] = df[col].fillna(fill)
                    else:
                        mode = df[col].mode()
                        df[col] = df[col].fillna(mode.iloc[0] if len(mode) else "unknown")
            centre = "median" if use_median else "mean"
            steps.append(
                f"Imputed {n_missing} missing values ({centre} for numeric, mode for categorical)"
            )

    # Outliers
    if clip_outliers:
        df, clipped = _clip_outliers(df, target_column)
        if clipped:
            steps.append(f"Clipped outliers beyond 1.5*IQR: {', '.join(clipped)}")

    # Encode categoricals (except the target, which models handle separately)
    if encode_categoricals:
        cat_cols = [
            c
            for c in df.select_dtypes(exclude=np.number).columns
            if c != target_column
        ]
        low_card = [c for c in cat_cols if df[c].nunique() <= 20]
        high_card = [c for c in cat_cols if df[c].nunique() > 20]
        if high_card:
            df = df.drop(columns=high_card)
            steps.append(
                f"Dropped high-cardinality text columns (>20 unique): {', '.join(high_card)}"
            )
        if low_card:
            df = pd.get_dummies(df, columns=low_card, drop_first=True, dtype=int)
            steps.append(f"One-hot encoded: {', '.join(low_card)}")

    # Scale numeric features
    if scaling and scaling != "None":
        from sklearn.preprocessing import MinMaxScaler, StandardScaler

        num_cols = [
            c
            for c in df.select_dtypes(include=np.number).columns
            if c != target_column
        ]
        if num_cols:
            if "Min-Max" in scaling:
                df[num_cols] = MinMaxScaler().fit_transform(df[num_cols])
                steps.append(f"Min-Max scaled {len(num_cols)} numeric columns to [0, 1]")
            else:
                df[num_cols] = StandardScaler().fit_transform(df[num_cols])
                steps.append(f"Standard-scaled {len(num_cols)} numeric columns (mean 0, std 1)")

    if not steps:
        steps.append("Data was already clean — no changes needed")

    df.to_csv(CLEANED_PATH, index=False)
    steps.append(f"Saved cleaned data ({len(df):,} rows x {df.shape[1]} cols) for modeling")
    return df, steps


def generate_preprocessing_code(
    missing_strategy: str,
    drop_duplicates: bool,
    encode_categoricals: bool,
    scaling: str,
    clip_outliers: bool,
    target_column: str | None,
) -> str:
    """Return equivalent standalone pandas code for the preprocessing performed."""
    code = [
        "import re",
        "import numpy as np",
        "import pandas as pd",
        "",
        "df = pd.read_csv('your_data.csv')",
        "",
        "# Tidy text cells and turn placeholder strings into real NaN",
        "MISSING = {'', 'nan', 'none', 'null', 'na', 'n/a', '-', '?', 'missing'}",
        "for col in df.select_dtypes(include='object').columns:",
        "    df[col] = df[col].str.strip()",
        "    mask = df[col].astype(str).str.strip().str.lower().isin(MISSING)",
        "    df.loc[mask, col] = np.nan",
        "",
        "# Convert numeric-looking text columns ('$1,234', '45%') to numbers",
        "for col in df.select_dtypes(include='object').columns:",
        "    num = pd.to_numeric(df[col].astype(str).str.replace(r'[\\s$€£,%]', '', regex=True),",
        "                        errors='coerce')",
        "    if df[col].notna().any() and num[df[col].notna()].notna().mean() >= 0.8:",
        "        df[col] = num",
        "",
        "# Drop fully-empty and constant columns",
        "df = df.dropna(axis=1, how='all')",
        "df = df.drop(columns=[c for c in df.columns if df[c].nunique(dropna=False) <= 1])",
    ]
    if drop_duplicates:
        code += ["", "# Remove duplicate rows", "df = df.drop_duplicates().reset_index(drop=True)"]
    if missing_strategy.startswith("Drop"):
        code += ["", "# Drop rows with missing values", "df = df.dropna().reset_index(drop=True)"]
    else:
        centre = "median" if "median" in missing_strategy.lower() else "mean"
        code += [
            "",
            f"# Impute missing values: {centre} for numeric, mode for categorical",
            "for col in df.columns:",
            "    if df[col].isna().any():",
            "        if pd.api.types.is_numeric_dtype(df[col]):",
            f"            df[col] = df[col].fillna(df[col].{centre}())",
            "        else:",
            "            df[col] = df[col].fillna(df[col].mode().iloc[0])",
        ]
    if clip_outliers:
        code += [
            "",
            "# Clip numeric outliers beyond 1.5*IQR",
            f"target = {target_column!r}",
            "for col in df.select_dtypes(include=np.number).columns:",
            "    if col == target:",
            "        continue",
            "    q1, q3 = df[col].quantile([0.25, 0.75])",
            "    iqr = q3 - q1",
            "    if iqr:",
            "        df[col] = df[col].clip(q1 - 1.5 * iqr, q3 + 1.5 * iqr)",
        ]
    if encode_categoricals:
        code += [
            "",
            "# One-hot encode low-cardinality categoricals (drop high-cardinality text)",
            f"target = {target_column!r}",
            "cat_cols = [c for c in df.select_dtypes(exclude=np.number).columns if c != target]",
            "df = df.drop(columns=[c for c in cat_cols if df[c].nunique() > 20])",
            "cat_cols = [c for c in df.select_dtypes(exclude=np.number).columns if c != target]",
            "df = pd.get_dummies(df, columns=cat_cols, drop_first=True, dtype=int)",
        ]
    if scaling and scaling != "None":
        scaler = "MinMaxScaler" if "Min-Max" in scaling else "StandardScaler"
        code += [
            "",
            f"# Scale numeric features with {scaler}",
            f"from sklearn.preprocessing import {scaler}",
            f"num_cols = [c for c in df.select_dtypes(include=np.number).columns if c != {target_column!r}]",
            f"df[num_cols] = {scaler}().fit_transform(df[num_cols])",
        ]
    code += ["", "df.to_csv('cleaned_data.csv', index=False)"]
    return "\n".join(code)