hierarchical_search.py

"""
hierarchical_search.py
=======================
Country-scale extension for the weather/crop-risk system. Does not modify
weather_forecast_env.py, zone_observation.py, crop_risk_scorer.py, or
gru_weather_policy.py.

Scaling problem: weather_forecast_env.py's action space is
Discrete(max_zones + 1), one zone per step -- unworkable at country scale
(20k+ zones), and gru_weather_policy.py's GRU input size is linear in
n_zones, baked into the weight matrices -- no single model works across
region sizes.

Fix: recursively bisect the map (split whichever axis -- lat or lon -- is
longer) and run the existing small env/model with exactly 2 "zones" (a
node's two children) at every level. The model never sees more than 2
zones regardless of country size. Binary branching also means this reuses
the exact model already trained/swept/fixed earlier in this project with
zero new training; a wider tree (e.g. province -> counties) would need a
fresh curriculum phase first. An administrative-hierarchy version (real
boundaries) is a natural upgrade once this scaffolding is validated, but
needs boundary data this repo doesn't have.

Tree depth maps to ForecastConfig.force_data_source (coarse/cheap ->
fine/expensive), so a clean branch prunes after one cheap fetch and
expensive fetches only happen for flagged nodes. Total fetches are capped
by `budget`, not by leaf-zone count.

Phase 1 (implemented, tested here): coarse-gate with
crop_risk_scorer.compute_risk_score() directly -- no RL/GPU in the loop.
Phase 2 (stub, NOT implemented -- see _decide_with_policy): swap in the
trained MaskablePPO model to make the recursion decision adaptively. This
is a real sim-to-real domain-transfer step, not validated here.
"""

from __future__ import annotations

import heapq
import itertools
import logging
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from typing import List, Optional, Tuple

import zone_observation as _zo

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

from zone_observation import DataSource, ForecastConfig, GeoPolygon, AlertLevel
from era5_data_pipeline import register_zone, fetch_episode_context
from crop_risk_scorer import compute_risk_score

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Depth -> data resolution
# ---------------------------------------------------------------------------

RESOLUTION_BY_DEPTH: List[DataSource] = [
    DataSource.ERA5_REANALYSIS,   # depth 0-1: country / province scale
    DataSource.ERA5_REANALYSIS,
    DataSource.SATELLITE_PRECIP,  # depth 2-3: county scale
    DataSource.SATELLITE_SOIL,
    DataSource.OPENMETEO_LIVE,    # depth 4+: finest available
]


def _resolution_for_depth(depth: int) -> DataSource:
    idx = min(depth, len(RESOLUTION_BY_DEPTH) - 1)
    return RESOLUTION_BY_DEPTH[idx]


# ---------------------------------------------------------------------------
# Spatial bisection tree
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class BBox:
    lat_min: float
    lat_max: float
    lon_min: float
    lon_max: float

    @property
    def centroid(self) -> Tuple[float, float]:
        return ((self.lat_min + self.lat_max) / 2.0, (self.lon_min + self.lon_max) / 2.0)

    def split(self) -> Tuple["BBox", "BBox"]:
        lat_span = self.lat_max - self.lat_min
        lon_span = self.lon_max - self.lon_min
        if lat_span >= lon_span:
            mid = (self.lat_min + self.lat_max) / 2.0
            return (
                BBox(self.lat_min, mid, self.lon_min, self.lon_max),
                BBox(mid, self.lat_max, self.lon_min, self.lon_max),
            )
        else:
            mid = (self.lon_min + self.lon_max) / 2.0
            return (
                BBox(self.lat_min, self.lat_max, self.lon_min, mid),
                BBox(self.lat_min, self.lat_max, mid, self.lon_max),
            )

    def to_polygon(self, zone_id: str) -> GeoPolygon:
        return GeoPolygon(
            zone_id=zone_id,
            vertices=[
                (self.lat_min, self.lon_min),
                (self.lat_min, self.lon_max),
                (self.lat_max, self.lon_max),
                (self.lat_max, self.lon_min),
            ],
        )


@dataclass(frozen=True)
class SpatialNode:
    node_id: str
    bbox: BBox
    depth: int

    def children(self) -> Tuple["SpatialNode", "SpatialNode"]:
        left, right = self.bbox.split()
        return (
            SpatialNode(f"{self.node_id}_0", left, self.depth + 1),
            SpatialNode(f"{self.node_id}_1", right, self.depth + 1),
        )


@dataclass
class ResolvedLeaf:
    node: SpatialNode
    max_risk: float
    alert_level: AlertLevel
    depth: int
    data_source: DataSource


