diff --git a/src/common.ts b/src/common.ts
index 35efece0..11949b8b 100644
--- a/src/common.ts
+++ b/src/common.ts
@@ -1,6 +1,7 @@
 import { HassEntity, MessageBase } from 'home-assistant-js-websocket';
 import {
   HomeAssistant,
+  computeStateDomain,
   handleActionConfig,
   hasAction,
   stateIcon,
@@ -400,6 +401,13 @@ export function refreshDynamicStateParameters(
   }
   params.title = params.title ?? (state?.attributes?.friendly_name || params.entity);
   params.icon = params.icon ?? stateIcon(state);
+
+  const domain = state ? computeStateDomain(state) : undefined;
+  params.data_domain =
+    params.state_color || (domain === 'light' && params.state_color !== false)
+      ? domain
+      : undefined;
+  params.data_state = computeActiveState(state);
   return params;
 }
 
diff --git a/src/components/menu.ts b/src/components/menu.ts
index 8f200a22..3996e6ca 100644
--- a/src/components/menu.ts
+++ b/src/components/menu.ts
@@ -231,9 +231,6 @@ export class FrigateCardMenu extends LitElement {
       button: true,
     };
 
-    const stateObj = stateParameters.entity ? this.hass?.states[stateParameters.entity] : undefined;
-    const domain = stateObj ? computeStateDomain(stateObj) : undefined;
-
     // =====================================================================================
     // For `data-domain` and `data-state`, see: See
     // https://github.com/home-assistant/frontend/blob/dev/src/components/entity/state-badge.ts#L54
