test_preprocess.py

"""
Test script to verify preprocessing pipeline on sample data.
"""

import polars as pl
import sys

def test_data_loading():
    """Test loading sample data."""
    print("="*60)
    print("TEST 1: Data Loading")
    print("="*60)

    try:
        # Load sample transaction data
        transactions = pl.scan_csv(
            "./preliminary_data/acct_transaction_sample.csv",
            schema_overrides={
                "from_acct": pl.Utf8,
                "from_acct_type": pl.Utf8,
                "to_acct": pl.Utf8,
                "to_acct_type": pl.Utf8,
                "is_self_txn": pl.Utf8,
                "txn_amt": pl.Float64,
                "txn_date": pl.Int64,
                "txn_time": pl.Utf8,
                "currency_type": pl.Utf8,
                "channel_type": pl.Utf8,
            }
        )

        alerts = pl.read_csv(
            "./preliminary_data/acct_alert.csv",
            schema_overrides={
                "acct": pl.Utf8,
                "event_date": pl.Int64
            }
        )

        df = transactions.collect()
        print(f"✓ Loaded {df.height} sample transactions")
        print(f"✓ Loaded {alerts.height} alert accounts")
        print(f"✓ Transaction columns: {df.columns}")

        return df, alerts

    except Exception as e:
        print(f"✗ Error loading data: {e}")
        return None, None


def test_esun_filtering(df):
    """Test E.SUN account filtering."""
    print("\n" + "="*60)
    print("TEST 2: E.SUN Account Filtering")
    print("="*60)

    try:
        # Get from_acct where type='01'
        from_accounts = (
            df.filter(pl.col("from_acct_type") == "01")
            .select("from_acct")
            .unique()
            .get_column("from_acct")
            .to_list()
        )

        # Get to_acct where type='01'
        to_accounts = (
            df.filter(pl.col("to_acct_type") == "01")
            .select("to_acct")
            .unique()
            .get_column("to_acct")
            .to_list()
        )

        esun_accounts = set(from_accounts + to_accounts)
        print(f"✓ Found {len(from_accounts)} unique from_acct with type='01'")
        print(f"✓ Found {len(to_accounts)} unique to_acct with type='01'")
        print(f"✓ Total unique E.SUN accounts: {len(esun_accounts)}")

        # Filter transactions
        df_filtered = df.filter(
            pl.col("from_acct").is_in(list(esun_accounts)) |
            pl.col("to_acct").is_in(list(esun_accounts))
        )

        print(f"✓ Kept {df_filtered.height}/{df.height} transactions ({df_filtered.height/df.height*100:.1f}%)")

        return df_filtered, esun_accounts

    except Exception as e:
        print(f"✗ Error filtering accounts: {e}")
        return None, None


