## Offline gun range: replay recorded or synthetic `WorldState` streams through
## the EXISTING `VirtualTracker`, exactly as ModularBot's live loop drives it —
## no Java server, no radar, no movement module and no real firing.
##
## Why this is definitionally the same metric as the online one: virtual-bullet
## fitness is already a pure function of (a stream of `WorldState`, a list of
## guns). The only thing the Java battle supplies is where the states come from.
## Guns keep their own internal history, so a `seq[WorldState]` replayed IN ORDER
## is a complete movement history.
##
## Fixture format (JSONL):
##   optional first line:
##     {"meta":{"adversary":"<name>","source":"synthetic|classic-robocode|live",
##              "perfect_info":false,"arena":{"w":800,"h":600},"note":"..."}}
##   then one line per tick, Tank Royale convention (0° = East, CCW +, degrees):
##     {"tick":<int>,"ex":<f>,"ey":<f>,"eh":<f>,"es":<f>,"ee":<f>,
##      "sx":<f>,"sy":<f>,"sh":<f>,"ss":<f>,"se":<f>}
##   `e*` is the enemy to predict, `s*` is the shooter.
##   Two optional per-tick extensions are written by the ModularBot recorder and
##   tolerated (defaulted) by every other producer:
##     "lst":<int>  enemyTracker.lastSeenTick at this tick (staleness fidelity)
##     "eid":<int>  enemy bot id (fitness bucket; defaults to 1)
##   A trailing live end marker is also optional:
##     {"end":{"enemy_died":<bool>,"ticks":<int>}}
##   It lets the replay reproduce the live resolver's final-tick behaviour.
##
## The replay never calls the gun selector, so it is RNG-free for every
## deterministic gun. Tsetlin is stochastic and is expected to differ.

import std/[json, os, strformat, math, tables, strutils, random]
import gun_interface
import virtual_bullets

export gun_interface, virtual_bullets

const
  DefaultArenaW* = 800.0
  DefaultArenaH* = 600.0
  DefaultEnemyId* = 1

type
  FixtureMeta* = object
    adversary*:   string
    source*:      string   ## "synthetic" | "classic-robocode" | "live"
    perfectInfo*: bool
    arenaW*, arenaH*: float
    note*:        string

  Fixture* = object
    meta*:    FixtureMeta
    states*:  seq[WorldState]
    lastSeen*: seq[int]    ## parallel to states; -1 = unknown (use state.tick)
    enemyId*: int
    enemyDied*: bool       ## target was dead at round end (recorder end marker)

  BinStat* = object
    shots*, hits*: int

  GunReport* = object
    ## Per-gun fitness over the whole fixture, computed with the SAME window
    ## semantics as ModularBot.onRoundEnded (min(count, WindowSize), iterate the
    ## ring slots 0..<n).
    name*:  string
    shots*, hits*: int
    bins*:  array[len(PowerBins), BinStat]

  GunDriver* = object
    ## Type-erased handle to one gun. `predictCb`/`resultCb` forward to the
    ## concrete gun inside a heap box so heterogeneous guns can live in a seq.
    name*:     string
    predictCb*: proc(state: WorldState, bulletSpeed: float): GunPrediction {.closure.}
    resultCb*:  proc(e: FeedbackEvent) {.closure.}
    readyCb*:   proc(): bool {.closure.}  ## nil => always ready (Tsetlin gate)

  GunBox[G] = ref object
    g: G

proc makeDriver*[G](name: string, gun: G): GunDriver =
  ## Wrap a concrete gun value in a type-erased driver.
  let box = GunBox[G](g: gun)
  result.name = name
  result.predictCb = proc(state: WorldState, bulletSpeed: float): GunPrediction =
    box.g.predict(state, bulletSpeed)
  result.resultCb = proc(e: FeedbackEvent) =
    box.g.onResult(e)
  when compiles(box.g.isWarmedUp()):
    result.readyCb = proc(): bool = box.g.isWarmedUp()
  else:
    result.readyCb = nil

