SzybciorSmartTraktor/main.py

import math
import random
import time
import pygame
import torch
import agent
import astar
import common
import field
import neural_network
import settings
import tree
import shutup

possibleFields = {
    'dirt': field.Dirt(),
    'grass': field.Grass(),
    'cobble': field.Cobble(),
    'sand': field.Sand(),
    'station': field.Station(),
    'carrot_empty': field.Carrot('carrot_empty'),
    'carrot_sow': field.Carrot('carrot_sow'),
    'carrot_watered': field.Carrot('carrot_watered'),
    'carrot_feritized': field.Carrot('carrot_feritized'),
    'potato_empty': field.Potato('potato_empty'),
    'potato_sow': field.Potato('potato_sow'),
    'potato_watered': field.Potato('potato_watered'),
    'potato_feritized': field.Potato('potato_feritized'),
    'wheat_empty': field.Wheat('wheat_empty'),
    'wheat_sow': field.Wheat('wheat_sow'),
    'wheat_watered': field.Wheat('wheat_watered'),
    'wheat_feritized': field.Wheat('wheat_feritized')
}
possibleFieldsWithPlants = [
    'carrot_empty',
    'carrot_sow',
    'carrot_watered',
    'carrot_feritized',
    'potato_empty',
    'potato_sow',
    'potato_watered',
    'potato_feritized',
    'wheat_empty',
    'wheat_sow',
    'wheat_watered',
    'wheat_feritized'
]

shutup.please()


def randomize_map():
    fields_array = []
    for i in possibleFields:
        fields_array.append(possibleFields[i].tile.object)
    field_array_big = []
    field_array_small = []
    field_array_big_2 = []
    field_array_big_3 = []
    field_array_small_2 = []
    field_array_small_3 = []
    width = settings.Field.horizontal_count()
    height = settings.Field.vertical_count()
    for i in range(width):
        for j in range(height):
            # k = random.choice(list(possibleFields.keys()))
            x = random.uniform(0, 100)
            if x <= 80:
                plant = random.choice(possibleFieldsWithPlants)
                field_array_small.append(possibleFields[plant].tile.object)
                field_array_small_2.append('dirt')
                field_array_small_3.append(plant)
            elif 80 < x <= 90:
                field_array_small.append(possibleFields['sand'].tile.object)
                field_array_small_2.append('sand')
                field_array_small_3.append('sand')
            elif 90 < x <= 100:
                field_array_small.append(possibleFields['grass'].tile.object)
                field_array_small_2.append('grass')
                field_array_small_3.append('grass')
        field_array_big.append(field_array_small)
        field_array_big_2.append(field_array_small_2)
        field_array_big_3.append(field_array_small_3)
        field_array_small = []
        field_array_small_2 = []
        field_array_small_3 = []

    for i in range(height):
        field_array_big[math.floor(width / 2)][i] = possibleFields['cobble'].tile.object
        field_array_big_2[math.floor(width / 2)][i] = 'cobble'
        field_array_big_3[math.floor(width / 2)][i] = 'cobble'
    for i in range(width):
        field_array_big[i][math.floor(height / 2)] = possibleFields['cobble'].tile.object
        field_array_big_2[i][math.floor(height / 2)] = 'cobble'
        field_array_big_3[i][math.floor(height / 2)] = 'cobble'
    field_array_big[0][0] = possibleFields['station'].tile.object
    field_array_big_2[0][0] = 'station'
    field_array_big_3[0][0] = 'station'
    return field_array_big, field_array_big_2, field_array_big_3


def get_plants_array(fields):
    field_array_small = []
    field_array_big = []
    for i in range(11):
        for j in range(11):
            if fields[i][j] == 'carrot_empty' or fields[i][j] == 'carrot_sow' or fields[i][j] == 'carrot_watered' or \
                    fields[i][j] == 'carrot_feritized':
                field_array_small.append('carrot')
            elif fields[i][j] == 'potato_empty' or fields[i][j] == 'potato_sow' or fields[i][j] == 'potato_watered' or \
                    fields[i][j] == 'potato_feritized':
                field_array_small.append('potato')
            elif fields[i][j] == 'wheat_empty' or fields[i][j] == 'wheat_sow' or fields[i][j] == 'wheat_watered' or \
                    fields[i][j] == 'wheat_feritized':
                field_array_small.append('wheat')
            else:
                field_array_small.append('none')
        field_array_big.append(field_array_small)
        field_array_small = []
    return field_array_big


