crop_risk_scorer.py

"""
crop_risk_scorer.py
===================
Deterministic, economics-calibrated crop risk scoring.
"""

from __future__ import annotations

import logging
import math
from dataclasses import dataclass, fields
from datetime import datetime, timedelta, timezone
from typing import Optional, Tuple

import zone_observation as _zo

assert _zo.SCHEMA_VERSION == 3, (
    f"crop_risk_scorer: zone_observation schema mismatch "
    f"(expected 3, got {_zo.SCHEMA_VERSION})"
)

from zone_observation import (
    AlertLevel,
    CropStage,
    DataSource,
    ForecastConfig,
    ForecastResult,
    RiskScore,
    ZoneObs,
    _clip,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Product alert thresholds (Indonesia scorecards — freeze candidates)
# ---------------------------------------------------------------------------

DEFAULT_DROUGHT_WARNING = 0.35
DEFAULT_DROUGHT_CRITICAL = 0.50
DEFAULT_FLOOD_WARNING = 0.25
DEFAULT_FLOOD_CRITICAL = 0.40


# ---------------------------------------------------------------------------
# RiskWeights
# ---------------------------------------------------------------------------

@dataclass
class RiskWeights:
    drought_obs_weight: float = 0.60
    drought_forecast_weight: float = 0.40

    flood_obs_weight: float = 0.55
    flood_forecast_weight: float = 0.45

    fungi_obs_weight: float = 0.70
    fungi_forecast_weight: float = 0.30

    supply_drought_weight: float = 0.40
    supply_flood_weight: float = 0.35
    supply_harvest_pressure_weight: float = 0.25

    quality_fungi_weight: float = 0.65
    quality_delay_weight: float = 0.35

    def __post_init__(self) -> None:
        for f in fields(self):
            v = getattr(self, f.name)
            if not (0.0 <= v <= 1.0):
                raise ValueError(f"{f.name}={v} outside [0,1]")

    def attach_to_config(self, config: ForecastConfig) -> ForecastConfig:
        config.risk_weights = self  # type: ignore
        return config

    @classmethod
    def from_config(cls, config: ForecastConfig) -> "RiskWeights":
        return getattr(config, "risk_weights", cls())


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _decayed_mean(seq, half_life_days: float = 5.0):
    """Apply exponential decay to forecast signals."""
    if not seq:
        return 0.0
    weights = [math.exp(-i / half_life_days) for i in range(len(seq))]
    denom = sum(weights)
    if denom == 0.0:
        return 0.0
    return sum(w * v for w, v in zip(weights, seq)) / denom


def _forecast_signal(forecast: ForecastResult, key: str, window: int) -> float:
    if not hasattr(forecast, key):
        raise AttributeError(f"ForecastResult missing required field: '{key}'")
    seq = getattr(forecast, key)
    if not seq:
        logger.warning(
            "_forecast_signal: '%s' is empty — scoring with 0.0 (check forecast pipeline)",
            key,
        )
        return 0.0
    return _decayed_mean(seq[:window])


# ---------------------------------------------------------------------------
# Risk components
# ---------------------------------------------------------------------------

def _drought_risk(obs: ZoneObs, forecast: ForecastResult, w: RiskWeights) -> float:
    return _clip(
        w.drought_obs_weight * obs.drought_signal()
        + w.drought_forecast_weight * _forecast_signal(forecast, "prob_drought_day", 14),
        0, 1,
    )


def _flood_risk(obs: ZoneObs, forecast: ForecastResult, w: RiskWeights) -> float:
    return _clip(
        w.flood_obs_weight * obs.flood_signal()
        + w.flood_forecast_weight * _forecast_signal(forecast, "prob_heavy_rain", 7),
        0, 1,
    )


def _fungi_risk(obs: ZoneObs, forecast: ForecastResult, w: RiskWeights) -> float:
    raw_forecast_term = _forecast_signal(forecast, "prob_high_humidity", 10)
    forecast_anomaly_adj = _clip(obs.rh_anomaly_idx / 3.0, -0.3, 0.3)
    forecast_term_adj = _clip(raw_forecast_term + forecast_anomaly_adj, 0, 1)
    return _clip(
        w.fungi_obs_weight * obs.fungi_risk_signal()
        + w.fungi_forecast_weight * forecast_term_adj,
        0, 1,
    )


# ---------------------------------------------------------------------------
# Harvest window
# ---------------------------------------------------------------------------

def _optimal_harvest_window(
    obs: ZoneObs,
    forecast: ForecastResult,
) -> Tuple[Optional[datetime], Optional[datetime]]:
    if obs.crop_stage not in (
        CropStage.GRAIN_FILLING,
        CropStage.MATURATION,
        CropStage.HARVEST,
    ):
        return None, None

    WINDOW = 5
    RAIN_THRESHOLD_MM  = 12.0
    HUMIDITY_THRESHOLD = 78.0

    for i in range(len(forecast.precip_mm) - WINDOW + 1):
        rain_ok = all(p < RAIN_THRESHOLD_MM  for p in forecast.precip_mm[i : i + WINDOW])
        hum_ok  = all(h < HUMIDITY_THRESHOLD for h in forecast.rh_mean_pct[i : i + WINDOW])
        if rain_ok and hum_ok:
            start = obs.valid_time + timedelta(days=i)
            return start, start + timedelta(days=WINDOW)

    # Fallback: use days_to_harvest if available
    if obs.days_to_harvest is not None:
        start = obs.valid_time + timedelta(days=obs.days_to_harvest)
        return start, start + timedelta(days=7)

    return None, None


# ---------------------------------------------------------------------------
# Harvest pressure (smoothed)
# ---------------------------------------------------------------------------

def _harvest_pressure(obs: ZoneObs) -> float:
    if obs.days_to_harvest is not None:
        return float(_clip(1.0 - obs.days_to_harvest / 10.0, 0.3, 1.0))
    # Proxy via crop stage when days_to_harvest is unavailable
    stage_pressure = {
        CropStage.GRAIN_FILLING: 0.8,
        CropStage.MATURATION:    1.0,
        CropStage.HARVEST:       1.0,
    }
    return stage_pressure.get(obs.crop_stage, 0.3)


# ---------------------------------------------------------------------------
# Confidence model
# ---------------------------------------------------------------------------

def _compute_confidence(obs: ZoneObs, forecast: ForecastResult) -> float:
    base = 0.6

    if obs.source.is_observational():
        base += 0.2

    if obs.source in (DataSource.SATELLITE_PRECIP, DataSource.SATELLITE_SOIL):
        base += 0.05  # direct retrieval, above the ERA5 obs bump already applied

    if forecast.source in (DataSource.SATELLITE_PRECIP, DataSource.SATELLITE_SOIL):
        base += 0.20
    elif forecast.source == DataSource.OPENMETEO_LIVE:
        base += 0.1
    elif forecast.source == DataSource.ERA5_REANALYSIS:
        base += 0.15
    else:
        base -= 0.1

    if not obs.has_reliable_ndvi():
        base *= 0.9
    if obs.quality_flag >= 2:
        base *= 0.8

    return float(_clip(base, 0.0, 1.0))


# ---------------------------------------------------------------------------
# Alert level
# ---------------------------------------------------------------------------

def _product_thresholds(cfg: ForecastConfig) -> Tuple[float, float, float, float]:
    return (
        float(getattr(cfg, "drought_warning_threshold", DEFAULT_DROUGHT_WARNING)),
        float(getattr(cfg, "drought_critical_threshold", DEFAULT_DROUGHT_CRITICAL)),
        float(getattr(cfg, "flood_warning_threshold", DEFAULT_FLOOD_WARNING)),
        float(getattr(cfg, "flood_critical_threshold", DEFAULT_FLOOD_CRITICAL)),
    )


def _alert_level(
    max_risk: float,
    cfg: ForecastConfig,
    drought_risk: float = 0.0,
    flood_risk: float = 0.0,
) -> AlertLevel:
    d_warn, d_crit, f_warn, f_crit = _product_thresholds(cfg)

    # --- Product path: hazard-specific WARNING / CRITICAL ---
    if drought_risk >= d_crit or flood_risk >= f_crit:
        return AlertLevel.CRITICAL
    if drought_risk >= d_warn or flood_risk >= f_warn:
        return AlertLevel.WARNING

    # --- Extreme max_risk fallback (legacy cuts; rarely hit after damping) ---
    if max_risk >= 0.85:
        return AlertLevel.CRITICAL
    if max_risk >= 0.65:
        return AlertLevel.WARNING

    # --- Economics ladder for ADVISORY / WATCH / NONE ---
    rational = cfg.rational_termination_threshold
    watch_threshold = rational / 2.0

    if rational >= 0.65:
        logger.warning(
            "_alert_level: rational_termination_threshold=%.4f has reached "
            "the legacy WARNING tier's fixed cutoff (0.65) -- ADVISORY's "
            "range has collapsed to near-zero width for this config. Check "
            "alert_value/false_alert_penalty/miss_penalty.", rational,
        )

    if max_risk >= rational:
        return AlertLevel.ADVISORY
    if max_risk >= watch_threshold:
        return AlertLevel.WATCH
    return AlertLevel.NONE


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

def compute_risk_score(
    obs:      ZoneObs,
    forecast: ForecastResult,
    config:   Optional[ForecastConfig] = None,
) -> RiskScore:
    cfg = config or ForecastConfig()
    w   = RiskWeights.from_config(cfg)

    drought = _drought_risk(obs, forecast, w)
    flood   = _flood_risk(obs, forecast, w)
    fungi   = _fungi_risk(obs, forecast, w)

    pressure = _harvest_pressure(obs)

    supply = _clip(
        w.supply_drought_weight * drought
        + w.supply_flood_weight * flood
        + w.supply_harvest_pressure_weight * pressure,
        0, 1,
    )

    delay_days   = forecast.max_consecutive_rain_days(15.0)
    delay_factor = _clip(delay_days / 10.0, 0, 1)

    quality = _clip(
        w.quality_fungi_weight * fungi
        + w.quality_delay_weight * delay_factor,
        0, 1,
    )

    confidence = _compute_confidence(obs, forecast)

    confidence_scale = 0.5 + 0.5 * confidence
    drought_adj = _clip(drought * confidence_scale, 0, 1)
    flood_adj   = _clip(flood   * confidence_scale, 0, 1)
    fungi_adj   = _clip(fungi   * confidence_scale, 0, 1)
    supply_adj  = _clip(supply  * confidence_scale, 0, 1)
    quality_adj = _clip(quality * confidence_scale, 0, 1)

    max_risk = max(drought_adj, flood_adj, fungi_adj, supply_adj, quality_adj)
    alert    = _alert_level(
        max_risk, cfg,
        drought_risk=drought_adj,
        flood_risk=flood_adj,
    )

    start, end = _optimal_harvest_window(obs, forecast)

    d_warn, d_crit, f_warn, f_crit = _product_thresholds(cfg)
    trigger = "none"
    if drought_adj >= d_crit:
        trigger = "drought_critical"
    elif flood_adj >= f_crit:
        trigger = "flood_critical"
    elif drought_adj >= d_warn:
        trigger = "drought_warning"
    elif flood_adj >= f_warn:
        trigger = "flood_warning"
    elif max_risk >= 0.65:
        trigger = "max_risk"
    elif max_risk >= cfg.rational_termination_threshold:
        trigger = "advisory_econ"

    action_notes = (
        f"drought={drought_adj:.2f} flood={flood_adj:.2f} fungi={fungi_adj:.2f} "
        f"supply={supply_adj:.2f} quality={quality_adj:.2f} "
        f"confidence={confidence:.2f} alert={alert.value} trigger={trigger}"
    )

    logger.debug(
        "RiskScore %s: %s", obs.zone_id, action_notes
    )

    return RiskScore(
        zone_id=obs.zone_id,
        scored_at=datetime.now(timezone.utc),
        supply_shortfall_prob=supply_adj,
        drought_risk=drought_adj,
        flood_risk=flood_adj,
        supply_risk_composite=supply_adj,
        fungi_contamination_prob=fungi_adj,
        harvest_delay_days=float(delay_days),
        quality_risk_composite=quality_adj,
        optimal_harvest_window_start=start,
        optimal_harvest_window_end=end,
        alert_level=alert,
        action_notes=action_notes,
        confidence=confidence,
    )