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

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}
import scala.annotation.tailrec


final case class InhibitingStrategyAdjuster(
  rate: Double
) extends StatefulAdjuster {
  private val logger = getLogger

  override type AdjusterState = InhibitorState

  implicit val ost: OptimizerStateType[InhibitorState] = InhibitorStateTypeclass

  private def increaseMoves[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, InhibitorState],
    movesToIncrease: Seq[Int],
    inhibitNotSelected: Boolean = true,
    rate: Double = this.rate
  )(
    implicit iterationInfo: IterationInfo
  ) : Boolean = {
    val is = mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.size)
    if(movesToIncrease.isEmpty) {
      false
    } else {
      movesToIncrease.foreach(idx => {
        val inc = Math.max(0, rate * (1 - is.inhibitingCoefficient(idx)))
        val (move, curr) = mixedStrategy.probabilities(idx)
        mixedStrategy.updateProb(idx, curr+inc) 
      })
      mixedStrategy.optimizerState = if(inhibitNotSelected) is.increaseInhibitorsNP(movesToIncrease, rate) else is
      mixedStrategy.normalize()
      true
    }
  }

  /**
    *  @return Czy była zmiana
    */
  override def inTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, InhibitorState],
    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 indicesToIncrease = mixedStrategy.probabilities.zipWithIndex
      .collect{case ((move, _), idx) if desiredAttackerPayoff.successors(move).payoff.distortedAttacker.attacker - optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker >= diff => (idx, desiredAttackerPayoff.successors(move).payoff.distortedAttacker.attacker - optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker, desiredAttackerPayoff.successors(move).payoff.real.defender)}// TODO Czy rozstrzyganie remisów powinno uwzględniać zniekształcenie?
    if(indicesToIncrease.nonEmpty && rng.nextDouble() < 0.7) {
      val is = mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.length)
      mixedStrategy.optimizerState = is
      val ord = Ordering.fromLessThan[(Int, Double, Double)]((p1, p2) =>
        if(is.inhibitingCoefficient(p1._1) ==
          is.inhibitingCoefficient(p2._1)) p1._3 > p2._3
        else is.inhibitingCoefficient(p1._1) > is.inhibitingCoefficient(p2._1)
      )
      val nnDiff = indicesToIncrease.filter(_._2 >= 0)
      if(nnDiff.nonEmpty && diff < 0) {
        val pDiff = nnDiff.filter(_._2 > 0)
        val best = if(pDiff.nonEmpty) pDiff.max(ord) else  nnDiff.max(ord)
        if(best._2 > 0) {
          val change = -diff/best._2
          val (move, prob) = mixedStrategy.probabilities(best._1)
          mixedStrategy.updateProb(best._1, prob+change)
          mixedStrategy.normalize()
          mixedStrategy.optimizerState = is.increaseInhibitorsNP(nnDiff.map(_._1), change)
          true
        } else {
          if(rng.nextDouble() > 0.0) {
            val movesToZero = mixedStrategy.probabilities.zipWithIndex.filter{case ((move, prob), idx) => desiredAttackerPayoff.successors(move).payoff.distortedAttacker.attacker - optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker < 0}
              .map{case ((move, prob), idx) => (move, idx)}
            movesToZero.foreach{case (move, idx) => mixedStrategy.updateProb(idx, 0.0)}
            mixedStrategy.optimizerState = is.increaseInhibitors(movesToZero.map(_._2), rate)
            if(mixedStrategy.probabilities.view.map(_._2).sum == 0) {
              nnDiff.foreach{case (idx, _, _) => mixedStrategy.updateProb(idx, mixedStrategy.probabilities(idx) match {case (move, pr) => rate*Math.max(0.05,(1-is.inhibitingCoefficient(idx)))})}
              //mixedStrategy.strategy.update(best._1, mixedStrategy.strategy(best._1) match {case (move, pr) => (move, 1)})
            }

            mixedStrategy.normalize()

            desiredAttackerPayoff.updateValues()
            optimalAttackerPayoff.updateValues()
            // require(desiredAttackerPayoff.payoff.attacker >= optimalAttackerPayoff.payoff.attacker)

            true
          } else {
            false
          }
        }
      } else {
        val ret = increaseMoves(mixedStrategy, indicesToIncrease.map(_._1), true);
        ret
      }
      true
    } else {
      false
    }
  }

   /**
    *  @return Czy była zmiana
    */
 override def outOfTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, InhibitorState],
    optimalAttackerPayoff: InnerResultNode[DM, DS]
  )(
   implicit rng: RandomGenerator,
    iterationInfo: IterationInfo
 ) : Boolean = {
   if(rng.nextDouble() > 0.2) {
     val payoff = optimalAttackerPayoff.payoff.distortedAttacker.attacker
     val indicesToIncrease = mixedStrategy.probabilities.zipWithIndex
     .collect{case((move, _), idx) if payoff > optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker => idx}

     if(indicesToIncrease.nonEmpty)
       increaseMoves(mixedStrategy, indicesToIncrease)
     else
       false
   } else {
     false
   }
 }


   /**
    *  @return Czy była zmiana
    */
 override def positiveAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM, DS, InhibitorState],
    gamePayoff: InnerResultNode[DM, DS]
  )(
   implicit rng: RandomGenerator,
    iterationInfo: IterationInfo
 ) : Boolean = {
   if(rng.nextDouble() > 0.4) {
   val payoff = gamePayoff.payoff.real.defender
   val indicesToIncrease = mixedStrategy.probabilities.zipWithIndex
     .collect{case ((move, _), idx) if payoff < gamePayoff.successors(move).payoff.real.defender => idx}

     if(indicesToIncrease.nonEmpty) {
       mixedStrategy.optimizerState = mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.length)
       logger.trace(s"inh: ${mixedStrategy.optimizerState.inhibitor}")
       val ii = indicesToIncrease.maxBy(-mixedStrategy.optimizerState.inhibitor(_))
       mixedStrategy.updateProb(ii, mixedStrategy.probabilities(ii) match {case (m, p) => p+rate*10})
       mixedStrategy.normalize()
//       mixedStrategy.optimizerState = mixedStrategy.optimizerState.increaseInhibitorsNP(List(ii), this.rate)
       true
//       increaseMoves(mixedStrategy, indicesToIncrease, rate = this.rate, inhibitNotSelected=true)
     } else {
       false
     }
   } else
       false
 }

}