# ---------------------------------------------------------------------------
# Node scoring (cheap, deterministic -- no RL, no GPU)
# ---------------------------------------------------------------------------
def _score_node(
    node: SpatialNode,
    date_range: Tuple[datetime, datetime],
    config_template: ForecastConfig,
) -> Tuple[float, AlertLevel, DataSource]:
    polygon = node.bbox.to_polygon(node.node_id)
    register_zone(polygon)

    cfg_dict = config_template.to_dict()
    if config_template.force_data_source is None:
        source = _resolution_for_depth(node.depth)
        cfg_dict["force_data_source"] = source.value
    else:
        source = config_template.force_data_source
    cfg = ForecastConfig.from_dict(cfg_dict)

    ctx = fetch_episode_context(node.node_id, date_range, cfg)
    risk = compute_risk_score(ctx.obs, ctx.forecast, cfg)
    return risk.supply_shortfall_prob, risk.alert_level, source


# ---------------------------------------------------------------------------
# Phase 1: best-first bisection search, coarse-gated by compute_risk_score
# ---------------------------------------------------------------------------
def hierarchical_search(
    root_bbox: BBox,
    date_range: Tuple[datetime, datetime],
    max_depth: int = 10,
    budget: int = 500,
    recurse_floor: AlertLevel = AlertLevel.WARNING,
    config_template: Optional[ForecastConfig] = None,
) -> List[ResolvedLeaf]:
    cfg_template = config_template or ForecastConfig()
    root = SpatialNode("root", root_bbox, depth=0)

    counter = itertools.count()  # heapq tie-breaker; heapq is min-heap, so risk is negated for best-first
    frontier: List[Tuple[float, int, SpatialNode]] = [(0.0, next(counter), root)]
    results: List[ResolvedLeaf] = []
    spent = 0
    failed = 0

    while frontier and spent < budget:
        _, _, node = heapq.heappop(frontier)
        spent += 1

        try:
            max_risk, alert, source = _score_node(node, date_range, cfg_template)
        except Exception as e:
            logger.warning("hierarchical_search: node %s failed to score: %s", node.node_id, e)
            failed += 1
            continue

        should_stop = (
            node.depth >= max_depth
            or alert.severity() < recurse_floor.severity()
        )

        if should_stop:
            results.append(ResolvedLeaf(node, max_risk, alert, node.depth, source))
            continue

        for child in node.children():
            heapq.heappush(frontier, (-max_risk, next(counter), child))

    if failed > 0 and failed == spent:
        raise RuntimeError(
            f"hierarchical_search: all {spent} scored nodes failed -- "
            f"check config_template/credentials before trusting any result "
            f"(including an empty one) from this run."
        )
    if failed > 0:
        logger.warning(
            "hierarchical_search: %d/%d node evaluations failed to score "
            "(see warnings above) -- results below exclude those nodes, "
            "they are neither 'resolved clean' nor 'pending'.", failed, spent,
        )
    if frontier:
        logger.info(
            "hierarchical_search: budget exhausted with %d nodes still pending "
            "(unexplored -- not the same as 'clean')", len(frontier),
        )

    return results


def _decide_with_policy(node: SpatialNode, children_contexts, model, env_factory):
    """Phase 2 STUB -- not implemented, not tested.

    Intended contract: node's two children become n_zones=2 "zones" for one
    WeatherForecastEnv episode. Seed belief_map from the already-fetched
    coarse compute_risk_score result per child (a real sim-to-real
    assumption, not validated here). "Inspect zone i" = fetch child i at
    the next-finer resolution and rescore; "terminate" = stop, uninspected
    children keep their coarse read. Recursion only happens for children
    the policy chose to inspect -- this is where RL adds value over Phase
    1's fixed recurse_floor: adaptive, learned budget allocation.

    Raises rather than returning a fake decision.
    """
    raise NotImplementedError(
        "Phase 2 (RL-driven recursion) is a design stub -- see docstring. "
        "Phase 1 (hierarchical_search, using compute_risk_score directly) "
        "is implemented and tested in this file."
    )


# ---------------------------------------------------------------------------
# Self-test: runs fully offline (SYNTHETIC data source, no credentials needed)
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO, format="%(asctime)s | %(levelname)s | %(message)s")

    indonesia_bbox = BBox(lat_min=-11.0, lat_max=6.0, lon_min=95.0, lon_max=141.0)

    date_range = (
        datetime(2026, 7, 1, tzinfo=timezone.utc),
        datetime(2026, 7, 31, tzinfo=timezone.utc),
    )

    cfg = ForecastConfig(force_data_source=DataSource.SYNTHETIC)

    results = hierarchical_search(
        indonesia_bbox, date_range,
        max_depth=8, budget=60,
        recurse_floor=AlertLevel.WATCH,
        config_template=cfg,
    )

    by_depth = {}
    for r in results:
        by_depth.setdefault(r.depth, []).append(r)

    print(f"\n{len(results)} leaves resolved out of a budget of 60 node evaluations")
    print(f"(worst case with no pruning at max_depth=8 would be up to {2**8} leaves)\n")
    for depth in sorted(by_depth):
        leaves = by_depth[depth]
        alerts = [l.alert_level.value for l in leaves]
        print(f"  depth {depth}: {len(leaves)} branches stopped here -- alerts: {alerts}")