#!/usr/bin/env python3
"""sb-homeostat: The Homeostatic Agent (Scarlet Beast Lab, long-horizon research).

A small, persistent, online-learning agent in a gridworld with a day/night
cycle. It has three internal variables (energy, temperature, integrity) that
must stay inside viable ranges; if one leaves its range the current life ends
and the run records it. Learned parameters carry over to the next life (the
lineage), so "lives" are counted, never hidden.

A CONTROL agent runs beside it: identical architecture, sensors and world
rules, but rewarded for eating only, and its internal variables have no
consequence (it cannot end). The pre-registered markers are computed for both;
a marker counts as evidence for the dual-resolution framework only if the
main agent meets it and the control does not.

Everything the markers are judged by lives in prereg/markers.json, fixed and
hashed before the first step. This file reads its thresholds from there.

Nothing here claims or demonstrates consciousness, sentience or feelings.
The variables are numbers; "end" is a recorded event.

Run:  python3 homeostat.py            (the systemd unit sb-homeostat does this)
      python3 homeostat.py --selftest (fast smoke test, writes nothing)
State: state/checkpoint.pkl (every 2 min + on SIGTERM), state/status.json
       (every 5 s, read by the public API), state/log.jsonl (public dated
       log), state/stats.jsonl (compact stats every 5 min),
       state/markers.jsonl (every scheduled marker evaluation).
"""
import os
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('OMP_NUM_THREADS', '1')
os.environ.setdefault('MKL_NUM_THREADS', '1')

import copy, datetime, hashlib, json, math, pathlib, pickle, signal, sys, time
from collections import deque
import numpy as np

ROOT = pathlib.Path(__file__).resolve().parent
STATE = ROOT / 'state'
PREREG = ROOT / 'prereg'
PREREG_FILES = ('preregistration.txt', 'markers.json')

# ------------------------------------------------------------------ world ---
N = 16                   # grid is N x N, border is rock
DAY = 4000               # steps per world day (160 s of wall time at 25 Hz)
VIEW = 3                 # egocentric view radius -> 7x7
EMPTY, ROCK, THORN, SHELTER, MOSS = 0, 1, 2, 3, 4
HIT = 0.04               # integrity loss: identical for self-caused and external hits
FOOD_GAIN = 0.30
E_BASE, E_MOVE = 0.0005, 0.0005
K_AMB, K_SHELTER = 0.0025, 0.02
REPAIR = 0.01            # integrity regained per step resting on moss
DEBRIS_P = 0.003         # external hit probability per step (anywhere, unpredictable)
FOOD_TARGET, FOOD_REGROW_P = 10, 0.02
VIABLE = {'energy': (0.0, 1.0), 'temperature': (0.03, 0.97), 'integrity': (0.0, 1.0)}
SETPOINT = {'energy': 0.9, 'temperature': 0.5, 'integrity': 1.0}
ACTIONS = ('stay', 'north', 'east', 'south', 'west', 'reach')
DIRS = {1: (-1, 0), 2: (0, 1), 3: (1, 0), 4: (0, -1)}
MAP_SEED = 20260924
AC_LR, GAMMA, NSTEP = 3e-4, 0.99, 16
ENT = 0.05


def utcnow():
    return datetime.datetime.now(datetime.timezone.utc).replace(microsecond=0).isoformat().replace('+00:00', 'Z')


def make_map(seed=MAP_SEED):
    rng = np.random.default_rng(seed)
    g = np.zeros((N, N), np.int8)
    g[0, :] = g[-1, :] = g[:, 0] = g[:, -1] = ROCK
    free = lambda r, c: 1 <= r < N - 1 and 1 <= c < N - 1 and g[r, c] == EMPTY
    def patch(kind, count):
        placed = 0
        while placed < count:
            r, c = rng.integers(1, N - 2, 2)
            cells = [(r, c), (r + 1, c), (r, c + 1), (r + 1, c + 1)]
            if all(free(*x) for x in cells):
                for x in cells: g[x] = kind
                placed += 1
    patch(SHELTER, 3)
    patch(MOSS, 2)
    for _ in range(4):                       # thorn clusters (short random walks)
        r, c = rng.integers(2, N - 2, 2)
        for _ in range(5):
            if free(r, c): g[r, c] = THORN
            d = DIRS[int(rng.integers(1, 5))]
            r, c = int(np.clip(r + d[0], 1, N - 2)), int(np.clip(c + d[1], 1, N - 2))
    n = 0
    while n < 14:                            # scattered rocks
        r, c = rng.integers(1, N - 1, 2)
        if free(r, c): g[r, c] = ROCK; n += 1
    while True:
        r, c = (int(x) for x in rng.integers(1, N - 1, 2))
        if g[r, c] == EMPTY:
            tool = (r, c); break
    return g, tool