# ── fixture I/O ───────────────────────────────────────────────────────────────

proc initMeta*(source = "synthetic", arenaW = DefaultArenaW,
               arenaH = DefaultArenaH): FixtureMeta =
  FixtureMeta(source: source, arenaW: arenaW, arenaH: arenaH)

proc initFixture*(states: seq[WorldState], meta: FixtureMeta,
                  lastSeen: seq[int] = @[], enemyId = DefaultEnemyId,
                  enemyDied = false): Fixture =
  Fixture(meta: meta, states: states, lastSeen: lastSeen, enemyId: enemyId,
          enemyDied: enemyDied)

proc stateToJson(ws: WorldState, lst: int, eid: int): JsonNode =
  result = %*{
    "tick": ws.tick,
    "ex": ws.enemyX, "ey": ws.enemyY,
    "eh": ws.enemyHeading, "es": ws.enemySpeed, "ee": ws.enemyEnergy,
    "sx": ws.selfX, "sy": ws.selfY,
    "sh": ws.selfHeading, "ss": ws.selfSpeed, "se": ws.selfEnergy,
  }
  if lst >= 0: result["lst"] = %lst
  if eid >= 0: result["eid"] = %eid

proc metaToJson(m: FixtureMeta): JsonNode =
  %*{"meta": {
    "adversary": m.adversary,
    "source": m.source,
    "perfect_info": m.perfectInfo,
    "arena": {"w": m.arenaW, "h": m.arenaH},
    "note": m.note,
  }}

proc stateFromJson(node: JsonNode, arenaW, arenaH: float,
                   enemyId: int): WorldState =
  let ex = node["ex"].getFloat()
  let ey = node["ey"].getFloat()
  result = WorldState(
    enemyX: ex, enemyY: ey,
    enemyHeading: node["eh"].getFloat(),
    enemySpeed: node["es"].getFloat(),
    enemyEnergy: node["ee"].getFloat(),
    selfX: node["sx"].getFloat(),
    selfY: node["sy"].getFloat(),
    selfHeading: node["sh"].getFloat(),
    selfRadarHeading: node["sh"].getFloat(),
    selfSpeed: node["ss"].getFloat(),
    selfEnergy: node["se"].getFloat(),
    arenaWidth: arenaW,
    arenaHeight: arenaH,
    tick: node["tick"].getInt(),
    enemies: @[EnemyInfo(id: enemyId, x: ex, y: ey,
                         heading: node["eh"].getFloat(),
                         speed: node["es"].getFloat(),
                         energy: node["ee"].getFloat())],
  )

proc parseMeta(node: JsonNode): FixtureMeta =
  result = initMeta()
  if node.hasKey("adversary"): result.adversary = node["adversary"].getStr()
  if node.hasKey("source"): result.source = node["source"].getStr()
  if node.hasKey("perfect_info"): result.perfectInfo = node["perfect_info"].getBool()
  if node.hasKey("note"): result.note = node["note"].getStr()
  if node.hasKey("arena"):
    let a = node["arena"]
    if a.hasKey("w"): result.arenaW = a["w"].getFloat()
    if a.hasKey("h"): result.arenaH = a["h"].getFloat()

proc saveFixture*(path: string, fx: Fixture) =
  ## Write a fixture as JSONL. Truncates the file. A trailing `end` marker
  ## records whether the target died at round end (live-resolver semantics).
  let f = open(path, fmWrite)
  defer: f.close()
  f.writeLine($metaToJson(fx.meta))
  for i, ws in fx.states:
    let lst = if i < fx.lastSeen.len: fx.lastSeen[i] else: -1
    f.writeLine($stateToJson(ws, lst, fx.enemyId))
  f.writeLine($(%*{"end": {"enemy_died": fx.enemyDied, "ticks": fx.states.len}}))

proc saveFixture*(path: string, states: seq[WorldState], meta: FixtureMeta,
                  lastSeen: seq[int] = @[], enemyId = DefaultEnemyId,
                  enemyDied = false) =
  saveFixture(path, initFixture(states, meta, lastSeen, enemyId, enemyDied))

