package pl.edu.pw.mini.sg.mcts

import scala.annotation.tailrec

import pl.edu.pw.mini.sg.rng.RandomGenerator
import pl.edu.pw.mini.sg.game.deterministic.single.SinglePlayerGame


final case class NestedMcsConfig (
  depth: Int,
  samples: Int
)

object NestedMcs {

  @tailrec
  def randomTrial[M, S, G <: SinglePlayerGame[M, S, G]](
    game: G,
    moves: Vector[M] = Vector()
  )(implicit rng: RandomGenerator) : (G, Vector[M]) = {
    if(game.possibleMoves.length == 0) (game, moves)
    else {
      val mm = randomMove[M, S, G](game)
      randomTrial[M, S, G](game.makeMove(mm), moves :+ mm)
    }
  }

  def randomMove[M, S, G <: SinglePlayerGame[M, S, G]](
    game: G,
  )(implicit rng: RandomGenerator) : M = {
    game.possibleMoves(rng.nextInt(game.possibleMoves.length))
  }

  def findBestPath[M, S, G <: SinglePlayerGame[M, S, G]](config: NestedMcsConfig)(
    game: G,
    stateEvaluator: G => Double,
    level: Int,
    moveSeq: Vector[M]
  )(implicit rng: RandomGenerator) : (Double, G, Vector[M]) = {
    if (game.possibleMoves.length == 0) {
      (stateEvaluator(game), game, moveSeq)
    } else {
      game.possibleMoves.map(mv => {
        if(level == config.depth) {
          val gg = game.makeMove(mv)
          val (ev, res, moves) = Range(0,config.samples).map(_ =>  randomTrial[M, S, G](gg)).map{case (a, b) => (stateEvaluator(a), a, b)}.maxBy(_._1)
          (ev, res, (moveSeq:+mv)++moves)
        } else {
          findBestPath[M,S,G](config)(game.makeMove(mv), stateEvaluator, level+1, moveSeq:+mv)
        }
      }).maxBy(_._1)
    }
  }
}
