package pl.edu.pw.mini.sg.gamegen

import argonaut.Parse
import java.util.UUID
import java.util.stream.Collectors

import scala.collection.JavaConverters._

import argonaut.JsonIdentity._
import better.files.File
import org.apache.commons.math3.random.{MersenneTwister, RandomDataGenerator, RandomGenerator}
import org.rogach.scallop.ScallopConf
import pl.edu.pw.mini.sg.game.graph.{DeterministicGraphGameConfig, DeterministicGraphGameConfigJson, Edge, GraphConfig, GraphGamePlugin}

object BuildingLikeGraphGeneratorDiversePenalty {
  def generateGame(config: BuildingLikeGraphGeneratorDiversePenaltyConfig)(
    implicit rng: RandomGenerator
  ) : (DeterministicGraphGameConfig, BuildingGraphStructure) = {
    val oldGen = new BuildingLikeGameGenerator(rng, config.toBLGGC)
    val structure = oldGen.generateRandomStructure()
    val nodes = {
      val nodesi = structure.getAllNodes.collect(Collectors.toList()).asScala.toVector
      val offices : Vector[Int] = nodesi.zipWithIndex.filter{case (nc, _) => structure.getNodeRole(nc) == NodeRole.OFFICE}.map(_._2)
      val start = rng.nextInt(offices.size)
      nodesi.updated(start, nodesi(0)).updated(0, nodesi(start))
    }

    val spawn = nodes.zipWithIndex.find { case (bc, idx) => structure.getNodeRole(bc) == NodeRole.ENTRANCE }.get._2

    val offices : Vector[Int] = nodes.zipWithIndex.filter{case (nc, idx) => idx > 0 && structure.getNodeRole(nc) == NodeRole.OFFICE}.map(_._2)
    val dg = new RandomDataGenerator(rng)
    val targetIds: Array[Int]  = dg.nextSample(offices.asJava, config.targetCount).map(_.asInstanceOf[java.lang.Integer]+0)

    val edges = (nodes.zipWithIndex).flatMap{ case (n1, i1) => (nodes.zipWithIndex).collect{
      case (n2, i2) if structure.areNodesConnected(n1, n2) => Edge(i1, i2)
    }}
    val graphConfig = GraphConfig(structure.getAllNodes.count().toInt, edges)
    val defenderPenalties = Vector.fill(config.targetCount)(
      config.defenderPenalty.random
    )
    val attackerRewards = Vector.fill(config.targetCount)(
      config.attackerReward.random
    )

    val defenderRewards = (nodes.zipWithIndex).map{case (bc, idx) =>
      if(targetIds.contains(idx)) config.defenderSuccessRewardTarget
      else config.defenderSuccessRewardNormal
    }
    val attackerPenalties = (nodes.zipWithIndex).map{case (bc, idx) =>
      if(targetIds.contains(idx)) config.attackerCaughtPenaltyTarget.random
      else config.attackerCaughtPenaltyNormal.random
    }

    (DeterministicGraphGameConfig(graphConfig, targetIds.toList, spawn, defenderRewards, attackerPenalties, defenderPenalties, attackerRewards, config.numberOfDefenders, config.rounds, UUID.randomUUID()), structure)
  }

  final class Options(args: Seq[String]) extends ScallopConf(args) {
    val generatorConfig = opt[String](name = "config", short='c', required=true).map(x => File(x))
    val outputPrefixes = opt[List[String]](name = "output", short = 'o', required=true)
    val structureSuffix = opt[String](name = "structure-sufix", short='s', default = Some(".structure.txt"))

    verify()
  }

  def main(args: Array[String]): Unit = {
    val options = new Options(args)
    Parse.decode[BuildingLikeGraphGeneratorDiversePenaltyConfig](options.generatorConfig().contentAsString) match {
      case Right(config) => {
      implicit val rng : RandomGenerator = new MersenneTwister();

      options.outputPrefixes().foreach(prefix => {
        val (game, structure) = generateGame(config)
      File(prefix+options.structureSuffix()).write(structure.toString())
      import DeterministicGraphGameConfigJson.deterministicGraphGameConfigCodec
      File(prefix+"."+GraphGamePlugin.fileExtension).write(game.asJson.spaces2)
      })
      }
      case Left(error) => println(error)
    }
  }
}