proc loadFixture*(path: string): Fixture =
  ## Parse a JSONL fixture. The meta line and the optional lst/eid fields are
  ## optional; arena defaults to 800x600, enemy id to 1.
  if not fileExists(path):
    raise newException(IOError, "fixture not found: " & path)
  result = initFixture(@[], initMeta())
  var arenaW = DefaultArenaW
  var arenaH = DefaultArenaH
  var enemyId = DefaultEnemyId
  for rawLine in lines(path):
    let line = rawLine.strip()
    if line.len == 0: continue
    let node = parseJson(line)
    if node.hasKey("meta"):
      result.meta = parseMeta(node["meta"])
      arenaW = result.meta.arenaW
      arenaH = result.meta.arenaH
      continue
    if node.hasKey("eid"): enemyId = node["eid"].getInt()
    if node.hasKey("end"):
      if node["end"].hasKey("enemy_died"):
        result.enemyDied = node["end"]["enemy_died"].getBool()
      continue
    result.states.add stateFromJson(node, arenaW, arenaH, enemyId)
    result.lastSeen.add (if node.hasKey("lst"): node["lst"].getInt() else: -1)
  result.enemyId = enemyId

# ── replay engine ─────────────────────────────────────────────────────────────

proc reportFor(tracker: VirtualTracker, drivers: seq[GunDriver],
               targetId: int): seq[GunReport] =
  ## Exactly the accounting ModularBot.onRoundEnded does.
  let fit = tracker.fitnessFor(targetId)
  for gi in 0..<drivers.len:
    var r = GunReport(name: drivers[gi].name)
    for binIdx in 0..<len(PowerBins):
      let fw = fit[gi].bins[binIdx]
      let n = min(fw.count, WindowSize)
      var hits = 0
      for k in 0..<n:
        if fw.hits[k]: inc hits
      r.bins[binIdx] = BinStat(shots: n, hits: hits)
      r.shots += n
      r.hits += hits
    result.add r

proc replayFixture*(fx: Fixture, drivers: seq[GunDriver],
                    targetId = -1, liveActual = false,
                    metric = ActiveMetric): seq[GunReport] =
  ## Drive a fresh `VirtualTracker` over the whole fixture, one tick at a time,
  ## in the same order the live loop uses:
  ##   1. predict(state, bulletSpeed(PowerBins[i])) for i = 0..3, per gun
  ##   2. spawnBullets for every gun (skipped for a not-yet-warm gun)
  ##   3. tickBullets -> onResult on the owning gun
  ## The gun selector (and therefore the 5th predict on the selected gun) is not
  ## replayed: every gun's per-tick caches are tick-guarded, so it is a no-op.
  ##
  ## `liveActual` reproduces two ModularBot ordering quirks discovered by the
  ## Task 3 acceptance test:
  ##  1. The live loop calls go() (which dispatches the NEXT tick's scan into
  ##     enemyTracker) BEFORE it builds the enemy table handed to tickBullets.
  ##     A fixture recorded at the WorldState construction site therefore
  ##     snapshots the position one tick before the one the live resolver used,
  ##     so the resolver reads the NEXT state's enemy pose/lastSeenTick.
  ##  2. If the target died during that final go(), the live aim block (spawn AND
  ##     resolution) is skipped entirely; the end marker tells us so and we drop
  ##     the final tick's resolutions. Synthetic fixtures leave liveActual false
  ##     (the state at the resolution tick is the ground truth).
  ##
  ## `metric` defaults to the process-wide `GUN_VBULLET_METRIC` switch read by
  ## virtual_bullets; pass it explicitly only to force a model in one process.
  let tid = if targetId >= 0: targetId else: fx.enemyId
  let skipFinal = liveActual and fx.enemyDied
  var tracker = initTracker(drivers.len, metric)
  for si in 0..<fx.states.len:
    let state = fx.states[si]
    for gi in 0..<drivers.len:
      var preds: array[len(PowerBins), GunPrediction]
      for i in 0..<len(PowerBins):
        preds[i] = drivers[gi].predictCb(state, bulletSpeed(PowerBins[i]))
      let ready = if drivers[gi].readyCb == nil: true else: drivers[gi].readyCb()
      if ready:
        tracker.spawnBullets(gi, preds, state, tid)

    let actIdx = if liveActual and si + 1 < fx.states.len: si + 1 else: si
    let act = fx.states[actIdx]
    var enemyPositions: Table[int, tuple[x, y: float, lastSeenTick: int, alive: bool]]
    var lst = act.tick
    if actIdx < fx.lastSeen.len and fx.lastSeen[actIdx] >= 0: lst = fx.lastSeen[actIdx]
    if act.enemies.len > 0:
      for e in act.enemies:
        enemyPositions[e.id] = (x: e.x, y: e.y, lastSeenTick: lst, alive: true)
    else:
      enemyPositions[tid] = (x: act.enemyX, y: act.enemyY, lastSeenTick: lst, alive: true)

    let dref = drivers
    if not (skipFinal and si == fx.states.len - 1):
      tracker.tickBullets(state, enemyPositions,
        proc(gunId: GunId, binIdx: int, e: FeedbackEvent) =
          dref[gunId].resultCb(e))

  result = reportFor(tracker, drivers, tid)

