DKR - Velocity Simulator

from enum import IntEnum

import pyqtgraph as pg

class Character(IntEnum):
    KRUNCH = 0,
    DIDDY = 1,
    BUMPER = 2,
    BANJO = 3,
    CONKER = 4,
    TIPTUP = 5,
    PIPSY = 6,
    TIMBER = 7,
    DRUMSTICK = 8,
    TICKTOCK = 9

class Kart:
    def __init__(self, character: Character):
        self.character = character
        self.velocity = 0.0
        self.throttle = 0.0
        self.brake = 0.0
        self.boost_timer = 30

    def update(self, framerate: int, is_holding_a: bool, is_holding_b: bool, is_using_boost: bool) -> None:
        """Update the kart's properties for a single frame"""

        self.update_throttle(is_holding_a)
        self.update_brake(is_holding_b)
        self.update_boost(framerate, is_using_boost)
        self.update_velocity(is_holding_a)

    def update_throttle(self, is_holding_a: bool) -> None:
        """Update the kart's throttle"""

        if self.throttle > 0.0:
            self.throttle -= 0.1

        if is_holding_a:
            self.throttle = 1.0

    def update_brake(self, is_holding_b: bool) -> None:
        """Update the kart's brakes"""

        if is_holding_b:
            if self.brake < 1.0:
                self.brake += 0.2
        else:
            if self.brake > 0.05:
                self.brake -= 0.1

    def update_boost(self, framerate: int, is_using_boost: bool) -> None:
        """Update the kart's boost timer"""

        self.boost_timer -= framerate

        if self.boost_timer < 0:
            self.boost_timer = 0

        if is_using_boost:
            self.boost_timer = 0
            self.throttle = 1.0

    def update_velocity(self, is_holding_a: bool) -> None:
        """Update the kart's velocity"""

        self.velocity -= self.calculate_drag_from_traction(is_holding_a)

        acceleration = self.calculate_interpolated_stats()

        self.velocity += self.calculate_thrust(acceleration)
        self.velocity -= self.calculate_drag_from_brake(acceleration)

        if self.velocity < 0.0:
            self.velocity = 0.0

    def calculate_drag_from_traction(self, is_holding_a: bool) -> float:
        """Return the drag calculated from traction"""

        if is_holding_a:
            return 1.0 * 0.004 * self.velocity * self.velocity

        return 8.0 * 0.004 * self.velocity

    def calculate_drag_from_brake(self, acceleration: float) -> float:
        """Return the drag calculated from braking"""

        return 0.32 * 1.7 * self.brake * acceleration

    def calculate_thrust(self, acceleration: float) -> float:
        """Return the thrust calculated from throttle"""

        if self.boost_timer > 0:
            return 2.0

        return 1.7 * self.throttle * acceleration

    def calculate_interpolated_stats(self) -> float:
        """Calculate the acceleration based on the character's stats"""

        character_acceleration = [
            [0.10, 0.10, 0.15, 0.12, 0.13, 0.14, 0.15, 0.18, 0.20, 0.26, 0.31, 0.37, 0.39, 0.39], # Krunch
            [0.10, 0.12, 0.15, 0.17, 0.19, 0.21, 0.25, 0.28, 0.29, 0.30, 0.32, 0.34, 0.34, 0.34], # Diddy
            [0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.24, 0.28, 0.29, 0.30, 0.32, 0.34, 0.35, 0.35], # Bumper
            [0.10, 0.12, 0.13, 0.15, 0.17, 0.19, 0.21, 0.23, 0.29, 0.30, 0.32, 0.34, 0.36, 0.36], # Banjo
            [0.10, 0.12, 0.14, 0.16, 0.18, 0.20, 0.24, 0.28, 0.29, 0.30, 0.32, 0.34, 0.35, 0.35], # Conker
            [0.10, 0.15, 0.20, 0.21, 0.24, 0.26, 0.28, 0.29, 0.30, 0.31, 0.32, 0.32, 0.34, 0.34], # Tiptup
            [0.10, 0.20, 0.22, 0.25, 0.26, 0.27, 0.28, 0.29, 0.29, 0.30, 0.30, 0.31, 0.34, 0.34], # Pipsy
            [0.10, 0.12, 0.15, 0.17, 0.19, 0.21, 0.25, 0.28, 0.29, 0.30, 0.32, 0.34, 0.35, 0.35], # Timber
            [0.10, 0.12, 0.14, 0.16, 0.17, 0.19, 0.21, 0.24, 0.26, 0.29, 0.31, 0.33, 0.41, 0.41], # Drumstick
            [0.10, 0.20, 0.22, 0.25, 0.26, 0.27, 0.28, 0.29, 0.30, 0.32, 0.35, 0.38, 0.43, 0.43]  # Ticktock
        ]

        idx = int(self.velocity)
        remainder = self.velocity - idx

        if idx > 12:
            idx = 12

        lower_acceleration = character_acceleration[self.character][idx]
        upper_acceleration = character_acceleration[self.character][idx + 1]

        #TODO: Add bananas + drifting

        return lower_acceleration * (1.0 - remainder) + upper_acceleration * remainder

def main():
    # Constants for the simulation
    FRAMERATE = 2 # 30 FPS (60 / 30 = 2)
    IS_HOLDING_A = True
    IS_HOLDING_B = False
    IS_USING_BOOST = False

    velocity_widget = pg.plot(title="DKR Velocity vs Time")
    velocity_widget.addLegend()
    velocity_widget.setLabel("left", "Velocity")
    velocity_widget.setLabel("bottom", "Vertical Interrupts")

    distance_widget = pg.plot(title="DKR Distance vs Time")
    distance_widget.addLegend()
    distance_widget.setLabel("left", "Distance")
    distance_widget.setLabel("bottom", "Vertical Interrupts")

    for character in list(Character):
        kart = Kart(character)
        distance = 0.0

        frames = []
        velocities = []
        distances = []

        # Go for 5 seconds at this framerate
        for frame in range(0, 5 * 60, FRAMERATE):
            kart.update(FRAMERATE, IS_HOLDING_A, IS_HOLDING_B, IS_USING_BOOST)
            distance += kart.velocity * FRAMERATE
            frames.append(frame)
            velocities.append(kart.velocity)
            distances.append(distance)

        # To differentiate colors, add this
        # symbol=int(character), symbolPen=int(character),
        velocity_widget.plot(frames, velocities, pen=int(character), name=character.name)
        distance_widget.plot(frames, distances, pen=int(character), name=character.name)

    pg.exec()

if __name__ == "__main__":
    main()