import math, options, sequtils, sets, sugar
import combinators, game, fixedseq


type
  ScoreSet* = array[Color, int]

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


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


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

  var diceRemaining: ColorStack
  diceRemaining.initFixedSeq
  for i, c in b.diceRolled:
    if not c: diceRemaining.add(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:
    echo "simulating ", outcomeStack[^1].len, " possible legs."
    var endStates: HashSet[Board]

    for o in outcomeStack[^1]:
      var o = o # make it mutable
      o.resetDice # o was describina an end-of-leg state, so dice were exhausted

      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.init
b.display(1, 5)

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

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

# var s: ColorStack
# s.initFixedSeq
# for i in 0..4:
#   s.add(Color(i))

# echo s
# echo s[2]

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