AI_PRO/TractorAction.py

from TractorLoad import TillageUnit


def action(field, tractor):
    if tractor.header and tractor.hitch == "Crop Trailer" and field.state == "toCut":
        return "toPlow"

    elif isinstance(tractor.hitch, TillageUnit) and tractor.hitch.load == "Nothing" and field.state == "toPlow":
        return "toWater"

    elif isinstance(tractor.hitch, TillageUnit) and tractor.hitch.load == "Water" and field.state == "toWater":
        return "toSeed"

    elif isinstance(tractor.hitch, TillageUnit) and tractor.hitch.load == "Seeds" and field.state == "toSeed":
        return "toFertilize"

    elif isinstance(tractor.hitch, TillageUnit) and tractor.hitch.load == "Fertilizer" and field.state == "toFertilize":
        return "toCut"


def chooseToolset(tractor, tillageUnit, set):
    if set == 0:
        tractor.hitch = tillageUnit
        tractor.hitch.load = "Nothing"
    if set == 1:
        tractor.hitch = tillageUnit
        tractor.hitch.load = "Water"
    if set == 2:
        tractor.hitch = tillageUnit
        tractor.hitch.load = "Seeds"
    if set == 3:
        tractor.hitch = tillageUnit
        tractor.hitch.load = "Fertilizer"
    if set == 4:
        tractor.header = True
        tractor.hitch = "Crop Trailer"

    return tractor, tillageUnit