## Unit guard for the energy-aware power policy (TR_POWER_*).
##
## The policy is a CAP on the gun's own preferred bin: at long range or low
## energy it trades single-hit damage for a faster bullet and more shots, and it
## withholds power 3.0 unless the gun's own chances for the chosen bin are above
## that gun's average. It must NEVER raise power, and `TR_POWER_POLICY=0` must
## reproduce the uncapped preference exactly (the control arm).
##
## Pure: no Java, no battle. Run:
##   nim c -r common_libs/tests/test_power_policy.nim
## and once with the flag off:
##   TR_POWER_POLICY=0 nim c -r common_libs/tests/test_power_policy.nim

import std/[strformat, math, tables]
import gun_harness/virtual_bullets
import gun_harness/selector

var failures = 0
proc check(name: string, ok: bool) =
  if ok: echo "PASS: ", name
  else: echo "FAIL: ", name; inc failures

proc recordHit(fw: var FitnessWindow, hit: bool) =
  fw.hits[fw.head] = hit
  fw.head = (fw.head + 1) mod WindowSize
  inc fw.count

proc seedWindow(t: var VirtualTracker, targetId, gunId, binIdx, hits, misses: int) =
  if targetId notin t.fitness:
    t.fitness[targetId] = newSeq[GunFitness](t.numGuns)
  var fw = addr t.fitness[targetId][gunId].bins[binIdx]
  for _ in 0..<hits:   recordHit(fw[], true)
  for _ in 0..<misses: recordHit(fw[], false)

# ── pure cap core ────────────────────────────────────────────────────────────

proc testFar() =
  let d = applyPowerPolicy(preferredPower = 3.0, dist = 250.0, selfEnergy = 100.0,
                           pEst = 0.9, pRef = 0.1, ramming = false, enabled = true)
  check "distance > TR_POWER_FAR_DIST -> cap 1.0 (far)",
        d.power == 1.0 and d.cap == 1.0 and d.reason == prFar

proc testLowEnergy() =
  let d = applyPowerPolicy(3.0, 100.0, 30.0, 0.9, 0.1, false, enabled = true)
  check "self energy < TR_POWER_LOW_ENERGY -> cap 1.0 (lowEnergy)",
        d.power == 1.0 and d.cap == 1.0 and d.reason == prLowEnergy

proc testBelowAverage() =
  # pEst == pRef is "not above average": withhold power 3.0 -> cap 2.0.
  let d = applyPowerPolicy(3.0, 100.0, 100.0, 0.2, 0.2, false, enabled = true)
  check "chances not above average -> cap 2.0 (belowAvg)",
        d.power == 2.0 and d.cap == 2.0 and d.reason == prBelowAvg
  # Strictly below also caps.
  let e = applyPowerPolicy(3.0, 100.0, 100.0, 0.1, 0.2, false, enabled = true)
  check "chances strictly below average -> cap 2.0 (belowAvg)",
        e.power == 2.0 and e.reason == prBelowAvg

proc testAboveAverageFull() =
  let d = applyPowerPolicy(3.0, 100.0, 100.0, 0.5, 0.2, false, enabled = true)
  check "above average + close + healthy -> cap 3.0 (full)",
        d.power == 3.0 and d.cap == 3.0 and d.reason == prFull

proc testRamExempt() =
  # Far, low energy, no chance data: still full power because we are ramming.
  let d = applyPowerPolicy(3.0, 500.0, 5.0, 0.0, 0.9, true)
  check "ramming exempts the caps even far + low energy",
        d.power == 3.0 and d.reason == prRam
  # Ram beats far even with the policy on.
  let e = applyPowerPolicy(3.0, 500.0, 5.0, 0.0, 0.9, true, enabled = true)
  check "ram exemption takes precedence over the far cap", e.power == 3.0

proc testCapNeverRaises() =
  # Preferred below every cap must pass through untouched.
  let a = applyPowerPolicy(1.0, 500.0, 5.0, 0.0, 0.9, false, enabled = true)
  check "far cap does not raise a preferred p=1.0", a.power == 1.0
  let b = applyPowerPolicy(1.5, 100.0, 100.0, 0.1, 0.2, false, enabled = true)
  check "mid cap does not raise a preferred p=1.5", b.power == 1.5
  let c = applyPowerPolicy(1.5, 100.0, 100.0, 0.9, 0.2, false, enabled = true)
  check "full cap does not raise a preferred p=1.5", c.power == 1.5

proc testPrecedence() =
  # far beats low energy, and low energy beats belowAvg.
  let a = applyPowerPolicy(3.0, 500.0, 5.0, 0.0, 0.9, false, enabled = true)
  check "far takes precedence over low energy", a.reason == prFar
  let b = applyPowerPolicy(3.0, 100.0, 5.0, 0.0, 0.9, false, enabled = true)
  check "low energy takes precedence over belowAvg", b.reason == prLowEnergy

proc testDisabledUncapped() =
  # enabled=false is the code path TR_POWER_POLICY=0 drives.
  for p in [1.0, 1.5, 2.0, 3.0]:
    let d = applyPowerPolicy(p, 500.0, 5.0, 0.0, 0.9, false, enabled = false)
    check fmt"policy off reproduces the uncapped preference p={p:.1f}",
          d.power == p and d.cap == 3.0

