package pl.edu.pw.mini.sg.game.util.attacker

import pl.edu.pw.mini.sg.game.boundedrationality.DistortedPayoff
import pl.edu.pw.mini.sg.util.StreamUtils
import scala.annotation.tailrec
import scala.math.Ordering

import pl.edu.pw.mini.sg.game.Payoff
import pl.edu.pw.mini.sg.game.deterministic.{AdvanceablePlayerVisibleState, BehavioralStrategy, DeterministicGame, PlayerVisibleState, PureStrategy}
import pl.edu.pw.mini.sg.game.deterministic.single.LinearizedAdvanceableGame
import pl.edu.pw.mini.sg.game.boundedrationality.BoundedRationality
import pl.edu.pw.mini.sg.game.boundedrationality.probabilitieschanger.{ProbabilitiesChangerLeaf, ProbabilitiesChangerRegular}
import scalaz.NonEmptyList
import DistortedPayoff.Distrotion

object OptimalAttacker {
  val EPS = 1e-3

  private implicit val ordering:  Ordering[DistortedPayoff] = Ordering.fromLessThan[DistortedPayoff]((a, b) =>
    if(Math.abs(a.distortedAttacker.attacker - b.distortedAttacker.attacker) < EPS)
      a.distortedAttacker.defender < b.distortedAttacker.defender
    else
      a.distortedAttacker.attacker < b.distortedAttacker.attacker
  )


  def browseLeaves[M, S <: AdvanceablePlayerVisibleState[M,S]](
    g: LinearizedAdvanceableGame[M, S]
  ) : Stream[LinearizedAdvanceableGame[M, S]]  = {
    if(g.possibleMoves.isEmpty) Stream(g)
    else {
      g.possibleMoves.toStream.flatMap(m => browseLeaves[M,S](g.makeMove(m)))
    }
  }

  def calculate[
    DM, AM, DS <: PlayerVisibleState[DM],
    AS <: AdvanceablePlayerVisibleState[AM, AS],
    G <: DeterministicGame[DM, AM, DS, AS, G]
  ](
    gameRoot: G,
    defenderStrategy: BehavioralStrategy[DM, DS]
  )(
    implicit boundedRationality : BoundedRationality
  ): (PureStrategy[AM, AS], DistortedPayoff) = {
    val lag = new LinearizedAdvanceableGame[AM, AS](NonEmptyList(gameRoot.attackerState), Nil, Nil)
    implicit val pairOrd : Ordering[(PureStrategy[AM, AS], DistortedPayoff)]  = Ordering.by(_._2)
    @tailrec
    def browseAttackerStrategies(
      queue: List[LinearizedAdvanceableGame[AM, AS]],
      currentBest : Option[(PureStrategy[AM, AS], DistortedPayoff)]
    ) : (PureStrategy[AM, AS], DistortedPayoff) = {
      queue match {
        case Nil => currentBest.get
        case lag :: tail => {
          if(lag.possibleMoves.isEmpty) {
            val ps = lag.pureStrategy
            val p = payoff(gameRoot, defenderStrategy, ps)
            browseAttackerStrategies(tail, implicitly[Ordering[Option[(PureStrategy[AM, AS], DistortedPayoff)]]].max(Some((ps,p)), currentBest))
          } else {
            browseAttackerStrategies(lag.possibleMoves.foldLeft(tail)((list, am) => (lag.makeMove(am))::list), currentBest)
          }
        }
      }
    }

    def browseDefenderStrategies(node: LinearizedAdvanceableGame[AM, AS]) : (PureStrategy[AM, AS], DistortedPayoff) = {
      @tailrec def chooseOptimalResponse(strategies: List[AM], currentBest: Option[(PureStrategy[AM, AS], DistortedPayoff)]) : Option[(PureStrategy[AM, AS], DistortedPayoff)] = {
        if(strategies.isEmpty) {
          currentBest
        } else {
          val r = browseDefenderStrategies(node.makeMove(strategies.head))
          if(currentBest.map(c => ordering.lt(c._2,r._2)).getOrElse(true))
            chooseOptimalResponse(strategies.tail, Some(r))
          else
            chooseOptimalResponse(strategies.tail, currentBest)
        }
      }
      chooseOptimalResponse(node.possibleMoves, None).getOrElse {
        val strategy = node.pureStrategy
        (strategy, payoff(gameRoot, defenderStrategy, strategy))
      }
    }

    browseAttackerStrategies(List(lag), None)
  }


