How to wait for NMI

; // Summary
; // A Very Basic NES ROM that displays 1 color: Blue

; ////////////////////////////////////////////////
; // Works fine!
.segment "INESHDR"
  .byt "NES",$1A
  .byt 1                ; 1 x 16kB PRG block.
  .byt 1                ; 1 x 8kB CHR block.
; ////////////////////////////////////////////////

.segment "VECTORS"
  .addr nmi, reset, irq

.segment "ZEROPAGE"
  retraces:        .res 1

PPUCTRL = $2000
PPUMASK = $2001
PPUSTATUS = $2002

.segment "CODE"
  .proc reset
    ; Disable interrupts:
    SEI

    ; Basic init:
    LDX #0
    STX PPUCTRL     ; General init state; NMIs (bit 7) disabled.
    STX PPUMASK     ; Disable rendering, i.e. turn off background & sprites.

    ; PPU Warmup (29658 Cycles)
    BIT PPUSTATUS   ; Clear the VPL Flag if it was set at reset time.
  vwait1:
    BIT PPUSTATUS   ; 27384 Cycles has passed.
    BPL vwait1
  vwait2:
    BIT PPUSTATUS   ; 57165 Cycles has passed.
    BPL vwait2

    ; Set stack pointer:
    ;DEX                ; X = $FF
    ;TXS

    LDA #%01000000 ; Intensify greens
    STA PPUMASK

    LDA #%10000000
    STA PPUCTRL

    LDX #60  ; one second
    JSR WaitXFrames

    LDA #%10000000 ; Intensify blues
    STA $2001

    Forever:
      JMP Forever     ; Infinite loop
  .endproc

  .proc nmi
    INC nmi
    RTI
  .endproc

  .proc irq
    RTI
  .endproc

  .proc WaitXFrames
    LDA retraces
    wait:       ; Wait for a frame
      CMP retraces  ; when the NMI handler changes the value
      BEQ wait      ; in this variable, one frame has elapsed
      DEX       ; X = X - 1
      BNE WaitXFrames  ; if X is not zero, keep going...
      RTS
  .endproc