impact_labels.py

"""
impact_labels.py
================
L1 impact / disaster labels for evaluation (point 1).

Single source of truth for offline impact events used by backtests and
research metrics. Does NOT invent live BNPB API access — only a file-backed
store with an explicit JSON schema.

Build order satisfied:
  entities → immutable RECORDED state → load/query contracts → tests in __main__
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import asdict, dataclass, field
from datetime import date, datetime, timezone
from enum import Enum
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple

logger = logging.getLogger(__name__)

IMPACT_SCHEMA_VERSION = 1

# ---------------------------------------------------------------------------
# Entities
# ---------------------------------------------------------------------------

class ImpactHazard(str, Enum):
    """Canonical hazard tags for L1 labels (disaster / insurance-shaped)."""

    DROUGHT = "drought"
    FLOOD = "flood"
    OTHER = "other"


class ImpactSource(str, Enum):
    """Provenance — not a trust ranking."""

    BNPB = "bnpb"
    PROVINCIAL = "provincial"
    AUTP_CLAIM = "autp_claim"
    MANUAL = "manual"
    SYNTHETIC_TEST = "synthetic_test"


@dataclass(frozen=True)
class ImpactEvent:
    """
    Immutable recorded impact fact for one zone-day (or multi-day span).

    State: RECORDED only. No transitions. Invalid rows are rejected at load.
    """

    event_id: str
    zone_id: str
    hazard: ImpactHazard
    start_date: date
    end_date: date
    source: ImpactSource
    confidence: float = 1.0  # [0, 1] label confidence, not model confidence
    notes: str = ""
    extras: Dict[str, Any] = field(default_factory=dict)
    schema_version: int = IMPACT_SCHEMA_VERSION

    def __post_init__(self) -> None:
        if not self.event_id or not str(self.event_id).strip():
            raise ValueError("ImpactEvent.event_id is required")
        if not self.zone_id or not str(self.zone_id).strip():
            raise ValueError("ImpactEvent.zone_id is required")
        if self.end_date < self.start_date:
            raise ValueError(
                f"ImpactEvent {self.event_id}: end_date < start_date"
            )
        if not (0.0 <= float(self.confidence) <= 1.0):
            raise ValueError(
                f"ImpactEvent {self.event_id}: confidence out of [0,1]"
            )
        object.__setattr__(self, "extras", dict(self.extras))

    def contains_date(self, d: date) -> bool:
        return self.start_date <= d <= self.end_date

    def to_dict(self) -> Dict[str, Any]:
        return {
            "event_id": self.event_id,
            "zone_id": self.zone_id,
            "hazard": self.hazard.value,
            "start_date": self.start_date.isoformat(),
            "end_date": self.end_date.isoformat(),
            "source": self.source.value,
            "confidence": self.confidence,
            "notes": self.notes,
            "extras": dict(self.extras),
            "schema_version": self.schema_version,
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ImpactEvent":
        if not isinstance(d, dict):
            raise ValueError("ImpactEvent.from_dict expects object")
        sv = int(d.get("schema_version", IMPACT_SCHEMA_VERSION))
        if sv != IMPACT_SCHEMA_VERSION:
            raise ValueError(f"unsupported impact schema_version={sv}")
        return cls(
            event_id=str(d["event_id"]).strip(),
            zone_id=str(d["zone_id"]).strip(),
            hazard=ImpactHazard(str(d["hazard"]).strip().lower()),
            start_date=_parse_date(d["start_date"]),
            end_date=_parse_date(d["end_date"]),
            source=ImpactSource(str(d.get("source", "manual")).strip().lower()),
            confidence=float(d.get("confidence", 1.0)),
            notes=str(d.get("notes", "")),
            extras=dict(d.get("extras") or {}),
            schema_version=sv,
        )


def _parse_date(v: Any) -> date:
    if isinstance(v, date) and not isinstance(v, datetime):
        return v
    if isinstance(v, datetime):
        return v.date()
    s = str(v).strip()
    return date.fromisoformat(s[:10])


def make_event_id(
    source: ImpactSource,
    zone_id: str,
    start: date,
    hazard: ImpactHazard,
    suffix: str = "",
) -> str:
    """Deterministic id helper for offline curation."""
    base = f"{source.value}:{zone_id}:{start.isoformat()}:{hazard.value}"
    if suffix:
        base = f"{base}:{suffix}"
    # sanitize
    return re.sub(r"[^a-zA-Z0-9:_\-\.]", "_", base)


# ---------------------------------------------------------------------------
# Result envelope (ERROR HANDLING RULE)
# ---------------------------------------------------------------------------

@dataclass
class ImpactResult:
    success: bool
    outcome_code: str
    data: Dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "outcome_code": self.outcome_code,
            "data": dict(self.data),
        }


# ---------------------------------------------------------------------------
# Store (SSOT for L1 labels in-process)
# ---------------------------------------------------------------------------

class ImpactLabelStore:
    """
    In-memory index of ImpactEvent, optionally loaded from JSON.

    Canonical data lives here for a process. UI/backtest must query this
    (or a reload from the same file) — not invent parallel event lists.
    """

    def __init__(self) -> None:
        self._by_id: Dict[str, ImpactEvent] = {}
        self._source_path: Optional[str] = None

    def __len__(self) -> int:
        return len(self._by_id)

    def get(self, event_id: str) -> Optional[ImpactEvent]:
        return self._by_id.get(event_id)

    def all_events(self) -> List[ImpactEvent]:
        return list(self._by_id.values())

    def clear(self) -> None:
        self._by_id.clear()
        self._source_path = None

    def upsert(self, event: ImpactEvent) -> None:
        """Replace or insert by event_id (load path only)."""
        self._by_id[event.event_id] = event

    def query(
        self,
        zone_id: Optional[str] = None,
        start: Optional[date] = None,
        end: Optional[date] = None,
        hazard: Optional[ImpactHazard] = None,
    ) -> List[ImpactEvent]:
        """
        Return events overlapping [start, end] (inclusive) for zone/hazard filters.
        If start/end omitted, no date filter.
        """
        out: List[ImpactEvent] = []
        for ev in self._by_id.values():
            if zone_id is not None and ev.zone_id != zone_id:
                continue
            if hazard is not None and ev.hazard != hazard:
                continue
            if start is not None and end is not None:
                # overlap test
                if ev.end_date < start or ev.start_date > end:
                    continue
            elif start is not None and ev.end_date < start:
                continue
            elif end is not None and ev.start_date > end:
                continue
            out.append(ev)
        out.sort(key=lambda e: (e.start_date, e.zone_id, e.event_id))
        return out

    def labels_for_day(
        self,
        zone_id: str,
        d: date,
    ) -> Tuple[bool, bool]:
        """
        Convenience for backtest: (event_drought, event_flood) on calendar day d.
        OTHER hazards do not set either flag.
        """
        drought = flood = False
        for ev in self.query(zone_id=zone_id, start=d, end=d):
            if ev.hazard == ImpactHazard.DROUGHT:
                drought = True
            elif ev.hazard == ImpactHazard.FLOOD:
                flood = True
        return drought, flood


def load_impact_events(path: str | Path) -> ImpactResult:
    """
    Contract: load_impact_events

    Purpose: Load L1 impact JSON into a new ImpactLabelStore.
    Allowed caller: offline jobs, backtest, tests.
    Preconditions: path readable; top-level object with "events" array
                   OR a raw array of event objects.
    Forbidden: partial silent accept of invalid rows without reporting.
    Writes: none on disk; returns store in data["store"].
    Side effects: none.
    Response outcome_codes:
      LOADED | LOADED_WITH_ERRORS | FILE_NOT_FOUND | INVALID_JSON | INVALID_SCHEMA
    """
    p = Path(path)
    if not p.is_file():
        return ImpactResult(
            False,
            "FILE_NOT_FOUND",
            {"path": str(p)},
        )
    try:
        raw = json.loads(p.read_text(encoding="utf-8"))
    except json.JSONDecodeError as e:
        return ImpactResult(
            False,
            "INVALID_JSON",
            {"path": str(p), "error": str(e)},
        )

    if isinstance(raw, dict):
        events_raw = raw.get("events")
        if events_raw is None:
            return ImpactResult(
                False,
                "INVALID_SCHEMA",
                {"path": str(p), "error": "missing 'events' array"},
            )
    elif isinstance(raw, list):
        events_raw = raw
    else:
        return ImpactResult(
            False,
            "INVALID_SCHEMA",
            {"path": str(p), "error": "root must be object or array"},
        )

    store = ImpactLabelStore()
    store._source_path = str(p)
    errors: List[Dict[str, Any]] = []
    loaded = 0
    for i, row in enumerate(events_raw):
        try:
            ev = ImpactEvent.from_dict(row)
            store.upsert(ev)
            loaded += 1
        except (KeyError, TypeError, ValueError) as e:
            errors.append({"index": i, "error": str(e)})

    if loaded == 0 and errors:
        return ImpactResult(
            False,
            "INVALID_SCHEMA",
            {"path": str(p), "events_loaded": 0, "errors": errors},
        )
    code = "LOADED_WITH_ERRORS" if errors else "LOADED"
    return ImpactResult(
        True,
        code,
        {
            "path": str(p),
            "events_loaded": loaded,
            "error_count": len(errors),
            "errors": errors,
            "store": store,
        },
    )


def empty_store() -> ImpactLabelStore:
    return ImpactLabelStore()


def write_impact_events_template(path: str | Path) -> None:
    """Write an example JSON file for curators (not used at runtime)."""
    example = {
        "schema_version": IMPACT_SCHEMA_VERSION,
        "description": (
            "L1 impact labels. Replace with BNPB/provincial/AUTP-curated rows. "
            "Dates are ISO YYYY-MM-DD. zone_id must match indonesia_zones ids."
        ),
        "events": [
            {
                "event_id": "bnpb:karawang_rice:2023-08-15:drought:example",
                "zone_id": "karawang_rice",
                "hazard": "drought",
                "start_date": "2023-08-15",
                "end_date": "2023-09-30",
                "source": "manual",
                "confidence": 0.7,
                "notes": "EXAMPLE ONLY — replace with curated labels",
                "extras": {},
            }
        ],
    }
    Path(path).write_text(json.dumps(example, indent=2), encoding="utf-8")


# ---------------------------------------------------------------------------
# Evaluation helper (read-only projection)
# ---------------------------------------------------------------------------

def match_prediction_to_impact(
    store: ImpactLabelStore,
    zone_id: str,
    d: date,
    pred_drought: bool,
    pred_flood: bool,
) -> Dict[str, Any]:
    """
    Day-level confusion ingredients against L1 labels.
    Does not invent labels when store is empty — both GT flags false.
    """
    gt_d, gt_f = store.labels_for_day(zone_id, d)
    return {
        "zone_id": zone_id,
        "date": d.isoformat(),
        "gt_drought": gt_d,
        "gt_flood": gt_f,
        "pred_drought": bool(pred_drought),
        "pred_flood": bool(pred_flood),
        "hit_drought": bool(pred_drought and gt_d),
        "hit_flood": bool(pred_flood and gt_f),
        "fa_drought": bool(pred_drought and not gt_d),
        "fa_flood": bool(pred_flood and not gt_f),
        "miss_drought": bool((not pred_drought) and gt_d),
        "miss_flood": bool((not pred_flood) and gt_f),
    }


# ---------------------------------------------------------------------------
# Self-test
# ---------------------------------------------------------------------------

def _self_test() -> None:
    import tempfile

    print("impact_labels.py self-test")

    # empty file missing
    r = load_impact_events("/no/such/impact_labels.json")
    assert r.success is False and r.outcome_code == "FILE_NOT_FOUND"
    print("  FILE_NOT_FOUND OK")

    with tempfile.TemporaryDirectory() as td:
        path = Path(td) / "labels.json"
        write_impact_events_template(path)
        r = load_impact_events(path)
        assert r.success and r.outcome_code == "LOADED"
        store: ImpactLabelStore = r.data["store"]
        assert len(store) == 1
        print("  LOADED template OK")

        # query
        evs = store.query(zone_id="karawang_rice", start=date(2023, 9, 1), end=date(2023, 9, 1))
        assert len(evs) == 1 and evs[0].hazard == ImpactHazard.DROUGHT
        d_flag, f_flag = store.labels_for_day("karawang_rice", date(2023, 9, 1))
        assert d_flag and not f_flag
        print("  query / labels_for_day OK")

        # malformed row
        bad = {
            "events": [
                {"event_id": "ok", "zone_id": "z", "hazard": "flood",
                 "start_date": "2022-01-01", "end_date": "2022-01-02", "source": "manual"},
                {"event_id": "bad", "zone_id": "z"},  # missing fields
            ]
        }
        bad_path = Path(td) / "bad.json"
        bad_path.write_text(json.dumps(bad), encoding="utf-8")
        r2 = load_impact_events(bad_path)
        assert r2.success and r2.outcome_code == "LOADED_WITH_ERRORS"
        assert r2.data["events_loaded"] == 1
        print("  LOADED_WITH_ERRORS OK")

        # match helper
        m = match_prediction_to_impact(
            store, "karawang_rice", date(2023, 9, 1),
            pred_drought=True, pred_flood=False,
        )
        assert m["hit_drought"] and not m["miss_drought"]
        print("  match_prediction_to_impact OK")

    # entity guards
    try:
        ImpactEvent(
            event_id="x", zone_id="z", hazard=ImpactHazard.FLOOD,
            start_date=date(2023, 2, 1), end_date=date(2023, 1, 1),
            source=ImpactSource.MANUAL,
        )
        raise AssertionError("should reject end < start")
    except ValueError:
        print("  date order guard OK")

    print("All impact_labels self-tests passed.")


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    _self_test()