sensorflow/seqeval/planner.py

"""Frozen stratified sampling planner (hybrid Neyman + risk floors + safety minimums).

The plan is built from the POPULATION ONLY (dimension codes, counts, and a
candidate-independent prior on baseline variance) and is persisted with its
seed, config and a content hash BEFORE any candidate outcome is observed.
This is the sampling-bias firewall: no outcome-dependent selection is possible
because the selection provably cannot depend on outcomes (the planner never
sees them), and the hash gives auditors a cheap way to verify the plan did not
change after results started arriving.

Allocation (hybrid, in order):
  1. Neyman baseline: n_h proportional to N_h * sigma_h, where sigma_h =
     sqrt(p_h (1 - p_h)) uses a PRIOR per-stratum baseline success rate
     (deterministic profile; never candidate data).
  2. Risk weighting: VRU-class and night/adverse-condition strata get a
     multiplicative boost (they are where regressions hurt).
  3. Floors: every stratum gets a mandatory minimum; pre-registered
     safety-critical strata get a much larger minimum sized for the target MDE.
  4. Truncation at stratum population size; whole CLUSTERS (container groups)
     are selected so the inference unit stays intact.

Because allocation is non-proportional, every sampled object carries the
Horvitz-Thompson style weight w_h = N_h / n_h; population-level estimates are
weight-corrected and therefore remain unbiased despite oversampling.

The plan also stores a per-stratum escalation RESERVE (the remaining cluster
units in shuffled order) so the controller can expand the sample without
breaking the frozen-selection property: the reserve order was fixed at plan
time too.
"""

from __future__ import annotations

import hashlib
import json
import os
from typing import Dict, List, Optional

import numpy as np

from sensorflow.megaeval import population as pop_mod
from sensorflow.seqeval import ledger as ledger_mod
from sensorflow.seqeval import paired as paired_mod

DIFFICULTY_BANDS = [0.30, 0.55]  # easy < 0.30 <= medium < 0.55 <= hard
BAND_NAMES = ["easy", "medium", "hard"]

# Candidate-independent prior baseline success rates used for Neyman sigma_h.
_PRIOR_BASE = 0.90
_PRIOR_NIGHT_DROP = 0.04
_PRIOR_VRU_DROP = 0.03


def stratum_key(cls_name: str, cond_name: str) -> str:
    return f"{cls_name}|{cond_name}"


def _prior_rate(cls_code: int, cond_code: int, condition_dim: str) -> float:
    p = _PRIOR_BASE
    if condition_dim == "lighting" and cond_code == 1:
        p -= _PRIOR_NIGHT_DROP
    if cls_code in paired_mod.VRU_CLASSES:
        p -= _PRIOR_VRU_DROP
    return p


def _risk_weight(cls_code: int, cond_code: int, condition_dim: str) -> float:
    w = 1.0
    if cls_code in paired_mod.VRU_CLASSES:
        w *= 1.5
    if condition_dim == "lighting" and cond_code == 1:
        w *= 1.3
    return w


def plan_path(plan_id: str) -> str:
    return os.path.join(ledger_mod.seq_root(), "plans", plan_id)


