package pl.edu.pw.mini.sg.mixed.adjuster.momentum

import pl.edu.pw.mini.sg.rng.RandomGenerator
import org.log4s._
import pl.edu.pw.mini.sg.game.deterministic.PlayerVisibleState
import pl.edu.pw.mini.sg.mixed.IterationInfo
import pl.edu.pw.mini.sg.mixed.adjuster.StatefulAdjuster
import pl.edu.pw.mini.sg.mixed.tree.{ InnerResultNode, MixedStrategyNode, OptimizerStateType }


case object CoordNormalizedMomentumAdjuster extends StatefulAdjuster {
  private val logger = getLogger
  override type AdjusterState = CoordNormalizedMomentum

  implicit val ost: OptimizerStateType[CoordNormalizedMomentum] = CoordNormalizedMomentumOptimizerStateType

  private def momentumAdjust[M](
    mixedStrategy: MixedStrategyNode[M, _, CoordNormalizedMomentum]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ): Boolean = {
    if(rng.nextDouble() > 0.3) {
      require(mixedStrategy.probabilities.length == mixedStrategy.optimizerState.momentum.length)
      mixedStrategy.probabilities.zip(mixedStrategy.optimizerState.normalized).zipWithIndex.map{
        case (((dm, prob), momentum), idx) => {
          if(momentum < 0) {
            val ch = Math.max(momentum, -prob)
            require(ch <= 0)
            mixedStrategy.updateProb(idx, prob+ch)
            -ch
          } else
              0.0
        }
      }.sum

      mixedStrategy.probabilities.zip(mixedStrategy.optimizerState.normalized).zipWithIndex.foreach{
        case (((dm, prob), momentum), idx) => {
          if(momentum > 0)
            mixedStrategy.updateProb(idx, prob + momentum)
        }
      }
      
      mixedStrategy.normalizeOrEqual()
      true
    } else false
  }

  /**
    *  @return Czy była zmiana
    */
  override def inTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, CoordNormalizedMomentum],
    desiredAttackerPayoff: InnerResultNode[DM, DS],
    optimalAttackerPayoff: InnerResultNode[DM, DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ) : Boolean = {
    val diff = desiredAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff.distortedAttacker.attacker
    def diffm(move: DM) = desiredAttackerPayoff.payoff(move).distortedAttacker.attacker - optimalAttackerPayoff.payoff(move).distortedAttacker.attacker

    mixedStrategy.optimizerState = mixedStrategy.optimizerState.accumulateAssessments(
      mixedStrategy.probabilities.view.map(p => diffm(p._1) - diff).toVector
    )
    momentumAdjust(mixedStrategy)
  }

   /**
    *  @return Czy była zmiana
    */
 override def outOfTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, AdjusterState],
    optimalAttackerPayoff: InnerResultNode[DM, DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
 ) : Boolean = {
   mixedStrategy.optimizerState = mixedStrategy.optimizerState.accumulateAssessments(
     mixedStrategy.probabilities.view.map(p =>
       optimalAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff(p._1).distortedAttacker.attacker).toVector
   )
   momentumAdjust(mixedStrategy)
 }


   /**
    *  @return Czy była zmiana
    */
 override def positiveAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, AdjusterState],
    gamePayoff: InnerResultNode[DM, DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
 ) : Boolean = {
   mixedStrategy.optimizerState = mixedStrategy.optimizerState.accumulateAssessments(
     mixedStrategy.probabilities.view.map(p =>
       gamePayoff.payoff(p._1).real.defender - gamePayoff.payoff.real.defender).map(_*10).toVector
   )
   momentumAdjust(mixedStrategy)
 }


}