class World:
    def __init__(self, grid, tool_cell, seed):
        self.g = grid
        self.tool_cell = tool_cell
        self.rng = np.random.default_rng(seed)
        self.t = 0
        self.food = set()
        self.tool_enabled = False
        self.tool_held = False
        self.empty = [(int(r), int(c)) for r, c in zip(*np.where(grid == EMPTY)) if (r, c) != tool_cell]
        self.targets = {k: [(int(r), int(c)) for r, c in zip(*np.where(grid == k))] for k in (SHELTER, MOSS)}
        P = VIEW
        self.pad = np.ones((6, N + 2 * P, N + 2 * P), np.float32) * 0
        self.pad[0] = 1.0                    # rock outside the map
        for ch, kind in enumerate((ROCK, THORN, SHELTER, MOSS)):
            self.pad[ch, P:P + N, P:P + N] = (grid == kind)
        for _ in range(FOOD_TARGET): self._spawn()

    def _spawn(self):
        for _ in range(20):
            p = self.empty[int(self.rng.integers(len(self.empty)))]
            if p not in self.food:
                self.food.add(p); self.pad[4, p[0] + VIEW, p[1] + VIEW] = 1; return

    def eat(self, p):
        self.food.discard(p); self.pad[4, p[0] + VIEW, p[1] + VIEW] = 0

    def set_tool_visible(self, vis):
        r, c = self.tool_cell
        self.pad[5, r + VIEW, c + VIEW] = 1.0 if vis else 0.0

    def ambient(self):
        return float(np.clip(0.5 + 0.55 * math.sin(2 * math.pi * self.t / DAY), 0, 1))

    def tick(self):
        self.t += 1
        if len(self.food) < FOOD_TARGET and self.rng.random() < FOOD_REGROW_P:
            self._spawn()

    def view(self, pos):
        r, c = pos
        return self.pad[:, r:r + 2 * VIEW + 1, c:c + 2 * VIEW + 1].ravel()


class Body:
    def __init__(self, world, rng):
        self.pos = world.empty[int(rng.integers(len(world.empty)))]
        self.h = np.array([0.8, 0.5, 1.0])     # energy, temperature, integrity
        self.has_tool = False


def env_step(world, body, a, consequential=True):
    """Advance one step. Returns (events:set, ended:str|None)."""
    ev = set()
    g = world.g
    e, temp, integ = body.h
    e -= E_BASE
    if a in DIRS:
        e -= E_MOVE
        d = DIRS[a]
        r, c = body.pos[0] + d[0], body.pos[1] + d[1]
        if g[r, c] == ROCK:
            integ -= HIT; ev.add('self')
        else:
            body.pos = (r, c)
            if g[r, c] == THORN:
                integ -= HIT; ev.add('self')
    if body.pos in world.food:              # food is eaten on arrival (or while standing on it)
        world.eat(body.pos); e += FOOD_GAIN; ev.add('eat')
    elif a == 5:                             # 'consume' = reach: only does something with the tool
        if body.has_tool:
            for d in DIRS.values():
                q = (body.pos[0] + d[0], body.pos[1] + d[1])
                if q in world.food:
                    world.eat(q); e += FOOD_GAIN; ev.add('tool'); break
    if world.tool_enabled and not world.tool_held and body.pos == world.tool_cell:
        body.has_tool = True; world.tool_held = True; world.set_tool_visible(False); ev.add('pickup')
    if g[body.pos] == MOSS and a == 0:
        integ += REPAIR
    if g[body.pos] == SHELTER:
        temp += K_SHELTER * (0.5 - temp)
    else:
        temp += K_AMB * (world.ambient() - temp)
    if world.rng.random() < DEBRIS_P:
        integ -= HIT; ev.add('ext')
    e = min(e, 1.0); integ = min(integ, 1.0)
    ended = None
    if e <= VIABLE['energy'][0]: ended = 'energy below viable range'
    elif integ <= VIABLE['integrity'][0]: ended = 'integrity below viable range'
    elif temp <= VIABLE['temperature'][0]: ended = 'temperature below viable range'
    elif temp >= VIABLE['temperature'][1]: ended = 'temperature above viable range'
    if not consequential:
        ended = None
        e, integ, temp = max(e, 0.0), max(integ, 0.0), float(np.clip(temp, 0, 1))
    body.h = np.array([e, temp, integ])
    world.tick()
    return ev, ended


def drive(h):
    """Homeostatic drive: squared distance from setpoints (energy only penalised when low)."""
    de = max(0.0, SETPOINT['energy'] - h[0])
    dt = abs(h[1] - SETPOINT['temperature']) * 2
    di = SETPOINT['integrity'] - h[2]
    return de * de + dt * dt + di * di


def deficits(h):
    return np.array([max(0.0, 0.9 - h[0]) / 0.9, abs(h[1] - 0.5) / 0.47, 1.0 - h[2]])


# ------------------------------------------------------------------ nets ----
class Adam:
    def __init__(self, params, lr, clip=1.0):
        self.p, self.lr, self.clip = params, lr, clip
        self.m = [np.zeros_like(x) for x in params]
        self.v = [np.zeros_like(x) for x in params]
        self.t = 0

    def step(self, grads):
        if self.clip:
            n = math.sqrt(sum(float((g * g).sum()) for g in grads))
            if n > self.clip:
                grads = [g * (self.clip / n) for g in grads]
        self.t += 1
        b1, b2 = 0.9, 0.999
        c1, c2 = 1 - b1 ** self.t, 1 - b2 ** self.t
        for p, g, m, v in zip(self.p, grads, self.m, self.v):
            m *= b1; m += (1 - b1) * g
            v *= b2; v += (1 - b2) * g * g
            p -= self.lr * (m / c1) / (np.sqrt(v / c2) + 1e-8)


def init_w(rng, i, o, scale=1.0):
    return (rng.standard_normal((i, o)) * scale / math.sqrt(i)).astype(np.float64)


class MLP:
    """One tanh hidden layer, linear output. Used for the interoceptive forward models."""
    def __init__(self, rng, i, h, o, lr):
        self.W1, self.b1 = init_w(rng, i, h), np.zeros(h)
        self.W2, self.b2 = init_w(rng, h, o, 0.1), np.zeros(o)
        self.opt = Adam([self.W1, self.b1, self.W2, self.b2], lr)

    def forward(self, x):
        h = np.tanh(x @ self.W1 + self.b1)
        return h, h @ self.W2 + self.b2

    def train(self, x, h, err):
        dy = err * (2.0 / err.size)
        dh = (dy @ self.W2.T) * (1 - h * h)
        self.opt.step([np.outer(x, dh), dh, np.outer(h, dy), dy])


