Add Screen PowerStatsProcessor

Bug: 333941740
Test: atest PowerStatsTestsRavenwood; atest PowerStatsTests
Flag: com.android.server.power.optimization.streamlined_misc_battery_stats
Change-Id: I49a78d3beb73e1b1b0380834bcd9ddb318186c6b

10 files changed, 862 insertions, 20 deletions

diff --git a/core/java/android/os/BatteryConsumer.java b/core/java/android/os/BatteryConsumer.java
index 2447ff93fdbc..000a537cf3c3 100644
--- a/core/java/android/os/BatteryConsumer.java
+++ b/core/java/android/os/BatteryConsumer.java
@@ -66,6 +66,7 @@ public abstract class BatteryConsumer {
             POWER_COMPONENT_WAKELOCK,
             POWER_COMPONENT_MEMORY,
             POWER_COMPONENT_PHONE,
+            POWER_COMPONENT_AMBIENT_DISPLAY,
             POWER_COMPONENT_IDLE,
             POWER_COMPONENT_REATTRIBUTED_TO_OTHER_CONSUMERS,
     })
diff --git a/services/core/java/com/android/server/am/BatteryStatsService.java b/services/core/java/com/android/server/am/BatteryStatsService.java
index 0f01fe044702..f88644a57a7c 100644
--- a/services/core/java/com/android/server/am/BatteryStatsService.java
+++ b/services/core/java/com/android/server/am/BatteryStatsService.java
@@ -124,6 +124,7 @@ import com.android.server.net.BaseNetworkObserver;
 import com.android.server.pm.UserManagerInternal;
 import com.android.server.power.optimization.Flags;
 import com.android.server.power.stats.AggregatedPowerStatsConfig;
+import com.android.server.power.stats.AmbientDisplayPowerStatsProcessor;
 import com.android.server.power.stats.AudioPowerStatsProcessor;
 import com.android.server.power.stats.BatteryExternalStatsWorker;
 import com.android.server.power.stats.BatteryStatsDumpHelperImpl;
@@ -142,6 +143,7 @@ import com.android.server.power.stats.PowerStatsExporter;
 import com.android.server.power.stats.PowerStatsScheduler;
 import com.android.server.power.stats.PowerStatsStore;
 import com.android.server.power.stats.PowerStatsUidResolver;
+import com.android.server.power.stats.ScreenPowerStatsProcessor;
 import com.android.server.power.stats.SystemServerCpuThreadReader.SystemServiceCpuThreadTimes;
 import com.android.server.power.stats.VideoPowerStatsProcessor;
 import com.android.server.power.stats.WifiPowerStatsProcessor;
@@ -488,6 +490,20 @@ public final class BatteryStatsService extends IBatteryStats.Stub
                 .setProcessor(
                         new CpuPowerStatsProcessor(mPowerProfile, mCpuScalingPolicies));
 
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .trackDeviceStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .trackUidStates(
+                        AggregatedPowerStatsConfig.STATE_POWER,
+                        AggregatedPowerStatsConfig.STATE_SCREEN)
+                .setProcessor(
+                        new ScreenPowerStatsProcessor(mPowerProfile));
+
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .setProcessor(new AmbientDisplayPowerStatsProcessor());
+
         config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO)
                 .trackDeviceStates(
                         AggregatedPowerStatsConfig.STATE_POWER,
@@ -638,6 +654,16 @@ public final class BatteryStatsService extends IBatteryStats.Stub
 
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_SCREEN,
                 Flags.streamlinedMiscBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_SCREEN,
+                Flags.streamlinedMiscBatteryStats());
+
+        mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                Flags.streamlinedMiscBatteryStats());
+        mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
+                BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                Flags.streamlinedMiscBatteryStats());
+
         mStats.setPowerStatsCollectorEnabled(BatteryConsumer.POWER_COMPONENT_MOBILE_RADIO,
                 Flags.streamlinedConnectivityBatteryStats());
         mBatteryUsageStatsProvider.setPowerStatsExporterEnabled(
diff --git a/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java
new file mode 100644
index 000000000000..a42929f6c508
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/AmbientDisplayPowerStatsProcessor.java
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.android.server.power.stats;
+
+import android.os.BatteryConsumer;
+import android.os.PersistableBundle;
+
+import com.android.internal.os.PowerStats;
+
+public class AmbientDisplayPowerStatsProcessor extends PowerStatsProcessor {
+    private final PowerStatsLayout mStatsLayout;
+    private final PowerStats.Descriptor mDescriptor;
+    private final long[] mTmpDeviceStats;
+    private PowerStats.Descriptor mScreenPowerStatsDescriptor;
+    private ScreenPowerStatsLayout mScreenPowerStatsLayout;
+    private long[] mTmpScreenStats;
+
+    public AmbientDisplayPowerStatsProcessor() {
+        mStatsLayout = new PowerStatsLayout();
+        mStatsLayout.addDeviceSectionPowerEstimate();
+        PersistableBundle extras = new PersistableBundle();
+        mStatsLayout.toExtras(extras);
+        mDescriptor = new PowerStats.Descriptor(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                mStatsLayout.getDeviceStatsArrayLength(), null, 0, 0, extras);
+        mTmpDeviceStats = new long[mDescriptor.statsArrayLength];
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+        stats.setPowerStatsDescriptor(mDescriptor);
+
+        PowerComponentAggregatedPowerStats screenStats =
+                stats.getAggregatedPowerStats().getPowerComponentStats(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN);
+        if (screenStats == null) {
+            return;
+        }
+
+        if (mScreenPowerStatsDescriptor == null) {
+            mScreenPowerStatsDescriptor = screenStats.getPowerStatsDescriptor();
+            if (mScreenPowerStatsDescriptor == null) {
+                return;
+            }
+
+            mScreenPowerStatsLayout = new ScreenPowerStatsLayout(mScreenPowerStatsDescriptor);
+            mTmpScreenStats = new long[mScreenPowerStatsDescriptor.statsArrayLength];
+        }
+
+        MultiStateStats.States[] deviceStateConfig = screenStats.getConfig().getDeviceStateConfig();
+
+        // Ambient display power estimates have already been calculated by the screen power stats
+        // processor. All that remains to be done is copy the estimates over.
+        MultiStateStats.States.forEachTrackedStateCombination(deviceStateConfig,
+                states -> {
+                    screenStats.getDeviceStats(mTmpScreenStats, states);
+                    double power =
+                            mScreenPowerStatsLayout.getScreenDozePowerEstimate(mTmpScreenStats);
+                    mStatsLayout.setDevicePowerEstimate(mTmpDeviceStats, power);
+                    stats.setDeviceStats(states, mTmpDeviceStats);
+                });
+    }
+}
diff --git a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
index ac6896696de6..a5e4cf5a0e33 100644
--- a/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
+++ b/services/core/java/com/android/server/power/stats/BatteryUsageStatsProvider.java
@@ -110,8 +110,13 @@ public class BatteryUsageStatsProvider {
                 if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_VIDEO)) {
                     mPowerCalculators.add(new VideoPowerCalculator(mPowerProfile));
                 }
