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

import pl.edu.pw.mini.sg.rng.RandomGenerator
import org.log4s._
import pl.edu.pw.mini.sg.mixed.IterationInfo
import pl.edu.pw.mini.sg.mixed.tree.InnerResultNode
import pl.edu.pw.mini.sg.mixed.tree.MixedStrategyNode
import pl.edu.pw.mini.sg.game.deterministic.PlayerVisibleState


final case class SimpleInTreeAdjuster (
  rate: Double,
  eps: Double
) extends StrategyAdjuster {
  private val logger = getLogger
    def adjust[DM, DS <: PlayerVisibleState[DM]](
      mixedStrategy: MixedStrategyNode[DM, DS, Unit],
      desiredAttackerPayoff: InnerResultNode[DM, DS],
      optimalAttackerPayoff: InnerResultNode[DM, DS],
      rng: RandomGenerator
    )(
      implicit     iterationInfo: IterationInfo
    ): Boolean = {
      if(desiredAttackerPayoff.payoff.distortedAttacker.attacker < optimalAttackerPayoff.payoff.distortedAttacker.attacker && rng.nextDouble() > 0.2) {
        val differences = mixedStrategy.probabilities.view.map{case (move, _) =>
          (move,
            desiredAttackerPayoff.successors(move).payoff.distortedAttacker.attacker - optimalAttackerPayoff.successors(move).payoff.distortedAttacker.attacker,
            //TODO: czy rozsztrzyganie remisów powinno używac zniekształceń?
            desiredAttackerPayoff.successors(move).payoff.distortedAttacker.defender
          )
        }

        logger.trace(s"Adjusting MS:${mixedStrategy.probabilities}, differences: ${differences.toVector}")

        val bestMove = differences.maxBy(triple => (triple._2, triple._3))

        if(bestMove._2 < 0) {
          logger.trace(s"Can't adjust in tree ${differences.toList}")
          false
        } else if(bestMove._2 <= 1e-5) {
          differences.zipWithIndex.foreach { case ((move, diff, _), idx) => {
            if(diff < 0)
              mixedStrategy.updateProb(idx, 0)
            else if(move == bestMove._1 && mixedStrategy.probabilities(idx)._2 == 0)
              mixedStrategy.updateProb(idx, rate)
          }}
          mixedStrategy.normalize()
          logger.trace("Update1")
          true
        } else {
          val currDiff = desiredAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff.distortedAttacker.attacker
          val bestDiff = bestMove._2
          require(currDiff < 0)
          require(bestDiff > 0)
          val change = -currDiff/bestDiff 
          val idx = mixedStrategy.probabilities.indexWhere { case (m, p) => m == bestMove._1 }
          val old = mixedStrategy.probabilities(idx)
          mixedStrategy.updateProb(idx, old._2+change)
          mixedStrategy.normalize()
          logger.trace(s"Update 2, cd=${currDiff}, bd=${bestDiff}, change=${change}")
          true
        }
      } else {
        false
      }


      // val nonNegativeDifferences = differences.zipWithIndex.collect{case ((dm, diff), ind) if diff >= 0 => (ind, diff)}.toList
      // if(nonNegativeDifferences.nonEmpty) {
      //   logger.trace(s"Adjusting started ${eps}, ${desiredAttackerPayoff.payoff.attacker}, ${optimalAttackerPayoff.payoff.attacker}")
      //   while(desiredAttackerPayoff.payoff.attacker + eps < optimalAttackerPayoff.payoff.attacker) {
      //     nonNegativeDifferences.foreach{case (ind, diff) =>
      //       val (move, prob) = mixedStrategy.strategy(ind)
      //       mixedStrategy.strategy.update(ind, (move,prob+rate*(rng.nextDouble())))
      //     }
      //     mixedStrategy.normalize()
      //     desiredAttackerPayoff.updateValues()
      //     optimalAttackerPayoff.updateValues()
      //     logger.trace(s"Adjusting: ${desiredAttackerPayoff.payoff.attacker - optimalAttackerPayoff.payoff.attacker}")
      //   }
      //   logger.trace("Adjusting finished")
      //   true
      // } else {
      //   logger.trace(s"Can't adjust in tree ${differences.toList}")
      //   false
      // }
    }
}
