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}

sealed trait Rank

final object Negative extends Rank
final object Positive extends Rank
final object Neutral extends Rank

object Rank {
  def apply(x: Double) = if(x < 0) Negative else if (x > 0) Positive else Neutral
}


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

  override type AdjusterState = InhibitorState
  implicit val ost: OptimizerStateType[InhibitorState] = InhibitorStateTypeclass

  private def adjust[DM, DS <: PlayerVisibleState[DM]](
    idxPositive: Seq[Int],
    idxNeutral: Seq[Int],
    idxNegative: Seq[Int],
    mixedStrategy: MixedStrategyNode[DM, DS, InhibitorState],
    rate: Double
  )(
    implicit iterationInfo: IterationInfo
  ) : Boolean  = {
    def probSum(idxs: Seq[Int]) : Double = idxs.view.map(mixedStrategy.probabilities(_)._2).sum
    val negativeSum: Double = probSum(idxNegative)
    val neutralSum: Double = probSum(idxNeutral)
    val is = mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.length)

    def increase(idxs: Seq[Int], rate: Double) : Unit = {
      val inhibitorsNorm=idxs.view.map(idx => 1 - is.inhibitingCoefficient(idx)).sum
      require(inhibitorsNorm >= 0, s"Required positive sum, got ${inhibitorsNorm}")
      if(inhibitorsNorm > 0)
        idxs.foreach(idx => mixedStrategy.updateProb(idx, mixedStrategy.probabilities(idx) match {case (m, p) => p+rate*(1-is.inhibitingCoefficient(idx))/inhibitorsNorm}))
      else {
        idxs.foreach(idx => mixedStrategy.updateProb(idx, mixedStrategy.probabilities(idx) match {case (m, p) => p+rate/idxs.length}))
      }
    }
    def decrease(idxs: Seq[Int], rate: Double) : Double = {
      idxs.map(idx => {
        val (m, pr) = mixedStrategy.probabilities(idx)
        val newPr = Math.max(0, pr-rate/idxs.length*(0.5*is.inhibitingCoefficient(idx)))
        mixedStrategy.updateProb(idx, newPr)
        pr-newPr
      }).sum
    }

    if(negativeSum > 0 && (idxPositive.length > 0 || idxNeutral.length > 0)) {
      val toInc = if(idxPositive.length > 0) idxPositive else idxNeutral
      val irMax = Math.min(0.9, toInc.view.map(i => is.inhibitingCoefficient(i)).max)
      val negativeReduced = decrease(idxNegative, rate*(1-irMax))
      increase(toInc, negativeReduced)
      mixedStrategy.normalize()
      mixedStrategy.optimizerState = is.increaseInhibitors(idxNegative, rate*(1-irMax))
      true
    } else if (idxPositive.length > 0 &&  neutralSum > 0) {
      val irMax = Math.min(0.9, idxPositive.view.map(i => is.inhibitingCoefficient(i)).max)
      val neutralReduces = decrease(idxNeutral, rate*(1-irMax))
      increase(idxPositive, neutralReduces)
      mixedStrategy.normalize()
      mixedStrategy.optimizerState = is.increaseInhibitors(idxNeutral, rate*(1-irMax))
      true
    } else {
      false
    }
  }

  private def assessment[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy : MixedStrategyNode[DM, DS, InhibitorState]
  )(
    moveAssesment:  DM => Double
  ) : List[(Int, Double)] = {
    mixedStrategy.probabilities.view.zipWithIndex.map{case ((dm, _), idx) => (idx, moveAssesment(dm))}.toList
  }

  private def groupAssessment(
    as: List[(Int, Double)]
  ) : Map[Rank, List[Int]] = as.groupBy(p => Rank(p._2)).mapValues(_.map(_._1))

  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 as = assessment(mixedStrategy)(move => (desiredAttackerPayoff.nextNode(move).payoff.distortedAttacker.attacker - optimalAttackerPayoff.nextNode(move).payoff.distortedAttacker.attacker) - diff)
    val indices = groupAssessment(as)
    
    val r = adjust(indices.getOrElse(Positive, Nil), indices.getOrElse(Neutral, Nil), indices.getOrElse(Negative, Nil), mixedStrategy, rate)
    mixedStrategy.optimizerState= mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.length).updated(as.map(p => (p._1, -p._2)))
    r
  }

  override def outOfTreeAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,InhibitorState],
    optimalAttackerPayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo): Boolean = {
    val as = assessment(mixedStrategy)(dm => optimalAttackerPayoff.payoff.distortedAttacker.attacker - optimalAttackerPayoff.payoff(dm).distortedAttacker.attacker)
    val indices = groupAssessment(as)
    
    val r = adjust(indices.getOrElse(Positive, Nil), indices.getOrElse(Neutral, Nil), indices.getOrElse(Negative, Nil), mixedStrategy, rate)
    mixedStrategy.optimizerState= mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.length).updated(as.map(p => (p._1, -p._2)))
    r
  }

  override def positiveAdjust[DM, DS <: PlayerVisibleState[DM]](
    mixedStrategy: MixedStrategyNode[DM,DS,InhibitorState],
    gamePayoff: InnerResultNode[DM,DS]
  )(
    implicit rng: RandomGenerator, iterationInfo: IterationInfo
  ): Boolean = {
    val as = assessment(mixedStrategy)(dm => gamePayoff.payoff(dm).real.defender - gamePayoff.payoff.real.defender)
    val indices = groupAssessment(as)
    val r = adjust(indices.getOrElse(Positive, Nil), indices.getOrElse(Neutral, Nil), indices.getOrElse(Negative, Nil), mixedStrategy, rate)
    mixedStrategy.optimizerState= mixedStrategy.optimizerState.extend(mixedStrategy.probabilities.length).updated(as.map(p => (p._1, -p._2)))
    r
  }
}
