"""Session 16 - Will this field need irrigation within the next 24 hours? (binary classifier)

Covers: feature selection, imbalanced classes (class_weight), TimeSeriesSplit cross-validation,
overfitting check (train vs test), comparison with a simple agronomic rule, permutation importance.
Optional: SHAP (pip install shap) for per-decision explanations.
Data: real KFS weather 2023-2024 + simulated irrigated field (FAO-56 water balance, common.py)
Output: outputs/irrigation_model.joblib - used by the Smart Farm web app (session 17).
"""
import joblib
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.inspection import permutation_importance
from sklearn.metrics import balanced_accuracy_score, classification_report, confusion_matrix, f1_score
from sklearn.model_selection import TimeSeriesSplit, cross_val_score

from common import (OUT, RAW_MM, REFILL_POINT, ROOT_DEPTH_M, THETA_FC, build_irrigation_dataset, load_hourly,
                    simulate_irrigated_field)

H = 24
df = simulate_irrigated_field(load_hourly())

X, y = build_irrigation_dataset(df, H)   # see common.py: features known NOW + label = irrigation within 24 h
print("Class balance (1 = irrigate within 24 h):", y.value_counts(normalize=True).round(3).to_dict())

X_tr, X_te = X.loc[:"2023-12-31"], X.loc["2024-01-01":]
y_tr, y_te = y.loc[X_tr.index], y.loc[X_te.index]

# ---- Baseline: simple threshold rule ("irrigate when VWC is close to the refill point")
warn_level = REFILL_POINT + 0.25 * RAW_MM / (1000 * ROOT_DEPTH_M)
rule = (X_te["soil_vwc"] < warn_level).astype(int)
print(f"Rule baseline (VWC < {warn_level:.3f})  F1={f1_score(y_te, rule):.3f}  "
      f"balanced acc={balanced_accuracy_score(y_te, rule):.3f}")

model = RandomForestClassifier(n_estimators=250, max_depth=12, min_samples_leaf=10,
                               class_weight="balanced", n_jobs=-1, random_state=0)
cv = cross_val_score(model, X_tr, y_tr, cv=TimeSeriesSplit(n_splits=5), scoring="f1", n_jobs=-1)
print("TimeSeriesSplit F1 (train year):", cv.round(3), "mean =", round(cv.mean(), 3))

model.fit(X_tr, y_tr)
pred = model.predict(X_te)
print(f"\nTrain F1={f1_score(y_tr, model.predict(X_tr)):.3f} | Test F1={f1_score(y_te, pred):.3f}"
      "   (a big gap would mean overfitting)")
print("Balanced accuracy:", round(balanced_accuracy_score(y_te, pred), 3))
print("Confusion matrix [[TN FP] [FN TP]]:\n", confusion_matrix(y_te, pred))
print(classification_report(y_te, pred, target_names=["wait", "irrigate"]))

imp = permutation_importance(model, X_te, y_te, n_repeats=5, random_state=0, n_jobs=-1, scoring="f1")
ranking = pd.Series(imp.importances_mean, index=X.columns).sort_values(ascending=False).round(4)
print("Permutation importance (drop in F1 when shuffled):\n", ranking.to_string())

joblib.dump({"model": model, "features": list(X.columns), "refill_point": REFILL_POINT,
             "field_capacity": THETA_FC, "horizon_h": H}, OUT / "irrigation_model.joblib", compress=3)
print("\nSaved ->", OUT / "irrigation_model.joblib")
print("Irrigation events in 2024 (simulated):", int((df.loc['2024', 'irrigation_mm'] > 0).sum()),
      "| water applied:", round(df.loc['2024', 'irrigation_mm'].sum()), "mm")

try:
    import shap  # noqa: F401
    print("SHAP available: shap.TreeExplainer(model).shap_values(X_te[:200]) in Jupyter.")
except ImportError:
    print("Tip: pip install shap  - to explain each individual irrigation decision.")