"""Session 15 - Machine Learning basics in agriculture.

Part A  Regression: predict the field soil moisture 6 hours ahead  (Linear vs Random Forest, MAE / RMSE)
Part B  Regression: predict TOMORROW's reference evapotranspiration ET0 from today's real weather
Part C  Classification: soil state (dry / moist / wet)  (Accuracy + report)

Data: real KFS-farm weather 2023-2024 (Open-Meteo) + simulated irrigated field (FAO-56, see common.py)
Time-based split: train = 2023, test = 2024  (never shuffle time series!)
"""
import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.linear_model import LinearRegression
from sklearn.metrics import accuracy_score, classification_report, mean_absolute_error, mean_squared_error

from common import KC, OUT, load_hourly, simulate_irrigated_field, soil_state

HORIZON = 6
df = simulate_irrigated_field(load_hourly())

# ---------------- Part A: field VWC in 6 h
feat = pd.DataFrame(index=df.index)
feat["vwc_now"] = df["field_vwc"]
for lag in (1, 3, 6, 12):
    feat[f"vwc_lag{lag}"] = df["field_vwc"].shift(lag)
feat["vwc_change_6h"] = df["field_vwc"].diff(6)
feat["temp"] = df["air_temp_c"]
feat["humidity"] = df["air_humidity_pct"]
feat["radiation"] = df["solar_radiation_wm2"]
feat["wind"] = df["wind_speed_kmh"]
feat["et0_next6h_proxy"] = df["et0_mm"].rolling(6).sum()   # last 6 h demand
feat["hour_sin"] = np.sin(2 * np.pi * df.index.hour / 24)
feat["hour_cos"] = np.cos(2 * np.pi * df.index.hour / 24)
feat["target"] = df["field_vwc"].shift(-HORIZON)
feat = feat.dropna()
# Irrigation events are planned by the farm manager, not by the weather -> exclude windows containing one
irrigated_soon = df["irrigation_mm"][::-1].rolling(HORIZON, min_periods=1).max()[::-1].shift(-1).reindex(feat.index) > 0
feat = feat[~irrigated_soon]

train, test = feat.loc[:"2023-12-31"], feat.loc["2024-01-01":]
X_cols = [c for c in feat.columns if c != "target"]
print(f"Part A - field VWC {HORIZON} h ahead | train={len(train)} test={len(test)}")
models = {"Persistence (baseline)": None, "LinearRegression": LinearRegression(),
          "RandomForest": RandomForestRegressor(n_estimators=200, min_samples_leaf=5, n_jobs=-1, random_state=0)}
for name, model in models.items():
    if model is None:
        pred = test["vwc_now"]
    else:
        model.fit(train[X_cols], train["target"])
        pred = model.predict(test[X_cols])
    print(f"  {name:24s} MAE={mean_absolute_error(test.target, pred):.4f}  "
          f"RMSE={np.sqrt(mean_squared_error(test.target, pred)):.4f} m3/m3")
joblib.dump({"model": models["RandomForest"], "features": X_cols, "horizon_h": HORIZON},
            OUT / "soil_moisture_6h_rf.joblib", compress=3)

# ---------------- Part B: tomorrow's ET0 (real data only)
daily = df.resample("D").agg({"et0_mm": "sum", "air_temp_c": ["mean", "max"], "air_humidity_pct": "mean",
                              "solar_radiation_wm2": "mean", "wind_speed_kmh": "mean"})
daily.columns = ["et0", "t_mean", "t_max", "rh", "rad", "wind"]
dfe = daily.add_suffix("_today")
dfe["et0_3d"] = daily["et0"].rolling(3).mean()
dfe["doy_sin"] = np.sin(2 * np.pi * daily.index.dayofyear / 365)
dfe["doy_cos"] = np.cos(2 * np.pi * daily.index.dayofyear / 365)
dfe["target"] = daily["et0"].shift(-1)
dfe = dfe.dropna()
e_tr, e_te = dfe.loc[:"2023-12-31"], dfe.loc["2024-01-01":]
e_cols = [c for c in dfe.columns if c != "target"]
print("\nPart B - tomorrow's ET0 (mm/day), real weather")
print(f"  {'Persistence (baseline)':24s} MAE={mean_absolute_error(e_te.target, e_te.et0_today):.3f} mm")
for name, m in {"LinearRegression": LinearRegression(),
                "RandomForest": RandomForestRegressor(n_estimators=300, min_samples_leaf=3, random_state=0)}.items():
    m.fit(e_tr[e_cols], e_tr.target)
    p_ = m.predict(e_te[e_cols])
    print(f"  {name:24s} MAE={mean_absolute_error(e_te.target, p_):.3f} mm  "
          f"RMSE={np.sqrt(mean_squared_error(e_te.target, p_)):.3f} mm")
joblib.dump({"model": m, "features": e_cols}, OUT / "et0_next_day_rf.joblib", compress=3)
etc = float(p_[-1]) * KC
print(f"  Tomorrow's crop water need (Kc={KC}): {etc:.2f} mm  ->  {etc * 4.2:.1f} m3 per feddan")

# ---------------- Part C: soil state classification
states = df["field_vwc"].apply(soil_state)
Xc = pd.DataFrame({"depletion_proxy_et0_72h": df["et0_mm"].rolling(72).sum(),
                   "hours_since_irrigation": df.groupby((df["irrigation_mm"] > 0).cumsum()).cumcount(),
                   "temp": df["air_temp_c"], "humidity": df["air_humidity_pct"],
                   "radiation": df["solar_radiation_wm2"]}).dropna()
yc = states.loc[Xc.index]
clf = RandomForestClassifier(n_estimators=200, min_samples_leaf=5, random_state=0, n_jobs=-1)
clf.fit(Xc.loc[:"2023-12-31"], yc.loc[:"2023-12-31"])
pc = clf.predict(Xc.loc["2024-01-01":])
print("\nPart C - soil state WITHOUT the soil sensor (weather + irrigation log only)")
print("  Accuracy:", round(accuracy_score(yc.loc["2024-01-01":], pc), 3))
print(classification_report(yc.loc["2024-01-01":], pc, zero_division=0))