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

import pl.edu.pw.mini.sg.game.boundedrationality.DistortedPayoff
import pl.edu.pw.mini.sg.game.deterministic.plugin.GameParams
import pl.edu.pw.mini.sg.mixed.response.{AttackerResponseFinder, HistorySkipConfig, OptimalAttackerResponse}
import pl.edu.pw.mini.sg.mixed.tree.InnerResultNode

import scala.annotation.tailrec
import pl.edu.pw.mini.sg.rng.RandomGenerator
import org.log4s._
import pl.edu.pw.mini.sg.game.Payoff
import pl.edu.pw.mini.sg.game.boundedrationality.BoundedRationality
import pl.edu.pw.mini.sg.game.deterministic.{AdvanceablePlayerVisibleState, DeterministicGame, PlayerVisibleState, PureStrategy}
import pl.edu.pw.mini.sg.game.deterministic.single.{DefenderGameView, LinearizedGameProvider}
import pl.edu.pw.mini.sg.game.util.attacker.OptimalAttacker
import pl.edu.pw.mini.sg.mcts.MctsAlgorithm
import pl.edu.pw.mini.sg.mcts.tree.{MctsNode, MctsTree}
import pl.edu.pw.mini.sg.mcts.tree.stat.MoveEstimate
import pl.edu.pw.mini.sg.mixed.tree.{CalculateResults, InnerResultNode, MixedStrategyNode}

import scalaz.NonEmptyList

sealed trait OracleResult[DM, DS <: PlayerVisibleState[DM]]
final case class FeasibleResult[DM, DS <: PlayerVisibleState[DM]](
  val result: MixedStrategyNode[DM, DS, Unit], val payoff: DistortedPayoff
) extends OracleResult[DM, DS]
final case class InfeasibleResult[DM, DS <: PlayerVisibleState[DM]](
  val result: MixedStrategyNode[DM, DS, Unit], val payoff: DistortedPayoff
) extends OracleResult[DM, DS]

object UctAttackerSampler {
  private val logger = getLogger

  val ATTACKER_EQUALS_EPS = 1e-6

  @tailrec
  def removePrefix[T](prefix: Vector[T], list: List[T]) : Option[List[T]] = {
    prefix match {
      case Vector() => Some(list)
      case x +: xs =>
        list match {
          case y::ys if y == x => removePrefix(xs, ys)
          case _ => None
        }
    }
  }

