import bitops, strutils, random


proc show(i: SomeInteger, bitlength = 16) =
  echo BiggestInt(i).toBin(bitlength)


type
  Color* = enum
    cRed, cGreen, cBlue, cYellow, cPurple

  ColorStack* = object
    data: uint16
    len*: uint8


const
  masks = [
    0'u16, # dummy value just to get the indices right
    0b0_000_000_000_000_111,
    0b0_000_000_000_111_111,
    0b0_000_000_111_111_111,
    0b0_000_111_111_111_111,
    0b0_111_111_111_111_111,
  ]
  allColors* = ColorStack(len: 5, data: 0b0_000_001_010_011_100)


template offset(s: ColorStack, idx: Natural): uint8 =
  # Compute the bit offset for a given index. 
  # Dependent on the stack's length.
  (s.len - 1 - idx.uint8) * 3
  # items are stored from left to right but right-aligned, so we
  # need to shift right to access anything other than the last


template offset(s: ColorStack, idx: BackwardsIndex): uint8 =
  # for backwards index, we are still shifting right but
  # the lower the index the less we have to shift
  (idx.uint8 - 1) * 3


proc add*(s: var ColorStack, c: Color) =
  # e.g. if stack is 0b0000000000000100:
  # and color is 0b00000011
  # shift: 0b0000000000100000
  # bitor: 0b0000000000100000 and 0b00000011
  # results in 0b0000000000100011
  s.data = (s.data shl 3).bitor(cast[uint8](c))
  inc s.len


proc high*(s: ColorStack): uint8 = s.len - 1

proc low*(s: ColorStack): uint8 = 0 # just... always 0, I guess


proc `[]`*(s: ColorStack, i: uint8 | BackwardsIndex): Color =
  # shift, then mask everything but the three rightmost bits
  result = Color(
    (s.data shr s.offset(i)) and masks[1]
  )


proc `[]=`*(s: var ColorStack, i: uint8 | BackwardsIndex, c: Color) =
  let offset = s.offset(i)
  s.data = (s.data and bitnot(masks[1] shl offset)) or (c.uint16 shl offset)


iterator items*(s: ColorStack): Color =
  # s.len is unsigned so it will wrap around if we do s.len - 1 in that case
  if s.len != 0:
    for i in countdown(s.len - 1, 0'u8):
      yield Color((s.data shr (i * 3)) and masks[1])


iterator pairs*(s: ColorStack): (uint8, Color) =
  var count = 0'u8
  for color in s:
    yield (count, color)
    inc count


proc find*(s: ColorStack, needle: Color): int8 =
  for i in 0'u8 .. s.high:
    if s[i] == needle:
      return i.int8
  return -1


proc moveSubstack*(src, dst: var ColorStack, startIdx: uint8) =
  if startIdx >= src.len:
    raise newException(IndexDefect, "index " & $startIdx & " is out of bounds.")
  # Moves a sub-stack from the top of src to the top of dst
  # shift the dst stack by the length of the substack to make room
  let nToMove = src.len - startIdx
  let shift = nToMove * 3
  dst.data = dst.data shl shift
  # then we mask the source data to present only the items 
  # being moved, and OR that with the shifted dst data
  dst.data = dst.data or (src.data and masks[nToMove])
  dst.len += nToMove
  # then we shift the source to get rid of the moved items
  src.data = src.data shr shift
  src.len -= nToMove


proc moveSubstackPre*(src, dst: var ColorStack, startIdx: uint8) =
  if startIdx >= src.len:
    raise newException(IndexDefect, "index " & $startIdx & " is out of bounds.")
  # Moves a sub-stack from the top of src to the bottom of dst
  let nToMove = src.len - startIdx
  # shift src to position the substack above its destination,
  # get rid of everything to the left of the substack,
  # and OR that with the existing dst data
  let newLen = dst.len + nToMove
  dst.data = dst.data or ( (src.data shl (dst.len * 3)) and masks[newLen] )
  dst.len = newLen
  # get rid of the substack we just moved
  src.data = src.data shr (nToMove * 3)
  src.len -= nToMove


proc swap*(s: var ColorStack, i1, i2: uint8) =
  # Swap the values at two indices in the stack
  if i1 == i2: return

  # i1 and i2 are unsigned, so we have to watch out for underflows
  let diff = if i1 > i2:
    (i1 - i2) * 3
  else:
    (i2 - i1) * 3

  # take masks[1] from above (rightmost position) and shift to position of i1.
  # then do the same for i2, and OR them together.
  let mask = (masks[1] shl s.offset(i1)) or (masks[1] shl s.offset(i2))
  # get rid of everything but the two values we're swapping
  let masked = s.data and mask
  # shift by the distance between values in both directions, combine, then mask
  let swapped = ((masked shl diff) or (masked shr diff)) and mask
  # finally, AND with the inverse of mask so that only the values being 
  # swapped are erased, and combine that with the swapped values
  s.data = (s.data and mask.bitnot) or swapped


proc shuffle*(r: var Rand, s: var ColorStack) =
  # Fisher-Yates shuffle
  for i in countdown(s.high, 1'u8):
    let j = r.rand(i).uint8
    if j != i:
      s.swap(i, j)


proc reverse*(s: var ColorStack, first, last: uint8) =
  var x = first
  var y = last
  while x < y:
    s.swap(x, y)
    inc x
    dec y


iterator asInt*(s: ColorStack): int8 =
  for i in 0'u8 .. s.high:
    yield int8(s[i]) # now we do have to convert


proc `$`*(s: ColorStack): string =
  result = "St@["
  for c in s:
    if result[^1] != '[':
      result.add(", ")
    result.add($c)
  result.add("]")


proc check(s: ColorStack) =
  # ensure length is accurate
  var d = s.data
  for i in 0'u8 .. 4'u8:
    if (d and masks[1]) > 4:
      raise newException(RangeDefect, "Value out of range.")
    if d > 0 and i >= s.len:
      raise newException(RangeDefect, "Invalid length.")
    else:
      d = d shr 3


when isMainModule:
  var one: ColorStack
  one.add(cRed)
  one.add(cGreen)
  one.add(cBlue)
  one.add(cYellow)
  one.add(cPurple)

  var two: ColorStack

  one.moveSubstack(two, 2)


  echo one, " ", one.len
  echo two, " ", two.len
  echo two.find(cRed)