train-journey-visuals/map_display.py

from typing import Optional

import pygame

from display import Display
from path import Path
from vec import Vec2


class MapDisplay(Display):
    PATH_WIDTH = 5
    SEGMENT_SIZE = 20
    MIN_SEGMENT_LENGTH = 5
    APP_NAME = "Train Journey - Map Display"

    def __init__(self,
                 surf: pygame.Surface,
                 min_lon: float,
                 max_lon: float,
                 min_lat: float,
                 max_lat: float,
                 cities: list[tuple[Vec2, str, str]]):

        super().__init__(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
        self._precomputeDisplayValues()
        self.cities: list[tuple[Vec2, str, str]] = cities

    def _precomputeDisplayValues(self) -> None:
        width, height = self.surf.get_size()
        self._lon_span = self.max_lon - self.min_lon
        self._lat_span = self.max_lat - self.min_lat
        r1 = width / self._lon_span
        r2 = height / self._lat_span

        self._r = min(r1, r2)

        self._ox = (width - self._lon_span * self._r) / 2
        self._oy = (height - self._lat_span * self._r) / 2

    def real_to_screen(self, lon: float, lat: float) -> tuple[float, float]:
        x = (lon - self.min_lon) * self._r + self._ox
        y = (lat - self.min_lat) * self._r + self._oy

        return x, self.surf.get_height() - y

    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)

    def draw_cities(self) -> None:
        for city in self.cities:
            self.draw_city(*city)

    def draw_city(self, pos: Vec2, name: str, label_side: str) -> None:
        pos2 = Vec2(*self.real_to_screen(pos.x, pos.y))

        pygame.draw.circle(self.surf, (180, 180, 180), (pos2.x, pos2.y), 10)

        label = self.font.render(name, True, (255, 255, 255))
        label_size = Vec2(*label.get_size())

        line_end = pos2
        label_pos = line_end - label_size / 2

        if label_side == "above":
            line_end -= Vec2(0, 20)
            label_pos = line_end - label_size.scale(Vec2(0.5, 1))

        elif label_side == "below":
            line_end += Vec2(0, 20)
            label_pos = line_end - label_size.scale(Vec2(0.5, 0))

        elif label_side == "left":
            line_end -= Vec2(20, 0)
            label_pos = line_end - label_size.scale(Vec2(1, 0.5))

        elif label_side == "right":
            line_end += Vec2(20, 0)
            label_pos = line_end - label_size.scale(Vec2(0, 0.5))

        pygame.draw.line(self.surf, (255, 255, 255), (pos2.x, pos2.y), (line_end.x, line_end.y))
        self.surf.blit(label, (label_pos.x, label_pos.y))