# ponytail: grid accumulator replaces ring buffer; add temporal features when needed (step 3)

import std/math

const
  VPERP_BINS* = 17    # -8 to +8 inclusive (integer speed units)
  DIST_BINS*  = 8     # distance bands
  DIST_BAND*  = 125.0 # pixels per band

type
  GridCell = object
    sumSin: float
    sumCos: float
    count:  float

  LeadGrid* = object
    cells: array[VPERP_BINS * DIST_BINS, GridCell]  # 17 × 8 = 136 cells

const BULLET_SPEED = 14.0  # power 2 bullet speed in pixels/tick

proc initLeadGrid*(): LeadGrid =
  for vBin in 0..<VPERP_BINS:
    for dBin in 0..<DIST_BINS:
      let vPerp = float(vBin) - 8.0
      let sinVal = clamp(vPerp / BULLET_SPEED, -1.0, 1.0)
      let offsetRad = arcsin(sinVal)
      let idx = vBin * DIST_BINS + dBin
      result.cells[idx].sumSin = 0.1 * sin(offsetRad)
      result.cells[idx].sumCos = 0.1 * cos(offsetRad)
      result.cells[idx].count = 0.1

proc cellIndex(vPerp: float, distance: float): int {.inline.} =
  let vBin = clamp(int(vPerp + 8.5), 0, VPERP_BINS - 1)
  let dBin = clamp(int(distance / DIST_BAND), 0, DIST_BINS - 1)
  result = vBin * DIST_BINS + dBin

proc forward*(grid: var LeadGrid, vPerp, distance: float): float =
  ## Returns circular mean offset in degrees, or -999.0 when no data.
  let idx = cellIndex(vPerp, distance)
  let vBin = clamp(int(vPerp + 8.5), 0, VPERP_BINS - 1)
  let dBin = clamp(int(distance / DIST_BAND), 0, DIST_BINS - 1)

  # Accumulate weighted contributions: direct cell (w=1), cardinal neighbors (w=0.5), diagonals (w=0.25)
  var sinSum = 0.0
  var cosSum = 0.0
  var totalW = 0.0
  for dv in -1 .. 1:
    for dd in -1 .. 1:
      let vb = vBin + dv
      let db = dBin + dd
      if vb < 0 or vb >= VPERP_BINS or db < 0 or db >= DIST_BINS: continue
      let c = grid.cells[vb * DIST_BINS + db]
      if c.count < 0.001: continue
      let w = if dv == 0 and dd == 0: 1.0
              elif dv == 0 or dd == 0: 0.5
              else: 0.25
      sinSum += w * c.sumSin / c.count
      cosSum += w * c.sumCos / c.count
      totalW += w
  if totalW == 0.0:
    return -999.0
  result = radToDeg(arctan2(sinSum, cosSum))

proc learn*(grid: var LeadGrid, vPerp, distance, correctOffset: float) =
  let idx = cellIndex(vPerp, distance)
  grid.cells[idx].sumSin += sin(degToRad(correctOffset))
  grid.cells[idx].sumCos += cos(degToRad(correctOffset))
  grid.cells[idx].count += 1.0

proc totalCount*(grid: LeadGrid): float =
  for c in grid.cells:
    result += c.count