drift_detector.py

"""
Feature + model score drift detection via Evidently.

PSI thresholds by feature tier (from PLAN.md):
  Score:  > 0.10 → alert
  Tier 1: > 0.15 → alert
  Tier 2: > 0.20 → alert
  Tier 3: > 0.25 → alert

Check schedule:
  Model score: every 1h
  Tier 1 features: every 6h
  Tier 2 features: daily
  Tier 3 features: weekly
"""
from __future__ import annotations

import os
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

import numpy as np
import polars as pl
import structlog
from prometheus_client import Gauge

log = structlog.get_logger(__name__)

# Prometheus gauges for drift metrics
PSI_GAUGE = Gauge(
    "fraud_feature_psi",
    "Population Stability Index per feature",
    ["feature_name", "tier"],
)
SCORE_PSI_GAUGE = Gauge("fraud_score_psi", "PSI of model fraud score output")


@dataclass
class DriftReport:
    feature_name: str
    tier: str
    psi: float
    alert: bool
    threshold: float


TIER_PSI_THRESHOLDS = {
    "score": 0.10,
    "tier1": 0.15,
    "tier2": 0.20,
    "tier3": 0.25,
}

TIER1_FEATURES = [
    "card_txn_count_1h", "card_txn_count_24h", "user_txn_count_1h",
    "amount_log", "hour_of_day", "mcc_risk_tier",
]
TIER2_FEATURES = [
    "merchant_fraud_rate_30d", "user_amount_avg_30d", "agent_risk_score_t2",
    "bin_fraud_rate_30d", "user_account_age_days",
]
TIER3_FEATURES = [
    "user_txn_count_30d", "agent_days_since_onboarding",
]


def compute_psi(
    reference: np.ndarray,
    current: np.ndarray,
    n_bins: int = 10,
    epsilon: float = 1e-6,
) -> float:
    """
    Population Stability Index.
    PSI < 0.10: no significant change.
    PSI 0.10-0.20: moderate shift, monitor.
    PSI > 0.20: significant shift, investigate.
    PSI > 0.25: major shift, retrain.
    """
    # Build bins from reference distribution
    bins = np.percentile(reference, np.linspace(0, 100, n_bins + 1))
    bins = np.unique(bins)
    if len(bins) < 2:
        return 0.0

    ref_counts, _ = np.histogram(reference, bins=bins)
    cur_counts, _ = np.histogram(current, bins=bins)

    ref_pct = ref_counts / (ref_counts.sum() + epsilon)
    cur_pct = cur_counts / (cur_counts.sum() + epsilon)

    ref_pct = np.clip(ref_pct, epsilon, None)
    cur_pct = np.clip(cur_pct, epsilon, None)

    psi = float(np.sum((cur_pct - ref_pct) * np.log(cur_pct / ref_pct)))
    return psi


def check_feature_drift(
    reference_df: pl.DataFrame,
    current_df: pl.DataFrame,
    features: list[str],
    tier: str,
) -> list[DriftReport]:
    threshold = TIER_PSI_THRESHOLDS.get(tier, 0.20)
    reports = []

    for feature in features:
        if feature not in reference_df.columns or feature not in current_df.columns:
            continue
        ref = reference_df[feature].drop_nulls().to_numpy().astype(float)
        cur = current_df[feature].drop_nulls().to_numpy().astype(float)

        if len(ref) < 100 or len(cur) < 10:
            continue

        psi = compute_psi(ref, cur)
        alert = psi > threshold

        PSI_GAUGE.labels(feature_name=feature, tier=tier).set(psi)

        if alert:
            log.warning("drift_alert", feature=feature, tier=tier, psi=round(psi, 4),
                        threshold=threshold)

        reports.append(DriftReport(
            feature_name=feature,
            tier=tier,
            psi=psi,
            alert=alert,
            threshold=threshold,
        ))

    return reports


def check_score_drift(
    reference_scores: np.ndarray,
    current_scores: np.ndarray,
) -> DriftReport:
    psi = compute_psi(reference_scores, current_scores)
    threshold = TIER_PSI_THRESHOLDS["score"]
    SCORE_PSI_GAUGE.set(psi)

    if psi > threshold:
        log.warning("score_drift_alert", psi=round(psi, 4), threshold=threshold)

    return DriftReport(
        feature_name="fraud_score",
        tier="score",
        psi=psi,
        alert=psi > threshold,
        threshold=threshold,
    )


RETRAIN_PSI_THRESHOLD = 0.25


def trigger_retrain(reason: str) -> None:
    """Write sentinel file; cron/systemd watches path and runs fraud-train."""
    flag_path = os.path.expanduser(
        os.getenv("RETRAIN_FLAG_PATH", "~/kg-data/retrain_requested.flag")
    )
    os.makedirs(os.path.dirname(flag_path), exist_ok=True)
    with open(flag_path, "w") as f:
        f.write(f"{int(time.time())}\t{reason}\n")
    log.warning("retrain_triggered", reason=reason, flag_path=flag_path)


def full_drift_report(
    reference_df: pl.DataFrame,
    current_df: pl.DataFrame,
    reference_scores: Optional[np.ndarray] = None,
    current_scores: Optional[np.ndarray] = None,
) -> list[DriftReport]:
    reports = []
    reports.extend(check_feature_drift(reference_df, current_df, TIER1_FEATURES, "tier1"))
    reports.extend(check_feature_drift(reference_df, current_df, TIER2_FEATURES, "tier2"))
    reports.extend(check_feature_drift(reference_df, current_df, TIER3_FEATURES, "tier3"))

    if reference_scores is not None and current_scores is not None:
        reports.append(check_score_drift(reference_scores, current_scores))

    alerts = [r for r in reports if r.alert]
    log.info("drift_report_complete", total=len(reports), alerts=len(alerts))

    critical = [r for r in alerts if r.psi >= RETRAIN_PSI_THRESHOLD]
    if critical:
        worst = max(critical, key=lambda r: r.psi)
        trigger_retrain(
            f"psi={worst.psi:.4f} feature={worst.feature_name} tier={worst.tier}"
        )

    return reports