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

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 }


final case object NormalizingAdjuster extends StatefulAdjuster {
  type AdjusterState = Normalizer

  implicit val ost: OptimizerStateType[Normalizer] = NormalizerOptimizerStateType

  private val logger = getLogger

  private def updateNode[DM](
    mixedStrategyNode: MixedStrategyNode[DM, _, AdjusterState],
    updates: Vector[Double]
  )(implicit  iterationInfo: IterationInfo) : Boolean = {
    mixedStrategyNode.optimizerState = mixedStrategyNode.optimizerState.registerUpdate(updates)
    val updated = mixedStrategyNode.probabilities.zip(mixedStrategyNode.optimizerState.normalizers.zip(updates))
      .map{case ((dm, prob), (norm, update)) =>{(dm, Math.max(0, prob+(update*norm)))}}
    updated.zipWithIndex.foreach{ case (v, i) => mixedStrategyNode.updateProb(i, v._2)}
    mixedStrategyNode.normalizeOrEqual()
    true
  }

  /**
    *  @return Czy była zmiana
    */
  override def inTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, Normalizer],
    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

    updateNode(mixedStrategy, mixedStrategy.probabilities.view.map(p => diffm(p._1) - diff).toVector)
  }

   /**
    *  @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 = {
   updateNode(mixedStrategy, mixedStrategy.probabilities.view.map(p =>
       optimalAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff(p._1).distortedAttacker.attacker).toVector)
 }


   /**
    *  @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 = {
   updateNode(mixedStrategy, mixedStrategy.probabilities.view.map(p =>
     gamePayoff.payoff(p._1).real.defender - gamePayoff.payoff.real.defender).toVector
   )
 }
}
