src/data_loader.py

"""
Data Loader Module - Fixed
- Config-driven paths
- Validates schema
- Labels both source AND target in fraud_flag
- Temporal features
- No Streamlit dependency
"""
import os
import functools
import pandas as pd
from datetime import datetime

from src.config_loader import get_config

REQUIRED_COLUMNS = ['Account', 'Account.1', 'Amount Paid', 'Timestamp', 'Payment Format', 'Is Laundering']
CHECKPOINT_EDGE_COLUMNS = ['source', 'target', 'amount', 'timestamp', 'payment_type', 'is_laundering']


def validate_schema(df: pd.DataFrame) -> None:
    missing = [c for c in REQUIRED_COLUMNS if c not in df.columns]
    if missing:
        raise ValueError(f"Raw dataset is missing required columns: {missing}")


def load_raw() -> pd.DataFrame:
    cfg = get_config()
    raw_path = cfg['data']['raw_path']
    df = pd.read_csv(
        raw_path,
        dtype={'Account': str, 'Account.1': str},
        parse_dates=['Timestamp'],
        low_memory=False
    )
    validate_schema(df)
    return df


def load_checkpoint_transactions() -> pd.DataFrame:
    """
    Recover transactions from the cached graph edge list when the raw CSV
    is not available on the deployment host.
    """
    project_root = os.path.dirname(os.path.dirname(__file__))
    checkpoint_path = os.path.join(project_root, 'data', 'checkpoints', 'graph_edges.parquet')
    if not os.path.exists(checkpoint_path):
        raise FileNotFoundError(checkpoint_path)

    df = pd.read_parquet(checkpoint_path)
    missing = [c for c in CHECKPOINT_EDGE_COLUMNS if c not in df.columns]
    if missing:
        raise ValueError(f"Checkpoint edge list is missing required columns: {missing}")

    df = df[CHECKPOINT_EDGE_COLUMNS].copy()
    df['source'] = df['source'].astype(str)
    df['target'] = df['target'].astype(str)
    df['amount'] = pd.to_numeric(df['amount'], errors='coerce')
    df['timestamp'] = pd.to_datetime(df['timestamp'], errors='coerce')
    df['payment_type'] = df['payment_type'].fillna('Unknown').astype(str)
    df['is_laundering'] = pd.to_numeric(df['is_laundering'], errors='coerce').fillna(0).astype(int)
    df = df.dropna(subset=['source', 'target', 'amount', 'timestamp'])
    return df.reset_index(drop=True)


def clean(df: pd.DataFrame) -> pd.DataFrame:
    cfg = get_config()
    df = df.rename(columns={
        'Account': 'source',
        'Account.1': 'target',
        'Amount Paid': 'amount',
        'Timestamp': 'timestamp',
        'Payment Format': 'payment_type',
        'Is Laundering': 'is_laundering',
        'From Bank': 'source_bank',
        'To Bank': 'target_bank',
    })
    df = df.dropna(subset=['source', 'target', 'amount', 'timestamp'])
    df = df[(df['source'] != df['target']) | (df['payment_type'] == 'Reinvestment')]
    df['is_laundering'] = df['is_laundering'].astype(int)
    df = df[df['amount'] > cfg['data']['min_amount_threshold']]
    df['timestamp'] = pd.to_datetime(df['timestamp'])
    return df.reset_index(drop=True)


def precompute_node_features(df: pd.DataFrame) -> pd.DataFrame:
    cfg = get_config()
    sent = df.groupby('source').agg(
        total_sent=('amount', 'sum'),
        count_sent=('source', 'count'),
        unique_recipients=('target', 'nunique'),
    )
    recv = df.groupby('target').agg(
        total_received=('amount', 'sum'),
        count_received=('target', 'count'),
        unique_senders=('source', 'nunique'),
    )
    node_features = sent.join(recv, how='outer').fillna(0)
    node_features.index.name = 'account'
    node_features['forward_ratio'] = node_features['total_sent'] / (node_features['total_received'] + 1)

    # FIX: label both source AND target of laundering transactions
    fraudulent_sources = set(df[df['is_laundering'] == 1]['source'].unique())
    fraudulent_targets = set(df[df['is_laundering'] == 1]['target'].unique())
    fraud_accounts = fraudulent_sources | fraudulent_targets
    node_features['fraud_flag'] = node_features.index.isin(fraud_accounts).astype(int)

    # Temporal features
    first_tx = df.groupby('source')['timestamp'].min().rename('first_tx')
    last_tx = df.groupby('source')['timestamp'].max().rename('last_tx')
    node_features = node_features.join(first_tx, how='left')
    node_features = node_features.join(last_tx, how='left')
    now = pd.Timestamp(datetime.now())
    node_features['account_age_days'] = (
        (node_features['last_tx'] - node_features['first_tx']).dt.days.fillna(0)
    )
    node_features['days_since_last_tx'] = (
        (now - node_features['last_tx']).dt.days.fillna(9999)
    )

    # Burst score: tx in last 3 days / total tx
    cutoff_3d = df['timestamp'].max() - pd.Timedelta(days=3)
    burst = df[df['timestamp'] >= cutoff_3d].groupby('source').size().rename('burst_count')
    node_features = node_features.join(burst, how='left')
    node_features['burst_score'] = node_features['burst_count'].fillna(0) / (node_features['count_sent'] + 1)

    node_features = node_features.drop(columns=['first_tx', 'last_tx', 'burst_count'], errors='ignore')

    os.makedirs(os.path.dirname(cfg['data']['processed_path']), exist_ok=True)
    node_features.to_parquet(cfg['data']['processed_path'])
    return node_features


@functools.lru_cache(maxsize=1)
def get_processed_data() -> tuple:
    cfg = get_config()
    tx_path = cfg['data']['transactions_path']
    nf_path = cfg['data']['processed_path']
    raw_path = cfg['data']['raw_path']

    if os.path.exists(tx_path):
        transactions_df = pd.read_parquet(tx_path)
    else:
        if os.path.exists(raw_path):
            raw_df = load_raw()
            transactions_df = clean(raw_df)
        else:
            transactions_df = load_checkpoint_transactions()
        os.makedirs(os.path.dirname(tx_path), exist_ok=True)
        transactions_df.to_parquet(tx_path)

    if os.path.exists(nf_path):
        node_features_df = pd.read_parquet(nf_path)
    else:
        node_features_df = precompute_node_features(transactions_df)

    return transactions_df, node_features_df