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

import pl.edu.pw.mini.sg.game.boundedrationality.DistortedPayoff
import pl.edu.pw.mini.sg.game.boundedrationality.probabilitieschanger.{ ProbabilitiesChangerLeaf, ProbabilitiesChangerRegular }
import pl.edu.pw.mini.sg.mixed.IterationInfo
import scala.language.existentials

import pl.edu.pw.mini.sg.rng.RandomGenerator
import pl.edu.pw.mini.sg.game.deterministic.single.{ GameWithPayoff }
import pl.edu.pw.mini.sg.game.Payoff
import pl.edu.pw.mini.sg.game.boundedrationality.BoundedRationality
import scala.collection.mutable.{Map => MMap}

import pl.edu.pw.mini.sg.game.deterministic.{PlayerVisibleState, TreeExpander}
import pl.edu.pw.mini.sg.game.deterministic.single.SinglePlayerGame

import DistortedPayoff.Distrotion

sealed trait GameResultNode[M, S] {
  def updateValues() : Unit
  def payoff: DistortedPayoff
  def state: S
}

// Skąd brać właściwe stany do rozwijania drzewa?
// Czy trzymać też atakującego??
final class InnerResultNode[M, S <: PlayerVisibleState[M]](
  val strategyNode: MixedStrategyNode[M, S, _],
  val playerState: S
)(
  implicit val boundedRationality: BoundedRationality
) extends GameResultNode[M, S] {
  var payoff : DistortedPayoff = DistortedPayoff.zero
  val successors = MMap[M, GameResultNode[M, S]]()

  def payoff(m : M) : DistortedPayoff = successors(m).payoff

  require(playerState.possibleMoves.nonEmpty)

  override def state = strategyNode.state
  override def updateValues() : Unit = {
    payoff = {
      boundedRationality.probabilitiesChanger match {
        case r : ProbabilitiesChangerRegular => r.changeProbability(
          strategyNode.probabilities, playerState.possibleMoves)
            .foldLeft(DistortedPayoff.zero){case (p, (move, prob)) =>  {
              val realProb = strategyNode.probabilities.find(_._1 == move).map(_._2).getOrElse(0.0)
              if (prob > 0 || realProb > 0)  p + successors(move).payoff*((realProb, prob)) else p}
            }
        case l : ProbabilitiesChangerLeaf => {
          val br = 1.0/playerState.possibleMoves.length.toDouble
          playerState.possibleMoves.foldLeft(DistortedPayoff.zero){case (p, move) => 
            {
              val prob = strategyNode.probabilities.find(_._1 == move).map(_._2).getOrElse(0.0)
              successors.get(move) match {
                case Some(irn: InnerResultNode[M,S]) if prob > 0 => p+irn.payoff*((prob, prob))
                case Some(ln: LeafNode[M,S]) => p+ln.payoff*((prob, l.leafAdjust(prob, br)))
                case _ => p
            }}}
        }
      }
    }
  }

  def nextNode(move: M) : GameResultNode[M, S] = {
    require(successors.contains(move))
    successors(move)
  }

  // Czy doszło do ekspansji + następnik
  def expandToLeaf[G <: SinglePlayerGame[M, S, G] with GameWithPayoff](
    move: M, stateAfterMove: G,
    treeExpander: TreeExpander[M, S, G]
  )(
    implicit rng: RandomGenerator,  iterationInfo: IterationInfo
  ): (GameResultNode[M, S], Boolean) = {
    successors.get(move) match {
      case Some(inr: InnerResultNode[M, S]) => (inr, inr.expandToLeaf(stateAfterMove, treeExpander))
      case Some(leaf: LeafNode[M, S]) => (leaf, false)
      case None => {
        val (msNext, expanded) = strategyNode.nextState(move, stateAfterMove.playerState, treeExpander, stateAfterMove)(None)
        if(expanded) {
          val succ = expandSinglePath(stateAfterMove, msNext)
          successors.put(move, succ)
          updateValues()
          (succ, true)
        } else {
          val (succ, ex) = if(stateAfterMove.possibleMoves.nonEmpty) {
            val r = new InnerResultNode[M, S](msNext, stateAfterMove.playerState)
            val e = r.expandToLeaf(stateAfterMove, treeExpander)
            if(successors.size == strategyNode.probabilities.length)
              updateValues()
            boundedRationality.probabilitiesChanger match {
              case l: ProbabilitiesChangerLeaf => {
                stateAfterMove.playerState.possibleMoves.foreach{m =>
                  val nextState = stateAfterMove.makeMove(m)
                  if(nextState.playerState.possibleMoves.isEmpty)
                    r.successors.put(m, new LeafNode(new Payoff(nextState.defenderPayoff, nextState.attackerPayoff), nextState.playerState))
                }
                updateValues()
              }
              case r: ProbabilitiesChangerRegular => ()
            }
            (r, e)
          } else {
            (new LeafNode[M, S](new Payoff(stateAfterMove.defenderPayoff, stateAfterMove.attackerPayoff),
              stateAfterMove.playerState), false)
          }
          successors.put(move,succ)
          if(successors.size == strategyNode.probabilities.length)
            updateValues()
          (succ, ex)
        }
      }
    }
  }

  // A co z ruchem??
  def expandToLeaf[G <: SinglePlayerGame[M, S, G] with GameWithPayoff](
    currentState: G,
    treeExpander: TreeExpander[M, S, G]
  )(
    implicit rng: RandomGenerator,  iterationInfo: IterationInfo
  ): Boolean = {
    if(strategyNode.probabilities.length > 0) {
      val ret = strategyNode.probabilities.map{case (dm, prob) => {
        expandToLeaf(dm, currentState.makeMove(dm), treeExpander)._2
      }}.contains(true)
      updateValues()
      ret
    } else {
      val (move, expanderState) = treeExpander.expand[M, S, G](currentState, treeExpander.zeroSubtreeData)
      strategyNode.addMove(move, 1.0)
      val nextState = currentState.makeMove(move)
      val (ms, x) = strategyNode.nextState(move, nextState.playerState, treeExpander, nextState)(Some(expanderState))
      require(x == true)
      successors.put(move, expandSinglePath(nextState, ms))
      updateValues()
      true
    }
  }

  // Z tego można zrobic tailrec. Ale to powinno byc maks 50 zejść więc po co. 
  private def expandSinglePath[G <: SinglePlayerGame[M, S, G] with GameWithPayoff](
    state: G,
    defenderStrategy: MixedStrategyNode[M, S, _]
  ) : GameResultNode[M, S] = {
    if(state.possibleMoves.nonEmpty) {
      val move = {
        require(defenderStrategy.probabilities.length == 1)
        defenderStrategy.probabilities.head._1
      }
      val stateNext = state.makeMove(move)
      val msNext = defenderStrategy.nextStateInTree(move, stateNext.playerState)
      val succ = new InnerResultNode(defenderStrategy, state.playerState)
      succ.successors.put(move, expandSinglePath(stateNext, msNext))
      succ.updateValues()
      succ
    } else {
      new LeafNode[M, S](new Payoff(state.defenderPayoff, state.attackerPayoff), state.playerState)
    }
  }

  def prune() : Unit = {
    val keysToPrune = successors.keys.filter(m => !strategyNode.probabilities.view.map(_._1).contains(m) &&
      ((successors(m), boundedRationality.probabilitiesChanger) match {
        case (x: LeafNode[M,S], y: ProbabilitiesChangerLeaf) => false
        case _ => true
      })).toList
    keysToPrune.foreach(successors.remove)
  }
}



final case class LeafNode[M, S](
  val  realPayoff: Payoff,
  override val state: S
)(
  implicit boundedRationality: BoundedRationality
) extends GameResultNode[M, S] {
  override def updateValues() : Unit = ()
  override def payoff: DistortedPayoff = realPayoff.toDistorted
}
