import pandas as pd
import torch
from torch.utils.data import TensorDataset, DataLoader
from sklearn.preprocessing import StandardScaler


def build_features(df: pd.DataFrame):
    """Prepare the processed PaySim frame for training.

    - Splits X/y and splits chronologically by step (no leakage)
    - Scales features (fit on train, transform test)
    - Returns train/test DataLoaders and the fitted scaler
    """
    X = df.drop(columns=['isFraud'])
    y = df['isFraud']

    train_mask = df['step'] < 600
    X_train, y_train = X[train_mask], y[train_mask]
    X_test,  y_test  = X[~train_mask], y[~train_mask]

    scaler = StandardScaler()
    X_train = scaler.fit_transform(X_train)   # fit on train only
    X_test  = scaler.transform(X_test)        # transform test with train stats

    def to_ds(Xa, ys):
        return TensorDataset(
            torch.tensor(Xa, dtype=torch.float32),
            torch.tensor(ys.to_numpy(), dtype=torch.float32).unsqueeze(1),
        )

    train_dataloader = DataLoader(to_ds(X_train, y_train), batch_size=32, shuffle=True)
    test_dataloader  = DataLoader(to_ds(X_test,  y_test),  batch_size=32, shuffle=False)

    return train_dataloader, test_dataloader, scaler