# ── formatting ────────────────────────────────────────────────────────────────

proc hitRate*(r: GunReport): float =
  if r.shots == 0: 0.0 else: r.hits.float / r.shots.float

proc formatReportRow*(r: GunReport): string =
  ## One deterministic line: name, hits/shots, %, per-bin hits/shots.
  var bins = ""
  for i in 0..<len(PowerBins):
    bins.add fmt"  p{PowerBins[i]:.1f}={r.bins[i].hits}/{r.bins[i].shots}"
  fmt"{r.name:<11} {r.hits:>5}/{r.shots:<5} {hitRate(r)*100.0:>6.1f}%{bins}"

proc formatReports*(fx: Fixture, reports: seq[GunReport]): string =
  result = fmt"# {fx.meta.adversary} (source={fx.meta.source}, ticks={fx.states.len}, enemyId={fx.enemyId})"
  result.add "\n"
  for r in reports:
    result.add formatReportRow(r) & "\n"

# ── synthetic trajectory generators (ground truth by construction) ────────────
#
# Coordinates: 0° = East, CCW positive. Self is stationary at (SelfX, SelfY) so
# that the trajectory is attributable solely to the enemy. Arena 800x600.

const
  SelfX = 200.0
  SelfY = 300.0
  SelfEnergy = 100.0

proc mkState(tick: int, ex, ey, eh, es, ee: float,
             arenaW = DefaultArenaW, arenaH = DefaultArenaH,
             selfX = SelfX, selfY = SelfY, selfEnergy = SelfEnergy,
             enemyId = DefaultEnemyId): WorldState =
  WorldState(
    enemyX: ex, enemyY: ey, enemyHeading: eh, enemySpeed: es, enemyEnergy: ee,
    selfX: selfX, selfY: selfY, selfSpeed: 0.0, selfHeading: 0.0,
    selfRadarHeading: 0.0, selfEnergy: selfEnergy,
    arenaWidth: arenaW, arenaHeight: arenaH, tick: tick,
    enemies: @[EnemyInfo(id: enemyId, x: ex, y: ey, heading: eh, speed: es, energy: ee)],
  )

proc finish(states: seq[WorldState], adversary, note: string): Fixture =
  var meta = initMeta(source = "synthetic")
  meta.adversary = adversary
  meta.note = note
  initFixture(states, meta)

proc synthesizeStationary*(ticks = 200, ex = 600.0, ey = 300.0): Fixture =
  var states: seq[WorldState]
  for t in 0..<ticks:
    states.add mkState(t, ex, ey, 0.0, 0.0, 100.0)
  finish(states, "stationary", "enemy fixed; any good gun scores ~100%")

