package pl.edu.pw.mini.sg.compactlp

import pl.edu.pw.mini.sg.solver._
import pl.edu.pw.mini.sg.compactlp.model._
import pl.edu.pw.mini.sg.compactlp.util.{Attacker, PureStrategy}
import squants.Time
import squants.time.Nanoseconds
import pl.edu.pw.mini.sg.util.TraversableUtils._

class CompactLPSolver(solver: Solver, strategyGenerator: StrategyGenerator, eps: Double) {

  private val solverModel = solver.createModel

  def solve(game: CompactGameWrapper, sparsity: Int, M: Double = 2.0, M2: Double = 3.0): (Time, Time, Double, Double, Double, Double, Option[Attacker], List[PureStrategy]) = {
    val startPreTime = System.nanoTime()
    //Najpierw zmienne, potem ograniczenia
    //Zmienne c
    val grid = (game.getZones(sparsity) cross (0 to game.getTimeSteps(sparsity))).map {
      case (p, t) => GridSpacetime(t, p)
    }.toSeq
    val strategiesWithPayoffs = game.getAttackerStrategiesWithPayoffs
    val gridStrategy = (grid cross strategiesWithPayoffs).map {
      case (g, s) => GridSpacetimeStrategy(g, s.attacker)
    }
    val cNames = gridStrategy.toArray
    val c = toMap(solverModel.addVariables(0.0, game.getDefenderCount, 0.0, Real, cNames.map("c_" + _.toString)), cNames)

    //zmienne f
    val fNames = gridStrategy.filterNot(g => g.gridSpacetime.time == game.getTimeSteps(sparsity)).flatMap {
      case GridSpacetimeStrategy(spaceTime, strategy) => {
        val neighbours = game.getGridNeighbours(spaceTime.position, sparsity)
        neighbours.map(n => FlowStrategy(Flow(spaceTime.position, n, spaceTime.time), strategy))
      }
    }.toArray

    val f = toMap(solverModel.addVariables(0.0, game.getDefenderCount, 0.0, Real, fNames.map("f_" + _.toString)), fNames)

    //zmienne q
    val qNames = strategiesWithPayoffs.map(_.attacker)
    val q = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Binary, qNames.map("q_" + _.toString)), qNames)

    //zmienna a
    //    val a = solverModel.addVariable(Double.NegativeInfinity, Double.PositiveInfinity, 0.0, Real, "a")
    val a = solverModel.addVariable(-1.0, 1.0, 0.0, Real, "a")

    //Zmienne pomocnicze dpp(c, F)
    val dppNames = (strategiesWithPayoffs cross strategiesWithPayoffs).map {
      case (strategyA, strategyB) => DoubleAttacker(strategyA.attacker, strategyB.attacker)
    }.toSeq
    val dpp = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Real, dppNames.map("dpp_" + _.toString)), dppNames)
    //Zmienne pomocnicze m1
    val m1 = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Binary, dppNames.map("m1_" + _.toString)), dppNames)
    val m2 = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Binary, dppNames.map("m2_" + _.toString)), dppNames)

    //zmienna pomocnicza dla min(qf, 1 - dpp(c, F))
    val UdMinNames = (strategiesWithPayoffs cross strategiesWithPayoffs).map {
      case (strategyA, strategyB) => DoubleAttacker(strategyA.attacker, strategyB.attacker)
    }.toSeq
    val UdMin = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Real, UdMinNames.map("UdMin_" + _.toString)), UdMinNames)
    val m3 = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Binary, UdMinNames.map("m3_" + _.toString)), UdMinNames)
    val m4 = toMap(solverModel.addVariables(0.0, 1.0, 0.0, Binary, UdMinNames.map("m4_" + _.toString)), UdMinNames)
    solverModel.update()

    //ograniczenie (1c)
    gridStrategy.filterNot(g => g.gridSpacetime.time == game.getTimeSteps(sparsity)).foreach {
      case gridSpacetimeStrategy@GridSpacetimeStrategy(GridSpacetime(time, position), strategy) => {
        val fLinExpr = solverModel.createLinExpr()
        val neighbours = game.getGridNeighbours(position, sparsity)
        fLinExpr.addTerms(1.0, neighbours.map(p => f(FlowStrategy(Flow(position, p, time), strategy))))
        val _c = c(gridSpacetimeStrategy)
        solverModel.addConstraint(fLinExpr, Equal, _c, "constr_1c_" + gridSpacetimeStrategy.toString)
      }
    }

    //ograniczenie (1d)
    gridStrategy.filterNot(g => g.gridSpacetime.time == 0).foreach {
      case gridSpacetimeStrategy@GridSpacetimeStrategy(GridSpacetime(time, position), strategy) => {
        val fLinExpr = solverModel.createLinExpr()
        val neighbours = game.getGridNeighbours(position, sparsity)
        fLinExpr.addTerms(1.0, neighbours.map(p => f(FlowStrategy(Flow(p, position, time - 1), strategy))))
        val _c = c(gridSpacetimeStrategy)
        solverModel.addConstraint(fLinExpr, Equal, _c, "constr_1c_" + gridSpacetimeStrategy.toString)
      }
    }

    //ograniczenie (1e)
    (Seq(0, game.getTimeSteps(sparsity)) cross strategiesWithPayoffs).foreach {
      case (time, strategy) =>
        val gridFiltered = gridStrategy.filter(g => g.gridSpacetime.time == time && g.strategy == strategy.attacker).toSeq
        val mLinExpr = solverModel.createLinExpr()
        mLinExpr.addTerms(1.0, gridFiltered.map(g => c(g)))
        val qLinExpr = solverModel.createLinExpr()
        qLinExpr.addTerm(game.getDefenderCount, q(strategy.attacker))
        solverModel.addConstraint(mLinExpr, Equal, qLinExpr, "constr_1e_" + time.toString + strategy.attacker.toString)
    }

    //Pozycje startowe
    val startingPositions = game.getDefenderStartingPositions(sparsity)
    val uniqueStartingPositions = startingPositions.distinct
    (uniqueStartingPositions cross strategiesWithPayoffs).map {
      case (position, strategy) =>
        val gridSpacetimeStrategy = GridSpacetimeStrategy(GridSpacetime(0, position), strategy.attacker)
        val covStart = c(gridSpacetimeStrategy)
        val qLinExpr = solverModel.createLinExpr()
        val qvar = q(strategy.attacker)
        qLinExpr.addTerm(startingPositions.count(p => p == position), qvar)
        solverModel.addConstraint(qLinExpr, Equal, covStart, "c_start_" + gridSpacetimeStrategy.toString)
    }

    //Suma q rowna 1
    val qSum = solverModel.createLinExpr()
    qSum.addTerms(1.0, q.values.toSeq)
    solverModel.addConstraint(qSum, Equal, 1.0, "constr_qSum")



    //Ograniczenie min dla dpp
    (strategiesWithPayoffs cross strategiesWithPayoffs).foreach {
      case (strategyWithPayoffsC, strategyWithPayoffsF) =>
        val dppLinExpr = solverModel.createLinExpr()
        val strategyC = strategyWithPayoffsC.attacker
        val strategyF = strategyWithPayoffsF.attacker
        val doubleStrategy = DoubleAttacker(strategyC, strategyF)

        if (sparsity == 1) {
          val neighbors = game.getNeighbours(strategyF)
          dppLinExpr.addTerms(1.0, neighbors.map {
            case gridSpaceTime => c(GridSpacetimeStrategy(gridSpaceTime, strategyC))
          })
        }
        else {
          val strategyFNeighbours = game.getNeighbours(strategyF)
          val flowVariables = f.keys.filter(_.strategy == strategyC).toSeq
          val covs = flowVariables.map {
            case FlowStrategy(flow, _) => game.getCoverage(flow, strategyFNeighbours, sparsity)
          }
          dppLinExpr.addTerms(covs, flowVariables.map(f(_)))
        }

        solverModel.addConstraint(dppLinExpr, GreaterEqual, dpp(doubleStrategy), "constr_dpp_" + doubleStrategy.toString)
        dppLinExpr.addTerm(-M2*game.getDefenderCount, m1(doubleStrategy))
        solverModel.addConstraint(dppLinExpr, LessEqual, dpp(doubleStrategy), "constr_dppm1_" + doubleStrategy.toString)

        val m2LinExpr = solverModel.createLinExpr()
        m2LinExpr.addConstant(1.0)
        m2LinExpr.addTerm(-M2, m2(doubleStrategy))
        solverModel.addConstraint(m2LinExpr, LessEqual, dpp(doubleStrategy), "constr_dppm2_" + doubleStrategy.toString)


        val mSumLinExpr = solverModel.createLinExpr()
        mSumLinExpr.addTerm(1.0, m1(doubleStrategy))
        mSumLinExpr.addTerm(1.0, m2(doubleStrategy))

        solverModel.addConstraint(mSumLinExpr, Equal, 1.0, "constr_m12Sum_" + doubleStrategy.toString)
    }

    //Ograniczenia dla min(qf, 1 - dpp(c, F))
    (strategiesWithPayoffs cross strategiesWithPayoffs).foreach {
      case (strategyWithPayoffsC, strategyWithPayoffsF) => {
        val strategyC = strategyWithPayoffsC.attacker
        val strategyF = strategyWithPayoffsF.attacker
        val doubleStrategy = DoubleAttacker(strategyC, strategyF)

        val udMin = UdMin(doubleStrategy)
        val qf = q(strategyC)
        solverModel.addConstraint(udMin, LessEqual, qf, "constr_min_q_" + doubleStrategy.toString)
        val m3LinExpr = solverModel.createLinExpr()
        m3LinExpr.addTerm(1.0, qf)
        m3LinExpr.addTerm(-M2, m3(doubleStrategy))
        solverModel.addConstraint(m3LinExpr, LessEqual, udMin, "constr_udMin_m3_" + doubleStrategy.toString)

        val dppLinExpr = solverModel.createLinExpr()
        val dppCf = dpp(doubleStrategy)
        dppLinExpr.addConstant(1.0)
        dppLinExpr.addTerm(-1.0, dppCf)
        solverModel.addConstraint(dppLinExpr, GreaterEqual, udMin, "constr_min_dpp_" + doubleStrategy.toString)
        dppLinExpr.addTerm(-M2, m4(doubleStrategy))
        solverModel.addConstraint(dppLinExpr, LessEqual, udMin, "constr_udMin_m4_" + doubleStrategy.toString)

        val mSumLinExpr = solverModel.createLinExpr()
        mSumLinExpr.addTerm(1.0, m3(doubleStrategy))
        mSumLinExpr.addTerm(1.0, m4(doubleStrategy))

        solverModel.addConstraint(mSumLinExpr, Equal, 1.0, "constr_m34Sum_" + doubleStrategy.toString)
      }
    }

    //Ograniczenia ze stala M
    strategiesWithPayoffs.foreach {
      case strategyWithPayoffs => {
        var aLinExpr = solverModel.createLinExpr()
        aLinExpr.addTerm(1.0, a)
        strategiesWithPayoffs.foreach(sp => aLinExpr -= attackerPayoffLinExpr(sp, strategyWithPayoffs, dpp, UdMin))
        solverModel.addConstraint(aLinExpr, GreaterEqual, 0.0, "constr_a0_" + strategyWithPayoffs.attacker.toString)

        val MLinExpr = solverModel.createLinExpr()
        MLinExpr.addConstant(M)
        MLinExpr.addTerm(-M, q(strategyWithPayoffs.attacker))
        solverModel.addConstraint(aLinExpr, LessEqual, MLinExpr, "constr_M_" + strategyWithPayoffs.attacker.toString)
      }
    }
    //Objective
    solverModel.setObjectiveType(Maximize)
    var objectiveLinExpr = solverModel.createLinExpr()
    strategiesWithPayoffs.foreach { sp => objectiveLinExpr += defenderPayoffLinExpr(sp, sp, dpp, UdMin) }
    solverModel.setObjective(objectiveLinExpr)

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

    solverModel.optimize()

    val strategyChosen = strategiesWithPayoffs.filter(s => q(s.attacker).value() > eps).head
    val attackTime = if (strategyChosen.attacker.startTime >= 0) Some(strategyChosen.attacker.startTime) else None
    val targetId = strategyChosen.attacker.target

    //Rozbicie f na strategie obroncy
    val flowValues = f.filter(_._1.strategy == strategyChosen.attacker).map {
      case (FlowStrategy(flow, _), v) => (flow, v.value())
    }
    val mixedStrategy = strategyGenerator.getStrategy(flowValues, grid, game, sparsity)

    //symulacja gry, aby stwierdzic faktyczna wartosc strategii
    val attacker = attackTime.map(a => Attacker(targetId, a))
    val payoffs = mixedStrategy.map(s => {
      val p = game.getPayoffs(attacker, s.positions)
      (s.probability * p._1, s.probability * p._2)
    })
    val (defenderPayoff, attackerPayoff) = (payoffs.map(_._1).sum, payoffs.map(_._2).sum)

    val solveTime = Nanoseconds(System.nanoTime() - startTime)

    if (math.abs(defenderPayoff - solverModel.getObjective) > eps)
      println(Console.RED + "Symulowana wyplata obroncy rozni sie od wyplaty od solwera")

    if (math.abs(attackerPayoff - a.value()) > eps)
      println(Console.RED + "Symulowana wyplata atakujacego rozni sie od wyplaty od solwera")

    print(Console.RESET)

    solverModel.close()

    (preprocessTime, solveTime, defenderPayoff, attackerPayoff, solverModel.getObjective, a.value(), attacker, mixedStrategy)
  }

  private def toMap[K, T](list: Seq[T], keys: Seq[K]): Map[K, T] = {
    keys.zip(list).toMap
  }

  private def attackerPayoffLinExpr(strategyC: AttackerWithPayoffs, strategyF: AttackerWithPayoffs, dpp: Map[DoubleAttacker, solver.VariableType], udMin: Map[DoubleAttacker, solver.VariableType]): solver.LinExprType = {
    val result = solverModel.createLinExpr()
    val doubleStrategy = DoubleAttacker(strategyC.attacker, strategyF.attacker)
    val udMinVar = udMin(doubleStrategy)
    val dppVar = dpp(doubleStrategy)
    result.addTerm(strategyF.attackerSuccessfulAttackPayoff, udMinVar)
    result.addTerm(strategyF.attackerUnsuccessfulAttackPayoff, dppVar)
    result
  }

  private def defenderPayoffLinExpr(strategyC: AttackerWithPayoffs, strategyF: AttackerWithPayoffs, dpp: Map[DoubleAttacker, solver.VariableType], udMin: Map[DoubleAttacker, solver.VariableType]): solver.LinExprType = {
    val result = solverModel.createLinExpr()
    val doubleStrategy = DoubleAttacker(strategyC.attacker, strategyF.attacker)
    val udMinVar = udMin(doubleStrategy)
    val dppVar = dpp(doubleStrategy)
    result.addTerm(strategyF.defenderSuccessfulAttackPayoff, udMinVar)
    result.addTerm(strategyF.defenderUnsuccessfulAttackPayoff, dppVar)
    result
  }

}
