02_train_model.py

"""
GridGuard AI — Hierarchical Energy-Theft Detection Model
=========================================================
Implements the 3-level detection logic described in the project pitch:

  Level 1  Transformer-level anomaly   -> which feeders have abnormal loss%
  Level 2  Neighbourhood peer anomaly  -> which households deviate from peers
            of similar income class / house size within their neighbourhood
  Level 3  Household behavioural anomaly -> consumption drop + Isolation
            Forest unsupervised score, combined into a final Fraud Risk Score

Output: data/household_risk_scores.csv  (used by the Gradio app)
        models/isolation_forest.joblib, models/scaler.joblib
        models/xgb_classifier.joblib (supervised — trained on ground truth,
        included to SHOW evaluation metrics in your presentation; in a real
        deployment you would not have these labels and would lean on the
        unsupervised + rules layers)
"""

import os
import numpy as np
import pandas as pd
import joblib
from sklearn.ensemble import IsolationForest
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
from sklearn.metrics import precision_score, recall_score, f1_score, roc_auc_score
import xgboost as xgb

BASE = os.path.dirname(__file__)
DATA_DIR = os.path.join(BASE, "data")
MODEL_DIR = os.path.join(BASE, "models")
os.makedirs(MODEL_DIR, exist_ok=True)

households = pd.read_csv(os.path.join(DATA_DIR, "households.csv"))
readings = pd.read_csv(os.path.join(DATA_DIR, "readings_monthly.csv"))
tx_monthly = pd.read_csv(os.path.join(DATA_DIR, "transformers.csv"))

N_MONTHS = readings["month_idx"].max() + 1
LAST3 = sorted(readings["month_idx"].unique())[-3:]
PRIOR3 = sorted(readings["month_idx"].unique())[:3]

# ---------------------------------------------------------------------------
# Household-level behavioural features
# ---------------------------------------------------------------------------
last3 = readings[readings.month_idx.isin(LAST3)].groupby("household_id")["metered_kwh"].mean().rename("avg_last3_kwh")
prior3 = readings[readings.month_idx.isin(PRIOR3)].groupby("household_id")["metered_kwh"].mean().rename("avg_prior3_kwh")

# linear trend slope of metered consumption over time (theft -> negative slope)
def slope(group):
    x = group["month_idx"].values
    y = group["metered_kwh"].values
    if len(x) < 2:
        return 0.0
    return np.polyfit(x, y, 1)[0]

trend = readings.groupby("household_id").apply(slope).rename("consumption_trend_slope")

feat = households.merge(last3, on="household_id").merge(prior3, on="household_id").merge(trend, on="household_id")
feat["drop_ratio"] = (feat["avg_prior3_kwh"] - feat["avg_last3_kwh"]) / feat["avg_prior3_kwh"].replace(0, np.nan)
feat["drop_ratio"] = feat["drop_ratio"].fillna(0)

# ---------------------------------------------------------------------------
# Level 2 — Neighbourhood peer comparison (z-score vs similar peers)
# Peer group = same neighbourhood_id + same income_class
# ---------------------------------------------------------------------------
peer_stats = feat.groupby(["neighbourhood_id", "income_class"])["avg_last3_kwh"].agg(["mean", "std"]).rename(
    columns={"mean": "peer_mean_kwh", "std": "peer_std_kwh"}
)
feat = feat.merge(peer_stats, on=["neighbourhood_id", "income_class"], how="left")
feat["peer_std_kwh"] = feat["peer_std_kwh"].replace(0, np.nan)
feat["peer_zscore"] = (feat["avg_last3_kwh"] - feat["peer_mean_kwh"]) / feat["peer_std_kwh"]
feat["peer_zscore"] = feat["peer_zscore"].fillna(0)

# ---------------------------------------------------------------------------
# Level 1 — Transformer-level loss flag (propagated down to its households)
# ---------------------------------------------------------------------------
latest_tx = tx_monthly[tx_monthly.month_idx == tx_monthly.month_idx.max()][["transformer_id", "loss_pct"]]
latest_tx = latest_tx.rename(columns={"loss_pct": "transformer_loss_pct"})
feat = feat.merge(latest_tx, on="transformer_id", how="left")
feat["transformer_flagged"] = feat["transformer_loss_pct"] > 0.08  # >8% loss = above normal technical loss

