4 files changed, 85 insertions, 10 deletions
diff --git a/src/assets.s b/src/assets.s
index 8e839d3..ffc6585 100644
--- a/src/assets.s
+++ b/src/assets.s
@@ -74,7 +74,30 @@
         dex
         bne @upper_title_bar_loop
 
-        ;; TODO: store back palettes after game over.
+        ;; Update 2nd palette for background. This is redundant upon entering
+        ;; the game, but it makes sense after a game over.
+        lda #$3F
+        sta PPU::ADDRESS
+        lda #$09
+        sta PPU::ADDRESS
+        lda #$28
+        sta PPU::DATA
+        lda #$2C
+        sta PPU::DATA
+        lda #$16
+        sta PPU::DATA
+
+        ;; Update 1st palette for foreground.
+        lda #$3F
+        sta PPU::ADDRESS
+        lda #$11
+        sta PPU::ADDRESS
+        lda #$30
+        sta PPU::DATA
+        lda #$10
+        sta PPU::DATA
+        lda #$30
+        sta PPU::DATA
 
         rts
     .endproc
@@ -82,11 +105,32 @@
     ;; Performs all the needed tricks in order to get the second nametable as
     ;; expected.
     .proc prepare_for_main_screen
-        ;; TODO: second palette for background should be:
-        ;; .byte $0F, $14, $2C, $28
-        ;;
-        ;; TODO: first palette for the foreground should be:
-        ;; .byte $0F, $16, $10, $30
+        bit PPU::STATUS
+
+        ;; Update 2nd palette for background.
+        lda #$3F
+        sta PPU::ADDRESS
+        lda #$09
+        sta PPU::ADDRESS
+        lda #$14
+        sta PPU::DATA
+        lda #$2C
+        sta PPU::DATA
+        lda #$28
+        sta PPU::DATA
+
+        ;; Update 1st palette for foreground.
+        lda #$3F
+        sta PPU::ADDRESS
+        lda #$11
+        sta PPU::ADDRESS
+        lda #$16
+        sta PPU::DATA
+        lda #$10
+        sta PPU::DATA
+        lda #$30
+        sta PPU::DATA
+
         rts
     .endproc
 
diff --git a/src/driver.s b/src/driver.s
index 8b7f8ce..bb0303f 100644
--- a/src/driver.s
+++ b/src/driver.s
@@ -9,7 +9,25 @@
     zp_player_timer = $30
     PLAYER_TIMER_VALUE = HZ * 2
 
+    ;; Switch from the title screen to the main street. Note that this function
+    ;; is to be called with the PPU disabled. If that's not the case, then it
+    ;; will set the proper values to disable it on the next `nmi` call and set
+    ;; the `title over` flag. With that, call again this function so the
+    ;; switching is actually performed.
     .proc switch
+        ;; Some things from here require the PPU to be disabled. Hence, if
+        ;; that's not the case, disable it now. The `ppu` and the `title over`
+        ;; flags are set as well.
+        lda PPU::zp_mask
+        beq @do_switch
+        lda #%01000100
+        ora Globals::zp_flags
+        sta Globals::zp_flags
+        lda #$00
+        sta PPU::zp_mask
+        rts
+
+    @do_switch:
         ;; Get the assets ready for the main screen. That is, make sure that the
         ;; palettes and such are as desired since the title screen needed
         ;; another setup.
@@ -19,6 +37,10 @@
         lda #%10001010
         sta PPU::zp_control
 
+        ;; Enable back the PPU.
+        lda #%00011110
+        sta PPU::zp_mask
+
         ;; Setup the player timer.
         .ifdef PARTIAL
             lda #1
@@ -28,10 +50,10 @@
         sta zp_player_timer
 
         ;; Mark the state of the game as "game". That is, the player has
-        ;; started. Also set the `ppu` flag so the PPU control update takes
-        ;; place.
+        ;; started. Also set the `ppu` flag and unset the `title over` one.
         lda #%01000001
         ora Globals::zp_flags
+        and #%11111011
         sta Globals::zp_flags
 
         rts
diff --git a/src/jetpac.s b/src/jetpac.s
index 66d5457..184374e 100644
--- a/src/jetpac.s
+++ b/src/jetpac.s
@@ -40,8 +40,8 @@
 .include "assets.s"
 .include "background.s"
 .include "player.s"
-.include "driver.s"
 .include "title.s"
+.include "driver.s"
 .include "vectors.s"
 
 .proc main
@@ -72,6 +72,7 @@
 
     ;; Enable back the PPU
     lda #%00011110
+    sta PPU::zp_mask
     sta PPU::MASK
 
 @main_game_loop:
@@ -85,9 +86,15 @@
     jmp @over
 
 @title_screen:
+    ;; If we are in a transitioning state, avoid the update from the title.
+    lda Globals::zp_flags
+    and #%00000100
+    bne @do_switch
+
     jsr Title::update
     beq @set_flags
 
+@do_switch:
     ;; Start was pressed by the player, switch to the main game.
     jsr Driver::switch
     jmp @set_flags
diff --git a/src/vectors.s b/src/vectors.s
index b06f404..b404d98 100644
--- a/src/vectors.s
+++ b/src/vectors.s
@@ -103,7 +103,9 @@
 
     bit PPU::STATUS
 
-    ;; Update the PPU control register with the shadowed value.
+    ;; Update the PPU control/mask registers with shadowed values.
+    lda PPU::zp_mask
+    sta PPU::MASK
     lda PPU::zp_control
     sta PPU::CONTROL