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

import pl.edu.pw.mini.sg.rng.RandomGenerator
import pl.edu.pw.mini.sg.mixed.IterationInfo
import pl.edu.pw.mini.sg.mixed.prune.PruneDecision
import scala.annotation.tailrec
import scala.collection.mutable.{ArrayBuffer, Map => MMap}

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

final class MixedStrategyNode[M, S <: PlayerVisibleState[M], T](
  val state: S,
  private val strategy : ArrayBuffer[(M, Double, IterationInfo)] = ArrayBuffer[(M, Double, IterationInfo)](),
  val succesors : MMap[M, MMap[S, MixedStrategyNode[M, S, T]]] = MMap[M, MMap[S, MixedStrategyNode[M, S, T]]]()
)(
  implicit val ost: OptimizerStateType[T]
) extends BehavioralStrategy[M, S] {
  var optimizerState: T = ost.zero


  override def nextState(move: M,reachedState: S): BehavioralStrategy[M,S] = {
    val ret = succesors.get(move).flatMap(_.get(reachedState))
    ret.foreach(x => require(x.state == reachedState))
    assert(ret.isEmpty == !successorsInTree.toList.contains((move, reachedState)), s"${move} ${reachedState} SIT:${successorsInTree.toList} SUC:${succesors} retstate: ${ret.get.state}")
    ret.getOrElse(new OutOfTreeNode[M, S](reachedState))
  }

  def nextStateInTree(move: M,reachedState: S): MixedStrategyNode[M, S, T] =
    succesors(move)(reachedState)


  def updateProb(move: M, newValue: Double)(implicit iterationInfo: IterationInfo) : Unit = {
    updateProb(strategy.zipWithIndex.find { case ((m, _, _), _) => move == m}.get._2, newValue)
  }

  def updateProb(idx: Int, newValue: Double)(implicit iterationInfo: IterationInfo) : Unit = {
    val (m, oldp, oldi) = strategy(idx)
    strategy.update(idx, (m, newValue, if(newValue == 0) oldi else iterationInfo))
  }

  def updateProbs(newValues: Seq[Double])(implicit iterationInfo: IterationInfo) : Unit = {
    strategy.zipWithIndex.zip(newValues).foreach{
      case (((m, _, oldi), idx), newVal) => strategy.update(idx, (m, newVal, if(newVal == 0) oldi else iterationInfo))
    }
  }

  def pruneMoves(pruneDecision: PruneDecision)(implicit current: IterationInfo) : Unit = {
    val toDelete = strategy.zipWithIndex.collect{case ((_, prob, lastNonZero), idx) if prob==0 && pruneDecision.decide(current, lastNonZero) => idx}.toVector
    optimizerState = ost.pruneMoves(optimizerState, toDelete.toSet)
    val movesToDelete = toDelete.map(i => strategy(i)._1).toSet
    movesToDelete.foreach(succesors.remove)
    toDelete.reverse.foreach(strategy.remove)
  }

  def pruneMove(move: M) : Unit = {
    val toDelete = strategy.zipWithIndex.collect{case ((m, prob, lastNonZero), idx) if prob==0 && m == move => idx}.toVector
    optimizerState = ost.pruneMoves(optimizerState, toDelete.toSet)
    val movesToDelete = toDelete.map(i => strategy(i)._1).toSet
    movesToDelete.foreach(succesors.remove)
    toDelete.reverse.foreach(strategy.remove)
  }

  def addMove(move: M, prob : Double = 0.0)(implicit iterationInfo : IterationInfo) : Unit = strategy.append((move, prob, iterationInfo))

  def nextState[G <: SinglePlayerGame[M, S, G]](
    move: M, reachedState: S, expander: TreeExpander[M, S, G], reachedFullState: => G
  )(
    expanderState : Option[expander.SubtreeData[M,S]] = None
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ): (MixedStrategyNode[M, S, T], Boolean) = {
    (succesors.get(move).flatMap(_.get(reachedState))) match {
      case Some(s) => (s, false)
      case None => {
        val newNode = new MixedStrategyNode[M, S, T](reachedState)
        succesors.getOrElseUpdate(move, MMap()).put(reachedState, newNode)
        if(strategy.find(_._1 == move).isEmpty) {
          strategy.append((move, 0, iterationInfo))
        }
        @tailrec
        def doExpansion(
          currentNode : MixedStrategyNode[M, S, T],
          currentState: G,  exState: expander.SubtreeData[M,S]
        ) : Unit = {
          if(currentState.possibleMoves.nonEmpty) {
            val (move, nextData) = expander.expand[M, S, G](currentState, exState)
            val nextState = currentState.makeMove(move)
            val node = new MixedStrategyNode[M, S, T](nextState.playerState)
            currentNode.strategy.clear()
            currentNode.strategy.append((move, 1.0, iterationInfo))
            currentNode.succesors.getOrElseUpdate(move, MMap()).put(node.state, node)
            doExpansion(node, nextState, nextData)
          } else {
            currentNode.strategy.clear()
          }
        }
        doExpansion(newNode, reachedFullState, expanderState.getOrElse(expander.zeroSubtreeData))
        (newNode, true)
      }
    }
  }

  override def probabilities: Seq[(M, Double)] = strategy.view.map{case (m, p, _) => (m, p)
  }

  def deepClone: MixedStrategyNode[M, S, Unit] = deepClone(MixedStrategyNode.unitST)

  def deepClone[OS](ost: OptimizerStateType[OS]) : MixedStrategyNode[M, S, OS] = {
    val ret = new MixedStrategyNode[M, S, OS](state)(ost)
    strategy.foreach(ret.strategy.append(_))
    succesors.foreach{case (key, value) => ret.succesors.put(key, MMap(value.toSeq.map{case (k,v) => (k, v.deepClone(ost))}:_*))}
    ret
  }

  def deepClonePreserveOS() : MixedStrategyNode[M, S, T] = {
    val ret = new MixedStrategyNode[M, S, T](state)(ost)
    strategy.foreach(ret.strategy.append(_))
    succesors.foreach{case (key, value) => ret.succesors.put(key,
      MMap(value.toSeq.map{case (k,v) => (k, v.deepClonePreserveOS)} :_*))}
    ret
  }

  def normalizeOrEqual()(implicit iterationInfo: IterationInfo) : Unit = {
    assert(strategy.view.map(_._2).forall(_>=0), s"no negative probabilities allowed!! strategy: ${strategy.view.map(_._2).toVector}")
    if(strategy.exists(_._2 > 0)) normalize()
    else strategy.view.zipWithIndex.foreach{case ((dm, prob, _), idx) =>
      strategy.update(idx, (dm, 1.0/strategy.length.toDouble, iterationInfo))
    }
  }

  def normalize(): Unit = {
    val sum = strategy.view.map(_._2).sum;
    require(sum > 0, s"Expected positive sum, got $sum")
    (0 to strategy.length-1).foreach{ idx => {
      val curr = strategy(idx)
      strategy.update(idx, (curr._1, curr._2/sum, curr._3))
    }}
  }

  def afterIteration() : Unit = {
    optimizerState = ost.afterIteration(optimizerState)
    succesors.values.foreach { _.values.foreach(_.afterIteration()) }
  }

  override def toString: String = s"st:${state} prob:${probabilities} ${succesors}"

  override def successorsInTree: Iterable[(M, S)] = succesors.toSeq.flatMap{case (k, v) => v.keys.map(vv => (k,vv)).toSeq}

  def updateBy(updates: Seq[Double])(implicit iterationInfo: IterationInfo) : Unit = {
    assert(updates.length == strategy.length)
    strategy.zipWithIndex.zip(updates).foreach { case (((dm, prob, oldi),idx), upd) =>
      strategy.update(idx, (dm, Math.max(0, prob+upd), if(prob+upd > 0) iterationInfo else oldi))
    }
    normalizeOrEqual()
  }

  def addSuccessor(msk: MoveStateKey[M, S], node: MixedStrategyNode[M, S, T]) : Unit = {
    assert(msk.state == node.state)
    succesors.getOrElseUpdate(msk.move, MMap()).put(msk.state,node)
  }

  def addEmptySuccessor(msk: MoveStateKey[M, S]) : Unit = {
    succesors.getOrElseUpdate(msk.move, MMap()).put(msk.state,
      new MixedStrategyNode(msk.state)
    )
  }
}


object MixedStrategyNode {
  implicit val unitST : OptimizerStateType[Unit] = new OptimizerStateType[Unit] {
    override val zero: Unit = ()
    override def afterIteration(prev: Unit) = prev
    override def pruneMoves(t: Unit, indices: Set[Int]) = t
  }
}


final class OutOfTreeNode[M, S <: PlayerVisibleState[M]](val state: S) extends BehavioralStrategy[M, S] {
  def nextState(move: M,reachedState: S): BehavioralStrategy[M,S] = new OutOfTreeNode[M, S](reachedState)
  def probabilities: List[(M, Double)] = {
    val l = state.possibleMoves.length
    state.possibleMoves.map(m => (m, 1.0/l.toDouble))
  }

  def successorsInTree: Iterable[(M, S)] = Seq()
}

