src/proxy_a_score.py

"""proxy_a_score — composite photosynthesis-health metric, 0..1.

LiCor (direct A measurement) is not available this season, so the control
loop and chatbot need a single scalar that says "how healthy is the canopy
right now" from the sensors we already collect. This module produces that
scalar by combining 5 already-streamed metrics on Crop_2Soil devices:

  - PRI (Photochemical Reflectance Index — light-use efficiency)
  - PSRI (Plant Senescence Reflectance Index — inverted)
  - leaf-air ΔT (transpiration cooling — inverted)
  - shallow soil moisture (water availability)
  - NDVI (greenness baseline)

Each metric is mapped linearly from its `stress`→`healthy` bound into
[0, 1] (clipped), then combined as a weighted average. Weights and bounds
live in `config/settings.py` (`PROXY_A_WEIGHTS`, `PROXY_A_BOUNDS`) so they
can be tuned after the null-intervention baseline period without code change.

Higher score = healthier / more productive canopy. The control loop should
prefer interventions that *don't* drop the score by more than ε.
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Optional

from config import settings


@dataclass
class ProxyAScore:
    """Composite photosynthesis-health metric, plus per-term breakdown.

    Attributes
    ----------
    score : float
        Weighted average of the per-metric sub-scores, in [0, 1].
        NaN if no metrics were available.
    contributions : dict[str, float]
        Per-metric sub-score in [0, 1] (only metrics that had data).
    weights_used : dict[str, float]
        Per-metric weight actually applied (renormalised over available
        metrics so missing inputs don't bias the score downward).
    """

    score: float
    contributions: Dict[str, float]
    weights_used: Dict[str, float]


def _linear_score(value: float, stress: float, healthy: float) -> float:
    """Map a raw measurement into [0, 1] given its stress/healthy bounds.

    Works whether `healthy > stress` (PRI, NDVI, soil — higher is better)
    or `healthy < stress` (PSRI, ΔT — lower is better).
    """
    if stress == healthy:
        return 0.5  # degenerate config, neutral
    raw = (value - stress) / (healthy - stress)
    return max(0.0, min(1.0, raw))


def compute_proxy_a(
    pri: Optional[float] = None,
    psri: Optional[float] = None,
    delta_t: Optional[float] = None,
    soil_moisture_pct: Optional[float] = None,
    ndvi: Optional[float] = None,
) -> ProxyAScore:
    """Compute the proxy-A score from any subset of the 5 input metrics.

    Missing metrics are dropped and the weights for the remaining metrics
    are renormalised so they sum to 1. If all metrics are missing, the
    score is NaN.
    """
    raw_inputs: Dict[str, Optional[float]] = {
        "pri":     pri,
        "psri":    psri,
        "delta_t": delta_t,
        "soil":    soil_moisture_pct,
        "ndvi":    ndvi,
    }

    contributions: Dict[str, float] = {}
    weights_used: Dict[str, float] = {}

    for key, val in raw_inputs.items():
        if val is None:
            continue
        b = settings.PROXY_A_BOUNDS.get(key)
        w = settings.PROXY_A_WEIGHTS.get(key, 0.0)
        if b is None or w <= 0:
            continue
        contributions[key] = _linear_score(val, b["stress"], b["healthy"])
        weights_used[key] = w

    if not weights_used:
        return ProxyAScore(score=float("nan"), contributions={}, weights_used={})

    total_w = sum(weights_used.values())
    renormalised = {k: w / total_w for k, w in weights_used.items()}
    score = sum(contributions[k] * renormalised[k] for k in contributions)
    return ProxyAScore(
        score=round(score, 3),
        contributions={k: round(v, 3) for k, v in contributions.items()},
        weights_used=renormalised,
    )


def compute_from_snapshot(snapshot: dict) -> ProxyAScore:
    """Convenience: pull the 5 inputs from a VineSnapshot.to_dict()."""
    return compute_proxy_a(
        pri=snapshot.get("treatment_crop_pri"),
        psri=snapshot.get("treatment_crop_psri"),
        delta_t=snapshot.get("treatment_air_leaf_delta_t"),
        soil_moisture_pct=snapshot.get("treatment_soil_moisture_pct"),
        ndvi=snapshot.get("treatment_crop_ndvi"),
    )


def slot_cost(
    *,
    energy_value_ils: float,
    proxy_a: ProxyAScore,
    proxy_a_floor: float = 0.5,
    proxy_a_weight: float = 5.0,
) -> float:
    """Slot-level cost for the tracker control loop.

    Combines instantaneous energy revenue with a penalty for any drop of
    `proxy_a.score` below `proxy_a_floor`. Higher = better.

      cost = energy_value_ils  −  proxy_a_weight × max(0, floor − score)

    Tune `proxy_a_weight` from the calibration period so that the penalty
    at the 95th-percentile stress reading equals roughly half a mid-day
    hour of energy revenue.
    """
    shortfall = max(0.0, proxy_a_floor - proxy_a.score)
    return round(energy_value_ils - proxy_a_weight * shortfall, 3)