package pl.edu.pw.mini.sg.stat

import argonaut.{ Json, JsonIdentity }
import better.files.File
import org.rogach.scallop.ScallopConf
import pl.edu.pw.mini.sg.game.deterministic.{ AdvanceablePlayerVisibleState, DeterministicGame, PlayerVisibleState }


final case class GameStats(
  defenderIS:  Int,
  attackerIS: Int,
  leaves: Int,
  leavesDepths: Map[Int, Int],
  nodes: Int,
  minPayoff: Double,
  maxPayoff: Double,
  pDiffMean: Double,
  pDiffSd: Double,
  pCorrelation: Double
)

object GameStat {
  private final class Options(args: Seq[String]) extends ScallopConf(args) {
    val games = opt[List[String]](name = "games", short = 'g', required=true).map(_.map(s => File.apply(s)))
    val jsonOutput = toggle(name="json-output", short='j', default = Some(false))

    verify()
  }

  def main(args: Array[String]) : Unit = {
    val gamePlugins = pl.edu.pw.mini.sg.game.deterministic.Plugins.defaultPlugins
    val options = new Options(args)

    options.games().foreach(game => {
      gamePlugins.find(plugin => plugin.fileExtension == game.extension(includeDot = false).getOrElse("")) match {
        case Some(pl) => {
          val gc = pl.readConfig(game.toJava).get
          val g = pl.singleCoordFactory.create(gc)
          val stats = gameStats[
            pl.singleCoordFactory.DefenderMove,
            pl.singleCoordFactory.AttackerMove,
            pl.singleCoordFactory.DefenderState,
            pl.singleCoordFactory.AttackerState,
            pl.singleCoordFactory.Game
          ](g)

          if(options.jsonOutput()) {
            val out = Json.obj(
              "defenderIs" -> Json.jNumber(stats.defenderIS),
              "attackerIs" -> Json.jNumber(stats.attackerIS),
              "leaves" -> Json.jNumber(stats.leaves),
              "nodes" -> Json.jNumber(stats.nodes),
              "leavesDepths" -> JsonIdentity.ToJsonIdentity(stats.leavesDepths.map{case (a,b) => (a.toString, b)}).asJson,
              "gameId" -> JsonIdentity.ToJsonIdentity(gc.id).asJson,
              "minPayoff" -> Json.jNumber(stats.minPayoff),
              "maxPayoff" -> Json.jNumber(stats.maxPayoff),
              "payoffDiffMean" -> Json.jNumber(stats.pDiffMean),
              "payoffDiffSd" -> Json.jNumber(stats.pDiffSd),
              "payoffCorrelation" -> Json.jNumber(stats.pCorrelation)
            )
            println(out.spaces2)
          } else {
            println(s"DefenderIS: ${stats.defenderIS}, AttackerIs: ${stats.attackerIS} leavesCount: ${stats.leaves}, nodes: ${stats.nodes}\n leavesDephts: ${stats.leavesDepths}")
            println(s"AttackerIs 2: ${calculateIsCount[pl.singleCoordFactory.AttackerMove, pl.singleCoordFactory.AttackerState](g.attackerState)}")
            println(s"Payoff stats: min: ${stats.minPayoff}, max: ${stats.maxPayoff}, diff: m=${stats.pDiffMean}, sd=${stats.pDiffSd}, correlation: ${stats.pCorrelation}")
          }
        }
        case None => println("No plugin found!")
      }
    })

  }


  def calculateIsCount[M, S <: AdvanceablePlayerVisibleState[M, S]](state: S) : Int = {
    if(state.possibleMoves.isEmpty) 0
    else 1+state.possibleMoves.flatMap(m => state.possibleContinuations(m).map(calculateIsCount[M,S])).sum
  }