def recognize_plants(plants_array, fields_for_astar, fields_for_movement, agent):
    checkpoint = torch.load(f'plants.model')
    model = neural_network.Net(num_classes=3)
    model.load_state_dict(checkpoint)
    model.eval()
    img = ''
    b=0
    j=0
    field_array_small = []
    field_array_big = []
    for i in range(11):
        field_array_small = []
        if b == 0:
            for j in range(11):
                if plants_array[j][i] == 'carrot':
                    img = 'assets/learning/test/carrot/' + str(random.randint(1, 200)) + '.jpg'
                    pred = neural_network.prediction(img, model)
                    show_plant_img(img)
                elif plants_array[j][i] == 'potato':
                    img = 'assets/learning/test/potato/' + str(random.randint(1, 200)) + '.jpg'
                    pred = neural_network.prediction(img, model)
                    show_plant_img(img)
                elif plants_array[j][i] == 'wheat':
                    img = 'assets/learning/test/wheat/' + str(random.randint(1, 200)) + '.jpg'
                    pred = neural_network.prediction(img, model)
                    show_plant_img(img)
                else:
                    pred = 'none'
                field_array_small.append(pred)
                print(i,',', j,'-',pred)
                agent_movement(['f'], agent, fields_for_movement, fields_for_astar)
            agent_movement(['r','f','r'], agent, fields_for_movement, fields_for_astar)
            field_array_big.append(field_array_small)
        else:
            for j in range(10,-1,-1):
                if plants_array[j][i] == 'carrot':
                    img = 'assets/learning/test/carrot/' + str(random.randint(1, 200)) + '.jpg'
                    pred = neural_network.prediction(img, model)
                    show_plant_img(img)
                elif plants_array[j][i] == 'potato':
                    img = 'assets/learning/test/potato/' + str(random.randint(1, 200)) + '.jpg'
                    pred = neural_network.prediction(img, model)
                    show_plant_img(img)
                elif plants_array[j][i] == 'wheat':
                    img = 'assets/learning/test/wheat/' + str(random.randint(1, 200)) + '.jpg'
                    pred = neural_network.prediction(img, model)
                    show_plant_img(img)
                else:
                    pred = 'none'
                field_array_small.append(pred)
                print(i,',', j,'-',pred)
                agent_movement(['f'], agent, fields_for_movement, fields_for_astar)
            field_array_small = field_array_small[::-1]
            field_array_big.append(field_array_small)
            agent_movement(['l','f','l'], agent, fields_for_movement, fields_for_astar)
        if b==0:
            b=1
        else:
            b=0
    correct = 0
    incorrect = 0
    for i in range(11):
        for j in range(11):
            if plants_array[i][j]=='none':
                continue
            else:
                if plants_array[i][j]==field_array_big[j][i]:
                    correct+=1
                else:
                    incorrect+=1
    print("Accuracy: ",correct/(correct+incorrect)*100,'%')



def read_img(agent, fields):
    window = common.get('window')
    current_field = fields[agent.x()][agent.y()]
    if current_field == possibleFields['grass'].tile.object:
        window.blit(possibleFields['grass'].block.object, (0, 0))
    elif current_field == possibleFields['dirt'].tile.object:
        window.blit(possibleFields['dirt'].block.object, (0, 0))
    elif current_field == possibleFields['sand'].tile.object:
        window.blit(possibleFields['sand'].block.object, (0, 0))
    elif current_field == possibleFields['cobble'].tile.object:
        window.blit(possibleFields['cobble'].block.object, (0, 0))
    elif current_field == possibleFields['station'].tile.object:
        window.blit(possibleFields['station'].block.object, (0, 0))

    pygame.display.update()
    pygame.time.delay(2000)
    common.set('state_imgShown', False)

def show_plant_img(img):
    window = common.get('window')
    image = pygame.image.load(img)
    image = pygame.transform.scale(image, (11*settings.Field.size(), 11*settings.Field.size()))
    window.blit(image,(0,0))
    pygame.display.update()
    pygame.time.delay(2000)

def agent_action(action: str):
    if action == 'open_window':
        common.set('state_imgShown', True)


