package pl.edu.pw.mini.sg.correlated

import pl.edu.pw.mini.sg.bosansky.{ExtendedMove, GameUtilsProvider, InformationSetId}
import pl.edu.pw.mini.sg.game.Payoff
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

class CorrelatedStrategiesSolver[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], restrictionGenerator: RestrictionGenerator){

  private val solverModel = solver.createModel

  private val eps = 1e-3

  def findStrategy(game: G, timeLimit: Option[Time]): (MixedStrategyNode[DM, DS, Unit], Payoff, Time) = {
    val startPreTime = System.nanoTime()
    val constraintBatchSize = 100000

    val paths = utils.getAllSequences(game).toList
    val attInformationSets = paths.view.filter(_._1.defenderState.possibleMoves.nonEmpty).flatMap{case (game, leader, follower) => game.defenderState.possibleMoves.map(dm => (game, ExtendedMove(dm, game.defenderState)::leader, follower)) }.groupBy(p => InformationSetId(p._3.toList, p._1.attackerState))
    val defInformationSets = paths.view.filterNot(_._1.isLeaf).map{case (g, l, _) => InformationSetId(l, g.defenderState)}.distinct.toList
    val leaves = paths.filter(n => n._1.isLeaf).groupBy(_._3).map{case (f, seq) => (f, seq.map{case (g, l, f) => (l, g)}.toMap)}
    val attackerPaths = paths.map(p => p._3).distinct.toArray
    val defenderPaths = paths.map(p => p._2).distinct.toArray
    val relevantPairs: List[(List[ExtendedMove[DM, DS]], List[ExtendedMove[AM, AS]])]  = {
      val rpp = paths.flatMap{case (g, l, f) => f.tails.map(ff => (l, ff))}.toSet ++ leaves.keys.toSet
      defenderPaths.flatMap(dp => attackerPaths.collect{case ap if rpp.contains((dp,ap))|| rpp.contains((dp, ap.tail)) => (dp, ap)}).toList
    }

    //p Variables
    val p = solverModel.addVariables(Array.fill(relevantPairs.length)(0.0),
      Array.fill(relevantPairs.length)(1.0),
      Array.fill(relevantPairs.length)(0.0),
      Array.fill(relevantPairs.length)(Real),
      Array.fill(relevantPairs.length)("")
    ).zip(relevantPairs).map(p => (p._2, p._1)).toMap

    //v sequence variables
    val vSeq = solverModel.addVariables(Array.fill(attackerPaths.length)(Double.NegativeInfinity),
      Array.fill(attackerPaths.length)(Double.PositiveInfinity),
      Array.fill(attackerPaths.length)(0.0),
      Array.fill(attackerPaths.length)(Real),
      Array.fill(attackerPaths.length)("")
    ).zip(attackerPaths).map(p => (p._2, p._1)).toMap

    //v inf set variables
    val vInfSet = {
      val vInfSequences = attInformationSets.toSeq.flatMap{
        case (infId, _) =>
          attackerPaths.map(ap => (infId, ap))
      }

      solverModel.addVariables(Array.fill(vInfSequences.size)(Double.NegativeInfinity),
        Array.fill(vInfSequences.size)(Double.PositiveInfinity),
        Array.fill(vInfSequences.size)(0.0),
        Array.fill(vInfSequences.size)(Real),
        Array.fill(vInfSequences.size)("")
      ).zip(vInfSequences).map(p => (p._2, p._1)).toMap
    }

    //objective
    def objective = {
      val obj = solverModel.createLinExpr()
      val ss = leaves.flatMap{case (f, ll) => ll.map{case (l, g) => (g, l, f)}}
      obj.addTerms(ss.map(_._1.defenderPayoff).toList, ss.map(l => p(l._2, l._3)).toList)
      obj
    }
    solverModel.setObjective(objective)
    solverModel.setObjectiveType(Maximize)
    //p constraint (3)
    solverModel.addConstraint(p(List(), List()), Equal, 1.0, "p")
    //p constraint (4)

    {
      val pDefConstr =
        relevantPairs.filter(l => l._1.nonEmpty).groupBy(l => (l._1.tail, l._1.head.precedingState, l._2)).filter { case ((prec, _, follower), _) => true || prec.length >= follower.length || prec.isEmpty
        }
          .map { case ((prec, is, f), moves) => {

            if (p.contains((prec, f))) Some {
              val lhs = solverModel.createLinExpr()
              lhs.addTerm(1.0, p((prec, f)))
              lhs.addTerms(-1.0, moves.map(m => p(m._1, f)))
              (lhs, (prec.map(_.move), f.map(_.move)))
            } else None
          }
          }.collect { case Some(x) => x }.toList

      solverModel.addConstraints(pDefConstr.map(_._1).toVector,
        Array.fill(pDefConstr.size)(Equal),
        Array.fill(pDefConstr.size)(0.0),
        pDefConstr.map(x => s"pDef ${x._2}").toArray)
    }

    //p constraint (5)
    {
      val pAttConstr = relevantPairs.filter(l => l._2.nonEmpty).groupBy(l => (l._2.tail, l._2.head.precedingState, l._1)).filter { case ((prec, _, leader), _) => true || prec.length >= leader.length || prec.isEmpty
      }.map { case ((prec, is, leader), moves) =>
        if (p.contains((leader, prec))) Some {
          val lhs = solverModel.createLinExpr()
          lhs.addTerm(1.0, p(leader, prec))
          lhs.addTerms(-1.0, moves.map(m => p(leader, m._2)))
          (lhs, (leader.map(_.move), prec.map(_.move)))
        } else None
      }.toList.collect { case Some(x) => x }
      solverModel.addConstraints(pAttConstr.map(_._1).toVector,
        Array.fill(pAttConstr.size)(Equal),
        Array.fill(pAttConstr.size)(0.0),
        pAttConstr.map(x => s"pAtt ${x._2}").toArray)
    }

    // v constraint (8)
    {
      val v8ConstrArgs = attInformationSets.flatMap {
        case (infId, _) =>
          val moves = infId.reachedState.possibleMoves.map(m => ExtendedMove(m, infId.reachedState))
          moves.map(m => {
            val lhs = solverModel.createLinExpr()
            lhs.addTerm(1.0, vSeq(m :: infId.precedingMoves))
            lhs.addTerm(-1.0, vInfSet(infId, m :: infId.precedingMoves))
            lhs
          })
      }
      solverModel.addConstraints(v8ConstrArgs.toVector,
        Array.fill(v8ConstrArgs.size)(Equal),
        Array.fill(v8ConstrArgs.size)(0.0),
        Array.fill(v8ConstrArgs.size)("")
      )
    }

    //v constraint (6)
    {
      val v6ConstrArgs = attackerPaths.map(path => {
        val lhs = solverModel.createLinExpr()
        lhs.addTerm(1.0, vSeq(path))
        val ll = leaves.get(path).map(_.toList).getOrElse(Nil)
        lhs.addTerms(ll.map(l => -l._2.attackerPayoff), ll.map(l => p(l._1, path)))
        lhs.addTerms(-1.0, attInformationSets.filter(_._1.precedingMoves == path).flatMap {
          case (infId, _) =>
            val moves = infId.reachedState.possibleMoves.map(m => ExtendedMove(m, infId.reachedState))
            moves.map(m => vSeq(m :: path))
        }.toSeq)
        lhs
      })
      solverModel.addConstraints(v6ConstrArgs.toVector,
        Array.fill(v6ConstrArgs.length)(Equal),
        Array.fill(v6ConstrArgs.length)(0.0),
        Array.fill(v6ConstrArgs.length)("")
      )
    }

    //v constraint (7)
    {
      val relAtt = relevantPairs.groupBy(_._2)
      val v7ConstrArgs = attInformationSets.flatMap {
        case (k, _) =>
          (List(k) ++ k.precedingMoves.tails.map(f =>
            InformationSetId(f.drop(1), f.headOption.map(_.precedingState).getOrElse(game.attackerState)))).flatMap(h => {
            val moves = k.reachedState.possibleMoves.map(m => ExtendedMove(m, k.reachedState))
            val attPaths = h.reachedState.possibleMoves.map(m => ExtendedMove(m, h.reachedState) :: h.precedingMoves)
            attPaths.flatMap(sigmaf => {
              moves.map(m => {
                val lhs = solverModel.createLinExpr()
                lhs.addTerm(1.0, vInfSet(k, sigmaf))
                val pSumArgs = relAtt(sigmaf).map {
                  case (dp, _) => (dp, leaves.get(m :: k.precedingMoves).flatMap(_.get(dp)))
                }.collect { case (dp, Some(x)) if x.isLeaf => (dp, x) }
                lhs.addTerms(pSumArgs.map(-_._2.attackerPayoff), pSumArgs.map(pArg => p(pArg._1, sigmaf)))
                lhs.addTerms(-1.0, attInformationSets.keys.filter(tt => tt.precedingMoves == m :: k.precedingMoves).map(is => vInfSet(is, sigmaf)).toSeq)
                lhs
              })
            })
          })
      }
      solverModel.addConstraints(v7ConstrArgs.toVector,
        Array.fill(v7ConstrArgs.size)(GreaterEqual),
        Array.fill(v7ConstrArgs.size)(0.0),
        Array.fill(v7ConstrArgs.size)("")
      )
    }


    val preprocessTime = Nanoseconds(System.nanoTime() - startPreTime)
    val solverStarted = Nanoseconds(System.nanoTime())
    def withinTimeLimit() : Boolean = {
      timeLimit.map(tl => {
        val time = Nanoseconds(System.nanoTime())
        val elapsed = time-solverStarted
        tl > elapsed
      }).getOrElse(true)
    }
    //compute SSE
    def computeSSE: (Double, Map[(List[ExtendedMove[DM, DS]], List[ExtendedMove[AM, AS]]), Double]) = {
      val modificationsQueue = mutable.PriorityQueue(
        Restriction[solver.LinExprType](Double.PositiveInfinity, List())
      )(Ordering.by(Restriction.restrictionOrder))
      var lowerBound: Double = Double.NegativeInfinity
      var correlationPlan = Map.empty[(List[ExtendedMove[DM, DS]], List[ExtendedMove[AM, AS]]), Double]

      while(modificationsQueue.nonEmpty && withinTimeLimit()) {
        val bestRestriction = modificationsQueue.dequeue()
        if(bestRestriction.upperBound < lowerBound)
          return (lowerBound, correlationPlan)
        val addedConstraints = solverModel.addConstraints(
          bestRestriction.constraints.map(_.linExpr).toVector,
          bestRestriction.constraints.map(_.constraintType).toArray,
          bestRestriction.constraints.map(_.rhs).toArray,
          Array.fill(bestRestriction.constraints.length)("")
        )

        if(timeLimit match {
          case None => solverModel.optimize().hasSolution
          case Some(limit) => {
            val lmt = limit - (Nanoseconds(System.nanoTime())-solverStarted)
            solverModel.optimize(lmt).hasSolution
          }
        }) {
          val defenderPayoff = solverModel.getObjective
          val pVals = p.map{case (k, v) => (k, v.value())}

          val inconsistentInformationSets = attInformationSets.filter {
            case (infId, states) =>
              val moves = infId.reachedState.possibleMoves.map(m => ExtendedMove(m, infId.reachedState))
              val mm = moves.filter(m => states.exists { case (_, defenderMoves, _) =>
                pVals.isDefinedAt((defenderMoves, m :: infId.precedingMoves)) && pVals((defenderMoves, m :: infId.precedingMoves)) > eps
              })
              mm.length > 1
          }.keys.toList
          if(inconsistentInformationSets.isEmpty) {
            if(defenderPayoff > lowerBound) {
              lowerBound = defenderPayoff
              correlationPlan = pVals
            }
          }
          else
            restrictionGenerator.addRestrictions(bestRestriction,
              modificationsQueue,
              inconsistentInformationSets,
              defenderPayoff,
              solverModel,
              p,
              relevantPairs)
        } else {
          println("Infeasible model!!")
        }
        addedConstraints.foreach(solverModel.removeConstraint)
      }
      (lowerBound, correlationPlan)
    }

    val (defenderPayoff, correlationPlan) = computeSSE
    val attackerPayoff = vSeq(List()).value()

    val root = new MixedStrategyNode[DM, DS, Unit](game.defenderState)(MixedStrategyNode.unitST)
    val queue = new mutable.Queue[(InformationSetId[DM, DS], MixedStrategyNode[DM, DS, Unit])]()
    val firstInfSet = InformationSetId(List.empty[ExtendedMove[DM, DS]], game.defenderState)
    queue.enqueue((firstInfSet, root))
    while(queue.nonEmpty) {
      val (infSetId, node) = queue.dequeue()
      val state = infSetId.reachedState
      val possibleMoves = state.possibleMoves.map(m => ExtendedMove(m, state))
      val curPVal = correlationPlan.find(pval => pval._2 > eps && pval._1._1 == infSetId.precedingMoves).get._2
      possibleMoves.foreach(m => {
        val moveList = m :: infSetId.precedingMoves
        val pNext = correlationPlan.find(pval => pval._2 > eps && pval._1._1 == moveList).map(_._2).getOrElse(0.0)
        if(pNext > eps) {
          val probability = pNext / curPVal
          node.addMove(m.move, probability)(IterationInfo.dummyIterationInfo)
          val possibleInfSets = defInformationSets.filter(i => i.precedingMoves == moveList)
          possibleInfSets.foreach(i => {
            val newNode = new MixedStrategyNode[DM, DS, Unit](i.reachedState)(MixedStrategyNode.unitST)
            node.addSuccessor(MoveStateKey(m.move, i.reachedState), newNode)
            queue.enqueue((i, newNode))
          })

        }
      })

    }


    solverModel.close()

    (root, new Payoff(defenderPayoff, attackerPayoff), preprocessTime)
  }

}
