src/data/loader.py

"""
Production data loading and validation module.

Handles reading transaction data from multiple sources (Parquet, CSV, Delta Lake)
with schema validation, data quality checks, and logging.
"""

import logging
from pathlib import Path
from typing import Optional

import numpy as np
import pandas as pd

logger = logging.getLogger(__name__)

EXPECTED_SCHEMA = {
    "transaction_id": "object",
    "timestamp": "datetime64[ns]",
    "cardholder_id": "object",
    "card_type": "object",
    "merchant_id": "object",
    "merchant_category_code": "object",
    "merchant_risk_score": "float64",
    "transaction_amount": "float64",
    "entry_mode": "object",
    "country_code": "object",
    "is_fraud": "int64",
}

REQUIRED_COLUMNS = list(EXPECTED_SCHEMA.keys())


class DataValidationError(Exception):
    """Raised when data fails validation checks."""


class TransactionDataLoader:
    """Production-grade data loader with validation and quality checks."""

    def __init__(self, data_path: str | Path):
        self.data_path = Path(data_path)
        if not self.data_path.exists():
            raise FileNotFoundError(f"Data file not found: {self.data_path}")

    def load(self, validate: bool = True, sample_frac: Optional[float] = None) -> pd.DataFrame:
        """Load transaction data from file.

        Args:
            validate: Whether to run schema and quality validation.
            sample_frac: Optional fraction to sample for development/testing.

        Returns:
            DataFrame with validated transaction data.
        """
        logger.info("Loading data from %s", self.data_path)

        if self.data_path.suffix == ".parquet":
            df = pd.read_parquet(self.data_path)
        elif self.data_path.suffix == ".csv":
            df = pd.read_csv(self.data_path, parse_dates=["timestamp"])
        else:
            raise ValueError(f"Unsupported file format: {self.data_path.suffix}")

        logger.info("Loaded %d rows, %d columns", len(df), len(df.columns))

        if sample_frac is not None:
            original_len = len(df)
            df = df.sample(frac=sample_frac, random_state=42).reset_index(drop=True)
            logger.info("Sampled %.1f%% — %d → %d rows", sample_frac * 100, original_len, len(df))

        if validate:
            self._validate_schema(df)
            self._validate_quality(df)

        return df

    def _validate_schema(self, df: pd.DataFrame) -> None:
        """Check that all required columns are present."""
        missing = set(REQUIRED_COLUMNS) - set(df.columns)
        if missing:
            raise DataValidationError(f"Missing required columns: {missing}")
        logger.info("Schema validation passed — all %d required columns present", len(REQUIRED_COLUMNS))

    def _validate_quality(self, df: pd.DataFrame) -> None:
        """Run data quality checks and log warnings."""
        issues = []

        # Check for nulls in critical columns
        critical_cols = ["transaction_id", "cardholder_id", "transaction_amount", "timestamp"]
        for col in critical_cols:
            null_count = df[col].isna().sum()
            if null_count > 0:
                issues.append(f"Column '{col}' has {null_count} null values")

        # Check for duplicate transaction IDs
        dup_count = df["transaction_id"].duplicated().sum()
        if dup_count > 0:
            issues.append(f"{dup_count} duplicate transaction IDs found")

        # Check for negative amounts
        neg_count = (df["transaction_amount"] < 0).sum()
        if neg_count > 0:
            issues.append(f"{neg_count} transactions with negative amounts")

        # Check fraud rate is reasonable (0.1% - 10%)
        fraud_rate = df["is_fraud"].mean()
        if fraud_rate < 0.001 or fraud_rate > 0.10:
            issues.append(f"Unusual fraud rate: {fraud_rate:.4%}")

        # Check timestamp range
        ts_range = (df["timestamp"].max() - df["timestamp"].min()).days
        if ts_range < 1:
            issues.append(f"Very narrow timestamp range: {ts_range} days")

        if issues:
            for issue in issues:
                logger.warning("Data quality issue: %s", issue)
        else:
            logger.info("Data quality validation passed")


def load_and_split(
    data_path: str | Path,
    test_size: float = 0.2,
    validation_size: float = 0.1,
    random_state: int = 42,
    stratify: bool = True,
) -> tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
    """Load data and split into train/validation/test sets.

    Uses time-based splitting for production realism — train on older data,
    test on newer data to prevent data leakage.

    Args:
        data_path: Path to the transaction data file.
        test_size: Fraction of data for the test set.
        validation_size: Fraction of data for the validation set.
        random_state: Random seed.
        stratify: Whether to use time-based stratification.

    Returns:
        Tuple of (train_df, val_df, test_df).
    """
    loader = TransactionDataLoader(data_path)
    df = loader.load(validate=True)

    # Time-based split to prevent data leakage
    df = df.sort_values("timestamp").reset_index(drop=True)
    n = len(df)

    train_end = int(n * (1 - test_size - validation_size))
    val_end = int(n * (1 - test_size))

    train_df = df.iloc[:train_end].copy()
    val_df = df.iloc[train_end:val_end].copy()
    test_df = df.iloc[val_end:].copy()

    logger.info(
        "Data split — Train: %d (%.1f%%), Val: %d (%.1f%%), Test: %d (%.1f%%)",
        len(train_df), len(train_df) / n * 100,
        len(val_df), len(val_df) / n * 100,
        len(test_df), len(test_df) / n * 100,
    )

    # Log fraud rates per split
    for name, split_df in [("Train", train_df), ("Val", val_df), ("Test", test_df)]:
        fraud_rate = split_df["is_fraud"].mean()
        logger.info("  %s fraud rate: %.4f%%", name, fraud_rate * 100)

    return train_df, val_df, test_df