package pl.edu.pw.mini.sg.colgen

import argonaut.{ Json, JsonIdentity }
import com.google.common.base.Stopwatch
import com.google.common.collect.ImmutableMap
import gurobi.{ GRB, GRBEnv, GRBException, GRBLinExpr, GRBModel, GRBVar }
import java.io.File
import java.util.concurrent.TimeUnit
import org.rogach.scallop.ScallopConf
import pl.edu.pw.mini.jk.sg.game.immutable.mutable.MutableWrapper
import pl.edu.pw.mini.jk.sg.gteval.GameMatrices
import pl.edu.pw.mini.jk.sg.opt.util.GameUtil
import pl.edu.pw.mini.sg.movingtargets.{ AttackDecision, MovingTargetsAttackerState }
import pl.edu.pw.mini.sg.movingtargets.coverage.{ CoverageDefenderMove, CoverageDefenderStateCached, CoverageImmutableGame }
import pl.edu.pw.mini.sg.movingtargets.{ AttackerDecision, MovingTargetsGamePlugin }
import pl.edu.pw.mini.sg.movingtargets.coverage.CoverageNormalizableFactory
import scala.collection.JavaConverters._
import scala.util.control.Breaks._

import scala.collection.mutable.{ArrayBuffer}
import squants.time.{ Milliseconds, Seconds }

object ColGen {
  def main(args: Array[String]) : Unit = {
    val options = new CmdOptions(args)

    val env = new GRBEnv()

    val gameConfig = MovingTargetsGamePlugin.fileParser(options.gameFile()).right.get
    val sw1 = Stopwatch.createStarted()
    val game = CoverageNormalizableFactory(gameConfig)
    val matrices = GameUtil.calculateGameMatrices[
      AttackerDecision, CoverageDefenderMove,
      MovingTargetsAttackerState, CoverageDefenderStateCached,
      MutableWrapper[AttackerDecision, CoverageDefenderMove,
        MovingTargetsAttackerState, CoverageDefenderStateCached,
        CoverageImmutableGame
      ]](game, gameConfig.roundLimit)


    sw1.stop()
    val sw2 = Stopwatch.createStarted()

    val attackerSequencesSorted = matrices.getAttackerSequences.asScala.sortBy(seq =>
      seq.asScala.collect{case d@AttackDecision(_) => d}.headOption
        .map{case AttackDecision(t) => gameConfig.targets(t).defenderUtilityUnuccessfulAttack}
    )
    val defenderSequencesSorted = matrices.getDefenderSequences().asScala.sortBy(seq =>
      -seq.asScala.map(cdm => cdm.coveredTargetIds.size).min
    )
    val columnsUsed: ArrayBuffer[Int] = ArrayBuffer()
    // (dr, AS, DS (prawd. dla tych strategii, gdzie jest true w wektorze))
    var bestStrategies : (Double, Int, List[Double], Vector[Int]) = (
      -1000.0, 0 , List(1), Vector(0))

    println(s"dseq: ${defenderSequencesSorted.length}, aseq: ${attackerSequencesSorted.length}, ${game.initialState.defenderVisibleState.getPossibleMoves}")

    var improved = true
    while(improved) {
      improved = defenderSequencesSorted.indices.filter(i => !columnsUsed.contains(i)).find(currentDs => {
        var im : Boolean = false
        breakable {
          for(
            currentAs <- attackerSequencesSorted.indices
          ) {
            val model = new GRBModel(env)
            columnsUsed+=currentDs

            val xVars: Vector[GRBVar]  = columnsUsed.map(ci => model.addVar(0, 1, matrices.getGamePayoff(defenderSequencesSorted(ci), attackerSequencesSorted(currentAs)).getDefenderPayoff, GRB.CONTINUOUS, s"x_$ci")).toVector
            model.update()
            val sumX = {
              val sum = new GRBLinExpr()
              xVars.foreach(sum.addTerm(1, _))
              model.addConstr(sum,GRB.EQUAL, 1, "sumX")
            }
            model.update()
            try {

              assert(model.getCoeff(sumX, xVars(0)) == 1.0)

              println(sumX.toString())

              val bestreward = new GRBLinExpr()
              xVars.zip(columnsUsed).foreach{case (xi,di) => bestreward.addTerm(
                matrices.getGamePayoff(defenderSequencesSorted(di), attackerSequencesSorted(currentAs)).getAttackerPayoff
                  , xi)}
//              bestreward.addConstant(1e-4)
              model.set(GRB.IntAttr.ModelSense, GRB.MAXIMIZE)

              val bestResp = attackerSequencesSorted.map(ar => {
                val areward = new GRBLinExpr()
                xVars.zip(columnsUsed).foreach{case (xi,di) => areward.addTerm(
                  matrices.getGamePayoff(defenderSequencesSorted(di), ar).getAttackerPayoff
                    , xi)}

                model.addConstr(bestreward, GRB.GREATER_EQUAL, areward, s"areward$ar")
              })
              model.optimize()
              model.update()
              val status = model.get(GRB.IntAttr.Status)
              if(status == ScalaHelper.INFEASIBLE) {
                columnsUsed.remove(columnsUsed.length-1)
              } else {
                assert(status == ScalaHelper.OPTIMAL)
                val objval: Double  = model.get(GRB.DoubleAttr.ObjVal)
                println(s"Objval: $objval")
                if(objval > bestStrategies._1) {
                  assert(columnsUsed.length == xVars.length)
                  bestStrategies = (objval, currentAs, xVars.map(x => x.get(GRB.DoubleAttr.X)).toList, columnsUsed.toVector)
                  im = true
                } else {
                  columnsUsed.remove(columnsUsed.length-1)
                }
              }
            } catch {
              case e: GRBException => {
                println(e.getMessage)
                println(e.getErrorCode)
                throw e;
              }
            }
            model.dispose()
          }
        }
        im

      }).nonEmpty
    }

    sw2.stop()
    env.dispose()

    System.out.println(s"Solution found, payoff: ${bestStrategies._1}")
    val resultDescription = game.defenderResultAsJson(bestStrategies._4.zip(bestStrategies._3).map{
      case (idx, prob) => (defenderSequencesSorted(idx), new java.lang.Double(prob))
    }.toMap.asJava, sw2.elapsed(TimeUnit.MILLISECONDS)/1000.0, "column generation LP", Json(),
      Milliseconds(sw1.elapsed(TimeUnit.MILLISECONDS)))

    import better.files._

    options.outputFile().toScala.write(resultDescription.toString)
  }
}

private[colgen] class CmdOptions(args: Array[String]) extends ScallopConf(args) {
  def gameFile = opt[File](name="game", short='g', required = true)
  def outputFile = opt[File](name="output", short='o', required = true)

  verify()
}
