package pl.edu.pw.mini.sg.movingtargets.aggregate

import argonaut.{ JsonIdentity }
import canes.game.TrainingResultJson
import canes.game.attacker.OptimalAttacker
import canes.game.{ AllUnitsPositions, MethodInfo, MixedStrategyEntry, PureDefenderStrategy, Strategy, StrategyJson, TrainingResult }
import pl.edu.pw.mini.jk.sg.game.NormalizableFactory
import pl.edu.pw.mini.jk.sg.game.immutable.mutable.MutableWrapper
import pl.edu.pw.mini.sg.movingtargets.MovingTargetsAttackerState
import pl.edu.pw.mini.sg.movingtargets.{ AttackerDecision, MovingTargetsGameConfig }
import squants.time.Time

import scala.collection.JavaConverters._

final class AggregatedNormalizableFactory(
  config: MovingTargetsGameConfig,
  aggregation: Int
) extends NormalizableFactory[
  AttackerDecision, AggregateMTGDefenderMove,
  MutableWrapper[
    AttackerDecision, AggregateMTGDefenderMove,
    MovingTargetsAttackerState, AggregateMTGDefenderState,
    AggregateMovingTargetsGame
  ]
]
{
  private val game = AggregateMovingTargetsGame(0, 0, AggregateMTGDefenderState.newGame(config, aggregation),
    MovingTargetsAttackerState.createState(false, config.targets.map(_.id).toList), None, Vector.fill(aggregation*aggregation)(false))
  def create: MutableWrapper[AttackerDecision,AggregateMTGDefenderMove,
    MovingTargetsAttackerState,AggregateMTGDefenderState,AggregateMovingTargetsGame] = new MutableWrapper(game)
  def defenderResultAsJson(strategy: java.util.Map[_ <: java.util.List[AggregateMTGDefenderMove], java.lang.Double],
    trainingTimeS: Double, methodName: String, methodConfig: argonaut.Json,
    preprocessingTime: Time): argonaut.Json = {
    val str = prepareDefenderStrategy(strategy)
    val optimalAttacker = OptimalAttacker.calculateStrategy(config, str)
    val payoffs = OptimalAttacker.getPayoff(config)(str, optimalAttacker)
    val tr = TrainingResult(str, trainingTimeS, payoffs.defenderPayoff, payoffs.attackerPayoff,
      MethodInfo.apply(methodName, methodConfig), Some(optimalAttacker), preprocessingTime.toSeconds
    )

    JsonIdentity.ToJsonIdentity(tr).asJson(TrainingResultJson.TrainingResultCodec)
  }

  def prepareDefenderStrategy(strategy: java.util.Map[_ <: java.util.List[AggregateMTGDefenderMove], java.lang.Double]): Strategy = {
        def sequenceToPureStrategies(sequence: java.util.List[AggregateMTGDefenderMove]) : Vector[PureDefenderStrategy] = {

          val aupVec: Vector[Vector[AllUnitsPositions]] = sequence.asScala.foldLeft((game.defenderVisibleState,Vector.fill(aggregation*aggregation)(Vector[AllUnitsPositions]()))) { case ((ds,vec), agm) => {
            val nextStep = ds.advance(agm)
            val vv = vec.zip(nextStep.states).map { case (vec, state) => vec:+AllUnitsPositions(state.currentPositions.view.map(config.toDoublePositions).toList) }
            (nextStep, vv)
          }}._2
          aupVec.map(PureDefenderStrategy.apply)
        }

    val msEntries = strategy.asScala.flatMap { case (str, prob) => {
      sequenceToPureStrategies(str).map(s => MixedStrategyEntry(prob/(aggregation*aggregation), s))
    }}

  Strategy(config.id,
    msEntries.groupBy(_.strategy).map { case (str, entries) =>  MixedStrategyEntry(entries.view.map(_.probability).sum, str)}.toVector
  )
  }

  def defenderStrategyToJson(strategy: java.util.Map[_ <: java.util.List[AggregateMTGDefenderMove], java.lang.Double]): argonaut.Json = {
    val strategyFin = prepareDefenderStrategy(strategy)
    JsonIdentity.ToJsonIdentity(strategyFin).asJson(StrategyJson.strategyCodec)
  }
  val denormalize: AggregatedNormalizableFactory = this
  def normalize: AggregatedNormalizableFactory = this
}