  def gameStats[DM, AM, DS <: PlayerVisibleState[DM], AS <: PlayerVisibleState[AM],
    G <: DeterministicGame[DM, AM, DS, AS, G]](game : G) = {
    final case class PartialResult(
      defenderIS: Set[(List[DM], List[DS])] = Set(),
      attackerIS: Set[(List[AM], List[AS])] = Set(),
      leaves: Int = 0,
      leavesDepths: Map[Int, Int] = Map(),
      nodes: Int = 0,
      maxPayoff: Double = Double.NegativeInfinity,
      minPayoff: Double = Double.PositiveInfinity,
      differenceMean: Double = 0,
      defenderMean: Double = 0,
      attackerMean: Double = 0,
      differenceS: Double = 0
    ) {
      def registerLeaf(depth: Int, node: G) = {
        val payoffDiff = node.defenderPayoff - node.attackerPayoff
        val newDifferenceMean = if(leaves == 0) payoffDiff else differenceMean+(payoffDiff-differenceMean)/(leaves+1)
        copy(leaves = leaves + 1, leavesDepths = leavesDepths.updated(depth,
        leavesDepths.get(depth).map(_+1).getOrElse(1)
      ), nodes = nodes+1, maxPayoff = Math.max(Math.max(node.defenderPayoff, node.attackerPayoff), maxPayoff),
        minPayoff = Math.min(minPayoff, Math.min(node.defenderPayoff, node.attackerPayoff)),
          differenceMean = newDifferenceMean,
          differenceS = differenceS + (payoffDiff-differenceMean)*(payoffDiff-newDifferenceMean),
          defenderMean = if(leaves == 0) node.defenderPayoff else defenderMean+(node.defenderPayoff-defenderMean)/(leaves+1),
          attackerMean = if(leaves == 0) node.attackerPayoff else attackerMean+(node.attackerPayoff-attackerMean)/(leaves+1)
          )
      }

      def registerNode(defenderHistory: (List[DM], List[DS]), attackerHistory: (List[AM], List[AS]),
        defenderIS: DS, attackerIS: AS) = copy(
        defenderIS = this.defenderIS + ((defenderHistory._1, defenderIS::defenderHistory._2)),
          attackerIS = this.attackerIS + ((attackerHistory._1, attackerIS::attackerHistory._2)),
          nodes = nodes+1
      )
    }

    val res = browseGameNodes[DM, AM, DS, AS, G](game, Nil, Nil, Nil, Nil, 0).foldLeft(PartialResult()){
      case (pr, (g, dm, ds, am, as, depth)) =>
        if(g.isLeaf) pr.registerLeaf(depth, g)
        else pr.registerNode((dm, ds), (am, as), g.defenderState, g.attackerState)
    }

    val (covariantion, varD, varA) = browseGameNodes[DM, AM, DS, AS, G](game, Nil, Nil, Nil, Nil, 0).foldLeft((0.0, 0.0, 0.0)){
      case ((cov, vd, va), (g, dm, ds, am, as, depth)) =>
        if(g.isLeaf) (
          cov+(g.defenderPayoff-res.defenderMean)*(g.attackerPayoff-res.attackerMean),
          vd+(g.defenderPayoff-res.defenderMean)*(g.defenderPayoff-res.defenderMean),
          va+(g.attackerPayoff-res.attackerMean)*(g.attackerPayoff-res.attackerMean),
        ) else (cov, vd, va)
    }

    val correlation =
      if (varA != 0) covariantion/Math.sqrt(varD)/Math.sqrt(varA)
      else 1

    GameStats(
      res.defenderIS.size, res.attackerIS.size,
      res.leaves, res.leavesDepths, res.nodes,
      res.minPayoff, res.maxPayoff,
      res.differenceMean,
      Math.sqrt(res.differenceS/(res.leaves-1)),
      correlation
    )
  }

  def browseGameNodes[DM, AM, DS <: PlayerVisibleState[DM], AS <: PlayerVisibleState[AM],
    G <: DeterministicGame[DM, AM, DS, AS, G]](
    game : G, dmh: List[DM], dsh: List[DS], amh: List[AM], ash: List[AS], depth: Int
  ) : Stream[(G, List[DM], List[DS], List[AM], List[AS], Int)]  = {
    (game, dmh, dsh, amh, ash, depth) #:: (
      if(game.isLeaf) Stream()
      else game.defenderState.possibleMoves.toStream.flatMap(dm =>
        game.attackerState.possibleMoves.toStream.flatMap(am =>
          browseGameNodes(game.makeMove(dm, am), dm::dmh, game.defenderState::dsh, am::amh, game.attackerState::ash, depth+1)
      ))
    )
  }
}
