## learned_surfer.nim — LEARNED movement: state-conditional wave danger.
##
## The novelty over every other mover in this repo (`strafe`, `tfil`,
## `wave_surfer`): the danger of a guess-factor bin is **learned per movement
## state** instead of being a single hand-tuned/global quantity. The learner is
## the counted SBC with global fractional decay from `common_libs/bitbrain`
## (jobs j102/j103: measured to forget a changed mapping and to produce true
## probabilities), so it can track an opponent that adapts to us.
##
## ── The design (and its measured basis) ─────────────────────────────────────
## 1. Waves: enemy fire is detected from the one-tick energy drop (the drop IS
##    the firepower), exactly as `wave_surfer.nim`/`strafe.nim` do. `WorldState`
##    carries no bullet bodies, so the wave frame is "origin = enemy position at
##    the fire tick, centre line = bearing from the origin to us at that tick".
##
## 2. Label. A wave RESOLVES at the nominal arrival tick
##    `ceil(startDist/speed)` — the time the bullet needs for the range at fire
##    time, known at fire time — and the label is the 31-bin guess factor of our
##    angular offset from the centre line at that tick (`gfToBin`, the same
##    31-bin quantisation `wave_surfer.nim` uses).
##
## 3. State = ONE coarse wave-relative state at the fire tick, 4 fields x 4
##    symbols = 256 states, NEVER a temporal window (measured dead:
##    `docs/state_window_gate.md`):
##      vlat  lateral velocity in the wave frame, px/tick  (signed)
##      dist  range to the enemy at the fire tick, px
##      room  directional wall room along the direction we are running, px
##      turn  our own signed heading change, deg/tick
##    `lat` (the perpendicular offset from the centre line) is deliberately NOT
##    a field: at the fire tick the centre line passes through us, so it is
##    identically 0. 256 states over ~55k recorded shots is ~150 observations
##    per state — no recurrence problem.
##
## 4. Learner = counted SBC (`initCountedSbc`), read with `inferProb` (the
##    per-cell posterior), interpolated with the global 31-bin histogram that
##    `wave_surfer.nim` used: `p = (posterior + alpha*global)/(1 + alpha)`. With
##    no data at a cell this is exactly the old global surfer.
##
## 5. Danger of a candidate escape direction = the predicted probability of the
##    GF bin we would arrive in, SUMMED over every live wave, plus a wall
##    penalty, a travel penalty and a reversal penalty. The safest reachable
##    bin wins, then the same perpendicular steering / wall escape / radial
##    blend the other movers use.
##
## MEASURED (offline gate, `common_libs/tests/learned_surfer_gate.py`, 70
## recorded battles, held out BY BATTLE, 3 seeds): the state-conditional model
## beats the global histogram and chance on held-out log-loss (4.927 vs 4.974
## vs 4.954 bits) in 63/63 held-out battles (sign-flip p = 5e-5), but the effect
## is TINY: top-1 3.93% (global 3.96%, chance 3.23%). Read the ledger section
## "Learned movement (SBC)" in `docs/movement_campaign.md` before trusting it.
##
## env knobs (all read by `loadLearnedEnv`):
##   TR_MOVEMENT              = learned   (selects this engine)
##   TR_LEARNED_DECAY_EVERY   learns between decay passes (default 128)
##   TR_LEARNED_DECAY_SHIFT   `c -= c shr shift`; 0 disables forgetting
##   TR_LEARNED_ALPHA         prior mix weight (default 5.0)
##   TR_LEARNED_TRAVEL        danger cost of travelling across the wave (0.01)
##   TR_LEARNED_REVERSAL      danger cost of flipping the strafe side (0.02)
##   TR_LEARNED_PREF_DIST     preferred engagement distance (400 px)
##   TR_LEARNED_DIST_BAND     deadband around it (50 px)
##   TR_LEARNED_RADIAL_FRAC   radial blend outside the band (0.35)
##   TR_LEARNED_WALL_MARGIN   wall margin (48 px)
##   TR_LEARNED_GLOBAL        =1: ignore the state (ABLATION, the same mover
##                            with a pure global histogram)
##   TR_LEARNED_LABEL         `histogram` (default) or `outcome` (job j130).
##                            `outcome` replaces the label (the GF bin we
##                            cross) with the OUTCOME: a counted SBC over
##                            (state, candidate bin) estimating P(hit | state,
##                            g), trained on the dense label `hit AND |g-b| <=
##                            window` (b = the bin the wave resolved at, window
##                            = the bot radius as an angle at the wave's
##                            distance). See docs/movement_campaign.md,
##                            "outcome label".
##   TR_LEARNED_REAL_EVENTS   =1: resolve a wave on the REAL bullet event
##                            (onHitByBullet / an enemy bullet intercepted by
##                            ours) using the exact origin->endpoint line and
##                            the real flight time, dropping the wave at once.
##                            Default off (arrival-deadline proxy). Enemy WALL
##                            hits are owner-private on server 0.35.5 and stay
##                            deadline misses - see the const-block note.
##   TR_LEARNED_LOG           per-decision log line