   // Czy propagować coś o atakującycm/obrońcy do drzewa?
  def findBehavioralStrategy[
    DM, AM,
    DS <: PlayerVisibleState[DM],
    AS <: AdvanceablePlayerVisibleState[AM, AS],
    G <: DeterministicGame[DM, AM, DS, AS, G]
  ](
    game: G, gameDescriptor: GameParams, rng: RandomGenerator,
    uctAttackerSamplerConfig: UctAttackerSamplerConfig
  ) : (MixedStrategyNode[DM, DS, Unit], DistortedPayoff) = {
    val linearizedAttackerGame = uctAttackerSamplerConfig.linearizedGameProvider.create[AM, AS](game.attackerState)
    val tree : MctsTree[AM, AS] = uctAttackerSamplerConfig.mctsConfig.treeFactory.create(
      linearizedAttackerGame.playerState, linearizedAttackerGame.possibleMoves)
    val infeasiblePayoff = (gameDescriptor.maxDefenderPayoff-gameDescriptor.minDefenderPayoff)*uctAttackerSamplerConfig.infeasiblePenalty+gameDescriptor.minDefenderPayoff
    var bestDS : Option[(MixedStrategyNode[DM, DS, Unit], DistortedPayoff)] = None

    def updateBestDS(str: MixedStrategyNode[DM, DS, Unit], payoff: DistortedPayoff) : Boolean = {
      bestDS match {
        case Some((_, bestPayoff)) if payoff.real.defender <= bestPayoff.real.defender => false
        case _ => bestDS = Some((str, payoff)); true
      }
    }

    def attackerStrategyEvaluation(prefix: Vector[AM])(attackerLeaf: uctAttackerSamplerConfig.linearizedGameProvider.Game[AM, AS]) : (Double, Option[(List[AM], Double)]) = {
      logger.info("Outer probe iteration")
      val desiredAttacker = attackerLeaf.pureStrategy
      val mixedMethodConfig = uctAttackerSamplerConfig.mixedMethodConfig
      implicit val br = mixedMethodConfig.boundedRationalityStrategy

      val defender: StrategyResult[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState] = defenderForGivenAttacker[DM, AM, DS, AS, G](mixedMethodConfig)(game, desiredAttacker)(new MixedStrategyNode[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState](game.defenderState)(mixedMethodConfig.strategyAdjuster.ost))(rng)

      maybeFeasibleToFeasibleStrategy[DM,AM,DS,AS,G](mixedMethodConfig)(game, desiredAttacker)(defender)(rng) match {
        case InfeasibleResult(str, payoff) => {
          if(updateBestDS(str.deepClone, payoff)) logger.debug(s"Updated best with infeasible")
          (infeasiblePayoff,
            uctAttackerSamplerConfig.intfeasibleAttackerPropagation match {
              case JustPropagation =>  removePrefix(prefix, 
                uctAttackerSamplerConfig.linearizedGameProvider.pureStrategyToLinearizedGameMoveSequence[AM, AS](
                  game.attackerState, OptimalAttacker.calculate[DM,AM,DS,AS,G](game, str)._1
                )).map(x => (x, payoff.real.defender))
              case RecalcPropagation => {
                val (ds, payoff, as) = defenderImproveUntilStableAttacker[DM, AM, DS, AS, G](uctAttackerSamplerConfig.mixedMethodConfig, uctAttackerSamplerConfig.linearizedGameProvider)(game, str)(rng)
                updateBestDS(ds.deepClone, payoff)
                removePrefix(prefix, as).map(x => (x, payoff.real.defender))
              }
              case NoPropagation => None
          })
        }
        case FeasibleResult(str, payoff) => {
          if(updateBestDS(str.deepClone, payoff)) logger.debug(s"Updated best with feasible")

          lazy val (oas, oap) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, str)
          (payoff.real.defender,
            {
              uctAttackerSamplerConfig.intfeasibleAttackerPropagation match {
                case JustPropagation if(oap.real.defender > payoff.real.defender) =>
                  updateBestDS(str.deepClone, oap)
                  removePrefix(prefix, uctAttackerSamplerConfig.linearizedGameProvider.pureStrategyToLinearizedGameMoveSequence[AM, AS](game.attackerState, oas)).map(x => 
                  (x, oap.real.defender))
                case RecalcPropagation if oap.real.defender > payoff.real.defender => {
                  val (ds, payoff, as) = defenderImproveUntilStableAttacker[DM, AM, DS, AS, G](uctAttackerSamplerConfig.mixedMethodConfig, uctAttackerSamplerConfig.linearizedGameProvider)(game, str)(rng)
                  updateBestDS(ds.deepClone, payoff)
                  removePrefix(prefix, as).map(x => (x, payoff.real.defender))
                }
                case _ => None
              }
            })
        }
      }
    }

    logger.info("Search started")
    if(!uctAttackerSamplerConfig.descent) {
      MctsAlgorithm.training(tree.root, linearizedAttackerGame, attackerStrategyEvaluation(Vector()), uctAttackerSamplerConfig.mctsConfig, rng)
    } else {
      @tailrec
      def samplePathEqually(tree : MctsNode[AM, AS], game: uctAttackerSamplerConfig.linearizedGameProvider.Game[AM, AS], prefix: Vector[AM]) : Unit = {
        if(game.possibleMoves.nonEmpty) {
          MctsAlgorithm.training(tree, game, attackerStrategyEvaluation(prefix), uctAttackerSamplerConfig.mctsConfig, rng)
          val MoveEstimate(m, _, _) = tree.estimates.maxBy(a => a.estimate)
          val nextState = game.makeMove(m)
          samplePathEqually(tree.successor(m, nextState.playerState).get, nextState, prefix :+ m)
        } else {
          ()
        }
      }
      samplePathEqually(tree.root, linearizedAttackerGame, Vector())
    }