def draw_window(agent, fields):
    window = common.get('window')
    rect = agent.rect
    for i in range(settings.Field.horizontal_count()):
        for j in range(settings.Field.vertical_count()):
            window.blit(fields[i][j], (i * settings.Field.size(), j * settings.Field.size()))
    if agent.direction == 2:
        AGENT_IMG = pygame.image.load('./assets/tracktor/tractor_east.png')
        AGENT = pygame.transform.scale(AGENT_IMG, (settings.Field.size(), settings.Field.size()))
        window.blit(AGENT, (rect.x, rect.y))
    elif agent.direction == 4:
        AGENT_IMG = pygame.image.load('./assets/tracktor/tractor_west.png')
        AGENT = pygame.transform.scale(AGENT_IMG, (settings.Field.size(), settings.Field.size()))
        window.blit(AGENT, (rect.x, rect.y))
    elif agent.direction == 1:
        AGENT_IMG = pygame.image.load('./assets/tracktor/tractor_north.png')
        AGENT = pygame.transform.scale(AGENT_IMG, (settings.Field.size(), settings.Field.size()))
        window.blit(AGENT, (rect.x, rect.y))
    elif agent.direction == 3:
        AGENT_IMG = pygame.image.load('./assets/tracktor/tractor_south.png')
        AGENT = pygame.transform.scale(AGENT_IMG, (settings.Field.size(), settings.Field.size()))
        window.blit(AGENT, (rect.x, rect.y))
    pygame.display.update()


