preprocessing.py

"""
preprocessing.py

Handles loading and preprocessing of the credit card transaction dataset.

The raw dataset (Kaggle "Credit Card Fraud Detection", ULB) contains:
    - Time    : seconds elapsed since the first transaction (raw, unscaled)
    - V1..V28 : PCA-transformed features (already zero-mean / unit-variance)
    - Amount  : transaction amount (raw, unscaled)
    - Class   : 0 = normal transaction, 1 = fraudulent transaction

Only 'Time' and 'Amount' need scaling since V1-V28 are already the output
of a PCA transformation performed by the original dataset authors.
"""

from typing import Tuple

import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler

TARGET_COLUMN = "Class"
COLUMNS_TO_SCALE = ["Time", "Amount"]


def load_dataset(path: str) -> pd.DataFrame:
    """Load the raw dataset from a CSV file."""
    df = pd.read_csv(path)
    return df


def get_feature_columns(df: pd.DataFrame) -> list:
    """Return the list of model input feature columns (everything but Class)."""
    return [c for c in df.columns if c != TARGET_COLUMN]


def split_dataset(
    df: pd.DataFrame, test_size: float = 0.2, random_state: int = 42
) -> Tuple[pd.DataFrame, pd.DataFrame]:
    """Stratified train/test split so both splits keep the original fraud ratio."""
    train_df, test_df = train_test_split(
        df,
        test_size=test_size,
        random_state=random_state,
        stratify=df[TARGET_COLUMN],
    )
    return train_df.reset_index(drop=True), test_df.reset_index(drop=True)


def fit_scaler(train_df: pd.DataFrame) -> StandardScaler:
    """Fit a StandardScaler on the 'Time' and 'Amount' columns of the training split only."""
    scaler = StandardScaler()
    scaler.fit(train_df[COLUMNS_TO_SCALE])
    return scaler


def apply_scaler(df: pd.DataFrame, scaler: StandardScaler) -> pd.DataFrame:
    """Apply an already-fitted scaler to 'Time' and 'Amount' columns, returning a new dataframe."""
    scaled_df = df.copy()
    scaled_df[COLUMNS_TO_SCALE] = scaler.transform(df[COLUMNS_TO_SCALE])
    return scaled_df


def preprocess(
    df: pd.DataFrame, test_size: float = 0.2, random_state: int = 42
):
    """
    Full preprocessing pipeline:
      1. Stratified train/test split
      2. Fit StandardScaler on training 'Time'/'Amount'
      3. Apply scaler to both splits
      4. Split into X (features) / y (labels)

    Returns:
        X_train, X_test, y_train, y_test, scaler, feature_columns
    """
    train_df, test_df = split_dataset(df, test_size=test_size, random_state=random_state)

    scaler = fit_scaler(train_df)
    train_scaled = apply_scaler(train_df, scaler)
    test_scaled = apply_scaler(test_df, scaler)

    feature_columns = get_feature_columns(df)

    X_train = train_scaled[feature_columns]
    y_train = train_scaled[TARGET_COLUMN]
    X_test = test_scaled[feature_columns]
    y_test = test_scaled[TARGET_COLUMN]

    return X_train, X_test, y_train, y_test, scaler, feature_columns