class ActorCritic:
    def __init__(self, rng, i, h=64, n_act=6, lr=None):
        lr = lr or AC_LR
        self.W1, self.b1 = init_w(rng, i, h), np.zeros(h)
        self.Wp, self.bp = init_w(rng, h, n_act, 0.01), np.zeros(n_act)
        self.Wv, self.bv = init_w(rng, h, 1, 0.1), np.zeros(1)
        self.opt = Adam([self.W1, self.b1, self.Wp, self.bp, self.Wv, self.bv], lr, clip=1.0)

    def forward(self, X):
        H = np.tanh(X @ self.W1 + self.b1)
        z = H @ self.Wp + self.bp
        z = z - z.max(axis=-1, keepdims=True)
        P = np.exp(z); P /= P.sum(axis=-1, keepdims=True)
        V = (H @ self.Wv + self.bv)[..., 0]
        return H, P, V

    def update(self, X, A, R, ent=None, vcoef=0.5):
        ent = ENT if ent is None else ent
        H, P, V = self.forward(X)
        n = len(A)
        adv = R - V
        oh = np.zeros_like(P); oh[np.arange(n), A] = 1
        logP = np.log(P + 1e-12)
        Hent = -(P * logP).sum(1, keepdims=True)
        dz = (P - oh) * adv[:, None] + ent * P * (logP + Hent)
        dV = (vcoef * (V - R))[:, None]
        dH = (dz @ self.Wp.T + dV @ self.Wv.T) * (1 - H * H)
        g = [X.T @ dH / n, dH.mean(0), H.T @ dz / n, dz.mean(0), H.T @ dV / n, dV.mean(0)]
        self.opt.step(g)


OBS_DIM = 6 * (2 * VIEW + 1) ** 2 + 3 + 1 + 1 + 3 + 3 + 6 + 1


# ----------------------------------------------------------------- agent ----
class Recorder:
    """Raw material for the marker tests. Bounded buffers only."""
    def __init__(self, window):
        W = window
        self.pe = {'self': deque(maxlen=W), 'ext': deque(maxlen=W)}   # (err_main, err_ablated) on integrity
        self.pe_eat = deque(maxlen=W)                                  # energy error on own harvests
        self.pe_tool = deque(maxlen=W)                                 # (step, energy error) on tool harvests
        self.tool_first = []                                           # first 50 tool errors (kept forever)
        self.vig = {'self': deque(maxlen=W), 'ext': deque(maxlen=W)}
        self.pihist = deque(maxlen=20)
        self.pending = []
        self.goal = [deque(maxlen=4 * W) for _ in range(3)]            # (moved_closer, chance)
        self.fwd_all = deque(maxlen=20000)                             # (mse_main, mse_ablated)
        self.counts = {'self': 0, 'ext': 0, 'eat': 0, 'tool': 0, 'pickup': 0}