import std/[math, os]
from std/strutils import parseFloat, parseInt, strip, toLowerAscii
import gun_harness/gun_interface
import movement_harness/movement_interface
import bitbrain/sbc

const
  LS_BINS   = 31           ## GF bins, identical to `wave_surfer.WS_BINS`
  LS_Q      = 4            ## symbols per field
  LS_NADE   = LS_Q * LS_Q  ## 16: (vlat, dist) on one axis, (room, turn) on the other
  DodgeTicks  = 15.0
  MaxBotSpeed = 8.0

## State-bin edges (4 symbols/field): the quantiles of the recorded corpus,
## derived by `common_libs/tests/learned_surfer_gate.py` (section E) and frozen
## here. Quantisation is the 25/50/75 % quantiles; a value is placed in the
## number of edges strictly below it.
const
  VlatEdges = [-6.736, 0.000, 6.753]
  DistEdges = [431.321, 487.612, 552.670]
  RoomEdges = [137.965, 206.589, 296.753]
  TurnEdges = [-0.142, 0.000, 0.105]

const
  LearnedDecayEveryEnv* = "TR_LEARNED_DECAY_EVERY"
  LearnedDecayShiftEnv* = "TR_LEARNED_DECAY_SHIFT"
  LearnedAlphaEnv*      = "TR_LEARNED_ALPHA"
  LearnedTravelEnv*     = "TR_LEARNED_TRAVEL"
  LearnedReversalEnv*   = "TR_LEARNED_REVERSAL"
  LearnedPrefDistEnv*   = "TR_LEARNED_PREF_DIST"
  LearnedDistBandEnv*   = "TR_LEARNED_DIST_BAND"
  LearnedRadialFracEnv* = "TR_LEARNED_RADIAL_FRAC"
  LearnedWallMarginEnv* = "TR_LEARNED_WALL_MARGIN"
  LearnedGlobalEnv*     = "TR_LEARNED_GLOBAL"
  LearnedLabelEnv*      = "TR_LEARNED_LABEL"
  LearnedRealEventsEnv* = "TR_LEARNED_REAL_EVENTS"
  LearnedLogEnv*        = "TR_LEARNED_LOG"

  ## Real-event matching (job j131).
  ## A wave is normally resolved on the nominal arrival tick
  ## `ceil(startDist/speed)`. With `TR_LEARNED_REAL_EVENTS=1` a wave is instead
  ## resolved by the REAL server event that carries the bullet's endpoint:
  ##   * `onHitByBullet`  -> the bullet HIT us; endpoint = our impact point;
  ##   * a bullet-vs-bullet intercept of an ENEMY bullet (our bullet hit theirs)
  ##     -> the enemy bullet's endpoint/heading are in the event.
  ## The bullet's raw straight line (origin at fire -> endpoint) then gives the
  ## EXACT GF bin and the real flight time (a cross-check on the energy-drop
  ## speed inference), and the wave is dropped immediately (no ghost build-up).
  ##
  ## ON THE RUNNING SERVER (0.35.5, verified from the server bytecode +
  ## `TurnToTickEventForBotMapper`) an ENEMY bullet that hits a WALL produces a
  ## `BulletHitWallEvent` only for the bullet's OWNER (`addPrivateBotEvent(
  ## bullet.botId, ...)`), and `bulletStates` is filtered to the bot's own
  ## bullets. So a wall HIT is NOT observable by the dodger; those waves fall
  ## back to the arrival deadline and are labelled a MISS. `resolveEnemyBullet`
  ## accepts a wall endpoint anyway so a future/other server can feed it.
  RealEventsMatchTol = 12.0  ## max |real flight - nominal| to accept a match, ticks
  RealEventsGrace     = 8     ## ticks past nominal before an unmatched wave resolves

  ## Number of joint (vlat,dist,room,turn) state codes = 4^4.
  LS_STATES = LS_Q * LS_Q * LS_Q * LS_Q
  ## Prior-mix weight for the 2-class outcome readout.
  OutcomePriorAlpha = 1.0

type
  LearnedLabelMode* = enum
    llHistogram   ## label = the GF bin the wave crossed us at (j128)
    llOutcome     ## label = the dense hit outcome (job j130)

var
  LearnedDecayEvery* = 128
  LearnedDecayShift* = 1
  LearnedAlpha*      = 5.0
  LearnedTravel*     = 0.01
  LearnedReversal*   = 0.02
  LearnedPrefDist*   = 400.0
  LearnedDistBand*   = 50.0
  LearnedRadialFrac* = 0.35
  LearnedWallMargin* = 48.0
  LearnedGlobal*     = false
  LearnedLabel*      = llHistogram
  LearnedRealEvents* = false
  LearnedLog*        = false

