diff --git a/powerlevel9k.zsh-theme b/powerlevel9k.zsh-theme
index 479c6880..7fb283dd 100644
--- a/powerlevel9k.zsh-theme
+++ b/powerlevel9k.zsh-theme
@@ -813,12 +813,14 @@ prompt_time() {
 
 # Vi Mode: show editing mode (NORMAL|INSERT)
 prompt_vi_mode() {
-  local mode="${${KEYMAP/vicmd/NORMAL}/(main|viins)/INSERT}"
-  if [[ "$mode" == "NORMAL" ]]; then
-    "$1_prompt_segment" "$0_NORMAL" "$DEFAULT_COLOR" "default" "$mode"
-  else
-    "$1_prompt_segment" "$0_INSERT" "$DEFAULT_COLOR" "blue" "$mode"
-  fi
+  case ${KEYMAP} in
+    main)
+      "$1_prompt_segment" "$0_INSERT" "$DEFAULT_COLOR" "blue" "INSERT"
+    ;;
+    *)
+      "$1_prompt_segment" "$0_NORMAL" "$DEFAULT_COLOR" "default" "NORMAL"
+    ;;
+  esac
 }
 
 # Virtualenv: current working virtualenv