Screen off animation

SysUI can now control the screen off animation as long as
config_dozeAfterScreenOff is set to false.

The current implementation collapses the notification shade and moves
the clock whenever the use is on the lock screen, or will fade the
scrims and show the clock when the keyguard is occluded.

Display state change (on, doze, doze_suspended) is delayed to let the
animations occur at 60Hz.

Test: atest packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
Test: atest packages/SystemUI/tests/src/com/android/systemui/doze/DozeWallpaperStateTest.java
Test: atest packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ScrimControllerTest.java
Test: atest tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
Test: atest tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
Test: atest packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java
Fixes: 73178121
Change-Id: Id5d964452f342d4e97bedf1084efa808604e602c
diff --git a/core/java/android/os/IPowerManager.aidl b/core/java/android/os/IPowerManager.aidl
index 1681f11..13e4e38 100644
--- a/core/java/android/os/IPowerManager.aidl
+++ b/core/java/android/os/IPowerManager.aidl
@@ -65,4 +65,7 @@
 
     // sets the attention light (used by phone app only)
     void setAttentionLight(boolean on, int color);
+
+    // controls whether PowerManager should doze after the screen turns off or not
+    void setDozeAfterScreenOff(boolean on);
 }
diff --git a/core/java/android/os/PowerManager.java b/core/java/android/os/PowerManager.java
index 66fa629..c00100b 100644
--- a/core/java/android/os/PowerManager.java
+++ b/core/java/android/os/PowerManager.java
@@ -1280,6 +1280,19 @@
     }
 
     /**
+     * If true, the doze component is not started until after the screen has been
+     * turned off and the screen off animation has been performed.
+     * @hide
+     */
+    public void setDozeAfterScreenOff(boolean dozeAfterScreenOf) {
+        try {
+            mService.setDozeAfterScreenOff(dozeAfterScreenOf);
+        } catch (RemoteException e) {
+            throw e.rethrowFromSystemServer();
+        }
+    }
+
+    /**
      * Returns the reason the phone was last shutdown. Calling app must have the
      * {@link android.Manifest.permission#DEVICE_POWER} permission to request this information.
      * @return Reason for shutdown as an int, {@link #SHUTDOWN_REASON_UNKNOWN} if the file could
diff --git a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
index 9d5fb52..d24675c 100644
--- a/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
+++ b/packages/SystemUI/src/com/android/keyguard/KeyguardUpdateMonitor.java
@@ -1726,6 +1726,7 @@
         callback.onPhoneStateChanged(mPhoneState);
         callback.onRefreshCarrierInfo();
         callback.onClockVisibilityChanged();
+        callback.onKeyguardVisibilityChangedRaw(mKeyguardIsVisible);
         for (Entry<Integer, SimData> data : mSimDatas.entrySet()) {
             final SimData state = data.getValue();
             callback.onSimStateChanged(state.subId, state.slotId, state.simState);
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java b/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
index 092f3d2..5bf62f6 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeFactory.java
@@ -22,8 +22,10 @@
 import android.hardware.Sensor;
 import android.hardware.SensorManager;
 import android.os.Handler;
+import android.os.PowerManager;
 
 import com.android.internal.hardware.AmbientDisplayConfiguration;
+import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
 import com.android.systemui.SystemUIApplication;
@@ -46,7 +48,7 @@
 
         DozeHost host = getHost(dozeService);
         AmbientDisplayConfiguration config = new AmbientDisplayConfiguration(context);
-        DozeParameters params = new DozeParameters(context);
+        DozeParameters params = DozeParameters.getInstance(context);
         Handler handler = new Handler();
         WakeLock wakeLock = new DelayedWakeLock(handler,
                 WakeLock.createPartial(context, "Doze"));
@@ -64,9 +66,9 @@
                 createDozeTriggers(context, sensorManager, host, alarmManager, config, params,
                         handler, wakeLock, machine),
                 createDozeUi(context, host, wakeLock, machine, handler, alarmManager, params),
-                new DozeScreenState(wrappedService, handler, params),
+                new DozeScreenState(wrappedService, handler, params, wakeLock),
                 createDozeScreenBrightness(context, wrappedService, sensorManager, host, handler),
-                new DozeWallpaperState(context)
+                new DozeWallpaperState(context, params)
         });
 
         return machine;
@@ -92,7 +94,8 @@
     private DozeMachine.Part createDozeUi(Context context, DozeHost host, WakeLock wakeLock,
             DozeMachine machine, Handler handler, AlarmManager alarmManager,
             DozeParameters params) {
-        return new DozeUi(context, alarmManager, machine, wakeLock, host, handler, params);
+        return new DozeUi(context, alarmManager, machine, wakeLock, host, handler, params,
+                KeyguardUpdateMonitor.getInstance(context));
     }
 
     public static DozeHost getHost(DozeService service) {
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
index 6ff8e3d..152b9fc 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeMachine.java
@@ -92,17 +92,17 @@
             switch (this) {
                 case UNINITIALIZED:
                 case INITIALIZED:
-                case DOZE:
+                case DOZE_REQUEST_PULSE:
+                    return parameters.shouldControlScreenOff() ? Display.STATE_ON
+                            : Display.STATE_OFF;
                 case DOZE_AOD_PAUSED:
+                case DOZE:
                     return Display.STATE_OFF;
                 case DOZE_PULSING:
                     return Display.STATE_ON;
                 case DOZE_AOD:
                 case DOZE_AOD_PAUSING:
                     return Display.STATE_DOZE_SUSPEND;
-                case DOZE_REQUEST_PULSE:
-                    return parameters.getDisplayNeedsBlanking() ? Display.STATE_OFF
-                            : Display.STATE_ON;
                 default:
                     return Display.STATE_UNKNOWN;
             }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
index 7d14564..f72bff5d6 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeScreenState.java
@@ -21,6 +21,7 @@
 import android.view.Display;
 
 import com.android.systemui.statusbar.phone.DozeParameters;
+import com.android.systemui.util.wakelock.WakeLock;
 
 /**
  * Controls the screen when dozing.
@@ -30,18 +31,27 @@
     private static final boolean DEBUG = DozeService.DEBUG;
     private static final String TAG = "DozeScreenState";
 
+    /**
+     * Delay entering low power mode when animating to make sure that we'll have
+     * time to move all elements into their final positions while still at 60 fps.
+     */
+    private static final int ENTER_DOZE_DELAY = 3000;
+
     private final DozeMachine.Service mDozeService;
     private final Handler mHandler;
     private final Runnable mApplyPendingScreenState = this::applyPendingScreenState;
     private final DozeParameters mParameters;
 
     private int mPendingScreenState = Display.STATE_UNKNOWN;
+    private boolean mWakeLockHeld;
+    private WakeLock mWakeLock;
 
     public DozeScreenState(DozeMachine.Service service, Handler handler,
-            DozeParameters parameters) {
+            DozeParameters parameters, WakeLock wakeLock) {
         mDozeService = service;
         mHandler = handler;
         mParameters = parameters;
+        mWakeLock = wakeLock;
     }
 
     @Override
@@ -69,12 +79,33 @@
             // that the screen turns on again before the navigation bar is hidden. To work around
             // that, wait for a traversal to happen before applying the initial screen state.
             mPendingScreenState = screenState;
