Lesson 25: Tower Defense Capstone

Work in order and run the game after every to-do; each one makes something new visible. For to-do 2, copy the idea of the while loop from the road section: step is a budget of pixels you spend leg by leg. For to-do 3, min() takes a generator and key=lambda e: e.remaining(). If gold explodes after a wave, to-do 4 is paying without changing self.phase. If nothing slows down, check that hit() calls apply_slow() before take_hit() can mark the enemy dead.

If your instructor hands you the lab file, you will see a few extra lines marked lab runtime near the top, plus an extra and frame_budget() condition on the main loop. They let the instructor's checker run the program automatically; when you run it yourself they do nothing. The game expects Kenney's sprites in an assets folder next to it, but it still runs (with plain circles) without them.

"""Tower Defense Capstone: Intermediate Lesson 25 (solution).

Build towers beside the road and stop six waves of enemies before they
reach the exit. Art: Kenney "Tower Defense (top-down)" pack, CC0.

Controls: 1-4 choose a tower, left-click to build or select, U upgrade,
S sell, Space start the next wave, P pause, F fast-forward, K screen
shake on/off, Esc deselect, R restart after the game ends.
"""
import math
import random
from dataclasses import dataclass
from enum import Enum
from pathlib import Path

import pygame


WIDTH, HEIGHT = 960, 640
TILE = 64
COLS, ROWS = 15, 9                  # the play area is 15 x 9 tiles (960 x 576)
HUD_Y = ROWS * TILE                 # the HUD bar fills the bottom 64 pixels
ASSETS = Path(__file__).parent / "assets"

START_GOLD = 120
START_LIVES = 20
SELL_REFUND = 0.7                   # selling returns 70% of what you spent
MAX_LEVEL = 3
HP_GROWTH = 0.15                    # enemies get 15% more health each wave


def clamp_color(r, g, b):
    """Keep every channel a whole number from 0 to 255."""
    return (max(0, min(255, int(r))), max(0, min(255, int(g))), max(0, min(255, int(b))))


# ---- The road: a list of grid cells, turned into pixel waypoints ----------------
ROAD_CELLS = [(-1, 2), (3, 2), (3, 6), (7, 6), (7, 2), (11, 2), (11, 7), (15, 7)]


def cell_center(col, row):
    return pygame.Vector2(col * TILE + TILE / 2, row * TILE + TILE / 2)


def road_tiles(cells):
    """Every on-screen grid cell the road passes through (no towers allowed there)."""
    tiles = set()
    for (c1, r1), (c2, r2) in zip(cells, cells[1:]):
        for c in range(min(c1, c2), max(c1, c2) + 1):
            for r in range(min(r1, r2), max(r1, r2) + 1):
                if 0 <= c < COLS and 0 <= r < ROWS:
                    tiles.add((c, r))
    return tiles


GROUND_ROUTE = [cell_center(c, r) for c, r in ROAD_CELLS]
AIR_ROUTE = [GROUND_ROUTE[0], GROUND_ROUTE[-1]]      # flyers cut straight across
ROAD = road_tiles(ROAD_CELLS)


# ---- Game data: what exists, as Enums and dataclasses ---------------------------
class TowerKind(Enum):
    ARROW = "arrow"
    CANNON = "cannon"
    FROST = "frost"
    LASER = "laser"


class EnemyKind(Enum):
    GRUNT = "grunt"
    RUNNER = "runner"
    TANK = "tank"
    FLYER = "flyer"


class Phase(Enum):
    BUILD = "build"                 # between waves: build and upgrade
    WAVE = "wave"                   # enemies are spawning or on the road
    WON = "won"
    LOST = "lost"


@dataclass
class TowerSpec:
    name: str
    cost: int
    damage: float
    range: float                    # pixels
    fire_rate: float                # shots per second
    shot_speed: float               # pixels per second; 0 means an instant hit
    hits_air: bool
    color: tuple
    sprite: str
    splash: float = 0.0             # blast radius in pixels (cannon only)
    slow_factor: float = 1.0        # 0.5 = half speed
    slow_time: float = 0.0          # seconds


@dataclass
class EnemySpec:
    name: str
    hp: float
    speed: float                    # pixels per second
    armor: float                    # subtracted from every hit (minimum 1 damage)
    reward: int
    flying: bool
    radius: float
    sprite: str


