Blame for scripts/collect_2026_training_data.py

"""Collect 2026-schema training data for the next-gen photosynthesis ML model.

Pulls 15-min telemetry from all Crop_2Soil devices (treatment + reference) and
the IMS station, derives Farquhar inputs from the 2026 fleet's reduced sensor
set, computes a Farquhar net-assimilation label (A_n) per row, and appends to
``Data/2026/sensor_history.parquet``.

Designed to be re-run periodically (weekly) to accumulate growing-season data.
By August/September 2026 the parquet should hold ≥3 months of growing-season
hours — enough to retrain Stage 2 ML on the 2026 schema (see
``context/3_todo.md`` § 12.6 ``bigml-train``).

Why we derive inputs vs measure them
------------------------------------
The 2026 fleet retired the on-site Air1 sensor that previously supplied PAR,
CO2, and VPD. We approximate:

  PAR (µmol/m²/s) ≈ 2.0 × GHI (W/m²)
      — daylight broadband-to-PAR conversion; ±10 % vs measured PAR
        (Akitsu et al. 2017). Worst at low sun angles.

  CO2 (ppm)       = 420 (assumed atmospheric)
      — Sde Boker is far from urban/forest. Hourly CO2 fluctuates by
        ~5 ppm; small effect on A via Ci.

  VPD (kPa)       = SVP(Tair) × (1 − RH/100)   (Tetens equation)
      — SVP(T) = 0.611 × exp(17.27·T / (T+237.3)).
        Within ~5 % of measured VPD over the Negev range.

  CWSI            = max(0, min(1, (Tleaf − Tair) / 15))
      — Crude but defensible until we add an empirical CWSI baseline.

Usage
-----
  python -m scripts.collect_2026_training_data
  python -m scripts.collect_2026_training_data --from 2026-05-01 --to 2026-05-18
"""

from __future__ import annotations

import argparse
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Dict, List

import numpy as np
import pandas as pd

_PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(_PROJECT_ROOT))

try:
    from dotenv import load_dotenv
    load_dotenv(_PROJECT_ROOT / ".env")
except ImportError:
    pass

from src.data.thingsboard_client import (
    ThingsBoardClient,
    DEVICE_REGISTRY,
    TREATMENT_DEVICES,
    REFERENCE_DEVICES,
    VineArea,
)
from src.models.farquhar_model import FarquharModel
from src.shading.solar_geometry import ShadowModel

_IMS_CSV = _PROJECT_ROOT / "Data" / "ims" / "ims_merged_15min.csv"
_OUT_DIR = _PROJECT_ROOT / "Data" / "2026"
_OUT_PARQUET = _OUT_DIR / "sensor_history.parquet"

# Telemetry keys we pull per Crop_2Soil device.
_DEVICE_KEYS: List[str] = [
    "leafTemperature", "ambientTemperatureIRT", "NDVI",
    "PRI", "PSRI", "SIPI", "GCI", "LCI", "DUVI",
    "soilTemperature", "soilMoisture", "soilBulkEC", "soilPoreWaterEC",
    "soilTemperature2", "soilMoisture2", "soilPoreWaterEC2",
]

_TRACKER_DEVICES = ["Tracker501", "Tracker502", "Tracker503", "Tracker509"]

_CO2_ASSUMED_PPM = 420.0
_PAR_FROM_GHI = 2.0   # µmol/m²/s per W/m²


# ---------------------------------------------------------------------------
# Derivations
# ---------------------------------------------------------------------------

def _saturation_vapour_pressure_kpa(t_c: pd.Series) -> pd.Series:
    """Tetens equation: SVP(T) in kPa given T in °C."""
    return 0.611 * np.exp(17.27 * t_c / (t_c + 237.3))


def derive_inputs(df: pd.DataFrame) -> pd.DataFrame:
    """Add PAR, VPD, CO2, CWSI columns derived from the 2026 sensor set.

    `par_umol_derived` is the **above-canopy** PAR estimate from IMS GHI;
    the treatment-area shading correction is applied separately in
    `apply_shading_correction()` so reference rows keep the open-sky PAR.
    """
    df = df.copy()
    df["par_umol_derived"] = (df["ghi_w_m2"].clip(lower=0) * _PAR_FROM_GHI).clip(lower=0, upper=3000)
    svp = _saturation_vapour_pressure_kpa(df["air_temperature_c"])
    df["vpd_kpa_derived"] = (svp * (1.0 - df["rh_percent"].clip(0, 100) / 100.0)).clip(lower=0, upper=10)
    df["co2_ppm_assumed"] = _CO2_ASSUMED_PPM
    # CWSI proxy from leaf-air ΔT (positive when leaf > air = stress)
    df["cwsi_proxy"] = ((df["leaf_temperature"] - df["air_temperature_c"]).clip(lower=0) / 15.0).clip(0, 1)
    return df


