## Unit guard for the energy-aware power policy (TR_POWER_*).
##
## The policy is a CAP on the gun's own preferred bin. Every rule can only LOWER
## power; `power = min(preference, min(all caps))`. The two energy-economy rules
## under test here are:
##
##   * the ENERGY SLOPE — a linear cap on OUR energy (`energySlopeCap`), replacing
##     the old hard cliff at `TR_POWER_LOW_ENERGY` (50);
##   * the FINISHING SLOPE — cap power at the smallest bullet that still removes
##     the ENEMY's remaining energy (`powerToKill`), because server 1.3.1 caps the
##     `bulletDamage` score at the energy actually removed, so overkill is wasted.
##
## `TR_POWER_POLICY=0` must reproduce the uncapped preference exactly (control).
##
## 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 close(a, b: float, eps = 1e-9): bool = abs(a - b) < eps

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)

# ── CHANGE 1: the energy slope (pure formula) ────────────────────────────────
# Defaults: capMax 3.0 at/above 80, capMin 0.5 at/below 20, linear between.

proc testEnergySlope() =
  # Above / at HI: no cap from this rule.
  check "energy slope: selfE=100 -> cap 3.0 (no cap)",
        close(energySlopeCap(100.0, 80.0, 20.0, 0.5, 3.0), 3.0)
  check "energy slope: selfE=80 (HI) -> cap 3.0",
        close(energySlopeCap(80.0, 80.0, 20.0, 0.5, 3.0), 3.0)
  # At/below LO: the floor cap.
  check "energy slope: selfE=20 (LO) -> cap 0.5",
        close(energySlopeCap(20.0, 80.0, 20.0, 0.5, 3.0), 0.5)
  check "energy slope: selfE=5 -> cap 0.5 (clamped at LO)",
        close(energySlopeCap(5.0, 80.0, 20.0, 0.5, 3.0), 0.5)
  # Linear interpolation in POWER.
  check "energy slope: selfE=50 -> cap 1.75 (halfway 0.5..3.0)",
        close(energySlopeCap(50.0, 80.0, 20.0, 0.5, 3.0), 1.75)
  check "energy slope: selfE=65 -> cap 2.375 (t=0.75)",
        close(energySlopeCap(65.0, 80.0, 20.0, 0.5, 3.0), 2.375)
  check "energy slope: selfE=35 -> cap 1.125 (t=0.25)",
        close(energySlopeCap(35.0, 80.0, 20.0, 0.5, 3.0), 1.125)
  # A non-default top: the interpolation endpoint is capMax, not hardcoded 3.0.
  check "energy slope: custom capMax=2.0 at HI",
        close(energySlopeCap(80.0, 80.0, 20.0, 0.5, 2.0), 2.0)

proc testEnergySlopeMonotone() =
  # The cap must be non-decreasing in our energy (more energy never caps lower).
  var mono = true
  var prev = -1.0
  for i in 0..20:
    let e = i.float * 5.0
    let c = energySlopeCap(e, 80.0, 20.0, 0.5, 3.0)
    if c < prev - 1e-12: mono = false
    prev = c
  check "energy slope: cap is non-decreasing in self energy", mono

# ── CHANGE 2: the finishing slope (pure formula) ─────────────────────────────

proc testPowerToKill() =
  # E <= 4: p = E/4, clamped to >= 0.1.
  check "finish: E=0.1 -> p=0.1 (clamp)",
        close(powerToKill(0.1), 0.1)
  check "finish: E=0.4 -> p=0.1 (clamp)",
        close(powerToKill(0.4), 0.1)
  check "finish: E=1 -> p=0.25",
        close(powerToKill(1.0), 0.25)
  check "finish: E=2 -> p=0.5",
        close(powerToKill(2.0), 0.5)
  check "finish: E=4 -> p=1.0 (branch seam)",
        close(powerToKill(4.0), 1.0)
  # E > 4: p = (E+2)/6.
  check "finish: E=6 -> p=4/3",
        close(powerToKill(6.0), 4.0 / 3.0)
  check "finish: E=10 -> p=2.0",
        close(powerToKill(10.0), 2.0)
  check "finish: E=16 -> p=3.0",
        close(powerToKill(16.0), 3.0)
  # E > 16: no cap (p=3.0 removes only 16).
  check "finish: E=17 -> p=3.0 (no cap)",
        close(powerToKill(17.0), 3.0)
  check "finish: E=20 -> p=3.0 (no cap)",
        close(powerToKill(20.0), 3.0)
  check "finish: E=100 -> p=3.0 (no cap)",
        close(powerToKill(100.0), 3.0)