proc getEnvFloat(name: string, default: float): float =
  let s = getEnv(name, "")
  if s.len == 0: return default
  try: result = parseFloat(s.strip())
  except ValueError: result = default

proc getEnvInt(name: string, default: int): int =
  let s = getEnv(name, "")
  if s.len == 0: return default
  try: result = parseInt(s.strip())
  except ValueError: result = default

proc envOn(name: string, default = false): bool =
  let s = getEnv(name, "").strip()
  if s.len == 0: return default
  s notin ["0", "false", "no", "off"]

proc loadLearnedEnv*() =
  ## Read the knobs; callable again after `putEnv` so a gate can sweep arms in
  ## one process.
  LearnedDecayEvery = max(0, getEnvInt(LearnedDecayEveryEnv, 128))
  LearnedDecayShift = max(0, getEnvInt(LearnedDecayShiftEnv, 1))
  LearnedAlpha      = max(0.0, getEnvFloat(LearnedAlphaEnv, 5.0))
  LearnedTravel     = max(0.0, getEnvFloat(LearnedTravelEnv, 0.01))
  LearnedReversal   = max(0.0, getEnvFloat(LearnedReversalEnv, 0.02))
  LearnedPrefDist   = max(1.0, getEnvFloat(LearnedPrefDistEnv, 400.0))
  LearnedDistBand   = max(0.0, getEnvFloat(LearnedDistBandEnv, 50.0))
  LearnedRadialFrac = clamp(getEnvFloat(LearnedRadialFracEnv, 0.35), 0.0, 1.0)
  LearnedWallMargin = max(0.0, getEnvFloat(LearnedWallMarginEnv, 48.0))
  LearnedGlobal     = envOn(LearnedGlobalEnv)
  LearnedRealEvents = envOn(LearnedRealEventsEnv)
  LearnedLog        = envOn(LearnedLogEnv)
  LearnedLabel      =
    case getEnv(LearnedLabelEnv, "").strip().toLowerAscii()
    of "outcome": llOutcome
    else: llHistogram

loadLearnedEnv()

# ── small helpers ───────────────────────────────────────────────────────────

proc wrapPi(x: float64): float64 {.inline.} =
  result = x
  while result >  PI: result -= 2.0*PI
  while result < -PI: result += 2.0*PI

proc wrap180(d: float): float {.inline.} =
  result = d
  while result >  180.0: result -= 360.0
  while result < -180.0: result += 360.0

proc mea(speed: float64): float64 {.inline.} =
  if speed <= 1e-9: return 0.0
  arcsin(min(MaxBotSpeed / speed, 1.0))

proc gfToBin(gf: float64): int {.inline.} =
  clamp(int(round((gf.clamp(-1.0, 1.0) + 1.0) * 0.5 * float64(LS_BINS - 1))),
        0, LS_BINS - 1)

proc binToGF(idx: int): float64 {.inline.} =
  float64(idx) / float64(LS_BINS - 1) * 2.0 - 1.0

proc code(value: float64, edges: array[3, float64]): int {.inline.} =
  result = 0
  for e in edges:
    if value > e: inc result

proc roomToWall(px, py, dx, dy, arenaW, arenaH: float64): float64 =
  ## Distance from (px,py) along unit (dx,dy) until leaving the arena, keeping
  ## the 18 px bot radius. Mirrors the offline gate.
  var t = Inf
  for i in 0..1:
    let p  = if i == 0: px else: py
    let d  = if i == 0: dx else: dy
    let lo = BotRadius
    let hi = (if i == 0: arenaW else: arenaH) - BotRadius
    if abs(d) > 1e-9:
      let cand = if d > 0.0: (hi - p) / d else: (lo - p) / d
      if cand < t: t = cand
  if t == Inf: return 0.0
  max(0.0, t)

# ── module types ────────────────────────────────────────────────────────────

type
  LSWave = object
    ownerId:          int       ## enemy that fired (energies are per-enemy)
    fireTick:         int       ## `ws.tick` at the fire tick (real flight time)
    originX, originY: float64
    bearing:          float64   ## enemy -> us at the fire tick (centre line)
    speed:            float64
    startDist:        float64
    power:            float64
    ticksLeft:        int       ## ticks until the nominal arrival
    fresh:            bool      ## created this tick: do not age it yet
    selfEnergyAtFire: float64   ## our energy at the fire tick (hit detection)
    stateRow:         int       ## (vlat, dist) code
    stateCol:         int       ## (room, turn) code

  LearnedSurferModule* = object
    sbc*:            Sbc
    outcome*:        Sbc         ## llOutcome: (state, candidate bin) -> hit/miss
    glob*:           array[LS_BINS, int]
    glc:             int        ## learns since the last global-histogram decay
    hitGlobal:       int        ## llOutcome: hits seen (for the prior)
    missGlobal:      int        ## llOutcome: non-hits seen (for the prior)
    scores:          seq[float]
    waves:           seq[LSWave]
    prevEnergy:      seq[tuple[id: int, energy: float]]
    strafeDir:       float64
    dir:             float64    ## direction commanded last tick (+-1)
    prevX, prevY:    float64    ## our position one tick ago
    prevHeading:     float64
    debugGraphics*:  bool
    decisions*:      int        ## decisions taken (diagnostic)
    resolvedReal*:   int        ## waves resolved by a real bullet event
    resolvedDead*:   int        ## waves resolved on the arrival deadline
    lastFlightErr*:  float      ## real flight - nominal flight, last resolution

