Merge changes I504d9c65,I45abcf87 into sc-v2-dev

* changes:
  Utilize measured energy for multidisplay power attribution
  Update Screen and AmbientDisplay Power Calculators for multi-display

9 files changed, 451 insertions, 107 deletions

diff --git a/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java b/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java
index 93baa193570d..94430704468f 100644
--- a/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java
+++ b/core/java/com/android/internal/os/AmbientDisplayPowerCalculator.java
@@ -31,12 +31,15 @@ import java.util.List;
  * Estimates power consumed by the ambient display
  */
 public class AmbientDisplayPowerCalculator extends PowerCalculator {
-    private final UsageBasedPowerEstimator mPowerEstimator;
+    private final UsageBasedPowerEstimator[] mPowerEstimators;
 
     public AmbientDisplayPowerCalculator(PowerProfile powerProfile) {
-        // TODO(b/200239964): update to support multidisplay.
-        mPowerEstimator = new UsageBasedPowerEstimator(
-                powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 0));
+        final int numDisplays = powerProfile.getNumDisplays();
+        mPowerEstimators = new UsageBasedPowerEstimator[numDisplays];
+        for (int display = 0; display < numDisplays; display++) {
+            mPowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, display));
+        }
     }
 
     /**
@@ -50,8 +53,8 @@ public class AmbientDisplayPowerCalculator extends PowerCalculator {
         final int powerModel = getPowerModel(measuredEnergyUC, query);
         final long durationMs = calculateDuration(batteryStats, rawRealtimeUs,
                 BatteryStats.STATS_SINCE_CHARGED);
-        final double powerMah = getMeasuredOrEstimatedPower(powerModel,
-                measuredEnergyUC, mPowerEstimator, durationMs);
+        final double powerMah = calculateTotalPower(powerModel, batteryStats, rawRealtimeUs,
+                measuredEnergyUC);
         builder.getAggregateBatteryConsumerBuilder(
                 BatteryUsageStats.AGGREGATE_BATTERY_CONSUMER_SCOPE_DEVICE)
                 .setUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY, durationMs)
@@ -71,9 +74,8 @@ public class AmbientDisplayPowerCalculator extends PowerCalculator {
         final long measuredEnergyUC = batteryStats.getScreenDozeMeasuredBatteryConsumptionUC();
         final long durationMs = calculateDuration(batteryStats, rawRealtimeUs, statsType);
         final int powerModel = getPowerModel(measuredEnergyUC);
-        final double powerMah = getMeasuredOrEstimatedPower(powerModel,
-                batteryStats.getScreenDozeMeasuredBatteryConsumptionUC(),
-                mPowerEstimator, durationMs);
+        final double powerMah = calculateTotalPower(powerModel, batteryStats, rawRealtimeUs,
+                measuredEnergyUC);
         if (powerMah > 0) {
             BatterySipper bs = new BatterySipper(BatterySipper.DrainType.AMBIENT_DISPLAY, null, 0);
             bs.usagePowerMah = powerMah;
@@ -86,4 +88,26 @@ public class AmbientDisplayPowerCalculator extends PowerCalculator {
     private long calculateDuration(BatteryStats batteryStats, long rawRealtimeUs, int statsType) {
         return batteryStats.getScreenDozeTime(rawRealtimeUs, statsType) / 1000;
     }
+
+    private double calculateTotalPower(@BatteryConsumer.PowerModel int powerModel,
+            BatteryStats batteryStats, long rawRealtimeUs, long consumptionUC) {
+        switch (powerModel) {
+            case BatteryConsumer.POWER_MODEL_MEASURED_ENERGY:
+                return uCtoMah(consumptionUC);
+            case BatteryConsumer.POWER_MODEL_POWER_PROFILE:
+            default:
+                return calculateEstimatedPower(batteryStats, rawRealtimeUs);
+        }
+    }
+
+    private double calculateEstimatedPower(BatteryStats batteryStats, long rawRealtimeUs) {
+        final int numDisplays = mPowerEstimators.length;
+        double power = 0;
+        for (int display = 0; display < numDisplays; display++) {
+            final long dozeTime = batteryStats.getDisplayScreenDozeTime(display, rawRealtimeUs)
+                    / 1000;
+            power += mPowerEstimators[display].calculatePower(dozeTime);
+        }
+        return power;
+    }
 }
diff --git a/core/java/com/android/internal/os/BatteryStatsImpl.java b/core/java/com/android/internal/os/BatteryStatsImpl.java
index 17fdb2e7054f..169eff009bff 100644
--- a/core/java/com/android/internal/os/BatteryStatsImpl.java
+++ b/core/java/com/android/internal/os/BatteryStatsImpl.java
@@ -903,6 +903,11 @@ public class BatteryStatsImpl extends BatteryStats {
          */
         public StopwatchTimer[] screenBrightnessTimers =
                 new StopwatchTimer[NUM_SCREEN_BRIGHTNESS_BINS];
