impact_labels.py

"""
impact_labels.py
================
"""

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):

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


class ImpactSource(str, Enum):

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


@dataclass(frozen=True)
class ImpactEvent:

    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:
    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:

    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:
        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]:
        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:
                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]:
        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:
    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:
    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]:
    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()