DKR - Optimal Tapping Sim

from copy import deepcopy

import pyqtgraph as pg

INITIAL_VELOCITY = 0.0
INITIAL_BOOST = 0
NUM_UPDATES = 120
VI_PER_VISUAL_FRAME = 2

class Kart:
    def __init__(self):
        self.velocity = INITIAL_VELOCITY
        self.throttle = 0.0
        self.distance = 0.0
        self.boost = INITIAL_BOOST
        self.inputs = []

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

        self.inputs.append(is_holding_a)

        self.update_throttle(is_holding_a)
        self.update_boost()
        self.update_velocity(is_holding_a)
        self.distance += self.velocity

    def update_boost(self) -> None:
        """Update the kart's boost"""

        self.boost -= VI_PER_VISUAL_FRAME

        if self.boost < 0:
            self.boost = 0

    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 or self.boost > 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)
        self.velocity += self.calculate_thrust(self.calculate_interpolated_stats())

        if self.velocity < 0.04:
            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_thrust(self, acceleration: float) -> float:
        """Return the thrust calculated from throttle"""

        if self.boost > 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"""

        stats = [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]

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

        if idx > 12:
            idx = 12

        lower_acceleration = stats[idx]
        upper_acceleration = stats[idx + 1]

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

def get_best_karts(karts: list[Kart]) -> list[Kart]:
    """Filter the karts down based on velocity and distance"""

    best_karts = []

    for throttle in set(kart.throttle for kart in karts):
        karts_at_throttle = [kart for kart in karts if kart.throttle == throttle]

        # If any other kart has both higher velocity AND a higher travelled distance at a given throttle, than the kart
        # is a poor performer and should not be included in the list.
        best_karts += [kart for kart in karts_at_throttle
                       if not any(k.velocity > kart.velocity and k.distance > kart.distance for k in karts_at_throttle)]

    return best_karts

def get_optimal_kart(frames: int) -> Kart:
    """Get the kart that would travel the furthest distance after the supplied number of visual frames"""

    karts = [Kart()]

    for _ in range(frames):
        copies_of_karts = [deepcopy(kart) for kart in karts]

        for kart in karts:
            kart.update(False)

        for kart in copies_of_karts:
            kart.update(True)

        karts += copies_of_karts

        karts = get_best_karts(karts)

    return max(karts, key=lambda k: k.distance)

def simulate_holding_a_kart(frames: int) -> tuple[list[float], list[float]]:
    """Return the simulated kart's distances and velocities of a player just holding A"""

    kart = Kart()

    distances = []
    velocities = []

    for _ in range(frames):
        distances.append(kart.distance)
        velocities.append(kart.velocity)

        kart.update(True)

    return distances, velocities

def simulate_best_kart(frames: int) -> tuple[list[float], list[float], list[bool]]:
    """Return the simulated kart's distances and velocities of a player performing optimally"""

    optimal_kart = get_optimal_kart(frames)
    #print(optimal_kart.distance)
    #print(*[f"{i}: {input}\n" for i, input in enumerate(optimal_kart.inputs, start=1)], end="")

    kart = Kart()

    distances = []
    velocities = []
    inputs = [1.0 if input_ else 0.0 for input_ in optimal_kart.inputs]

    for i in range(frames):
        distances.append(kart.distance)
        velocities.append(kart.velocity)

        kart.update(optimal_kart.inputs[i])

    return distances, velocities, inputs

def main():
    velocity_widget = pg.plot(title="DKR Velocity vs Time")
    velocity_widget.addLegend()
    velocity_widget.setLabel("left", "Velocity")
    velocity_widget.setLabel("bottom", "Visual Frames")

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

    inputs_widget = pg.plot(title="DKR Inputs vs Time")
    inputs_widget.addLegend()
    inputs_widget.setLabel("left", "Inputs")
    inputs_widget.setLabel("bottom", "Visual Frames")

    frames = list(range(NUM_UPDATES))

    distances, velocities = simulate_holding_a_kart(NUM_UPDATES)
    velocity_widget.plot(frames, velocities)
    distance_widget.plot(frames, distances)

    distances, velocities, inputs = simulate_best_kart(NUM_UPDATES)
    velocity_widget.plot(frames, velocities)
    distance_widget.plot(frames, distances)
    inputs_widget.plot(frames, inputs)

    pg.exec()

if __name__ == "__main__":
    main()