TOWER_SPECS = {
    TowerKind.ARROW: TowerSpec("Arrow", 50, 10, 150, 2.0, 520, True, (120, 220, 120),
                               "towerDefense_tile249.png"),
    TowerKind.CANNON: TowerSpec("Cannon", 90, 28, 130, 0.6, 320, False, (240, 110, 90),
                                "towerDefense_tile250.png", splash=60),
    TowerKind.FROST: TowerSpec("Frost", 70, 3, 120, 1.0, 420, False, (130, 190, 255),
                               "towerDefense_tile226.png", slow_factor=0.5, slow_time=1.5),
    TowerKind.LASER: TowerSpec("Laser", 120, 6, 170, 5.0, 0, True, (255, 230, 90),
                               "towerDefense_tile203.png"),
}

ENEMY_SPECS = {
    EnemyKind.GRUNT: EnemySpec("Grunt", 40, 60, 0, 5, False, 14, "towerDefense_tile245.png"),
    EnemyKind.RUNNER: EnemySpec("Runner", 25, 110, 0, 6, False, 12, "towerDefense_tile247.png"),
    EnemyKind.TANK: EnemySpec("Tank", 160, 38, 4, 15, False, 22, "towerDefense_tile268.png"),
    EnemyKind.FLYER: EnemySpec("Flyer", 50, 75, 0, 10, True, 18, "towerDefense_tile270.png"),
}

# Each wave is a list of groups: (enemy kind, how many, seconds between spawns).
WAVES = [
    [(EnemyKind.GRUNT, 8, 0.9)],
    [(EnemyKind.GRUNT, 10, 0.7), (EnemyKind.RUNNER, 6, 0.5)],
    [(EnemyKind.RUNNER, 12, 0.45), (EnemyKind.FLYER, 3, 1.0)],
    [(EnemyKind.TANK, 4, 1.6), (EnemyKind.GRUNT, 10, 0.6)],
    [(EnemyKind.FLYER, 8, 0.9), (EnemyKind.GRUNT, 8, 0.6)],
    [(EnemyKind.TANK, 6, 1.3), (EnemyKind.RUNNER, 12, 0.4), (EnemyKind.FLYER, 6, 0.8)],
]

KEY_TO_KIND = {pygame.K_1: TowerKind.ARROW, pygame.K_2: TowerKind.CANNON,
               pygame.K_3: TowerKind.FROST, pygame.K_4: TowerKind.LASER}
BUTTONS = {kind: pygame.Rect(8 + i * 132, HUD_Y + 8, 124, 48) for i, kind in enumerate(TowerKind)}


def wave_bonus(number):
    return 20 + 5 * number


# ---- Enemies ------------------------------------------------------------------
class Enemy:
    def __init__(self, kind, wave_number):
        self.kind = kind
        self.spec = ENEMY_SPECS[kind]
        self.route = AIR_ROUTE if self.spec.flying else GROUND_ROUTE
        self.route_length = sum(a.distance_to(b) for a, b in zip(self.route, self.route[1:]))
        self.pos = pygame.Vector2(self.route[0])
        self.leg = 0                        # index of the waypoint we are walking away from
        self.traveled = 0.0                 # pixels walked along the route
        self.heading = pygame.Vector2(1, 0)
        self.max_hp = self.spec.hp * (1 + HP_GROWTH * (wave_number - 1))
        self.hp = self.max_hp
        self.slow_timer = 0.0
        self.slow_factor = 1.0
        self.flash = 0.0                    # seconds of white hit-flash left
        self.alive = True

    def remaining(self):
        """Pixels left to the exit: the smaller, the more dangerous."""
        return self.route_length - self.traveled

    def current_speed(self):
        return self.spec.speed * (self.slow_factor if self.slow_timer > 0 else 1.0)

    def velocity(self):
        return self.heading * self.current_speed()

    def apply_slow(self, factor, seconds):
        """The strongest slow wins; hitting again refreshes the timer (never stacks)."""
        self.slow_factor = min(self.slow_factor, factor) if self.slow_timer > 0 else factor
        self.slow_timer = max(self.slow_timer, seconds)

    def take_hit(self, damage):
        """Apply one hit. Returns True only on the hit that kills."""
        if not self.alive:
            return False
        self.hp -= max(1.0, damage - self.spec.armor)
        self.flash = 0.08
        if self.hp <= 0:
            self.alive = False
            return True
        return False

    def update(self, dt):
        """Walk along the route. Returns True on the frame the enemy reaches the exit."""
        if not self.alive:
            return False
        self.slow_timer = max(0.0, self.slow_timer - dt)
        self.flash = max(0.0, self.flash - dt)
        step = self.current_speed() * dt
        while step > 0 and self.leg < len(self.route) - 1:
            target = self.route[self.leg + 1]
            to_target = target - self.pos
            dist = to_target.length()
            if dist > 0:
                self.heading = to_target / dist
            if step < dist:
                self.pos += self.heading * step
                self.traveled += step
                step = 0
            else:                           # reach the corner, keep the leftover step
                self.pos.update(target)
                self.traveled += dist
                step -= dist
                self.leg += 1
        if self.leg >= len(self.route) - 1:
            self.alive = False
            return True
        return False