def generate_movement(fields_for_astar, fields_with_plants, fields_for_movement, agent, t):
    fields_to_sow = []
    fields_to_water = []
    fields_to_feritize = []
    fields_to_harvest = []
    width = settings.Field.horizontal_count()
    height = settings.Field.vertical_count()
    k = 0
    for i in range(width):
        for j in range(height):
            if fields_with_plants[i][j] == 'potato_empty' or fields_with_plants[i][j] == 'carrot_empty' or \
                    fields_with_plants[i][j] == 'wheat_empty':
                tab = [i, j]
                fields_to_sow.append(tab)
            elif fields_with_plants[i][j] == 'potato_sow' or fields_with_plants[i][j] == 'carrot_sow' or \
                    fields_with_plants[i][j] == 'wheat_sow':
                tab = [i, j]
                fields_to_water.append(tab)
            elif fields_with_plants[i][j] == 'potato_watered' or fields_with_plants[i][j] == 'carrot_watered' or \
                    fields_with_plants[i][j] == 'wheat_watered':
                tab = [i, j]
                fields_to_feritize.append(tab)
            elif fields_with_plants[i][j] == 'potato_feritized' or fields_with_plants[i][j] == 'carrot_feritized' or \
                    fields_with_plants[i][j] == 'wheat_feritized':
                tab = [i, j]
                fields_to_harvest.append(tab)
    while True:
        cords = agent.coordinates()
        x = cords['x']
        y = cords['y']
        dir = agent.get_direction()
        fuel = agent.get_tank_capacity()
        water = agent.get_water()
        feritizer = agent.get_feritizer()
        seeds = agent.get_seeds()
        carrots = agent.get_carrots()
        potatoes = agent.get_potatoes()
        wheat = agent.get_wheat()
        decision = tree.make_decision(t, fuel, water, feritizer, carrots, potatoes, wheat, x, y, seeds)
        print("fuel:", fuel, "water:", water, "feritizer:", feritizer, "seeds:", seeds, "carrots:", carrots,
              "potatoes:", potatoes, "wheat:", wheat)
        if k == 0 and len(fields_to_harvest) > 0 and decision == False:  # harvest
            field_to_visit, l = get_closest_field(fields_to_harvest, x, y, dir, fields_for_astar)
            x1, y1 = field_to_visit[0], field_to_visit[1]
            del fields_to_harvest[l]
            if fields_with_plants[x1][y1] == 'wheat_feritized':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_sow.append(field_to_visit)
                fields_with_plants[x1][y1] = 'wheat_empty'
                fields_for_movement[x1][y1] = possibleFields['wheat_empty'].tile.object
                draw_window(agent, fields_for_movement)
                wheat += 1
                agent.set_wheat(wheat)
            elif fields_with_plants[x1][y1] == 'carrot_feritized':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_sow.append(field_to_visit)
                fields_with_plants[x1][y1] = 'carrot_empty'
                fields_for_movement[x1][y1] = possibleFields['carrot_empty'].tile.object
                draw_window(agent, fields_for_movement)
                carrots += 1
                agent.set_carrots(carrots)
            elif fields_with_plants[x1][y1] == 'potato_feritized':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_sow.append(field_to_visit)
                fields_with_plants[x1][y1] = 'potato_empty'
                fields_for_movement[x1][y1] = possibleFields['potato_empty'].tile.object
                draw_window(agent, fields_for_movement)
                potatoes += 1
                agent.set_potatoes(potatoes)
            if len(fields_to_harvest) == 0:
                k += 1
        elif k == 1 and len(fields_to_water) > 0 and decision == False:  # water
            field_to_visit, l = get_closest_field(fields_to_water, x, y, dir, fields_for_astar)
            del fields_to_water[l]
            x1, y1 = field_to_visit[0], field_to_visit[1]
            if fields_with_plants[x1][y1] == 'wheat_sow':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_feritize.append(field_to_visit)
                fields_with_plants[x1][y1] = 'wheat_watered'
                fields_for_movement[x1][y1] = possibleFields['wheat_watered'].tile.object
                draw_window(agent, fields_for_movement)
                water -= 1
                agent.set_water(water)
            elif fields_with_plants[x1][y1] == 'carrot_sow':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_feritize.append(field_to_visit)
                fields_with_plants[x1][y1] = 'carrot_watered'
                fields_for_movement[x1][y1] = possibleFields['carrot_watered'].tile.object
                draw_window(agent, fields_for_movement)
                water -= 1
                agent.set_water(water)
            elif fields_with_plants[x1][y1] == 'potato_sow':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_feritize.append(field_to_visit)
                fields_with_plants[x1][y1] = 'potato_watered'
                fields_for_movement[x1][y1] = possibleFields['potato_watered'].tile.object
                draw_window(agent, fields_for_movement)
                water -= 1
                agent.set_water(water)
            if len(fields_to_water) == 0:
                k += 1
        elif k == 2 and len(fields_to_feritize) > 0 and decision == False:  # feritize
            field_to_visit, l = get_closest_field(fields_to_feritize, x, y, dir, fields_for_astar)
            del fields_to_feritize[l]
            x1, y1 = field_to_visit[0], field_to_visit[1]
            if fields_with_plants[x1][y1] == 'wheat_watered':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_harvest.append(field_to_visit)
                fields_with_plants[x1][y1] = 'wheat_feritized'
                fields_for_movement[x1][y1] = possibleFields['wheat_feritized'].tile.object
                draw_window(agent, fields_for_movement)
                feritizer -= 1
                agent.set_feritizer(feritizer)
            elif fields_with_plants[x1][y1] == 'carrot_watered':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_harvest.append(field_to_visit)
                fields_with_plants[x1][y1] = 'carrot_feritized'
                fields_for_movement[x1][y1] = possibleFields['carrot_feritized'].tile.object
                draw_window(agent, fields_for_movement)
                feritizer -= 1
                agent.set_feritizer(feritizer)
            elif fields_with_plants[x1][y1] == 'potato_watered':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_harvest.append(field_to_visit)
                fields_with_plants[x1][y1] = 'potato_feritized'
                fields_for_movement[x1][y1] = possibleFields['potato_feritized'].tile.object
                draw_window(agent, fields_for_movement)
                feritizer -= 1
                agent.set_feritizer(feritizer)
            if len(fields_to_feritize) == 0:
                k += 1
        elif k == 3 and len(fields_to_sow) > 0 and decision == False:  # sow
            field_to_visit, l = get_closest_field(fields_to_sow, x, y, dir, fields_for_astar)
            del fields_to_sow[l]
            x1, y1 = field_to_visit[0], field_to_visit[1]
            if fields_with_plants[x1][y1] == 'wheat_empty':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_water.append(field_to_visit)
                fields_with_plants[x1][y1] = 'wheat_sow'
                fields_for_movement[x1][y1] = possibleFields['wheat_sow'].tile.object
                draw_window(agent, fields_for_movement)
                seeds -= 1
                agent.set_seeds(seeds)
            elif fields_with_plants[x1][y1] == 'carrot_empty':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_water.append(field_to_visit)
                fields_with_plants[x1][y1] = 'carrot_sow'
                fields_for_movement[x1][y1] = possibleFields['carrot_sow'].tile.object
                draw_window(agent, fields_for_movement)
                seeds -= 1
                agent.set_seeds(seeds)
            elif fields_with_plants[x1][y1] == 'potato_empty':
                state = astar.State(dir, x, y)
                move_list = (
                    astar.graphsearch([], astar.f, [], [field_to_visit[0], field_to_visit[1]], state, fields_for_astar,
                                      astar.succ))
                agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
                fields_to_water.append(field_to_visit)
                fields_with_plants[x1][y1] = 'potato_sow'
                fields_for_movement[x1][y1] = possibleFields['potato_sow'].tile.object
                draw_window(agent, fields_for_movement)
                seeds -= 1
                agent.set_seeds(seeds)
            if len(fields_to_sow) == 0:
                k = 0
        elif decision:
            print("Going back to base")
            state = astar.State(dir, x, y)
            move_list = (
                astar.graphsearch([], astar.f, [], [0, 0], state, fields_for_astar,
                                  astar.succ))
            agent_movement(move_list, agent, fields_for_movement, fields_for_astar)
            draw_window(agent, fields_for_movement)
            fuel = agent.get_tank_capacity()
            water = agent.get_water()
            feritizer = agent.get_feritizer()
            seeds = agent.get_seeds()
            carrots = agent.get_carrots()
            potatoes = agent.get_potatoes()
            wheat = agent.get_wheat()
            print("fuel:", fuel, "water:", water, "feritizer:", feritizer, "seeds:", seeds, "carrots:", carrots,
                  "potatoes:", potatoes, "wheat:", wheat)
            agent.set_tank_capacity(1500)
            agent.set_water(50)
            agent.set_feritizer(50)
            agent.set_seeds(50)
            agent.set_carrots(0)
            agent.set_potatoes(0)
            agent.set_wheat(0)
            print("I'm in base")
            time.sleep(10)


