import std/math
import binary_encoding

const
  N_SECTORS      = 8
  N_BANDS        = 3
  LEARNING_RATE  = 0.2
  BAND_THR_LO    = 33.0
  BAND_THR_HI    = 66.0

type
  ResidualTable* = object
    corrections*: array[N_SECTORS * N_BANDS, tuple[cx, cy: float]]

proc initPredictor*(): ResidualTable = discard  # zero-init is correct

proc getSector*(headingSin, headingCos: int): int =
  let s = headingSin.float / 199.0 * 2.0 - 1.0
  let c = headingCos.float / 199.0 * 2.0 - 1.0
  var deg = radToDeg(arctan2(s, c))
  if deg < 0.0: deg += 360.0
  int(deg / 45.0) mod N_SECTORS

proc getBand*(distance: int): int =
  if distance.float < BAND_THR_LO: 0
  elif distance.float < BAND_THR_HI: 1
  else: 2

proc predict*(table: ResidualTable,
              f0, f1: array[NUM_FRAME_FIELDS, int],
              power: float): tuple[predX, predY: float] =
  # field indices: 2=distance, 4=heading_sin, 5=heading_cos, 6=enemy_x, 7=enemy_y
  let bulletSpd    = 20.0 - 3.0 * power
  let distancePx   = f0[2].float / 99.0 * 1414.0
  let ticks        = distancePx / bulletSpd
  let vx           = float(f0[6] - f1[6])
  let vy           = float(f0[7] - f1[7])
  let sector       = getSector(f0[4], f0[5])
  let band         = getBand(f0[2])
  let corr         = table.corrections[sector * N_BANDS + band]
  result.predX = f0[6].float + vx * ticks + corr.cx
  result.predY = f0[7].float + vy * ticks + corr.cy

proc sectorBand*(f0: array[NUM_FRAME_FIELDS, int]): tuple[sector, band: int] =
  (getSector(f0[4], f0[5]), getBand(f0[2]))

proc learn*(table: var ResidualTable, sector, band: int,
            residualX, residualY: float) =
  table.corrections[sector * N_BANDS + band].cx += LEARNING_RATE * residualX
  table.corrections[sector * N_BANDS + band].cy += LEARNING_RATE * residualY