package pl.edu.pw.mini.sg.bosansky

import pl.edu.pw.mini.sg.game.Payoff
import pl.edu.pw.mini.sg.game.boundedrationality.payoffchanger.{ NeutralPayoffChanger, PayoffChanger }
import pl.edu.pw.mini.sg.game.deterministic.{AdvanceablePlayerVisibleState, DeterministicGame, PlayerVisibleState}
import pl.edu.pw.mini.sg.game.util.MoveStateKey
import pl.edu.pw.mini.sg.mixed.IterationInfo
import pl.edu.pw.mini.sg.mixed.tree.MixedStrategyNode
import pl.edu.pw.mini.sg.solver._
import pl.edu.pw.mini.sg.util.StreamUtils
import squants.Time
import squants.time.Nanoseconds

import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer

class DeterministicGameSolver[DM, AM, DS <: PlayerVisibleState[DM], AS <: AdvanceablePlayerVisibleState[AM, AS], G <: DeterministicGame[DM, AM, DS, AS, G]]
    (solver: Solver, utils : GameUtilsProvider[DM, AM, DS, AS, G]) {

  private val solverModel = solver.createModel

  import pl.edu.pw.mini.sg.util.StreamUtils.myStreamLength

  //TODO: refactor and extract a parameters object
  // Should those parameters be in the class or in here?
  // TODO: adjust M value for anchoring!!
  def findStrategy(game: G, M: Double = 1.0, timeout: Option[Time] = None,
    payoffChanger: PayoffChanger = NeutralPayoffChanger, anchoringAlpha: Option[Double] = None
  ): (MixedStrategyNode[DM, DS, Unit], Payoff, Time) = {
    val startPreTime = System.nanoTime()
   
    val constraintBatchSize = 100000
    def paths = utils.getAllSequences(game)
    def leaves = paths.filter(n => n._1.isLeaf)
    val attInformationSets = utils.getAttackerInformationSets(paths)
    val defInformationSets = utils.getDefenderInformationSets(paths)
    val attackerPaths = paths.map(p => p._3).distinct.toArray

    val leavesCount = myStreamLength(leaves)
    //Dla kazdego liscia z dodajemy zmienna p(z)
    val p = solverModel.addVariables(Array.fill(leavesCount)(0.0),
      Array.fill(leavesCount)(1.0),
      Array.fill(leavesCount)(0.0),
      Array.fill(leavesCount)(Real),
      Array.fill(leavesCount)("")
    )

    def createSumExpr = {
      val sump = solverModel.createLinExpr()
      sump.addTerms(Array.fill(leaves.length)(1.0), p)
      sump
    }

    //Prawd. sumuja sie do jedynki - wzor (23)
    solverModel.addConstraint(createSumExpr, Equal, 1.0, "peq1")

    def objectiveXX = {
      val objective = solverModel.createLinExpr()
      objective.addTerms(leaves.map {
        case (state, _, _) => state.defenderPayoff //W tym miejscu ewentualnie mozna dodac nature
      }.toArray, p)
      objective
    }

    //objective to suma p(z) razy wyplata za z dla obroncy - wzor (16)
    solverModel.setObjective(objectiveXX)
    solverModel.setObjectiveType(Maximize)

    //Zmienne symulujace plan realizacji dla obroncy - wzor (25)
    val r1Names =  defInformationSets.map(k => k.precedingMoves).toArray
    val r1 = solverModel.addVariables(Array.fill(r1Names.length)(0.0),
      Array.fill(r1Names.length)(1.0),
      Array.fill(r1Names.length)(0.0),
      Array.fill(r1Names.length)(Real),
      Array.fill(r1Names.length)("")).zip(r1Names).map(v => (v._2, v._1)).toMap
    solverModel.addConstraint(r1(List()), Equal, 1.0, "cr1") //wzor (18)

    //Zmienne symulujace plan realizacji dla atakujacego - wzor (24)
    val r2Names = attInformationSets.map(k => k.precedingMoves).toArray
    val r2 = solverModel.addVariables(Array.fill(r2Names.length)(0.0),
      Array.fill(r2Names.length)(1.0),
      Array.fill(r2Names.length)(0.0),
      Array.fill(r2Names.length)(Binary),
      Array.fill(r2Names.length)("")).zip(r2Names).map(v => (v._2, v._1)).toMap
    solverModel.addConstraint(r2(List()), Equal, 1.0, "cr2") //wzor (18)

    //Zmienne slack
    val slackNames = attackerPaths.toArray
    val slack = solverModel.addVariables(Array.fill(slackNames.length)(0.0),
      Array.fill(slackNames.length)(Double.PositiveInfinity),
      Array.fill(slackNames.length)(0.0),
      Array.fill(slackNames.length)(Real),
      Array.fill(slackNames.length)("")).zip(slackNames).map(v => (v._2, v._1)).toMap

    //Zerujemy zmienne slack dla strategii faktycznie realizowanych przez atakujacego - wzor (20)
    attackerPaths.foreach(
      attackerMoves => {
        val r = r2(attackerMoves)
        val s = slack(attackerMoves)
        val expr = solverModel.createLinExpr()
        expr.addConstant(M)
        expr.addTerm(-M, r)
        solverModel.addConstraint(expr, GreaterEqual, s, "s")
      }
    )

    {
      StreamUtils.streamGrouped(StreamUtils.streamZip(p.toStream, leaves).map{case (pVar, (_,defenderMoves, attackerMoves)) => {
          val leDef = solverModel.createLinExpr()
          leDef.addTerm(1, pVar)
          leDef.addTerm(-1, r1(defenderMoves))
          val leAtt = solverModel.createLinExpr()
          leAtt.addTerm(1, pVar)
          leAtt.addTerm(-1, r2(attackerMoves))
          (leDef, leAtt)}},constraintBatchSize).foreach{ batch => {
          val (linexprDef, linexprAtt) = batch.toVector.unzip
          val types = Array.fill[ConstraintType](linexprDef.length)(LessEqual)
          val rhs = Array.fill(linexprDef.length)(0.0)
          val names = Array.fill(linexprDef.length)("")
          solverModel.addConstraints(linexprDef, types, rhs, names)
          solverModel.addConstraints(linexprAtt, types, rhs, names)
        }}
    }

    //Ograniczenie, ze wartosc planu realizacji musi byc rowna sumie wartosci nastepnikow - wzor (19)
    defInformationSets.foreach {
      id =>
        val r = r1(id.precedingMoves)
        //Possible moves from one inf. set are the same for every state
        val defenderState = id.reachedState
        val moves = defenderState.possibleMoves.map(m => ExtendedMove(m, defenderState))
        if (moves.nonEmpty) {
          val sumR = solverModel.createLinExpr()
          sumR.addTerms(Array.fill(moves.length)(1.0),
            moves.map(m => r1(m :: id.precedingMoves)))
          solverModel.addConstraint(sumR, Equal, r, "r1")
        }
    }

    attInformationSets.foreach {
      id =>
        val r = r2(id.precedingMoves)
        //Possible moves from one inf. set are the same for every state
        val attackerState = id.reachedState
        val moves = attackerState.possibleMoves.map(m => ExtendedMove(m, attackerState))
        if (moves.nonEmpty) {
          val sumR = solverModel.createLinExpr()
          sumR.addTerms(Array.fill(moves.length)(1.0),
            moves.map(m => r2(m :: id.precedingMoves)))
          solverModel.addConstraint(sumR, Equal, r, "r2")
        }
    }

    val attInfSetsWhereLastActionWasTaken  = (attInformationSets.flatMap {
      id =>
        val attackerState = id.reachedState
        val moves = attackerState.possibleMoves.map(m => ExtendedMove(m, attackerState))
        moves.map(m => (m :: id.precedingMoves, id))
    } ++ List((List(), attInformationSets.find(_.precedingMoves == List()).get))).toMap

    //Zmienne v oraz ograniczenie dla nich - wzor (17)
    def formula17Args = {
      val vNames = attInfSetsWhereLastActionWasTaken.values.toArray
      solverModel.addVariables(Array.fill(vNames.length)(Double.NegativeInfinity),
      Array.fill(vNames.length)(Double.PositiveInfinity),
      Array.fill(vNames.length)(0.0),
      Array.fill(vNames.length)(Real),
      Array.fill(vNames.length)("")).zip(vNames).map(v => (v._2, v._1)).toMap
    }
    val v = formula17Args

    attInfSetsWhereLastActionWasTaken.foreach {
      case (attackerMoves, infSet) =>

      val vInf = v(infSet)
      val s = slack(attackerMoves)
      val linExpr = solverModel.createLinExpr()
      linExpr.addTerm(1.0, s)

        linExpr.addTerms(1.0, attInformationSets.view.filter(k => k.precedingMoves == attackerMoves && v.contains(k)).map(i => v(i)).toSeq)


        val vILeaves = {
          val vl = ArrayBuffer[(G, List[ExtendedMove[DM, DS]], List[ExtendedMove[AM, AS]])]()
          leaves.foreach{l => if(l._3 == attackerMoves) vl+=l}
          vl.toArray
        }
      if(vILeaves.nonEmpty)
        linExpr.addTerms( vILeaves.map(l => payoffChanger.changeExpectedPayoff(new Payoff(l._1.defenderPayoff,l._1.attackerPayoff)).attacker),
          vILeaves.map(l => r1(l._2)))

        anchoringAlpha match {
          case Some(alpha) => 
            linExpr.addTerms( vILeaves.map(l => (1.0/l._2.head.precedingState.possibleMoves.length.toDouble)*alpha*payoffChanger.changeExpectedPayoff(new Payoff(l._1.defenderPayoff,l._1.attackerPayoff)).attacker),
              vILeaves.map(l => r1(l._2.tail)))
          case None => ()
        }

      if(attackerMoves.nonEmpty)
        solverModel.addConstraint(linExpr, Equal, vInf, "v")
    }

    val preprocessTime = Nanoseconds(System.nanoTime() - startPreTime)


    val solutionFound = timeout match {
      case None => solverModel.optimize() == Feasible
      case Some(tm) => solverModel.optimize(tm).hasSolution
    }

    if(solutionFound) {
      val result = (getStrategyFromModel(game, r1, defInformationSets.toStream), getPayoffFromModel(solverModel, v), preprocessTime)
      solverModel.close()
      result
    } else {
      throw new RuntimeException("No solution found!")
    }
  }

  private def getStrategyFromModel(initialGame: G, realizationPlanDefender: Map[List[ExtendedMove[DM, DS]], Solver#VariableType], defenderInfSets: Stream[InformationSetId[DM, DS]]): MixedStrategyNode[DM, DS, Unit] = {
    val root = new MixedStrategyNode[DM, DS, Unit](initialGame.defenderState)(MixedStrategyNode.unitST)
    val queue = new mutable.Queue[(InformationSetId[DM, DS], MixedStrategyNode[DM, DS, Unit])]()
    val firstInfSet = defenderInfSets.filter(_.precedingMoves.isEmpty).head
    queue.enqueue((firstInfSet, root))
    while(queue.nonEmpty) {
      val (infSetId, node) = queue.dequeue()
      val state = infSetId.reachedState
      val r = realizationPlanDefender(infSetId.precedingMoves).value()

      val possibleMoves = state.possibleMoves.map(m => ExtendedMove(m, state))
      possibleMoves.foreach(em => {
        val moveList = em :: infSetId.precedingMoves
        val rNext = realizationPlanDefender(moveList).value()
          if(rNext != 0) {
            val probability = rNext / r
            val possibleInfSets = defenderInfSets.filter(i => i.precedingMoves == moveList)
            node.addMove(em.move, probability)(IterationInfo(0,0))
            possibleInfSets.foreach(i => {
              val newNode = new MixedStrategyNode[DM, DS, Unit](i.reachedState)(MixedStrategyNode.unitST)
              node.succesors.getOrElseUpdate(em.move, scala.collection.mutable.Map()) += (i.reachedState -> newNode)
              queue.enqueue((i, newNode))
            })
          }
      })
    }
  root
}


  private def getPayoffFromModel(solverModel: Solver#ModelType, attackerExpectedPayoffs: Map[InformationSetId[AM, AS], Solver#VariableType]): Payoff = {
    new Payoff(solverModel.getObjective, attackerExpectedPayoffs.filterKeys(_.precedingMoves == List()).head._2.value())
  }

}
