train-journey-visuals/speed_map_display.py

import pygame

from map_display import MapDisplay
from path import Path
from vec import Vec2


class SpeedMapDisplay(MapDisplay):
    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

    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