zsh color prompt based on vi mode

.zshrc

# enable colors
autoload -U colors && colors

# vi mode
bindkey -v

# default color (white)
prompt_color="\$reset_color"

# set color based on current mode
# insert mode: white
# normal mode: blue
function zle-line-init zle-keymap-select {
    prompt_color="${${KEYMAP/vicmd/\$fg[blue]}/(main|viins)/\$reset_color}"
    PS1="%{$prompt_color%}%n%\@%m%  %~ %%%{$reset_color%} "
    zle reset-prompt
}

zle -N zle-line-init
zle -N zle-keymap-select

export KEYTIMEOUT=2