package pl.edu.pw.mini.sg.mixed.expansion

import pl.edu.pw.mini.sg.game.Payoff
import pl.edu.pw.mini.sg.game.boundedrationality.BoundedRationality
import pl.edu.pw.mini.sg.game.deterministic.PureStrategy
import pl.edu.pw.mini.sg.game.deterministic.single.DummyPureStrategy
import pl.edu.pw.mini.sg.game.util.MoveStateKey
import pl.edu.pw.mini.sg.mcts.MctsAlgorithm

import pl.edu.pw.mini.sg.rng.RandomGenerator
import pl.edu.pw.mini.sg.util.Logging
import pl.edu.pw.mini.sg.game.deterministic.PlayerVisibleState
import pl.edu.pw.mini.sg.game.deterministic.single.{GameWithPayoff, SinglePlayerGame}
import pl.edu.pw.mini.sg.mcts.MctsParameters
import pl.edu.pw.mini.sg.mixed.expander.NegativeAttacker
import pl.edu.pw.mini.sg.mixed.{BothTreesPosition, IterationInfo}
import pl.edu.pw.mini.sg.mixed.expander.{MctsExpander, NegativeAttacker, PositiveDefender}
import pl.edu.pw.mini.sg.mixed.tree.{InnerResultNode, MixedStrategyNode}
import pl.edu.pw.mini.sg.game.deterministic.DeterministicGame
import scala.annotation.tailrec



