package pl.edu.pw.mini.sg.solver.gurobisolver

import gurobi._
import pl.edu.pw.mini.sg.solver.{Binary, Feasible, Infeasible, Integer, Maximize, Minimize, ObjectiveType, Real, SolverModel, VariableType, _}
import squants.time.Time

class GurobiSolverModel(model: GRBModel) extends SolverModel[GurobiVariable, GRBLinExpr, GurobiLinExpr, GurobiConstraint] {

  override def addConstraints(expr: Vector[GurobiLinExpr], constraintType: Array[ConstraintType],
    rhs: Array[Double], name: Array[String]) : Array[GurobiConstraint] = {
    model.addConstrs(expr.view.map(_.inner).toArray, constraintType.map(constraintTypeToChar _), rhs, name).map(new GurobiConstraint(_))
  }

  override def addVariable(lowerBound: Double, upperBound: Double, obj: Double, variableType: VariableType, name: String): GurobiVariable = {
    new GurobiVariable(model.addVar(lowerBound, upperBound, obj, variableTypeToChar(variableType), name))
  }

  override def addVariables(lowerBounds: Seq[Double], upperBounds: Seq[Double], objs: Seq[Double], variableTypes: Seq[VariableType], names: Seq[String]): Array[GurobiVariable] = {
    model.addVars(lowerBounds.toArray, upperBounds.toArray, objs.toArray, variableTypes.map(variableTypeToChar).toArray, names.toArray).map(v => new GurobiVariable(v))
  }

  override def addVariables(lowerBounds: Double, upperBounds: Double, objs: Double, variableTypes: VariableType, names: Seq[String]): Array[GurobiVariable] = {
    model.addVars(Array.fill(names.length)(lowerBounds),
      Array.fill(names.length)(upperBounds),
      Array.fill(names.length)(objs),
      Array.fill(names.length)(variableTypes).map(variableTypeToChar), names.toArray).map(v => new GurobiVariable(v))
  }

  override def createLinExpr(): GurobiLinExpr = {
    new GurobiLinExpr(new GRBLinExpr())
  }

  override def addConstraint(lhsExpr: GurobiLinExpr,
                             constraintType: ConstraintType,
                             rhsExpr: GurobiLinExpr,
                             name: String): GurobiConstraint = {
    new GurobiConstraint(model.addConstr(lhsExpr.inner, constraintTypeToChar(constraintType), rhsExpr.inner, name))
  }



  override def addConstraint(lhsExpr: GurobiLinExpr,
                             constraintType: ConstraintType,
                             rhsVariable: GurobiVariable,
                             name: String): GurobiConstraint = {
    new GurobiConstraint(model.addConstr(lhsExpr.inner, constraintTypeToChar(constraintType), rhsVariable.inner, name))
  }

  override def addConstraint(lhsExpr: GurobiLinExpr,
                             constraintType: ConstraintType,
                             rhs: Double,
                             name: String): GurobiConstraint = {
    new GurobiConstraint(model.addConstr(lhsExpr.inner, constraintTypeToChar(constraintType), rhs, name))
  }

  override def addConstraint(lhsVariable: GurobiVariable, constraintType: ConstraintType, rhsVariable: GurobiVariable, name: String): GurobiConstraint = {
    new GurobiConstraint(model.addConstr(lhsVariable.inner, constraintTypeToChar(constraintType), rhsVariable.inner, name))
  }

  override def addConstraint(lhsVariable: GurobiVariable, constraintType: ConstraintType, rhs: Double, name: String): GurobiConstraint = {
    new GurobiConstraint(model.addConstr(lhsVariable.inner, constraintTypeToChar(constraintType), rhs, name))
  }

  override def removeConstraint(constraint: GurobiConstraint): Unit = {
    model.remove(constraint.inner)
  }

  override def setObjective(solverExpr: GurobiLinExpr): Unit = {
    model.setObjective(solverExpr.inner)
  }

  override def setObjectiveType(objectiveType: ObjectiveType): Unit = {
    model.set(GRB.IntAttr.ModelSense, objectiveTypeToInt(objectiveType))
  }

  override def getObjective: Double = {
    model.get(GRB.DoubleAttr.ObjVal)
  }

  override def optimize(): ModelSolutionType = {
    model.optimize()
//    model.write("test2.lp")
    if(model.get(GRB.IntAttr.Status) == GRB.INFEASIBLE) {
      println("The model is infeasible; computing IIS")
      model.computeIIS()
      System.out.println("\nThe following constraint(s) " + "cannot be satisfied:")
      for (c <- model.getConstrs) {
        if (c.get(GRB.IntAttr.IISConstr) == 1) println(c.get(GRB.StringAttr.ConstrName))
      }
      return Infeasible
    }
    Feasible
  }

  def optimize(timeout: Time) : ModelWithTimeoutSolutionType = {
    println(s"Running with time limit $timeout")
    model.set(GRB.DoubleParam.TimeLimit, timeout.toSeconds)
    model.optimize()
    model.set(GRB.DoubleParam.TimeLimit, Double.PositiveInfinity)
    model.get(GRB.IntAttr.Status) match {
      case GRB.INFEASIBLE => {
        println("The model is infeasible; computing IIS")
        model.computeIIS()
        System.out.println("\nThe following constraint(s) " + "cannot be satisfied:")
        for (c <- model.getConstrs) {
          if (c.get(GRB.IntAttr.IISConstr) == 1) println(c.get(GRB.StringAttr.ConstrName))
        }
        return Infeasible
      }
      case GRB.OPTIMAL => Feasible
      case GRB.TIME_LIMIT => {
        if(model.get(GRB.IntAttr.SolCount) > 0)
          TimeoutSolution
        else
          TimeoutNoSolution
      }
      case code => throw new RuntimeException(s"Unexpected gurobi status code $code")
    }
  }

  override def update(): Unit = {
    model.update()
  }

  override def close(): Unit = {
    model.dispose()
  }

  private def variableTypeToChar(variableType: VariableType): Char = {
    variableType match {
      case Binary => GRB.BINARY
      case Integer => GRB.INTEGER
      case Real => GRB.CONTINUOUS
    }
  }

  private def constraintTypeToChar(constraintType: ConstraintType): Char = {
    constraintType match {
      case Equal => GRB.EQUAL
      case GreaterEqual => GRB.GREATER_EQUAL
      case LessEqual => GRB.LESS_EQUAL
    }
  }

  private def objectiveTypeToInt(objectiveType: ObjectiveType) : Int = {
    objectiveType match {
      case Maximize => GRB.MAXIMIZE
      case Minimize => GRB.MINIMIZE
    }
  }

}
