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


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

  Tile = enum
    tBackward = -1,
    tForward = 1

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

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

  ScoreSet = array[Color, int]

  LegResults = tuple[scores: ScoreSet, endStates: HashSet[Board]]

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


const allDice = @[cRed, cGreen, cBlue, cYellow, cPurple]


proc update(scores: var ScoreSet, toAdd: ScoreSet) =
  for i, s in toAdd:
    scores[i] += s


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


proc hash(b: Board): Hash =
  var h: Hash = 0
  # there could be a tile anywhere so we have to check all squares
  for i, sq in b.squares:
    if sq.camels.len > 0 or sq.tile.isSome:
      h = h !& i
      if sq.tile.isSome:
        h = h !& int(sq.tile.get) * 10 # so it isn't confused with a camel
      else:
        for c in sq.camels:
          h = h !& int(c)
  result = !$h


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 setState(b: var Board; 
              camels: openArray[tuple[c: Color, p: int]];
              tiles: openArray[tuple[t: Tile, p: int]]) =
  for (color, dest) in camels: # note that `camels` is ordered, as this determines stacking
    b[dest].camels.add(color)
    b.camels[color] = dest

  for (tile, dest) in tiles:
    b[dest].tile = some(tile)

  let leadCamel = b[max(b.camels)].camels[^1] # top camel in the last currently-occupied space
  b.leader = some(leadCamel)


proc advance(b: var Board, die: Die) =
  let
    (color, roll) = die
    startPos = b.camels[color]
  var endPos = startPos + roll
  
  if endPos > 16: # camel has passed the finish line
    b.leader = some(b[startPos].camels[^1])
    b.gameOver = true
    return

  var prepend = false
  if b[endPos].tile.isSome: # adjust position (and possibly stacking) to account for tile
    let t = b[endPos].tile.get
    endPos += int(t)
    if t == tBackward: prepend = true

  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
      # if we are stacking on or moving past the previous leader
      if endPos >= b.camels[b.leader.get]:
        b.leader = some(b[endPos].camels[^1])
      break # breaking the outer loop here, not the inner - but only conditionally! gah!
  b.diceRolled[color] = true


proc projectLeg(b: Board): LegResults =
  var scores: ScoreSet
  var endStates: HashSet[Board]

  let diceRemaining = collect(newSeq):
    for i, c in b.diceRolled:
      if not c: i

  for future in possibleFutures(diceRemaining):
    var prediction = b # make a copy
    for dieRoll in future:
      prediction.advance(dieRoll)
    inc scores[prediction.leader.get]
    # deduplicate results
    endStates.incl(prediction)
  result = (scores, endStates)


proc projectOutcomes(b: Board, maxDepth = 1): ScoreSet =
  var outcomeStack = @[ [b].toHashSet ]
  for depth in 1..maxDepth:
    var endStates: HashSet[Board]
    for o in outcomeStack[^1]:
      let projection = o.projectLeg
      result.update(projection[0])
      endStates.incl(projection[1])
      stdout.write("simulated: " & $result.sum & "\r")

    outcomeStack.add(endStates)
  echo "\nDistinct end states: ", outcomeStack.mapIt(it.len).sum


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

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

let r = b.projectOutcomes(2)
let total = r.sum
for i, c in r:
  echo Color(i), ": ", (100 * c / total).round(2), "% (", c, " / ", total, ")"