def fetch_tracker_angles_15min(client: ThingsBoardClient,
                               start: datetime, end: datetime) -> pd.Series:
    """Pull 15-min mean tracker angle across the 4 trackers."""
    step = timedelta(days=7)
    per_tracker = {}
    for name in _TRACKER_DEVICES:
        frames = []
        cursor = start
        while cursor < end:
            cursor_end = min(cursor + step, end)
            try:
                df = client.get_timeseries(
                    name, ["angle"],
                    start=cursor, end=cursor_end,
                    limit=10_000, interval_ms=900_000, agg="AVG",
                )
            except Exception:
                df = pd.DataFrame()
            if not df.empty:
                frames.append(df)
            cursor = cursor_end
        if frames:
            tdf = pd.concat(frames).sort_index()
            tdf.index = pd.to_datetime(tdf.index, utc=True).floor("15min")
            tdf = tdf[~tdf.index.duplicated(keep="last")]
            per_tracker[name] = tdf["angle"].rename(name)
    if not per_tracker:
        return pd.Series(dtype=float, name="tracker_angle_mean")
    wide = pd.concat(per_tracker.values(), axis=1)
    return wide.mean(axis=1, skipna=True).rename("tracker_angle_mean")


def apply_shading_correction(df: pd.DataFrame,
                             tracker_angles: pd.Series) -> pd.DataFrame:
    """Compute treatment-area fruiting-zone PAR with panel shading.

    For each treatment-area timestep, runs `ShadowModel.project_shadow()`
    with the actual tracker angle and replaces `par_umol_derived` with
    the fruiting-zone mean PAR (averaged across the row's horizontal
    positions in the fruiting vertical zone). Reference rows are
    untouched (open sky = above-canopy PAR).

    Adds:
      par_factor_treatment    : ratio of corrected to open-sky PAR
      tracker_angle_mean      : mean across the 4 trackers (NaN if missing)
    """
    df = df.copy()
    df["par_factor_treatment"] = 1.0

    # Align tracker telemetry to every row's timestamp.
    df["tracker_angle_mean"] = tracker_angles.reindex(
        df.index, method="nearest", tolerance=pd.Timedelta("15min"),
    ).values

    treat_mask = (df["area"] == "treatment").to_numpy()
    if not treat_mask.any():
        return df

    sm = ShadowModel()
    from config.settings import FRUITING_ZONE_INDEX
    fz_idx = FRUITING_ZONE_INDEX

    # Compute solar position once per unique timestamp (treatment + reference
    # share timestamps, so use the unique set to avoid duplicate work).
    unique_ts = pd.DatetimeIndex(df.index.unique())
    sun = sm.get_solar_position(unique_ts)
    sun_lookup = {
        ts: (float(sun.loc[ts, "solar_elevation"]), float(sun.loc[ts, "solar_azimuth"]))
        for ts in unique_ts
    }

    # Vector pass over treatment rows via positional indices to avoid the
    # duplicate-index gotcha with `.at[]`.
    treat_positions = np.where(treat_mask)[0]
    ghi_col = df.columns.get_loc("ghi_w_m2")
    par_col = df.columns.get_loc("par_umol_derived")
    tilt_col = df.columns.get_loc("tracker_angle_mean")
    factor_col = df.columns.get_loc("par_factor_treatment")

    for pos in treat_positions:
        ts = df.index[pos]
        ghi = df.iat[pos, ghi_col]
        if pd.isna(ghi) or ghi <= 0:
            continue
        elev, azim = sun_lookup[ts]
        if elev <= 2.0:
            continue
        tilt = df.iat[pos, tilt_col]
        if pd.isna(tilt):
            tilt = float(sm.compute_tracker_tilt(azim, elev)["tracker_theta"])
        total_par = float(ghi * _PAR_FROM_GHI)
        try:
            mask = sm.project_shadow(elev, azim, float(tilt))
            par_grid = sm.compute_par_distribution(
                total_par, mask,
                solar_elevation=elev, solar_azimuth=azim, tracker_tilt=float(tilt),
            )
            factor = float(par_grid[fz_idx, :].mean()) / total_par
        except Exception:
            factor = 1.0
        df.iat[pos, factor_col] = factor
        df.iat[pos, par_col] = df.iat[pos, par_col] * factor

    return df


