Untitled

mail@pastecode.io avatar
unknown
python
2 years ago
4.1 kB
3
Indexable
Never
import time
import numpy as np
from pyglet.math import Vec2

class FlowField(object):
    # Const neighbours vectors
    DIR_VECTORS = [[1, 0], [-1, 0], [0, 1], [0, -1]]
    FULL_DIR_VECTORS = [[1, 0], [-1, 0], [0, 1], [0, -1],
    [1, 1], [-1, -1], [1, -1], [-1, 1]]

    def __init__(self, cost_map: np.array, map_size: Vec2):
        self.cost_map = cost_map
        self.size = map_size
        self.node_map = self.generate_node_map()

    def generate_node_map(self):
        nodes = np.empty((self.size.x, self.size.y), dtype=object)

        for x in range(self.size.x):
            for y in range(self.size.y):
                n = Node(x, y, self.cost_map[x, y])
                nodes[x, y] = n
        return nodes

    def reset_nodes(self):
        for x in range(self.size.x):
            for y in range(self.size.y):
                self.node_map[x,y].best_cost = 65536
                self.node_map[x,y].cost = self.cost_map[x, y]

    def calculate_integration_field(self, target_x: int, target_y: int):
        self.start = time.time()

        # Create a new, empty list
        open_node_list = []

        self.reset_nodes()

        # Set the target node cost and best cost to 0
        self.node_map[target_x, target_y].cost = 0
        self.node_map[target_x, target_y].best_cost = 0
        self.node_map[target_x, target_y].best_direction = Vec2(0, 0)

        # Add the node goal to the open list
        open_node_list.append(self.node_map[target_x, target_y])

        counter = 0

        while len(open_node_list) > 0:
            # Get the next node and removes it from the list
            current_node = open_node_list.pop(0)

            # Get the neighbours cells
            neighbours = self.get_node_neighbours(current_node.position, self.DIR_VECTORS)
            counter = counter + 1

            for node in neighbours:
                if node.cost == 255:
                    continue
                if node.cost + current_node.best_cost < node.best_cost:
                    if node not in open_node_list:
                        open_node_list.append(node)
                    node.best_cost = current_node.best_cost + node.cost
                    self.node_map[node.position.x, node.position.y].best_cost = node.best_cost

        print(counter)
        self.create_vector_field()

    def create_vector_field(self):
        counter = 0
        for x in range(self.size.x):
            for y in range(self.size.y):
                # Get the neighbours cells
                node = self.node_map[x,y]
                neighbours = self.get_node_neighbours(node.position, self.FULL_DIR_VECTORS)
                best_cost = node.best_cost

                counter = counter + 1

                for neighbour in neighbours:
                    if neighbour.best_cost < best_cost:
                        best_cost = neighbour.best_cost
                        node.best_direction = Vec2(neighbour.position.x - node.position.x, neighbour.position.y - node.position.y)

        end = time.time()
        print("Vector Field Generated in: " + str(end - self.start))
        print(counter)

    def get_node_neighbours(self, current_node_pos: Vec2, directions: list):
        neighbours = []

        for i in range(len(directions)):
            n_pos = [current_node_pos.x + directions[i][0], current_node_pos.y + directions[i][1]]
            if 0 <= n_pos[0] < self.size.x:
                if 0 <= n_pos[1] < self.size.y:
                    neighbours.append(self.node_map[n_pos[0], n_pos[1]])

        return neighbours

class Node(object):

    def __init__(self, pos_x: int, pos_y: int, cost: int):
        self.position = Vec2(pos_x, pos_y)
        self.cost = cost
        self.best_cost = 65536
        self.best_direction = Vec2(0, 0)

    def increase_cost(self, amount: int):
        # If the cost of the cell is already at the maximum value, return.
        if self.cost == 255:
            return
        # Clamp the value
        if amount + self.cost > 255:
            self.cost = 255
        else:
            self.cost += amount