proc testEnvFlag() =
  # The flag is read once at module init, so the active branch is selected by
  # the process environment. Run the test twice (default and TR_POWER_POLICY=0).
  if PowerPolicyEnabled:
    check "TR_POWER_POLICY default (on): far shot is capped to 1.0",
          applyPowerPolicy(3.0, 300.0, 100.0, 0.5, 0.2, false).power == 1.0
  else:
    check "TR_POWER_POLICY=0: far shot is NOT capped (control arm)",
          applyPowerPolicy(3.0, 300.0, 100.0, 0.5, 0.2, false).power == 3.0

proc testBinIndex() =
  check "binIndexForPower maps the shipped bins exactly",
        binIndexForPower(1.0) == 0 and binIndexForPower(1.5) == 1 and
        binIndexForPower(2.0) == 2 and binIndexForPower(3.0) == 3
  check "binIndexForPower snaps a non-bin cap to the highest bin not above it",
        binIndexForPower(2.5) == 2 and binIndexForPower(0.5) == 0

proc testReasonNames() =
  check "reason names match the documented log vocabulary",
        powerReasonName(prFar) == "far" and
        powerReasonName(prLowEnergy) == "lowEnergy" and
        powerReasonName(prBelowAvg) == "belowAvg" and
        powerReasonName(prFull) == "full" and
        powerReasonName(prRam) == "ram"

# ── integration through the tracker (per-target fitness + empty bins) ─────────
# These depend on the process-wide flag, so each branch asserts the behaviour of
# the mode it is actually running in.

proc testTrackerAboveAverage() =
  var t = initTracker(1)
  # bin3 is far above this gun's aggregate (50% vs 20%): above average.
  seedWindow(t, 7, 0, 0, 10, 90)
  seedWindow(t, 7, 0, 1, 10, 90)
  seedWindow(t, 7, 0, 2, 10, 90)
  seedWindow(t, 7, 0, 3, 50, 50)
  let (g, _, p, d) = t.selectShotPolicy(7, tick = 0, dist = 100.0,
                                        selfEnergy = 100.0, ramming = false)
  if PowerPolicyEnabled:
    check "tracker: above-average bin + close + healthy -> power 3.0",
          g == 0 and p == 3.0 and d.reason == prFull
  else:
    check "tracker (control): above-average bin still fires power 3.0",
          g == 0 and p == 3.0

proc testTrackerBelowAverage() =
  var t = initTracker(1)
  for b in 0..<len(PowerBins):
    seedWindow(t, 7, 0, b, 20, 80)  # all bins equal: chosen bin not above mean
  let (_, _, p, d) = t.selectShotPolicy(7, tick = 0, dist = 100.0,
                                        selfEnergy = 100.0, ramming = false)
  if PowerPolicyEnabled:
    check "tracker: flat gun -> cap 2.0 (belowAvg)", p == 2.0 and d.reason == prBelowAvg
  else:
    check "tracker (control): flat gun keeps its uncapped preference (power 3.0)",
          p == 3.0

proc testTrackerFarAndLowEnergy() =
  var t = initTracker(1)
  seedWindow(t, 7, 0, 3, 50, 50)
  let (_, _, pFar, dFar) = t.selectShotPolicy(7, 0, dist = 250.0,
                                              selfEnergy = 100.0, ramming = false)
  let (_, _, pLow, dLow) = t.selectShotPolicy(7, 0, dist = 100.0,
                                              selfEnergy = 30.0, ramming = false)
  if PowerPolicyEnabled:
    check "tracker: far -> power 1.0", pFar == 1.0 and dFar.reason == prFar
    check "tracker: low energy -> power 1.0", pLow == 1.0 and dLow.reason == prLowEnergy
  else:
    check "tracker (control): far does NOT cap (uncapped preference)", pFar == 3.0
    check "tracker (control): low energy does NOT cap (uncapped preference)", pLow == 3.0

proc testTrackerColdAndEmptyBin() =
  var t = initTracker(1)
  let (_, _, pCold, dCold) = t.selectShotPolicy(7, 0, dist = 100.0,
                                                selfEnergy = 100.0, ramming = false)
  if PowerPolicyEnabled:
    # Cold gun: no data at all -> pEst <= pRef vacuously -> mid cap, but the
    # preferred bin is already 1.0, so the fired power stays 1.0.
    check "tracker: cold gun gets the mid cap but keeps its p=1.0 preference",
          pCold == 1.0 and dCold.reason == prBelowAvg
  else:
    check "tracker (control): cold gun keeps its p=1.0 preference", pCold == 1.0

proc testTrackerRamExempt() =
  var t = initTracker(1)
  seedWindow(t, 7, 0, 3, 50, 50)
  let (_, _, p, d) = t.selectShotPolicy(7, 0, dist = 500.0,
                                        selfEnergy = 5.0, ramming = true)
  check "tracker: ramming exempts the caps", p == 3.0 and d.reason == prRam

# ── driver ───────────────────────────────────────────────────────────────────

testFar()
testLowEnergy()
testBelowAverage()
testAboveAverageFull()
testRamExempt()
testCapNeverRaises()
testPrecedence()
testDisabledUncapped()
testEnvFlag()
testBinIndex()
testReasonNames()
testTrackerAboveAverage()
testTrackerBelowAverage()
testTrackerFarAndLowEnergy()
testTrackerColdAndEmptyBin()
testTrackerRamExempt()

if failures > 0:
  echo "\n", failures, " check(s) FAILED"
  quit(1)
echo "\nAll power-policy checks passed."