proc synthesizeConstantVelocity*(ticks = 150, ex = 100.0, ey = 300.0,
                                 heading = 0.0, speed = 4.0): Fixture =
  var states: seq[WorldState]
  var x = ex
  var y = ey
  let hr = degToRad(heading)
  for t in 0..<ticks:
    states.add mkState(t, x, y, heading, speed, 100.0)
    x += cos(hr) * speed
    y += sin(hr) * speed
  finish(states, "constant-velocity",
    "straight line, no walls reached; Linear/HeadOn should lead the target")

proc synthesizeCircular*(ticks = 220, ex = 400.0, ey = 300.0,
                         heading = 0.0, speed = 6.0, turnDeg = 3.0): Fixture =
  var states: seq[WorldState]
  var x = ex
  var y = ey
  var h = heading
  for t in 0..<ticks:
    states.add mkState(t, x, y, h, speed, 100.0)
    let hr = degToRad(h)
    x += cos(hr) * speed
    y += sin(hr) * speed
    h += turnDeg
  finish(states, "circular",
    fmt"constant turn {turnDeg} deg/tick; Circular/Accel should fit")

proc synthesizeWallBounce*(ticks = 240, ex = 100.0, ey = 100.0,
                           heading = 45.0, speed = 6.0): Fixture =
  var states: seq[WorldState]
  var x = ex
  var y = ey
  var h = heading
  let m = BotRadius
  for t in 0..<ticks:
    states.add mkState(t, x, y, h, speed, 100.0)
    let hr = degToRad(h)
    var nx = x + cos(hr) * speed
    var ny = y + sin(hr) * speed
    if nx < m: nx = m; h = 180.0 - h
    elif nx > DefaultArenaW - m: nx = DefaultArenaW - m; h = 180.0 - h
    if ny < m: ny = m; h = -h
    elif ny > DefaultArenaH - m: ny = DefaultArenaH - m; h = -h
    x = nx
    y = ny
  finish(states, "wall-bounce", "specular reflection off all four walls")

proc synthesizeOscillator*(ticks = 240, ex = 200.0, ey = 300.0,
                           speed = 4.0, period = 30): Fixture =
  var states: seq[WorldState]
  var x = ex
  var y = ey
  for t in 0..<ticks:
    let phase = (t div period) mod 2
    let h = if phase == 0: 0.0 else: 180.0
    let hr = degToRad(h)
    states.add mkState(t, x, y, h, speed, 100.0)
    x += cos(hr) * speed
    y += sin(hr) * speed
  finish(states, "oscillator", fmt"east for {period}, then west for {period}, repeat")

proc synthesizeRandomWalk*(ticks = 260, ex = 400.0, ey = 300.0,
                           speed = 4.0, seed = 20250920,
                           maxTurn = 15.0): Fixture =
  var rng = initRand(seed)
  var states: seq[WorldState]
  var x = ex
  var y = ey
  var h = 0.0
  let m = BotRadius
  for t in 0..<ticks:
    states.add mkState(t, x, y, h, speed, 100.0)
    h += rng.rand(-maxTurn .. maxTurn)
    let hr = degToRad(h)
    var nx = x + cos(hr) * speed
    var ny = y + sin(hr) * speed
    if nx < m or nx > DefaultArenaW - m: h = 180.0 - h
    if ny < m or ny > DefaultArenaH - m: h = -h
    nx = clamp(nx, m, DefaultArenaW - m)
    ny = clamp(ny, m, DefaultArenaH - m)
    x = nx
    y = ny
  finish(states, "random-walk",
    fmt"seeded (seed={seed}), +-{maxTurn} deg/tick heading jitter")