  def payoff[
    DM, AM,
    DS <: PlayerVisibleState[DM],
    AS <: PlayerVisibleState[AM],
    G <: DeterministicGame[DM, AM, DS, AS, G]
  ](
    gameState: G,
    defenderStrategy: BehavioralStrategy[DM, DS],
    attackerStrategy: PureStrategy[AM, AS]
  )(
    implicit boundedRationality: BoundedRationality
  ) : DistortedPayoff = {
    final case class TraverseData(
      gameNode: G,
      defenderStrategy: BehavioralStrategy[DM, DS],
      attackerStrategy: PureStrategy[AM, AS],
      multiplierWithoutLast: Double,
      lastProb: Double,
      parentBranchingCoeff: Double
    )


    boundedRationality.probabilitiesChanger match {
      case r : ProbabilitiesChangerRegular => {
        @tailrec def calculateSum(queue: List[(G, BehavioralStrategy[DM, DS], PureStrategy[AM, AS], Double, Double)],
          currentSum: DistortedPayoff) : DistortedPayoff = {
          queue match {
            case (game, ds, as, multiplierReal, multiplierDistorted) :: tail => {
              if(game.isLeaf) calculateSum(tail, currentSum+(
                (new Payoff(game.defenderPayoff, game.attackerPayoff).toDistorted*((multiplierReal, multiplierDistorted))
                )))
              else {
                val newQueue = r
                  .changeProbability(ds.probabilities, game.defenderState.possibleMoves)
                  .foldLeft(tail){case (list, (dm, prob)) => {
                    require(prob >= 0, s"negative probability: $prob")
                    if (prob > 0) {
                      val nextState = game.makeMove(dm, as.move)
                      val nextDS = ds.nextState(dm, nextState.defenderState)
                      val nextAS = as.nextState(nextState.attackerState)
                      (nextState, nextDS, nextAS, ds.probabilities.find(_._1==dm).map(_._2).getOrElse(0.0)
                        *multiplierReal, prob*multiplierDistorted) :: list
                    } else {
                      list
                    }}}
                        
                calculateSum(newQueue, currentSum)
              }
            }
            case Nil => currentSum
          }
        }

        calculateSum(List((gameState, defenderStrategy, attackerStrategy, 1.0, 1.0)), DistortedPayoff.zero)
      }
      case l : ProbabilitiesChangerLeaf => {
        @tailrec def calculateSum(queue: List[TraverseData],
          currentSum: DistortedPayoff) : DistortedPayoff = {
          queue match {
            case td :: tail => {
              if(td.gameNode.isLeaf) calculateSum(tail, currentSum+(
                (new Payoff(td.gameNode.defenderPayoff, td.gameNode.attackerPayoff).toDistorted*
                  ((td.lastProb*td.multiplierWithoutLast, l.leafAdjust(td.lastProb, td.parentBranchingCoeff)*td.multiplierWithoutLast))
                )))
              else {
                val mwl = td.multiplierWithoutLast*td.lastProb
                val branching = 1.0/td.gameNode.defenderState.possibleMoves.length.toDouble
                val newQueue =   if(mwl > 0)
                  td.gameNode.defenderState.possibleMoves.map(m =>
                    (m, td.defenderStrategy.probabilities.find(m == _._1).map(_._2).getOrElse(0.0)))
                    .foldLeft(tail){case (list, (dm, prob)) => {
                      val nextState = td.gameNode.makeMove(dm, td.attackerStrategy.move)
                      val nextDS = td.defenderStrategy.nextState(dm, nextState.defenderState)
                      val nextAS = td.attackerStrategy.nextState(nextState.attackerState)
                      TraverseData(nextState, nextDS, nextAS, mwl, prob, branching) :: list
                    }} else tail
                        
                calculateSum(newQueue, currentSum)
              }
            }
            case Nil => currentSum
          }
        }
        calculateSum(List(TraverseData(gameState, defenderStrategy, attackerStrategy, 1.0, 1.0, 1.0)), DistortedPayoff.zero)
      }
    }
  }

}