proc resetRound*(m: var LearnedSurferModule) =
  ## Per-ROUND reset: the waves and the smoothed global prior are per round, but
  ## the counted SBC memory is deliberately NOT wiped — the whole point of the
  ## counted+decay mode is that the memory is bounded and decays on its own, so
  ## it survives a round boundary without becoming a battle-long static average
  ## (the j115 defect). Use `resetBattle` for a hard wipe.
  m.waves      = @[]
  m.prevEnergy = @[]
  m.strafeDir  = 1.0
  m.dir        = 1.0
  m.prevX      = 0.0
  m.prevY      = 0.0
  m.prevHeading = 0.0
  m.decisions  = 0
  m.resolvedReal = 0
  m.resolvedDead = 0
  m.lastFlightErr = 0.0
  for i in 0..<LS_BINS: m.glob[i] = 0
  m.glc = 0
  m.hitGlobal = 0
  m.missGlobal = 0

proc initLearnedSurfer*(): LearnedSurferModule =
  result.debugGraphics = false
  result.sbc = initCountedSbc(LS_NADE, LS_BINS,
                              max(0, LearnedDecayEvery),
                              max(0, LearnedDecayShift))
  result.sbc.decayEvery = max(0, LearnedDecayEvery)
  result.sbc.decayShift = max(0, LearnedDecayShift)
  # The outcome model learns LS_BINS samples per resolved wave, so its decay
  # period is scaled to the same WAVE cadence as the histogram's (decayEvery
  # waves), keeping the two labels' forgetting comparable.
  block:
    let oe = if LearnedDecayEvery > 0: LearnedDecayEvery * LS_BINS else: 0
    result.outcome = initCountedSbc(LS_STATES, 2, oe, max(0, LearnedDecayShift))
    result.outcome.decayEvery = oe
    result.outcome.decayShift = max(0, LearnedDecayShift)
  result.scores = newSeq[float](LS_BINS)
  result.resetRound()

proc resetBattle*(m: var LearnedSurferModule) =
  ## Hard wipe of the learned memory (round the SBC's counters and `clear`).
  m.sbc.clear()
  m.outcome.clear()
  m.resetRound()

proc clearGraphics*(m: var LearnedSurferModule) {.inline.} = discard
proc removeBulletNear*(m: var LearnedSurferModule, x, y: float) {.inline.} = discard
proc liveWaves*(m: LearnedSurferModule): int {.inline.} = m.waves.len

# ── the learner ─────────────────────────────────────────────────────────────

proc prior(m: LearnedSurferModule, k: int): float =
  var tot = 0
  for g in m.glob: tot += g
  (float(m.glob[k]) + 1.0) / (float(tot) + float(LS_BINS))

proc learnWave(m: var LearnedSurferModule, w: LSWave, bin: int, hit: bool) =
  ## One resolved wave -> training samples.
  ##   llHistogram: one sample, (state) -> (GF bin at arrival).
  ##   llOutcome:   LS_BINS samples, (state, candidate g) -> hit, where the label
  ##                is `hit and |g - bin| <= window` (the dense counterfactual:
  ##                would this wave have hit us at candidate g?). `window` is
  ##                the bot radius as an angle at the wave's distance, in GF
  ##                bins — the angular/physical tolerance of the body. The SBC
  ##                applies its own global decay; the hit/miss counters below
  ##                are aged on the same schedule.
  case LearnedLabel
  of llHistogram:
    discard m.sbc.learn([w.stateRow.int32], [w.stateCol.int32], bin)
    if m.glob[bin] < 255: inc m.glob[bin]
    inc m.glc
    if LearnedDecayEvery > 0 and LearnedDecayShift > 0 and
       m.glc >= LearnedDecayEvery:
      for k in 0..<LS_BINS:
        m.glob[k] = m.glob[k] - (m.glob[k] shr LearnedDecayShift)
      m.glc = 0
  of llOutcome:
    let maxA = mea(w.speed)
    let wbin =
      if maxA > 1e-9:
        (arcsin(min(BotRadius / max(w.startDist, 1.0), 1.0)) / maxA) *
          float(LS_BINS - 1) / 2.0
      else: 0.0
    let sc = if LearnedGlobal: 0 else: w.stateRow * LS_Q * LS_Q + w.stateCol
    for g in 0..<LS_BINS:
      let lab = if hit and abs(float(g - bin)) <= wbin: 1 else: 0
      discard m.outcome.learn([sc.int32], [g.int32], lab)
      inc m.glc
    if hit: inc m.hitGlobal else: inc m.missGlobal
    if LearnedDecayEvery > 0 and LearnedDecayShift > 0 and
       m.glc >= LearnedDecayEvery * LS_BINS:
      m.hitGlobal = m.hitGlobal - (m.hitGlobal shr LearnedDecayShift)
      m.missGlobal = m.missGlobal - (m.missGlobal shr LearnedDecayShift)
      m.glc = 0

