import input_type
import map_type
import turtle
import turtle_type

import math
import sequtils
import strutils
import strformat
import terminal
import algorithm

const blueCutoff = 0.05
const cyanCutoff = 0.1
const greenCutoff = 0.205
const yellowCutoff = 0.335
const redCutoff = 0.505
const magentaCutoff = 1.0

proc setHotnessBackground(hotness: float, maxHotness: float) =
  if hotness == 0:
    stdout.setBackgroundColor(bgDefault)
  else:
    let hotnessFrac = hotness / maxHotness
    if hotnessFrac < blueCutoff:
      stdout.setBackgroundColor(bgBlue)
    elif hotnessFrac < cyanCutoff:
      stdout.setBackgroundColor(bgCyan)
    elif hotnessFrac < greenCutoff:
      stdout.setBackgroundColor(bgGreen)
    elif hotnessFrac < yellowCutoff:
      stdout.setBackgroundColor(bgYellow)
    elif hotnessFrac < redCutoff:
      stdout.setBackgroundColor(bgRed)
    else:
      stdout.setBackgroundColor(bgMagenta)

proc printHotmapLegend(maxHotness: float, x,y: int) =
  template line(num: int) =  
    setCursorPos(x, y + num - 1)
  template showColor(col: BackgroundColor) =
    stdout.setBackgroundColor(col)
    stdout.write(" ")
    stdout.setBackgroundColor(bgBlack)
    stdout.write(" ")
  
  proc writeNumber(num: float) =
    stdout.write(fmt"{num.int()}")

  line(1)
  showColor(bgDefault)
  writeNumber(0.0)
  
  line(2)
  showColor(bgBlue)
  writeNumber(blueCutoff * maxHotness)

  line(3)
  showColor(bgCyan)
  writeNumber(cyanCutoff * maxHotness)

  line(4)
  showColor(bgGreen)
  writeNumber(greenCutoff * maxHotness)

  line(5)
  showColor(bgYellow)
  writeNumber(yellowCutoff * maxHotness)

  line(6)
  showColor(bgRed)
  writeNumber(redCutoff * maxHotness)

  line(7)
  showColor(bgMagenta)
  writeNumber(magentaCutoff * maxHotness)
  
  

proc displayGeneration*(generationNumber: int, timeSeconds: float, seed: int, input: Input, map: Map, turtles: var seq[Turtle]) =
  
  # SCREEN INIT
  stdout.setBackgroundColor(bgBlack)
  stdout.setForegroundColor(fgWhite)
  eraseScreen()
  setCursorPos(0, 0)

  # STATISTICS
  let scoreCount = turtles.len()
  var scoreSum = 0
  var stepsSum = 0
  for turtle in turtles:
    scoreSum += turtle.score
    stepsSum += turtle.steps
  
  let scoreAvg = scoreSum.float() / scoreCount.float()
  let stepsAvg = stepsSum.float() / scoreCount.float()

  let variance = turtles.mapIt((it.score.float() - scoreAvg)^2).sum() / scoreCount.float()
  let stddev = variance.sqrt()

  # PRINT BASIC INFO
  echo &"Generation number: {generationNumber}"
  echo &"Computation time since last stop: {timeSeconds} s"
  echo &"Seed: {seed}"
  echo &"Population: {turtles.len()} turtles"
  echo &"Mean score: {scoreAvg} (in {stepsAvg} steps)"
  echo &"Standard deviation: {stddev}"
  echo &"Top score: {turtles[0].score} (in {turtles[0].steps} steps)"

  # PERCENTILS
  turtles.sort(compareTurtles)
  var percentils: seq[int] = @[]
  for i in 1..9:
    let index = (input.population div 10) * i
    percentils.add(turtles[index].score)
  let percentilsString = percentils.join(", ")
  echo &"Score of every 10%th percentil: {percentilsString}"

  # CHOOSE NOTABLE
  var top10Paths: seq[Turtle] = @[]
  var notables: seq[int] = @[]
  var i = 0
  while true:
    if i > turtles.high():
      break

    if top10Paths.len() == 0 or top10Paths[top10Paths.high()].score != turtles[i].score:
      top10Paths.add(turtles[i])
      notables.add(i)
    
    if top10Paths.len() >= 10:
      break

    inc i
  
  let notableString = notables.join(", ")
  echo &"Notable turtles: {notableString}"

  # HOTNESS MAP
  var hotness: seq[seq[int]] = @[]
  var top10Hot: seq[seq[int]] = @[]
  var topHot: seq[seq[int]] = @[]
  for i in 0..input.size-1:
    hotness.add(@[])
    top10Hot.add(@[])
    topHot.add(@[])
    for j in 0..input.size-1:
      hotness[i].add(0)
      top10Hot[i].add(0)
      topHot[i].add(0)
  
  template populateHotness(turtles: seq[Turtle], hotSeq: var seq[seq[int]], maxHot: var int) =
    for turtle in turtles:
      let path = turtle.path
      for i in 0..path.xs.high():
        let x = path.xs[i]
        let y = path.ys[i]
        hotSeq[y][x].inc()
        if hotSeq[y][x] > maxHot:
          maxHot = hotSeq[y][x]

  var maxHotness = 0
  populateHotness(turtles, hotness, maxHotness)
  var maxTop10Hotness = 0
  populateHotness(top10Paths, top10Hot, maxTop10Hotness)
  var maxTopHotness = 0
  populateHotness(@[turtles[0]], topHot, maxTopHotness)

  # DISPLAY HOTNESS MAP
  echo "Hotmap of All/Notable/Best Turtles"

  const yOffset = 10
  let xOffset = input.size + 2

  for y in 0..input.size-1:
    for x in 0..input.size-1:
      let hotness = hotness[y][x]
      let top10hotness = top10Hot[y][x]
      let topHotness = topHot[y][x]
      let symbol = map[y][x].tileToString()

      # hotmap
      setHotnessBackground(hotness.float(), maxHotness.float())
      setCursorPos(x, y+yOffset)
      stdout.write(symbol)

      # top10 map
      setHotnessBackground(top10hotness.float(), maxTop10Hotness.float())
      setCursorPos(x+xOffset, y+yOffset)
      stdout.write(symbol)

      # top map
      setHotnessBackground(topHotness.float(), maxTopHotness.float())
      setCursorPos(x+2*xOffset, y+yOffset)
      stdout.write(symbol)

    
    stdout.setBackgroundColor(bgBlack)
    stdout.write("\n")

  echo ""

  # display legend
  let legendYOffset = yOffset + 1 + input.size
  printHotmapLegend(maxHotness.float(), 0, legendYOffset)
  printHotmapLegend(maxTop10Hotness.float(), xOffset, legendYOffset)
  printHotmapLegend(maxTopHotness.float(), 2*xOffset, legendYOffset)
  echo ""
  echo ""
  # FINISH
  echo "Press enter to continue."
  discard stdin.readLine()
  stdout.setBackgroundColor(bgDefault)
  eraseScreen()
  setCursorPos(1,1)
  echo ""