def test_feature_engineering(df, alerts):
    """Test feature engineering."""
    print("\n" + "="*60)
    print("TEST 3: Feature Engineering")
    print("="*60)

    try:
        # Convert time to seconds
        def time_to_seconds(time_str):
            try:
                h, m, s = map(int, time_str.split(':'))
                return h * 3600 + m * 60 + s
            except:
                return 0

        df = df.with_columns([
            pl.col("txn_time").map_elements(time_to_seconds, return_dtype=pl.Int64).alias("txn_time_seconds")
        ])
        print("✓ Converted txn_time to seconds")

        # Create labels
        alert_set = set(alerts.get_column("acct").to_list())
        df = df.with_columns([
            pl.col("from_acct").map_elements(lambda x: 1 if x in alert_set else 0, return_dtype=pl.Int64).alias("label")
        ])

        num_alerts = df.filter(pl.col("label") == 1).get_column("from_acct").n_unique()
        num_normal = df.filter(pl.col("label") == 0).get_column("from_acct").n_unique()
        print(f"✓ Created labels: {num_alerts} alert accounts, {num_normal} normal accounts")

        # One-hot encode is_self_txn
        df = df.with_columns([
            (pl.col("is_self_txn") == "Y").cast(pl.Int64).alias("is_self_txn_Y"),
            (pl.col("is_self_txn") == "N").cast(pl.Int64).alias("is_self_txn_N"),
            (pl.col("is_self_txn") == "UNK").cast(pl.Int64).alias("is_self_txn_UNK"),
        ])
        print("✓ One-hot encoded is_self_txn")

        # Binary encode acct_type
        df = df.with_columns([
            (pl.col("from_acct_type") == "02").cast(pl.Int64).alias("from_acct_type_binary"),
            (pl.col("to_acct_type") == "02").cast(pl.Int64).alias("to_acct_type_binary"),
        ])
        print("✓ Binary encoded acct_type")

        # Sort and calculate delta features
        df = df.sort(["from_acct", "txn_date", "txn_time_seconds"])

        df = df.with_columns([
            pl.col("txn_date").diff().over("from_acct").fill_null(0).alias("delta_days"),
            pl.col("txn_time_seconds").diff().over("from_acct").fill_null(0).alias("delta_time_raw"),
        ])

        df = df.with_columns([
            pl.when(pl.col("delta_time_raw") < 0)
            .then(pl.col("delta_days") * 86400 + pl.col("delta_time_raw"))
            .otherwise(pl.col("delta_time_raw"))
            .alias("delta_time_seconds")
        ])
        print("✓ Calculated delta time features")

        # Z-score normalization
        amt_mean = df.get_column("txn_amt").mean()
        amt_std = df.get_column("txn_amt").std()

        df = df.with_columns([
            ((pl.col("txn_amt") - amt_mean) / amt_std).alias("txn_amt_normalized")
        ])
        print(f"✓ Normalized txn_amt (mean={amt_mean:.2f}, std={amt_std:.2f})")

        print(f"✓ Final feature shape: {df.shape}")
        print(f"✓ Feature columns: {len(df.columns)}")

        return df

    except Exception as e:
        print(f"✗ Error in feature engineering: {e}")
        import traceback
        traceback.print_exc()
        return None


def test_time_split(df):
    """Test time-based split."""
    print("\n" + "="*60)
    print("TEST 4: Time-Based Split")
    print("="*60)

    try:
        train_df = df.filter(pl.col("txn_date") <= 90)
        val_df = df.filter(pl.col("txn_date") > 90)

        print(f"✓ Train set: {train_df.height} transactions (days 1-90)")
        print(f"✓ Val set: {val_df.height} transactions (days 91+)")

        train_alerts = train_df.filter(pl.col("label") == 1).get_column("from_acct").n_unique()
        val_alerts = val_df.filter(pl.col("label") == 1).get_column("from_acct").n_unique()

        print(f"✓ Train alert accounts: {train_alerts}")
        print(f"✓ Val alert accounts: {val_alerts}")

        return True

    except Exception as e:
        print(f"✗ Error in time split: {e}")
        return False


def main():
    print("\n" + "="*60)
    print("PREPROCESSING PIPELINE TEST")
    print("="*60 + "\n")

    # Test 1: Load data
    df, alerts = test_data_loading()
    if df is None:
        print("\n✗ FAILED: Could not load data")
        sys.exit(1)

    # Test 2: Filter E.SUN accounts
    df_filtered, esun_accounts = test_esun_filtering(df)
    if df_filtered is None:
        print("\n✗ FAILED: Could not filter E.SUN accounts")
        sys.exit(1)

    # Test 3: Feature engineering
    df_features = test_feature_engineering(df_filtered, alerts)
    if df_features is None:
        print("\n✗ FAILED: Could not engineer features")
        sys.exit(1)

    # Test 4: Time-based split
    success = test_time_split(df_features)
    if not success:
        print("\n✗ FAILED: Could not perform time split")
        sys.exit(1)

    print("\n" + "="*60)
    print("✓ ALL TESTS PASSED!")
    print("="*60)
    print("\nThe preprocessing pipeline is working correctly.")
    print("You can now run the full preprocessing with:")
    print("  python data_preprocess.py")


if __name__ == "__main__":
    main()