def compute_farquhar_a(df: pd.DataFrame) -> pd.Series:
    """Run the Semillon Farquhar model row-by-row over the derived inputs."""
    fm = FarquharModel()
    out = np.full(len(df), np.nan)
    for i, row in enumerate(df.itertuples(index=False)):
        # Skip rows missing any input
        par = getattr(row, "par_umol_derived", None)
        tleaf = getattr(row, "leaf_temperature", None)
        tair = getattr(row, "air_temperature_c", None)
        vpd = getattr(row, "vpd_kpa_derived", None)
        cwsi = getattr(row, "cwsi_proxy", 0.0)
        if any(v is None or (isinstance(v, float) and np.isnan(v))
               for v in (par, tleaf, tair, vpd)):
            continue
        try:
            out[i] = fm.calc_photosynthesis(par, tleaf, _CO2_ASSUMED_PPM, vpd, tair, cwsi)
        except Exception:
            continue
    return pd.Series(out, index=df.index, name="a_farquhar_umol")


# ---------------------------------------------------------------------------
# TB fetch helpers
# ---------------------------------------------------------------------------

def _fetch_device_history(client: ThingsBoardClient,
                          name: str,
                          start: datetime,
                          end: datetime) -> pd.DataFrame:
    """Pull 15-min telemetry for one Crop_2Soil device over [start, end]."""
    step = timedelta(days=7)
    frames = []
    cursor = start
    while cursor < end:
        cursor_end = min(cursor + step, end)
        try:
            df = client.get_timeseries(
                name, _DEVICE_KEYS,
                start=cursor, end=cursor_end,
                limit=10_000, interval_ms=900_000, agg="AVG",
            )
        except Exception:
            df = pd.DataFrame()
        if not df.empty:
            frames.append(df)
        cursor = cursor_end
    if not frames:
        return pd.DataFrame()
    out = pd.concat(frames).sort_index()
    out.index = pd.to_datetime(out.index, utc=True).floor("15min")
    out = out[~out.index.duplicated(keep="last")]
    return out


def fetch_area_history(client: ThingsBoardClient,
                       area: VineArea,
                       start: datetime,
                       end: datetime) -> pd.DataFrame:
    """Aggregate (mean across devices) per area for the requested window."""
    names = TREATMENT_DEVICES if area == VineArea.TREATMENT else REFERENCE_DEVICES
    per_device: Dict[str, pd.DataFrame] = {}
    for name in names:
        df = _fetch_device_history(client, name, start, end)
        if not df.empty:
            per_device[name] = df

    if not per_device:
        return pd.DataFrame()

    # Stack by device, average columns per timestamp
    stacked = pd.concat(per_device.values(), keys=per_device.keys(),
                        names=["device", "timestamp_utc"])
    averaged = stacked.groupby(level="timestamp_utc").mean(numeric_only=True)

    # Rename camelCase TB keys → snake_case for the parquet schema
    averaged = averaged.rename(columns={
        "leafTemperature":          "leaf_temperature",
        "ambientTemperatureIRT":    "ambient_temp_irt",
        "NDVI": "ndvi", "PRI": "pri", "PSRI": "psri",
        "SIPI": "sipi", "GCI": "gci", "LCI": "lci", "DUVI": "duvi",
        "soilTemperature":          "soil_temp_shallow_c",
        "soilTemperature2":         "soil_temp_deep_c",
        "soilMoisture":             "soil_moisture_shallow_pct",
        "soilMoisture2":            "soil_moisture_deep_pct",
        "soilBulkEC":               "soil_bulk_ec",
        "soilPoreWaterEC":          "soil_pore_water_ec_shallow",
        "soilPoreWaterEC2":         "soil_pore_water_ec_deep",
    })
    averaged["area"] = area.value
    return averaged