    def verifyPayoffs(output: (MixedStrategyNode[DM, DS, Unit], DistortedPayoff))(implicit boundedRationality: BoundedRationality) : Unit = {
      val (str, pa) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, output._1)
      println(s"Found: ${output._2}, calculated: $pa")
      require(Math.abs(pa.distortedAttacker.attacker - output._2.distortedAttacker.attacker) < 1e-2, s"To large payoff gap returned: ${output} verified: ${pa}")
    }

    verifyPayoffs(bestDS.get)(uctAttackerSamplerConfig.mixedMethodConfig.boundedRationalityStrategy)
    bestDS.map{case (a, b) => (a, b)}.get
  }

  @tailrec
  def defenderImproveUntilStableAttacker[
    DM, AM,
    DS <: PlayerVisibleState[DM],
    AS <: AdvanceablePlayerVisibleState[AM, AS],
    G <: DeterministicGame[DM, AM, DS, AS, G]
  ](
    mixedMethodConfig : MixedMethodConfig,
    lgp: LinearizedGameProvider
  )(
    game: G,
    defenderStrategy : MixedStrategyNode[DM, DS, _]
  )(
    implicit rng: RandomGenerator
  ) : (MixedStrategyNode[DM, DS, Unit], DistortedPayoff, List[AM]) = {
    implicit val br = (mixedMethodConfig.boundedRationalityStrategy)
    val (attacker, attackerPayoff) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, defenderStrategy)
    maybeFeasibleToFeasibleStrategy[DM, AM, DS, AS, G](mixedMethodConfig)(game, attacker)(defenderForGivenAttacker[DM, AM, DS, AS, G](mixedMethodConfig)(game, attacker
    )(defenderStrategy.deepClone(mixedMethodConfig.strategyAdjuster.ost))) match {
      case FeasibleResult(strategy, payoff) => {
        val (optA, oap) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, strategy)
        if(oap.distortedAttacker.attacker - 1e-3 > payoff.distortedAttacker.attacker)
          defenderImproveUntilStableAttacker[DM, AM, DS, AS, G](
            mixedMethodConfig,lgp)(game, strategy.deepClone(mixedMethodConfig.strategyAdjuster.ost))
        else
          (strategy, payoff, lgp.pureStrategyToLinearizedGameMoveSequence[AM, AS](game.attackerState,optA))
      }
      case InfeasibleResult(strategy, payoff) => defenderImproveUntilStableAttacker[DM, AM, DS, AS, G](
        mixedMethodConfig,lgp)(game, strategy.deepClone(mixedMethodConfig.strategyAdjuster.ost))
    }
  }

  def maybeFeasibleToFeasibleStrategy[
    DM, AM,
    DS <: PlayerVisibleState[DM],
    AS <: AdvanceablePlayerVisibleState[AM, AS],
    G <: DeterministicGame[DM, AM, DS, AS, G]
  ](
    mixedMethodConfig: MixedMethodConfig
  )(
    game: G,
    desiredAttacker: PureStrategy[AM, AS]
  )(
    defenderStrategy: StrategyResult[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState]
  )(
    implicit rng: RandomGenerator
  ): OracleResult[DM, DS] = {
    implicit val br = mixedMethodConfig.boundedRationalityStrategy
    defenderStrategy match {
      case FeasibleStrategy(strategy, payoff) => FeasibleResult(strategy.deepClone, payoff)
      case InfeasibleStrategy(strategy, payoff) => InfeasibleResult(strategy.deepClone, payoff)
      case MaybeFeasibleStrategy(strategy, payoffUB, iterationInfo) => {
        val (exactAttacker, payoff) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, strategy)
        if(payoff.distortedAttacker.attacker - ATTACKER_EQUALS_EPS < payoffUB.distortedAttacker.attacker)
          FeasibleResult(strategy.deepClone, payoff)
        else {
          val desiredAttackerPayoff = new InnerResultNode(strategy, game.defenderState)(mixedMethodConfig.boundedRationalityStrategy)
          CalculateResults.calculate(game, strategy, desiredAttacker, desiredAttackerPayoff, mixedMethodConfig.expansionDecision.positiveExpander, true)
          MixedMethod.feasibilityLoop(mixedMethodConfig)(game, strategy, desiredAttackerPayoff, desiredAttacker, OptimalAttackerResponse(HistorySkipConfig.noSkip), StableAttackerFinder.zeroState, iterationInfo.improvementIteration+1) match {
            case Left(strategy) => InfeasibleResult(strategy.deepClone,
              OptimalAttacker.calculate[DM, AM, DS, AS, G](game, strategy)._2)
            case Right((strategy, payoff, _)) => FeasibleResult(strategy.deepClone, payoff)
          }
        }
      }
    }
  }

  def defenderForGivenAttacker[
    DM, AM,
    DS <: PlayerVisibleState[DM],
    AS <: AdvanceablePlayerVisibleState[AM, AS],
    G <: DeterministicGame[DM, AM, DS, AS, G]
  ](
    mixedMethodConfig: MixedMethodConfig
  )(
    game: G,
    desiredAttacker: PureStrategy[AM, AS]
  )(
    defenderStrategy: MixedStrategyNode[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState] =
      new MixedStrategyNode[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState](game.defenderState)(mixedMethodConfig.strategyAdjuster.ost)
  )(
    implicit rng: RandomGenerator
  ): StrategyResult[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState] = {
    implicit val br = mixedMethodConfig.boundedRationalityStrategy
      val desiredAttackerPayoff = new InnerResultNode(defenderStrategy, game.defenderState)(mixedMethodConfig.boundedRationalityStrategy)
      CalculateResults.calculate(game, defenderStrategy, desiredAttacker, desiredAttackerPayoff, mixedMethodConfig.expansionDecision.positiveExpander, true)
      desiredAttackerPayoff.expandToLeaf(
        new DefenderGameView[DM, AM, DS, AS, G](game, desiredAttacker),
        mixedMethodConfig.expansionDecision.initialExpander
      )(rng, IterationInfo(-1, -1))

    import mixedMethodConfig.strategyAdjuster.AdjusterState
      @tailrec
    def iterativeImprovement(allIterations: Int, poorImprovementIterations: Int, bestDefenderStrategy: Option[(MixedStrategyNode[DM, DS, mixedMethodConfig.strategyAdjuster.AdjusterState], DistortedPayoff, Boolean)], afState: StableAttackerFinderState[AM, AS]) : StrategyResult[DM, DS,AdjusterState] = {
      val currentIteration = mixedMethodConfig.allRefinementIterations - allIterations
      implicit val ii : IterationInfo = IterationInfo(-1, currentIteration)
        if(allIterations <= 0 || poorImprovementIterations >= mixedMethodConfig.poorImprovementIterations)
          bestDefenderStrategy.map{case ((a, b, c)) => StrategyResult.feasibleResult[DM, DS, AdjusterState](a,b, c)}.get
        else {
          val attackerOracle: AttackerResponseFinder = mixedMethodConfig.attackerOracleProvider.attackerResponse(mixedMethodConfig.allRefinementIterations - allIterations)
          MixedMethod.strategyImprovementIteration(mixedMethodConfig)(
            game, desiredAttacker, desiredAttackerPayoff, defenderStrategy, afState, attackerOracle, currentIteration)(rng) match {
            case Left(str) => bestDefenderStrategy.map{ case (a,b,c) => StrategyResult.feasibleResult[DM, DS, AdjusterState](a,b,c)}
                .getOrElse({
                  val (oa, payoff) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, str)
                  InfeasibleStrategy(str, payoff)
                })
            case res@Right((ds, payoff, state)) => {
              logger.debug(s"Found: ${payoff}")
              bestDefenderStrategy match {
                case Some((_, bestPayoff, _)) if bestPayoff.real.defender >= payoff.real.defender => 
                  iterativeImprovement(allIterations,
                    poorImprovementIterations + 1,
                    bestDefenderStrategy, state)
                
                case _ => iterativeImprovement(allIterations - 1,
                  if(payoff.real.defender - bestDefenderStrategy.map(_._2.real.defender).getOrElse(0.0) < mixedMethodConfig.poorImprovementEps)
                      poorImprovementIterations + 1
                    else 0
                      , Some((ds, payoff, attackerOracle.isExact)), state)
              }
            }
          }
        }
      }

      val firstResult = {
          val (response, responsePayoff) = OptimalAttacker.calculate[DM, AM, DS, AS, G](game, defenderStrategy)
          if(Math.abs(responsePayoff.distortedAttacker.attacker - desiredAttackerPayoff.payoff.distortedAttacker.attacker) < ATTACKER_EQUALS_EPS)
            Some((defenderStrategy.deepClone(mixedMethodConfig.strategyAdjuster.ost), responsePayoff, true))
          else
            None
        }

      {
        iterativeImprovement(
          mixedMethodConfig.allRefinementIterations, 0, firstResult,
          StableAttackerFinder.zeroState
        ) match {
          case r@InfeasibleStrategy(str, payoff) => {
            logger.debug("Strategy not found")
            r
          }
          case som@FeasibleStrategy(str, payoff) => {
            logger.debug(s"Strategy found with payoff ${payoff}")
            som
          }
          case ma@MaybeFeasibleStrategy(str, payoff, _) => {
            logger.debug(s"Strategy (inexact) found with payoff ${payoff}")
            ma
          }
        }
      } 
  }
}
