Advanced constraint programming modeling

advancedModelingWeek1.py

from ortools.sat.python import cp_model

''' course: https://www.coursera.org/learn/advanced-modeling#about '''


def theCavalryWedgeProblem():

    model = cp_model.CpModel()
    horseList = ["H1", "H2", "H3", "H4", "H5", "H6", "H7", "H8", "H9", "H10"]
    riderList = ["R1", "R2", "R3", "R4", "R5", "R6", "R7", "R8", "R9", "R10", "R11"]
    speedList = [10, 9, 8, 7, 6, 5, 7, 4, 3, 2]
    enduranceList = [8, 4, 3, 2, 6, 4, 2, 6, 7, 5, 3]
    strengthList = [5, 2, 8, 9, 4, 2, 1, 3, 4, 5, 9]
    compatibleList = [[0,1,1,0,0,0,0,0,0,0,1], [0,0,0,0,1,1,0,0,0,0,0], [0,0,0,0,0,0,0,1,0,0,0],
                      [1,0,0,0,1,0,0,0,0,0,0], [0,0,0,1,0,0,0,0,0,0,0], [0,1,0,0,0,0,1,0,0,0,0],
                      [1,0,1,0,0,0,0,0,0,0,0], [1,0,0,0,0,0,0,0,1,1,0], [0,0,1,0,0,0,0,0,0,0,1],
                      [0,0,0,0,0,0,1,0,1,1,0]]
    sizeOfWedge = 5 # should be odd
    assert sizeOfWedge%2 == 1, "sizeOfWedge (currentValue:" +str(sizeOfWedge)+" must be odd"
    numOfHorses = len(horseList)
    numOfRiders = len(riderList)
    maxSpeed = max(speedList)
    maxEndurance = max(enduranceList)
    maxStrength = max(strengthList)
    compatibleListFlattened = sum(compatibleList, [])


    positionHorseVarList = []
    positionRiderVarList = []
    for position in range(sizeOfWedge):
        positionHorseVarList.append(model.NewIntVar(0, numOfHorses-1, "HorseForPosition"+str(position)))
        positionRiderVarList.append(model.NewIntVar(0, numOfRiders-1, "RiderForPosition"+str(position)))


    model.AddAllDifferent(positionHorseVarList)
    model.AddAllDifferent(positionRiderVarList)


    speedVarList = []
    enduranceVarList = []
    strengthVarList = []
    for position in range(sizeOfWedge):
        speedVarList.append(model.NewIntVar(0, maxSpeed, "Speed"+str(position)))
        enduranceVarList.append(model.NewIntVar(0, maxEndurance, "Endurance"+str(position)))
        strengthVarList.append(model.NewIntVar(0, maxStrength, "Strength"+str(position)))
        model.AddElement(positionHorseVarList[position], speedList, speedVarList[position])
        model.AddElement(positionRiderVarList[position], enduranceList, enduranceVarList[position])
        model.AddElement(positionRiderVarList[position], strengthList, strengthVarList[position])


    for position in range(sizeOfWedge//2):
        print("speed[" + str(position) + "] < speed[" + str(position + 1) + "]")
        model.Add(speedVarList[position] < speedVarList[position+1])
        print("endur[" + str(position) + "] < endur[" + str(position + 1) + "]")
        model.Add(enduranceVarList[position] < enduranceVarList[position+1])

    for position in range(sizeOfWedge//2, sizeOfWedge-1):
        print("speed[" + str(position) + "] > speed[" + str(position + 1) + "]")
        model.Add(speedVarList[position] > speedVarList[position+1])
        print("endur[" + str(position) + "] > endur[" + str(position + 1) + "]")
        model.Add(enduranceVarList[position] > enduranceVarList[position+1])


    compatibilityVarList = []
    for position in range(sizeOfWedge):
        indexVar = model.NewIntVar(0, numOfRiders*numOfHorses, "indexVar"+str(position))
        compatibilityVarList.append(model.NewBoolVar("compatibility" + str(position)))
        model.Add(indexVar == (numOfRiders * positionHorseVarList[position]) + positionRiderVarList[position])
        model.AddElement(indexVar, compatibleListFlattened, compatibilityVarList[position])
        model.Add(compatibilityVarList[position] == 1)

    model.Maximize(sum(strengthVarList))


    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Maximum strength: %i' % solver.ObjectiveValue(), "\n")

        for position in range(sizeOfWedge):
            horseIndex = solver.Value(positionHorseVarList[position])
            riderIndex = solver.Value(positionRiderVarList[position])
            print("Position: "+str(position)+ " - Horse: "+str(horseList[horseIndex])+" - Rider: "+str(riderList[riderIndex])+" - Compatibility: "+str(solver.Value(compatibilityVarList[position])))
    print("----------------------------")



#----------------------------------------------------------------------------------------------------------------------#



def pickTheEscapePathProblem():

    model = cp_model.CpModel()

    pathTimeList = [2, 4, 4, 4, 14, 10, 8, 14, 12, 10]
    escapeGuardTimeList = [10, 8, 2, 4, 0, 0, 0, 0, 0, 0]

    NUM_OF_WEATHERS = 3 # 0 = storm, 1 = rain, 2 = sun
    numOfPaths = len(pathTimeList)
    maxTimeForPath = max(pathTimeList)
    maxTimeForEscapingGuard = max(escapeGuardTimeList)

    weatherVar = model.NewIntVar(1, NUM_OF_WEATHERS-1, "Weather") # can't take storm because no solution
    pathVar = model.NewIntVar(0, numOfPaths-1, "Path")
    timePathOnWeatherVar = model.NewIntVar(0, maxTimeForPath, "TimeOnPathOnWeather")
    timeEscapingGuardOnWeatherVar = model.NewIntVar(0, maxTimeForEscapingGuard * NUM_OF_WEATHERS-1, "TimeEscapingGuardOnWeather")

    guardTimeVar = model.NewIntVar(0, maxTimeForEscapingGuard, "GuardTime")
    model.AddElement(pathVar, escapeGuardTimeList, guardTimeVar)
    model.AddMultiplicationEquality(timeEscapingGuardOnWeatherVar, [weatherVar, guardTimeVar])

    pathTimeVar = model.NewIntVar(0, maxTimeForPath, "PathTime")
    model.AddElement(pathVar, pathTimeList, pathTimeVar)
    model.AddDivisionEquality(timePathOnWeatherVar, pathTimeVar, weatherVar)

    ObjVar = model.NewIntVar(0, maxTimeForPath + (maxTimeForEscapingGuard * NUM_OF_WEATHERS-1), "ObjVar")
    model.Add(ObjVar == timeEscapingGuardOnWeatherVar + timePathOnWeatherVar)
    model.Minimize(ObjVar)


    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print("Path: "+str(solver.Value(pathVar))+ ", MinimumTime: "+str(solver.Value(ObjVar)))
    print("----------------------------")




#----------------------------------------------------------------------------------------------------------------------#



def twoGeneralsProblem():

    model = cp_model.CpModel()

    heroList = ["hero1", "hero2", "hero3"]
    spotList = ["Baihui", "Danzhong", "Quchi", "Tianshu", "Yongquan"]
    damageList = [[9, 7, 6, 5, 8], [8, 2, 5, 1, 4], [4, 3, 7, 2, 5]]

    damageListFlattened = sum(damageList, [])
    maxDamage = max(damageListFlattened)
    numOfSpots = len(spotList)
    numOfHeroes = len(heroList)
    lowerPositionsAfterIndex = 2

    position1VarList = []
    position2VarList = []
    for hero in heroList:
        position1VarList.append(model.NewIntVar(1, numOfSpots-1, "position1"+str(hero)))  # position1 cant be equal to Baihui
        position2VarList.append(model.NewIntVar(0, numOfSpots - 1, "position2"+str(hero)))

    model.AddAllDifferent(position1VarList)
    model.AddAllDifferent(position2VarList)

    for i in range(numOfHeroes):
        strikedLowerPosFirstVar = model.NewBoolVar("StrikedLowerPositionFirst")
        model.Add(position1VarList[i] > lowerPositionsAfterIndex).OnlyEnforceIf(strikedLowerPosFirstVar)
        model.Add(position1VarList[i] <= lowerPositionsAfterIndex).OnlyEnforceIf(strikedLowerPosFirstVar.Not())
        model.Add(position2VarList[i] <= lowerPositionsAfterIndex).OnlyEnforceIf(strikedLowerPosFirstVar)

    damage1VarList = []
    index1VarList = []
    damage2VarList = []
    index2VarList = []
    for i, hero in enumerate(heroList):
        damage1VarList.append(model.NewIntVar(0, maxDamage, "Damage1"+str(hero)))
        index1VarList.append(model.NewIntVar(0, numOfSpots*numOfHeroes, "Index1"+str(hero)))
        model.Add(index1VarList[i] == (numOfSpots * i) + position1VarList[i])
        model.AddElement(index1VarList[i], damageListFlattened, damage1VarList[i])

        damage2VarList.append(model.NewIntVar(0, maxDamage, "Damage2" + str(hero)))
        index2VarList.append(model.NewIntVar(0, numOfSpots*numOfHeroes, "Index2"+str(hero)))
        model.Add(index2VarList[i] == (numOfSpots * i) + position2VarList[i])
        model.AddElement(index2VarList[i], damageListFlattened, damage2VarList[i])

    model.Maximize(sum(damage1VarList) + sum(damage2VarList))


    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Maximum damage: %i' % solver.ObjectiveValue(), "\n")

        for i, hero in enumerate(heroList):
            print("Hero: "+str(hero)+ " - position1: "+spotList[solver.Value(position1VarList[i])] +", position2: "+spotList[solver.Value(position2VarList[i])])
    print("----------------------------")



#----------------------------------------------------------------------------------------------------------------------#



def planningEscapeProblem():

    model = cp_model.CpModel()

    guardList = [0, 8, 3, 5, 0, 3, 9, 8, 0, 4, 0, 0, 10, 8, 14, 8]
    edgeList = [[0, 1], [0, 9], [1, 2], [1, 6], [2, 3],
                [2, 7], [2, 10], [3, 4], [3, 7], [3, 15],
                [4, 15], [5, 9], [5, 10], [6, 7], [6, 11],
                [7, 9], [7, 11], [8, 8], [8, 14], [8, 15],
                [10, 13], [11, 12], [11, 13], [11, 14], [12, 14],
                [12, 15], [13, 14]]

    undirectionalEdgeList = []
    for edgePair in edgeList:
        undirectionalEdgeList.append(edgePair)
        undirectionalEdgeList.append([edgePair[1], edgePair[0]])

    REST = 3
    START = 0
    DESTINATION = 8
    numOfNodes = len(guardList)
    MAX_STEP = 10

    stepVar = model.NewIntVar(0, MAX_STEP-1, "stepVar")

    pathVarList = []
    for step in range(MAX_STEP):
        pathVarList.append(model.NewIntVar(0, numOfNodes-1, "NodeForPath"+str(step)))

    model.Add(pathVarList[0] == START)

    for i in range(MAX_STEP-1):
        model.AddAllowedAssignments([pathVarList[i], pathVarList[i+1]], undirectionalEdgeList)


    stepVarLEVarList = []
    for i in range(MAX_STEP):
        stepVarLEVarList.append(model.NewBoolVar("stepVarLE"+str(i)))
        model.Add(i >= stepVar).OnlyEnforceIf(stepVarLEVarList[i])
        model.Add(i < stepVar).OnlyEnforceIf(stepVarLEVarList[i].Not())
        model.Add(pathVarList[i] == DESTINATION).OnlyEnforceIf(stepVarLEVarList[i])


    guardVarList = []
    hasNoGuardVarList = []
    for i in range(MAX_STEP):
        guardVarList.append(model.NewIntVar(0, max(guardList), "NumOfGuardsAt" + str(i)))
        model.AddElement(pathVarList[i], guardList, guardVarList[i])

        hasNoGuardVarList.append(model.NewBoolVar("HasNoGuard"+str(i)))
        model.Add(guardVarList[i] == 0).OnlyEnforceIf(hasNoGuardVarList[i])
        model.Add(guardVarList[i] != 0).OnlyEnforceIf(hasNoGuardVarList[i].Not())


    for i in range(MAX_STEP - REST):
        model.AddLinearConstraint(sum([hasNoGuardVarList[i], hasNoGuardVarList[i+1], hasNoGuardVarList[i+2]]), 1, REST)

    model.Minimize(sum(guardVarList))

    solver = cp_model.CpSolver()
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Minimum guards encountered : %i' % solver.ObjectiveValue(), "\n")

        solList = []
        for i in range(MAX_STEP):
            solList.append(solver.Value(pathVarList[i]))
        print("Path =", solList)
    print("----------------------------")



#----------------------------------------------------------------------------------------------------------------------#




class VarArraySolutionPrinter(cp_model.CpSolverSolutionCallback):
    """Print intermediate solutions."""

    def __init__(self, variables):
        cp_model.CpSolverSolutionCallback.__init__(self)
        self.__variables = variables
        self.__solution_count = 0

    def on_solution_callback(self):
        self.__solution_count += 1
        for v in self.__variables:
            print('%s=%i' % (v, self.Value(v)), end=' ')
        print()

    def solution_count(self):
        return self.__solution_count




def paradeProblem():

    model = cp_model.CpModel()

    STRAW_SOLDIER_OR_EMPTY = 0
    heightList = [6, 7, 9, 13, 12, 9, 8, 10, 11, 8, 3, 4, 6, 6, 7]
    numOfSoldiers = len(heightList)
    MAX_ROWS = 8
    MAX_COLUMNS = 8

    positionVarList = []
    positionDict = {}
    counter = 0
    for row in range(MAX_ROWS):
        for col in range(MAX_COLUMNS):
            positionVarList.append(model.NewIntVar(0, numOfSoldiers-1, "positionFor"+str(row)+str(col)))
            positionDict[(row, col)] = counter
            counter += 1

    rowVar = model.NewIntVar(0, MAX_ROWS-1, "RowVar")
    colVar = model.NewIntVar(0, MAX_COLUMNS-1, "ColumnVar")


    heightVarList = []
    for row in range(MAX_ROWS):
        for col in range(MAX_COLUMNS):
            heightVarList.append(model.NewIntVar(0, max(heightList), "HeightAt"+str(row)+str(col)))
            model.AddElement(positionVarList[positionDict[(row, col)]], heightList, heightVarList[-1])


    for soldier in range(1, numOfSoldiers): # 0 is straw soldier
        isSoldierVarList = []
        for row in range(MAX_ROWS):
            for col in range(MAX_COLUMNS):
                isSoldierVarList.append(model.NewBoolVar("isSoldier"+str(soldier)+"at"+str(row)+str(col)))
                model.Add(positionVarList[positionDict[(row, col)]] == soldier).OnlyEnforceIf(isSoldierVarList[-1])
                model.Add(positionVarList[positionDict[(row, col)]] != soldier).OnlyEnforceIf(isSoldierVarList[-1].Not())
        model.Add(sum(isSoldierVarList) == 1)


    for row in range(MAX_ROWS):
        for col in range(MAX_COLUMNS):
            isEmptyVar1 = model.NewBoolVar("IsEmpty1"+str(row)+str(col))
            model.Add(row > rowVar).OnlyEnforceIf(isEmptyVar1)
            model.Add(row <= rowVar).OnlyEnforceIf(isEmptyVar1.Not())
            model.Add(positionVarList[positionDict[(row, col)]] == STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(isEmptyVar1)

            isEmptyVar2 = model.NewBoolVar("IsEmpty2" + str(row) + str(col))
            model.Add(col > colVar).OnlyEnforceIf(isEmptyVar2)
            model.Add(col <= colVar).OnlyEnforceIf(isEmptyVar2.Not())
            model.Add(positionVarList[positionDict[(row, col)]] == STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(isEmptyVar2)


    isStrawSoldierVarList = []
    for row in range(MAX_ROWS):
        for col in range(MAX_COLUMNS):

            isLERowVar = model.NewBoolVar("isLERowVar"+str(row)+str(col))
            isLEColVar = model.NewBoolVar("isLEColVar"+str(row)+str(col))
            equalsStrawOrEmptyVar = model.NewBoolVar("equalsStrawOrEmpty"+str(row)+str(col))
            isStrawSoldierVarList.append(model.NewBoolVar("isStrawSoldier"+str(row)+str(col)))

            model.AddBoolAnd([isLERowVar, isLEColVar, equalsStrawOrEmptyVar]).OnlyEnforceIf(isStrawSoldierVarList[-1])
            model.AddBoolOr([isLERowVar.Not(), isLEColVar.Not(), equalsStrawOrEmptyVar.Not(), isStrawSoldierVarList[-1]])

            model.Add(row <= rowVar).OnlyEnforceIf(isLERowVar)
            model.Add(row > rowVar).OnlyEnforceIf(isLERowVar.Not())
            model.Add(col <= colVar).OnlyEnforceIf(isLEColVar)
            model.Add(col > colVar).OnlyEnforceIf(isLEColVar.Not())
            model.Add(positionVarList[positionDict[(row,col)]] == STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(equalsStrawOrEmptyVar)
            model.Add(positionVarList[positionDict[(row, col)]] != STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(equalsStrawOrEmptyVar.Not())

            hasRealSoldierToTheRightVar = model.NewBoolVar("hasRealSoldierToTheRight"+str(row)+str(col))
            if col < MAX_COLUMNS-1:
                model.Add(positionVarList[positionDict[(row, col+1)]] != STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(hasRealSoldierToTheRightVar)

            hasRealSoldierToTheLeftVar = model.NewBoolVar("hasRealSoldierToTheLeft" + str(row) + str(col))
            if col > 0:
                model.Add(positionVarList[positionDict[(row, col-1)]] != STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(hasRealSoldierToTheLeftVar)

            hasRealSoldierToTheBackVar = model.NewBoolVar("hasRealSoldierToTheBack"+str(row)+str(col))
            if row < MAX_ROWS-1:
                model.Add(positionVarList[positionDict[(row+1, col)]] != STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(hasRealSoldierToTheBackVar)

            hasRealSoldierToTheFrontVar = model.NewBoolVar("hasRealSoldierToTheFront"+str(row)+str(col))
            if row > 0:
                model.Add(positionVarList[positionDict[(row-1, col)]] != STRAW_SOLDIER_OR_EMPTY).OnlyEnforceIf(hasRealSoldierToTheFrontVar)

            model.AddBoolOr([hasRealSoldierToTheRightVar, hasRealSoldierToTheLeftVar,
                             hasRealSoldierToTheBackVar, hasRealSoldierToTheFrontVar]).OnlyEnforceIf(isStrawSoldierVarList[-1])


    for row in range(1, MAX_ROWS):
        for col in range(MAX_COLUMNS):
            varList = []
            for row2 in range(row):
                varList.append(heightVarList[positionDict[(row2, col)]])

            maxVar = model.NewIntVar(0, max(heightList), "maxVar"+str(row)+str(col))
            model.AddMaxEquality(maxVar, varList)
            indexOfList = positionDict[(row, col)]
            model.Add(maxVar > heightVarList[indexOfList]).OnlyEnforceIf(isStrawSoldierVarList[indexOfList])


    objVar = model.NewIntVar(0, MAX_ROWS*MAX_COLUMNS, "ObjVar")
    model.AddMultiplicationEquality(objVar, [rowVar, colVar])
    # model.Maximize(objVar) # Slow to solve till optimality.


    solList = [7, 9, 13, 12, 10, 8, 9, 6,
               6, 6, 6, 6, 6, 6, 6, 6,
               6, 11, 6, 6, 6, 8, 6, 6,
               6, 6, 6, 3, 6, 6, 6, 6,
               4, 6, 6, 6, 6, 6, 6, 6,
               6, 6, 6, 6, 7, 6, 6, 6,
               6, 6, 6, 6, 6, 6, 6, 6,
               6, 6, 6, 6, 6, 6, 6, 6]
    for i, heightVar in enumerate(heightVarList):
        model.AddHint(heightVar, solList[i])


    solver = cp_model.CpSolver()
    # solution_printer = VarArraySolutionPrinter(heightVarList)
    # status = solver.SearchForAllSolutions(model, solution_printer)
    status = solver.Solve(model)

    if status == cp_model.INFEASIBLE:
        print("Infeasible model")

    if status == cp_model.OPTIMAL:
        print('Max row * column : %i' % solver.ObjectiveValue(), "\n")

        for row in range(MAX_ROWS):
            for col in range(MAX_COLUMNS):
                index = positionDict[(row, col)]
                height = solver.Value(heightVarList[index])
                if solver.Value(positionVarList[index]) == STRAW_SOLDIER_OR_EMPTY:
                    print(" . ", end=" ")
                else:
                    print(" "+str(height)+" ", end=" ")
            print()
        print(solver.Value(rowVar), "rowVar")
        print(solver.Value(colVar), "colVar")
    print("----------------------------")



#----------------------------------------------------------------------------------------------------------------------#





if __name__ == "__main__":

    theCavalryWedgeProblem()
    pickTheEscapePathProblem()
    twoGeneralsProblem()
    planningEscapeProblem()
    paradeProblem()