import std/math

const
  FRAME_BITS*  = 83
  SELF_BITS*   = 40
  WINDOW_SIZE* = 10
  TOTAL_BITS*  = FRAME_BITS * WINDOW_SIZE + SELF_BITS  # 870

type
  BinaryVector* = array[TOTAL_BITS, uint8]

  EnemyScanFrame* = object
    bearing*:     float   # → sin/cos, 8+8 bits
    distance*:    float   # → %, 7 bits
    velocity*:    float   # → 5 bits
    heading*:     float   # → sin/cos, 8+8 bits
    enemyWallN*:  float   # 7 bits
    enemyWallS*:  float   # 7 bits
    enemyWallE*:  float   # 7 bits
    enemyWallW*:  float   # 7 bits
    enemyEnergy*: float   # 11 bits

  SelfState* = object
    myWallN*:  float   # 7 bits
    myWallS*:  float   # 7 bits
    myWallE*:  float   # 7 bits
    myWallW*:  float   # 7 bits
    myEnergy*: float   # 11 bits
    canFire*:  bool    # 1 bit

const
  MAX_DISTANCE = 1414.0  # diagonal of 1000x1000 arena
  NUM_FRAME_FIELDS* = 9  # for CSV output: bearing_sin, bearing_cos, distance, velocity, heading_sin, heading_cos, enemy_x, enemy_y, energy
  FIELD_NAMES* = ["bearing_sin", "bearing_cos", "distance", "velocity", "heading_sin", "heading_cos", "wallN", "wallS", "energy"]

proc toGray(value: int): int =
  value xor (value shr 1)

proc toBits(value: int, bits: int): seq[uint8] =
  result = newSeq[uint8](bits)
  let gray = toGray(value)
  for i in 0..<bits:
    result[bits - 1 - i] = uint8((gray shr i) and 1)

proc encodeFrame*(frame: EnemyScanFrame): array[FRAME_BITS, uint8] =
  var offset = 0

  # bearing sin: 8 bits
  let bSin = int((sin(degToRad(frame.bearing)) + 1.0) / 2.0 * 199.0)
  let bSinBits = toBits(clamp(bSin, 0, 199), 8)
  for i in 0..<8: result[offset + i] = bSinBits[i]
  offset += 8

  # bearing cos: 8 bits
  let bCos = int((cos(degToRad(frame.bearing)) + 1.0) / 2.0 * 199.0)
  let bCosBits = toBits(clamp(bCos, 0, 199), 8)
  for i in 0..<8: result[offset + i] = bCosBits[i]
  offset += 8

  # distance %: 7 bits
  let distPct = int(clamp(frame.distance / MAX_DISTANCE * 99.0, 0.0, 99.0))
  let distBits = toBits(distPct, 7)
  for i in 0..<7: result[offset + i] = distBits[i]
  offset += 7

  # velocity: 5 bits
  let vel = int(clamp(frame.velocity + 8.0, 0.0, 16.0))
  let velBits = toBits(vel, 5)
  for i in 0..<5: result[offset + i] = velBits[i]
  offset += 5

  # heading sin: 8 bits
  let hSin = int((sin(degToRad(frame.heading)) + 1.0) / 2.0 * 199.0)
  let hSinBits = toBits(clamp(hSin, 0, 199), 8)
  for i in 0..<8: result[offset + i] = hSinBits[i]
  offset += 8

  # heading cos: 8 bits
  let hCos = int((cos(degToRad(frame.heading)) + 1.0) / 2.0 * 199.0)
  let hCosBits = toBits(clamp(hCos, 0, 199), 8)
  for i in 0..<8: result[offset + i] = hCosBits[i]
  offset += 8

  # enemy wall distances: 4×7 bits
  for wall in [frame.enemyWallN, frame.enemyWallS, frame.enemyWallE, frame.enemyWallW]:
    let pct = int(clamp(wall / 1000.0 * 99.0, 0.0, 99.0))
    let bits = toBits(pct, 7)
    for i in 0..<7: result[offset + i] = bits[i]
    offset += 7

  # enemy energy: 11 bits
  let enE = int(clamp(frame.enemyEnergy * 10.0, 0.0, 1500.0))
  let enEBits = toBits(enE, 11)
  for i in 0..<11: result[offset + i] = enEBits[i]
  # offset += 11  # last field, no need

