train-journey-visuals/map_display.py

from typing import Optional

import pygame

from path import Path


class MapDisplay:
    PATH_WIDTH = 5

    def __init__(self, surf: pygame.Surface, min_lat: float, max_lat: float, min_lon: float, max_lon: float):
        self.surf: pygame.Surface = surf
        self.min_lat: float = min_lat
        self.max_lat: float = max_lat
        self.min_lon: float = min_lon
        self.max_lon: float = max_lon

    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, 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:
        raise NotImplementedError