-            if (!messagePending) {
-                mHandler.post(mApplyPendingScreenState);
+
+            // Delay screen state transitions even longer while animations are running.
+            boolean shouldDelayTransition = newState == DozeMachine.State.DOZE_AOD
+                    && mParameters.shouldControlScreenOff();
+
+            if (!mWakeLockHeld && shouldDelayTransition) {
+                mWakeLockHeld = true;
+                mWakeLock.acquire();
             }
-            return;
+
+            if (!messagePending) {
+                if (DEBUG) {
+                    Log.d(TAG, "Display state changed to " + screenState + " delayed by "
+                            + (shouldDelayTransition ? ENTER_DOZE_DELAY : 1));
+                }
+
+                if (shouldDelayTransition) {
+                    mHandler.postDelayed(mApplyPendingScreenState, ENTER_DOZE_DELAY);
+                } else {
+                    mHandler.post(mApplyPendingScreenState);
+                }
+            } else if (DEBUG) {
+                Log.d(TAG, "Pending display state change to " + screenState);
+            }
+        } else {
+            applyScreenState(screenState);
         }
-        applyScreenState(screenState);
     }
 
     private void applyPendingScreenState() {
@@ -87,6 +118,10 @@
             if (DEBUG) Log.d(TAG, "setDozeScreenState(" + screenState + ")");
             mDozeService.setDozeScreenState(screenState);
             mPendingScreenState = Display.STATE_UNKNOWN;
+            if (mWakeLockHeld) {
+                mWakeLockHeld = false;
+                mWakeLock.release();
+            }
         }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
index 75f1b50..778e630 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeUi.java
@@ -25,6 +25,9 @@
 import android.text.format.Formatter;
 import android.util.Log;
 
+import com.android.internal.annotations.VisibleForTesting;
+import com.android.keyguard.KeyguardUpdateMonitor;
+import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.statusbar.phone.DozeParameters;
 import com.android.systemui.util.AlarmTimeout;
 import com.android.systemui.util.wakelock.WakeLock;
@@ -44,22 +47,46 @@
     private final WakeLock mWakeLock;
     private final DozeMachine mMachine;
     private final AlarmTimeout mTimeTicker;
-    private final boolean mCanAnimateWakeup;
+    private final boolean mCanAnimateTransition;
+    private final DozeParameters mDozeParameters;
+
+    private boolean mKeyguardShowing;
+    private final KeyguardUpdateMonitorCallback mKeyguardVisibilityCallback =
+            new KeyguardUpdateMonitorCallback() {
+
+                @Override
+                public void onKeyguardVisibilityChanged(boolean showing) {
+                    mKeyguardShowing = showing;
+                    updateAnimateScreenOff();
+                }
+            };
 
     private long mLastTimeTickElapsed = 0;
 
     public DozeUi(Context context, AlarmManager alarmManager, DozeMachine machine,
             WakeLock wakeLock, DozeHost host, Handler handler,
-            DozeParameters params) {
+            DozeParameters params, KeyguardUpdateMonitor keyguardUpdateMonitor) {
         mContext = context;
         mMachine = machine;
         mWakeLock = wakeLock;
         mHost = host;
         mHandler = handler;
-        mCanAnimateWakeup = !params.getDisplayNeedsBlanking();
-
+        mCanAnimateTransition = !params.getDisplayNeedsBlanking();
+        mDozeParameters = params;
         mTimeTicker = new AlarmTimeout(alarmManager, this::onTimeTick, "doze_time_tick", handler);
-        mHost.setAnimateScreenOff(params.getCanControlScreenOffAnimation());
+        keyguardUpdateMonitor.registerCallback(mKeyguardVisibilityCallback);
+    }
+
+    /**
+     * Decide if we're taking over the screen-off animation
+     * when the device was configured to skip doze after screen off.
+     */
+    private void updateAnimateScreenOff() {
+        if (mCanAnimateTransition) {
+            final boolean controlScreenOff = mDozeParameters.getAlwaysOn() && mKeyguardShowing;
+            mDozeParameters.setControlScreenOffAnimation(controlScreenOff);
+            mHost.setAnimateScreenOff(controlScreenOff);
+        }
     }
 
     private void pulseWhileDozing(int reason) {
@@ -118,7 +145,7 @@
                 // Keep current state.
                 break;
             default:
-                mHost.setAnimateWakeup(mCanAnimateWakeup);
+                mHost.setAnimateWakeup(mCanAnimateTransition && mDozeParameters.getAlwaysOn());
                 break;
         }
     }
@@ -170,4 +197,9 @@
 
         scheduleTimeTick();
     }