@@ -247,12 +244,8 @@ export class FrigateCardMenu extends LitElement {
     //   (`data-state`). This looks up a CSS style in `menu.scss`.
 
     return html` <ha-icon-button
-      data-domain=${ifDefined(
-        stateParameters.state_color || (domain === "light" && stateParameters.state_color !== false)
-          ? domain
-          : undefined
-      )}
-      data-state=${stateObj ? computeActiveState(stateObj) : ""}
+      data-domain=${ifDefined(stateParameters.data_domain)}
+      data-state=${ifDefined(stateParameters.data_state)}
       class="${classMap(classes)}"
       style="${styleMap(stateParameters.style || {})}"
       .actionHandler=${actionHandler({
diff --git a/src/components/submenu.ts b/src/components/submenu.ts
index 2893d14a..23cb33de 100644
--- a/src/components/submenu.ts
+++ b/src/components/submenu.ts
@@ -1,11 +1,12 @@
 import { CSSResultGroup, html, LitElement, TemplateResult, unsafeCSS } from 'lit';
 import { customElement, property } from 'lit/decorators';
-import { HomeAssistant } from 'custom-card-helpers';
+import { computeActiveState, frigateCardHasAction, refreshDynamicStateParameters } from '../common.js';
+import { computeStateDomain, HomeAssistant } from 'custom-card-helpers';
+import { ifDefined } from 'lit/directives/if-defined';
 import { styleMap } from 'lit/directives/style-map';
 
 import { ExtendedHomeAssistant, MenuSubmenu, MenuSubmenuItem } from '../types.js';
 import { actionHandler } from '../action-handler-directive.js';
-import { frigateCardHasAction, refreshDynamicStateParameters } from '../common.js';
 
 import submenuStyle from '../scss/submenu.scss';
 import type { Corner } from "@material/mwc-menu";
@@ -46,6 +47,8 @@ export class FrigateCardSubmenu extends LitElement {
         ${stateParameters.title || ''}
         ${stateParameters.icon
           ? html` <ha-icon
+              data-domain=${ifDefined(stateParameters.data_domain)}
+              data-state=${ifDefined(stateParameters.data_state)}
               style="${styleMap(stateParameters.style || {})}"
               slot="graphic"
               icon="${stateParameters.icon}"
diff --git a/src/scss/button.scss b/src/scss/button.scss
index 3504aab3..3577c19b 100644
--- a/src/scss/button.scss
+++ b/src/scss/button.scss
@@ -8,67 +8,12 @@ ha-icon-button.button {
   /* Buttons can always be clicked */
   pointer-events: auto;
   opacity: 0.9;
-
 }
+
 // =====================================================================================
 // From: https://github.com/home-assistant/frontend/blob/dev/src/common/style/icon_color_css.ts
-// Modified for `ha-icon-button` instead of `ha-state-icon`.
+// Modified version for ha-icon and ha-icon-button.
 // =====================================================================================
-ha-icon-button[data-domain="alert"][data-state="on"],
-ha-icon-button[data-domain="automation"][data-state="on"],
-ha-icon-button[data-domain="binary_sensor"][data-state="on"],
-ha-icon-button[data-domain="calendar"][data-state="on"],
-ha-icon-button[data-domain="camera"][data-state="streaming"],
-ha-icon-button[data-domain="cover"][data-state="open"],
-ha-icon-button[data-domain="fan"][data-state="on"],
-ha-icon-button[data-domain="humidifier"][data-state="on"],
-ha-icon-button[data-domain="light"][data-state="on"],
-ha-icon-button[data-domain="input_boolean"][data-state="on"],
-ha-icon-button[data-domain="lock"][data-state="unlocked"],
-ha-icon-button[data-domain="media_player"][data-state="on"],
-ha-icon-button[data-domain="media_player"][data-state="paused"],
-ha-icon-button[data-domain="media_player"][data-state="playing"],
-ha-icon-button[data-domain="script"][data-state="on"],
-ha-icon-button[data-domain="sun"][data-state="above_horizon"],
-ha-icon-button[data-domain="switch"][data-state="on"],
-ha-icon-button[data-domain="timer"][data-state="active"],
-ha-icon-button[data-domain="vacuum"][data-state="cleaning"],
-ha-icon-button[data-domain="group"][data-state="on"],
-ha-icon-button[data-domain="group"][data-state="home"],
-ha-icon-button[data-domain="group"][data-state="open"],
-ha-icon-button[data-domain="group"][data-state="locked"],
-ha-icon-button[data-domain="group"][data-state="problem"] {
-  color: var(--paper-item-icon-active-color, #fdd835);
-}
-
-ha-icon-button[data-domain="climate"][data-state="cooling"] {
-  color: var(--cool-color, var(--state-climate-cool-color));
-}
-
-ha-icon-button[data-domain="climate"][data-state="heating"] {
-  color: var(--heat-color, var(--state-climate-heat-color));
-}
-
-ha-icon-button[data-domain="climate"][data-state="drying"] {
-  color: var(--dry-color, var(--state-climate-dry-color));
-}
-
-ha-icon-button[data-domain="alarm_control_panel"] {
-  color: var(--alarm-color-armed, var(--label-badge-red));
-}
-ha-icon-button[data-domain="alarm_control_panel"][data-state="disarmed"] {
-  color: var(--alarm-color-disarmed, var(--label-badge-green));
-}
-ha-icon-button[data-domain="alarm_control_panel"][data-state="pending"],
-ha-icon-button[data-domain="alarm_control_panel"][data-state="arming"] {
-  color: var(--alarm-color-pending, var(--label-badge-yellow));
-  animation: pulse 1s infinite;
-}
-ha-icon-button[data-domain="alarm_control_panel"][data-state="triggered"] {
-  color: var(--alarm-color-triggered, var(--label-badge-red));
-  animation: pulse 1s infinite;
-}
-
 @keyframes pulse {
   0% {
     opacity: 1;
@@ -81,12 +26,72 @@ ha-icon-button[data-domain="alarm_control_panel"][data-state="triggered"] {
   }
 }
 
-ha-icon-button[data-domain="plant"][data-state="problem"],
-ha-icon-button[data-domain="zwave"][data-state="dead"] {
-  color: var(--state-icon-error-color);
+@mixin state-icon($element) {
+  #{$element}[data-domain="alert"][data-state="on"],
+  #{$element}[data-domain="automation"][data-state="on"],
+  #{$element}[data-domain="binary_sensor"][data-state="on"],
+  #{$element}[data-domain="calendar"][data-state="on"],
+  #{$element}[data-domain="camera"][data-state="streaming"],
+  #{$element}[data-domain="cover"][data-state="open"],
+  #{$element}[data-domain="fan"][data-state="on"],
+  #{$element}[data-domain="humidifier"][data-state="on"],
+  #{$element}[data-domain="light"][data-state="on"],
+  #{$element}[data-domain="input_boolean"][data-state="on"],
+  #{$element}[data-domain="lock"][data-state="unlocked"],
+  #{$element}[data-domain="media_player"][data-state="on"],
+  #{$element}[data-domain="media_player"][data-state="paused"],
+  #{$element}[data-domain="media_player"][data-state="playing"],
+  #{$element}[data-domain="script"][data-state="on"],
+  #{$element}[data-domain="sun"][data-state="above_horizon"],
+  #{$element}[data-domain="switch"][data-state="on"],
+  #{$element}[data-domain="timer"][data-state="active"],
+  #{$element}[data-domain="vacuum"][data-state="cleaning"],
+  #{$element}[data-domain="group"][data-state="on"],
+  #{$element}[data-domain="group"][data-state="home"],
+  #{$element}[data-domain="group"][data-state="open"],
+  #{$element}[data-domain="group"][data-state="locked"],
+  #{$element}[data-domain="group"][data-state="problem"] {
+    color: var(--paper-item-icon-active-color, #fdd835);
+  }
+
+  #{$element}[data-domain="climate"][data-state="cooling"] {
+    color: var(--cool-color, var(--state-climate-cool-color));
+  }
+
+  #{$element}[data-domain="climate"][data-state="heating"] {
+    color: var(--heat-color, var(--state-climate-heat-color));
+  }
+
+  #{$element}[data-domain="climate"][data-state="drying"] {
+    color: var(--dry-color, var(--state-climate-dry-color));
+  }
+
+  #{$element}[data-domain="alarm_control_panel"] {
+    color: var(--alarm-color-armed, var(--label-badge-red));
+  }
+  #{$element}[data-domain="alarm_control_panel"][data-state="disarmed"] {
+    color: var(--alarm-color-disarmed, var(--label-badge-green));
+  }
+  #{$element}[data-domain="alarm_control_panel"][data-state="pending"],
+  #{$element}[data-domain="alarm_control_panel"][data-state="arming"] {
+    color: var(--alarm-color-pending, var(--label-badge-yellow));
+    animation: pulse 1s infinite;
+  }
+  #{$element}[data-domain="alarm_control_panel"][data-state="triggered"] {
+    color: var(--alarm-color-triggered, var(--label-badge-red));
+    animation: pulse 1s infinite;
+  }
+
+  #{$element}[data-domain="plant"][data-state="problem"],
+  #{$element}[data-domain="zwave"][data-state="dead"] {
+    color: var(--state-icon-error-color);
+  }
+
+  /* Color the icon if unavailable */
+  #{$element}[data-state="unavailable"] {
+    color: var(--state-unavailable-color);
+  }
 }
 
-/* Color the icon if unavailable */
-ha-icon-button[data-state="unavailable"] {
-  color: var(--state-unavailable-color);
-}
\ No newline at end of file
+@include state-icon("ha-icon");
+@include state-icon("ha-icon-button");
\ No newline at end of file
diff --git a/src/types.ts b/src/types.ts
index 1aea91e3..ee024c16 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -842,6 +842,8 @@ export interface StateParameters {
   title?: string | null;
   state_color?: boolean;
   style?: StyleInfo;
+  data_domain?: string;
+  data_state?: string;
 }
 
 export interface FrigateCardMediaPlayer {