"""Shared helpers for the KFS Smart Farm AI projects (sessions 13-19).

Two data sources are used:
1) REAL weather & soil reanalysis (Open-Meteo, ERA5-Land based) for the KFS farm location, 2023-2024.
   Its topsoil moisture describes a NON-irrigated grid cell (mean ~0.12 m3/m3).
2) A SIMULATED irrigated field driven by that real weather through the FAO-56 soil water balance
   (Allen et al., 1998 — FAO Irrigation & Drainage Paper 56). This gives realistic irrigation cycles
   for the irrigation-decision models. It is clearly labelled "simulated" everywhere.
"""
from pathlib import Path

import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parent
DATA = ROOT.parent.parent / "datasets"          # public/agritech/datasets
OUT = ROOT / "outputs"
OUT.mkdir(exist_ok=True)

# ---- Real reanalysis (rain-fed cell): data-driven threshold for "dry spells" in session 13
RAINFED_DRY = 0.110          # ~20th percentile of the 0-7 cm VWC at the farm grid cell

# ---- Simulated irrigated clay-loam field (Nile Delta) — replace with the farm soil-lab values
THETA_FC = 0.36              # field capacity (m3/m3)
THETA_WP = 0.20              # permanent wilting point (m3/m3)
ROOT_DEPTH_M = 0.40          # effective root zone (m)
P_DEPLETION = 0.50           # FAO-56 "p": fraction of TAW the crop can use without stress
TAW_MM = round(1000 * (THETA_FC - THETA_WP) * ROOT_DEPTH_M, 2)   # total available water = 64 mm
RAW_MM = round(P_DEPLETION * TAW_MM, 2)                   # readily available water = 32 mm
REFILL_POINT = round(THETA_FC - RAW_MM / (1000 * ROOT_DEPTH_M), 4)  # 0.28 m3/m3 -> irrigate below this
WET_POINT = THETA_FC - 0.01                               # >= 0.35 -> "wet"
KC = 1.05                    # single crop coefficient (e.g. a vegetable crop at mid-season)


def load_hourly() -> pd.DataFrame:
    """Real hourly reanalysis data (Open-Meteo) for the KFS farm location, 2023-2024."""
    df = pd.read_csv(DATA / "kfs_farm_hourly_2023_2024.csv", parse_dates=["time"])
    return df.set_index("time").sort_index()


def load_daily() -> pd.DataFrame:
    return pd.read_csv(DATA / "kfs_farm_daily_2023_2024.csv", parse_dates=["date"]).set_index("date")


def simulate_irrigated_field(df: pd.DataFrame, irrigate_hour: int = 6, seed: int = 0) -> pd.DataFrame:
    """FAO-56 root-zone depletion bucket model, hourly.

    Dr(t) = Dr(t-1) + ETc(t) - P(t) - I(t)   (bounded 0..TAW)
    Irrigation (refill to field capacity) is applied at `irrigate_hour` when Dr >= RAW.
    Returns the original frame + depletion_mm, irrigation_mm and a noisy "sensor" VWC.
    """
    rng = np.random.default_rng(seed)
    etc = (df["et0_mm"] * KC).to_numpy()
    rain = df["precipitation_mm"].to_numpy()
    hours = df.index.hour.to_numpy()
    dr = np.zeros(len(df)); irr = np.zeros(len(df))
    d = RAW_MM * 0.3
    for i in range(len(df)):
        d = min(max(d + etc[i] - rain[i], 0.0), TAW_MM)
        if hours[i] == irrigate_hour and d >= RAW_MM:
            irr[i] = d * 1.1              # 10 % leaching fraction
            d = 0.0
        dr[i] = d
    out = df.copy()
    out["depletion_mm"] = dr.round(2)
    out["irrigation_mm"] = irr.round(2)
    vwc = THETA_FC - dr / (1000 * ROOT_DEPTH_M)
    out["field_vwc"] = (vwc + rng.normal(0, 0.003, len(df))).round(4)   # capacitive sensor noise
    return out


def soil_state(vwc: float) -> str:
    """Agronomic class of a volumetric water content reading (irrigated field)."""
    if vwc < REFILL_POINT:
        return "dry"
    if vwc >= WET_POINT:
        return "wet"
    return "moist"


def build_irrigation_dataset(field: pd.DataFrame, horizon_h: int = 24):
    """Features a farm node + weather station know NOW, and the label
    'an irrigation event happens within the next `horizon_h` hours' (sessions 16-17)."""
    X = pd.DataFrame(index=field.index)
    X["soil_vwc"] = field["field_vwc"]
    X["vwc_drop_24h"] = -field["field_vwc"].diff(24)
    X["air_temp"] = field["air_temp_c"]
    X["air_humidity"] = field["air_humidity_pct"]
    X["soil_temp"] = field["soil_temp_c"]
    X["et0_24h"] = field["et0_mm"].rolling(24).sum()
    X["rain_24h"] = field["precipitation_mm"].rolling(24).sum()
    X["radiation"] = field["solar_radiation_wm2"]
    X["wind"] = field["wind_speed_kmh"]
    X["hour"] = field.index.hour
    irr = (field["irrigation_mm"] > 0).astype(int)
    y = irr[::-1].rolling(horizon_h, min_periods=1).max()[::-1].shift(-1)
    data = X.assign(y=y).dropna()
    return data.drop(columns="y"), data["y"].astype(int)