# ---- Towers and projectiles -----------------------------------------------------
class Tower:
    def __init__(self, kind, cell):
        self.kind = kind
        self.spec = TOWER_SPECS[kind]
        self.cell = cell
        self.pos = cell_center(*cell)
        self.level = 1
        self.spent = self.spec.cost
        self.cooldown = 0.0                 # seconds until the next shot
        self.angle = -90.0                  # degrees; -90 points up the screen

    def damage(self):
        return self.spec.damage * (1 + 0.5 * (self.level - 1))

    def range(self):
        return self.spec.range * (1 + 0.1 * (self.level - 1))

    def upgrade_cost(self):
        return None if self.level >= MAX_LEVEL else self.spec.cost * self.level

    def can_hit(self, enemy):
        if not enemy.alive or (enemy.spec.flying and not self.spec.hits_air):
            return False
        return self.pos.distance_to(enemy.pos) <= self.range()

    def pick_target(self, enemies):
        """Target "first": the enemy in range with the least road left."""
        return min((e for e in enemies if self.can_hit(e)), key=lambda e: e.remaining(), default=None)

    def update(self, dt, enemies):
        """Count the cooldown down every frame; return a target when it is time to fire."""
        self.cooldown = max(0.0, self.cooldown - dt)
        target = self.pick_target(enemies)
        if target is None:
            return None
        aim = target.pos - self.pos
        if aim.length_squared() > 0:
            self.angle = math.degrees(math.atan2(aim.y, aim.x))
        if self.cooldown > 0:
            return None
        self.cooldown = 1 / self.spec.fire_rate
        return target


class Projectile:
    """Bullets chase a live target; cannon shells fly to a fixed point and explode."""

    def __init__(self, tower, target):
        self.spec = tower.spec
        self.damage = tower.damage()
        self.pos = pygame.Vector2(tower.pos)
        self.target = target
        self.alive = True
        self.aim = None
        if self.spec.splash > 0:            # lead the target: aim where it will be
            flight_time = self.pos.distance_to(target.pos) / self.spec.shot_speed
            self.aim = target.pos + target.velocity() * flight_time

    def update(self, dt, game):
        if self.aim is None and not self.target.alive:
            self.alive = False              # target died or escaped: the bullet fizzles
            return
        goal = self.aim if self.aim is not None else self.target.pos
        to_goal = goal - self.pos
        step = self.spec.shot_speed * dt
        if to_goal.length() <= step:
            self.pos.update(goal)
            self.alive = False
            game.projectile_hit(self)
        else:
            self.pos += to_goal.normalize() * step


# ---- Juice ----------------------------------------------------------------------
class Particle:
    def __init__(self, pos, vel, color, life):
        self.pos = pygame.Vector2(pos)
        self.vel = pygame.Vector2(vel)
        self.color = color
        self.life = self.max_life = life

    def update(self, dt):
        self.life -= dt
        self.pos += self.vel * dt
        self.vel *= math.exp(-4 * dt)       # frame-rate independent drag


