import cpuinfo, math, options, random, tables
import combinators, game, faststack, fixedseq


type
  ScoreSet* = array[Color, int]
  WinPercents* = array[Color, float]

  ScoreSpread = object
    lo*: array[Color, float]
    hi*: array[Color, float]

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


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


proc display*(scores: ScoreSet) =
  let total = scores.sum
  for color, score in scores:
    let line = $color & ": " & $round(100 * scores[color] / total, 2) & '%'
    stdout.writeLine(line)
    stdout.flushFile()
    # echo color, ": ", round(100 * scores[color] / total, 2), '%'


proc percents*(scores: ScoreSet): WinPercents =
  let total = scores.sum
  for c, score in scores:
    result[c] = score / total


# ======================
# Single-leg simulations
# ======================

iterator legEndStates(b: Board): Board =
  var diceRemaining: FixedSeq[5, Color, int8]
  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 so we can mutate
    for dieRoll in future:
      prediction.advance(dieRoll)
    yield prediction


proc getLegScores*(b: Board): ScoreSet =
  for prediction in b.legEndStates:
    inc result[prediction.leader.get]


# =====================
# Full-game simulations
# =====================

# get rid of this later
import strutils
proc showSpaces*(b: Board; start, stop: Natural): string =
  let numSpaces = stop - start + 1
  let width = 4 * numSpaces - 1
  var lines: array[7, string]
  # start by building up an empty board
  for i in 0 .. 6: # gotta initialize the strings
    lines[i] = newString(width)
    for c in lines[i].mitems:
      c = ' '
  # fill in the dividers
  lines[5] = repeat("=== ", numSpaces - 1)
  lines[5].add("===")

  # now populate the board
  for sp in 0 ..< numSpaces:
    # fill in the square numbers
    let squareNum = sp + start
    let cellMid = 4 * sp + 1
    for i, chr in $squareNum:
      lines[6][cellMid + i] = chr

    # fill in the camel stacks
    for i, color in b.squares[squareNum].camels:
      let lineNum = 4 - i # lines go to 6, but bottom 2 are reserved
      let repr = '|' & color.abbrev & '|'
      for j, chr in repr:
        lines[lineNum][cellMid - 1 + j] = chr

  result = lines.join("\n")
# get rid of this later

proc randomGame*(b: Board, r: var Rand): Color =
  var projection = b
  while true:
    for roll in randomFuture(projection.diceRemaining, r):
      projection.advance(roll)
      if projection.gameOver:
        return projection.leader.get
    projection.resetDice()


proc randomGamesWorker(b: Board, count: Natural, r: var Rand): ScoreSet =
  for i in 1 .. count:
    let winner = b.randomGame(r)
    inc result[winner]


# =======================
# Multithreading nonsense
# =======================

type WorkerArgs = object
  board: Board
  count: Natural
  seed: int64


# have to do this at the module level so it can be shared
var gamesChannel: Channel[ScoreSet]
gamesChannel.open()


proc randomGamesThread(args: WorkerArgs) =
  var r = initRand(args.seed)
  let scores = randomGamesWorker(args.board, args.count, r)
  gamesChannel.send(scores)


proc randomGames*(b: Board, count: Natural, parallel = true, numThreads = 0): ScoreSet =
  randomize()

  if not parallel:
    var r = initRand(rand(int64))
    return randomGamesWorker(b, count, r)

  let numThreads =
    if numThreads == 0:
      countProcessors()
    else:
      numThreads

  var workers = newSeq[Thread[WorkerArgs]](numThreads)
  for i, w in workers.mpairs:
    var numGames = int(floor(count / numThreads))
    if i < (count mod numThreads):
      numGames += 1
    let args = WorkerArgs(board: b, count: numGames, seed: rand(int64))

    createThread(w, randomGamesThread, args)

  for i in 1 .. numThreads:
    let scores = gamesChannel.recv()
    result.update(scores)


proc randomSpread*(b: Board; nTests, nSamples: Natural): ScoreSpread =
  for s in result.lo.mitems:
    s = 1

  for i in 0 ..< nTests:
    let scores = b.randomGames(nSamples)
    let total = scores.sum
    for color, score in scores:
      let pct = score / total
      if pct < result.lo[color]:
        result.lo[color] = pct
      if pct > result.hi[color]:
        result.hi[color] = pct