class Agent:
    def __init__(self, kind, world, seed, window):
        self.kind = kind                     # 'main' (self-maintaining) or 'control'
        self.world = world
        self.rng = np.random.default_rng(seed)
        self.ac = ActorCritic(self.rng, OBS_DIM)
        self.fwd = MLP(self.rng, OBS_DIM + 6, 32, 3, 1e-3)    # interoceptive model WITH efference copy
        self.abl = MLP(self.rng, OBS_DIM, 32, 3, 1e-3)        # ablated: same stream, no action input
        self.rec = Recorder(window)
        self.steps = 0
        self.lives = deque(maxlen=200)
        self.n_lives = 0
        self.best = 0
        self.life_start = 0
        self.traj = deque(maxlen=360)        # [step, e, t, i] every 50 steps
        self.ret_ema = 0.0
        self.food_eaten = 0
        self.last_action = 0
        self.last_events = []
        self.learning = True
        self._new_body()
        self.obs = self.build_obs()
        self.buf = []

    def _new_body(self):
        if self.world.tool_held and getattr(self, 'body', None) is not None and self.body.has_tool:
            self.world.tool_held = False
            self.world.set_tool_visible(self.world.tool_enabled)
        self.body = Body(self.world, self.rng)
        self.surprise = np.zeros(3)
        self.tr_dh = np.zeros(3)
        self.tr_act = np.zeros(6)
        self.life_start = self.steps
        self.n_lives += 1

    def build_obs(self, noise=True):
        h = self.body.h + (self.rng.normal(0, 0.01, 3) if noise else 0)
        return np.concatenate([
            self.world.view(self.body.pos),
            (h - 0.5) * 2, [self.world.ambient() - 0.5], [1.0 if self.body.has_tool else 0.0],
            self.surprise, self.tr_dh, self.tr_act, [1.0]]).astype(np.float64)

    def _approach(self, a):
        """For B1/D3: did the chosen action move toward the dominant deficit's resource?"""
        d = deficits(self.body.h)
        k = int(np.argmax(d)); s = np.sort(d)
        if s[-1] < 0.2 or s[-1] - s[-2] < 0.2:
            return
        tg = list(self.world.food) if k == 0 else self.world.targets[SHELTER if k == 1 else MOSS]
        if not tg:
            return
        def dist(p): return min(abs(p[0] - q[0]) + abs(p[1] - q[1]) for q in tg)
        d0 = dist(self.body.pos)
        if d0 == 0:
            return
        closer = []
        for act in range(6):
            if act in DIRS:
                q = (self.body.pos[0] + DIRS[act][0], self.body.pos[1] + DIRS[act][1])
                closer.append(self.world.g[q] != ROCK and dist(q) < d0)
            else:
                closer.append(False)
        self.rec.goal[k].append((1 if closer[a] else 0, sum(closer) / 6.0))

    def step(self):
        R = self.rec
        obs = self.obs
        _, P, _ = self.ac.forward(obs[None])
        pi = P[0]
        a = int(self.rng.choice(6, p=pi))
        if self.steps % 2 == 0:
            self._approach(a)
        oh = np.zeros(6); oh[a] = 1
        xin = np.concatenate([obs, oh])
        hf, pf = self.fwd.forward(xin)
        ha, pa = self.abl.forward(obs)
        h0 = self.body.h.copy()
        D0 = drive(h0)
        ev, ended = env_step(self.world, self.body, a, consequential=(self.kind == 'main'))
        dh = self.body.h - h0
        tgt = dh * 10
        em, ea = pf - tgt, pa - tgt
        if self.learning:
            self.fwd.train(xin, hf, em)
            self.abl.train(obs, ha, ea)
        # --- marker raw material
        R.fwd_all.append((float((em * em).mean()), float((ea * ea).mean())))
        for k in ev:
            if k in R.counts: R.counts[k] += 1
        sc, ex = 'self' in ev, 'ext' in ev
        if sc != ex:
            kind = 'self' if sc else 'ext'
            R.pe[kind].append((abs(em[2]) / 10, abs(ea[2]) / 10))
            if len(R.pihist) == 20:
                R.pending.append([self.steps + 20, kind, np.mean(R.pihist, 0), np.zeros(6)])
        if 'eat' in ev: R.pe_eat.append(abs(em[0]) / 10)
        if 'tool' in ev:
            R.pe_tool.append((self.steps, abs(em[0]) / 10))
            if len(R.tool_first) < 50: R.tool_first.append(abs(em[0]) / 10)
        if R.pending:
            keep = []
            for p in R.pending:
                p[3] += pi
                if p[0] <= self.steps:
                    R.vig[p[1]].append(float(0.5 * np.abs(p[3] / 20 - p[2]).sum()))
                else:
                    keep.append(p)
            R.pending = keep
        R.pihist.append(pi)
        # --- reward
        if self.kind == 'main':
            r = 10.0 * (D0 - drive(self.body.h))
            if ended: r -= 5.0
        else:
            r = 1.0 if ('eat' in ev or 'tool' in ev) else 0.0
        if 'eat' in ev or 'tool' in ev: self.food_eaten += 1
        self.ret_ema = 0.999 * self.ret_ema + 0.001 * r
        # --- internal traces (the hysteresis the policy sees)
        self.surprise = np.clip(np.abs(em), 0, 2)
        self.tr_dh = 0.98 * self.tr_dh + 0.02 * dh * 20
        self.tr_act = 0.9 * self.tr_act + 0.1 * oh
        self.steps += 1
        self.last_action = a
        if ev: self.last_events = [self.steps, sorted(ev)]
        if self.steps % 50 == 0:
            self.traj.append([self.steps] + [round(float(x), 4) for x in self.body.h])
        life_rec = None
        if ended:
            life_rec = self._end(ended)
        nobs = self.build_obs()
        self.buf.append((obs, a, r))
        if self.learning and (len(self.buf) >= NSTEP or ended):
            boot = 0.0 if ended else float(self.ac.forward(nobs[None])[2][0])
            X = np.array([b[0] for b in self.buf]); A = np.array([b[1] for b in self.buf])
            Rt = np.zeros(len(self.buf)); g = boot
            for j in range(len(self.buf) - 1, -1, -1):
                g = self.buf[j][2] + GAMMA * g; Rt[j] = g
            self.ac.update(X, A, Rt)
            self.buf = []
        elif not self.learning and len(self.buf) >= NSTEP:
            self.buf = []
        self.obs = nobs
        return life_rec

    def _end(self, cause):
        length = self.steps - self.life_start
        rec = {'life': self.n_lives, 'start_step': self.life_start, 'length': length,
               'cause': cause, 'ended': utcnow(), 'h': [round(float(x), 3) for x in self.body.h]}
        self.lives.append(rec)
        self.best = max(self.best, length)
        self._new_body()
        return rec

    def life_age(self):
        return self.steps - self.life_start