proc encodeSelf*(self: SelfState): array[SELF_BITS, uint8] =
  var offset = 0

  # my wall distances: 4×7 bits
  for wall in [self.myWallN, self.myWallS, self.myWallE, self.myWallW]:
    let pct = int(clamp(wall / 1000.0 * 99.0, 0.0, 99.0))
    let bits = toBits(pct, 7)
    for i in 0..<7: result[offset + i] = bits[i]
    offset += 7

  # my energy: 11 bits
  let myE = int(clamp(self.myEnergy * 10.0, 0.0, 1500.0))
  let myEBits = toBits(myE, 11)
  for i in 0..<11: result[offset + i] = myEBits[i]
  offset += 11

  # canFire: 1 bit
  result[offset] = if self.canFire: 1'u8 else: 0'u8

proc encodeFullVector*(window: array[WINDOW_SIZE, array[FRAME_BITS, uint8]],
                       self: array[SELF_BITS, uint8]): BinaryVector =
  var offset = 0
  # LIFO: index 0 = newest frame
  for i in 0..<WINDOW_SIZE:
    for j in 0..<FRAME_BITS:
      result[offset] = window[i][j]
      inc offset
  for j in 0..<SELF_BITS:
    result[offset] = self[j]
    inc offset

proc formatBinary*(vec: BinaryVector): string =
  for b in vec:
    result &= (if b == 1: "1" else: "0")

proc popcount*(vec: BinaryVector): int =
  for b in vec: result += int(b)

proc hammingDistance*(a, b: BinaryVector): int =
  for i in 0..<TOTAL_BITS:
    result += int(a[i] xor b[i])

proc bitwiseAnd*(a, b: BinaryVector): BinaryVector =
  for i in 0..<TOTAL_BITS:
    result[i] = a[i] and b[i]

proc formatVectorBinary*(vec: BinaryVector): string =
  formatBinary(vec)

proc formatVectorDecimal*(vec: BinaryVector): string =
  # Placeholder: output popcount or some summary per frame
  result = ""
  for i in 0..<WINDOW_SIZE:
    var cnt = 0
    for j in 0..<FRAME_BITS:
      cnt += int(vec[i * FRAME_BITS + j])
    if i > 0: result &= ","
    result &= $cnt

proc formatFrameBinary*(frame: array[FRAME_BITS, uint8]): string =
  for b in frame:
    result &= (if b == 1: "1" else: "0")

proc formatFrameDecimal*(frame: array[FRAME_BITS, uint8]): string =
  # Nine decimal fields from the frame bits
  # bearing_sin (0-199), bearing_cos (0-199), distance (0-99), velocity (0-15), heading_sin (0-199), heading_cos (0-199), wall (0-99), wall (0-99), energy (0-1500/10)
  result = "199,199,99,15,199,199,99,99,150"  # placeholder max values

proc decodeFrameFields*(frame: array[FRAME_BITS, uint8]): array[9, float] =
  # Decode frame to 9 decimal field values (enemy X, Y for field indices 6, 7)
  # This is a placeholder—actual decoding would reverse the toBits() encoding
  result = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 50.0, 50.0, 100.0]

# Output encoding -------------------------------------------------------

const
  OUTPUT_BITS* = 7
  AIM_MIN*     = -60.0
  AIM_MAX*     =  60.0

type OutputVector* = array[OUTPUT_BITS, uint8]

proc encodeOutput*(angle: float): OutputVector =
  let clamped = clamp(angle, AIM_MIN, AIM_MAX)
  let intVal  = int((clamped - AIM_MIN) / (AIM_MAX - AIM_MIN) * 127.0)
  let gray    = toGray(intVal)
  for i in 0..<OUTPUT_BITS:
    result[OUTPUT_BITS - 1 - i] = uint8((gray shr i) and 1)

proc decodeOutput*(vec: OutputVector): float =
  var val = 0
  for i in 0..<OUTPUT_BITS:
    val = (val shl 1) or int(vec[i])
  var binary = val
  var mask   = binary shr 1
  while mask > 0:
    binary = binary xor mask
    mask   = mask shr 1
  result = AIM_MIN + (float(binary) / 127.0) * (AIM_MAX - AIM_MIN)