diff --git a/services/core/java/com/android/server/wm/AppWindowToken.java b/services/core/java/com/android/server/wm/AppWindowToken.java
index fc76102..e57fea3 100644
--- a/services/core/java/com/android/server/wm/AppWindowToken.java
+++ b/services/core/java/com/android/server/wm/AppWindowToken.java
@@ -679,11 +679,12 @@
         removed = true;
         stopFreezingScreen(true, true);
 
-        if (mService.mFocusedApp == this) {
-            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Removing focused app token:" + this);
-            mService.mFocusedApp = null;
+        final DisplayContent dc = getDisplayContent();
+        if (dc.mFocusedApp == this) {
+            if (DEBUG_FOCUS_LIGHT) Slog.v(TAG_WM, "Removing focused app token:" + this
+                   + " displayId=" + dc.getDisplayId());
+            dc.setFocusedApp(null);
             mService.updateFocusedWindowLocked(UPDATE_FOCUS_NORMAL, true /*updateInputWindows*/);
-            getDisplayContent().getInputMonitor().setFocusedAppLw(null);
         }
 
         if (!delayed) {
@@ -1064,6 +1065,18 @@
         }
     }
 
+    @Override
+    void onDisplayChanged(DisplayContent dc) {
+        DisplayContent prevDc = mDisplayContent;
+        super.onDisplayChanged(dc);
+        if (prevDc != null && prevDc.mFocusedApp == this) {
+            prevDc.setFocusedApp(null);
+            if (dc.getTopStack().getTopChild().getTopChild() == this) {
+                dc.setFocusedApp(this);
+            }
+        }
+    }
+
     /**
      * Freezes the task bounds. The size of this task reported the app will be fixed to the bounds
      * freezed by {@link Task#prepareFreezingBounds} until {@link #unfreezeBounds} gets called, even