train-journey-visuals/speed_map_display.py

from typing import Optional

import pygame

from gps_loader import GPSLoader
from map_display import MapDisplay
from path import Path
from units import Unit
from vec import Vec2


class SpeedMapDisplay(MapDisplay):
    APP_NAME = "Train Journey - Speed 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]],
                 min_speed_col: tuple[int, int, int],
                 max_speed_col: tuple[int, int, int],
                 segment_threshold: float):
        super().__init__(surf, min_lon, max_lon, min_lat, max_lat, cities)
        self.min_speed_col: tuple[int, int, int] = min_speed_col
        self.max_speed_col: tuple[int, int, int] = max_speed_col
        self.segment_threshold: float = segment_threshold

        self._path: Optional[Path] = None

    def draw_path(self, path: Path) -> None:
        min_speed = min(path.extra_data)
        max_speed = max(path.extra_data)

        colors = list(map(lambda s: self.interpolate_color(s, min_speed, max_speed), path.extra_data))

        self.draw_colored_path(path, colors)

        in_segment = False
        start_i = 0
        for i, speed in enumerate(path.extra_data):
            if speed >= self.segment_threshold:
                if not in_segment:
                    in_segment = True
                    start_i = i

            elif in_segment:
                in_segment = False
                self.draw_segment(path, start_i, i)

    def interpolate_color(self, speed: float, min_speed: float, max_speed: float) -> tuple[int, int, int]:
        r_span = self.max_speed_col[0] - self.min_speed_col[0]
        g_span = self.max_speed_col[1] - self.min_speed_col[1]
        b_span = self.max_speed_col[2] - self.min_speed_col[2]

        f = (speed - min_speed) / (max_speed - min_speed)
        r = int(r_span * f + self.min_speed_col[0])
        g = int(g_span * f + self.min_speed_col[1])
        b = int(b_span * f + self.min_speed_col[2])

        return r, g, b

    def render(self) -> None:
        self.surf.fill((0, 0, 0))
        self._tooltip_surf.fill((0, 0, 0, 0))
        self.draw_cities()

        if self._path is not None:
            self.draw_path(self._path)
            mpos = Vec2(*pygame.mouse.get_pos())
            self.tooltip_nearest(mpos)
            self.surf.blit(self._tooltip_surf, (0, 0))

    def set_path(self, path: Path) -> None:
        self._path = path

    def tooltip_nearest(self, mpos: Vec2) -> None:
        closest = None
        closest_i = 0
        min_dist = 0

        for i, pt in enumerate(self._path.points):
            pt = Vec2(*self.real_to_screen(pt.x, pt.y))
            dist = (pt - mpos).mag
            if i == 0 or dist < min_dist:
                closest = pt
                closest_i = i
                min_dist = dist

        pt = closest
        speed = self._path.extra_data[closest_i]
        col = (200, 150, 50)
        pygame.draw.circle(self.surf, col, (pt.x, pt.y), 2 * self.PATH_WIDTH, 2)
        pygame.draw.line(self.surf, col, (pt.x, pt.y), (mpos.x, mpos.y), 2)
        txt = self.font.render(f"{speed:.1f} km/h", True, (150, 200, 255))
        txt_size = Vec2(*txt.get_size())
        txt_pos = mpos - txt_size.scale(Vec2(0.5, 1))
        pygame.draw.rect(self._tooltip_surf, (0, 0, 0, 150), (txt_pos.x, txt_pos.y, txt_size.x, txt_size.y))
        self._tooltip_surf.blit(txt, (txt_pos.x, txt_pos.y))


if __name__ == '__main__':
    name = "data_28-03"
    cities = [
        (Vec2(7.359119, 46.227302), "Sion", "above"),
        (Vec2(7.079001, 46.105981), "Martigny", "below"),
        (Vec2(7.001849, 46.216559), "Saint-Maurice", "right")
    ]

    data = GPSLoader.load_data(f"data/{name}.csv")
    speeds = list(map(lambda s: s.convert(Unit.KM_H).value, data["speeds"]))
    path = Path(data["points"], speeds)

    pygame.init()

    win = pygame.display.set_mode([600, 600])
    display = SpeedMapDisplay(
        win,
        6.971094, 7.430611,
        46.076312, 46.253036,
        cities,
        (255, 0, 0), (0, 255, 0), 155)

    display.set_path(path)

    display.mainloop()