proc resolveWaveIdx(m: var LearnedSurferModule, idx, bin: int, hit: bool,
                    currentTick: int) =
  ## One live wave -> one training sample + removal. Shared by the arrival
  ## deadline (unobserved wall misses) and the real-event path, so a wave is
  ## ALWAYS trained and dropped exactly once - no ghost accumulation.
  let w = m.waves[idx]
  let nom = w.startDist / max(w.speed, 1e-9)
  let realFlight = float(currentTick - w.fireTick)
  m.lastFlightErr = realFlight - nom
  if LearnedLog:
    echo "[learned] resolve bin=", bin, " state=", w.stateRow, "/",
         w.stateCol, " d=", w.startDist.int, " e=", w.originX.int, ",",
         w.originY.int, " flight=", realFlight.int, " nominal=", nom.int,
         " hit=", hit
  m.learnWave(w, bin, hit)
  m.waves.del(idx)

proc missileLineBin(w: LSWave, x, y, headingRad: float): int =
  ## GF bin of the bullet's real straight line through `(x,y)` (the endpoint),
  ## falling back to the real heading when the endpoint is degenerate. This is
  ## the EXACT geometry: origin at the fire tick + real endpoint, no timing
  ## guess.
  let maxA = mea(w.speed)
  if maxA < 1e-9: return gfToBin(0.0)
  let ex = x - w.originX
  let ey = y - w.originY
  let lineDir =
    if hypot(ex, ey) > 1.0: arctan2(ey, ex)
    else: headingRad
  gfToBin(clamp(wrapPi(lineDir - w.bearing) / maxA, -1.0, 1.0))

proc resolveEnemyBullet*(m: var LearnedSurferModule, x, y, headingRad: float,
                         ownerId, currentTick: int, hit: bool): bool =
  ## Resolve (and DROP) the live wave matching a REAL enemy-bullet event.
  ##
  ## `x,y`        the bullet's real endpoint (our impact point for a HIT, the
  ##              wall point for a wall hit, the intercept point for a
  ##              bullet-vs-bullet hit),
  ## `headingRad` the bullet's real heading (fallback when the endpoint is
  ##              degenerate),
  ## `hit`        true only for a HIT on us.
  ## The exact straight line origin->endpoint sets the label's GF bin and the
  ## real flight time `currentTick - fireTick` is recorded, which cross-checks
  ## the energy-drop speed inference. No-op unless `TR_LEARNED_REAL_EVENTS=1`.
  result = false
  if not LearnedRealEvents or m.waves.len == 0: return
  # The wave whose nominal arrival is closest to now is the one this bullet
  # belongs to; ownerId disambiguates when several enemies are firing.
  var best = -1
  var bestKey = Inf
  for i in 0..<m.waves.len:
    let w = m.waves[i]
    if ownerId >= 0 and w.ownerId != ownerId: continue
    let key = abs(float(currentTick - w.fireTick) -
                  w.startDist / max(w.speed, 1e-9))
    if key < bestKey:
      bestKey = key
      best = i
  if best < 0:  # no wave from that enemy: fall back to time-only matching
    for i in 0..<m.waves.len:
      let w = m.waves[i]
      let key = abs(float(currentTick - w.fireTick) -
                    w.startDist / max(w.speed, 1e-9))
      if key < bestKey:
        bestKey = key
        best = i
  if best < 0 or bestKey > RealEventsMatchTol: return
  let w = m.waves[best]
  let bin = missileLineBin(w, x, y, headingRad)
  m.resolveWaveIdx(best, bin, hit, currentTick)
  inc m.resolvedReal
  result = true