proc synthesizeDecelBeforeTurn*(ticks = 260, ex = 100.0, ey = 300.0,
                                cruise = 6.0, cruiseTicks = 45,
                                turnDeg = 45.0): Fixture =
  ## Known rule: accelerate to cruise, hold, brake to a full stop, pivot
  ## `turnDeg` degrees while stopped, then accelerate again. The stop is the
  ## tell-tale StopShot is meant to catch.
  var states: seq[WorldState]
  var x = ex
  var y = ey
  var h = 0.0
  var v = 0.0
  var phase = "accel"
  var phaseT = 0
  for t in 0..<ticks:
    states.add mkState(t, x, y, h, v, 100.0)
    case phase
    of "accel":
      v = min(cruise, v + 1.0)
      if v >= cruise:
        phase = "cruise"; phaseT = 0
    of "cruise":
      inc phaseT
      if phaseT >= cruiseTicks: phase = "brake"
    of "brake":
      v = max(0.0, v - 3.0)
      if v <= 0.0:
        v = 0.0
        phase = "pivot"; phaseT = 0
    of "pivot":
      if phaseT < 3:
        h += turnDeg
      inc phaseT
      if phaseT >= 4:
        phase = "accel"
    else: discard
    let hr = degToRad(h)
    x += cos(hr) * v
    y += sin(hr) * v
    x = clamp(x, BotRadius, DefaultArenaW - BotRadius)
    y = clamp(y, BotRadius, DefaultArenaH - BotRadius)
  finish(states, "decel-before-turn",
    "cruise -> full stop -> pivot 3x45deg -> accelerate; tests StopShot")

proc synthesizeEnergyThresholdTurner*(ticks = 200, ex = 100.0, ey = 300.0,
                                      e0 = 50.0, decay = 0.5, threshold = 30.0,
                                      speed = 5.0, hardTurnDeg = 20.0): Fixture =
  ## The falsifiable one. RULE (known by construction, stated in `note`):
  ##   energy(t) = max(5, e0 - decay*t)
  ##   while energy >= threshold: heading constant (straight, predictable)
  ##   while energy <  threshold: heading += hardTurnDeg each tick (hard turn)
  ## A learner that finds the rule should switch from straight-line prediction
  ## to turn prediction exactly at the tick energy crosses the threshold.
  var states: seq[WorldState]
  var x = ex
  var y = ey
  var h = 0.0
  let m = BotRadius
  for t in 0..<ticks:
    let e = max(5.0, e0 - decay * t.float)
    states.add mkState(t, x, y, h, speed, e)
    if e < threshold: h += hardTurnDeg
    let hr = degToRad(h)
    var nx = x + cos(hr) * speed
    var ny = y + sin(hr) * speed
    if nx < m or nx > DefaultArenaW - m: h = 180.0 - h
    if ny < m or ny > DefaultArenaH - m: h = -h
    nx = clamp(nx, m, DefaultArenaW - m)
    ny = clamp(ny, m, DefaultArenaH - m)
    x = nx
    y = ny
  var meta = initMeta(source = "synthetic")
  meta.adversary = "energy-threshold-turner"
  meta.note = fmt"RULE: straight while energy>={threshold}; hard {hardTurnDeg} deg/tick turn below {threshold}; energy={e0}-{decay}*t (floor 5)"
  initFixture(states, meta)

# Named fixtures for the demo runner.
const SyntheticFixtureNames* = [
  "stationary",
  "constant-velocity",
  "circular",
  "wall-bounce",
  "oscillator",
  "random-walk",
  "decel-before-turn",
  "energy-threshold-turner",
]

proc synthesizeByName*(name: string): Fixture =
  case name
  of "stationary": synthesizeStationary()
  of "constant-velocity": synthesizeConstantVelocity()
  of "circular": synthesizeCircular()
  of "wall-bounce": synthesizeWallBounce()
  of "oscillator": synthesizeOscillator()
  of "random-walk": synthesizeRandomWalk()
  of "decel-before-turn": synthesizeDecelBeforeTurn()
  of "energy-threshold-turner": synthesizeEnergyThresholdTurner()
  else:
    raise newException(ValueError, "unknown synthetic fixture: " & name)