src/dataset.py

"""Dataset loading and preprocessing for Banco Ripley credit-scoring.

Handles data ingestion, dtype casting, null imputation and ordinal encoding
of categorical variables.  Produces a single cleaned CSV ready for feature
engineering and modelling.
"""

from pathlib import Path

import pandas as pd
from loguru import logger
from tqdm import tqdm

# ── Column-type registry ─────────────────────────────────────────────────────

IDENTIFIER_COLS: list[str] = ["serie"]
TEMPORAL_COLS: list[str] = ["PERIODO"]
TARGET_COL: str = "predictiva_early"

BINARY_FLAG_COLS: list[str] = [
    "FLAG_ENTIDAD_PRINCIPAL", "FLAG_TC_MODELOS", "FLAG_MES",
    "MARCA_HP", "MARCA_SEG_VIDA", "MARCA_DIF", "MARCA_CONV",
    "MARCA_DIF3", "MARCA_SEG_VIDA3", "MARCA_GAR3",
    "MARCA_DIF6", "MARCA_SEG_VIDA6", "MARCA_GAR6", "MARCA_HIP6",
    "MARCA_DIF12", "MARCA_SEG_VIDA12", "MARCA_GAR12", "MARCA_HIP12",
    "MARCA_CONV3", "MARCA_HIP3",
    "MARCA_CONV6",
    "MARCA_CONV12",
    "MARCA_GAR", "FLAG_TENENCIA_VEHICULAR",
]

CATEGORICAL_COLS: list[str] = [
    "GENERO", "GRADO_INSTRUCCION", "DEPARTAMENTO",
    "PROVINCIA", "DISTRITO", "SITUACION_LABORAL", "ESTADO_CIVIL",
]

DISCRETE_COLS: list[str] = [
    "EDAD", "MESES_ANT_RCC", "CONTAR_COMP",
    "MAX_ATRASO3", "MAX_ATRASO6", "MAX_ATRASO12",
    "NMES_UMORA",
    "Max_AumKP", "Max_AumMORA", "Max_AumKT", "Max_DismDTOTAL",
    "Max_AumMES", "Max_DismMES", "Max_DismCONS", "Max_AumCONS",
    "Max_AumDTOTAL", "Max_DismKP", "Max_DismKT",
]

# ── Public helpers ────────────────────────────────────────────────────────────


def load_raw(input_path: Path) -> pd.DataFrame:
    """Load raw CSV from disk.

    Args:
        input_path: Absolute path to the raw CSV file.

    Returns:
        Raw ``DataFrame`` with all original columns preserved.
    """
    logger.info(f"Loading raw data from {input_path}")
    df = pd.read_csv(input_path, low_memory=False, encoding="latin1")
    # Normalize whitespace-only strings (e.g. " ", "  ") to proper NaN so
    # that type casting, null detection and numeric operations work correctly.
    df = df.replace(r"^\s*$", pd.NA, regex=True)
    logger.info(f"Loaded {len(df):,} rows × {df.shape[1]} columns")
    return df


def cast_column_types(df: pd.DataFrame) -> pd.DataFrame:
    """Cast columns to correct dtypes based on the data dictionary.

    Conversion rules:

    * Binary flags → ``int8`` (0 / 1).
    * Categoricals → ``category``.
    * Discrete integers → nullable ``Int16``.
    * ``PERIODO`` → ``str`` (year-month identifier).
    * Target → ``int8``.

    Args:
        df: Raw ``DataFrame``.

    Returns:
        ``DataFrame`` with corrected dtypes.
    """
    df = df.copy()

    for col in BINARY_FLAG_COLS:
        if col in df.columns:
            # -1 sentinel: distinguishes unknown from confirmed absence (0)
            df[col] = pd.to_numeric(df[col], errors="coerce").fillna(-1).astype("int8")

    for col in CATEGORICAL_COLS:
        if col in df.columns:
            df[col] = df[col].astype("category")

    for col in DISCRETE_COLS:
        if col in df.columns:
            df[col] = pd.to_numeric(df[col], errors="coerce").astype("Int16")

    if "PERIODO" in df.columns:
        df["PERIODO"] = df["PERIODO"].astype(str)

    if TARGET_COL in df.columns:
        # Invert encoding: original 1=Bueno → 0, 0=Malo/moroso → 1
        # Convention: positive class (1) = risk event (morosidad)
        df[TARGET_COL] = (1 - pd.to_numeric(df[TARGET_COL], errors="coerce")).astype("int8")

    return df


def impute_nulls(df: pd.DataFrame) -> pd.DataFrame:
    """Impute missing values with sentinel-based strategies for tree models.

    Sentinel values keep missingness as a learnable signal rather than
    conflating it with observed zero/mean values.

    Strategy by column type:

    * Binary flags  → ``-1``       (distinguishes unknown from confirmed 0/1).
    * Categoricals  → ``"Otros"``  (new sentinel category).
    * Discrete ints → ``-999``     (clearly out-of-range, fits ``Int16``).
    * Continuous    → ``-999999``  (clearly out-of-range sentinel).

    Args:
        df: ``DataFrame`` after :func:`cast_column_types`.

    Returns:
        ``DataFrame`` with no remaining null values in feature columns.
    """
    df = df.copy()
    null_counts = df.isnull().sum()
    cols_with_nulls = null_counts[null_counts > 0].index.tolist()
    logger.info(f"Columns with nulls before imputation: {len(cols_with_nulls)}")

    for col in tqdm(cols_with_nulls, desc="Imputing nulls"):
        if col in BINARY_FLAG_COLS:
            df[col] = df[col].fillna(-1)
        elif col in CATEGORICAL_COLS:
            # Sentinel category keeps missingness as a learnable signal for tree models
            df[col] = df[col].cat.add_categories("Otros").fillna("Otros")
        elif col in DISCRETE_COLS:
            # -999 fits within Int16 range and acts as a clear sentinel
            df[col] = pd.to_numeric(df[col], errors="coerce").fillna(-999).astype("Int16")
        else:
            # -999999 as sentinel for continuous features
            df[col] = pd.to_numeric(df[col], errors="coerce").fillna(-999999)

    return df


def encode_categoricals(df: pd.DataFrame) -> pd.DataFrame:
    """Ordinal-encode categorical columns for tree-based models.

    Each category is mapped to its integer code (pandas ``cat.codes``).
    Unknown / unseen categories receive code ``-1`` automatically.

    Args:
        df: ``DataFrame`` with imputed values.

    Returns:
        ``DataFrame`` where categorical columns are stored as ``int16``.
    """
    df = df.copy()
    for col in CATEGORICAL_COLS:
        if col in df.columns:
            df[col] = df[col].cat.codes.astype("int16")
    return df


def get_feature_cols(df: pd.DataFrame) -> list[str]:
    """Return feature column names, excluding IDs, temporal cols and target.

    Args:
        df: Processed ``DataFrame``.

    Returns:
        Ordered list of column names suitable for model input.
    """
    exclude = set(IDENTIFIER_COLS + TEMPORAL_COLS + [TARGET_COL])
    return [c for c in df.columns if c not in exclude]


if __name__ == "__main__":
    app()