def get_closest_field(fields_to_visit, x, y, dir, fields_for_astar):
    j = 0
    closest_field = fields_to_visit[0]
    state = astar.State(dir, x, y)
    move_list_prev = (astar.graphsearch([], astar.f, [], [closest_field[0], closest_field[1]], state, fields_for_astar,
                                        astar.succ))
    for i in range(1, len(fields_to_visit)):
        curr_field = fields_to_visit[i]
        state = astar.State(dir, x, y)
        move_list = (astar.graphsearch([], astar.f, [], [curr_field[0], curr_field[1]], state, fields_for_astar,
                                       astar.succ))
        if len(move_list) < len(move_list_prev):
            closest_field = curr_field
            move_list_prev = move_list
            j = i
    return closest_field, j


def agent_movement(move_list, agent, fields, fields_2):
    for action in move_list:
        if action == 'l' or action == 'r':
            agent_action(agent.rotate(action))
        elif action == 'f':
            agent_action(agent.move(action))
            cords = agent.coordinates()
            x = cords['x']
            y = cords['y']
            fuel = agent.get_tank_capacity()
            if fields_2[x][y] == 'dirt':
                fuel -= 5
            elif fields_2[x][y] == 'cobble':
                fuel -= 1
            elif fields_2[x][y] == 'grass':
                fuel -= 3
            elif fields_2[x][y] == 'sand':
                fuel -= 10
            agent.set_tank_capacity(fuel)
        draw_window(agent, fields)
        time.sleep(0.5)


common = common.Instance()
agent = agent.Instance(1500, 2, 50, 50, 50, 0, 0, 0)


def main():
    common.set('game_running', True)
    common.set('state_imgShown', False)
    common.set(
        'window',
        pygame.display.set_mode((
            settings.Pygame.width(),
            settings.Pygame.height())))

    pygame.display.set_caption(settings.Pygame.display_name())
    fields, fields_2, fields_3 = randomize_map()
    plants_array = get_plants_array(fields_3)
    # neural_network.learn()
    recognize_plants(plants_array, fields_2, fields,agent)
    x = True
    t = tree.treelearn()
    while common.get('game_running'):
        pygame.time.Clock().tick(settings.Pygame.fps())
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                common.set('game_running', False)
        if common.get('state_imgShown'):
            read_img(agent, fields)
        else:
            draw_window(agent, fields)
        if x:
            generate_movement(fields_2, fields_3, fields, agent, t)
        x = False

        agent_action(agent.rotate(None))
        agent_action(agent.move(None))

    pygame.quit()


if __name__ == "__main__":
    main()