def build_plan(meta: Dict, target_n: int, seed: int,
               condition_dim: str = "lighting",
               safety_primaries: Optional[List[str]] = None,
               min_per_stratum: int = 200,
               safety_floor: int = 2500,
               persist: bool = True) -> Dict:
    """Build (and persist) a frozen stratified cluster-sampling plan.

    Takes NO candidate/model arguments by design — see module docstring.
    """
    safety_primaries = list(safety_primaries or [])
    frame = paired_mod.load_frame(meta)
    cls_vocab = pop_mod.DIMENSIONS["class"]
    cond_vocab = pop_mod.DIMENSIONS[condition_dim]
    n_cond = len(cond_vocab)
    strata_ids = frame["class"].astype(np.int64) * n_cond + frame[condition_dim].astype(np.int64)
    n_strata = len(cls_vocab) * n_cond

    counts = np.bincount(strata_ids, minlength=n_strata)
    keys = [stratum_key(cls_vocab[s // n_cond], cond_vocab[s % n_cond])
            for s in range(n_strata)]

    # --- 1-2. Neyman shares with risk multipliers (prior variance, no outcomes)
    shares = np.zeros(n_strata)
    for s in range(n_strata):
        if counts[s] == 0:
            continue
        p = _prior_rate(s // n_cond, s % n_cond, condition_dim)
        sigma = np.sqrt(p * (1 - p))
        shares[s] = counts[s] * sigma * _risk_weight(s // n_cond, s % n_cond, condition_dim)
    shares = shares / shares.sum() if shares.sum() > 0 else shares

    # --- 3. floors and safety minimums, truncated at N_h
    alloc = np.zeros(n_strata, dtype=np.int64)
    for s in range(n_strata):
        if counts[s] == 0:
            continue
        n_h = int(round(target_n * shares[s]))
        n_h = max(n_h, min(min_per_stratum, int(counts[s])))
        if keys[s] in safety_primaries:
            n_h = max(n_h, min(safety_floor, int(0.85 * counts[s])))
        alloc[s] = min(n_h, int(counts[s]))

    # --- 4. whole-cluster selection in seeded shuffled order + reserve
    rng = np.random.default_rng([int(seed), int(meta["seed"]), 20260812])
    strata = {}
    arrays = {}
    for s in range(n_strata):
        if counts[s] == 0:
            continue
        rows = np.where(strata_ids == s)[0]
        containers = frame["container_id"][rows]
        order = np.argsort(containers, kind="stable")
        rows, containers = rows[order], containers[order]
        uniq, start = np.unique(containers, return_index=True)
        bounds = np.append(start, rows.size)
        cluster_order = rng.permutation(uniq.size)

        sel_ids, sel_bounds = [], [0]
        res_ids, res_bounds = [], [0]
        taken = 0
        for ci in cluster_order:
            members = frame["object_id"][rows[bounds[ci]:bounds[ci + 1]]]
            if taken < alloc[s]:
                sel_ids.append(members)
                sel_bounds.append(sel_bounds[-1] + members.size)
                taken += members.size
            else:
                res_ids.append(members)
                res_bounds.append(res_bounds[-1] + members.size)
        sel = np.concatenate(sel_ids) if sel_ids else np.empty(0, dtype=np.int64)
        res = np.concatenate(res_ids) if res_ids else np.empty(0, dtype=np.int64)
        arrays[f"s{s}_ids"] = sel
        arrays[f"s{s}_bounds"] = np.asarray(sel_bounds, dtype=np.int64)
        arrays[f"s{s}_res_ids"] = res
        arrays[f"s{s}_res_bounds"] = np.asarray(res_bounds, dtype=np.int64)
        strata[str(s)] = {
            "key": keys[s],
            "class": cls_vocab[s // n_cond],
            "condition": cond_vocab[s % n_cond],
            "N": int(counts[s]),
            "allocated": int(sel.size),
            "clusters": len(sel_bounds) - 1,
            "reserve": int(res.size),
            "weight": round(float(counts[s]) / max(sel.size, 1), 4),
            "safety_primary": keys[s] in safety_primaries,
        }

    config = {
        "target_n": int(target_n),
        "condition_dim": condition_dim,
        "safety_primaries": safety_primaries,
        "min_per_stratum": int(min_per_stratum),
        "safety_floor": int(safety_floor),
        "difficulty_bands": DIFFICULTY_BANDS,
        "allocation": "neyman(prior sigma) * risk_weights, floors, whole clusters",
    }
    h = hashlib.sha256()
    h.update(json.dumps({"dataset": paired_mod.dataset_fingerprint(meta),
                         "seed": int(seed), "config": config}, sort_keys=True).encode())
    for s in sorted(int(k) for k in strata):
        h.update(arrays[f"s{s}_ids"].tobytes())
        h.update(arrays[f"s{s}_res_ids"].tobytes())
    plan_hash = h.hexdigest()[:24]

    plan = {
        "plan_id": f"plan-{plan_hash[:10]}",
        "plan_hash": plan_hash,
        "population_id": meta["population_id"],
        "dataset_fingerprint": paired_mod.dataset_fingerprint(meta),
        "seed": int(seed),
        "created_at": ledger_mod.now_iso(),
        "config": config,
        "population_objects": int(meta["num_objects"]),
        "total_allocated": int(sum(v["allocated"] for v in strata.values())),
        "strata": strata,
        "frozen_before_candidate_outcomes": True,
    }
    if persist:
        d = plan_path(plan["plan_id"])
        os.makedirs(d, exist_ok=True)
        with open(os.path.join(d, "plan.json"), "w") as f:
            json.dump(plan, f)
        np.savez_compressed(os.path.join(d, "units.npz"), **arrays)
    plan["_arrays"] = arrays  # in-memory only; not serialized
    return plan


def load_plan(plan_id: str) -> Optional[Dict]:
    d = plan_path(plan_id)
    meta_path = os.path.join(d, "plan.json")
    if not os.path.exists(meta_path):
        return None
    with open(meta_path) as f:
        plan = json.load(f)
    with np.load(os.path.join(d, "units.npz")) as z:
        plan["_arrays"] = {k: z[k] for k in z.files}
    return plan


def stratum_clusters(plan: Dict, s: int, reserve: bool = False) -> List[np.ndarray]:
    """Ordered list of cluster-unit object_id arrays for stratum s."""
    suffix = "_res" if reserve else ""
    ids = plan["_arrays"].get(f"s{s}{suffix}_ids")
    bounds = plan["_arrays"].get(f"s{s}{suffix}_bounds")
    if ids is None or bounds is None or ids.size == 0:
        return []
    return [ids[bounds[i]:bounds[i + 1]] for i in range(bounds.size - 1)]