import math, options, tables, sequtils, sets, sugar
import combinators


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

  Tile = enum
    tForward = -1,
    tBackward = 1

  Square = object
    camels: seq[Color]
    tile: Option[Tile]

  Die = tuple[color: Color, value: int]

  ScoreSet = array[Color, int]

  Board = object
    squares: array[1..16, Square]
    camels: array[Color, range[1..16]]
    leader: Color
    gameOver: bool


proc `[]`[T](b: var Board, idx: T): var Square =
  b.squares[idx]


proc display(b: Board, start, stop: int) =
  for i in start..stop:
    let sq = b.squares[i]
    let lead = $i & ": "
    if sq.tile.isSome:
      echo lead, sq.tile.get
    else:
      echo lead, sq.camels
  echo ""


proc setLeader(b: var Board) =
  b.leader = b[b.camels.max].camels[^1] # top camel in the last currently-occupied space


proc setState[T](b: var Board, state: T) =
  for (color, dest) in state: # note that `state` is ordered, as this determines stacking
    b[dest].camels.add(color)
    b.camels[color] = dest
  b.setLeader


proc advance(b: var Board, die: Die) =
  let
    (color, roll) = die
    startPos = b.camels[color]
  var prepend = false
  var endPos = startPos + roll
  
  if endPos > 16: # camel has passed the finish line
    b.leader = color
    b.gameOver = true
    return
  if b[endPos].tile.isSome:
    endPos += int(b[endPos].tile.get)
    if b[endPos].tile.get == tBackward: prepend = true # if moving backward, we will prepend camel to the seq

  for i, c in b[startPos].camels:
    if c == color:
      let subStack = b[startPos].camels[i .. ^1]
      if prepend:
        b[endPos].camels.insert(subStack)
      else:
        b[endPos].camels.add(subStack)

      b[startPos].camels[i .. ^1] = @[]
      for moved in subStack:
        b.camels[moved] = endPos
      b.setLeader
      break # breaking the outer loop here, not the inner - but only conditionally! gah!


type CamelPositions = seq[tuple[square: int, camels: seq[Color]]]

proc simulateLeg(b: Board, diceRemaining: seq[Color]): ScoreSet =
  var endStates: HashSet[Board]
  for future in possibleFutures(diceRemaining):
    var prediction = b # make a copy
    for dieRoll in future:
      prediction.advance(dieRoll)
    inc result[prediction.leader]
    # deduplicate results
    endStates.incl(prediction) # need to add hash implementation for Board
    #[
    var positions: CamelPositions
    for i in prediction.camels.min .. prediction.camels.max:
      positions.add((i, prediction.squares[i].camels))
    endStates.incl(positions)
    ]#
  echo "Distinct end states: ", endStates.len


var b: Board
b.display(1, 5)

b.setState({cGreen: 4, cYellow: 3, cPurple: 4, cBlue: 3, cRed: 2})
b.display(1, 5)

b.advance((cRed, 2))
b.display(1, 5)

let r = b.simulateLeg(@[cRed, cGreen, cBlue, cYellow, cPurple])
let total = r.sum
for i, c in r:
  echo Color(i), ": ", (100 * c / total).round(2), "% (", c, " / ", total, ")"