proc predictHit*(m: LearnedSurferModule, row, col, g: int): float =
  ## P(hit | state, candidate bin g) — the `outcome` danger (lower = safer),
  ## from the 2-class counted SBC read out with the per-cell posterior and
  ## blended with the global hit rate.
  let t = m.hitGlobal + m.missGlobal
  let prior = if t == 0: 0.5 else: float(m.hitGlobal) / float(t)
  let sc = if LearnedGlobal: 0 else: row * LS_Q * LS_Q + col
  let c0 = m.outcome.countAt(sc, g, 0)
  let c1 = m.outcome.countAt(sc, g, 1)
  let n = c0 + c1
  if n == 0: return prior
  (float(c1) + OutcomePriorAlpha * prior) / (float(n) + OutcomePriorAlpha)

proc predictState*(m: var LearnedSurferModule, row, col: int,
                  p: var array[LS_BINS, float]) =
  ## p(bin | state) — the counted SBC's per-cell posterior interpolated with the
  ## global histogram. With no data at the cell this is exactly the old surfer.
  ## `p` is filled with the prior FIRST, so the fallback costs nothing.
  var tot = 0
  for g in m.glob: tot += g
  let den = float(tot) + float(LS_BINS)
  for k in 0..<LS_BINS: p[k] = (float(m.glob[k]) + 1.0) / den
  if LearnedGlobal: return
  for k in 0..<LS_BINS: m.scores[k] = 0.0
  m.sbc.inferProb([row.int32], [col.int32], m.scores)
  var st = 0.0
  for k in 0..<LS_BINS: st += m.scores[k]
  if st <= 0.0: return
  for k in 0..<LS_BINS:
    p[k] = (m.scores[k] + LearnedAlpha * p[k]) / (1.0 + LearnedAlpha)

# ── fire detection + state ──────────────────────────────────────────────────

proc prevEnergyGet(m: LearnedSurferModule, id: int): float =
  for e in m.prevEnergy:
    if e.id == id: return e.energy
  100.0

proc prevEnergySet(m: var LearnedSurferModule, id: int, energy: float) =
  for i in 0..<m.prevEnergy.len:
    if m.prevEnergy[i].id == id:
      m.prevEnergy[i].energy = energy
      return
  m.prevEnergy.add((id: id, energy: energy))

proc detectFire(m: var LearnedSurferModule, id: int, ex, ey, eenergy: float,
                ws: WorldState) =
  ## One enemy's energy sample. A plausible one-tick firepower drop IS a wave;
  ## the state is the wave-relative state at THIS tick (the fire tick).
  let prev = m.prevEnergyGet(id)
  let drop = prev - eenergy
  m.prevEnergySet(id, eenergy)
  if drop < 0.09 or drop > 3.01: return

  let botX = ws.selfX
  let botY = ws.selfY
  let bspeed = 20.0 - 3.0 * drop
  let d      = hypot(botX - ex, botY - ey)
  let bearing = arctan2(botY - ey, botX - ex)      # centre line
  let ux = cos(bearing)
  let uy = sin(bearing)

  # vlat: lateral velocity in the wave frame. The centre line passes through us
  # at this tick, so lat(now) == 0 and vlat == -lat(prev).
  let dxp = m.prevX - ex
  let dyp = m.prevY - ey
  var vlat = -(dxp * (-uy) + dyp * ux)
  if m.prevX == 0.0 and m.prevY == 0.0 and m.prevHeading == 0.0: vlat = 0.0
  let roomDx = if vlat >= 0.0: -uy else: uy
  let roomDy = if vlat >= 0.0:  ux else: -ux
  let room   = roomToWall(botX, botY, roomDx, roomDy, ws.arenaWidth, ws.arenaHeight)
  let turn   = wrap180(float(ws.selfHeading) - float(m.prevHeading))

  let row = code(vlat, VlatEdges) * LS_Q + code(d, DistEdges)
  let col = code(room, RoomEdges) * LS_Q + code(turn, TurnEdges)

  m.waves.add LSWave(
    ownerId: id, fireTick: ws.tick,
    originX: ex, originY: ey, bearing: bearing, speed: bspeed,
    startDist: d, power: drop,
    ticksLeft: max(1, int(ceil(d / max(bspeed, 1e-9)))),
    fresh: true,
    selfEnergyAtFire: ws.selfEnergy,
    stateRow: row, stateCol: col,
  )

# ── the mover ───────────────────────────────────────────────────────────────

