cleaner.py

"""
Data cleaning pipeline.

Steps: dedup → timezone normalisation → test-transaction filter →
       amount sanity → label maturity filter → schema coercion.
"""
from __future__ import annotations

import structlog

import polars as pl

log = structlog.get_logger(__name__)


def dedup(df: pl.DataFrame, key: str = "transaction_id") -> pl.DataFrame:
    """Drop duplicate transaction IDs, keeping first occurrence."""
    before = len(df)
    if key not in df.columns:
        return df
    df = df.unique(subset=[key], keep="first")
    dropped = before - len(df)
    if dropped:
        log.info("dedup", dropped=dropped)
    return df


def normalise_timestamps(df: pl.DataFrame, ts_col: str = "timestamp_utc") -> pl.DataFrame:
    """Parse timestamp to UTC datetime; drop rows with null timestamps."""
    if ts_col not in df.columns:
        return df
    df = df.with_columns(
        pl.col(ts_col).str.to_datetime(strict=False, time_unit="us").alias(ts_col)
    )
    before = len(df)
    df = df.drop_nulls(subset=[ts_col])
    dropped = before - len(df)
    if dropped:
        log.warning("timestamp_nulls_dropped", count=dropped)
    return df


def filter_test_transactions(df: pl.DataFrame) -> pl.DataFrame:
    """
    Remove known test transactions.
    Heuristics: amount exactly ₹1.00, ₹0.01, or transaction_id starts with 'test_'.
    """
    before = len(df)
    filters = []
    if "amount_inr" in df.columns:
        filters.append(~pl.col("amount_inr").is_in([1.0, 0.01]))
    if "transaction_id" in df.columns:
        filters.append(~pl.col("transaction_id").str.starts_with("test_"))

    if filters:
        combined = filters[0]
        for f in filters[1:]:
            combined = combined & f
        df = df.filter(combined)

    dropped = before - len(df)
    if dropped:
        log.info("test_transactions_removed", count=dropped)
    return df


def filter_amount_outliers(
    df: pl.DataFrame,
    col: str = "amount_inr",
    min_amount: float = 0.01,
    max_amount: float = 10_000_000.0,
) -> pl.DataFrame:
    """Drop transactions with amount outside plausible range."""
    if col not in df.columns:
        return df
    before = len(df)
    df = df.filter(pl.col(col).is_between(min_amount, max_amount))
    dropped = before - len(df)
    if dropped:
        log.info("amount_outliers_dropped", col=col, count=dropped)
    return df


def filter_label_maturity(
    df: pl.DataFrame,
    ts_col: str = "timestamp_utc",
    cutoff_days: int = 14,
) -> pl.DataFrame:
    """
    Exclude transactions from the last `cutoff_days` days for training.
    Labels still arriving; including them causes false negatives in training set.
    Only applies when df has is_fraud column (training mode).
    """
    if "is_fraud" not in df.columns or ts_col not in df.columns:
        return df

    max_ts = df[ts_col].max()
    if max_ts is None:
        return df

    cutoff = max_ts - pl.duration(days=cutoff_days)
    before = len(df)
    df = df.filter(pl.col(ts_col) < cutoff)
    dropped = before - len(df)
    if dropped:
        log.info("label_maturity_filter", cutoff_days=cutoff_days, rows_dropped=dropped)
    return df


def coerce_schema(df: pl.DataFrame) -> pl.DataFrame:
    """Cast columns to expected dtypes; fill nulls in numeric columns with 0."""
    float_cols = ["amount_inr", "amount_usd", "biometric_quality_score", "session_duration_seconds"]
    int_cols = ["is_fraud", "biometric_attempt_count", "bureau_pull_count_30d"]

    for col in float_cols:
        if col in df.columns:
            df = df.with_columns(
                pl.col(col).cast(pl.Float64, strict=False).fill_null(0.0)
            )
    for col in int_cols:
        if col in df.columns:
            df = df.with_columns(
                pl.col(col).cast(pl.Int8, strict=False).fill_null(0)
            )
    return df


def clean(
    df: pl.DataFrame,
    cutoff_days: int = 14,
    is_training: bool = True,
) -> pl.DataFrame:
    """Full cleaning pipeline."""
    df = dedup(df)
    df = normalise_timestamps(df)
    df = filter_test_transactions(df)
    df = filter_amount_outliers(df)
    if is_training:
        df = filter_label_maturity(df, cutoff_days=cutoff_days)
    df = coerce_schema(df)

    log.info(
        "cleaning_complete",
        rows=len(df),
        fraud_rate=round(df["is_fraud"].mean(), 4) if "is_fraud" in df.columns else None,
    )
    return df