def load_ims_15min(start: datetime, end: datetime) -> pd.DataFrame:
    """Load IMS CSV, filter to window, leave at 15-min resolution."""
    if not _IMS_CSV.exists():
        raise FileNotFoundError(f"{_IMS_CSV} not found. Run scripts.download_ims_data first.")
    df = pd.read_csv(_IMS_CSV)
    df["timestamp_utc"] = pd.to_datetime(df["timestamp_utc"], utc=True)
    df = df.set_index("timestamp_utc").sort_index()
    df = df.loc[start:end]
    return df


# ---------------------------------------------------------------------------
# Top-level pipeline
# ---------------------------------------------------------------------------

def build_dataset(start: datetime, end: datetime) -> pd.DataFrame:
    client = ThingsBoardClient()

    print(f"Fetching treatment area ({len(TREATMENT_DEVICES)} devices) ...")
    treat = fetch_area_history(client, VineArea.TREATMENT, start, end)
    print(f"  {len(treat):,} 15-min rows")

    print(f"Fetching reference area ({len(REFERENCE_DEVICES)} devices) ...")
    ref = fetch_area_history(client, VineArea.REFERENCE, start, end)
    print(f"  {len(ref):,} 15-min rows")

    print("Loading IMS weather ...")
    ims = load_ims_15min(start, end)
    print(f"  {len(ims):,} 15-min rows")

    print("Fetching tracker angles for shading correction ...")
    tracker_series = fetch_tracker_angles_15min(client, start, end)
    print(f"  {len(tracker_series):,} 15-min tracker rows")

    frames = []
    for label, area_df in [("treatment", treat), ("reference", ref)]:
        if area_df.empty:
            continue
        # Inner-join with IMS so every row has weather context
        joined = area_df.join(ims, how="inner")
        joined = derive_inputs(joined)
        frames.append(joined)

    if not frames:
        raise RuntimeError("No data assembled — TB or IMS returned nothing.")

    combined = pd.concat(frames)
    combined.index.name = "timestamp_utc"

    # Apply panel shading correction to treatment rows only.
    combined = apply_shading_correction(combined, tracker_series)

    # Recompute Farquhar label after PAR correction.
    combined["a_farquhar_umol"] = compute_farquhar_a(combined)
    return combined


def main() -> None:
    p = argparse.ArgumentParser(description="Collect 2026-schema training data.")
    p.add_argument("--from", dest="from_date", default=None,
                   help="Start date (UTC, inclusive). Default: 30 days ago.")
    p.add_argument("--to", dest="to_date", default=None,
                   help="End date (UTC, inclusive). Default: today.")
    p.add_argument("--mode", choices=["append", "replace"], default="append",
                   help="append (default) merges into existing parquet on timestamp+area; "
                        "replace overwrites the file.")
    args = p.parse_args()

    end = (datetime.fromisoformat(args.to_date).replace(tzinfo=timezone.utc)
           if args.to_date else datetime.now(tz=timezone.utc))
    start = (datetime.fromisoformat(args.from_date).replace(tzinfo=timezone.utc)
             if args.from_date else end - timedelta(days=30))

    new = build_dataset(start, end)
    print(f"\nCollected {len(new):,} new rows  range: {new.index.min()} → {new.index.max()}")

    _OUT_DIR.mkdir(parents=True, exist_ok=True)
    if args.mode == "append" and _OUT_PARQUET.exists():
        existing = pd.read_parquet(_OUT_PARQUET)
        combined = pd.concat([existing.reset_index(), new.reset_index()])
        combined = combined.drop_duplicates(subset=["timestamp_utc", "area"], keep="last")
        combined = combined.set_index("timestamp_utc").sort_index()
        print(f"Merged with existing {len(existing):,} rows → {len(combined):,} total")
    else:
        combined = new

    combined.to_parquet(_OUT_PARQUET)
    print(f"Wrote → {_OUT_PARQUET}")

    # Validation summary
    a = combined["a_farquhar_umol"].dropna()
    if not a.empty:
        sun = combined[combined["ghi_w_m2"] > 100]["a_farquhar_umol"].dropna()
        print(f"\nFarquhar A summary:")
        print(f"  total non-null rows : {len(a):,}")
        print(f"  range               : {a.min():.2f}  →  {a.max():.2f} µmol CO2/m²/s")
        print(f"  daytime (GHI>100) mean : {sun.mean():.2f}  median : {sun.median():.2f}")


if __name__ == "__main__":
    main()