final case class MctsNodeExpansion(
  mctsConfig: MctsParameters,
  expandWorsePayoff: Boolean,
  expansionProbability: Double
) extends NodeExpansionDecision with Logging {

  val initialExpander: MctsExpander = MctsExpander(mctsConfig, NegativeAttacker)
  val negativeExpander: MctsExpander = MctsExpander(mctsConfig, NegativeAttacker)
  val positiveExpander: MctsExpander = MctsExpander(mctsConfig, PositiveDefender)


  def inTreeExpansion[DM, AM, DS <: PlayerVisibleState[DM], AS <: PlayerVisibleState[AM], G <: DeterministicGame[DM, AM, DS, AS, G]](
    bothTreesPosition: BothTreesPosition[DM,AM,DS,AS,G],
    defenderMixedStrategy: MixedStrategyNode[DM, DS, _],
    desiredAttackerPayoff: InnerResultNode[DM,DS],
    undesiredAttackerPayoff: InnerResultNode[DM,DS],
    noSubtreeExpanded: Boolean
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo, br: BoundedRationality
  ) : Boolean = {
    if(defenderMixedStrategy.probabilities.length < bothTreesPosition.desiredAttacker.possibleMoves.length && rng.nextDouble() < expansionProbability) {
      val doubleGame : PayoffDifferenceGame[DM, AM, DS, AS, G] = PayoffDifferenceGameBoth[DM, AM, DS, AS, G](bothTreesPosition.desiredState, bothTreesPosition.optimalState, bothTreesPosition.desiredAttacker.attackerStrategy, bothTreesPosition.optimalAttacker.attackerStrategy)
      val root = mctsConfig.treeFactory.create(doubleGame.playerState, doubleGame.possibleMoves.filter(m => defenderMixedStrategy.probabilities.find(_._1 == m.m).isEmpty))
      val (_, bestMoves, bestPayoff): (DS, List[MoveType[DM]], Double) = MctsAlgorithm.training(root.root, doubleGame, (g : GameWithPayoff) => (g.attackerPayoff, None), mctsConfig, rng)
      if(bestPayoff > desiredAttackerPayoff.payoff.distortedAttacker.attacker - undesiredAttackerPayoff.payoff.distortedAttacker.attacker || expandWorsePayoff) {
        val desiredMoves: List[DM] = bestMoves.reverse.collect{case Both(m) => m case Desired(m) => m}
        val betterMoves: List[DM] = bestMoves.reverse.collect{case Both(m) => m case Better(m) => m}
        require(desiredMoves.head == betterMoves.head)
        defenderMixedStrategy.addMove(desiredMoves.head, 0.0)
        val nextDesiredState = bothTreesPosition.desiredState.makeMove(desiredMoves.head, bothTreesPosition.desiredAttacker.attackerStrategy.move)
        val nextDesiredStrategy = bothTreesPosition.desiredAttacker.attackerStrategy.nextState(nextDesiredState.attackerState)
        val nextBetterState = bothTreesPosition.optimalState.makeMove(desiredMoves.head, bothTreesPosition.optimalAttacker.attackerStrategy.move)
        val nextBetterStrategy = bothTreesPosition.optimalAttacker.attackerStrategy.nextState(nextBetterState.attackerState)

        @tailrec
        def expandByList(currentState: G, attackerStrategy: PureStrategy[AM, AS],
          currentMS: MixedStrategyNode[DM, DS, _], moves: List[DM]) : Unit = {
          moves match {
            case Nil => require(currentState.isLeaf)
            case m::ms => {
              if(currentMS.probabilities.isEmpty)
                currentMS.addMove(m, 1.0)
              val nextState = currentState.makeMove(m, attackerStrategy.move)
              val nextStrategy = attackerStrategy.nextState(nextState.attackerState)
              if(currentMS.succesors.get(m).flatMap(_.get(nextState.defenderState)).isEmpty) {
                currentMS.addEmptySuccessor(new MoveStateKey(m, nextState.defenderState))
              }
              val nextMs = currentMS.nextStateInTree(m, nextState.defenderState)
              expandByList(nextState, nextStrategy, nextMs, ms)
            }
          }
        }



        defenderMixedStrategy.addEmptySuccessor(new MoveStateKey(desiredMoves.head, nextDesiredState.defenderState))
        if(nextDesiredState.defenderState != nextBetterState.defenderState)
          defenderMixedStrategy.addEmptySuccessor(new MoveStateKey(desiredMoves.head, nextBetterState.defenderState))
        expandByList(nextBetterState, nextBetterStrategy,
          defenderMixedStrategy.nextStateInTree(desiredMoves.head, nextBetterState.defenderState), betterMoves.tail)
        expandByList(nextDesiredState, nextDesiredStrategy,
          defenderMixedStrategy.nextStateInTree(desiredMoves.head, nextDesiredState.defenderState), desiredMoves.tail)

        desiredAttackerPayoff.expandToLeaf(bothTreesPosition.desiredAttacker, DummyExpander)
        undesiredAttackerPayoff.expandToLeaf(bothTreesPosition.optimalAttacker, DummyExpander)
        require(desiredAttackerPayoff.successors.contains(desiredMoves.head))
        true
      } else {
        false
      }
    } else {
      false
    }
  }

  @tailrec
  private def expandByList[DM, DS <: PlayerVisibleState[DM], G <: SinglePlayerGame[DM, DS, G]](currentState: G,
    currentMS: MixedStrategyNode[DM, DS, _], moves: List[DM])(
    implicit iterationInfo: IterationInfo
  ): Unit = {
    moves match {
      case Nil => require(currentState.possibleMoves.isEmpty)
      case m::ms => {
        if(currentMS.probabilities.isEmpty)
          currentMS.addMove(m, 1.0)
        val nextState = currentState.makeMove(m)
        if(currentMS.succesors.get(m).flatMap(_.get(nextState.playerState)).isEmpty) {
          currentMS.addEmptySuccessor(new MoveStateKey(m, nextState.playerState))
        }
        val nextMs = currentMS.nextStateInTree(m, nextState.playerState)
        expandByList(nextState, nextMs, ms)
      }
    }
  }



  def outTreeExpansion[DM, DS <: PlayerVisibleState[DM], G <: SinglePlayerGame[DM,DS,G] with GameWithPayoff](
    gameState: G, defenderMixedStrategy: MixedStrategyNode[DM, DS, _],
    undesiredAttackerPayoff: InnerResultNode[DM,DS],
    noSubtreeExpanded: Boolean
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo, br: BoundedRationality
  ) : Boolean = {
    if(defenderMixedStrategy.probabilities.length < gameState.possibleMoves.length && rng.nextDouble() < expansionProbability) {
      val root = mctsConfig.treeFactory.create(gameState.playerState,
        gameState.playerState.possibleMoves.filter(m => defenderMixedStrategy.probabilities.find(_._1 == m).isEmpty))
      val (_, bb, bestPayoff) = MctsAlgorithm.training(root.root, gameState,
        (g: G) => (-br.payoffChanger.changeExpectedPayoff(new Payoff(g.defenderPayoff, g.attackerPayoff)).attacker, None),
        mctsConfig, rng
      )
      val bestMoves = bb.reverse

      if(bestPayoff > - undesiredAttackerPayoff.payoff.distortedAttacker.attacker || expandWorsePayoff) {
        val nextState = gameState.makeMove(bestMoves.head)
        defenderMixedStrategy.addMove(bestMoves.head, 0)
        defenderMixedStrategy.addEmptySuccessor(new MoveStateKey(bestMoves.head, nextState.playerState))
        expandByList(nextState, defenderMixedStrategy.nextStateInTree(bestMoves.head, nextState.playerState), bestMoves.tail)
        undesiredAttackerPayoff.expandToLeaf(gameState, DummyExpander)
        true
      } else {
        false
      }
    } else {
      false
    }
  }

  def positivePassExpansion[DM, DS <: PlayerVisibleState[DM], G <: SinglePlayerGame[DM,DS,G] with GameWithPayoff](
    gameState: G,
    defenderMixedStrategy: MixedStrategyNode[DM, DS, _],
    desiredAttackerPayoff: InnerResultNode[DM,DS],
    noSubtreeExpanded: Boolean
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ): Boolean =  {
    if(defenderMixedStrategy.probabilities.length < gameState.possibleMoves.length && rng.nextDouble() < expansionProbability) {
      val root = mctsConfig.treeFactory.create(gameState.playerState,
        gameState.playerState.possibleMoves.filter(m => defenderMixedStrategy.probabilities.find(_._1 == m).isEmpty))
      val (_, bb, bestPayoff) = MctsAlgorithm.training(root.root, gameState,
        (g: G) => (g.defenderPayoff, None),
        mctsConfig, rng
      )
      val bestMoves = bb.reverse

      if(bestPayoff > desiredAttackerPayoff.payoff.real.defender || expandWorsePayoff) {
        val nextState = gameState.makeMove(bestMoves.head)
        defenderMixedStrategy.addMove(bestMoves.head, 0)
        defenderMixedStrategy.addEmptySuccessor(new MoveStateKey(bestMoves.head, nextState.playerState))
        expandByList(nextState, defenderMixedStrategy.nextStateInTree(bestMoves.head, nextState.playerState), bestMoves.tail)
        desiredAttackerPayoff.expandToLeaf(gameState, DummyExpander)
        true
      } else {
        false
      }
    } else {
      false
    }
  }
}
