src/preprocess_bd_structured.py

import numpy as np
import pandas as pd
from sklearn.compose import ColumnTransformer
from sklearn.impute import SimpleImputer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import OneHotEncoder

from .text_cleaning import clean_text


def fill_blank_or_nan(df: pd.DataFrame, col: str, value: str) -> None:
    if col not in df.columns:
        return
    s = df[col].astype(str).str.strip()
    missing_mask = (
        df[col].isna() | s.eq("") | s.str.lower().isin(["nan", "none", "null"])
    )
    df.loc[missing_mask, col] = value


def preprocess_bd(
    bd_df: pd.DataFrame, drop_cols: list[str], defaults: dict
) -> pd.DataFrame:
    bd = bd_df.copy()
    bd.columns = [str(c).strip() for c in bd.columns]
    bd = bd.loc[:, ~bd.columns.duplicated()]

    bd = bd.drop(columns=drop_cols, errors="ignore")

    for c in bd.columns:
        if bd[c].dtype == "object":
            bd[c] = bd[c].astype(str).str.strip()

    numeric_candidates = [
        "age",
        "temperature",
        "feels_like",
        "temp_min",
        "temp_max",
        "air_pressure",
        "air_humidity",
        "wind_speed",
        "wind_deg",
        "clouds_sky",
    ]
    for c in numeric_candidates:
        if c in bd.columns:
            bd[c] = pd.to_numeric(bd[c], errors="coerce")

    if "suicide_date" in bd.columns:
        bd["suicide_date"] = pd.to_datetime(bd["suicide_date"], errors="coerce")
        bd["month"] = bd["suicide_date"].dt.month
        bd["dayofweek"] = bd["suicide_date"].dt.dayofweek
        bd["year"] = bd["suicide_date"].dt.year

    if "time" in bd.columns:
        bd["time"] = bd["time"].astype(str).str.lower().str.strip()

    for col in ["reason", "reason_description"]:
        if col in bd.columns:
            bd[col] = bd[col].apply(clean_text)

    for col, val in defaults.items():
        fill_blank_or_nan(bd, col, val)

    return bd


def build_profile_matrix(profile_df: pd.DataFrame):
    num_cols = [
        c for c in profile_df.columns if pd.api.types.is_numeric_dtype(profile_df[c])
    ]
    cat_cols = [c for c in profile_df.columns if c not in num_cols]

    numeric_pipe = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="median")),
        ]
    )

    categorical_pipe = Pipeline(
        steps=[
            ("imputer", SimpleImputer(strategy="most_frequent")),
            ("onehot", OneHotEncoder(handle_unknown="ignore")),
        ]
    )

    preprocessor = ColumnTransformer(
        transformers=[
            ("num", numeric_pipe, num_cols),
            ("cat", categorical_pipe, cat_cols),
        ],
        remainder="drop",
    )

    X_profile = preprocessor.fit_transform(profile_df)
    return X_profile, preprocessor, num_cols, cat_cols