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
import pl.edu.pw.mini.sg.mcts.tree.MctsNode


object MctsAlgorithm {
  def training[M,S, G <: SinglePlayerGame[M,S, G]](
    root: MctsNode[M,S],
    rootState: G,
    leafEvaluationFunction: G => (Double, Option[(List[M], Double)]),
    parameters: MctsParameters,
    rng: RandomGenerator
  ) : (S, List[M], Double) = {
    require(root.state == rootState.playerState)

    // Potrzebuję móc zakazać pewnych ruchów (wystarczy tylko w korzeniu).
    // Jak to zrobić?
    // 1) Moge zrobić nakładke na grę - jak duży narzut na wydajność??
    // 2) 
    // Czy powinienem dodawać nowe węzły, gdy selection powie, że tak? 
    @tailrec
    def selectionPhase(state: G, node: MctsNode[M, S], path: List[(MctsNode[M, S], M)]) : (List[(MctsNode[M, S], M)], MctsNode[M, S], G) = {
      parameters.mctsSelection(node) match {
        case Some(m) => {
          val newState = state.makeMove(m)
          val newPath = (node, m) :: path
          val newNode = node.successor(m, newState.playerState) match {
            case Some(n) => n
            case None => node.addSuccessor(m, newState.playerState, newState.possibleMoves.toVector)
          }

          selectionPhase(newState, newNode, newPath)
        }
        case None => (path, node, state)
      }
    }

    def expansionPhase(state: G, node: MctsNode[M, S], move: M) : (List[(MctsNode[M, S], M)], MctsNode[M, S], G) = {
      val newState = state.makeMove(move)
      val newNode = node.addSuccessor(move, newState.playerState, newState.possibleMoves.toVector)
      (List((node, move)), newNode, newState)
    }

    @tailrec
    def mcPhase(state: G, moves: List[M]) : (G, List[M]) = {
      state.possibleMoves match {
        case Nil => (state, moves)
        case _ => {
          val move = parameters.mcPhaseMove(state, rng)
          mcPhase(state.makeMove(move), move::moves)
        }
      }
    }

    def simulation() : (G, List[M], Double) = {
      val (selectionPath, selLeaf, selLeafState) = selectionPhase(rootState, root, Nil)
      val (expandedPath, expNode, expLeafState) =
        if(selLeafState.possibleMoves.isEmpty)
          (Nil, selLeaf, selLeafState)
        else 
          expansionPhase(selLeafState, selLeaf, parameters.mcPhaseMove.apply(selLeafState, selLeaf.unvisitedMoves, rng))
      val (mcLeafState, moves) = mcPhase(expLeafState, Nil)

      val (reward, forcedUpdate) = leafEvaluationFunction(mcLeafState)
      parameters.propagation.propagate(reward, expandedPath ++ selectionPath)
      forcedUpdate match {
        case None => ()
        case Some((moves, payoff)) => forcedSimulation(rootState, moves, root, payoff)
      }
      (mcLeafState, moves ++ expandedPath.map(_._2) ++ selectionPath.map(_._2), reward)
    }

    def forcedSimulation(game: G, moves: List[M], root: MctsNode[M, S], result: Double) : Unit = {
      val nodesToUpdate = getNodeList(game, root, moves, List())
      parameters.propagation.propagate(result, nodesToUpdate)
    }

    @tailrec
    def getNodeList(game: G, root: MctsNode[M, S], moves: List[M], result: List[(MctsNode[M,S], M)]): List[(MctsNode[M, S], M)] =
      moves match {
        case m::tail => {
          val next = game.makeMove(m)
          root.successor(m, next.playerState) match {
            case None => root.addSuccessor(m, next.playerState, next.possibleMoves.toVector); (root,m)::result
            case Some(n) => getNodeList(next, n, tail, (root,m)::result)
          }
        }
        case List() => result
      }

    val ord = Ordering.fromLessThan[(G, List[M], Double)](_._3 < _._3)

    @tailrec
    def training(prevRes: Option[(G, List[M], Double)]) : Option[(G, List[M], Double)] = {
      if(root.visitCounter >= parameters.simulationsInRoot) {
        prevRes
      } else {
        val newRes = simulation()
        training(prevRes.map(ord.max(_, newRes)).orElse(Some(newRes)))
      }
    }

    val t  = training(None).getOrElse(simulation())
    (t._1.playerState, t._2, t._3)
  }
}
