train-journey-visuals/map_display.py

60 lines
1.9 KiB
Python

from typing import Optional
import pygame
from path import Path
from vec import Vec2
class MapDisplay:
PATH_WIDTH = 5
SEGMENT_SIZE = 20
MIN_SEGMENT_LENGTH = 5
def __init__(self,
surf: pygame.Surface,
min_lon: float,
max_lon: float,
min_lat: float,
max_lat: float):
self.surf: pygame.Surface = surf
self.min_lon: float = min_lon
self.max_lon: float = max_lon
self.min_lat: float = min_lat
self.max_lat: float = max_lat
def real_to_screen(self, lon: float, lat: float) -> tuple[float, float]:
x = (lon - self.min_lon) / (self.max_lon - self.min_lon)
y = (lat - self.min_lat) / (self.max_lat - self.min_lat)
w, h = self.surf.get_size()
return x * w, h - y * h
def draw_path(self, path: Path) -> None:
self.draw_colored_path(path, None)
def draw_colored_path(self, path: Path, colors: Optional[list[tuple[int, int, int]]] = None) -> None:
for i, pt in enumerate(path.points):
lon = pt.x
lat = pt.y
x, y = self.real_to_screen(lon, lat)
col = (255, 255, 255) if colors is None else colors[i]
pygame.draw.circle(self.surf, col, (x, y), self.PATH_WIDTH)
def draw_segment(self, path: Path, start_i: int, end_i: int) -> None:
if end_i - start_i < self.MIN_SEGMENT_LENGTH:
return
points = []
for i in range(start_i, end_i):
pt = path.points[i]
pt = Vec2(*self.real_to_screen(pt.x, pt.y))
n = path.normals[i]
n = Vec2(n.x, -n.y)
pt1 = pt + n * self.SEGMENT_SIZE
pt2 = pt - n * self.SEGMENT_SIZE
points.insert(0, (pt1.x, pt1.y))
points.append((pt2.x, pt2.y))
pygame.draw.lines(self.surf, (255, 255, 255), True, points)