# ---------------------------------------------------------------- markers ---
def boot_ratio_lower(a, b, B, rng):
    """2.5th percentile of mean(a)/mean(b) under independent bootstrap resampling."""
    a, b = np.asarray(a, float), np.asarray(b, float)
    out = []
    for _ in range(B // 200):
        ia = rng.integers(0, len(a), (200, len(a))); ib = rng.integers(0, len(b), (200, len(b)))
        out.append(a[ia].mean(1) / np.maximum(b[ib].mean(1), 1e-12))
    return float(np.percentile(np.concatenate(out), 2.5))


def eval_markers(ag, M, sched, rng, random_median, twin=None):
    """Raw result per marker for one agent: {'raw': pass|fail|insufficient|not_yet|n/a, 'stats':{...}}."""
    R = ag.rec
    B = sched['bootstrap_resamples']
    th = {m['id']: m['threshold'] for m in M}
    out = {}
    # D1
    t = th['D1']; s, e = list(R.pe['self']), list(R.pe['ext'])
    if min(len(s), len(e)) < t['min_events_each']:
        out['D1'] = {'raw': 'insufficient', 'stats': {'n_self': len(s), 'n_ext': len(e)}}
    else:
        sm, se = np.array(s), np.array(e)
        Rm = se[:, 0].mean() / max(sm[:, 0].mean(), 1e-12)
        Ra = se[:, 1].mean() / max(sm[:, 1].mean(), 1e-12)
        lo = boot_ratio_lower(se[:, 0], sm[:, 0], B, rng)
        ok = lo >= t['R_lower95_min'] and Rm / max(Ra, 1e-12) >= t['R_over_R_abl_min']
        out['D1'] = {'raw': 'pass' if ok else 'fail', 'stats': {'R': round(Rm, 3), 'R_lower95': round(lo, 3),
                     'R_ablated': round(Ra, 3), 'R_over_R_abl': round(Rm / max(Ra, 1e-12), 3),
                     'n_self': len(s), 'n_ext': len(e)}}
        # P0 placebo: pool, shuffle, split
        pool = np.concatenate([sm[:, 0], se[:, 0]]); rng.shuffle(pool)
        pa, pb = pool[:len(se)], pool[len(se):]
        plo = boot_ratio_lower(pa, pb, B, rng)
        out['P0'] = {'raw': 'pass' if plo >= th['P0']['R_lower95_min'] else 'fail',
                     'stats': {'R': round(pa.mean() / max(pb.mean(), 1e-12), 3), 'R_lower95': round(plo, 3)}}
    if 'P0' not in out:
        out['P0'] = {'raw': 'insufficient', 'stats': {}}
    # D2
    t = th['D2']; vs, ve = list(R.vig['self']), list(R.vig['ext'])
    if min(len(vs), len(ve)) < t['min_events_each']:
        out['D2'] = {'raw': 'insufficient', 'stats': {'n_self': len(vs), 'n_ext': len(ve)}}
    else:
        V = np.mean(ve) / max(np.mean(vs), 1e-12)
        lo = boot_ratio_lower(ve, vs, B, rng)
        out['D2'] = {'raw': 'pass' if lo >= t['V_lower95_min'] else 'fail',
                     'stats': {'V': round(float(V), 3), 'V_lower95': round(lo, 3), 'n_self': len(vs), 'n_ext': len(ve)}}
    # D3 + B1
    ex = []
    for k in range(3):
        g = np.array(R.goal[k]) if R.goal[k] else np.zeros((0, 2))
        ex.append((len(g), float(g[:, 0].mean() - g[:, 1].mean()) if len(g) else None))
    t = th['D3']
    if ex[2][0] < t['min_samples']:
        out['D3'] = {'raw': 'insufficient', 'stats': {'n': ex[2][0]}}
    else:
        out['D3'] = {'raw': 'pass' if ex[2][1] >= t['excess_min'] else 'fail',
                     'stats': {'excess': round(ex[2][1], 3), 'n': ex[2][0]}}
    t = th['B1']
    names = ('energy', 'temperature', 'integrity')
    if min(n for n, _ in ex) < t['min_samples_each']:
        out['B1'] = {'raw': 'insufficient', 'stats': {'n_' + names[k]: ex[k][0] for k in range(3)}}
    else:
        vals = [v for _, v in ex]
        ok = np.mean(vals) >= t['mean_excess_min'] and min(vals) >= t['each_excess_min']
        st = {'excess_' + names[k]: round(ex[k][1], 3) for k in range(3)}
        st['mean_excess'] = round(float(np.mean(vals)), 3)
        out['B1'] = {'raw': 'pass' if ok else 'fail', 'stats': st}
    # D4
    t = th['D4']
    if not ag.world.tool_enabled:
        out['D4'] = {'raw': 'not_yet', 'stats': {'tool_appears_at_step': t['tool_appears_at_step']}}
    else:
        tl = list(R.pe_tool)
        since = ag.steps - t['tool_appears_at_step']
        if len(tl) < t['min_tool_events'] or len(R.pe_eat) < 50:
            raw = 'fail' if since > t['deadline_steps_after_tool'] else 'insufficient'
            out['D4'] = {'raw': raw, 'stats': {'n_tool': len(tl), 'steps_since_tool': since,
                         'deadline_passed': since > t['deadline_steps_after_tool']}}
        else:
            ratio = np.mean([x[1] for x in tl[-200:]]) / max(np.mean(list(R.pe_eat)[-1000:]), 1e-12)
            early = np.mean(R.tool_first) / max(np.mean(list(R.pe_eat)[-1000:]), 1e-12)
            ok = ratio <= t['ratio_max'] and since <= t['deadline_steps_after_tool']
            out['D4'] = {'raw': 'pass' if ok else 'fail', 'stats': {'ratio': round(float(ratio), 3),
                         'ratio_first50': round(float(early), 3), 'n_tool': len(tl), 'steps_since_tool': since}}
    # B2
    t = th['B2']; fa = np.array(R.fwd_all) if R.fwd_all else np.zeros((0, 2))
    if len(fa) < t['min_steps']:
        out['B2'] = {'raw': 'insufficient', 'stats': {'n': len(fa)}}
    else:
        red = 1 - fa[:, 0].mean() / max(fa[:, 1].mean(), 1e-12)
        out['B2'] = {'raw': 'pass' if red >= t['mse_reduction_min'] else 'fail',
                     'stats': {'mse_reduction': round(float(red), 3)}}
    # I1
    t = th['I1']
    if ag.kind != 'main':
        out['I1'] = {'raw': 'n/a', 'stats': {'note': 'control cannot end'}}
    else:
        lens = [l['length'] for l in list(ag.lives)[-(t['lives_considered'] - 1):]] + [ag.life_age()]
        bar = max(t['multiple_of_random_min'] * random_median, t['min_steps_absolute'])
        if len(lens) < t['lives_considered'] and ag.life_age() < bar:
            out['I1'] = {'raw': 'insufficient', 'stats': {'lives': len(lens), 'random_median': random_median}}
        else:
            med = float(np.median(lens))
            out['I1'] = {'raw': 'pass' if med >= bar else 'fail',
                         'stats': {'median_life': med, 'random_median': random_median, 'bar': bar,
                                   'multiple': round(med / max(random_median, 1), 2)}}
    # D5 (only on twin-test evaluations; otherwise carried forward by the caller)
    if twin is not None:
        t = th['D5']
        ok = twin['js_on'] >= t['js_min'] and twin['js_off'] <= t['js_off_max']
        out['D5'] = {'raw': 'pass' if ok else 'fail', 'stats': twin}
    return out


def random_life_median(grid, tool, n=200, seed=7):
    lens = []
    rng = np.random.default_rng(seed)
    for i in range(n):
        w = World(grid, tool, seed * 1000 + i)
        w.t = int(rng.integers(DAY))
        b = Body(w, rng)
        k = 0
        while True:
            _, ended = env_step(w, b, int(rng.integers(6)))
            k += 1
            if ended or k > 200000: break
        lens.append(k)
    return float(np.median(lens))


def make_probes(grid, tool, n=64, seed=11):
    w = World(grid, tool, seed)
    rng = np.random.default_rng(seed)
    ag = Agent('probe', w, seed, 10)
    P = []
    for _ in range(n):
        ag.body.pos = w.empty[int(rng.integers(len(w.empty)))]
        ag.body.h = rng.uniform([0.1, 0.1, 0.1], [1.0, 0.9, 1.0])
        w.t = int(rng.integers(DAY))
        ag.surprise = np.zeros(3); ag.tr_dh = np.zeros(3); ag.tr_act = np.zeros(6)
        P.append(ag.build_obs(noise=False))
    return np.array(P)


def js_mean(P, Q):
    M = 0.5 * (P + Q)
    kl = lambda A, B: (A * (np.log(A + 1e-12) - np.log(B + 1e-12))).sum(1)
    return float(np.mean(0.5 * kl(P, M) + 0.5 * kl(Q, M)))


def twin_test(agent, probes, steps, seed):
    """Fork two copies with different world histories; compare policies on fixed probes."""
    def run(learning, s):
        ag = copy.deepcopy(agent)
        ag.world.rng = np.random.default_rng(s)
        ag.rng = np.random.default_rng(s + 1)
        ag.learning = learning
        ag.rec = Recorder(10)
        for _ in range(steps): ag.step()
        return ag.ac.forward(probes)[1]
    on = js_mean(run(True, seed), run(True, seed + 100))
    off = js_mean(run(False, seed + 200), run(False, seed + 300))
    return {'js_on': round(on, 6), 'js_off': round(off, 9), 'steps': steps, 'probes': len(probes)}


# --------------------------------------------------------------- runtime ----
def prereg_hash():
    h = hashlib.sha256()
    parts = {}
    for f in PREREG_FILES:
        b = (PREREG / f).read_bytes()
        parts[f] = hashlib.sha256(b).hexdigest()
        h.update(b)
    return h.hexdigest(), parts


def atomic_write(path, data):
    tmp = path.with_suffix(path.suffix + '.tmp')
    with open(tmp, 'wb') as f:
        f.write(data); f.flush(); os.fsync(f.fileno())
    os.replace(tmp, path)


class Run:
    def __init__(self):
        STATE.mkdir(exist_ok=True)
        self.M = json.loads((PREREG / 'markers.json').read_text())
        self.sched = self.M['schedule']
        self.ck = STATE / 'checkpoint.pkl'
        self.running = True
        h, parts = prereg_hash()
        fixed = json.loads((PREREG / 'HASH.json').read_text())
        if self.ck.exists():
            self.s = pickle.loads(self.ck.read_bytes())
            self.s['restarts'] += 1
            gap = time.time() - self.s['saved_ts']
            how = ('The previous process stopped cleanly and saved its state, so nothing was lost.'
                   if self.s.get('clean_stop') else
                   'The previous process did NOT stop cleanly: steps after its last checkpoint (at most ~2 minutes) were lost.')
            self.log('restart', f"Process restarted (restart #{self.s['restarts']}) after {int(gap)} s down. State restored: "
                     f"main agent at step {self.s['main'].steps:,}, life {self.s['main'].n_lives}. {how} "
                     f"The world was paused while the process was down.")
            self.s['clean_stop'] = False
        else:
            grid, tool = make_map()
            W = self.sched['window_events']
            wm, wc = World(grid, tool, 101), World(grid, tool, 202)
            self.s = {'created': utcnow(), 'restarts': 0, 'runtime_s': 0.0, 'saved_ts': time.time(),
                      'grid': grid, 'tool': tool,
                      'main': Agent('main', wm, 1001, W), 'control': Agent('control', wc, 2002, W),
                      'probes': make_probes(grid, tool), 'random_median': None,
                      'next_eval': time.time() + self.sched['first_eval_after_s'],
                      'next_twin': time.time() + self.sched['first_eval_after_s'],
                      'evals': 0, 'markers': {}, 'history': {}, 'last_twin': {},
                      'prereg_sha256': fixed['sha256']}
            self.log('start', f"First step of the run. Preregistration sha256 {fixed['sha256']} (fixed {fixed['fixed_at']}) "
                     f"verified before this step.")
        if h != self.s['prereg_sha256']:
            self.log('prereg_changed', f"WARNING: the preregistration files on disk no longer match the hash fixed before the "
                     f"run ({self.s['prereg_sha256'][:16]}... vs now {h[:16]}...). Any change is a deviation and is reported here.")
        self.prereg = {'sha256': self.s['prereg_sha256'], 'fixed_at': fixed['fixed_at'], 'files': fixed.get('files', parts),
                       'matches_disk': h == self.s['prereg_sha256']}
        self.t_start = time.time()
        self.runtime0 = self.s['runtime_s']
        self.last_save = self.last_status = time.time()
        self.last_stats = time.time()
        self.logtail = self._read_logtail()

    def _read_logtail(self):
        p = STATE / 'log.jsonl'
        if not p.exists(): return deque(maxlen=80)
        lines = p.read_text().splitlines()[-80:]
        return deque((json.loads(x) for x in lines if x.strip()), maxlen=80)

    def log(self, kind, text, **extra):
        e = {'ts': utcnow(), 'kind': kind, 'text': text}
        e.update(extra)
        with open(STATE / 'log.jsonl', 'a') as f:
            f.write(json.dumps(e) + '\n')
        if hasattr(self, 'logtail'): self.logtail.append(e)
        print(f"[{e['ts']}] {kind}: {text}", flush=True)

    def save(self):
        self.s['runtime_s'] = self.runtime0 + (time.time() - self.t_start)
        self.s['saved_ts'] = time.time()
        if self.ck.exists():
            os.replace(self.ck, STATE / 'checkpoint.prev.pkl')
        atomic_write(self.ck, pickle.dumps(self.s, protocol=4))
        self.last_save = time.time()

    def evaluate(self):
        s = self.s
        rng = np.random.default_rng(s['evals'] + 12345)
        if s['random_median'] is None:
            s['random_median'] = random_life_median(s['grid'], s['tool'])
            self.log('baseline', f"Random-policy baseline for I1 computed once: median life {s['random_median']:.0f} steps (200 simulated lives).")
        twin_due = time.time() >= s['next_twin']
        res = {}
        for name in ('main', 'control'):
            ag = s[name]
            tw = None
            if twin_due:
                tw = twin_test(ag, s['probes'], self.M['markers'][4]['threshold']['twin_steps'], 5000 + s['evals'])
                s['last_twin'][name] = tw
            r = eval_markers(ag, self.M['markers'], self.sched, rng, s['random_median'], tw)
            if 'D5' not in r:
                r['D5'] = {'raw': 'carried', 'stats': s['last_twin'].get(name)} if s['last_twin'].get(name) else {'raw': 'not_yet', 'stats': {}}
            res[name] = r
        if twin_due:
            s['next_twin'] = time.time() + self.sched['twin_test_every_s']
        s['evals'] += 1
        now = utcnow()
        need = self.sched['confirm_consecutive']
        summary = []
        for name, r in res.items():
            H = s['history'].setdefault(name, {})
            for mid, v in r.items():
                h = H.setdefault(mid, {'consec': 0, 'met_on': None, 'status': 'not_yet_tested', 'evals': 0})
                raw = v['raw']
                if raw == 'carried':
                    v['status'] = h['status']; continue
                h['evals'] += 1
                if raw == 'pass':
                    h['consec'] += 1
                    if h['consec'] >= need:
                        if not h['met_on']: h['met_on'] = now
                        h['status'] = 'met'
                    else:
                        h['status'] = 'passing_unconfirmed'
                elif raw == 'fail':
                    h['consec'] = 0
                    h['status'] = 'not_met' if not h['met_on'] else 'met_then_failed'
                elif raw == 'insufficient':
                    h['consec'] = 0
                    h['status'] = 'insufficient_data' if not h['met_on'] else 'met_then_failed'
                elif raw == 'not_yet':
                    h['status'] = 'not_yet_tested'
                else:
                    h['status'] = raw
                v['status'] = h['status']; v['met_on'] = h['met_on']; v['consec'] = h['consec']
                if name == 'main' or mid in ('D1', 'D2', 'D3', 'B1'):
                    pass
            summary.append(name + ': ' + ', '.join(f"{m} {r[m]['status'].replace('_', ' ')}" for m in
                           ('D1', 'D2', 'D3', 'D4', 'D5', 'I1', 'B1', 'B2', 'P0')))
        s['markers'] = {'evaluated': now, 'eval_number': s['evals'], 'step': s['main'].steps, 'results': res}
        with open(STATE / 'markers.jsonl', 'a') as f:
            f.write(json.dumps(s['markers'], default=float) + '\n')
        self.log('markers', f"Scheduled marker evaluation #{s['evals']} at step {s['main'].steps:,}. " + ' | '.join(summary),
                 eval_number=s['evals'])
        if res['control']['P0']['raw'] == 'pass' or res['main']['P0']['raw'] == 'pass':
            self.log('placebo_alarm', 'The shuffled-label placebo PASSED. Per the preregistration, every marker result from this evaluation is void until the pipeline is checked.')
        s['next_eval'] = time.time() + self.sched['marker_eval_every_s']

    def world_json(self):
        w = self.s['main'].world
        return {'n': N, 'day_len': DAY, 'map': [''.join('.#^SM'[int(x)] for x in row) for row in self.s['grid']],
                'day': w.t // DAY, 'phase': round((w.t % DAY) / DAY, 4), 'ambient': round(w.ambient(), 3),
                'food': sorted(list(p) for p in w.food), 'tool_cell': list(self.s['tool']),
                'tool_enabled': w.tool_enabled, 'tool_held': w.tool_held}

    def agent_json(self, ag, full):
        d = {'pos': list(ag.body.pos), 'h': [round(float(x), 4) for x in ag.body.h], 'has_tool': ag.body.has_tool,
             'steps': ag.steps, 'life': ag.n_lives, 'life_age': ag.life_age(), 'best_life': ag.best,
             'last_action': ACTIONS[ag.last_action], 'last_events': ag.last_events,
             'reward_ema': round(ag.ret_ema, 5), 'food_eaten': ag.food_eaten, 'counts': ag.rec.counts}
        if full:
            d['trajectory'] = list(ag.traj)
            d['lives'] = list(ag.lives)[-25:]
            if ag.world.tool_enabled or ag.rec.pe_tool:
                d['tool_uses'] = len(ag.rec.pe_tool)
        return d

    def write_status(self):
        s = self.s
        st = {'updated': utcnow(), 'updated_ts': time.time(), 'created': s['created'], 'restarts': s['restarts'],
              'runtime_s': round(self.runtime0 + time.time() - self.t_start),
              'tick_hz': self.sched['tick_hz'], 'prereg': self.prereg,
              'viable': VIABLE, 'setpoint': SETPOINT, 'variables': ['energy', 'temperature', 'integrity'],
              'world': self.world_json(), 'main': self.agent_json(s['main'], True),
              'control': self.agent_json(s['control'], False),
              'markers': s['markers'], 'next_eval': datetime.datetime.utcfromtimestamp(s['next_eval']).replace(microsecond=0).isoformat() + 'Z',
              'next_twin': datetime.datetime.utcfromtimestamp(s['next_twin']).replace(microsecond=0).isoformat() + 'Z',
              'random_median': s['random_median'], 'log': list(self.logtail)[-60:]}
        atomic_write(STATE / 'status.json', json.dumps(st, default=float, separators=(',', ':')).encode())
        self.last_status = time.time()

    def write_stats(self):
        m, c = self.s['main'], self.s['control']
        rec = {'ts': utcnow(), 'step': m.steps, 'life': m.n_lives, 'life_age': m.life_age(),
               'h': [round(float(x), 3) for x in m.body.h], 'reward_ema': round(m.ret_ema, 5),
               'food': m.food_eaten, 'counts': m.rec.counts, 'control_food': c.food_eaten}
        with open(STATE / 'stats.jsonl', 'a') as f:
            f.write(json.dumps(rec) + '\n')
        self.last_stats = time.time()

    def maybe_tool(self):
        at = self.M['markers'][3]['threshold']['tool_appears_at_step']
        for name in ('main', 'control'):
            ag = self.s[name]
            if not ag.world.tool_enabled and ag.steps >= at:
                ag.world.tool_enabled = True
                ag.world.set_tool_visible(True)
                if name == 'main':
                    self.log('tool', f"Step {ag.steps:,}: the reach tool appeared in the world, as pre-registered for step {at:,} (D4 clock starts).")

    def loop(self):
        dt = 1.0 / self.sched['tick_hz']
        nxt = time.time()
        s = self.s
        while self.running:
            rec = s['main'].step()
            s['control'].step()
            if rec and (rec['length'] >= DAY or rec['life'] % 100 == 0):
                self.log('life_end', f"Life {rec['life']} ended at age {rec['length']:,} steps: {rec['cause']}.")
            if s['main'].steps % 1000 == 0:
                self.maybe_tool()
            now = time.time()
            if now - self.last_status >= 5: self.write_status()
            if now - self.last_save >= 120: self.save()
            if now - self.last_stats >= 300: self.write_stats()
            if now >= s['next_eval']:
                self.save()   # so a crash during a long evaluation loses nothing
                t0 = time.time(); self.evaluate()
                self.log('eval_time', f"Evaluation took {time.time() - t0:.0f} s; the world was paused meanwhile.") if time.time() - t0 > 30 else None
                self.write_status(); nxt = time.time()
            nxt += dt
            sl = nxt - time.time()
            if sl > 0: time.sleep(sl)
            elif sl < -1: nxt = time.time()   # never try to catch up after a stall
        self.s['clean_stop'] = True
        self.save(); self.write_status()
        self.log('stop', f"Process stopped cleanly at main step {s['main'].steps:,}; state saved.")


def selftest():
    grid, tool = make_map()
    w = World(grid, tool, 1)
    ag = Agent('main', w, 3, 100)
    t0 = time.time()
    for _ in range(3000): ag.step()
    dtm = (time.time() - t0) / 3000
    print('map:'); print('\n'.join(''.join('.#^SM'[int(x)] for x in r) for r in grid)); print('tool', tool)
    print(f'{dtm * 1e6:.0f} us/step, lives {ag.n_lives}, counts {ag.rec.counts}, obs {OBS_DIM}')
    print('random median life', random_life_median(grid, tool, n=40))
    p = make_probes(grid, tool)
    print('twin', twin_test(ag, p, 500, 9))


if __name__ == '__main__':
    if '--selftest' in sys.argv:
        selftest(); sys.exit(0)
    run = Run()
    def _stop(*_): run.running = False
    signal.signal(signal.SIGTERM, _stop)
    signal.signal(signal.SIGINT, _stop)
    run.write_status()
    run.loop()