-                mPowerCalculators.add(new ScreenPowerCalculator(mPowerProfile));
-                mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
+                if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_SCREEN)) {
+                    mPowerCalculators.add(new ScreenPowerCalculator(mPowerProfile));
+                }
+                if (!mPowerStatsExporterEnabled.get(
+                        BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY)) {
+                    mPowerCalculators.add(new AmbientDisplayPowerCalculator(mPowerProfile));
+                }
                 mPowerCalculators.add(new IdlePowerCalculator(mPowerProfile));
                 if (!mPowerStatsExporterEnabled.get(BatteryConsumer.POWER_COMPONENT_ANY)) {
                     mPowerCalculators.add(new CustomEnergyConsumerPowerCalculator(mPowerProfile));
diff --git a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
index 8384b2b8db82..6820197fa0f2 100644
--- a/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
+++ b/services/core/java/com/android/server/power/stats/PowerComponentAggregatedPowerStats.java
@@ -163,15 +163,14 @@ class PowerComponentAggregatedPowerStats {
         }
     }
 
-    void setDeviceStats(@AggregatedPowerStatsConfig.TrackedState int[] states, long[] values) {
+    void setDeviceStats(int[] states, long[] values) {
         if (mDeviceStats == null) {
             createDeviceStats(0);
         }
         mDeviceStats.setStats(states, values);
     }
 
-    void setUidStats(int uid, @AggregatedPowerStatsConfig.TrackedState int[] states,
-            long[] values) {
+    void setUidStats(int uid, int[] states, long[] values) {
         UidStats uidStats = getUidStats(uid);
         uidStats.stats.setStats(states, values);
     }
diff --git a/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
index dfc8daa15c37..7d7b3c2fa3c5 100644
--- a/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
+++ b/services/core/java/com/android/server/power/stats/PowerStatsProcessor.java
@@ -220,8 +220,7 @@ public abstract class PowerStatsProcessor {
         }
 
         @Nullable
-        public DeviceStateEstimation getDeviceStateEstimate(
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        public DeviceStateEstimation getDeviceStateEstimate(int[] stateValues) {
             String label = concatLabels(mConfig.getDeviceStateConfig(), stateValues);
             for (int i = 0; i < deviceStateEstimations.size(); i++) {
                 DeviceStateEstimation deviceStateEstimation = this.deviceStateEstimations.get(i);
@@ -233,8 +232,7 @@ public abstract class PowerStatsProcessor {
         }
 
         public CombinedDeviceStateEstimate getCombinedDeviceStateEstimate(
-                MultiStateStats.States[] deviceStates,
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+                MultiStateStats.States[] deviceStates, int[] stateValues) {
             String label = concatLabels(deviceStates, stateValues);
             for (int i = 0; i < combinedDeviceStateEstimations.size(); i++) {
                 CombinedDeviceStateEstimate cdse = combinedDeviceStateEstimations.get(i);
@@ -275,12 +273,10 @@ public abstract class PowerStatsProcessor {
 
     protected static class DeviceStateEstimation {
         public final String id;
-        @AggregatedPowerStatsConfig.TrackedState
         public final int[] stateValues;
         public Object intermediates;
 
-        public DeviceStateEstimation(MultiStateStats.States[] config,
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        public DeviceStateEstimation(MultiStateStats.States[] config, int[] stateValues) {
             id = concatLabels(config, stateValues);
             this.stateValues = stateValues;
         }
@@ -288,11 +284,12 @@ public abstract class PowerStatsProcessor {
 
     protected static class CombinedDeviceStateEstimate {
         public final String id;
+        public final int[] stateValues;
         public List<DeviceStateEstimation> deviceStateEstimations = new ArrayList<>();
         public Object intermediates;
 
-        public CombinedDeviceStateEstimate(MultiStateStats.States[] config,
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        public CombinedDeviceStateEstimate(MultiStateStats.States[] config, int[] stateValues) {
+            this.stateValues = Arrays.copyOf(stateValues, stateValues.length);
             id = concatLabels(config, stateValues);
         }
     }
@@ -310,19 +307,16 @@ public abstract class PowerStatsProcessor {
     }
 
     protected static class UidStateProportionalEstimate {
-        @AggregatedPowerStatsConfig.TrackedState
         public final int[] stateValues;
         public Object intermediates;
 
-        protected UidStateProportionalEstimate(
-                @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+        protected UidStateProportionalEstimate(int[] stateValues) {
             this.stateValues = stateValues;
         }
     }
 
     @NonNull
-    private static String concatLabels(MultiStateStats.States[] config,
-            @AggregatedPowerStatsConfig.TrackedState int[] stateValues) {
+    private static String concatLabels(MultiStateStats.States[] config, int[] stateValues) {
         List<String> labels = new ArrayList<>();
         for (int state = 0; state < config.length; state++) {
             if (config[state] != null && config[state].isTracked()) {
@@ -334,7 +328,6 @@ public abstract class PowerStatsProcessor {
         return labels.toString();
     }
 
-    @AggregatedPowerStatsConfig.TrackedState
     private static int[][] getAllTrackedStateCombinations(MultiStateStats.States[] states) {
         List<int[]> combinations = new ArrayList<>();
         MultiStateStats.States.forEachTrackedStateCombination(states, stateValues -> {
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
index 24fee9ee6bbf..f134aa81057d 100644
--- a/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsLayout.java
@@ -16,9 +16,12 @@
 
 package com.android.server.power.stats;
 
+import android.annotation.NonNull;
 import android.os.BatteryStats;
 import android.os.PersistableBundle;
 
+import com.android.internal.os.PowerStats;
+
 /**
  * Captures the positions and lengths of sections of the stats array, such as time-in-state,
  * power usage estimates etc.
@@ -28,14 +31,23 @@ public class ScreenPowerStatsLayout extends PowerStatsLayout {
     private static final String EXTRA_DEVICE_SCREEN_ON_DURATION_POSITION = "dsd";
     private static final String EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS = "dbd";
     private static final String EXTRA_DEVICE_DOZE_DURATION_POSITION = "ddd";
+    private static final String EXTRA_DEVICE_DOZE_POWER_POSITION = "ddp";
     private static final String EXTRA_UID_FOREGROUND_DURATION = "uf";
 
     private int mDisplayCount;
     private int mDeviceScreenOnDurationPosition;
     private int[] mDeviceBrightnessDurationPositions;
     private int mDeviceScreenDozeDurationPosition;
+    private int mDeviceScreenDozePowerPosition;
     private int mUidTopActivityTimePosition;
 
+    ScreenPowerStatsLayout() {
+    }
+
+    ScreenPowerStatsLayout(@NonNull PowerStats.Descriptor descriptor) {
+        super(descriptor);
+    }
+
     void addDeviceScreenUsageDurationSection(int displayCount) {
         mDisplayCount = displayCount;
         mDeviceScreenOnDurationPosition = addDeviceSection(displayCount, "on");
@@ -47,6 +59,13 @@ public class ScreenPowerStatsLayout extends PowerStatsLayout {
         mDeviceScreenDozeDurationPosition = addDeviceSection(displayCount, "doze");
     }
 
+    @Override
+    public void addDeviceSectionPowerEstimate() {
+        super.addDeviceSectionPowerEstimate();
+        // Used by AmbientDisplayPowerStatsProcessor
+        mDeviceScreenDozePowerPosition = addDeviceSection(1, "doze-power", FLAG_HIDDEN);
+    }
+
     public int getDisplayCount() {
         return mDisplayCount;
     }
@@ -94,6 +113,20 @@ public class ScreenPowerStatsLayout extends PowerStatsLayout {
         return stats[mDeviceScreenDozeDurationPosition + display];
     }
 
+    /**
+     * Stores estimated power in the doze (ambient) state.
+     */
+    public void setScreenDozePowerEstimate(long[] stats, double power) {
+        stats[mDeviceScreenDozePowerPosition] = (long) (power * MILLI_TO_NANO_MULTIPLIER);
+    }
+
+    /**
+     * Retrieves estimated power in the doze (ambient) state.
+     */
+    public double getScreenDozePowerEstimate(long[] stats) {
+        return stats[mDeviceScreenDozePowerPosition] / MILLI_TO_NANO_MULTIPLIER;
+    }
+
     void addUidTopActivitiyDuration() {
         mUidTopActivityTimePosition = addUidSection(1, "top");
     }
@@ -120,6 +153,7 @@ public class ScreenPowerStatsLayout extends PowerStatsLayout {
         extras.putIntArray(EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS,
                 mDeviceBrightnessDurationPositions);
         extras.putInt(EXTRA_DEVICE_DOZE_DURATION_POSITION, mDeviceScreenDozeDurationPosition);
+        extras.putInt(EXTRA_DEVICE_DOZE_POWER_POSITION, mDeviceScreenDozePowerPosition);
         extras.putInt(EXTRA_UID_FOREGROUND_DURATION, mUidTopActivityTimePosition);
     }
 
@@ -131,6 +165,7 @@ public class ScreenPowerStatsLayout extends PowerStatsLayout {
         mDeviceBrightnessDurationPositions = extras.getIntArray(
                 EXTRA_DEVICE_BRIGHTNESS_DURATION_POSITIONS);
         mDeviceScreenDozeDurationPosition = extras.getInt(EXTRA_DEVICE_DOZE_DURATION_POSITION);
+        mDeviceScreenDozePowerPosition = extras.getInt(EXTRA_DEVICE_DOZE_POWER_POSITION);
         mUidTopActivityTimePosition = extras.getInt(EXTRA_UID_FOREGROUND_DURATION);
     }
 }
diff --git a/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java b/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java
new file mode 100644
index 000000000000..e203e4a4175b
--- /dev/null
+++ b/services/core/java/com/android/server/power/stats/ScreenPowerStatsProcessor.java
@@ -0,0 +1,238 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static android.os.BatteryConsumer.PROCESS_STATE_ANY;
+
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_AMBIENT;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL;
+import static com.android.internal.os.PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_PROCESS_STATE;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import android.os.BatteryStats;
+import android.util.Slog;
+
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+
+import java.util.ArrayList;
+import java.util.List;
+
+public class ScreenPowerStatsProcessor extends PowerStatsProcessor {
+    private static final String TAG = "ScreenPowerStatsProcessor";
+    private final int mDisplayCount;
+    private final UsageBasedPowerEstimator[] mScreenOnPowerEstimators;
+    private final UsageBasedPowerEstimator[] mScreenDozePowerEstimators;
+    private final UsageBasedPowerEstimator[][] mScreenBrightnessLevelPowerEstimators;
+    private PowerStats.Descriptor mLastUsedDescriptor;
+    private ScreenPowerStatsLayout mStatsLayout;
+    private PowerEstimationPlan mPlan;
+    private long[] mTmpDeviceStatsArray;
+    private long[] mTmpUidStatsArray;
+
+    private static class Intermediates {
+        public double power;
+    }
+
+    public ScreenPowerStatsProcessor(PowerProfile powerProfile) {
+        mDisplayCount = powerProfile.getNumDisplays();
+        mScreenOnPowerEstimators = new UsageBasedPowerEstimator[mDisplayCount];
+        mScreenDozePowerEstimators = new UsageBasedPowerEstimator[mDisplayCount];
+        mScreenBrightnessLevelPowerEstimators = new UsageBasedPowerEstimator[mDisplayCount][];
+        for (int display = 0; display < mDisplayCount; display++) {
+            mScreenOnPowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, display));
+
+            double averagePowerFullBrightness = powerProfile.getAveragePowerForOrdinal(
+                    POWER_GROUP_DISPLAY_SCREEN_FULL, display);
+            mScreenBrightnessLevelPowerEstimators[display] =
+                    new UsageBasedPowerEstimator[BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS];
+            for (int bin = 0; bin < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; bin++) {
+                // For example, if the number of bins is 3, the corresponding averages
+                // are calculated as 0.5 * full, 1.5 * full, 2.5 * full
+                final double binPowerMah = averagePowerFullBrightness * (bin + 0.5)
+                        / BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS;
+                mScreenBrightnessLevelPowerEstimators[display][bin] =
+                        new UsageBasedPowerEstimator(binPowerMah);
+            }
+
+            mScreenDozePowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, display));
+        }
+    }
+
+    private boolean unpackPowerStatsDescriptor(PowerStats.Descriptor descriptor) {
+        if (descriptor == null) {
+            return false;
+        }
+
+        if (descriptor.equals(mLastUsedDescriptor)) {
+            return true;
+        }
+
+        mLastUsedDescriptor = descriptor;
+        mStatsLayout = new ScreenPowerStatsLayout(descriptor);
+        if (mStatsLayout.getDisplayCount() != mDisplayCount) {
+            Slog.e(TAG, "Incompatible number of displays: " + mStatsLayout.getDisplayCount()
+                    + ", expected: " + mDisplayCount);
+            return false;
+        }
+
+        mTmpDeviceStatsArray = new long[descriptor.statsArrayLength];
+        mTmpUidStatsArray = new long[descriptor.uidStatsArrayLength];
+        return true;
+    }
+
+    @Override
+    void finish(PowerComponentAggregatedPowerStats stats, long timestampMs) {
+        if (!unpackPowerStatsDescriptor(stats.getPowerStatsDescriptor())) {
+            return;
+        }
+
+        if (mPlan == null) {
+            mPlan = new PowerEstimationPlan(stats.getConfig());
+        }
+
+        computeDevicePowerEstimates(stats);
+        combineDeviceStateEstimates();
+
+        List<Integer> uids = new ArrayList<>();
+        stats.collectUids(uids);
+
+        if (!uids.isEmpty()) {
+            computeUidPowerEstimates(stats, uids);
+        }
+    }
+
+    private void computeDevicePowerEstimates(PowerComponentAggregatedPowerStats stats) {
+        for (int i = mPlan.deviceStateEstimations.size() - 1; i >= 0; i--) {
+            DeviceStateEstimation estimation = mPlan.deviceStateEstimations.get(i);
+            if (!stats.getDeviceStats(mTmpDeviceStatsArray, estimation.stateValues)) {
+                continue;
+            }
+
+            if (estimation.stateValues[STATE_SCREEN] == SCREEN_STATE_ON) {
+                double power;
+                if (mStatsLayout.getEnergyConsumerCount() > 0) {
+                    power = uCtoMah(mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, 0));
+                } else {
+                    power = 0;
+                    for (int display = 0; display < mStatsLayout.getDisplayCount(); display++) {
+                        power += computeDisplayPower(mTmpDeviceStatsArray, display);
+                    }
+                }
+                mStatsLayout.setDevicePowerEstimate(mTmpDeviceStatsArray, power);
+                Intermediates intermediates = new Intermediates();
+                intermediates.power = power;
+                estimation.intermediates = intermediates;
+            } else {
+                double power = 0;
+                if (mStatsLayout.getEnergyConsumerCount() > 0) {
+                    power = uCtoMah(mStatsLayout.getConsumedEnergy(mTmpDeviceStatsArray, 0));
+                } else {
+                    for (int display = 0; display < mStatsLayout.getDisplayCount(); display++) {
+                        power += mScreenDozePowerEstimators[display].calculatePower(
+                                mStatsLayout.getScreenDozeDuration(mTmpDeviceStatsArray, display));
+                    }
+                }
+                mStatsLayout.setScreenDozePowerEstimate(mTmpDeviceStatsArray, power);
+            }
+
+            stats.setDeviceStats(estimation.stateValues, mTmpDeviceStatsArray);
+        }
+    }
+
+    private double computeDisplayPower(long[] stats, int display) {
+        double power = mScreenOnPowerEstimators[display]
+                .calculatePower(mStatsLayout.getScreenOnDuration(stats, display));
+        for (int bin = 0; bin < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; bin++) {
+            power += mScreenBrightnessLevelPowerEstimators[display][bin]
+                    .calculatePower(mStatsLayout.getBrightnessLevelDuration(stats, display, bin));
+        }
+        return power;
+    }
+
+    /**
+     * Combine power estimates before distributing them proportionally to UIDs.
+     */
+    private void combineDeviceStateEstimates() {
+        for (int i = mPlan.combinedDeviceStateEstimations.size() - 1; i >= 0; i--) {
+            CombinedDeviceStateEstimate cdse = mPlan.combinedDeviceStateEstimations.get(i);
+            List<DeviceStateEstimation> deviceStateEstimations = cdse.deviceStateEstimations;
+            double power = 0;
+            for (int j = deviceStateEstimations.size() - 1; j >= 0; j--) {
+                DeviceStateEstimation dse = deviceStateEstimations.get(j);
+                Intermediates intermediates = (Intermediates) dse.intermediates;
+                if (intermediates != null) {
+                    power += intermediates.power;
+                }
+            }
+            if (power != 0) {
+                Intermediates cdseIntermediates = new Intermediates();
+                cdseIntermediates.power = power;
+                cdse.intermediates = cdseIntermediates;
+            }
+        }
+    }
+
+    private void computeUidPowerEstimates(PowerComponentAggregatedPowerStats stats,
+            List<Integer> uids) {
+        int[] uidStateValues = new int[stats.getConfig().getUidStateConfig().length];
+        uidStateValues[STATE_SCREEN] = SCREEN_STATE_ON;
+        uidStateValues[STATE_PROCESS_STATE] = PROCESS_STATE_ANY;
+
+        for (int i = mPlan.uidStateEstimates.size() - 1; i >= 0; i--) {
+            UidStateEstimate uidStateEstimate = mPlan.uidStateEstimates.get(i);
+            Intermediates intermediates =
+                    (Intermediates) uidStateEstimate.combinedDeviceStateEstimate.intermediates;
+            int[] deviceStateValues = uidStateEstimate.combinedDeviceStateEstimate
+                    .stateValues;
+            if (deviceStateValues[STATE_SCREEN] != SCREEN_STATE_ON
+                    || intermediates == null) {
+                continue;
+            }
+
+            uidStateValues[STATE_POWER] = deviceStateValues[STATE_POWER];
+
+            long totalTopActivityDuration = 0;
+            for (int j = uids.size() - 1; j >= 0; j--) {
+                int uid = uids.get(j);
+                if (stats.getUidStats(mTmpUidStatsArray, uid, uidStateValues)) {
+                    totalTopActivityDuration +=
+                            mStatsLayout.getUidTopActivityDuration(mTmpUidStatsArray);
+                }
+            }
+
+            if (totalTopActivityDuration == 0) {
+                return;
+            }
+
+            for (int j = uids.size() - 1; j >= 0; j--) {
+                int uid = uids.get(j);
+                if (stats.getUidStats(mTmpUidStatsArray, uid, uidStateValues)) {
+                    long duration = mStatsLayout.getUidTopActivityDuration(mTmpUidStatsArray);
+                    double power = intermediates.power * duration / totalTopActivityDuration;
+                    mStatsLayout.setUidPowerEstimate(mTmpUidStatsArray, power);
+                    stats.setUidStats(uid, uidStateValues, mTmpUidStatsArray);
+                }
+            }
+        }
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java
new file mode 100644
index 000000000000..8d2849bdfe73
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/AmbientDisplayPowerStatsProcessorTest.java
@@ -0,0 +1,183 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.when;
+
+import android.hardware.power.stats.EnergyConsumerType;
+import android.os.BatteryConsumer;
+import android.os.Handler;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.function.IntSupplier;
+
+public class AmbientDisplayPowerStatsProcessorTest {
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setNumDisplays(2)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 0, 180.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 1, 360.0);
+
+    private static final double PRECISION = 0.1;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+
+    private final Injector mInjector = new Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return new PowerStatsUidResolver();
+        }
+
+        @Override
+        public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
+            return 0;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> VOLTAGE_MV;
+        }
+
+        @Override
+        public int getDisplayCount() {
+            return 2;
+        }
+
+        @Override
+        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+            return mScreenUsageTimeRetriever;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @Test
+    public void processPowerStats() {
+        PowerComponentAggregatedPowerStats stats = collectAndAggregatePowerStats();
+
+        assertPowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 16.2);
+        assertPowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 5.4);
+        assertPowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_ON, 0);
+        assertPowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_ON, 0);
+    }
+
+    private PowerComponentAggregatedPowerStats collectAndAggregatePowerStats() {
+        ScreenPowerStatsProcessor screenPowerStatsProcessor =
+                new ScreenPowerStatsProcessor(mStatsRule.getPowerProfile());
+        AmbientDisplayPowerStatsProcessor ambientDisplayPowerStatsProcessor =
+                new AmbientDisplayPowerStatsProcessor();
+
+        AggregatedPowerStatsConfig config = new AggregatedPowerStatsConfig();
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                .trackUidStates(STATE_POWER, STATE_SCREEN)
+                .setProcessor(screenPowerStatsProcessor);
+        config.trackPowerComponent(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY,
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                .setProcessor(ambientDisplayPowerStatsProcessor);
+
+        AggregatedPowerStats stats = new AggregatedPowerStats(config);
+
+        stats.setDeviceState(STATE_POWER, POWER_STATE_OTHER, 0);
+        stats.setDeviceState(STATE_SCREEN, SCREEN_STATE_OTHER, 0);
+
+        ScreenPowerStatsCollector collector = new ScreenPowerStatsCollector(mInjector);
+        collector.setEnabled(true);
+
+        when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                .thenReturn(new int[0]);
+
+        stats.addPowerStats(collector.collectStats(), 1000);
+
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
+                .thenReturn(60_000L);
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(1))
+                .thenReturn(120_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(0))
+                .thenReturn(180_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(1))
+                .thenReturn(240_000L);
+        stats.setDeviceState(STATE_POWER, POWER_STATE_BATTERY, 101_000);
+        stats.setDeviceState(STATE_SCREEN, SCREEN_STATE_ON, 401_000);
+
+        // Slightly larger than 600_000 total screen time, to simulate a sight race
+        // between state changes and power stats collection
+        stats.addPowerStats(collector.collectStats(), 612_000);
+
+        stats.finish(612_000);
+
+        return stats.getPowerComponentStats(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY);
+    }
+
+    private void assertPowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats,
+            int powerState, int screenState, double expectedPowerEstimate) {
+        PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
+        PowerStatsLayout layout = new PowerStatsLayout(descriptor);
+        long[] stats = new long[descriptor.statsArrayLength];
+        aggregatedStats.getDeviceStats(stats, new int[]{powerState, screenState});
+        assertThat(layout.getDevicePowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedPowerEstimate);
+    }
+}
diff --git a/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java
new file mode 100644
index 000000000000..9fde61a51e75
--- /dev/null
+++ b/services/tests/powerstatstests/src/com/android/server/power/stats/ScreenPowerStatsProcessorTest.java
@@ -0,0 +1,287 @@
+/*
+ * Copyright (C) 2024 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.android.server.power.stats;
+
+import static android.os.BatteryConsumer.PROCESS_STATE_ANY;
+
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_BATTERY;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.POWER_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_ON;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.SCREEN_STATE_OTHER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_POWER;
+import static com.android.server.power.stats.AggregatedPowerStatsConfig.STATE_SCREEN;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+
+import android.hardware.power.stats.EnergyConsumerType;
+import android.os.BatteryConsumer;
+import android.os.BatteryStats;
+import android.os.Handler;
+import android.os.Process;
+import android.platform.test.ravenwood.RavenwoodRule;
+
+import com.android.internal.os.Clock;
+import com.android.internal.os.PowerProfile;
+import com.android.internal.os.PowerStats;
+import com.android.server.power.stats.ScreenPowerStatsCollector.Injector;
+
+import org.junit.Before;
+import org.junit.Rule;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.util.function.IntSupplier;
+
+public class ScreenPowerStatsProcessorTest {
+
+    @Rule(order = 0)
+    public final RavenwoodRule mRavenwood = new RavenwoodRule.Builder()
+            .setProvideMainThread(true)
+            .build();
+
+    @Rule(order = 1)
+    public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
+            .setNumDisplays(2)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 0, 180.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_AMBIENT, 1, 360.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON, 0, 480.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON, 1, 720.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_FULL, 0, 4800.0)
+            .setAveragePowerForOrdinal(PowerProfile.POWER_GROUP_DISPLAY_SCREEN_ON, 1, 7200.0)
+            .initMeasuredEnergyStatsLocked();
+
+    private static final double PRECISION = 0.1;
+    private static final int APP_UID1 = Process.FIRST_APPLICATION_UID + 42;
+    private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 101;
+    private static final int VOLTAGE_MV = 3500;
+
+    @Mock
+    private PowerStatsCollector.ConsumedEnergyRetriever mConsumedEnergyRetriever;
+    @Mock
+    private ScreenPowerStatsCollector.ScreenUsageTimeRetriever mScreenUsageTimeRetriever;
+
+    private final Injector mInjector = new Injector() {
+        @Override
+        public Handler getHandler() {
+            return mStatsRule.getHandler();
+        }
+
+        @Override
+        public Clock getClock() {
+            return mStatsRule.getMockClock();
+        }
+
+        @Override
+        public PowerStatsUidResolver getUidResolver() {
+            return new PowerStatsUidResolver();
+        }
+
+        @Override
+        public long getPowerStatsCollectionThrottlePeriod(String powerComponentName) {
+            return 0;
+        }
+
+        @Override
+        public PowerStatsCollector.ConsumedEnergyRetriever getConsumedEnergyRetriever() {
+            return mConsumedEnergyRetriever;
+        }
+
+        @Override
+        public IntSupplier getVoltageSupplier() {
+            return () -> VOLTAGE_MV;
+        }
+
+        @Override
+        public int getDisplayCount() {
+            return 2;
+        }
+
+        @Override
+        public ScreenPowerStatsCollector.ScreenUsageTimeRetriever getScreenUsageTimeRetriever() {
+            return mScreenUsageTimeRetriever;
+        }
+    };
+
+    @Before
+    public void setup() {
+        MockitoAnnotations.initMocks(this);
+    }
+
+    @Test
+    public void processPowerStats_powerProfile() {
+        PowerComponentAggregatedPowerStats stats = collectAndAggregatePowerStats(false);
+
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_ON, 195.5, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0, 0.6);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_ON, 97.8, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0, 0);
+
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_ON, 78.2);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_ON, 39.1);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_ON, 117.3);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_ON, 58.7);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+    }
+
+    @Test
+    public void processPowerStats_energyConsumer() {
+        PowerComponentAggregatedPowerStats stats = collectAndAggregatePowerStats(true);
+
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_ON, 261.9, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0, 7.2);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_ON, 130.9, 0);
+        assertDevicePowerEstimate(stats, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0, 0);
+
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_ON, 104.8);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_ON, 52.4);
+        assertUidPowerEstimate(stats, APP_UID1, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_ON, 157.1);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_BATTERY, SCREEN_STATE_OTHER, 0);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_ON, 78.6);
+        assertUidPowerEstimate(stats, APP_UID2, POWER_STATE_OTHER, SCREEN_STATE_OTHER, 0);
+    }
+
+    private PowerComponentAggregatedPowerStats collectAndAggregatePowerStats(
+            boolean energyConsumer) {
+        ScreenPowerStatsProcessor processor =
+                new ScreenPowerStatsProcessor(mStatsRule.getPowerProfile());
+
+        PowerComponentAggregatedPowerStats aggregatedStats = createAggregatedPowerStats(processor);
+
+        ScreenPowerStatsCollector collector = new ScreenPowerStatsCollector(mInjector);
+        collector.setEnabled(true);
+
+        if (energyConsumer) {
+            when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                    .thenReturn(new int[]{77});
+            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
+                    .thenReturn(new long[]{10_000});
+        } else {
+            when(mConsumedEnergyRetriever.getEnergyConsumerIds(EnergyConsumerType.DISPLAY))
+                    .thenReturn(new int[0]);
+        }
+
+        doAnswer(inv -> {
+            ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
+                    inv.getArgument(0);
+            callback.onUidTopActivityTime(APP_UID1, 1000);
+            callback.onUidTopActivityTime(APP_UID2, 2000);
+            return null;
+        }).when(mScreenUsageTimeRetriever).retrieveTopActivityTimes(any(
+                ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback.class));
+
+        aggregatedStats.addPowerStats(collector.collectStats(), 1000);
+
+        if (energyConsumer) {
+            // 400 mAh represented as microWattSeconds
+            long energyUws = 400L * 3600 * VOLTAGE_MV;
+            when(mConsumedEnergyRetriever.getConsumedEnergyUws(new int[]{77}))
+                    .thenReturn(new long[]{10_000 + energyUws});
+        }
+
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(0))
+                .thenReturn(60_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_DARK))
+                .thenReturn(10_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_MEDIUM))
+                .thenReturn(20_000L);
+        when(mScreenUsageTimeRetriever.getBrightnessLevelTimeMs(0,
+                BatteryStats.SCREEN_BRIGHTNESS_BRIGHT))
+                .thenReturn(30_000L);
+        when(mScreenUsageTimeRetriever.getScreenOnTimeMs(1))
+                .thenReturn(120_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(0))
+                .thenReturn(180_000L);
+        when(mScreenUsageTimeRetriever.getScreenDozeTimeMs(1))
+                .thenReturn(240_000L);
+        doAnswer(inv -> {
+            ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback callback =
+                    inv.getArgument(0);
+            callback.onUidTopActivityTime(APP_UID1, 3000);
+            callback.onUidTopActivityTime(APP_UID2, 5000);
+            return null;
+        }).when(mScreenUsageTimeRetriever).retrieveTopActivityTimes(any(
+                ScreenPowerStatsCollector.ScreenUsageTimeRetriever.Callback.class));
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_BATTERY, 201_000);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_OTHER, 601_000);
+
+        // Slightly larger than 600_000 total screen time, to simulate a sight race
+        // between state changes and power stats collection
+        aggregatedStats.addPowerStats(collector.collectStats(), 612_000);
+
+        aggregatedStats.getConfig().getProcessor().finish(aggregatedStats, 180_000);
+        return aggregatedStats;
+    }
+
+    private static PowerComponentAggregatedPowerStats createAggregatedPowerStats(
+            ScreenPowerStatsProcessor processor) {
+        AggregatedPowerStatsConfig.PowerComponent config =
+                new AggregatedPowerStatsConfig.PowerComponent(
+                        BatteryConsumer.POWER_COMPONENT_SCREEN)
+                        .trackDeviceStates(STATE_POWER, STATE_SCREEN)
+                        .trackUidStates(STATE_POWER, STATE_SCREEN)
+                        .setProcessor(processor);
+
+        PowerComponentAggregatedPowerStats aggregatedStats =
+                new PowerComponentAggregatedPowerStats(
+                        new AggregatedPowerStats(mock(AggregatedPowerStatsConfig.class)), config);
+
+        aggregatedStats.setState(STATE_POWER, POWER_STATE_OTHER, 0);
+        aggregatedStats.setState(STATE_SCREEN, SCREEN_STATE_ON, 0);
+
+        return aggregatedStats;
+    }
+
+    private void assertDevicePowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats,
+            int powerState, int screenState, double expectedScreenPowerEstimate,
+            double expectedDozePowerEstimate) {
+        PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
+        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout(descriptor);
+        long[] stats = new long[descriptor.statsArrayLength];
+        aggregatedStats.getDeviceStats(stats, new int[]{powerState, screenState});
+        assertThat(layout.getDevicePowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedScreenPowerEstimate);
+        assertThat(layout.getScreenDozePowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedDozePowerEstimate);
+    }
+
+    private void assertUidPowerEstimate(PowerComponentAggregatedPowerStats aggregatedStats, int uid,
+            int powerState, int screenState, double expectedScreenPowerEstimate) {
+        PowerStats.Descriptor descriptor = aggregatedStats.getPowerStatsDescriptor();
+        ScreenPowerStatsLayout layout = new ScreenPowerStatsLayout(descriptor);
+        long[] stats = new long[descriptor.uidStatsArrayLength];
+        aggregatedStats.getUidStats(stats, uid,
+                new int[]{powerState, screenState, PROCESS_STATE_ANY});
+        assertThat(layout.getUidPowerEstimate(stats)).isWithin(PRECISION)
+                .of(expectedScreenPowerEstimate);
+    }
+}