+        /**
+         * Per display screen state the last time {@link #updateDisplayMeasuredEnergyStatsLocked}
+         * was called.
+         */
+        public int screenStateAtLastEnergyMeasurement = Display.STATE_UNKNOWN;
 
         DisplayBatteryStats(Clocks clocks, TimeBase timeBase) {
             screenOnTimer = new StopwatchTimer(clocks, null, -1, null,
@@ -929,6 +934,8 @@ public class BatteryStatsImpl extends BatteryStats {
 
     DisplayBatteryStats[] mPerDisplayBatteryStats;
 
+    private int mDisplayMismatchWtfCount = 0;
+
     boolean mInteractive;
     StopwatchTimer mInteractiveTimer;
 
@@ -1073,8 +1080,6 @@ public class BatteryStatsImpl extends BatteryStats {
     @GuardedBy("this")
     @VisibleForTesting
     protected @Nullable MeasuredEnergyStats mGlobalMeasuredEnergyStats;
-    /** Last known screen state. Needed for apportioning display energy. */
-    int mScreenStateAtLastEnergyMeasurement = Display.STATE_UNKNOWN;
     /** Bluetooth Power calculator for attributing measured bluetooth charge consumption to uids */
     @Nullable BluetoothPowerCalculator mBluetoothPowerCalculator = null;
     /** Cpu Power calculator for attributing measured cpu charge consumption to uids */
@@ -12924,22 +12929,43 @@ public class BatteryStatsImpl extends BatteryStats {
      * is always 0 when the screen is not "ON" and whenever the rail energy is 0 (if supported).
      * To the extent that those assumptions are violated, the algorithm will err.
      *
-     * @param chargeUC amount of charge (microcoulombs) used by Display since this was last called.
-     * @param screenState screen state at the time this data collection was scheduled
+     * @param chargesUC amount of charge (microcoulombs) used by each Display since this was last
+     *                 called.
+     * @param screenStates each screen state at the time this data collection was scheduled
      */
     @GuardedBy("this")
-    public void updateDisplayMeasuredEnergyStatsLocked(long chargeUC, int screenState,
+    public void updateDisplayMeasuredEnergyStatsLocked(long[] chargesUC, int[] screenStates,
             long elapsedRealtimeMs) {
-        if (DEBUG_ENERGY) Slog.d(TAG, "Updating display stats: " + chargeUC);
+        if (DEBUG_ENERGY) Slog.d(TAG, "Updating display stats: " + Arrays.toString(chargesUC));
         if (mGlobalMeasuredEnergyStats == null) {
             return;
         }
 
-        final @StandardPowerBucket int powerBucket =
-                MeasuredEnergyStats.getDisplayPowerBucket(mScreenStateAtLastEnergyMeasurement);
-        mScreenStateAtLastEnergyMeasurement = screenState;
+        final int numDisplays;
+        if (mPerDisplayBatteryStats.length == screenStates.length) {
+            numDisplays = screenStates.length;
+        } else {
+            // if this point is reached, it will be reached every display state change.
+            // Rate limit the wtf logging to once every 100 display updates.
+            if (mDisplayMismatchWtfCount++ % 100 == 0) {
+                Slog.wtf(TAG, "Mismatch between PowerProfile reported display count ("
+                        + mPerDisplayBatteryStats.length
+                        + ") and PowerStatsHal reported display count (" + screenStates.length
+                        + ")");
+            }
+            // Keep the show going, use the shorter of the two.
+            numDisplays = mPerDisplayBatteryStats.length < screenStates.length
+                    ? mPerDisplayBatteryStats.length : screenStates.length;
+        }
+
+        final int[] oldScreenStates = new int[numDisplays];
+        for (int i = 0; i < numDisplays; i++) {
+            final int screenState = screenStates[i];
+            oldScreenStates[i] = mPerDisplayBatteryStats[i].screenStateAtLastEnergyMeasurement;
+            mPerDisplayBatteryStats[i].screenStateAtLastEnergyMeasurement = screenState;
+        }
 
-        if (!mOnBatteryInternal || chargeUC <= 0) {
+        if (!mOnBatteryInternal) {
             // There's nothing further to update.
             return;
         }
@@ -12954,17 +12980,31 @@ public class BatteryStatsImpl extends BatteryStats {
             return;
         }
 
-        mGlobalMeasuredEnergyStats.updateStandardBucket(powerBucket, chargeUC);
+        long totalScreenOnChargeUC = 0;
+        for (int i = 0; i < numDisplays; i++) {
+            final long chargeUC = chargesUC[i];
+            if (chargeUC <= 0) {
+                // There's nothing further to update.
+                continue;
+            }
+
+            final @StandardPowerBucket int powerBucket =
+                    MeasuredEnergyStats.getDisplayPowerBucket(oldScreenStates[i]);
+            mGlobalMeasuredEnergyStats.updateStandardBucket(powerBucket, chargeUC);
+            if (powerBucket == MeasuredEnergyStats.POWER_BUCKET_SCREEN_ON) {
+                totalScreenOnChargeUC += chargeUC;
+            }
+        }
 
         // Now we blame individual apps, but only if the display was ON.
-        if (powerBucket != MeasuredEnergyStats.POWER_BUCKET_SCREEN_ON) {
+        if (totalScreenOnChargeUC <= 0) {
             return;
         }
         // TODO(b/175726779): Consider unifying the code with the non-rail display power blaming.
 
         // NOTE: fg time is NOT pooled. If two uids are both somehow in fg, then that time is
         // 'double counted' and will simply exceed the realtime that elapsed.
-        // If multidisplay becomes a reality, this is probably more reasonable than pooling.
+        // TODO(b/175726779): collect per display uid visibility for display power attribution.
 
         // Collect total time since mark so that we can normalize power.
         final SparseDoubleArray fgTimeUsArray = new SparseDoubleArray();
@@ -12977,7 +13017,8 @@ public class BatteryStatsImpl extends BatteryStats {
             if (fgTimeUs == 0) continue;
             fgTimeUsArray.put(uid.getUid(), (double) fgTimeUs);
         }
-        distributeEnergyToUidsLocked(powerBucket, chargeUC, fgTimeUsArray, 0);
+        distributeEnergyToUidsLocked(MeasuredEnergyStats.POWER_BUCKET_SCREEN_ON,
+                totalScreenOnChargeUC, fgTimeUsArray, 0);
     }
 
     /**
@@ -14883,7 +14924,12 @@ public class BatteryStatsImpl extends BatteryStats {
     public void initMeasuredEnergyStatsLocked(@Nullable boolean[] supportedStandardBuckets,
             String[] customBucketNames) {
         boolean supportedBucketMismatch = false;
-        mScreenStateAtLastEnergyMeasurement = mScreenState;
+
+        final int numDisplays = mPerDisplayBatteryStats.length;
+        for (int i = 0; i < numDisplays; i++) {
+            final int screenState = mPerDisplayBatteryStats[i].screenState;
+            mPerDisplayBatteryStats[i].screenStateAtLastEnergyMeasurement = screenState;
+        }
 
         if (supportedStandardBuckets == null) {
             if (mGlobalMeasuredEnergyStats != null) {
diff --git a/core/java/com/android/internal/os/PowerCalculator.java b/core/java/com/android/internal/os/PowerCalculator.java
index 4979ecbae8cb..93d562c571f8 100644
--- a/core/java/com/android/internal/os/PowerCalculator.java
+++ b/core/java/com/android/internal/os/PowerCalculator.java
@@ -133,32 +133,6 @@ public abstract class PowerCalculator {
     }
 
     /**
-     * Returns either the measured energy converted to mAh or a usage-based estimate.
-     */
-    protected static double getMeasuredOrEstimatedPower(@BatteryConsumer.PowerModel int powerModel,
-            long measuredEnergyUC, UsageBasedPowerEstimator powerEstimator, long durationMs) {
-        switch (powerModel) {
-            case BatteryConsumer.POWER_MODEL_MEASURED_ENERGY:
-                return uCtoMah(measuredEnergyUC);
-            case BatteryConsumer.POWER_MODEL_POWER_PROFILE:
-            default:
-                return powerEstimator.calculatePower(durationMs);
-        }
-    }
-
-    /**
-     * Returns either the measured energy converted to mAh or a usage-based estimate.
-     */
-    protected static double getMeasuredOrEstimatedPower(
-            long measuredEnergyUC, UsageBasedPowerEstimator powerEstimator, long durationMs) {
-        if (measuredEnergyUC != BatteryStats.POWER_DATA_UNAVAILABLE) {
-            return uCtoMah(measuredEnergyUC);
-        } else {
-            return powerEstimator.calculatePower(durationMs);
-        }
-    }
-
-    /**
      * Prints formatted amount of power in milli-amp-hours.
      */
     public static void printPowerMah(PrintWriter pw, double powerMah) {
diff --git a/core/java/com/android/internal/os/ScreenPowerCalculator.java b/core/java/com/android/internal/os/ScreenPowerCalculator.java
index 72ad4e72707a..2b634598bbbc 100644
--- a/core/java/com/android/internal/os/ScreenPowerCalculator.java
+++ b/core/java/com/android/internal/os/ScreenPowerCalculator.java
@@ -44,8 +44,8 @@ public class ScreenPowerCalculator extends PowerCalculator {
     // Minimum amount of time the screen should be on to start smearing drain to apps
     public static final long MIN_ACTIVE_TIME_FOR_SMEARING = 10 * DateUtils.MINUTE_IN_MILLIS;
 
-    private final UsageBasedPowerEstimator mScreenOnPowerEstimator;
-    private final UsageBasedPowerEstimator mScreenFullPowerEstimator;
+    private final UsageBasedPowerEstimator[] mScreenOnPowerEstimators;
+    private final UsageBasedPowerEstimator[] mScreenFullPowerEstimators;
 
     private static class PowerAndDuration {
         public long durationMs;
@@ -53,11 +53,16 @@ public class ScreenPowerCalculator extends PowerCalculator {
     }
 
     public ScreenPowerCalculator(PowerProfile powerProfile) {
-        // TODO(b/200239964): update to support multidisplay.
-        mScreenOnPowerEstimator = new UsageBasedPowerEstimator(
-                powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 0));
-        mScreenFullPowerEstimator = new UsageBasedPowerEstimator(
-                powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 0));
+        final int numDisplays = powerProfile.getNumDisplays();
+        mScreenOnPowerEstimators = new UsageBasedPowerEstimator[numDisplays];
+        mScreenFullPowerEstimators = new UsageBasedPowerEstimator[numDisplays];
+        for (int display = 0; display < numDisplays; display++) {
+            mScreenOnPowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, display));
+            mScreenFullPowerEstimators[display] = new UsageBasedPowerEstimator(
+                    powerProfile.getAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL,
+                            display));
+        }
     }
 
     @Override
@@ -172,7 +177,7 @@ public class ScreenPowerCalculator extends PowerCalculator {
             case BatteryConsumer.POWER_MODEL_POWER_PROFILE:
             default:
                 totalPowerAndDuration.powerMah = calculateTotalPowerFromBrightness(batteryStats,
-                        rawRealtimeUs, statsType, totalPowerAndDuration.durationMs);
+                        rawRealtimeUs);
         }
     }
 
@@ -194,19 +199,25 @@ public class ScreenPowerCalculator extends PowerCalculator {
         return batteryStats.getScreenOnTime(rawRealtimeUs, statsType) / 1000;
     }
 
-    private double calculateTotalPowerFromBrightness(BatteryStats batteryStats, long rawRealtimeUs,
-            int statsType, long durationMs) {
-        double power = mScreenOnPowerEstimator.calculatePower(durationMs);
-        for (int i = 0; i < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; i++) {
-            final long brightnessTime =
-                    batteryStats.getScreenBrightnessTime(i, rawRealtimeUs, statsType) / 1000;
-            final double binPowerMah = mScreenFullPowerEstimator.calculatePower(brightnessTime)
-                    * (i + 0.5f) / BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS;
-            if (DEBUG && binPowerMah != 0) {
-                Slog.d(TAG, "Screen bin #" + i + ": time=" + brightnessTime
-                        + " power=" + formatCharge(binPowerMah));
+    private double calculateTotalPowerFromBrightness(BatteryStats batteryStats,
+            long rawRealtimeUs) {
+        final int numDisplays = mScreenOnPowerEstimators.length;
+        double power = 0;
+        for (int display = 0; display < numDisplays; display++) {
+            final long displayTime = batteryStats.getDisplayScreenOnTime(display, rawRealtimeUs)
+                    / 1000;
+            power += mScreenOnPowerEstimators[display].calculatePower(displayTime);
+            for (int bin = 0; bin < BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS; bin++) {
+                final long brightnessTime = batteryStats.getDisplayScreenBrightnessTime(display,
+                        bin, rawRealtimeUs) / 1000;
+                final double binPowerMah = mScreenFullPowerEstimators[display].calculatePower(
+                        brightnessTime) * (bin + 0.5f) / BatteryStats.NUM_SCREEN_BRIGHTNESS_BINS;
+                if (DEBUG && binPowerMah != 0) {
+                    Slog.d(TAG, "Screen bin #" + bin + ": time=" + brightnessTime
+                            + " power=" + formatCharge(binPowerMah));
+                }
+                power += binPowerMah;
             }
-            power += binPowerMah;
         }
         return power;
     }
diff --git a/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
index e95f6c294633..130f552f6e3a 100644
--- a/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/AmbientDisplayPowerCalculatorTest.java
@@ -38,26 +38,28 @@ public class AmbientDisplayPowerCalculatorTest {
 
     @Rule
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
-            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 0, 10.0);
+            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 0, 10.0)
+            .setNumDisplays(1);
 
     @Test
     public void testMeasuredEnergyBasedModel() {
         mStatsRule.initMeasuredEnergyStatsLocked();
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
-        stats.updateDisplayMeasuredEnergyStatsLocked(300_000_000, Display.STATE_ON, 0);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300_000_000},
+                new int[]{Display.STATE_ON}, 0);
 
         stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
                 30 * MINUTE_IN_MS);
 
-        stats.updateDisplayMeasuredEnergyStatsLocked(200_000_000, Display.STATE_DOZE,
-                30 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{200_000_000},
+                new int[]{Display.STATE_DOZE}, 30 * MINUTE_IN_MS);
 
         stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
                 120 * MINUTE_IN_MS);
 
-        stats.updateDisplayMeasuredEnergyStatsLocked(100_000_000, Display.STATE_OFF,
-                120 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100_000_000},
+                new int[]{Display.STATE_OFF}, 120 * MINUTE_IN_MS);
 
         AmbientDisplayPowerCalculator calculator =
                 new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
@@ -75,6 +77,67 @@ public class AmbientDisplayPowerCalculatorTest {
     }
 
     @Test
+    public void testMeasuredEnergyBasedModel_multiDisplay() {
+        mStatsRule.initMeasuredEnergyStatsLocked()
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 1, 20.0)
+                .setNumDisplays(2);
+        BatteryStatsImpl stats = mStatsRule.getBatteryStats();
+
+
+        final int[] screenStates = new int[] {Display.STATE_OFF, Display.STATE_OFF};
+
+        stats.noteScreenStateLocked(0, screenStates[0], 0, 0, 0);
+        stats.noteScreenStateLocked(1, screenStates[1], 0, 0, 0);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300, 400}, screenStates, 0);
+
+        // Switch display0 to doze
+        screenStates[0] = Display.STATE_DOZE;
+        stats.noteScreenStateLocked(0, screenStates[0], 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
+                30 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{200, 300},
+                screenStates, 30 * MINUTE_IN_MS);
+
+        // Switch display1 to doze
+        screenStates[1] = Display.STATE_DOZE;
+        stats.noteScreenStateLocked(1, Display.STATE_DOZE, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS,
+                90 * MINUTE_IN_MS);
+        // 100,000,000 uC should be attributed to display 0 doze here.
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100_000_000, 700_000_000},
+                screenStates, 90 * MINUTE_IN_MS);
+
+        // Switch display0 to off
+        screenStates[0] = Display.STATE_OFF;
+        stats.noteScreenStateLocked(0, screenStates[0], 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
+                120 * MINUTE_IN_MS);
+        // 40,000,000 and 70,000,000 uC should be attributed to display 0 and 1 doze here.
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{40_000_000, 70_000_000},
+                screenStates, 120 * MINUTE_IN_MS);
+
+        // Switch display1 to off
+        screenStates[1] = Display.STATE_OFF;
+        stats.noteScreenStateLocked(1, screenStates[1], 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS,
+                150 * MINUTE_IN_MS);
+        stats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100, 90_000_000}, screenStates,
+                150 * MINUTE_IN_MS);
+        // 90,000,000 uC should be attributed to display 1 doze here.
+
+        AmbientDisplayPowerCalculator calculator =
+                new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(calculator);
+
+        BatteryConsumer consumer = mStatsRule.getDeviceBatteryConsumer();
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(120 * MINUTE_IN_MS);
+        // 100,000,000 + 40,000,000 + 70,000,000 + 90,000,000 uC / 1000 (micro-/milli-) / 3600
+        // (seconds/hour) = 27.777778 mAh
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isWithin(PRECISION).of(83.33333);
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+    }
+
+    @Test
     public void testPowerProfileBasedModel() {
         BatteryStatsImpl stats = mStatsRule.getBatteryStats();
 
@@ -96,4 +159,36 @@ public class AmbientDisplayPowerCalculatorTest {
         assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
     }
+
+    @Test
+    public void testPowerProfileBasedModel_multiDisplay() {
+        mStatsRule.setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_AMBIENT, 1, 20.0)
+                .setNumDisplays(2);
+
+        BatteryStatsImpl stats = mStatsRule.getBatteryStats();
+
+        stats.noteScreenStateLocked(1, Display.STATE_OFF, 0, 0, 0);
+        stats.noteScreenStateLocked(0, Display.STATE_DOZE, 30 * MINUTE_IN_MS, 30 * MINUTE_IN_MS,
+                30 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_DOZE, 90 * MINUTE_IN_MS, 90 * MINUTE_IN_MS,
+                90 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(0, Display.STATE_OFF, 120 * MINUTE_IN_MS, 120 * MINUTE_IN_MS,
+                120 * MINUTE_IN_MS);
+        stats.noteScreenStateLocked(1, Display.STATE_OFF, 150 * MINUTE_IN_MS, 150 * MINUTE_IN_MS,
+                150 * MINUTE_IN_MS);
+
+        AmbientDisplayPowerCalculator calculator =
+                new AmbientDisplayPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(BatteryUsageStatsRule.POWER_PROFILE_MODEL_ONLY, calculator);
+
+        BatteryConsumer consumer = mStatsRule.getDeviceBatteryConsumer();
+        // Duration should only be the union of
+        assertThat(consumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(120 * MINUTE_IN_MS);
+        assertThat(consumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isWithin(PRECISION).of(35.0);
+        assertThat(consumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_AMBIENT_DISPLAY))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+    }
 }
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
index 4adc09d72e94..3e2885a74287 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryStatsNoteTest.java
@@ -998,7 +998,7 @@ public class BatteryStatsNoteTest extends TestCase {
         bi.initMeasuredEnergyStats(new String[]{"FOO", "BAR"});
 
         clocks.realtime = 0;
-        int screen = Display.STATE_OFF;
+        int[] screen = new int[]{Display.STATE_OFF};
         boolean battery = false;
 
         final int uid1 = 10500;
@@ -1008,35 +1008,35 @@ public class BatteryStatsNoteTest extends TestCase {
         long globalDoze = 0;
 
         // Case A: uid1 off, uid2 off, battery off, screen off
-        bi.updateTimeBasesLocked(battery, screen, clocks.realtime*1000, 0);
+        bi.updateTimeBasesLocked(battery, screen[0], clocks.realtime * 1000, 0);
         bi.setOnBatteryInternal(battery);
-        bi.updateDisplayMeasuredEnergyStatsLocked(500_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{500_000}, screen, clocks.realtime);
         checkMeasuredCharge("A", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case B: uid1 off, uid2 off, battery ON,  screen off
         clocks.realtime += 17;
         battery = true;
-        bi.updateTimeBasesLocked(battery, screen, clocks.realtime*1000, 0);
+        bi.updateTimeBasesLocked(battery, screen[0], clocks.realtime * 1000, 0);
         bi.setOnBatteryInternal(battery);
         clocks.realtime += 19;
-        bi.updateDisplayMeasuredEnergyStatsLocked(510_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{510_000}, screen, clocks.realtime);
         checkMeasuredCharge("B", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case C: uid1 ON,  uid2 off, battery on,  screen off
         clocks.realtime += 18;
         setFgState(uid1, true, bi);
         clocks.realtime += 18;
-        bi.updateDisplayMeasuredEnergyStatsLocked(520_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{520_000}, screen, clocks.realtime);
         checkMeasuredCharge("C", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case D: uid1 on,  uid2 off, battery on,  screen ON
         clocks.realtime += 17;
-        screen = Display.STATE_ON;
-        bi.updateDisplayMeasuredEnergyStatsLocked(521_000, screen, clocks.realtime);
+        screen[0] = Display.STATE_ON;
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{521_000}, screen, clocks.realtime);
         blame1 += 0; // Screen had been off during the measurement period
         checkMeasuredCharge("D.1", uid1, blame1, uid2, blame2, globalDoze, bi);
         clocks.realtime += 101;
-        bi.updateDisplayMeasuredEnergyStatsLocked(530_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{530_000}, screen, clocks.realtime);
         blame1 += 530_000;
         checkMeasuredCharge("D.2", uid1, blame1, uid2, blame2, globalDoze, bi);
 
@@ -1044,33 +1044,33 @@ public class BatteryStatsNoteTest extends TestCase {
         clocks.realtime += 20;
         setFgState(uid2, true, bi);
         clocks.realtime += 40;
-        bi.updateDisplayMeasuredEnergyStatsLocked(540_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{540_000}, screen, clocks.realtime);
         // In the past 60ms, sum of fg is 20+40+40=100ms. uid1 is blamed for 60/100; uid2 for 40/100
         blame1 += 540_000 * (20 + 40) / (20 + 40 + 40);
-        blame2 += 540_000 * ( 0 + 40) / (20 + 40 + 40);
+        blame2 += 540_000 * (0 + 40) / (20 + 40 + 40);
         checkMeasuredCharge("E", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case F: uid1 on,  uid2 OFF, battery on,  screen on
         clocks.realtime += 40;
         setFgState(uid2, false, bi);
         clocks.realtime += 120;
-        bi.updateDisplayMeasuredEnergyStatsLocked(550_000, screen, clocks.realtime);
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{550_000}, screen, clocks.realtime);
         // In the past 160ms, sum f fg is 200ms. uid1 is blamed for 40+120 of it; uid2 for 40 of it.
         blame1 += 550_000 * (40 + 120) / (40 + 40 + 120);
-        blame2 += 550_000 * (40 + 0  ) / (40 + 40 + 120);
+        blame2 += 550_000 * (40 + 0) / (40 + 40 + 120);
         checkMeasuredCharge("F", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case G: uid1 on,  uid2 off,  battery on, screen DOZE
         clocks.realtime += 5;
-        screen = Display.STATE_DOZE;
-        bi.updateDisplayMeasuredEnergyStatsLocked(570_000, screen, clocks.realtime);
+        screen[0] = Display.STATE_DOZE;
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{570_000}, screen, clocks.realtime);
         blame1 += 570_000; // All of this pre-doze time is blamed on uid1.
         checkMeasuredCharge("G", uid1, blame1, uid2, blame2, globalDoze, bi);
 
         // Case H: uid1 on,  uid2 off,  battery on, screen ON
         clocks.realtime += 6;
-        screen = Display.STATE_ON;
-        bi.updateDisplayMeasuredEnergyStatsLocked(580_000, screen, clocks.realtime);
+        screen[0] = Display.STATE_ON;
+        bi.updateDisplayMeasuredEnergyStatsLocked(new long[]{580_000}, screen, clocks.realtime);
         blame1 += 0; // The screen had been doze during the energy period
         globalDoze += 580_000;
         checkMeasuredCharge("H", uid1, blame1, uid2, blame2, globalDoze, bi);
diff --git a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java
index ab38f017936d..ac87806b1639 100644
--- a/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java
+++ b/core/tests/coretests/src/com/android/internal/os/BatteryUsageStatsRule.java
@@ -118,6 +118,12 @@ public class BatteryUsageStatsRule implements TestRule {
         return this;
     }
 
+    public BatteryUsageStatsRule setNumDisplays(int value) {
+        when(mPowerProfile.getNumDisplays()).thenReturn(value);
+        mBatteryStats.setDisplayCountLocked(value);
+        return this;
+    }
+
     /** Call only after setting the power profile information. */
     public BatteryUsageStatsRule initMeasuredEnergyStatsLocked() {
         return initMeasuredEnergyStatsLocked(new String[0]);
diff --git a/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java b/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
index 73f4eb2d2be8..eee5d57c7bc6 100644
--- a/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
+++ b/core/tests/coretests/src/com/android/internal/os/ScreenPowerCalculatorTest.java
@@ -42,11 +42,13 @@ public class ScreenPowerCalculatorTest {
     private static final int APP_UID2 = Process.FIRST_APPLICATION_UID + 43;
     private static final long MINUTE_IN_MS = 60 * 1000;
     private static final long MINUTE_IN_US = 60 * 1000 * 1000;
+    private static final long HOUR_IN_MS = 60 * MINUTE_IN_MS;
 
     @Rule
     public final BatteryUsageStatsRule mStatsRule = new BatteryUsageStatsRule()
             .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 0, 36.0)
-            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 0, 48.0);
+            .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 0, 48.0)
+            .setNumDisplays(1);
 
     @Test
     public void testMeasuredEnergyBasedModel() {
@@ -54,12 +56,13 @@ public class ScreenPowerCalculatorTest {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
 
         batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(0, Display.STATE_ON, 0);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{0},
+                new int[]{Display.STATE_ON}, 0);
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
                 0, 0);
 
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(200_000_000, Display.STATE_ON,
-                15 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{200_000_000},
+                new int[]{Display.STATE_ON}, 15 * MINUTE_IN_MS);
 
         setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
@@ -67,16 +70,16 @@ public class ScreenPowerCalculatorTest {
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
                 20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
 
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(300_000_000, Display.STATE_ON,
-                60 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300_000_000},
+                new int[]{Display.STATE_ON}, 60 * MINUTE_IN_MS);
 
         batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
         setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
                 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
 
-        batteryStats.updateDisplayMeasuredEnergyStatsLocked(100_000_000, Display.STATE_DOZE,
-                120 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{100_000_000},
+                new int[]{Display.STATE_DOZE}, 120 * MINUTE_IN_MS);
 
         mStatsRule.setTime(120 * MINUTE_IN_US, 120 * MINUTE_IN_US);
 
@@ -129,6 +132,104 @@ public class ScreenPowerCalculatorTest {
                 .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
     }
 
+
+    @Test
+    public void testMeasuredEnergyBasedModel_multiDisplay() {
+        mStatsRule.initMeasuredEnergyStatsLocked()
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 1, 60.0)
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 1, 100.0)
+                .setNumDisplays(2);
+
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        final int[] screenStates = new int[]{Display.STATE_ON, Display.STATE_OFF};
+
+        batteryStats.noteScreenStateLocked(0, screenStates[0], 0, 0, 0);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], 0, 0, 0);
+        batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true, 0, 0);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{300, 400}, screenStates, 0);
+
+        batteryStats.noteScreenBrightnessLocked(0, 100, 5 * MINUTE_IN_MS, 5 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(0, 200, 10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
+
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+
+        screenStates[0] = Display.STATE_OFF;
+        screenStates[1] = Display.STATE_ON;
+        batteryStats.noteScreenStateLocked(0, screenStates[0],
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, screenStates[1], 80 * MINUTE_IN_MS,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{600_000_000, 500},
+                screenStates, 80 * MINUTE_IN_MS);
+
+        batteryStats.noteScreenBrightnessLocked(1, 25, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 250, 86 * MINUTE_IN_MS, 86 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 75, 98 * MINUTE_IN_MS, 98 * MINUTE_IN_MS);
+
+        screenStates[1] = Display.STATE_OFF;
+        batteryStats.noteScreenStateLocked(1, screenStates[1], 110 * MINUTE_IN_MS,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+        batteryStats.updateDisplayMeasuredEnergyStatsLocked(new long[]{700, 800_000_000},
+                screenStates, 110 * MINUTE_IN_MS);
+
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+
+        mStatsRule.setTime(120 * MINUTE_IN_US, 120 * MINUTE_IN_US);
+
+        ScreenPowerCalculator calculator =
+                new ScreenPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(calculator);
+
+        BatteryConsumer deviceConsumer = mStatsRule.getDeviceBatteryConsumer();
+        assertThat(deviceConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        // (600000000 + 800000000) uAs * (1 mA / 1000 uA) * (1 h / 3600 s)  = 166.66666 mAh
+        assertThat(deviceConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(388.88888);
+        assertThat(deviceConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+        UidBatteryConsumer uid1 = mStatsRule.getUidBatteryConsumer(APP_UID1);
+        assertThat(uid1.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(20 * MINUTE_IN_MS);
+
+        // Uid1 ran for 20 out of 80 min during the first Display update.
+        // It also ran for 5 out of 45 min during the second Display update:
+        // Uid1 charge = 20 / 80 * 600000000 mAs = 41.66666 mAh
+        assertThat(uid1.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(41.66666);
+        assertThat(uid1.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+        UidBatteryConsumer uid2 = mStatsRule.getUidBatteryConsumer(APP_UID2);
+        assertThat(uid2.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(90 * MINUTE_IN_MS);
+
+        // Uid2 ran for 60 out of 80 min during the first Display update.
+        // It also ran for all of the second Display update:
+        // Uid1 charge = 60 / 80 * 600000000 + 800000000 mAs = 347.22222 mAh
+        assertThat(uid2.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(347.22222);
+        assertThat(uid2.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+        BatteryConsumer appsConsumer = mStatsRule.getAppsBatteryConsumer();
+        assertThat(appsConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        assertThat(appsConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(388.88888);
+        assertThat(appsConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_MEASURED_ENERGY);
+
+    }
+
     @Test
     public void testPowerProfileBasedModel() {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
@@ -197,6 +298,95 @@ public class ScreenPowerCalculatorTest {
                 .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
     }
 
+
+    @Test
+    public void testPowerProfileBasedModel_multiDisplay() {
+        mStatsRule.setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_ON, 1, 60.0)
+                .setAveragePowerForOrdinal(POWER_GROUP_DISPLAY_SCREEN_FULL, 1, 100.0)
+                .setNumDisplays(2);
+
+        BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
+
+        batteryStats.noteScreenStateLocked(0, Display.STATE_ON, 0, 0, 0);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, 0, 0, 0);
+        batteryStats.noteScreenBrightnessLocked(0, 255, 0, 0);
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_TOP, true,
+                0, 0);
+
+        batteryStats.noteScreenBrightnessLocked(0, 100, 5 * MINUTE_IN_MS, 5 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(0, 200, 10 * MINUTE_IN_MS, 10 * MINUTE_IN_MS);
+
+        setProcState(APP_UID1, ActivityManager.PROCESS_STATE_CACHED_EMPTY, false,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP, true,
+                20 * MINUTE_IN_MS, 20 * MINUTE_IN_MS);
+
+        batteryStats.noteScreenStateLocked(0, Display.STATE_OFF,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_ON, 80 * MINUTE_IN_MS,
+                80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 20, 80 * MINUTE_IN_MS, 80 * MINUTE_IN_MS);
+
+        batteryStats.noteScreenBrightnessLocked(1, 250, 86 * MINUTE_IN_MS, 86 * MINUTE_IN_MS);
+        batteryStats.noteScreenBrightnessLocked(1, 75, 98 * MINUTE_IN_MS, 98 * MINUTE_IN_MS);
+        batteryStats.noteScreenStateLocked(1, Display.STATE_OFF, 110 * MINUTE_IN_MS,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+
+        setProcState(APP_UID2, ActivityManager.PROCESS_STATE_TOP_SLEEPING, false,
+                110 * MINUTE_IN_MS, 110 * MINUTE_IN_MS);
+
+        mStatsRule.setTime(120 * MINUTE_IN_US, 120 * MINUTE_IN_US);
+        ScreenPowerCalculator calculator =
+                new ScreenPowerCalculator(mStatsRule.getPowerProfile());
+
+        mStatsRule.apply(BatteryUsageStatsRule.POWER_PROFILE_MODEL_ONLY, calculator);
+
+        BatteryConsumer deviceConsumer = mStatsRule.getDeviceBatteryConsumer();
+        assertThat(deviceConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        // First display consumed 92 mAh.
+        // Second display ran for 0.5 hours at a base drain rate of 60 mA.
+        // 6 minutes (0.1 hours) spent in the first brightness level which drains an extra 10 mA.
+        // 12 minutes (0.2 hours) spent in the fifth brightness level which drains an extra 90 mA.
+        // 12 minutes (0.2 hours) spent in the second brightness level which drains an extra 30 mA.
+        // 92 + 60 * 0.5 + 10 * 0.1 + 90 * 0.2 + 30 * 0.2 = 147
+        assertThat(deviceConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(147);
+        assertThat(deviceConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+        UidBatteryConsumer uid1 = mStatsRule.getUidBatteryConsumer(APP_UID1);
+        assertThat(uid1.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(20 * MINUTE_IN_MS);
+
+        // Uid1 took 20 out of the total of 110 min of foreground activity
+        // Uid1 charge = 20 / 110 * 147.0 = 23.0 mAh
+        assertThat(uid1.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(26.72727);
+        assertThat(uid1.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+        UidBatteryConsumer uid2 = mStatsRule.getUidBatteryConsumer(APP_UID2);
+        assertThat(uid2.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(90 * MINUTE_IN_MS);
+
+        // Uid2 took 90 out of the total of 110 min of foreground activity
+        // Uid2 charge = 90 / 110 * 92.0 = 69.0 mAh
+        assertThat(uid2.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(120.272727);
+        assertThat(uid2.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+        BatteryConsumer appsConsumer = mStatsRule.getAppsBatteryConsumer();
+        assertThat(appsConsumer.getUsageDurationMillis(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(110 * MINUTE_IN_MS);
+        assertThat(appsConsumer.getConsumedPower(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isWithin(PRECISION).of(147);
+        assertThat(appsConsumer.getPowerModel(BatteryConsumer.POWER_COMPONENT_SCREEN))
+                .isEqualTo(BatteryConsumer.POWER_MODEL_POWER_PROFILE);
+
+    }
+
     private void setProcState(int uid, int procState, boolean resumed, long realtimeMs,
             long uptimeMs) {
         BatteryStatsImpl batteryStats = mStatsRule.getBatteryStats();
diff --git a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
index 6374d34294d5..cf4c8a356662 100644
--- a/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
+++ b/services/core/java/com/android/server/am/BatteryExternalStatsWorker.java
@@ -669,12 +669,10 @@ class BatteryExternalStatsWorker implements BatteryStatsImpl.ExternalStatsSync {
             if (measuredEnergyDeltas != null) {
                 final long[] displayChargeUC = measuredEnergyDeltas.displayChargeUC;
                 if (displayChargeUC != null && displayChargeUC.length > 0) {
-                    // TODO (b/194107383): pass all display ordinals to mStats with
-                    //  displayScreenStates
-                    final long primaryDisplayChargeUC = displayChargeUC[0];
-                    // If updating, pass in what BatteryExternalStatsWorker thinks screenState is.
-                    mStats.updateDisplayMeasuredEnergyStatsLocked(primaryDisplayChargeUC,
-                            screenState, elapsedRealtime);
+                    // If updating, pass in what BatteryExternalStatsWorker thinks
+                    // displayScreenStates is.
+                    mStats.updateDisplayMeasuredEnergyStatsLocked(displayChargeUC,
+                            displayScreenStates, elapsedRealtime);
                 }
 
                 final long gnssChargeUC = measuredEnergyDeltas.gnssChargeUC;