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

import pl.edu.pw.mini.sg.rng.RandomGenerator
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, MinimalChangeAdjuster}
import pl.edu.pw.mini.sg.mixed.tree.{ InnerResultNode, MixedStrategyNode, OptimizerStateType }


final case class DampingRateAdjuster(
  initialRate: Double,
  intraIterationDamping: Double,
  interIterationDamping: Double
) extends StatefulAdjuster {
  override type AdjusterState = DampingState 

  implicit override val ost: OptimizerStateType[DampingState] =
    new DampingStateTypeclass(initialRate, interIterationDamping)

  override def inTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,DampingState],
    desiredAttackerPayoff: InnerResultNode[DM,DS],
    optimalAttackerPayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator,
    iterationInfo: IterationInfo
  ): Boolean = {
    val diff = desiredAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff.distortedAttacker.attacker
    val upd = mixedStrategy.probabilities.view.zipWithIndex.collect{
      case ((move, prob), idx) if(
        desiredAttackerPayoff.successors(move).payoff.distortedAttacker.attacker -
          optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker >= Math.max(diff,0)
      ) =>
        {
        ((move, prob+mixedStrategy.optimizerState.currentRate), idx)
        }
    }.toList

    upd.foreach{case (nv, idx) => if(rng.nextDouble() > 0.5) mixedStrategy.updateProb(idx, nv._2)}
    mixedStrategy.optimizerState = mixedStrategy.optimizerState.dampCurrent(intraIterationDamping)
    mixedStrategy.normalize()
    upd.nonEmpty    
  }

  override def outOfTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,DampingState],
    optimalAttackerPayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ): Boolean = {
    val payoff = optimalAttackerPayoff.payoff.distortedAttacker.attacker

    val upd = mixedStrategy.probabilities.view.zipWithIndex.collect{
      case ((move, prob), idx) if(optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker <= payoff) =>
        {
        ((move, prob+mixedStrategy.optimizerState.currentRate), idx)
        }
    }.toList

    upd.foreach{case (nv, idx) => if(rng.nextDouble() > 0.5) mixedStrategy.updateProb(idx, nv._2)}
    mixedStrategy.optimizerState = mixedStrategy.optimizerState.dampCurrent(intraIterationDamping)
    mixedStrategy.normalize()
    upd.nonEmpty
  }


  override def positiveAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,DampingState],
    gamePayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ): Boolean = {
    val payoff = gamePayoff.payoff.real.defender

    val upd = mixedStrategy.probabilities.view.zipWithIndex.collect{
      case ((move, prob), idx) if(gamePayoff.successors(move).payoff.real.defender >= payoff) =>
        {
        ((move, prob+mixedStrategy.optimizerState.currentRate), idx)
        }
    }.toList

    upd.foreach{case (nv, idx) => if(rng.nextDouble() > 0.5) mixedStrategy.updateProb(idx, nv._2)}
    mixedStrategy.optimizerState = mixedStrategy.optimizerState.dampCurrent(intraIterationDamping)
    mixedStrategy.normalize()
    upd.nonEmpty
  }
}