proc testPowerToKillCovers() =
  # The formula must be a TRUE inverse for the range it can cover (E <= 16);
  # above 16 even p=3.0 removes only 16, so there is no cap to find.
  var covers = true
  for i in 0..159:
    let e = 0.1 + i.float * 0.1   # 0.1 .. 16.0
    if bulletDamageAtPower(powerToKill(e)) < e - 1e-9: covers = false
  check "finish: powerToKill(E) always removes at least E (E in 0.1..16.0)",
        covers

proc testPowerToKillMonotone() =
  var mono = true
  var prev = -1.0
  for i in 0..200:
    let e = i.float * 0.1
    let p = powerToKill(e)
    if p < prev - 1e-12: mono = false
    prev = p
  check "finish: powerToKill is non-decreasing in enemy energy", mono

# ── the combined 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 testEnergySlopeCap() =
  # Close, above average, healthy enemy: only the energy slope binds.
  let d = applyPowerPolicy(3.0, 100.0, 50.0, 0.9, 0.1, false,
                           enemyEnergy = 100.0, enabled = true)
  check "selfE=50 -> cap 1.75 (energySlope)",
        close(d.power, 1.75) and close(d.cap, 1.75) and d.reason == prEnergySlope
  let e = applyPowerPolicy(3.0, 100.0, 20.0, 0.9, 0.1, false,
                           enemyEnergy = 100.0, enabled = true)
  check "selfE=20 -> cap 0.5 (energySlope)",
        close(e.power, 0.5) and e.reason == prEnergySlope
  let f = applyPowerPolicy(3.0, 100.0, 80.0, 0.9, 0.1, false,
                           enemyEnergy = 100.0, enabled = true)
  check "selfE=80 -> no energy cap, full power", f.power == 3.0

proc testFinishingCap() =
  # Close, healthy self, above average, low enemy: only finishing binds.
  let d = applyPowerPolicy(3.0, 100.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 2.0, enabled = true, finishKill = true)
  check "enemyE=2 -> smallest killing bullet p=0.5 (finishKill)",
        close(d.power, 0.5) and close(d.cap, 0.5) and d.reason == prFinishKill
  let e = applyPowerPolicy(3.0, 100.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 10.0, enabled = true, finishKill = true)
  check "enemyE=10 -> p=2.0 (finishKill)",
        close(e.power, 2.0) and e.reason == prFinishKill
  let f = applyPowerPolicy(3.0, 100.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 16.0, enabled = true, finishKill = true)
  check "enemyE=16 -> p=3.0 (finish cap does not bite)", f.power == 3.0
  let g = applyPowerPolicy(3.0, 100.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 20.0, enabled = true, finishKill = true)
  check "enemyE=20 -> no finish cap", g.power == 3.0
  # Disabled finishing leaves power alone.
  let h = applyPowerPolicy(3.0, 100.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 2.0, enabled = true, finishKill = false)
  check "enemyE=2 with finishKill=false -> uncapped", h.power == 3.0
  # A dead/unknown target (energy 0) must NOT collapse power to 0.1.
  let i = applyPowerPolicy(3.0, 100.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 0.0, enabled = true, finishKill = true)
  check "enemyE=0 (dead) -> finish rule skipped, no collapse", i.power == 3.0

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
  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 self energy, low enemy energy, no chance data: still full power.
  let d = applyPowerPolicy(3.0, 500.0, 5.0, 0.0, 0.9, true,
                           enemyEnergy = 1.0, enabled = true, finishKill = true)
  check "ramming exempts far + energy slope + finishing",
        d.power == 3.0 and d.reason == prRam

proc testCapNeverRaises() =
  # Preferred below every cap must pass through untouched.
  # Caps may LOWER a preference (the new slope does exactly that here) but must
  # never raise it: power is always min(preference, cap).
  let a = applyPowerPolicy(1.0, 500.0, 5.0, 0.0, 0.9, false,
                           enemyEnergy = 1.0, enabled = true, finishKill = true)
  check "all caps never raise a preferred p=1.0 (power = min(pref, cap))",
        a.power <= 1.0 and close(a.power, min(1.0, a.cap))
  let b = applyPowerPolicy(1.5, 100.0, 100.0, 0.1, 0.2, false,
                           enemyEnergy = 20.0, enabled = true, finishKill = true)
  check "belowAvg cap does not raise a preferred p=1.5", b.power == 1.5
  # Exhaustive: for every cap combination, power <= preference.
  var never = true
  for pref in [0.5, 1.0, 1.5, 2.0, 3.0]:
    for selfE in [5.0, 20.0, 50.0, 80.0, 100.0]:
      for enemyE in [0.0, 0.5, 2.0, 10.0, 20.0]:
        for d in [50.0, 250.0]:
          let r = applyPowerPolicy(pref, d, selfE, 0.1, 0.2, false,
                                   enemyEnergy = enemyE, enabled = true,
                                   finishKill = true)
          if r.power > pref + 1e-12: never = false
  check "cap never raises power (exhaustive sweep)", never