proc computeMove*(m: var LearnedSurferModule, ws: WorldState): MoveCommand =
  let botX = ws.selfX
  let botY = ws.selfY

  # ── 1. fire detection (every alive enemy, per-enemy energy) ───────────────
  var seen = 0
  for ei in ws.enemies:
    inc seen
    m.detectFire(ei.id, ei.x, ei.y, ei.energy, ws)
  if seen == 0 and (ws.enemyX != 0.0 or ws.enemyY != 0.0):
    m.detectFire(-1, ws.enemyX, ws.enemyY, ws.enemyEnergy, ws)

  # ── 2. advance + resolve waves; every resolution is a training sample ────
  var i = 0
  while i < m.waves.len:
    if m.waves[i].fresh:
      m.waves[i].fresh = false            # created this tick: not one tick old yet
    else:
      dec m.waves[i].ticksLeft
    # With real events on, wait `RealEventsGrace` ticks past the nominal arrival
    # so a late HitByBullet can still claim the wave; a wave no event claims is
    # a WALL MISS (the wall event is owner-private - see the const block).
    let deadline = if LearnedRealEvents: -RealEventsGrace else: 0
    if m.waves[i].ticksLeft <= deadline:
      let w = m.waves[i]
      let maxA = mea(w.speed)
      if maxA >= 1e-9:
        let off = wrapPi(arctan2(botY - w.originY, botX - w.originX) - w.bearing)
        let bin = gfToBin(clamp(off / maxA, -1.0, 1.0))
        # llOutcome label: did THIS wave hit us? The energy drop is the proxy;
        # with real events on the HIT is taken from `onHitByBullet` instead, so
        # an unmatched wave is a wall MISS.
        let hit = (not LearnedRealEvents) and
                  ws.selfEnergy < w.selfEnergyAtFire - 0.01
        m.resolveWaveIdx(i, bin, hit, ws.tick)
        inc m.resolvedDead
      else:
        m.waves.del(i)
      continue
    inc i

  # ── 3. danger of every candidate bin, summed over every live wave ────────
  #   llHistogram: precompute the predicted arrival-bin distribution per wave.
  #   llOutcome:   danger(wave, g) = P(hit | wave.state, g), computed on demand.
  var ps: seq[array[LS_BINS, float]]
  if LearnedLabel == llHistogram:
    ps.setLen(m.waves.len)
    for j in 0..<m.waves.len:
      m.predictState(m.waves[j].stateRow, m.waves[j].stateCol, ps[j])

  # nearest wave = largest radius/distance ratio (the one about to arrive)
  var nearest = -1
  var bestRatio = -1.0
  for j in 0..<m.waves.len:
    let dd = max(1e-6, hypot(botX - m.waves[j].originX, botY - m.waves[j].originY))
    let done = 1.0 - float(m.waves[j].ticksLeft) /
               float(max(1, int(ceil(m.waves[j].startDist /
                                      max(m.waves[j].speed, 1e-9)))))
    let ratio = done / dd
    if ratio > bestRatio:
      bestRatio = ratio
      nearest = j

  var perpAngle = 0.0

  if nearest >= 0:
    let w = m.waves[nearest]
    let dx = botX - w.originX
    let dy = botY - w.originY
    let d  = max(1.0, hypot(dx, dy))
    let toBot = arctan2(dy, dx)
    let maxA  = mea(w.speed)
    let curGF = if maxA >= 1e-9:
                  clamp(wrapPi(toBot - w.bearing) / maxA, -1.0, 1.0)
                else: 0.0
    let dodgeDist = max(MaxBotSpeed, ws.selfSpeed) * DodgeTicks

    var bestBin    = gfToBin(curGF)
    var bestDanger = Inf
    var bestDirs: array[LS_BINS, float]

    for j in 0..<LS_BINS:
      let gfJ    = binToGF(j)
      let angleJ = w.bearing + gfJ * maxA
      let px = w.originX + cos(angleJ) * d
      let py = w.originY + sin(angleJ) * d
      var ux = px - botX
      var uy = py - botY
      let ul = hypot(ux, uy)
      if ul < 1e-6:
        ux = cos(perpAngle); uy = sin(perpAngle)
      else:
        ux /= ul; uy /= ul
      let futureX = botX + ux * dodgeDist
      let futureY = botY + uy * dodgeDist
      bestDirs[j] = if gfJ >= curGF: 1.0 else: -1.0

      var danger =
        if LearnedLabel == llHistogram: ps[nearest][j]
        else: m.predictHit(w.stateRow, w.stateCol, j)
      # every OTHER live wave contributes its own predicted mass at the bin the
      # candidate direction would land in for THAT wave.
      for k in 0..<m.waves.len:
        if k == nearest: continue
        let w2 = m.waves[k]
        let d2 = max(1.0, hypot(futureX - w2.originX, futureY - w2.originY))
        let mA2 = mea(w2.speed)
        if mA2 < 1e-9: continue
        let off2 = wrapPi(arctan2(futureY - w2.originY, futureX - w2.originX) -
                          w2.bearing)
        let b2 = gfToBin(clamp(off2 / mA2, -1.0, 1.0))
        if LearnedLabel == llHistogram: danger += ps[k][b2]
        else: danger += m.predictHit(w2.stateRow, w2.stateCol, b2)

      let wallHit = futureX < LearnedWallMargin or
                    futureX > ws.arenaWidth  - LearnedWallMargin or
                    futureY < LearnedWallMargin or
                    futureY > ws.arenaHeight - LearnedWallMargin
      if wallHit: danger *= 5.0
      danger += LearnedTravel * abs(gfJ - curGF)
      if bestDirs[j] != m.dir: danger += LearnedReversal

      if danger < bestDanger:
        bestDanger = danger
        bestBin    = j

    let bestGF = binToGF(bestBin)
    m.strafeDir = bestDirs[bestBin]
    m.dir       = m.strafeDir
    inc m.decisions
    if LearnedLog:
      var pAvg, pSafe, pCur: float
      if LearnedLabel == llHistogram:
        pAvg = 0.0
        for k in 0..<LS_BINS: pAvg += ps[nearest][k] / float(LS_BINS)
        pSafe = ps[nearest][bestBin]
        pCur  = ps[nearest][gfToBin(curGF)]
      else:
        pAvg = 0.0
        for k in 0..<LS_BINS:
          pAvg += m.predictHit(w.stateRow, w.stateCol, k) / float(LS_BINS)
        pSafe = m.predictHit(w.stateRow, w.stateCol, bestBin)
        pCur  = m.predictHit(w.stateRow, w.stateCol, gfToBin(curGF))
      echo "[learned] wave d=", d.int, " curGF=", (curGF * 100.0).int,
           " bestGF=", (bestGF * 100.0).int, " dir=", m.strafeDir.int,
           " pSafe=", (pSafe * 1000.0).int,
           " pCur=", (pCur * 1000.0).int,
           " pAvg=", (pAvg * 1000.0).int,
           " state=", m.waves[nearest].stateRow, "/", m.waves[nearest].stateCol
    # perpendicular in the wave's own frame: +90 deg from origin->bot increases
    # the GF (CCW), -90 decreases it.
    perpAngle = if m.strafeDir >= 0.0: toBot + PI * 0.5
                else:                   toBot - PI * 0.5
  elif ws.enemyX != 0.0 or ws.enemyY != 0.0:
    let toBot = arctan2(botY - ws.enemyY, botX - ws.enemyX)
    perpAngle = if m.strafeDir >= 0.0: toBot + PI * 0.5
                else:                   toBot - PI * 0.5
  else:
    perpAngle = degToRad(ws.selfHeading)

  # ── 4. never drive into a wall ───────────────────────────────────────────
  let nearLeft   = botX < LearnedWallMargin
  let nearRight  = botX > ws.arenaWidth  - LearnedWallMargin
  let nearBottom = botY < LearnedWallMargin
  let nearTop    = botY > ws.arenaHeight - LearnedWallMargin
  if nearLeft or nearRight or nearBottom or nearTop:
    let px = cos(perpAngle)
    let py = sin(perpAngle)
    if (nearLeft   and px < 0.0) or (nearRight  and px > 0.0) or
       (nearBottom and py < 0.0) or (nearTop    and py > 0.0):
      m.strafeDir = -m.strafeDir
      m.dir       = m.strafeDir
      perpAngle   = perpAngle + PI
    let escapeAngle = arctan2(ws.arenaHeight * 0.5 - botY,
                              ws.arenaWidth * 0.5 - botX)
    let ex = cos(escapeAngle) + cos(perpAngle)
    let ey = sin(escapeAngle) + sin(perpAngle)
    perpAngle = arctan2(ey, ex)

  # ── 5. distance control outside the deadband ────────────────────────────
  let enemyDist = hypot(ws.enemyX - botX, ws.enemyY - botY)
  let distErr   = enemyDist - LearnedPrefDist
  let radialFrac =
    if   distErr >  LearnedDistBand:  LearnedRadialFrac     # too far -> approach
    elif distErr < -LearnedDistBand: -LearnedRadialFrac     # too close -> retreat
    else: 0.0
  if abs(radialFrac) > 1e-9:
    let radialAngle = arctan2(ws.enemyY - botY, ws.enemyX - botX) +
                      (if radialFrac < 0.0: PI else: 0.0)
    let rx = cos(perpAngle) * (1.0 - abs(radialFrac)) +
             cos(radialAngle) * abs(radialFrac)
    let ry = sin(perpAngle) * (1.0 - abs(radialFrac)) +
             sin(radialAngle) * abs(radialFrac)
    perpAngle = arctan2(ry, rx)

  # ── 6. steer the body, full speed ───────────────────────────────────────
  let desiredDeg = radToDeg(perpAngle)
  var delta = desiredDeg - ws.selfHeading
  while delta >  180.0: delta -= 360.0
  while delta < -180.0: delta += 360.0
  let goForward = abs(delta) <= 90.0
  if not goForward:
    delta = if delta >= 0.0: delta - 180.0 else: delta + 180.0

  m.prevX = botX
  m.prevY = botY
  m.prevHeading = float(ws.selfHeading)

  (speed: (if goForward: 8.0 else: -8.0),
   turnRate: delta.clamp(-10.0, 10.0))