107 lines
2.9 KiB
Python
107 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import struct
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Optional
|
|
|
|
import numpy as np
|
|
import pygame
|
|
|
|
from src.vec import Vec
|
|
|
|
if TYPE_CHECKING:
|
|
from src.car import Car
|
|
from src.game import Game
|
|
|
|
|
|
def iter_unpack(format, data):
|
|
nbr_bytes = struct.calcsize(format)
|
|
return struct.unpack(format, data[:nbr_bytes]), data[nbr_bytes:]
|
|
|
|
|
|
@dataclass
|
|
class Snapshot:
|
|
controls: tuple[bool, bool, bool, bool] = (False, False, False, False)
|
|
position: Vec = field(default_factory=Vec)
|
|
direction: Vec = field(default_factory=Vec)
|
|
speed: float = 0
|
|
raycast_distances: list[float] | tuple[float, ...] = field(
|
|
default_factory=list)
|
|
image: Optional[np.ndarray] = None
|
|
|
|
def pack(self):
|
|
data: bytes = b""
|
|
data += struct.pack(">BBBB", *self.controls)
|
|
data += struct.pack(
|
|
">fffff",
|
|
self.position.x,
|
|
self.position.y,
|
|
self.direction.x,
|
|
self.direction.y,
|
|
self.speed,
|
|
)
|
|
|
|
nbr_raycasts: int = len(self.raycast_distances)
|
|
data += struct.pack(f">B{nbr_raycasts}f",
|
|
nbr_raycasts, *self.raycast_distances)
|
|
|
|
if self.image is not None:
|
|
data += struct.pack(">II",
|
|
self.image.shape[0], self.image.shape[1])
|
|
data += self.image.tobytes()
|
|
else:
|
|
data += struct.pack(">II", 0, 0)
|
|
|
|
return data
|
|
|
|
@staticmethod
|
|
def unpack(data: bytes) -> Snapshot:
|
|
controls, data = iter_unpack(">BBBB", data)
|
|
(x, y, dx, dy, s), data = iter_unpack(">fffff", data)
|
|
position = Vec(x, y)
|
|
direction = Vec(dx, dy)
|
|
speed = s
|
|
|
|
(nbr_raycasts,), data = iter_unpack(">B", data)
|
|
raycast_distances, data = iter_unpack(f">{nbr_raycasts}f", data)
|
|
|
|
(h, w), data = iter_unpack(">II", data)
|
|
|
|
if h * w > 0:
|
|
image = np.frombuffer(data, np.uint8).reshape(h, w, 3)
|
|
else:
|
|
image = None
|
|
|
|
return Snapshot(
|
|
controls=controls,
|
|
position=position,
|
|
direction=direction,
|
|
speed=speed,
|
|
raycast_distances=raycast_distances,
|
|
image=image,
|
|
)
|
|
|
|
@staticmethod
|
|
def from_car(car: Car) -> Snapshot:
|
|
return Snapshot(
|
|
controls=(
|
|
car.forward,
|
|
car.backward,
|
|
car.left,
|
|
car.right
|
|
),
|
|
position=car.pos.copy(),
|
|
direction=car.direction.copy(),
|
|
speed=car.speed,
|
|
raycast_distances=car.rays.copy(),
|
|
image=None
|
|
)
|
|
|
|
def apply(self, car: Car):
|
|
car.pos = self.position.copy()
|
|
car.direction = self.direction.copy()
|
|
car.speed = 0
|
|
|
|
def add_image(self, game: Game):
|
|
self.image = pygame.surfarray.array3d(game.game_surf)
|