# ---------------------------------------------------------------------------
# Level 3 — Unsupervised anomaly score (Isolation Forest)
# ---------------------------------------------------------------------------
unsup_cols = ["drop_ratio", "peer_zscore", "consumption_trend_slope", "transformer_loss_pct"]
X = feat[unsup_cols].fillna(0).values
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)

iso = IsolationForest(n_estimators=300, contamination=0.08, random_state=42)
iso.fit(X_scaled)
# decision_function: higher = more normal. Flip & rescale to 0-1 "anomaly score"
raw_score = -iso.decision_function(X_scaled)
feat["anomaly_score"] = (raw_score - raw_score.min()) / (raw_score.max() - raw_score.min())

# ---------------------------------------------------------------------------
# Combine into a single Fraud Risk Score (0-100), used for triage
# Weighted blend: behavioural anomaly + peer deviation + transformer context
# ---------------------------------------------------------------------------
norm_drop = np.clip(feat["drop_ratio"], 0, 1)
norm_peer = np.clip(-feat["peer_zscore"] / 3, 0, 1)          # big NEGATIVE z-score (below peers) = risky
tx_boost = feat["transformer_flagged"].astype(float) * 0.15

feat["fraud_risk_score"] = np.clip(
    100 * (0.45 * feat["anomaly_score"] + 0.30 * norm_drop + 0.25 * norm_peer + tx_boost), 0, 100
)

feat["risk_tier"] = pd.cut(
    feat["fraud_risk_score"], bins=[-1, 40, 70, 100], labels=["Low", "Medium", "High"]
)

# ---------------------------------------------------------------------------
# Supervised model — trained on ground truth, FOR EVALUATION / DEMO ONLY
# (shows precision/recall you can quote in the presentation)
# ---------------------------------------------------------------------------
y = feat["is_fraud_ground_truth"].astype(int)
X_train, X_test, y_train, y_test = train_test_split(X_scaled, y, test_size=0.3, stratify=y, random_state=42)

xgb_clf = xgb.XGBClassifier(
    n_estimators=150, max_depth=3, learning_rate=0.08,
    scale_pos_weight=(y_train == 0).sum() / max((y_train == 1).sum(), 1),
    eval_metric="logloss", random_state=42,
)
xgb_clf.fit(X_train, y_train)
y_pred = xgb_clf.predict(X_test)
y_proba = xgb_clf.predict_proba(X_test)[:, 1]

print("=== Supervised model evaluation (ground-truth labels, demo only) ===")
print(f"Precision : {precision_score(y_test, y_pred):.3f}")
print(f"Recall    : {recall_score(y_test, y_pred):.3f}")
print(f"F1        : {f1_score(y_test, y_pred):.3f}")
print(f"ROC-AUC   : {roc_auc_score(y_test, y_proba):.3f}")

# Unsupervised-only evaluation (what you'd realistically have in production, no labels)
high_risk = feat["risk_tier"] == "High"
print("\n=== Unsupervised risk-tier vs ground truth (sanity check) ===")
print(f"Households flagged High risk : {high_risk.sum()} ({high_risk.mean()*100:.1f}% of base)")
print(f"Of those, actually fraud      : {feat.loc[high_risk, 'is_fraud_ground_truth'].mean()*100:.1f}%")
print(f"Fraud households caught in High/Medium tier: "
      f"{feat.loc[feat.is_fraud_ground_truth, 'risk_tier'].isin(['High','Medium']).mean()*100:.1f}%")

# ---------------------------------------------------------------------------
# Save artifacts
# ---------------------------------------------------------------------------
feat_out = feat.drop(columns=["fraud_start_month_idx", "bypass_severity", "is_genuine_low_consumer", "is_lifestyle_change"])  # keep ground truth flag only, hide leak cols
feat_out.to_csv(os.path.join(DATA_DIR, "household_risk_scores.csv"), index=False)
tx_monthly.to_csv(os.path.join(DATA_DIR, "transformers.csv"), index=False)

joblib.dump(iso, os.path.join(MODEL_DIR, "isolation_forest.joblib"))
joblib.dump(scaler, os.path.join(MODEL_DIR, "scaler.joblib"))
joblib.dump(xgb_clf, os.path.join(MODEL_DIR, "xgb_classifier.joblib"))
joblib.dump(unsup_cols, os.path.join(MODEL_DIR, "feature_cols.joblib"))

print(f"\nSaved: {os.path.join(DATA_DIR, 'household_risk_scores.csv')}")
print(f"Saved models to: {MODEL_DIR}")