package pl.edu.pw.mini.sg.correlated
import pl.edu.pw.mini.sg.bosansky.{ExtendedMove, InformationSetId}
import pl.edu.pw.mini.sg.game.deterministic.{AdvanceablePlayerVisibleState, PlayerVisibleState}
import pl.edu.pw.mini.sg.solver._

import scala.collection.mutable

case class MilpRestrictionGenerator(informationSetChooser: InformationSetChooser) extends RestrictionGenerator {

  private val eps = 1e-3

  override def addRestrictions[DM,
    AM,
    DS <: PlayerVisibleState[DM],
    AS <: AdvanceablePlayerVisibleState[AM, AS],
    V <: SolverVariable,
    I,
    L <: SolverLinExpr[V, I],
    C <: SolverConstraint](currentBestRestriction: Restriction[L],
                           currentRestrictions: mutable.PriorityQueue[Restriction[L]],
                           inconsistentInformationSets: List[InformationSetId[AM, AS]],
                           value: Double,
                           solverModel: SolverModel[V, I, L, C],
                           p: Map[(List[ExtendedMove[DM, DS]], List[ExtendedMove[AM, AS]]), V],
                           relevantPairs: List[(List[ExtendedMove[DM, DS]], List[ExtendedMove[AM, AS]])]): Unit = {

    val upperBound: Double = Double.PositiveInfinity
    val bestRestrictionConstraints = currentBestRestriction.constraints
    val chosenInfSets = informationSetChooser.chooseSets(inconsistentInformationSets)
    println(inconsistentInformationSets)
    println(chosenInfSets)
    val ncc = chosenInfSets.flatMap(
      i => {
        val moves = i.reachedState.possibleMoves.map(m => ExtendedMove(m, i.reachedState))
        val chosenMoves = moves
        val binaryVars = chosenMoves.zip(solverModel.addVariables(0.0, 1.0, 0.0, Binary, Array.fill(chosenMoves.length)(""))).toMap
        val binarySumExpr = solverModel.createLinExpr()
        binarySumExpr.addTerms(1.0, binaryVars.values.toSeq)

        val newConstraints = ConstraintParameters(binarySumExpr, Equal, 1.0) :: chosenMoves.flatMap(m => {
          relevantPairs.filter{case (_, attMoves) => attMoves == m :: i.precedingMoves}.map{ case (defMoves, attMoves) => {
            val linExpr = solverModel.createLinExpr()
            linExpr.addTerm(1.0, p(defMoves, attMoves))
            linExpr.addTerm(-1.0, binaryVars(m))
            ConstraintParameters(linExpr, LessEqual, 0.0)
          }}
        })

        println(s"ncl: ${newConstraints.length}")
        newConstraints
      })
    currentRestrictions.enqueue(Restriction(upperBound, bestRestrictionConstraints ++ ncc))
  }
}