# ---- The game: all rules, no drawing --------------------------------------------
class Game:
    def __init__(self, seed=None):
        self.rng = random.Random(seed)
        self.gold = START_GOLD
        self.lives = START_LIVES
        self.phase = Phase.BUILD
        self.wave_number = 0
        self.spawn_queue = []
        self.spawn_timer = 0.0
        self.towers = {}                    # (col, row) -> Tower
        self.enemies = []
        self.projectiles = []
        self.particles = []
        self.beams = []                     # [start, end, color, seconds left]
        self.popups = []                    # [text, position, seconds left]
        self.selected_kind = TowerKind.ARROW
        self.selected_tower = None
        self.paused = False
        self.speed = 1
        self.shake_on = True
        self.trauma = 0.0
        self.shake = pygame.Vector2()
        self.kills = 0
        self.message = "Build towers, then press Space to start wave 1"

    # -- building --
    def can_build(self, cell):
        col, row = cell
        return 0 <= col < COLS and 0 <= row < ROWS and cell not in ROAD and cell not in self.towers

    def build(self, kind, cell):
        spec = TOWER_SPECS[kind]
        if not self.can_build(cell):
            self.message = "You can't build there"
            return None
        if self.gold < spec.cost:
            self.message = f"{spec.name} costs {spec.cost} gold"
            return None
        self.gold -= spec.cost
        tower = Tower(kind, cell)
        self.towers[cell] = tower
        self.burst(tower.pos, (255, 215, 90), 10)
        return tower

    def upgrade(self, tower):
        cost = tower.upgrade_cost()
        if cost is None or self.gold < cost:
            return False
        self.gold -= cost
        tower.spent += cost
        tower.level += 1
        self.popup(f"Level {tower.level}", tower.pos)
        return True

    def sell(self, tower):
        refund = int(tower.spent * SELL_REFUND)
        self.gold += refund
        del self.towers[tower.cell]
        if self.selected_tower is tower:
            self.selected_tower = None
        self.popup(f"+{refund}", tower.pos)
        return refund

    def click(self, pos):
        x, y = pos
        if y >= HUD_Y:
            for kind, rect in BUTTONS.items():
                if rect.collidepoint(pos):
                    self.selected_kind = kind
                    self.selected_tower = None
            return
        cell = (int(x // TILE), int(y // TILE))
        if cell in self.towers:
            self.selected_tower = self.towers[cell]
        else:
            self.selected_tower = None
            self.build(self.selected_kind, cell)

    # -- waves --
    def start_wave(self):
        if self.phase != Phase.BUILD or self.wave_number >= len(WAVES):
            return False
        self.wave_number += 1
        self.spawn_queue = [(kind, gap) for kind, count, gap in WAVES[self.wave_number - 1]
                            for _ in range(count)]
        self.spawn_timer = 0.0
        self.phase = Phase.WAVE
        self.message = f"Wave {self.wave_number} incoming!"
        return True

    def finish_wave(self):
        """Runs once per wave: the phase change stops it from paying again."""
        bonus = wave_bonus(self.wave_number)
        self.gold += bonus
        if self.wave_number >= len(WAVES):
            self.phase = Phase.WON
            self.message = "Every wave defeated! Press R to play again"
        else:
            self.phase = Phase.BUILD
            self.message = f"Wave cleared: +{bonus} gold. Space starts the next one"

    # -- combat --
    def hit(self, enemy, damage, spec):
        if not enemy.alive:
            return
        if spec.slow_time > 0:
            enemy.apply_slow(spec.slow_factor, spec.slow_time)
        if enemy.take_hit(damage):
            self.gold += enemy.spec.reward
            self.kills += 1
            self.popup(f"+{enemy.spec.reward}", enemy.pos)
            self.burst(enemy.pos, (230, 90, 70), 12)

    def projectile_hit(self, projectile):
        if projectile.spec.splash > 0:      # explode where the shell lands
            radius = projectile.spec.splash
            self.burst(projectile.pos, (255, 170, 60), 20)
            self.add_trauma(0.2)
            for enemy in self.enemies:
                if enemy.alive and not enemy.spec.flying:
                    d = enemy.pos.distance_to(projectile.pos)
                    if d <= radius:
                        self.hit(enemy, projectile.damage * (1 - 0.5 * d / radius), projectile.spec)
        else:
            self.hit(projectile.target, projectile.damage, projectile.spec)

    # -- effects --
    def burst(self, pos, color, count):
        for _ in range(count):
            angle = self.rng.uniform(0, math.tau)
            speed = self.rng.uniform(40, 160)
            vel = (math.cos(angle) * speed, math.sin(angle) * speed)
            self.particles.append(Particle(pos, vel, color, self.rng.uniform(0.3, 0.6)))

    def popup(self, text, pos):
        self.popups.append([text, pygame.Vector2(pos), 0.8])

    def add_trauma(self, amount):
        self.trauma = min(1.0, self.trauma + amount)

    def update_effects(self, dt):
        for p in self.particles:
            p.update(dt)
        self.particles = [p for p in self.particles if p.life > 0]
        for beam in self.beams:
            beam[3] -= dt
        self.beams = [b for b in self.beams if b[3] > 0]
        for popup in self.popups:
            popup[1].y -= 30 * dt
            popup[2] -= dt
        self.popups = [p for p in self.popups if p[2] > 0]
        self.trauma = max(0.0, self.trauma - 1.5 * dt)
        amount = self.trauma ** 2 * 10 if self.shake_on else 0
        self.shake.update(self.rng.uniform(-1, 1) * amount, self.rng.uniform(-1, 1) * amount)

    # -- the frame --
    def update(self, dt):
        """dt is real seconds. Pause and fast-forward happen here, so every timer obeys them."""
        dt = 0.0 if self.paused else dt * self.speed
        self.update_effects(dt)
        if self.phase != Phase.WAVE:
            return

        # 1. Spawn on a timer that counts seconds of game time.
        self.spawn_timer -= dt
        while self.spawn_queue and self.spawn_timer <= 0:
            kind, gap = self.spawn_queue.pop(0)
            self.enemies.append(Enemy(kind, self.wave_number))
            self.spawn_timer += gap

        # 2. Move enemies; the ones that reach the exit cost a life.
        for enemy in self.enemies:
            if enemy.update(dt):
                self.lives -= 1
                self.add_trauma(0.5)

        # 3. Towers aim and fire.
        for tower in self.towers.values():
            target = tower.update(dt, self.enemies)
            if target is None:
                continue
            if tower.spec.shot_speed == 0:  # laser: instant hit, still limited by fire_rate
                self.beams.append([tower.pos.copy(), target.pos.copy(), tower.spec.color, 0.08])
                self.hit(target, tower.damage(), tower.spec)
            else:
                self.projectiles.append(Projectile(tower, target))

        # 4. Move shots, then drop everything that is finished.
        for projectile in self.projectiles:
            projectile.update(dt, self)
        self.projectiles = [p for p in self.projectiles if p.alive]
        self.enemies = [e for e in self.enemies if e.alive]

        # 5. Win or lose the wave.
        if self.lives <= 0:
            self.lives = 0
            self.phase = Phase.LOST
            self.message = "The base has fallen. Press R to try again"
        elif not self.spawn_queue and not self.enemies:
            self.finish_wave()


# ---- Drawing --------------------------------------------------------------------
def load_sprite(filename, color):
    """Load a Kenney sprite, or draw a stand-in if the file is missing."""
    path = ASSETS / filename
    if path.exists():
        return pygame.image.load(path).convert_alpha()
    image = pygame.Surface((TILE, TILE), pygame.SRCALPHA)
    pygame.draw.circle(image, color, (TILE // 2, TILE // 2), TILE // 4)
    return image


def load_art():
    art = {"grass": load_sprite("towerDefense_tile024.png", (60, 150, 70)),
           "road": load_sprite("towerDefense_tile093.png", (170, 120, 70)),
           "base": load_sprite("towerDefense_tile181.png", (150, 160, 170)),
           "bullet": load_sprite("towerDefense_tile272.png", (255, 220, 80)),
           "shell": load_sprite("towerDefense_tile275.png", (200, 200, 200))}
    for kind, spec in TOWER_SPECS.items():
        art[kind] = load_sprite(spec.sprite, spec.color)
    for kind, spec in ENEMY_SPECS.items():
        art[kind] = load_sprite(spec.sprite, (220, 80, 80))
    return art


def build_background(art):
    """Draw the grass and the road once; each frame just blits this Surface."""
    background = pygame.Surface((WIDTH, HUD_Y))
    for col in range(COLS):
        for row in range(ROWS):
            tile = art["road"] if (col, row) in ROAD else art["grass"]
            background.blit(tile, (col * TILE, row * TILE))
    return background


_rotation_cache = {}


def rotated(image, degrees):
    """Rotate in 5-degree steps and remember the result (at most 72 per image)."""
    step = int(round(degrees / 5) * 5) % 360
    key = (id(image), step)
    if key not in _rotation_cache:
        _rotation_cache[key] = pygame.transform.rotate(image, -step)
    return _rotation_cache[key]


_circle_cache = {}


def range_circle(radius, color):
    key = (int(radius), color)
    if key not in _circle_cache:
        size = int(radius) * 2 + 2
        circle = pygame.Surface((size, size), pygame.SRCALPHA)
        pygame.draw.circle(circle, (*color, 40), (size // 2, size // 2), int(radius))
        pygame.draw.circle(circle, (*color, 160), (size // 2, size // 2), int(radius), 2)
        _circle_cache[key] = circle
    return _circle_cache[key]


def draw(screen, play, game, art, fonts, background):
    """Draw the world on `play`, then blit it shaken onto the screen, then the HUD."""
    play.blit(background, (0, 0))

    # Range preview under the mouse, or around the selected tower.
    mx, my = pygame.mouse.get_pos()
    hover = (mx // TILE, my // TILE)
    if game.selected_tower is not None:
        t = game.selected_tower
        circle = range_circle(t.range(), (255, 255, 255))
        play.blit(circle, circle.get_rect(center=t.pos))
    elif my < HUD_Y and game.can_build(hover):
        spec = TOWER_SPECS[game.selected_kind]
        circle = range_circle(spec.range, spec.color)
        play.blit(circle, circle.get_rect(center=cell_center(*hover)))

    for tower in game.towers.values():
        play.blit(art["base"], art["base"].get_rect(center=tower.pos))
        turret = rotated(art[tower.kind], tower.angle + 90)   # turret art points up
        play.blit(turret, turret.get_rect(center=tower.pos))
        for i in range(tower.level):
            pygame.draw.circle(play, (255, 215, 90), (tower.pos.x - 12 + i * 12, tower.pos.y + 24), 4)

    for enemy in sorted(game.enemies, key=lambda e: e.spec.flying):  # flyers on top
        image = rotated(art[enemy.kind], math.degrees(math.atan2(enemy.heading.y, enemy.heading.x)))
        if enemy.flash > 0:
            image = image.copy()
            image.fill((120, 120, 120), special_flags=pygame.BLEND_RGB_ADD)
        if enemy.slow_timer > 0:
            pygame.draw.circle(play, (130, 190, 255), enemy.pos, enemy.spec.radius + 4, 2)
        play.blit(image, image.get_rect(center=enemy.pos))
        if enemy.hp < enemy.max_hp:
            bar = pygame.Rect(0, 0, 32, 5)
            bar.midbottom = (enemy.pos.x, enemy.pos.y - enemy.spec.radius - 4)
            pygame.draw.rect(play, (60, 20, 20), bar)
            bar.width = max(0, int(32 * enemy.hp / enemy.max_hp))
            pygame.draw.rect(play, (90, 220, 90), bar)

    for p in game.projectiles:
        image = art["shell"] if p.spec.splash > 0 else art["bullet"]
        play.blit(image, image.get_rect(center=p.pos))
        if p.spec.slow_time > 0:
            pygame.draw.circle(play, (130, 190, 255), p.pos, 5)
    for start, end, color, _ in game.beams:
        pygame.draw.line(play, color, start, end, 3)
    for p in game.particles:
        t = max(0.0, p.life / p.max_life)
        r, g, b = p.color
        pygame.draw.circle(play, clamp_color(r * t + 40, g * t + 40, b * t + 40), p.pos, 2 + 3 * t)
    for text, pos, _ in game.popups:
        label = fonts["small"].render(text, True, (255, 230, 120))
        play.blit(label, label.get_rect(center=pos))

    screen.fill((0, 0, 0))
    screen.blit(play, game.shake)
    draw_hud(screen, game, fonts)


def draw_hud(screen, game, fonts):
    pygame.draw.rect(screen, (28, 32, 44), (0, HUD_Y, WIDTH, HEIGHT - HUD_Y))
    for i, (kind, rect) in enumerate(BUTTONS.items()):
        spec = TOWER_SPECS[kind]
        chosen = kind == game.selected_kind and game.selected_tower is None
        pygame.draw.rect(screen, (60, 70, 96) if chosen else (40, 46, 62), rect, border_radius=6)
        pygame.draw.rect(screen, spec.color, rect, 2, border_radius=6)
        affordable = game.gold >= spec.cost
        color = (235, 235, 235) if affordable else (130, 130, 130)
        screen.blit(fonts["small"].render(f"{i + 1} {spec.name}", True, color), (rect.x + 8, rect.y + 6))
        screen.blit(fonts["small"].render(f"{spec.cost} gold", True, color), (rect.x + 8, rect.y + 26))

    stats = f"Gold {game.gold}   Lives {game.lives}   Wave {game.wave_number}/{len(WAVES)}"
    if game.speed > 1:
        stats += "   x2"
    if game.paused:
        stats += "   PAUSED"
    screen.blit(fonts["big"].render(stats, True, (240, 240, 240)), (540, HUD_Y + 8))
    line = game.message
    if game.selected_tower is not None:
        t = game.selected_tower
        cost = t.upgrade_cost()
        up = f"U: upgrade ({cost})" if cost is not None else "max level"
        line = f"{t.spec.name} level {t.level}   {up}   S: sell (+{int(t.spent * SELL_REFUND)})"
    screen.blit(fonts["small"].render(line, True, (200, 210, 230)), (540, HUD_Y + 38))

    if game.phase in (Phase.WON, Phase.LOST):
        shade = pygame.Surface((WIDTH, HUD_Y), pygame.SRCALPHA)
        shade.fill((0, 0, 0, 150))
        screen.blit(shade, (0, 0))
        title = "Victory!" if game.phase == Phase.WON else "Game Over"
        label = fonts["huge"].render(title, True, (255, 230, 120))
        screen.blit(label, label.get_rect(center=(WIDTH / 2, HUD_Y / 2)))


# ---- Input and the main loop ------------------------------------------------------
def handle_key(game, key):
    if key in KEY_TO_KIND:
        game.selected_kind = KEY_TO_KIND[key]
        game.selected_tower = None
    elif key == pygame.K_SPACE:
        game.start_wave()
    elif key == pygame.K_p:
        game.paused = not game.paused
    elif key == pygame.K_f:
        game.speed = 2 if game.speed == 1 else 1
    elif key == pygame.K_k:
        game.shake_on = not game.shake_on       # accessibility: screen shake on/off
    elif key == pygame.K_ESCAPE:
        game.selected_tower = None
    elif game.selected_tower is not None:
        if key == pygame.K_u:
            game.upgrade(game.selected_tower)
        elif key == pygame.K_s:
            game.sell(game.selected_tower)


def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Tower Defense Capstone")
    clock = pygame.time.Clock()
    fonts = {"small": pygame.font.Font(None, 22), "big": pygame.font.Font(None, 28),
             "huge": pygame.font.Font(None, 80)}
    art = load_art()
    background = build_background(art)
    play = pygame.Surface((WIDTH, HUD_Y))   # the world is drawn here, then shaken
    game = Game(seed=1)

    running = True
    while running:
        dt = min(clock.tick(60) / 1000, 0.05)   # cap long frames (window drags)
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == pygame.K_r and game.phase in (Phase.WON, Phase.LOST):
                    game = Game(seed=1)
                else:
                    handle_key(game, event.key)
            elif event.type == pygame.MOUSEBUTTONDOWN and event.button == 1:
                game.click(event.pos)

        game.update(dt)
        draw(screen, play, game, art, fonts, background)
        pygame.display.flip()

    pygame.quit()
    print(f"Reached wave {game.wave_number} with {game.lives} lives, "
          f"{game.gold} gold and {game.kills} kills.")


if __name__ == "__main__":
    main()