+
+    @VisibleForTesting
+    KeyguardUpdateMonitorCallback getKeyguardCallback() {
+        return mKeyguardVisibilityCallback;
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/doze/DozeWallpaperState.java b/packages/SystemUI/src/com/android/systemui/doze/DozeWallpaperState.java
index 5156272..9d110fb 100644
--- a/packages/SystemUI/src/com/android/systemui/doze/DozeWallpaperState.java
+++ b/packages/SystemUI/src/com/android/systemui/doze/DozeWallpaperState.java
@@ -26,7 +26,6 @@
 import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.keyguard.KeyguardUpdateMonitorCallback;
 import com.android.systemui.statusbar.phone.DozeParameters;
-import com.android.systemui.statusbar.phone.StatusBarKeyguardViewManager;
 
 import java.io.PrintWriter;
 
@@ -43,10 +42,10 @@
     private boolean mIsAmbientMode;
     private final DozeParameters mDozeParameters;
 
-    public DozeWallpaperState(Context context) {
+    public DozeWallpaperState(Context context, DozeParameters dozeParameters) {
         this(IWallpaperManager.Stub.asInterface(
                 ServiceManager.getService(Context.WALLPAPER_SERVICE)),
-                new DozeParameters(context), KeyguardUpdateMonitor.getInstance(context));
+                dozeParameters, KeyguardUpdateMonitor.getInstance(context));
     }
 
     @VisibleForTesting
@@ -80,7 +79,7 @@
 
         final boolean animated;
         if (isAmbientMode) {
-            animated = mDozeParameters.getCanControlScreenOffAnimation() && !mKeyguardVisible;
+            animated = mDozeParameters.shouldControlScreenOff();
         } else {
             animated = !mDozeParameters.getDisplayNeedsBlanking();
         }
@@ -88,8 +87,10 @@
         if (isAmbientMode != mIsAmbientMode) {
             mIsAmbientMode = isAmbientMode;
             try {
-                Log.i(TAG, "AoD wallpaper state changed to: " + mIsAmbientMode
-                        + ", animated: " + animated);
+                if (DEBUG) {
+                    Log.i(TAG, "AOD wallpaper state changed to: " + mIsAmbientMode
+                            + ", animated: " + animated);
+                }
                 mWallpaperManagerService.setInAmbientMode(mIsAmbientMode, animated);
             } catch (RemoteException e) {
                 // Cannot notify wallpaper manager service, but it's fine, let's just skip it.
diff --git a/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java b/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java
index 951c0ea..59c7f23 100644
--- a/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java
+++ b/packages/SystemUI/src/com/android/systemui/keyguard/WakefulnessLifecycle.java
@@ -41,21 +41,33 @@
     }
 
     public void dispatchStartedWakingUp() {
+        if (getWakefulness() == WAKEFULNESS_WAKING) {
+            return;
+        }
         setWakefulness(WAKEFULNESS_WAKING);
         dispatch(Observer::onStartedWakingUp);
     }
 
     public void dispatchFinishedWakingUp() {
+        if (getWakefulness() == WAKEFULNESS_AWAKE) {
+            return;
+        }
         setWakefulness(WAKEFULNESS_AWAKE);
         dispatch(Observer::onFinishedWakingUp);
     }
 
     public void dispatchStartedGoingToSleep() {
+        if (getWakefulness() == WAKEFULNESS_GOING_TO_SLEEP) {
+            return;
+        }
         setWakefulness(WAKEFULNESS_GOING_TO_SLEEP);
         dispatch(Observer::onStartedGoingToSleep);
     }
 
     public void dispatchFinishedGoingToSleep() {
+        if (getWakefulness() == WAKEFULNESS_ASLEEP) {
+            return;
+        }
         setWakefulness(WAKEFULNESS_ASLEEP);
         dispatch(Observer::onFinishedGoingToSleep);
     }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
index 4f09133..473e290 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/NotificationShelf.java
@@ -177,6 +177,9 @@
             mShelfState.yTranslation = Math.max(Math.min(viewEnd, maxShelfEnd) - mShelfState.height,
                     getFullyClosedTranslation());
             mShelfState.zTranslation = ambientState.getBaseZHeight();
+            if (mAmbientState.isDark()) {
+                mShelfState.yTranslation = mAmbientState.getDarkTopPadding();
+            }
             float openedAmount = (mShelfState.yTranslation - getFullyClosedTranslation())
                     / (getIntrinsicHeight() * 2);
             openedAmount = Math.min(1.0f, openedAmount);
@@ -343,7 +346,7 @@
         float maxTop = row.getTranslationY();
         StatusBarIconView icon = row.getEntry().expandedIcon;
         float shelfIconPosition = getTranslationY() + icon.getTop() + icon.getTranslationY();
-        if (shelfIconPosition < maxTop) {
+        if (shelfIconPosition < maxTop && !mAmbientState.isDark()) {
             int top = (int) (maxTop - shelfIconPosition);
             Rect clipRect = new Rect(0, top, icon.getWidth(), Math.max(top, icon.getHeight()));
             icon.setClipBounds(clipRect);
@@ -354,7 +357,7 @@
 
     private void updateContinuousClipping(final ExpandableNotificationRow row) {
         StatusBarIconView icon = row.getEntry().expandedIcon;
-        boolean needsContinuousClipping = ViewState.isAnimatingY(icon);
+        boolean needsContinuousClipping = ViewState.isAnimatingY(icon) && !mAmbientState.isDark();
         boolean isContinuousClipping = icon.getTag(TAG_CONTINUOUS_CLIPPING) != null;
         if (needsContinuousClipping && !isContinuousClipping) {
             ViewTreeObserver.OnPreDrawListener predrawListener =
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
index fb3adf4..07b79a2 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeParameters.java
@@ -17,6 +17,7 @@
 package com.android.systemui.statusbar.phone;
 
 import android.content.Context;
+import android.os.PowerManager;
 import android.os.SystemProperties;
 import android.os.UserHandle;
 import android.provider.Settings;
@@ -24,6 +25,7 @@
 import android.util.MathUtils;
 import android.util.SparseBooleanArray;
 
+import com.android.internal.annotations.VisibleForTesting;
 import com.android.internal.hardware.AmbientDisplayConfiguration;
 import com.android.systemui.Dependency;
 import com.android.systemui.R;
@@ -36,19 +38,35 @@
     private static final int MAX_DURATION = 60 * 1000;
     public static final String DOZE_SENSORS_WAKE_UP_FULLY = "doze_sensors_wake_up_fully";
 
+    private static IntInOutMatcher sPickupSubtypePerformsProxMatcher;
+    private static DozeParameters sInstance;
+
     private final Context mContext;
     private final AmbientDisplayConfiguration mAmbientDisplayConfiguration;
+    private final PowerManager mPowerManager;
 
-    private static IntInOutMatcher sPickupSubtypePerformsProxMatcher;
     private final AlwaysOnDisplayPolicy mAlwaysOnPolicy;
 
     private boolean mDozeAlwaysOn;
+    private boolean mControlScreenOffAnimation;
 
-    public DozeParameters(Context context) {
+    public static DozeParameters getInstance(Context context) {
+        if (sInstance == null) {
+            sInstance = new DozeParameters(context);
+        }
+        return sInstance;
+    }
+
+    @VisibleForTesting
+    protected DozeParameters(Context context) {
         mContext = context;
         mAmbientDisplayConfiguration = new AmbientDisplayConfiguration(mContext);
         mAlwaysOnPolicy = new AlwaysOnDisplayPolicy(context);
 
+        mControlScreenOffAnimation = !getDisplayNeedsBlanking();
+        mPowerManager = mContext.getSystemService(PowerManager.class);
+        mPowerManager.setDozeAfterScreenOff(!mControlScreenOffAnimation);
+
         Dependency.get(TunerService.class).addTunable(this, Settings.Secure.DOZE_ALWAYS_ON,
                 Settings.Secure.ACCESSIBILITY_DISPLAY_INVERSION_ENABLED);
     }
@@ -165,15 +183,21 @@
                 com.android.internal.R.bool.config_displayBlanksAfterDoze);
     }
 
-    /**
-     * Whether we can implement our own screen off animation or if we need
-     * to rely on DisplayPowerManager to dim the display.
-     *
-     * @return {@code true} if SystemUI can control the screen off animation.
-     */
-    public boolean getCanControlScreenOffAnimation() {
-        return !mContext.getResources().getBoolean(
-                com.android.internal.R.bool.config_dozeAfterScreenOff);
+    public boolean shouldControlScreenOff() {
+        return mControlScreenOffAnimation;
+    }
+
+    public void setControlScreenOffAnimation(boolean controlScreenOffAnimation) {
+        if (mControlScreenOffAnimation == controlScreenOffAnimation) {
+            return;
+        }
+        mControlScreenOffAnimation = controlScreenOffAnimation;
+        getPowerManager().setDozeAfterScreenOff(!controlScreenOffAnimation);
+    }
+
+    @VisibleForTesting
+    protected PowerManager getPowerManager() {
+        return mPowerManager;
     }
 
     private boolean getBoolean(String propName, int resId) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeScrimController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeScrimController.java
index 1011383..afd64f3 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeScrimController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/DozeScrimController.java
@@ -78,9 +78,10 @@
         }
     };
 
-    public DozeScrimController(ScrimController scrimController, Context context) {
+    public DozeScrimController(ScrimController scrimController, Context context,
+            DozeParameters dozeParameters) {
         mScrimController = scrimController;
-        mDozeParameters = new DozeParameters(context);
+        mDozeParameters = dozeParameters;
     }
 
     public void setDozing(boolean dozing) {
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelView.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelView.java
index 2711d7a..798a63b 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelView.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/NotificationPanelView.java
@@ -2197,14 +2197,6 @@
         return (1 - t) * start + t * end;
     }
 
-    public void setDozing(boolean dozing, boolean animate) {
-        if (dozing == mDozing) return;
-        mDozing = dozing;
-        if (mStatusBarState == StatusBarState.KEYGUARD) {
-            updateDozingVisibilities(animate);
-        }
-    }
-
     private void updateDozingVisibilities(boolean animate) {
         if (mDozing) {
             mKeyguardStatusBar.setVisibility(View.INVISIBLE);
@@ -2600,11 +2592,16 @@
         }
     }
 
-    public void setDark(boolean dark, boolean animate) {
-        float darkAmount = dark ? 1 : 0;
-        if (mDarkAmount == darkAmount) {
-            return;
+    public void setDozing(boolean dozing, boolean animate) {
+        if (dozing == mDozing) return;
+        mDozing = dozing;
+
+        if (mStatusBarState == StatusBarState.KEYGUARD
+                || mStatusBarState == StatusBarState.SHADE_LOCKED) {
+            updateDozingVisibilities(animate);
         }
+
+        final float darkAmount = dozing ? 1 : 0;
         if (mDarkAnimator != null && mDarkAnimator.isRunning()) {
             if (animate && mDarkAmountTarget == darkAmount) {
                 return;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimController.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimController.java
index 739d8d5..1fb4e13 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimController.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimController.java
@@ -285,10 +285,11 @@
             // with too many things at this case, in order to not skip the initial frames.
             mScrimInFront.postOnAnimationDelayed(this::scheduleUpdate, 16);
             mAnimationDelay = StatusBar.FADE_KEYGUARD_START_DELAY;
-        } else if (!mDozeParameters.getAlwaysOn() && oldState == ScrimState.AOD) {
-            // Execute first frame immediately when display was completely off.
-            // Scheduling a frame isn't enough because the system may aggressively enter doze,
-            // delaying callbacks or never triggering them until the power button is pressed.
+        } else if (!mDozeParameters.getAlwaysOn() && oldState == ScrimState.AOD
+                || (mState == ScrimState.AOD && !mDozeParameters.getDisplayNeedsBlanking())) {
+            // Scheduling a frame isn't enough when:
+            //  • Leaving doze and we need to modify scrim color immediately
+            //  • ColorFade will not kick-in and scrim cannot wait for pre-draw.
             onPreDraw();
         } else {
             scheduleUpdate();
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimState.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimState.java
index 55e8714..5b734eb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimState.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/ScrimState.java
@@ -105,8 +105,7 @@
         @Override
         public void prepare(ScrimState previousState) {
             final boolean alwaysOnEnabled = mDozeParameters.getAlwaysOn();
-            final boolean wasPulsing = previousState == ScrimState.PULSING;
-            mBlankScreen = wasPulsing && !mCanControlScreenOff;
+            mBlankScreen = mDisplayRequiresBlanking;
             mCurrentBehindAlpha = mWallpaperSupportsAmbientMode
                     && !mKeyguardUpdateMonitor.hasLockscreenWallpaper() ? 0f : 1f;
             mCurrentInFrontAlpha = alwaysOnEnabled ? mAodFrontScrimAlpha : 1f;
@@ -114,7 +113,7 @@
             mCurrentBehindTint = Color.BLACK;
             // DisplayPowerManager will blank the screen for us, we just need
             // to set our state.
-            mAnimateChange = mCanControlScreenOff;
+            mAnimateChange = !mDisplayRequiresBlanking;
         }
 
         @Override
@@ -178,7 +177,6 @@
     ScrimView mScrimBehind;
     DozeParameters mDozeParameters;
     boolean mDisplayRequiresBlanking;
-    boolean mCanControlScreenOff;
     boolean mWallpaperSupportsAmbientMode;
     KeyguardUpdateMonitor mKeyguardUpdateMonitor;
     int mIndex;
@@ -192,7 +190,6 @@
         mScrimBehind = scrimBehind;
         mDozeParameters = dozeParameters;
         mDisplayRequiresBlanking = dozeParameters.getDisplayNeedsBlanking();
-        mCanControlScreenOff = dozeParameters.getCanControlScreenOffAnimation();
         mKeyguardUpdateMonitor = KeyguardUpdateMonitor.getInstance(scrimInFront.getContext());
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
index 7422a43..8a92702 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBar.java
@@ -906,7 +906,8 @@
                     if (mStatusBarWindowManager != null) {
                         mStatusBarWindowManager.setScrimsVisibility(scrimsVisible);
                     }
-                }, new DozeParameters(mContext), mContext.getSystemService(AlarmManager.class));
+                }, DozeParameters.getInstance(mContext),
+                mContext.getSystemService(AlarmManager.class));
         if (mScrimSrcModeEnabled) {
             Runnable runnable = () -> {
                 boolean asSrc = mBackdrop.getVisibility() != View.VISIBLE;
@@ -918,7 +919,8 @@
         }
         mHeadsUpManager.addListener(mScrimController);
         mStackScroller.setScrimController(mScrimController);
-        mDozeScrimController = new DozeScrimController(mScrimController, context);
+        mDozeScrimController = new DozeScrimController(mScrimController, context,
+                DozeParameters.getInstance(context));
 
         // Other icons
         mVolumeComponent = getComponent(VolumeComponent.class);
@@ -1867,7 +1869,7 @@
 
     @Override
     public boolean isDozing() {
-        return mDozing;
+        return mDozing && mStackScroller.isFullyDark();
     }
 
     @Override
@@ -3817,13 +3819,16 @@
     private void updateDozingState() {
         Trace.traceCounter(Trace.TRACE_TAG_APP, "dozing", mDozing ? 1 : 0);
         Trace.beginSection("StatusBar#updateDozingState");
+
+        boolean sleepingFromKeyguard =
+                mStatusBarKeyguardViewManager.isGoingToSleepVisibleNotOccluded();
         boolean animate = (!mDozing && mDozeServiceHost.shouldAnimateWakeup())
-                || (mDozing && mDozeServiceHost.shouldAnimateScreenOff());
-        mNotificationPanel.setDozing(mDozing, animate);
+                || (mDozing && mDozeServiceHost.shouldAnimateScreenOff() && sleepingFromKeyguard);
+
         mStackScroller.setDark(mDozing, animate, mWakeUpTouchLocation);
         mDozeScrimController.setDozing(mDozing);
         mKeyguardIndicationController.setDozing(mDozing);
-        mNotificationPanel.setDark(mDozing, animate);
+        mNotificationPanel.setDozing(mDozing, animate);
         updateQsExpansionEnabled();
         mViewHierarchyManager.updateRowStates();
         Trace.endSection();
@@ -4739,6 +4744,7 @@
             if (mDozingRequested) {
                 mDozingRequested = false;
                 DozeLog.traceDozing(mContext, mDozing);
+                mWakefulnessLifecycle.dispatchStartedWakingUp();
                 updateDozing();
             }
         }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
index b26b7c9..8b23fdb 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarKeyguardViewManager.java
@@ -95,6 +95,7 @@
     protected boolean mLastRemoteInputActive;
     private boolean mLastDozing;
     private int mLastFpMode;
+    private boolean mGoingToSleepVisibleNotOccluded;
 
     private OnDismissAction mAfterKeyguardGoneAction;
     private final ArrayList<Runnable> mAfterKeyguardGoneRunnables = new ArrayList<>();
@@ -262,11 +263,16 @@
         }
     }
 
+    public boolean isGoingToSleepVisibleNotOccluded() {
+        return mGoingToSleepVisibleNotOccluded;
+    }
+
     public void onStartedGoingToSleep() {
-        // TODO: remove
+        mGoingToSleepVisibleNotOccluded = isShowing() && !isOccluded();
     }
 
     public void onFinishedGoingToSleep() {
+        mGoingToSleepVisibleNotOccluded = false;
         mBouncer.onScreenTurnedOff();
     }
 
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowManager.java b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowManager.java
index defb46c..309a1a7 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowManager.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/phone/StatusBarWindowManager.java
@@ -73,7 +73,7 @@
         mWindowManager = (WindowManager) context.getSystemService(Context.WINDOW_SERVICE);
         mActivityManager = ActivityManager.getService();
         mKeyguardScreenRotation = shouldEnableKeyguardScreenRotation();
-        mDozeParameters = new DozeParameters(mContext);
+        mDozeParameters = DozeParameters.getInstance(mContext);
         mScreenBrightnessDoze = mDozeParameters.getScreenBrightnessDoze();
     }
 
@@ -141,11 +141,9 @@
             mLpChanged.privateFlags &= ~WindowManager.LayoutParams.PRIVATE_FLAG_KEYGUARD;
         }
 
-        final boolean showWallpaperOnAod = mDozeParameters.getAlwaysOn() &&
-                state.wallpaperSupportsAmbientMode &&
-                state.scrimsVisibility != ScrimController.VISIBILITY_FULLY_OPAQUE;
-        if (state.keyguardShowing && !state.backdropShowing &&
-                (!state.dozing || showWallpaperOnAod)) {
+        final boolean scrimsOccludingWallpaper =
+                state.scrimsVisibility == ScrimController.VISIBILITY_FULLY_OPAQUE;
+        if (state.keyguardShowing && !state.backdropShowing && !scrimsOccludingWallpaper) {
             mLpChanged.flags |= WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
         } else {
             mLpChanged.flags &= ~WindowManager.LayoutParams.FLAG_SHOW_WALLPAPER;
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/stack/AmbientState.java b/packages/SystemUI/src/com/android/systemui/statusbar/stack/AmbientState.java
index 0f637fb..7c1c566 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/stack/AmbientState.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/stack/AmbientState.java
@@ -70,6 +70,8 @@
     private int mIntrinsicPadding;
     private int mExpandAnimationTopChange;
     private ExpandableNotificationRow mExpandingNotification;
+    private boolean mFullyDark;
+    private int mDarkTopPadding;
 
     public AmbientState(Context context) {
         reload(context);
@@ -409,4 +411,26 @@
     public int getExpandAnimationTopChange() {
         return mExpandAnimationTopChange;
     }
+
+    /**
+     * {@see isFullyDark}
+     */
+    public void setFullyDark(boolean fullyDark) {
+        mFullyDark = fullyDark;
+    }
+
+    /**
+     * @return {@code true } when shade is completely dark: in AOD or ambient display.
+     */
+    public boolean isFullyDark() {
+        return mFullyDark;
+    }
+
+    public void setDarkTopPadding(int darkTopPadding) {
+        mDarkTopPadding = darkTopPadding;
+    }
+
+    public int getDarkTopPadding() {
+        return mDarkTopPadding;
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/stack/NotificationStackScrollLayout.java b/packages/SystemUI/src/com/android/systemui/statusbar/stack/NotificationStackScrollLayout.java
index a85f4e2..1e1378e 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/stack/NotificationStackScrollLayout.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/stack/NotificationStackScrollLayout.java
@@ -400,7 +400,7 @@
     private int mSidePaddings;
     private final int mSeparatorWidth;
     private final int mSeparatorThickness;
-    private final Rect mTmpRect = new Rect();
+    private final Rect mBackgroundAnimationRect = new Rect();
     private int mClockBottom;
     private int mAntiBurnInOffsetX;
 
@@ -515,26 +515,29 @@
         final int darkBottom = darkTop + mSeparatorThickness;
 
         if (mAmbientState.hasPulsingNotifications()) {
-            // TODO draw divider between notification and shelf
-        } else if (mAmbientState.isDark()) {
+            // No divider, we have a notification icon instead
+        } else if (mAmbientState.isFullyDark()) {
             // Only draw divider on AOD if we actually have notifications
             if (mFirstVisibleBackgroundChild != null) {
                 canvas.drawRect(darkLeft, darkTop, darkRight, darkBottom, mBackgroundPaint);
             }
-            setClipBounds(null);
         } else {
             float animProgress = Interpolators.FAST_OUT_SLOW_IN
                     .getInterpolation(1f - mDarkAmount);
             float sidePaddingsProgress = Interpolators.FAST_OUT_SLOW_IN
                     .getInterpolation((1f - mDarkAmount) * 2);
-            mTmpRect.set((int) MathUtils.lerp(darkLeft, lockScreenLeft, sidePaddingsProgress),
+            mBackgroundAnimationRect.set(
+                    (int) MathUtils.lerp(darkLeft, lockScreenLeft, sidePaddingsProgress),
                     (int) MathUtils.lerp(darkTop, lockScreenTop, animProgress),
                     (int) MathUtils.lerp(darkRight, lockScreenRight, sidePaddingsProgress),
                     (int) MathUtils.lerp(darkBottom, lockScreenBottom, animProgress));
-            canvas.drawRoundRect(mTmpRect.left, mTmpRect.top, mTmpRect.right, mTmpRect.bottom,
-                    mCornerRadius, mCornerRadius, mBackgroundPaint);
-            setClipBounds(animProgress == 1 ? null : mTmpRect);
+            if (!mAmbientState.isDark() || mFirstVisibleBackgroundChild != null) {
+                canvas.drawRoundRect(mBackgroundAnimationRect.left, mBackgroundAnimationRect.top,
+                        mBackgroundAnimationRect.right, mBackgroundAnimationRect.bottom,
+                        mCornerRadius, mCornerRadius, mBackgroundPaint);
+            }
         }
+        updateClipping();
     }
 
     private void updateBackgroundDimming() {
@@ -693,7 +696,7 @@
         if (mPulsing) {
             mTopPadding = mClockBottom;
         } else {
-            mTopPadding = mAmbientState.isDark() ? mDarkTopPadding : mRegularTopPadding;
+            mTopPadding = (int) MathUtils.lerp(mRegularTopPadding, mDarkTopPadding, mDarkAmount);
         }
         mAmbientState.setLayoutHeight(getLayoutHeight());
         updateAlgorithmLayoutMinHeight();
@@ -820,6 +823,7 @@
         if (mRegularTopPadding != topPadding) {
             mRegularTopPadding = topPadding;
             mDarkTopPadding = topPadding + mDarkSeparatorPadding;
+            mAmbientState.setDarkTopPadding(mDarkTopPadding);
             updateAlgorithmHeightAndPadding();
             updateContentHeight();
             if (animate && mAnimationsEnabled && mIsExpanded) {
@@ -883,13 +887,17 @@
     }
 
     public void updateClipping() {
+        boolean animatingClipping = mDarkAmount > 0 && mDarkAmount < 1;
         boolean clipped = mRequestedClipBounds != null && !mInHeadsUpPinnedMode
                 && !mHeadsUpAnimatingAway;
         if (mIsClipped != clipped) {
             mIsClipped = clipped;
             updateFadingState();
         }
-        if (clipped) {
+
+        if (animatingClipping) {
+            setClipBounds(mBackgroundAnimationRect);
+        } else if (clipped) {
             setClipBounds(mRequestedClipBounds);
         } else {
             setClipBounds(null);
@@ -2075,7 +2083,7 @@
         float previousPaddingAmount = 0.0f;
         int numShownItems = 0;
         boolean finish = false;
-        int maxDisplayedNotifications = mAmbientState.isDark()
+        int maxDisplayedNotifications = mAmbientState.isFullyDark()
                 ? (hasPulsingNotifications() ? 1 : 0)
                 : mMaxDisplayedNotifications;
 
@@ -2085,7 +2093,7 @@
                     && !expandableView.hasNoContentHeight()) {
                 boolean limitReached = maxDisplayedNotifications != -1
                         && numShownItems >= maxDisplayedNotifications;
-                boolean notificationOnAmbientThatIsNotPulsing = mAmbientState.isDark()
+                boolean notificationOnAmbientThatIsNotPulsing = mAmbientState.isFullyDark()
                         && hasPulsingNotifications()
                         && expandableView instanceof ExpandableNotificationRow
                         && !isPulsing(((ExpandableNotificationRow) expandableView).getEntry());
@@ -2168,7 +2176,7 @@
 
     private void updateBackground() {
         // No need to update the background color if it's not being drawn.
-        if (!mShouldDrawNotificationBackground || mAmbientState.isDark()) {
+        if (!mShouldDrawNotificationBackground || mAmbientState.isFullyDark()) {
             return;
         }
 
@@ -3298,7 +3306,7 @@
                             .animateY(mShelf));
             ev.darkAnimationOriginIndex = mDarkAnimationOriginIndex;
             mAnimationEvents.add(ev);
-            startBackgroundFadeIn();
+            startBackgroundFade();
         }
         mDarkNeedsAnimation = false;
     }
@@ -3889,7 +3897,6 @@
         requestChildrenUpdate();
         applyCurrentBackgroundBounds();
         updateWillNotDraw();
-        updateContentHeight();
         updateAntiBurnInTranslation();
         notifyHeightChangeListener(mShelf);
     }
@@ -3910,6 +3917,11 @@
 
     private void setDarkAmount(float darkAmount) {
         mDarkAmount = darkAmount;
+        final boolean fullyDark = darkAmount == 1;
+        if (mAmbientState.isFullyDark() != fullyDark) {
+            mAmbientState.setFullyDark(fullyDark);
+            updateContentHeight();
+        }
         updateBackgroundDimming();
     }
 
@@ -3917,8 +3929,9 @@
         return mDarkAmount;
     }
 
-    private void startBackgroundFadeIn() {
-        ObjectAnimator fadeAnimator = ObjectAnimator.ofFloat(this, DARK_AMOUNT, mDarkAmount, 0f);
+    private void startBackgroundFade() {
+        ObjectAnimator fadeAnimator = ObjectAnimator.ofFloat(this, DARK_AMOUNT, mDarkAmount,
+                mAmbientState.isDark() ? 1f : 0);
         fadeAnimator.setDuration(StackStateAnimator.ANIMATION_DURATION_WAKEUP);
         fadeAnimator.setInterpolator(Interpolators.ALPHA_IN);
         fadeAnimator.start();
@@ -4502,6 +4515,10 @@
                 mAmbientState.getScrollY()));
     }
 
+    public boolean isFullyDark() {
+        return mAmbientState.isFullyDark();
+    }
+
     /**
      * A listener that is notified when the empty space below the notifications is clicked on
      */
diff --git a/packages/SystemUI/src/com/android/systemui/statusbar/stack/StackScrollAlgorithm.java b/packages/SystemUI/src/com/android/systemui/statusbar/stack/StackScrollAlgorithm.java
index 51737a8..f86e9fd 100644
--- a/packages/SystemUI/src/com/android/systemui/statusbar/stack/StackScrollAlgorithm.java
+++ b/packages/SystemUI/src/com/android/systemui/statusbar/stack/StackScrollAlgorithm.java
@@ -184,7 +184,7 @@
     private void updateDimmedActivatedHideSensitive(AmbientState ambientState,
             StackScrollState resultState, StackScrollAlgorithmState algorithmState) {
         boolean dimmed = ambientState.isDimmed();
-        boolean dark = ambientState.isDark();
+        boolean dark = ambientState.isFullyDark();
         boolean hideSensitive = ambientState.isHideSensitive();
         View activatedChild = ambientState.getActivatedChild();
         int childCount = algorithmState.visibleChildren.size();
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
index 8c4fd73..6683636 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeScreenStateTest.java
@@ -28,6 +28,9 @@
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.os.Looper;
@@ -37,6 +40,7 @@
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.phone.DozeParameters;
+import com.android.systemui.util.wakelock.WakeLock;
 import com.android.systemui.utils.os.FakeHandler;
 
 import org.junit.Before;
@@ -54,14 +58,17 @@
     FakeHandler mHandlerFake;
     @Mock
     DozeParameters mDozeParameters;
+    @Mock
+    WakeLock mWakeLock;
 
     @Before
     public void setUp() throws Exception {
         MockitoAnnotations.initMocks(this);
         when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(true);
+        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
         mServiceFake = new DozeServiceFake();
         mHandlerFake = new FakeHandler(Looper.getMainLooper());
-        mScreen = new DozeScreenState(mServiceFake, mHandlerFake, mDozeParameters);
+        mScreen = new DozeScreenState(mServiceFake, mHandlerFake, mDozeParameters, mWakeLock);
     }
 
     @Test
@@ -142,4 +149,23 @@
         assertFalse(mServiceFake.screenStateSet);
     }
 
+    @Test
+    public void test_holdsWakeLockWhenGoingToLowPowerDelayed() {
+        // Transition to low power mode will be delayed to let
+        // animations play at 60 fps.
+        when(mDozeParameters.shouldControlScreenOff()).thenReturn(true);
+        mHandlerFake.setMode(QUEUEING);
+
+        mScreen.transitionTo(UNINITIALIZED, INITIALIZED);
+        mHandlerFake.dispatchQueuedMessages();
+        reset(mWakeLock);
+
+        mScreen.transitionTo(INITIALIZED, DOZE_AOD);
+        verify(mWakeLock).acquire();
+        verify(mWakeLock, never()).release();
+
+        mHandlerFake.dispatchQueuedMessages();
+        verify(mWakeLock).release();
+    }
+
 }
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
index 75ade9d..0d8d952 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeUiTest.java
@@ -28,15 +28,20 @@
 import static org.mockito.ArgumentMatchers.eq;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import android.app.AlarmManager;
 import android.os.Handler;
 import android.os.HandlerThread;
+import android.os.PowerManager;
 import android.support.test.filters.SmallTest;
 import android.support.test.runner.AndroidJUnit4;
 
+import com.android.keyguard.KeyguardUpdateMonitor;
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.phone.DozeParameters;
 import com.android.systemui.util.wakelock.WakeLockFake;
@@ -46,33 +51,39 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 @RunWith(AndroidJUnit4.class)
 @SmallTest
 public class DozeUiTest extends SysuiTestCase {
 
+    @Mock
     private AlarmManager mAlarmManager;
+    @Mock
     private DozeMachine mMachine;
+    @Mock
+    private DozeParameters mDozeParameters;
+    @Mock
+    private KeyguardUpdateMonitor mKeyguardUpdateMonitor;
+    @Mock
+    private DozeHost mHost;
     private WakeLockFake mWakeLock;
-    private DozeHostFake mHost;
     private Handler mHandler;
     private HandlerThread mHandlerThread;
     private DozeUi mDozeUi;
 
     @Before
     public void setUp() throws Exception {
+        MockitoAnnotations.initMocks(this);
+
         mHandlerThread = new HandlerThread("DozeUiTest");
         mHandlerThread.start();
-        mAlarmManager = mock(AlarmManager.class);
-        mMachine = mock(DozeMachine.class);
         mWakeLock = new WakeLockFake();
-        mHost = new DozeHostFake();
         mHandler = mHandlerThread.getThreadHandler();
-        DozeParameters params = mock(DozeParameters.class);
-        when(params.getCanControlScreenOffAnimation()).thenReturn(true);
-        when(params.getDisplayNeedsBlanking()).thenReturn(false);
 
-        mDozeUi = new DozeUi(mContext, mAlarmManager, mMachine, mWakeLock, mHost, mHandler, params);
+        mDozeUi = new DozeUi(mContext, mAlarmManager, mMachine, mWakeLock, mHost, mHandler,
+                mDozeParameters, mKeyguardUpdateMonitor);
     }
 
     @After
@@ -96,18 +107,69 @@
     }
 
     @Test
-    public void propagatesAnimateScreenOff() {
-        Assert.assertTrue("animateScreenOff should be true", mHost.animateScreenOff);
+    public void propagatesAnimateScreenOff_noAlwaysOn() {
+        reset(mHost);
+        when(mDozeParameters.getAlwaysOn()).thenReturn(false);
+        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
 
-        DozeParameters params = mock(DozeParameters.class);
-        new DozeUi(mContext, mAlarmManager, mMachine, mWakeLock, mHost, mHandler, params);
-        Assert.assertFalse("animateScreenOff should be false", mHost.animateScreenOff);
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
+        verify(mHost).setAnimateScreenOff(eq(false));
     }
 
     @Test
-    public void transitionSetsAnimateWakeup() {
-        mHost.animateWakeup = false;
+    public void propagatesAnimateScreenOff_alwaysOn() {
+        reset(mHost);
+        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
+        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
+
+        // Take over when the keyguard is visible.
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
+        verify(mHost).setAnimateScreenOff(eq(true));
+
+        // Do not animate screen-off when keyguard isn't visible - PowerManager will do it.
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
+        verify(mHost).setAnimateScreenOff(eq(false));
+    }
+
+    @Test
+    public void neverAnimateScreenOff_whenNotSupported() {
+        // Re-initialize DozeParameters saying that the display requires blanking.
+        reset(mDozeParameters);
+        reset(mHost);
+        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(true);
+        mDozeUi = new DozeUi(mContext, mAlarmManager, mMachine, mWakeLock, mHost, mHandler,
+                mDozeParameters, mKeyguardUpdateMonitor);
+
+        // Never animate if display doesn't support it.
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
+        verify(mHost, never()).setAnimateScreenOff(eq(false));
+    }
+
+    @Test
+    public void transitionSetsAnimateWakeup_alwaysOn() {
+        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
+        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
         mDozeUi.transitionTo(UNINITIALIZED, DOZE);
-        Assert.assertTrue("animateScreenOff should be true", mHost.animateWakeup);
+        verify(mHost).setAnimateWakeup(eq(true));
+    }
+
+    @Test
+    public void keyguardVisibility_changesControlScreenOffAnimation() {
+        // Pre-condition
+        reset(mDozeParameters);
+        when(mDozeParameters.getAlwaysOn()).thenReturn(true);
+        when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
+
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(false);
+        verify(mDozeParameters).setControlScreenOffAnimation(eq(false));
+        mDozeUi.getKeyguardCallback().onKeyguardVisibilityChanged(true);
+        verify(mDozeParameters).setControlScreenOffAnimation(eq(true));
+    }
+
+    @Test
+    public void transitionSetsAnimateWakeup_noAlwaysOn() {
+        mDozeUi.transitionTo(UNINITIALIZED, DOZE);
+        verify(mHost).setAnimateWakeup(eq(false));
     }
 }
\ No newline at end of file
diff --git a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeWallpaperStateTest.java b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeWallpaperStateTest.java
index 2705bca..5c80bb5 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/doze/DozeWallpaperStateTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/doze/DozeWallpaperStateTest.java
@@ -24,7 +24,6 @@
 import static org.mockito.Mockito.when;
 
 import android.app.IWallpaperManager;
-import android.os.Handler;
 import android.os.RemoteException;
 import android.support.test.filters.SmallTest;
 
@@ -37,7 +36,6 @@
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
 import org.mockito.Mock;
-import org.mockito.Mockito;
 import org.mockito.MockitoAnnotations;
 
 @RunWith(JUnit4.class)
@@ -77,7 +75,7 @@
     public void testAnimates_whenSupported() throws RemoteException {
         // Pre-conditions
         when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(false);
-        when(mDozeParameters.getCanControlScreenOffAnimation()).thenReturn(true);
+        when(mDozeParameters.shouldControlScreenOff()).thenReturn(true);
         when(mDozeParameters.getAlwaysOn()).thenReturn(true);
 
         mDozeWallpaperState.transitionTo(DozeMachine.State.UNINITIALIZED,
@@ -92,8 +90,8 @@
     public void testDoesNotAnimate_whenNotSupported() throws RemoteException {
         // Pre-conditions
         when(mDozeParameters.getDisplayNeedsBlanking()).thenReturn(true);
-        when(mDozeParameters.getCanControlScreenOffAnimation()).thenReturn(false);
         when(mDozeParameters.getAlwaysOn()).thenReturn(true);
+        when(mDozeParameters.shouldControlScreenOff()).thenReturn(false);
 
         mDozeWallpaperState.transitionTo(DozeMachine.State.UNINITIALIZED,
                 DozeMachine.State.DOZE_AOD);
diff --git a/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java b/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java
index e15e0b4..9eba9b8 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/keyguard/WakefulnessLifecycleTest.java
@@ -18,6 +18,7 @@
 
 import static org.junit.Assert.assertEquals;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.verifyNoMoreInteractions;
 
@@ -99,6 +100,23 @@
     }
 
     @Test
+    public void doesNotDispatchTwice() throws Exception {
+        mWakefulness.dispatchStartedWakingUp();
+        mWakefulness.dispatchStartedWakingUp();
+        mWakefulness.dispatchFinishedWakingUp();
+        mWakefulness.dispatchFinishedWakingUp();
+        mWakefulness.dispatchStartedGoingToSleep();
+        mWakefulness.dispatchStartedGoingToSleep();
+        mWakefulness.dispatchFinishedGoingToSleep();
+        mWakefulness.dispatchFinishedGoingToSleep();
+
+        verify(mWakefulnessObserver, times(1)).onStartedGoingToSleep();
+        verify(mWakefulnessObserver, times(1)).onFinishedGoingToSleep();
+        verify(mWakefulnessObserver, times(1)).onStartedWakingUp();
+        verify(mWakefulnessObserver, times(1)).onFinishedWakingUp();
+    }
+
+    @Test
     public void dump() throws Exception {
         mWakefulness.dump(null, new PrintWriter(new ByteArrayOutputStream()), new String[0]);
     }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
index e89ff97..550a35d 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeParametersTest.java
@@ -16,11 +16,15 @@
 
 package com.android.systemui.statusbar.phone;
 
+import android.content.Context;
+import android.os.PowerManager;
 import android.support.test.runner.AndroidJUnit4;
 import android.test.suitebuilder.annotation.SmallTest;
 
 import com.android.systemui.SysuiTestCase;
 import com.android.systemui.statusbar.phone.DozeParameters.IntInOutMatcher;
+
+import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 
@@ -28,6 +32,14 @@
 import static junit.framework.Assert.assertFalse;
 import static junit.framework.Assert.fail;
 
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
 @SmallTest
 @RunWith(AndroidJUnit4.class)
 public class DozeParametersTest extends SysuiTestCase {
@@ -186,4 +198,38 @@
         }
     }
 
+    @Test
+    public void test_setControlScreenOffAnimation_setsDozeAfterScreenOff_false() {
+        TestableDozeParameters dozeParameters = new TestableDozeParameters(getContext());
+        PowerManager mockedPowerManager = dozeParameters.getPowerManager();
+        dozeParameters.setControlScreenOffAnimation(true);
+        reset(mockedPowerManager);
+        dozeParameters.setControlScreenOffAnimation(false);
+        verify(mockedPowerManager).setDozeAfterScreenOff(eq(true));
+    }
+
+    @Test
+    public void test_setControlScreenOffAnimation_setsDozeAfterScreenOff_true() {
+        TestableDozeParameters dozeParameters = new TestableDozeParameters(getContext());
+        PowerManager mockedPowerManager = dozeParameters.getPowerManager();
+        dozeParameters.setControlScreenOffAnimation(false);
+        reset(mockedPowerManager);
+        dozeParameters.setControlScreenOffAnimation(true);
+        verify(dozeParameters.getPowerManager()).setDozeAfterScreenOff(eq(false));
+    }
+
+    private class TestableDozeParameters extends DozeParameters {
+        private PowerManager mPowerManager;
+
+        TestableDozeParameters(Context context) {
+            super(context);
+            mPowerManager = mock(PowerManager.class);
+        }
+
+        @Override
+        protected PowerManager getPowerManager() {
+            return mPowerManager;
+        }
+    }
+
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
index ca2f713..203ebe6 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/DozeScrimControllerTest.java
@@ -34,18 +34,23 @@
 import org.junit.Before;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
 
 @RunWith(AndroidTestingRunner.class)
 @TestableLooper.RunWithLooper
 @SmallTest
 public class DozeScrimControllerTest extends SysuiTestCase {
 
+    @Mock
     private ScrimController mScrimController;
+    @Mock
+    private DozeParameters mDozeParameters;
     private DozeScrimController mDozeScrimController;
 
     @Before
     public void setup() {
-        mScrimController = mock(ScrimController.class);
+        MockitoAnnotations.initMocks(this);
         // Make sure callbacks will be invoked to complete the lifecycle.
         doAnswer(invocationOnMock -> {
             ScrimController.Callback callback = invocationOnMock.getArgument(1);
@@ -56,7 +61,8 @@
         }).when(mScrimController).transitionTo(any(ScrimState.class),
                 any(ScrimController.Callback.class));
 
-        mDozeScrimController = new DozeScrimController(mScrimController, getContext());
+        mDozeScrimController = new DozeScrimController(mScrimController, getContext(),
+                mDozeParameters);
         mDozeScrimController.setDozing(true);
     }
 
diff --git a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ScrimControllerTest.java b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ScrimControllerTest.java
index d32c9a8..22e4980 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ScrimControllerTest.java
+++ b/packages/SystemUI/tests/src/com/android/systemui/statusbar/phone/ScrimControllerTest.java
@@ -349,7 +349,7 @@
     }
 
     @Test
-    public void testWillHideAoDWallpaper() {
+    public void testWillHideAodWallpaper() {
         mScrimController.setWallpaperSupportsAmbientMode(true);
         mScrimController.transitionTo(ScrimState.AOD);
         verify(mAlarmManager).setExact(anyInt(), anyLong(), any(), any(), any());
@@ -540,6 +540,7 @@
 
         private FakeHandler mHandler;
         private boolean mAnimationCancelled;
+        boolean mOnPreDrawCalled;
 
         SynchronousScrimController(LightBarController lightBarController,
                 ScrimView scrimBehind, ScrimView scrimInFront, View headsUpScrim,
@@ -550,6 +551,12 @@
             mHandler = new FakeHandler(Looper.myLooper());
         }
 
+        @Override
+        public boolean onPreDraw() {
+            mOnPreDrawCalled = true;
+            return super.onPreDraw();
+        }
+
         void finishAnimationsImmediately() {
             boolean[] animationFinished = {false};
             setOnAnimationFinished(()-> animationFinished[0] = true);
diff --git a/services/core/java/com/android/server/power/PowerManagerService.java b/services/core/java/com/android/server/power/PowerManagerService.java
index 055e6ea..40a94a7 100644
--- a/services/core/java/com/android/server/power/PowerManagerService.java
+++ b/services/core/java/com/android/server/power/PowerManagerService.java
@@ -69,7 +69,6 @@
 import android.service.vr.IVrManager;
 import android.service.vr.IVrStateCallbacks;
 import android.util.KeyValueListParser;
-import android.util.MathUtils;
 import android.util.PrintWriterPrinter;
 import android.util.Slog;
 import android.util.SparseArray;
@@ -408,7 +407,7 @@
     private boolean mDreamsActivateOnDockSetting;
 
     // True if doze should not be started until after the screen off transition.
-    private boolean mDozeAfterScreenOffConfig;
+    private boolean mDozeAfterScreenOff;
 
     // The minimum screen off timeout, in milliseconds.
     private long mMinimumScreenOffTimeoutConfig;
@@ -896,7 +895,7 @@
                 com.android.internal.R.integer.config_dreamsBatteryLevelMinimumWhenNotPowered);
         mDreamsBatteryLevelDrainCutoffConfig = resources.getInteger(
                 com.android.internal.R.integer.config_dreamsBatteryLevelDrainCutoff);
-        mDozeAfterScreenOffConfig = resources.getBoolean(
+        mDozeAfterScreenOff = resources.getBoolean(
                 com.android.internal.R.bool.config_dozeAfterScreenOff);
         mMinimumScreenOffTimeoutConfig = resources.getInteger(
                 com.android.internal.R.integer.config_minimumScreenOffTimeout);
@@ -2507,7 +2506,7 @@
             if ((mWakeLockSummary & WAKE_LOCK_DOZE) != 0) {
                 return DisplayPowerRequest.POLICY_DOZE;
             }
-            if (mDozeAfterScreenOffConfig) {
+            if (mDozeAfterScreenOff) {
                 return DisplayPowerRequest.POLICY_OFF;
             }
             // Fall through and preserve the current screen policy if not configured to
@@ -3094,6 +3093,12 @@
         light.setFlashing(color, Light.LIGHT_FLASH_HARDWARE, (on ? 3 : 0), 0);
     }
 
+    private void setDozeAfterScreenOffInternal(boolean on) {
+        synchronized (mLock) {
+            mDozeAfterScreenOff = on;
+        }
+    }
+
     private void boostScreenBrightnessInternal(long eventTime, int uid) {
         synchronized (mLock) {
             if (!mSystemReady || mWakefulness == WAKEFULNESS_ASLEEP
@@ -3372,7 +3377,7 @@
             pw.println("  mDreamsEnabledSetting=" + mDreamsEnabledSetting);
             pw.println("  mDreamsActivateOnSleepSetting=" + mDreamsActivateOnSleepSetting);
             pw.println("  mDreamsActivateOnDockSetting=" + mDreamsActivateOnDockSetting);
-            pw.println("  mDozeAfterScreenOffConfig=" + mDozeAfterScreenOffConfig);
+            pw.println("  mDozeAfterScreenOff=" + mDozeAfterScreenOff);
             pw.println("  mLowPowerModeSetting=" + mLowPowerModeSetting);
             pw.println("  mAutoLowPowerModeConfigured=" + mAutoLowPowerModeConfigured);
             pw.println("  mAutoLowPowerModeSnoozing=" + mAutoLowPowerModeSnoozing);
@@ -3656,7 +3661,7 @@
                     mDreamsActivateOnDockSetting);
             proto.write(
                     PowerServiceSettingsAndConfigurationDumpProto.IS_DOZE_AFTER_SCREEN_OFF_CONFIG,
-                    mDozeAfterScreenOffConfig);
+                    mDozeAfterScreenOff);
             proto.write(
                     PowerServiceSettingsAndConfigurationDumpProto.IS_LOW_POWER_MODE_SETTING,
                     mLowPowerModeSetting);
@@ -4603,6 +4608,19 @@
         }
 
         @Override // Binder call
+        public void setDozeAfterScreenOff(boolean on) {
+            mContext.enforceCallingOrSelfPermission(
+                    android.Manifest.permission.DEVICE_POWER, null);
+
+            final long ident = Binder.clearCallingIdentity();
+            try {
+                setDozeAfterScreenOffInternal(on);
+            } finally {
+                Binder.restoreCallingIdentity(ident);
+            }
+        }
+
+        @Override // Binder call
         public void boostScreenBrightness(long eventTime) {
             if (eventTime > SystemClock.uptimeMillis()) {
                 throw new IllegalArgumentException("event time must not be in the future");