package pl.edu.pw.mini.sg.game.results

import argonaut.{ CodecJson, DecodeJson, DecodeResult, EncodeJson, HCursor, Json }
import java.util.UUID
import pl.edu.pw.mini.sg.game.Payoff
import pl.edu.pw.mini.sg.game.boundedrationality.DistortedPayoff
import pl.edu.pw.mini.sg.game.deterministic.json.DeterministicGameJsonDecoders
import pl.edu.pw.mini.sg.game.deterministic.plugin.GamePluginDecoding
import pl.edu.pw.mini.sg.game.deterministic.{BehavioralStrategy, PlayerVisibleState}
import squants.time.{Seconds, Time}

final case class TrainingResult[DM, DS, MI](
  payoff: Payoff,
  defenderStrategy: BehavioralStrategy[DM, DS],
  methodInfo: MI,
  trainingTime: Time,
  preprocessTime: Time,
  cpuInfo: String,
  game: GameDescriptor,
  distortedPayoff : Option[Payoff]
)

object TrainingResult {
  val PAYOFF_KEY = "payoff"
  val DEFENDER_STRATEGY_KEY = "defenderStrategy"
  val METHOD_INFO_KEY = "methodInfo"
  val TRAINING_TIME_KEY = "trainingTime"
  val PREPROCESS_TIME_KEY = "preprocessTime"
  val CPUINFO_KEY = "cpuInfo"
  val GAME_KEY = "game"
  val DISTORTED_PAYOFF_KEY = "distortedPayoff"

  val secondsCodec : CodecJson[Time] = CodecJson.derived[Double].xmap[Time](Seconds(_))(_.toSeconds)


  def apply[DM, DS <: PlayerVisibleState[DM], MI](
    payoff: DistortedPayoff,
    defenderStrategy: BehavioralStrategy[DM, DS],
    methodInfo: MI,
    trainingTime: Time,
    preprocessTime: Time,
    cpuInfo: String,
    game: GameDescriptor
  ) : TrainingResult[DM, DS, MI] = apply[DM, DS, MI](payoff.real, defenderStrategy, methodInfo, trainingTime, preprocessTime,
    cpuInfo, game, Some(payoff.distortedAttacker))

  implicit def trainingResultCodec[DM, DS, MI](implicit dmEncode: EncodeJson[DM],
                                                                         dsEncode: EncodeJson[DS], miEncode: EncodeJson[MI]
                                                                        ) : EncodeJson[TrainingResult[DM, DS, MI]] = EncodeJson.jencode8L(
    (TrainingResult.unapply[DM, DS, MI]_).andThen(_.get)
  )(PAYOFF_KEY, DEFENDER_STRATEGY_KEY, METHOD_INFO_KEY, TRAINING_TIME_KEY,
    PREPROCESS_TIME_KEY, CPUINFO_KEY, GAME_KEY, DISTORTED_PAYOFF_KEY
  )(
    Payoff.payoffCodec,
    BehavioralStrategy.strategyEncode,
    miEncode,
    secondsCodec,
    secondsCodec,
    EncodeJson.StringEncodeJson,
    GameDescriptor.gameDescriptorCodec,
    EncodeJson.OptionEncodeJson(Payoff.payoffCodec)
  )

  implicit def behavioralStrategyDecode[DM, DS <: PlayerVisibleState[DM]](
    implicit dm: DecodeJson[DM], ds: DecodeJson[DS]
  ): DecodeJson[ImmutableBehavioralStrategy[DM, DS]] = {
    def dec(hc: HCursor) : DecodeResult[ImmutableBehavioralStrategy[DM, DS]] = for {
      state <- (hc --\ BehavioralStrategy.STATE_KEY).as[DS]
      probabilities <- (hc --\ BehavioralStrategy.PROBABILITIES_KEY).as[Vector[(DM, Double)]]
      successors <- (hc --\ BehavioralStrategy.SUCCESSORS_KEY).as[Vector[(DM, ImmutableBehavioralStrategy[DM, DS])]]
    } yield (ImmutableBehavioralStrategy(state, probabilities, successors.map{ case (dm, str) =>
        ((dm, str.state), str)}.toMap))

    DecodeJson(dec)
  }


  def trainingResultDecode(dec : GamePluginDecoding)(gameConfigs : Map[UUID, dec.GameConfig]) :
      DecodeJson[TrainingResult[dec.singleCoordFactory.DefenderMove, dec.singleCoordFactory.DefenderState, Json]] =
    DecodeJson(hc => {
      for {
        payoff <- (hc --\ PAYOFF_KEY).as[Payoff]
        dp <- (hc --\ DISTORTED_PAYOFF_KEY).as[Payoff] ||| DecodeResult.ok(new Payoff(Double.NaN, Double.NaN))
        gameDescriptor <- (hc --\ GAME_KEY).as[GameDescriptor]
        decoders <- gameConfigs.get(gameDescriptor.id).map(c => dec.singleCoordFactory.jsonDecoders(c)) match {
          case None => DecodeResult.fail(s"Unknown game id: ${gameDescriptor.id}", hc.history)
          case Some(d) => DecodeResult.ok(d)
        }
        defenderStrategy <- (hc --\ DEFENDER_STRATEGY_KEY).as(behavioralStrategyDecode[dec.singleCoordFactory.DefenderMove, dec.singleCoordFactory.DefenderState](decoders.defenderMoveDecode, decoders.defenderStateDecode))
        methodInfo <- (hc --\ METHOD_INFO_KEY).as[Json]
        trainingTime <- (hc --\ TRAINING_TIME_KEY).as[Time](secondsCodec)
        preprocessTime <- (hc --\ PREPROCESS_TIME_KEY).as[Time](secondsCodec)
        cpuinfo <- (hc --\ CPUINFO_KEY).as[String]
      } yield (TrainingResult(new DistortedPayoff(payoff, dp), defenderStrategy, methodInfo, trainingTime, preprocessTime, cpuinfo, gameDescriptor))
    })
}