proc testPrecedence() =
  # far (1.0) beats the energy slope (1.75 at selfE=50).
  let a = applyPowerPolicy(3.0, 500.0, 50.0, 0.0, 0.9, false,
                           enemyEnergy = 100.0, enabled = true)
  check "far takes precedence over the energy slope", a.reason == prFar
  # the energy slope (1.75) beats belowAvg (2.0).
  let b = applyPowerPolicy(3.0, 100.0, 50.0, 0.0, 0.9, false,
                           enemyEnergy = 100.0, enabled = true)
  check "energy slope takes precedence over belowAvg",
        b.reason == prEnergySlope
  # finishing (0.5 at enemyE=2) is smaller than far (1.0), so it wins.
  let c = applyPowerPolicy(3.0, 500.0, 100.0, 0.9, 0.1, false,
                           enemyEnergy = 2.0, enabled = true, finishKill = true)
  check "finishing beats far when it is the smaller cap",
        close(c.power, 0.5) and c.reason == prFinishKill

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,
                             enemyEnergy = 1.0, enabled = false, finishKill = true)
    check fmt"policy off reproduces the uncapped preference p={p:.1f}",
          d.power == p and d.cap == 3.0 and d.reason == prFull

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
    check "TR_POWER_POLICY default (on): low self energy is capped by the slope",
          applyPowerPolicy(3.0, 100.0, 20.0, 0.5, 0.2, false).power == 0.5
  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
    check "TR_POWER_POLICY=0: low self energy is NOT capped (control arm)",
          applyPowerPolicy(3.0, 100.0, 20.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(prEnergySlope) == "energySlope" and
        powerReasonName(prFinishKill) == "finishKill" 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, enemyEnergy = 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 testTrackerEnergySlope() =
  var t = initTracker(1)
  seedWindow(t, 7, 0, 3, 50, 50)
  let (_, _, p, d) = t.selectShotPolicy(7, tick = 0, dist = 100.0,
                                        selfEnergy = 50.0, enemyEnergy = 100.0,
                                        ramming = false)
  if PowerPolicyEnabled:
    check "tracker: selfE=50 -> power 1.75 (energySlope)",
          close(p, 1.75) and d.reason == prEnergySlope
  else:
    check "tracker (control): selfE=50 keeps the uncapped preference (3.0)",
          p == 3.0

proc testTrackerFinishing() =
  var t = initTracker(1)
  seedWindow(t, 7, 0, 3, 50, 50)
  let (_, _, p, d) = t.selectShotPolicy(7, tick = 0, dist = 100.0,
                                        selfEnergy = 100.0, enemyEnergy = 2.0,
                                        ramming = false)
  if PowerPolicyEnabled and PowerFinishKill:
    check "tracker: enemyE=2 -> power 0.5 (finishKill)",
          close(p, 0.5) and d.reason == prFinishKill
  else:
    check "tracker (control/finish off): enemyE=2 keeps preference (3.0)",
          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, enemyEnergy = 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 testTrackerFar() =
  var t = initTracker(1)
  seedWindow(t, 7, 0, 3, 50, 50)
  let (_, _, pFar, dFar) = t.selectShotPolicy(7, 0, dist = 250.0,
                                              selfEnergy = 100.0, enemyEnergy = 100.0,
                                              ramming = false)
  if PowerPolicyEnabled:
    check "tracker: far -> power 1.0", pFar == 1.0 and dFar.reason == prFar
  else:
    check "tracker (control): far does NOT cap (uncapped preference)", pFar == 3.0

proc testTrackerColdAndEmptyBin() =
  var t = initTracker(1)
  let (_, _, pCold, dCold) = t.selectShotPolicy(7, 0, dist = 100.0,
                                                selfEnergy = 100.0, enemyEnergy = 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, enemyEnergy = 1.0,
                                        ramming = true)
  check "tracker: ramming exempts the caps", p == 3.0 and d.reason == prRam

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

testEnergySlope()
testEnergySlopeMonotone()
testPowerToKill()
testPowerToKillCovers()
testPowerToKillMonotone()
testFar()
testEnergySlopeCap()
testFinishingCap()
testBelowAverage()
testAboveAverageFull()
testRamExempt()
testCapNeverRaises()
testPrecedence()
testDisabledUncapped()
testEnvFlag()
testBinIndex()
testReasonNames()
testTrackerAboveAverage()
testTrackerEnergySlope()
testTrackerFinishing()
testTrackerBelowAverage()
testTrackerFar()
testTrackerColdAndEmptyBin()
testTrackerRamExempt()

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