diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..bff2d76
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+*.iml
diff --git a/Android.bp b/Android.bp
index b4c1b32..4dbe399 100644
--- a/Android.bp
+++ b/Android.bp
@@ -82,6 +82,7 @@
         "androidx.lifecycle_lifecycle-viewmodel-ktx",
         "androidx.recyclerview_recyclerview",
         "SystemUICustomizationLib",
+        "hilt_android",
     ],
 
     jni_libs: [
diff --git a/OWNERS b/OWNERS
new file mode 100644
index 0000000..04cedef
--- /dev/null
+++ b/OWNERS
@@ -0,0 +1,2 @@
+#Refer the WallpaperPicker2/ThemePicker owners here
+include platform/packages/apps/WallpaperPicker2:/OWNERS
diff --git a/TEST_MAPPING b/TEST_MAPPING
index 72b9f3a..6345ffb 100644
--- a/TEST_MAPPING
+++ b/TEST_MAPPING
@@ -1,22 +1,6 @@
 {
   "presubmit": [
     {
-      "name": "WallpaperPicker2Tests",
-      "options": [
-        {
-          "exclude-annotation": "org.junit.Ignore"
-        }
-      ]
-    },
-    {
-      "name": "ThemePickerTests",
-      "options": [
-        {
-          "exclude-annotation": "org.junit.Ignore"
-        }
-      ]
-    },
-    {
       "name": "WallpaperPickerGoogleTests",
       "options": [
         {
diff --git a/res/drawable/carousel_item_card_background.xml b/res/drawable/carousel_item_card_background.xml
index 946d279..4197698 100644
--- a/res/drawable/carousel_item_card_background.xml
+++ b/res/drawable/carousel_item_card_background.xml
@@ -16,6 +16,6 @@
 -->
 <shape xmlns:android="http://schemas.android.com/apk/res/android"
     android:shape="rectangle">
-    <solid android:color="@color/color_surface" />
+    <solid android:color="@color/system_surface_container_highest" />
     <corners android:radius="16dp" />
 </shape>
\ No newline at end of file
diff --git a/res/drawable/clock_color_size_button_background.xml b/res/drawable/clock_color_size_button_background.xml
index 08d1c45..09b43c8 100644
--- a/res/drawable/clock_color_size_button_background.xml
+++ b/res/drawable/clock_color_size_button_background.xml
@@ -19,7 +19,7 @@
     <item android:id="@android:id/mask">
         <shape android:shape="rectangle">
             <corners android:radius="@dimen/clock_color_size_button_corner_radius" />
-            <solid android:color="?android:colorControlHighlight" />
+            <solid android:color="@color/clock_color_and_size_button_background" />
         </shape>
     </item>
 
diff --git a/res/drawable/color_overflow.xml b/res/drawable/color_overflow.xml
index 90c2684..1ad29fc 100644
--- a/res/drawable/color_overflow.xml
+++ b/res/drawable/color_overflow.xml
@@ -25,10 +25,11 @@
             android:innerRadius="@dimen/component_color_overflow_small_radius_default"
             android:thickness="-1dp"
             android:useLevel="false">
-            <solid android:color="@color/color_surface_variant"/>
+            <solid android:color="@color/system_outline"/>
         </shape>
     </item>
     <item
         android:drawable="@drawable/ic_more_horiz"
+        android:tint="@color/system_on_surface"
         android:gravity="center"/>
 </layer-list>
\ No newline at end of file
diff --git a/res/drawable/picker_section_icon_background.xml b/res/drawable/picker_section_icon_background.xml
new file mode 100644
index 0000000..2a124e0
--- /dev/null
+++ b/res/drawable/picker_section_icon_background.xml
@@ -0,0 +1,30 @@
+<?xml version="1.0" encoding="utf-8"?><!--
+     Copyright (C) 2023 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.
+-->
+<ripple xmlns:android="http://schemas.android.com/apk/res/android"
+    android:color="?android:attr/colorControlHighlight">
+    <item android:id="@android:id/mask">
+        <shape android:shape="rectangle">
+            <solid android:color="@android:color/white"/>
+            <corners android:radius="@dimen/option_tile_radius" />
+        </shape>
+    </item>
+    <item>
+        <shape android:shape="rectangle">
+            <solid android:color="@color/picker_section_icon_background"/>
+            <corners android:radius="@dimen/option_tile_radius" />
+        </shape>
+    </item>
+</ripple>
diff --git a/res/drawable/top_connected_section_background.xml b/res/drawable/top_connected_section_background.xml
deleted file mode 100644
index f96824b..0000000
--- a/res/drawable/top_connected_section_background.xml
+++ /dev/null
@@ -1,29 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!--
-  ~ Copyright (C) 2023 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.
-  ~
-  -->
-
-<shape
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:shape="rectangle">
-
-    <solid android:color="@color/color_surface" />
-
-    <corners
-        android:topLeftRadius="32dp"
-        android:topRightRadius="32dp"
-        android:bottomLeftRadius="8dp"
-        android:bottomRightRadius="8dp" />
-</shape>
\ No newline at end of file
diff --git a/res/layout/clock_color_option.xml b/res/layout/clock_color_option.xml
index 786d633..e5aa3d6 100644
--- a/res/layout/clock_color_option.xml
+++ b/res/layout/clock_color_option.xml
@@ -29,7 +29,7 @@
         android:layout_width="match_parent"
         android:layout_height="wrap_content"
         android:layout_marginTop="@dimen/option_bottom_margin"
-        android:textColor="@color/text_color_primary"
+        android:textColor="@color/system_on_surface"
         android:maxLines="2"
         android:hyphenationFrequency="normal"
         android:ellipsize="end"
diff --git a/res/layout/clock_size_radio_button_group.xml b/res/layout/clock_size_radio_button_group.xml
index 2aa78d2..56d0ab7 100644
--- a/res/layout/clock_size_radio_button_group.xml
+++ b/res/layout/clock_size_radio_button_group.xml
@@ -68,11 +68,21 @@
             android:layout_marginEnd="8dp"
             android:clickable="false" />
 
-        <TextView
-            style="@style/SectionTitleTextStyle"
-            android:layout_width="wrap_content"
+        <LinearLayout
+            android:layout_width="match_parent"
             android:layout_height="wrap_content"
-            android:layout_gravity="center_vertical"
-            android:text="@string/clock_size_small" />
+            android:orientation="vertical">
+            <TextView
+                style="@style/SectionTitleTextStyle"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:text="@string/clock_size_small" />
+
+            <TextView
+                style="@style/SectionSubtitleTextStyle"
+                android:layout_width="wrap_content"
+                android:layout_height="wrap_content"
+                android:text="@string/clock_size_small_description" />
+        </LinearLayout>
     </LinearLayout>
 </LinearLayout>
\ No newline at end of file
diff --git a/res/layout/color_section_view.xml b/res/layout/color_section_view.xml
deleted file mode 100644
index 6095c9b..0000000
--- a/res/layout/color_section_view.xml
+++ /dev/null
@@ -1,38 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-     Copyright (C) 2022 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.
--->
-<com.android.customization.picker.color.ColorSectionView
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="match_parent"
-    android:layout_height="wrap_content"
-    android:paddingBottom="@dimen/section_bottom_padding"
-    android:orientation="vertical">
-
-    <FrameLayout
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:paddingHorizontal="@dimen/separated_tabs_horizontal_padding">
-        <include layout="@layout/separated_tabs" />
-    </FrameLayout>
-
-    <androidx.viewpager2.widget.ViewPager2
-        android:id="@+id/color_section_view_pager"
-        android:layout_width="match_parent"
-        android:layout_height="wrap_content"
-        android:layout_marginTop="@dimen/color_options_container_top_margin"
-        android:clipChildren="false"
-        android:clipToPadding="false"/>
-</com.android.customization.picker.color.ColorSectionView>
diff --git a/res/layout/fragment_clock_settings.xml b/res/layout/fragment_clock_settings.xml
index bb4a0c2..5736265 100644
--- a/res/layout/fragment_clock_settings.xml
+++ b/res/layout/fragment_clock_settings.xml
@@ -19,7 +19,8 @@
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:orientation="vertical"
-    android:clipChildren="false">
+    android:clipChildren="false"
+    android:transitionGroup="true">
 
     <FrameLayout
         android:id="@+id/section_header_container"
diff --git a/res/layout/fragment_color_picker.xml b/res/layout/fragment_color_picker.xml
index 7cf94d6..ebff6a4 100644
--- a/res/layout/fragment_color_picker.xml
+++ b/res/layout/fragment_color_picker.xml
@@ -101,7 +101,7 @@
             android:id="@+id/color_type_tab_subhead"
             android:layout_width="match_parent"
             android:layout_height="wrap_content"
-            android:textColor="@color/text_color_secondary"
+            android:textColor="@color/system_on_surface_variant"
             android:textSize="12sp"
             android:gravity="center"
             android:paddingHorizontal="16dp"/>
diff --git a/res/layout/fragment_grid.xml b/res/layout/fragment_grid.xml
index 4f0aaef..8c97d45 100644
--- a/res/layout/fragment_grid.xml
+++ b/res/layout/fragment_grid.xml
@@ -88,6 +88,24 @@
 
         </FrameLayout>
 
+        <Button
+            android:id="@+id/apply_button"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_horizontal"
+            android:visibility="gone"
+            android:enabled="false"
+            style="@style/ActionPrimaryButton"
+            android:text="@string/apply_btn" />
+
+        <TextView
+            android:id="@+id/apply_button_note"
+            android:layout_width="wrap_content"
+            android:layout_height="wrap_content"
+            android:layout_gravity="center_horizontal"
+            android:textSize="@dimen/grid_apply_button_note_text_size"
+            android:visibility="gone"
+            android:text="@string/apply_grid_btn_note" />
     </LinearLayout>
 
 </LinearLayout>
diff --git a/res/layout/fragment_lock_screen_quick_affordances.xml b/res/layout/fragment_lock_screen_quick_affordances.xml
index 36d1697..ee259f1 100644
--- a/res/layout/fragment_lock_screen_quick_affordances.xml
+++ b/res/layout/fragment_lock_screen_quick_affordances.xml
@@ -19,7 +19,8 @@
     android:layout_width="match_parent"
     android:layout_height="match_parent"
     android:orientation="vertical"
-    android:clipChildren="false">
+    android:clipChildren="false"
+    android:transitionGroup="true">
 
     <FrameLayout
         android:id="@+id/section_header_container"
diff --git a/res/layout/grid_option_2.xml b/res/layout/grid_option_2.xml
index a8b453a..efed018 100644
--- a/res/layout/grid_option_2.xml
+++ b/res/layout/grid_option_2.xml
@@ -61,7 +61,7 @@
         android:id="@id/text"
         android:layout_width="wrap_content"
         android:layout_height="wrap_content"
-        android:textColor="@color/text_color_primary"
+        android:textColor="@color/system_on_surface"
         android:singleLine="true"
         android:ellipsize="end"
         android:text="Placeholder for stable size calculation, please do not remove."
diff --git a/res/layout/grid_section_view.xml b/res/layout/grid_section_view.xml
index a90ee64..a7b2b1f 100644
--- a/res/layout/grid_section_view.xml
+++ b/res/layout/grid_section_view.xml
@@ -49,7 +49,8 @@
         android:layout_height="@dimen/option_tile_width"
         android:scaleType="center"
         android:src="@drawable/ic_grid_24px"
-        android:background="@drawable/option_border_color"
+        android:tint="@color/system_on_surface"
+        android:background="@drawable/picker_section_icon_background"
         android:contentDescription="@string/gird_picker_entry_content_description" />
 
 </com.android.customization.picker.grid.GridSectionView>
\ No newline at end of file
diff --git a/res/layout/keyguard_quick_affordance.xml b/res/layout/keyguard_quick_affordance.xml
index a86489a..d008ac1 100644
--- a/res/layout/keyguard_quick_affordance.xml
+++ b/res/layout/keyguard_quick_affordance.xml
@@ -50,7 +50,7 @@
             android:layout_width="@dimen/keyguard_quick_affordance_icon_size"
             android:layout_height="@dimen/keyguard_quick_affordance_icon_size"
             android:layout_gravity="center"
-            android:tint="@color/text_color_primary" />
+            android:tint="@color/system_on_surface" />
 
     </FrameLayout>
 
@@ -63,7 +63,7 @@
         android:layout_width="wrap_content"
         android:layout_height="wrap_content"
         android:gravity="center_horizontal"
-        android:textColor="@color/text_color_primary"
+        android:textColor="@color/system_on_surface"
         android:maxLines="2"
         android:hyphenationFrequency="normal"
         android:ellipsize="end"
diff --git a/res/layout/keyguard_quick_affordance_section_view.xml b/res/layout/keyguard_quick_affordance_section_view.xml
index 7c5babb..3cda20f 100644
--- a/res/layout/keyguard_quick_affordance_section_view.xml
+++ b/res/layout/keyguard_quick_affordance_section_view.xml
@@ -48,7 +48,7 @@
         android:layout_width="@dimen/option_tile_width"
         android:layout_height="@dimen/option_tile_width"
         android:orientation="horizontal"
-        android:background="@drawable/option_border_color"
+        android:background="@drawable/picker_section_icon_background"
         android:importantForAccessibility="noHideDescendants"
         android:gravity="center"
         android:divider="@drawable/horizontal_divider_14dp"
@@ -60,14 +60,14 @@
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
             android:visibility="gone"
-            android:tint="@color/text_color_primary" />
+            android:tint="@color/system_on_surface" />
 
         <ImageView
             android:id="@+id/icon_2"
             android:layout_width="wrap_content"
             android:layout_height="wrap_content"
             android:visibility="gone"
-            android:tint="@color/text_color_primary" />
+            android:tint="@color/system_on_surface" />
 
     </LinearLayout>
 
diff --git a/res/layout/single_clock_view.xml b/res/layout/single_clock_view.xml
deleted file mode 100644
index ffc6039..0000000
--- a/res/layout/single_clock_view.xml
+++ /dev/null
@@ -1,27 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?><!--
-     Copyright (C) 2023 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.
--->
-<com.android.wallpaper.picker.FixedWidthDisplayRatioFrameLayout
-    xmlns:android="http://schemas.android.com/apk/res/android"
-    android:layout_width="@dimen/screen_preview_width"
-    android:layout_height="wrap_content"
-    android:layout_gravity="center"
-    android:clipChildren="false">
-    <com.android.customization.picker.clock.ui.view.ClockHostView
-        android:id="@+id/single_clock_host_view"
-        android:layout_width="match_parent"
-        android:layout_height="match_parent"
-        android:layout_gravity="center" />
-</com.android.wallpaper.picker.FixedWidthDisplayRatioFrameLayout>
\ No newline at end of file
diff --git a/res/values-h800dp/dimens.xml b/res/values-h800dp/dimens.xml
deleted file mode 100644
index 36d0d62..0000000
--- a/res/values-h800dp/dimens.xml
+++ /dev/null
@@ -1,9 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-
-<resources>
-    <!--
-    Dimension for clock translationY when swiping clock carousel in picker,
-    copied from sysui resources
-    -->
-    <dimen name="keyguard_large_clock_top_margin">-112dp</dimen>
-</resources>
\ No newline at end of file
diff --git a/res/values-w300dp/integers.xml b/res/values-w300dp/integers.xml
new file mode 100644
index 0000000..3bd7f4b
--- /dev/null
+++ b/res/values-w300dp/integers.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2019 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.
+-->
+<resources>
+    <integer name="color_section_num_columns">5</integer>
+</resources>
diff --git a/res/values-w400dp/integers.xml b/res/values-w400dp/integers.xml
new file mode 100644
index 0000000..0bb42e3
--- /dev/null
+++ b/res/values-w400dp/integers.xml
@@ -0,0 +1,19 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!--
+     Copyright (C) 2019 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.
+-->
+<resources>
+    <integer name="color_section_num_columns">6</integer>
+</resources>
diff --git a/res/values/dimens.xml b/res/values/dimens.xml
index a3e0d4a..aa6c477 100644
--- a/res/values/dimens.xml
+++ b/res/values/dimens.xml
@@ -61,6 +61,7 @@
     <!-- Dimensions for the grid options -->
     <dimen name="grid_options_container_bottom_margin">@dimen/bottom_actions_height</dimen>
     <dimen name="grid_options_container_horizontal_margin">24dp</dimen>
+    <dimen name="grid_apply_button_note_text_size">11sp</dimen>
 
     <dimen name="card_title_text_size">16sp</dimen>
     <dimen name="card_header_icon_size">32dp</dimen>
@@ -170,14 +171,4 @@
     <dimen name="clock_color_size_button_corner_radius">16dp</dimen>
     <dimen name="clock_color_size_button_icon_size">20dp</dimen>
     <dimen name="clock_color_size_button_icon_margin_end">12dp</dimen>
-
-    <!--
-    Dimension for clock translationY when swiping clock carousel in picker,
-    copied from sysui resources
-    -->
-    <dimen name="keyguard_large_clock_top_margin">-60dp</dimen>
-    <!-- Dimension for the clock view, copied from sysui resources. -->
-    <dimen name="small_clock_height">114dp</dimen>
-    <dimen name="small_clock_padding_top">28dp</dimen>
-    <dimen name="clock_padding_start">28dp</dimen>
 </resources>
diff --git a/res/values/strings.xml b/res/values/strings.xml
index 680947d..9b55965 100755
--- a/res/values/strings.xml
+++ b/res/values/strings.xml
@@ -90,6 +90,9 @@
     <!-- Title of a radio button to apply clock size small. [CHAR LIMIT=15] -->
     <string name="clock_size_small">Small</string>
 
+    <!-- Description of a radio button to apply clock size small. [CHAR LIMIT=NONE] -->
+    <string name="clock_size_small_description">A small clock shows in the corner of your screen</string>
+
     <!-- Title of a section of the customization picker where the user can select a Grid size for
         the home screen. [CHAR LIMIT=15] -->
     <string name="grid_title">App grid</string>
@@ -175,6 +178,15 @@
     <!--Title for a grid option, describing the number of columns and rows, eg: 4x4 [CHAR LIMIT=10] -->
     <string name="grid_title_pattern"><xliff:g name="num_cols" example="1">%1$d</xliff:g>x<xliff:g name="num_rows" example="1">%2$d</xliff:g></string>
 
+    <!-- Label of what will happen when user tap on apply button to change grid. [CHAR LIMIT=50] -->
+    <string name="apply_grid_btn_note">Changing grid size will reload workspace and may take a few seconds.</string>
+
+    <!-- Toast of reloading workspace with new grid. -->
+    <string name="toast_of_changing_grid">Reloading workspace with %1$s grid</string>
+
+    <!-- Toast of failure to reload workspace with new grid. -->
+    <string name="toast_of_failure_to_change_grid">Failed to reload workspace with %1$s grid</string>
+
     <!-- Message shown when a theme has been applied successfully in the system [CHAR LIMIT=NONE] -->
     <string name="applied_theme_msg">Style set successfully</string>
 
diff --git a/res/values/styles.xml b/res/values/styles.xml
index c70b596..b2cd71d 100644
--- a/res/values/styles.xml
+++ b/res/values/styles.xml
@@ -48,6 +48,7 @@
         <item name="colorSurface">@color/design_default_color_surface</item>
         <item name="colorOnSurface">@color/design_default_color_on_surface</item>
         <item name="android:textAllCaps">false</item>
+        <item name="motionDurationLong1">500</item>
     </style>
 
     <!-- Snackbar margin -->
@@ -114,7 +115,7 @@
 
     <style name="SectionTitleTextStyle">
         <item name="android:fontFamily">@*android:string/config_headlineFontFamily</item>
-        <item name="android:textColor">?android:textColorPrimary</item>
+        <item name="android:textColor">@color/system_on_surface</item>
         <item name="android:textDirection">locale</item>
         <item name="android:textSize">20sp</item>
         <item name="android:lineHeight">24dp</item>
@@ -122,7 +123,7 @@
 
     <style name="SectionSubtitleTextStyle" parent="SectionTitleTextStyle">
         <item name="android:fontFamily">@*android:string/config_bodyFontFamily</item>
-        <item name="android:textColor">?android:textColorSecondary</item>
+        <item name="android:textColor">@color/system_on_surface_variant</item>
         <item name="android:textSize">14sp</item>
         <item name="android:lineHeight">16sp</item>
     </style>
diff --git a/src/com/android/customization/model/CustomizationManager.java b/src/com/android/customization/model/CustomizationManager.java
index 104cc83..e6d3872 100644
--- a/src/com/android/customization/model/CustomizationManager.java
+++ b/src/com/android/customization/model/CustomizationManager.java
@@ -72,6 +72,9 @@
      */
     void apply(T option, Callback callback);
 
+    /** Preview the given option without committing the change. */
+    default void preview(T option) {}
+
     /**
      * Loads the available options for the type of Customization managed by this class, calling the
      * given callback when done.
diff --git a/src/com/android/customization/model/color/ColorProvider.kt b/src/com/android/customization/model/color/ColorProvider.kt
index 63d298d..74a4051 100644
--- a/src/com/android/customization/model/color/ColorProvider.kt
+++ b/src/com/android/customization/model/color/ColorProvider.kt
@@ -16,6 +16,7 @@
 package com.android.customization.model.color
 
 import android.app.WallpaperColors
+import android.app.WallpaperManager
 import android.content.Context
 import android.content.res.ColorStateList
 import android.content.res.Resources
@@ -38,7 +39,6 @@
 import com.android.systemui.monet.ColorScheme
 import com.android.systemui.monet.Style
 import com.android.wallpaper.R
-import com.android.wallpaper.compat.WallpaperManagerCompat
 import com.android.wallpaper.module.InjectorProvider
 import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.Dispatchers
@@ -130,9 +130,9 @@
     private fun isLockScreenWallpaperLastApplied(): Boolean {
         // The WallpaperId increases every time a new wallpaper is set, so the larger wallpaper id
         // is the most recently set wallpaper
-        val manager = InjectorProvider.getInjector().getWallpaperManagerCompat(mContext)
-        return manager.getWallpaperId(WallpaperManagerCompat.FLAG_LOCK) >
-            manager.getWallpaperId(WallpaperManagerCompat.FLAG_SYSTEM)
+        val manager = WallpaperManager.getInstance(mContext)
+        return manager.getWallpaperId(WallpaperManager.FLAG_LOCK) >
+            manager.getWallpaperId(WallpaperManager.FLAG_SYSTEM)
     }
 
     private fun loadSeedColors(
@@ -432,7 +432,8 @@
             val extractor = ColorBundlePreviewExtractor(mContext)
             val bundles: MutableList<ColorOption> = ArrayList()
 
-            val bundleNames = getItemsFromStub(COLOR_BUNDLES_ARRAY_NAME)
+            val bundleNames =
+                if (isAvailable) getItemsFromStub(COLOR_BUNDLES_ARRAY_NAME) else emptyArray()
             // Color option index value starts from 1.
             var index = 1
             val maxPresetColors = if (themeStyleEnabled) bundleNames.size else MAX_PRESET_COLORS
diff --git a/src/com/android/customization/model/color/ColorSectionController.java b/src/com/android/customization/model/color/ColorSectionController.java
deleted file mode 100644
index be051ac..0000000
--- a/src/com/android/customization/model/color/ColorSectionController.java
+++ /dev/null
@@ -1,538 +0,0 @@
-/*
- * Copyright (C) 2022 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.customization.model.color;
-
-import static android.view.View.IMPORTANT_FOR_ACCESSIBILITY_NO;
-import static android.view.View.VISIBLE;
-
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_HOME;
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_LOCK;
-import static com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_PRESET;
-import static com.android.customization.widget.OptionSelectorController.CheckmarkStyle.CENTER;
-
-import android.app.Activity;
-import android.app.WallpaperColors;
-import android.content.Context;
-import android.os.Bundle;
-import android.os.Handler;
-import android.os.SystemClock;
-import android.stats.style.StyleEnums;
-import android.text.TextUtils;
-import android.util.Log;
-import android.view.LayoutInflater;
-import android.view.View;
-import android.view.ViewGroup;
-import android.view.accessibility.AccessibilityNodeInfo;
-import android.widget.FrameLayout;
-
-import androidx.annotation.Nullable;
-import androidx.lifecycle.LifecycleOwner;
-import androidx.recyclerview.widget.RecyclerView;
-import androidx.viewpager2.widget.ViewPager2;
-
-import com.android.customization.model.CustomizationManager;
-import com.android.customization.model.theme.OverlayManagerCompat;
-import com.android.customization.module.CustomizationInjector;
-import com.android.customization.module.ThemesUserEventLogger;
-import com.android.customization.picker.color.ColorSectionView;
-import com.android.customization.widget.OptionSelectorController;
-import com.android.wallpaper.R;
-import com.android.wallpaper.model.CustomizationSectionController;
-import com.android.wallpaper.model.WallpaperColorsViewModel;
-import com.android.wallpaper.module.InjectorProvider;
-import com.android.wallpaper.widget.PageIndicator;
-import com.android.wallpaper.widget.SeparatedTabLayout;
-
-import com.google.common.collect.Iterables;
-import com.google.common.collect.Lists;
-
-import java.util.ArrayList;
-import java.util.List;
-import java.util.Locale;
-import java.util.Optional;
-
-/**
- * Color section view's controller for the logic of color customization.
- */
-public class ColorSectionController implements CustomizationSectionController<ColorSectionView> {
-
-    private static final String TAG = "ColorSectionController";
-    private static final String KEY_COLOR_TAB_POSITION = "COLOR_TAB_POSITION";
-    private static final String KEY_COLOR_PAGE_POSITION = "COLOR_PAGE_POSITION";
-    private static final String ID_VIEWPAGER = "ColorSectionController_colorSectionViewPager";
-    private static final String ID_ITEMVIEW = "ColorSectionController_itemView";
-    private static final String ID_CONTAINER = "ColorSectionController_container";
-    private static final long MIN_COLOR_APPLY_PERIOD = 500L;
-
-    private static final int WALLPAPER_TAB_INDEX = 0;
-    private static final int PRESET_TAB_INDEX = 1;
-
-    private final ThemesUserEventLogger mEventLogger;
-    private final ColorCustomizationManager mColorManager;
-    private final WallpaperColorsViewModel mWallpaperColorsViewModel;
-    private final LifecycleOwner mLifecycleOwner;
-    private final ColorSectionAdapter mColorSectionAdapter = new ColorSectionAdapter();
-
-    private List<ColorOption> mWallpaperColorOptions = new ArrayList<>();
-    private List<ColorOption> mPresetColorOptions = new ArrayList<>();
-    private ViewPager2 mColorSectionViewPager;
-    private ColorOption mSelectedColor;
-    private SeparatedTabLayout mTabLayout;
-    @Nullable private WallpaperColors mHomeWallpaperColors;
-    @Nullable private WallpaperColors mLockWallpaperColors;
-    // Uses a boolean value to indicate whether wallpaper color is ready because WallpaperColors
-    // maybe be null when it's ready.
-    private boolean mHomeWallpaperColorsReady;
-    private boolean mLockWallpaperColorsReady;
-    private Optional<Integer> mTabPositionToRestore = Optional.empty();
-    private Optional<Integer>[] mPagePositionToRestore =
-            new Optional[]{Optional.empty(), Optional.empty()};
-    private long mLastColorApplyingTime = 0L;
-    private ColorSectionView mColorSectionView;
-
-    private static int getNumPages(int optionsPerPage, int totalOptions) {
-        return (int) Math.ceil((float) totalOptions / optionsPerPage);
-    }
-
-    public ColorSectionController(Activity activity, WallpaperColorsViewModel viewModel,
-            LifecycleOwner lifecycleOwner, @Nullable Bundle savedInstanceState) {
-        CustomizationInjector injector = (CustomizationInjector) InjectorProvider.getInjector();
-        mEventLogger = (ThemesUserEventLogger) injector.getUserEventLogger(activity);
-        mColorManager = ColorCustomizationManager.getInstance(activity,
-                new OverlayManagerCompat(activity));
-        mWallpaperColorsViewModel = viewModel;
-        mLifecycleOwner = lifecycleOwner;
-
-        if (savedInstanceState != null) {
-            if (savedInstanceState.containsKey(KEY_COLOR_TAB_POSITION)) {
-                mTabPositionToRestore = Optional.of(
-                        savedInstanceState.getInt(KEY_COLOR_TAB_POSITION));
-            }
-
-            for (int i = 0; i < mPagePositionToRestore.length; i++) {
-                String keyColorPage = getPagePositionKey(i);
-                if (keyColorPage != null && savedInstanceState.containsKey(keyColorPage)) {
-                    setPagePositionToRestore(i, savedInstanceState.getInt(keyColorPage));
-                }
-            }
-        }
-    }
-
-    private String getPagePositionKey(int index) {
-        return String.format(Locale.US, "%s_%d", KEY_COLOR_PAGE_POSITION, index);
-    }
-
-    private void setPagePositionToRestore(int pagePositionKeyIndex, int pagePosition) {
-        if (pagePositionKeyIndex >= 0 && pagePositionKeyIndex < mPagePositionToRestore.length) {
-            mPagePositionToRestore[pagePositionKeyIndex] = Optional.of(pagePosition);
-        }
-    }
-
-    private int getPagePositionToRestore(int pagePositionKeyIndex, int defaultPagePosition) {
-        if (pagePositionKeyIndex >= 0 && pagePositionKeyIndex < mPagePositionToRestore.length) {
-            return mPagePositionToRestore[pagePositionKeyIndex].orElse(defaultPagePosition);
-        }
-        return 0;
-    }
-
-    @Override
-    public boolean isAvailable(@Nullable Context context) {
-        return context != null && ColorUtils.isMonetEnabled(context) && mColorManager.isAvailable();
-    }
-
-    @Override
-    public ColorSectionView createView(Context context) {
-        mColorSectionView = (ColorSectionView) LayoutInflater.from(context).inflate(
-                R.layout.color_section_view, /* root= */ null);
-        mColorSectionViewPager = mColorSectionView.findViewById(R.id.color_section_view_pager);
-        mColorSectionViewPager.setAccessibilityDelegate(createAccessibilityDelegate(ID_VIEWPAGER));
-        mColorSectionViewPager.setAdapter(mColorSectionAdapter);
-        mColorSectionViewPager.setUserInputEnabled(false);
-        if (ColorProvider.themeStyleEnabled) {
-            mColorSectionViewPager.setImportantForAccessibility(IMPORTANT_FOR_ACCESSIBILITY_NO);
-        }
-        mTabLayout = mColorSectionView.findViewById(R.id.separated_tabs);
-        mColorSectionAdapter.setNumColors(context.getResources().getInteger(
-                R.integer.options_grid_num_columns));
-        // TODO(b/202145216): Use just 2 views when tapping either button on top.
-        mTabLayout.setViewPager(mColorSectionViewPager);
-
-        mWallpaperColorsViewModel.getHomeWallpaperColorsLiveData().observe(mLifecycleOwner,
-                homeColors -> {
-                    mHomeWallpaperColors = homeColors;
-                    mHomeWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        mWallpaperColorsViewModel.getLockWallpaperColorsLiveData().observe(mLifecycleOwner,
-                lockColors -> {
-                    mLockWallpaperColors = lockColors;
-                    mLockWallpaperColorsReady = true;
-                    maybeLoadColors();
-                });
-        return mColorSectionView;
-    }
-
-    @Override
-    public void onSaveInstanceState(Bundle savedInstanceState) {
-        if (mColorSectionViewPager != null) {
-            savedInstanceState.putInt(KEY_COLOR_TAB_POSITION,
-                    mColorSectionViewPager.getCurrentItem());
-
-            for (int i = 0; i < mPagePositionToRestore.length; i++) {
-                savedInstanceState.putInt(getPagePositionKey(i), getPagePositionToRestore(i, 0));
-            }
-        }
-    }
-
-    private void maybeLoadColors() {
-        if (mHomeWallpaperColorsReady && mLockWallpaperColorsReady) {
-            mColorManager.setWallpaperColors(mHomeWallpaperColors, mLockWallpaperColors);
-            loadColorOptions(/* reload= */ false);
-        }
-    }
-
-    private void loadColorOptions(boolean reload) {
-        mColorManager.fetchOptions(new CustomizationManager.OptionsFetchedListener<ColorOption>() {
-            @Override
-            public void onOptionsLoaded(List<ColorOption> options) {
-                List<ColorOption> wallpaperColorOptions = new ArrayList<>();
-                List<ColorOption> presetColorOptions = new ArrayList<>();
-                for (ColorOption option : options) {
-                    if (option instanceof ColorSeedOption) {
-                        wallpaperColorOptions.add(option);
-                    } else if (option instanceof ColorBundle) {
-                        presetColorOptions.add(option);
-                    }
-                }
-                mWallpaperColorOptions = wallpaperColorOptions;
-                mPresetColorOptions = presetColorOptions;
-                mSelectedColor = findActiveColorOption(mWallpaperColorOptions,
-                        mPresetColorOptions);
-                mTabLayout.post(()-> setUpColorViewPager());
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                if (throwable != null) {
-                    Log.e(TAG, "Error loading theme bundles", throwable);
-                }
-            }
-        }, reload);
-    }
-
-    private void setUpColorViewPager() {
-        mColorSectionAdapter.notifyDataSetChanged();
-
-        if (mTabLayout != null && mTabLayout.getTabCount() == 0) {
-            mTabLayout.addTab(mTabLayout.newTab().setText(R.string.wallpaper_color_tab),
-                    WALLPAPER_TAB_INDEX);
-            mTabLayout.addTab(mTabLayout.newTab().setText(R.string.preset_color_tab),
-                    PRESET_TAB_INDEX);
-        }
-
-        if (mWallpaperColorOptions.isEmpty()) {
-            // Select preset tab and disable wallpaper tab.
-            mTabLayout.getTabAt(WALLPAPER_TAB_INDEX).view.setEnabled(false);
-            mColorSectionViewPager.setCurrentItem(PRESET_TAB_INDEX, /* smoothScroll= */ false);
-            return;
-        }
-
-        mColorSectionViewPager.setCurrentItem(
-                mTabPositionToRestore.orElseGet(
-                        () -> COLOR_SOURCE_PRESET.equals(mColorManager.getCurrentColorSource())
-                                ? PRESET_TAB_INDEX
-                                : WALLPAPER_TAB_INDEX),
-                /* smoothScroll= */ false);
-
-        // Disable "wallpaper colors" and "basic colors" swiping for new color style.
-        mColorSectionViewPager.setUserInputEnabled(!ColorProvider.themeStyleEnabled);
-    }
-
-    private void setupColorPages(ViewPager2 container, int colorsPerPage, int sectionPosition,
-            List<ColorOption> options, PageIndicator pageIndicator) {
-        container.setAdapter(new ColorPageAdapter(options, /* pageEnabled= */ true,
-                colorsPerPage));
-        if (ColorProvider.themeStyleEnabled) {
-            // Update page index to show selected items.
-            int selectedIndex = options.indexOf(mSelectedColor);
-            if (colorsPerPage != 0) {
-                int pageIndex = selectedIndex / colorsPerPage;
-                int position = getPagePositionToRestore(sectionPosition, pageIndex);
-                container.setCurrentItem(position, /* smoothScroll= */ false);
-            }
-            pageIndicator.setNumPages(getNumPages(colorsPerPage, options.size()));
-            registerOnPageChangeCallback(sectionPosition, container, pageIndicator);
-        }
-    }
-
-    private void registerOnPageChangeCallback(int sectionPosition, ViewPager2 container,
-            PageIndicator pageIndicator) {
-        container.registerOnPageChangeCallback(new ViewPager2.OnPageChangeCallback() {
-            @Override
-            public void onPageSelected(int position) {
-                super.onPageSelected(position);
-                if (mColorSectionViewPager.getCurrentItem() == sectionPosition) {
-                    pageIndicator.setLocation(getPagePosition(pageIndicator, position));
-                    setPagePositionToRestore(sectionPosition, position);
-                }
-            }
-
-            @Override
-            public void onPageScrolled(int position, float positionOffset,
-                    int positionOffsetPixels) {
-                super.onPageScrolled(position, positionOffset, positionOffsetPixels);
-                if (mColorSectionViewPager.getCurrentItem() == sectionPosition) {
-                    pageIndicator.setLocation(getPagePosition(pageIndicator, position));
-                    setPagePositionToRestore(sectionPosition, position);
-                }
-            }
-
-            private int getPagePosition(PageIndicator pageIndicator, int position) {
-                return pageIndicator.isLayoutRtl() ? pageIndicator.getChildCount() - 1 - position
-                        : position;
-            }
-        });
-    }
-
-    private void setupColorOptions(RecyclerView container, List<ColorOption> colorOptions,
-            boolean pageEnabled, int index, int colorsPerPage) {
-        int totalSize = colorOptions.size();
-        if (totalSize == 0) {
-            return;
-        }
-
-        List<ColorOption> subOptions;
-        if (pageEnabled && ColorProvider.themeStyleEnabled) {
-            subOptions = colorOptions.subList(colorsPerPage * index,
-                    Math.min(colorsPerPage * (index + 1), totalSize));
-        } else {
-            subOptions = colorOptions;
-        }
-
-        OptionSelectorController<ColorOption> adaptiveController = new OptionSelectorController<>(
-                container, subOptions, /* useGrid= */ true, CENTER);
-        adaptiveController.initOptions(mColorManager);
-        setUpColorOptionsController(adaptiveController);
-    }
-
-    private ColorOption findActiveColorOption(List<ColorOption> wallpaperColorOptions,
-            List<ColorOption> presetColorOptions) {
-        ColorOption activeColorOption = null;
-        for (ColorOption colorOption : Lists.newArrayList(
-                Iterables.concat(wallpaperColorOptions, presetColorOptions))) {
-            if (colorOption.isActive(mColorManager)) {
-                activeColorOption = colorOption;
-                break;
-            }
-        }
-        // Use the first one option by default. This should not happen as above should have an
-        // active option found.
-        if (activeColorOption == null) {
-            activeColorOption = wallpaperColorOptions.isEmpty()
-                    ? presetColorOptions.get(0)
-                    : wallpaperColorOptions.get(0);
-        }
-        return activeColorOption;
-    }
-
-    private void setUpColorOptionsController(
-            OptionSelectorController<ColorOption> optionSelectorController) {
-        if (mSelectedColor != null && optionSelectorController.containsOption(mSelectedColor)) {
-            optionSelectorController.setSelectedOption(mSelectedColor);
-        }
-
-        optionSelectorController.addListener(selectedOption -> {
-            ColorOption selectedColor = (ColorOption) selectedOption;
-            if (mSelectedColor.equals(selectedColor)) {
-                return;
-            }
-            mSelectedColor = (ColorOption) selectedOption;
-            // Post with delay for color option to run ripple.
-            new Handler().postDelayed(()-> applyColor(mSelectedColor), /* delayMillis= */ 100);
-        });
-    }
-
-    private void applyColor(ColorOption colorOption) {
-        if (SystemClock.elapsedRealtime() - mLastColorApplyingTime < MIN_COLOR_APPLY_PERIOD) {
-            return;
-        }
-        mLastColorApplyingTime = SystemClock.elapsedRealtime();
-        mColorManager.apply(colorOption, new CustomizationManager.Callback() {
-            @Override
-            public void onSuccess() {
-                mColorSectionView.announceForAccessibility(
-                        mColorSectionView.getContext().getString(R.string.color_changed));
-                mEventLogger.logColorApplied(getColorAction(colorOption), colorOption);
-            }
-
-            @Override
-            public void onError(@Nullable Throwable throwable) {
-                Log.w(TAG, "Apply theme with error: " + throwable);
-            }
-        });
-    }
-
-    private int getColorAction(ColorOption colorOption) {
-        int action = StyleEnums.DEFAULT_ACTION;
-        boolean isForBoth = mLockWallpaperColors == null || mLockWallpaperColors.equals(
-                mHomeWallpaperColors);
-
-        if (TextUtils.equals(colorOption.getSource(), COLOR_SOURCE_PRESET)) {
-            action = StyleEnums.COLOR_PRESET_APPLIED;
-        } else if (isForBoth) {
-            action = StyleEnums.COLOR_WALLPAPER_HOME_LOCK_APPLIED;
-        } else {
-            switch (colorOption.getSource()) {
-                case COLOR_SOURCE_HOME:
-                    action = StyleEnums.COLOR_WALLPAPER_HOME_APPLIED;
-                    break;
-                case COLOR_SOURCE_LOCK:
-                    action = StyleEnums.COLOR_WALLPAPER_LOCK_APPLIED;
-                    break;
-            }
-        }
-        return action;
-    }
-
-    private View.AccessibilityDelegate createAccessibilityDelegate(String id) {
-        return new View.AccessibilityDelegate() {
-            @Override
-            public void onInitializeAccessibilityNodeInfo(View host, AccessibilityNodeInfo info) {
-                super.onInitializeAccessibilityNodeInfo(host, info);
-                info.setUniqueId(id);
-            }
-        };
-    }
-
-    private class ColorSectionAdapter extends
-            RecyclerView.Adapter<ColorSectionAdapter.ColorPageViewHolder> {
-
-        private final int mItemCounts = new int[]{WALLPAPER_TAB_INDEX, PRESET_TAB_INDEX}.length;
-        private int mNumColors;
-
-        @Override
-        public int getItemCount() {
-            return mItemCounts;
-        }
-
-        @Override
-        public void onBindViewHolder(ColorPageViewHolder viewHolder, int position) {
-            switch (position) {
-                case WALLPAPER_TAB_INDEX:
-                    setupColorPages(viewHolder.mContainer, mNumColors, position,
-                            mWallpaperColorOptions, viewHolder.mPageIndicator);
-                    break;
-                case PRESET_TAB_INDEX:
-                    setupColorPages(viewHolder.mContainer, mNumColors, position,
-                            mPresetColorOptions, viewHolder.mPageIndicator);
-                    break;
-                default:
-                    break;
-            }
-        }
-
-        @Override
-        public ColorPageViewHolder onCreateViewHolder(ViewGroup viewGroup, int viewType) {
-            return new ColorPageViewHolder(LayoutInflater.from(viewGroup.getContext()).inflate(
-                    viewType, viewGroup, false));
-        }
-
-        @Override
-        public int getItemViewType(int position) {
-            return R.layout.color_pages_view;
-        }
-
-        public void setNumColors(int numColors) {
-            mNumColors = numColors;
-        }
-
-        private class ColorPageViewHolder extends RecyclerView.ViewHolder {
-            private ViewPager2 mContainer;
-            private PageIndicator mPageIndicator;
-
-            ColorPageViewHolder(View itemView) {
-                super(itemView);
-                mContainer = itemView.findViewById(R.id.color_page_container);
-                // Correct scrolling goes under collapsing toolbar while scrolling oclor options.
-                mContainer.getChildAt(0).setNestedScrollingEnabled(false);
-                mPageIndicator = itemView.findViewById(R.id.color_page_indicator);
-                if (ColorProvider.themeStyleEnabled) {
-                    mPageIndicator.setVisibility(VISIBLE);
-                }
-                itemView.setAccessibilityDelegate(createAccessibilityDelegate(ID_ITEMVIEW));
-                mContainer.setAccessibilityDelegate(createAccessibilityDelegate(ID_CONTAINER));
-            }
-        }
-    }
-
-    private class ColorPageAdapter extends
-            RecyclerView.Adapter<ColorPageAdapter.ColorOptionViewHolder> {
-
-        private final boolean mPageEnabled;
-        private final List<ColorOption> mColorOptions;
-        private final int mColorsPerPage;
-
-        private ColorPageAdapter(List<ColorOption> colorOptions, boolean pageEnabled,
-                int colorsPerPage) {
-            mPageEnabled = pageEnabled;
-            mColorOptions = colorOptions;
-            mColorsPerPage = colorsPerPage;
-        }
-
-        @Override
-        public int getItemCount() {
-            if (!mPageEnabled || !ColorProvider.themeStyleEnabled) {
-                return 1;
-            }
-            // Color page size.
-            return getNumPages(mColorsPerPage, mColorOptions.size());
-        }
-
-        @Override
-        public void onBindViewHolder(ColorOptionViewHolder viewHolder, int position) {
-            setupColorOptions(viewHolder.mContainer, mColorOptions, mPageEnabled, position,
-                    mColorsPerPage);
-        }
-
-        @Override
-        public ColorOptionViewHolder onCreateViewHolder(ViewGroup viewGroup, int viewType) {
-            return new ColorOptionViewHolder(
-                    LayoutInflater.from(viewGroup.getContext()).inflate(viewType, viewGroup,
-                            false));
-        }
-
-        @Override
-        public int getItemViewType(int position) {
-            return R.layout.color_options_view;
-        }
-
-        private class ColorOptionViewHolder extends RecyclerView.ViewHolder {
-            private RecyclerView mContainer;
-
-            ColorOptionViewHolder(View itemView) {
-                super(itemView);
-                mContainer = itemView.findViewById(R.id.color_option_container);
-                // Sets layout with margins to separate color options.
-                final FrameLayout.LayoutParams layoutParams = new FrameLayout.LayoutParams(
-                        mContainer.getLayoutParams());
-                final int margin = itemView.getContext().getResources().getDimensionPixelSize(
-                        R.dimen.section_horizontal_padding);
-                layoutParams.setMargins(margin, /* top= */ 0, margin, /* bottom= */ 0);
-                mContainer.setLayoutParams(layoutParams);
-            }
-        }
-    }
-}
diff --git a/src/com/android/customization/model/color/ColorUtils.kt b/src/com/android/customization/model/color/ColorUtils.kt
deleted file mode 100644
index 9d925e1..0000000
--- a/src/com/android/customization/model/color/ColorUtils.kt
+++ /dev/null
@@ -1,68 +0,0 @@
-/*
- * Copyright (C) 2022 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.customization.model.color
-
-import android.content.Context
-import android.content.pm.PackageManager
-import android.content.res.Resources
-import android.os.SystemProperties
-import android.util.Log
-import androidx.annotation.ColorInt
-
-/** Utility to wrap Monet's color extraction */
-object ColorUtils {
-    private const val TAG = "ColorUtils"
-    private const val MONET_FLAG = "flag_monet"
-    private var sSysuiRes: Resources? = null
-    private var sFlagId = 0
-
-    /** Returns true if color extraction is enabled in systemui. */
-    @JvmStatic
-    fun isMonetEnabled(context: Context): Boolean {
-        var monetEnabled = SystemProperties.getBoolean("persist.systemui.flag_monet", false)
-        if (!monetEnabled) {
-            if (sSysuiRes == null) {
-                try {
-                    val pm = context.packageManager
-                    val sysUIInfo =
-                        pm.getApplicationInfo(
-                            "com.android.systemui",
-                            PackageManager.GET_META_DATA or PackageManager.MATCH_SYSTEM_ONLY
-                        )
-                    if (sysUIInfo != null) {
-                        sSysuiRes = pm.getResourcesForApplication(sysUIInfo)
-                    }
-                } catch (e: PackageManager.NameNotFoundException) {
-                    Log.w(TAG, "Couldn't read color flag, skipping section", e)
-                }
-            }
-            if (sFlagId == 0) {
-                sFlagId =
-                    if (sSysuiRes == null) 0
-                    else sSysuiRes!!.getIdentifier(MONET_FLAG, "bool", "com.android.systemui")
-            }
-            if (sFlagId > 0) {
-                monetEnabled = sSysuiRes!!.getBoolean(sFlagId)
-            }
-        }
-        return monetEnabled
-    }
-
-    @JvmStatic
-    fun toColorString(@ColorInt color: Int): String {
-        return String.format("%06X", 0xFFFFFF and color)
-    }
-}
diff --git a/src/com/android/customization/model/color/WallpaperColorResources.java b/src/com/android/customization/model/color/WallpaperColorResources.java
deleted file mode 100644
index dc3b903..0000000
--- a/src/com/android/customization/model/color/WallpaperColorResources.java
+++ /dev/null
@@ -1,52 +0,0 @@
-/*
- * Copyright (C) 2022 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.customization.model.color;
-
-import android.app.WallpaperColors;
-import android.content.Context;
-import android.util.SparseIntArray;
-import android.widget.RemoteViews.ColorResources;
-
-import com.android.systemui.monet.ColorScheme;
-import com.android.systemui.monet.TonalPalette;
-
-/** A class to override colors in a {@link Context} with wallpaper colors. */
-public class WallpaperColorResources {
-
-    private final SparseIntArray mColorOverlay = new SparseIntArray();
-
-    public WallpaperColorResources(WallpaperColors wallpaperColors) {
-        ColorScheme wallpaperColorScheme = new ColorScheme(wallpaperColors, /* darkTheme= */ false);
-        addOverlayColor(wallpaperColorScheme.getNeutral1(), android.R.color.system_neutral1_10);
-        addOverlayColor(wallpaperColorScheme.getNeutral2(), android.R.color.system_neutral2_10);
-        addOverlayColor(wallpaperColorScheme.getAccent1(), android.R.color.system_accent1_10);
-        addOverlayColor(wallpaperColorScheme.getAccent2(), android.R.color.system_accent2_10);
-        addOverlayColor(wallpaperColorScheme.getAccent3(), android.R.color.system_accent3_10);
-    }
-
-    /** Applies the wallpaper color resources to the {@code context}. */
-    public void apply(Context context) {
-        ColorResources.create(context, mColorOverlay).apply(context);
-    }
-
-    private void addOverlayColor(TonalPalette colorSchemehue, int firstResourceColorId) {
-        int resourceColorId = firstResourceColorId;
-        for (int color : colorSchemehue.getAllShades()) {
-            mColorOverlay.put(resourceColorId, color);
-            resourceColorId++;
-        }
-    }
-}
diff --git a/src/com/android/customization/model/grid/GridOption.java b/src/com/android/customization/model/grid/GridOption.java
index 19c5d4f..a5307c9 100644
--- a/src/com/android/customization/model/grid/GridOption.java
+++ b/src/com/android/customization/model/grid/GridOption.java
@@ -48,19 +48,19 @@
         }
     };
 
-    private final String mTitle;
-    private final boolean mIsCurrent;
     private final String mIconShapePath;
     private final GridTileDrawable mTileDrawable;
+    public final String title;
     public final String name;
     public final int rows;
     public final int cols;
     public final Uri previewImageUri;
     public final int previewPagesCount;
+    private boolean mIsCurrent;
 
     public GridOption(String title, String name, boolean isCurrent, int rows, int cols,
             Uri previewImageUri, int previewPagesCount, String iconShapePath) {
-        mTitle = title;
+        this.title = title;
         mIsCurrent = isCurrent;
         mIconShapePath = iconShapePath;
         mTileDrawable = new GridTileDrawable(rows, cols, mIconShapePath);
@@ -71,8 +71,12 @@
         this.previewPagesCount = previewPagesCount;
     }
 
+    public void setIsCurrent(boolean isCurrent) {
+        mIsCurrent = isCurrent;
+    }
+
     protected GridOption(Parcel in) {
-        mTitle = in.readString();
+        title = in.readString();
         mIsCurrent = in.readByte() != 0;
         mIconShapePath = in.readString();
         name = in.readString();
@@ -85,7 +89,7 @@
 
     @Override
     public String getTitle() {
-        return mTitle;
+        return title;
     }
 
     @Override
@@ -139,7 +143,7 @@
 
     @Override
     public void writeToParcel(Parcel parcel, int i) {
-        parcel.writeString(mTitle);
+        parcel.writeString(title);
         parcel.writeByte((byte) (mIsCurrent ? 1 : 0));
         parcel.writeString(mIconShapePath);
         parcel.writeString(name);
@@ -154,7 +158,7 @@
         return String.format(
                 "GridOption{mTitle='%s', mIsCurrent=%s, mTileDrawable=%s, name='%s', rows=%d, "
                         + "cols=%d, previewImageUri=%s, previewPagesCount=%d}\n",
-                mTitle, mIsCurrent, mTileDrawable, name, rows, cols, previewImageUri,
+                title, mIsCurrent, mTileDrawable, name, rows, cols, previewImageUri,
                 previewPagesCount);
     }
 }
diff --git a/src/com/android/customization/model/grid/GridOptionsManager.java b/src/com/android/customization/model/grid/GridOptionsManager.java
index b7ee37f..78dbb5b 100644
--- a/src/com/android/customization/model/grid/GridOptionsManager.java
+++ b/src/com/android/customization/model/grid/GridOptionsManager.java
@@ -99,6 +99,11 @@
     }
 
     @Override
+    public void preview(GridOption option) {
+        mProvider.updateView();
+    }
+
+    @Override
     public void fetchOptions(OptionsFetchedListener<GridOption> callback, boolean reload) {
         sExecutorService.submit(() -> {
             List<GridOption> gridOptions = mProvider.fetch(reload);
diff --git a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
index 4e775c6..e71cca9 100644
--- a/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
+++ b/src/com/android/customization/model/grid/LauncherGridOptionsProvider.java
@@ -33,6 +33,7 @@
 
 import com.android.customization.model.ResourceConstants;
 import com.android.wallpaper.R;
+import com.android.wallpaper.config.BaseFlags;
 import com.android.wallpaper.util.PreviewUtils;
 
 import java.util.ArrayList;
@@ -57,12 +58,14 @@
 
     private final Context mContext;
     private final PreviewUtils mPreviewUtils;
+    private final boolean mIsGridApplyButtonEnabled;
     private List<GridOption> mOptions;
     private OptionChangeLiveData mLiveData;
 
     public LauncherGridOptionsProvider(Context context, String authorityMetadataKey) {
         mPreviewUtils = new PreviewUtils(context, authorityMetadataKey);
         mContext = context;
+        mIsGridApplyButtonEnabled = BaseFlags.get().isGridApplyButtonEnabled(context);
     }
 
     boolean areGridsAvailable() {
@@ -117,9 +120,14 @@
         mPreviewUtils.renderPreview(bundle, callback);
     }
 
+    void updateView() {
+        mLiveData.postValue(new Object());
+    }
+
     int applyGrid(String name) {
         ContentValues values = new ContentValues();
         values.put("name", name);
+        values.put("enable_apply_button", mIsGridApplyButtonEnabled);
         return mContext.getContentResolver().update(mPreviewUtils.getUri(DEFAULT_GRID), values,
                 null, null);
     }
@@ -153,6 +161,14 @@
             mContentObserver = new ContentObserver(handler) {
                 @Override
                 public void onChange(boolean selfChange) {
+                    // If grid apply button is enabled, user has previewed the grid before applying
+                    // the grid change. Thus there is no need to preview again (which will cause a
+                    // blank preview as launcher's is loader thread is busy reloading workspace)
+                    // after applying grid change. Thus we should ignore ContentObserver#onChange
+                    // from launcher
+                    if (BaseFlags.get().isGridApplyButtonEnabled(context.getApplicationContext())) {
+                        return;
+                    }
                     postValue(new Object());
                 }
             };
diff --git a/src/com/android/customization/model/grid/data/repository/GridRepository.kt b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
index 9a3be0c..4379dad 100644
--- a/src/com/android/customization/model/grid/data/repository/GridRepository.kt
+++ b/src/com/android/customization/model/grid/data/repository/GridRepository.kt
@@ -19,6 +19,7 @@
 
 import androidx.lifecycle.asFlow
 import com.android.customization.model.CustomizationManager
+import com.android.customization.model.CustomizationManager.Callback
 import com.android.customization.model.grid.GridOption
 import com.android.customization.model.grid.GridOptionsManager
 import com.android.customization.model.grid.shared.model.GridOptionItemModel
@@ -38,12 +39,17 @@
     suspend fun isAvailable(): Boolean
     fun getOptionChanges(): Flow<Unit>
     suspend fun getOptions(): GridOptionItemsModel
+    fun getSelectedOption(): GridOption?
+    fun applySelectedOption(callback: Callback)
+    fun clearSelectedOption()
+    fun isSelectedOptionApplied(): Boolean
 }
 
 class GridRepositoryImpl(
     private val applicationScope: CoroutineScope,
     private val manager: GridOptionsManager,
     private val backgroundDispatcher: CoroutineDispatcher,
+    private val isGridApplyButtonEnabled: Boolean,
 ) : GridRepository {
 
     override suspend fun isAvailable(): Boolean {
@@ -55,6 +61,10 @@
 
     private val selectedOption = MutableStateFlow<GridOption?>(null)
 
+    private var appliedOption: GridOption? = null
+
+    override fun getSelectedOption() = selectedOption.value
+
     override suspend fun getOptions(): GridOptionItemsModel {
         return withContext(backgroundDispatcher) {
             suspendCancellableCoroutine { continuation ->
@@ -62,7 +72,14 @@
                     object : CustomizationManager.OptionsFetchedListener<GridOption> {
                         override fun onOptionsLoaded(options: MutableList<GridOption>?) {
                             val optionsOrEmpty = options ?: emptyList()
-                            selectedOption.value = optionsOrEmpty.find { it.isActive(manager) }
+                            // After Apply Button is added, we will rely on onSelected() method
+                            // to update selectedOption.
+                            if (!isGridApplyButtonEnabled || selectedOption.value == null) {
+                                selectedOption.value = optionsOrEmpty.find { it.isActive(manager) }
+                            }
+                            if (isGridApplyButtonEnabled && appliedOption == null) {
+                                appliedOption = selectedOption.value
+                            }
                             continuation.resume(
                                 GridOptionItemsModel.Loaded(
                                     optionsOrEmpty.map { option -> toModel(option) }
@@ -105,22 +122,59 @@
     private suspend fun onSelected(option: GridOption) {
         withContext(backgroundDispatcher) {
             suspendCancellableCoroutine { continuation ->
-                manager.apply(
-                    option,
-                    object : CustomizationManager.Callback {
-                        override fun onSuccess() {
-                            continuation.resume(true)
-                        }
+                if (isGridApplyButtonEnabled) {
+                    selectedOption.value?.setIsCurrent(false)
+                    selectedOption.value = option
+                    selectedOption.value?.setIsCurrent(true)
+                    manager.preview(option)
+                    continuation.resume(true)
+                } else {
+                    manager.apply(
+                        option,
+                        object : CustomizationManager.Callback {
+                            override fun onSuccess() {
+                                continuation.resume(true)
+                            }
 
-                        override fun onError(throwable: Throwable?) {
-                            continuation.resume(false)
-                        }
-                    },
-                )
+                            override fun onError(throwable: Throwable?) {
+                                continuation.resume(false)
+                            }
+                        },
+                    )
+                }
             }
         }
     }
 
+    override fun applySelectedOption(callback: Callback) {
+        val option = getSelectedOption()
+        manager.apply(
+            option,
+            if (isGridApplyButtonEnabled) {
+                object : Callback {
+                    override fun onSuccess() {
+                        callback.onSuccess()
+                        appliedOption = option
+                    }
+
+                    override fun onError(throwable: Throwable?) {
+                        callback.onError(throwable)
+                    }
+                }
+            } else callback
+        )
+    }
+
+    override fun clearSelectedOption() {
+        if (!isGridApplyButtonEnabled) {
+            return
+        }
+        selectedOption.value?.setIsCurrent(false)
+        selectedOption.value = null
+    }
+
+    override fun isSelectedOptionApplied() = selectedOption.value?.name == appliedOption?.name
+
     private fun GridOption?.key(): String? {
         return if (this != null) "${cols}x${rows}" else null
     }
diff --git a/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
index cdb679d..7abd605 100644
--- a/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
+++ b/src/com/android/customization/model/grid/domain/interactor/GridInteractor.kt
@@ -17,6 +17,8 @@
 
 package com.android.customization.model.grid.domain.interactor
 
+import com.android.customization.model.CustomizationManager
+import com.android.customization.model.grid.GridOption
 import com.android.customization.model.grid.data.repository.GridRepository
 import com.android.customization.model.grid.shared.model.GridOptionItemModel
 import com.android.customization.model.grid.shared.model.GridOptionItemsModel
@@ -73,6 +75,16 @@
         }
     }
 
+    fun getSelectOptionNonSuspend(): GridOption? = repository.getSelectedOption()
+
+    fun clearSelectedOption() = repository.clearSelectedOption()
+
+    fun isSelectedOptionApplied() = repository.isSelectedOptionApplied()
+
+    fun applySelectedOption(callback: CustomizationManager.Callback) {
+        repository.applySelectedOption(callback)
+    }
+
     private suspend fun reload(): GridOptionItemsModel {
         val model = repository.getOptions()
         return if (model is GridOptionItemsModel.Loaded) {
diff --git a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
index 78536ca..56fe425 100644
--- a/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
+++ b/src/com/android/customization/model/grid/ui/binder/GridScreenBinder.kt
@@ -18,6 +18,7 @@
 package com.android.customization.model.grid.ui.binder
 
 import android.view.View
+import android.widget.Button
 import android.widget.ImageView
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleOwner
@@ -41,6 +42,8 @@
         lifecycleOwner: LifecycleOwner,
         backgroundDispatcher: CoroutineDispatcher,
         onOptionsChanged: () -> Unit,
+        isGridApplyButtonEnabled: Boolean,
+        onOptionApplied: () -> Unit,
     ) {
         val optionView: RecyclerView = view.requireViewById(R.id.options)
         optionView.layoutManager =
@@ -57,8 +60,8 @@
                 backgroundDispatcher = backgroundDispatcher,
                 foregroundTintSpec =
                     OptionItemBinder.TintSpec(
-                        selectedColor = view.context.getColor(R.color.text_color_primary),
-                        unselectedColor = view.context.getColor(R.color.text_color_secondary),
+                        selectedColor = view.context.getColor(R.color.system_on_surface),
+                        unselectedColor = view.context.getColor(R.color.system_on_surface),
                     ),
                 bindIcon = { foregroundView: View, gridIcon: GridIconViewModel ->
                     val imageView = foregroundView as? ImageView
@@ -67,6 +70,13 @@
             )
         optionView.adapter = adapter
 
+        if (isGridApplyButtonEnabled) {
+            val applyButton: Button = view.requireViewById(R.id.apply_button)
+            applyButton.visibility = View.VISIBLE
+            view.requireViewById<View>(R.id.apply_button_note).visibility = View.VISIBLE
+            applyButton.setOnClickListener { onOptionApplied() }
+        }
+
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
                 launch {
diff --git a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
index 6ab561f..9e99efe 100644
--- a/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
+++ b/src/com/android/customization/model/grid/ui/fragment/GridFragment2.kt
@@ -18,14 +18,24 @@
 package com.android.customization.model.grid.ui.fragment
 
 import android.os.Bundle
+import android.util.Log
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import android.widget.Button
+import android.widget.Toast
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
+import androidx.transition.Transition
+import androidx.transition.doOnStart
+import com.android.customization.model.CustomizationManager.Callback
+import com.android.customization.model.grid.domain.interactor.GridInteractor
 import com.android.customization.model.grid.ui.binder.GridScreenBinder
 import com.android.customization.model.grid.ui.viewmodel.GridScreenViewModel
 import com.android.customization.module.ThemePickerInjector
 import com.android.wallpaper.R
+import com.android.wallpaper.config.BaseFlags
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.module.InjectorProvider
@@ -38,9 +48,13 @@
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.suspendCancellableCoroutine
 
+private val TAG = GridFragment2::class.java.simpleName
+
 @OptIn(ExperimentalCoroutinesApi::class)
 class GridFragment2 : AppbarFragment() {
 
+    private lateinit var gridInteractor: GridInteractor
+
     override fun onCreateView(
         inflater: LayoutInflater,
         container: ViewGroup?,
@@ -54,6 +68,8 @@
             )
         setUpToolbar(view)
 
+        val isGridApplyButtonEnabled = BaseFlags.get().isGridApplyButtonEnabled(requireContext())
+
         val injector = InjectorProvider.getInjector() as ThemePickerInjector
 
         val wallpaperInfoFactory = injector.getCurrentWallpaperInfoFactory(requireContext())
@@ -61,10 +77,12 @@
             bindScreenPreview(
                 view,
                 wallpaperInfoFactory,
-                injector.getWallpaperInteractor(requireContext())
+                injector.getWallpaperInteractor(requireContext()),
+                injector.getGridInteractor(requireContext())
             )
 
         val viewModelFactory = injector.getGridScreenViewModelFactory(requireContext())
+        gridInteractor = injector.getGridInteractor(requireContext())
         GridScreenBinder.bind(
             view = view,
             viewModel =
@@ -80,11 +98,50 @@
                     bindScreenPreview(
                         view,
                         wallpaperInfoFactory,
-                        injector.getWallpaperInteractor(requireContext())
+                        injector.getWallpaperInteractor(requireContext()),
+                        gridInteractor,
                     )
+                if (isGridApplyButtonEnabled) {
+                    val applyButton: Button = view.requireViewById(R.id.apply_button)
+                    applyButton.isEnabled = !gridInteractor.isSelectedOptionApplied()
+                }
+            },
+            isGridApplyButtonEnabled = isGridApplyButtonEnabled,
+            onOptionApplied = {
+                gridInteractor.applySelectedOption(
+                    object : Callback {
+                        override fun onSuccess() {
+                            Toast.makeText(
+                                    context,
+                                    getString(
+                                        R.string.toast_of_changing_grid,
+                                        gridInteractor.getSelectOptionNonSuspend()?.title
+                                    ),
+                                    Toast.LENGTH_SHORT
+                                )
+                                .show()
+                            val applyButton: Button = view.requireViewById(R.id.apply_button)
+                            applyButton.isEnabled = false
+                        }
+
+                        override fun onError(throwable: Throwable?) {
+                            val errorMsg =
+                                getString(
+                                    R.string.toast_of_failure_to_change_grid,
+                                    gridInteractor.getSelectOptionNonSuspend()?.title
+                                )
+                            Toast.makeText(context, errorMsg, Toast.LENGTH_SHORT).show()
+                            Log.e(TAG, errorMsg, throwable)
+                        }
+                    }
+                )
             }
         )
 
+        (returnTransition as? Transition)?.doOnStart {
+            view.requireViewById<View>(R.id.preview).isVisible = false
+        }
+
         return view
     }
 
@@ -92,10 +149,19 @@
         return getString(R.string.grid_title)
     }
 
+    override fun getToolbarColorId(): Int {
+        return android.R.color.transparent
+    }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
+
     private fun bindScreenPreview(
         view: View,
         wallpaperInfoFactory: CurrentWallpaperInfoFactory,
         wallpaperInteractor: WallpaperInteractor,
+        gridInteractor: GridInteractor
     ): ScreenPreviewBinder.Binding {
         return ScreenPreviewBinder.bind(
             activity = requireActivity(),
@@ -111,6 +177,11 @@
                                         R.string.grid_control_metadata_name,
                                     ),
                         ),
+                    initialExtrasProvider = {
+                        val bundle = Bundle()
+                        bundle.putString("name", gridInteractor.getSelectOptionNonSuspend()?.name)
+                        bundle
+                    },
                     wallpaperInfoProvider = {
                         suspendCancellableCoroutine { continuation ->
                             wallpaperInfoFactory.createCurrentWallpaperInfos(
@@ -129,4 +200,11 @@
             onWallpaperPreviewDirty = { activity?.recreate() },
         )
     }
+
+    override fun onBackPressed(): Boolean {
+        if (BaseFlags.get().isGridApplyButtonEnabled(requireContext())) {
+            gridInteractor.clearSelectedOption()
+        }
+        return super.onBackPressed()
+    }
 }
diff --git a/src/com/android/customization/model/theme/ThemeBundledWallpaperInfo.java b/src/com/android/customization/model/theme/ThemeBundledWallpaperInfo.java
index d3935e8..4f0cd6c 100644
--- a/src/com/android/customization/model/theme/ThemeBundledWallpaperInfo.java
+++ b/src/com/android/customization/model/theme/ThemeBundledWallpaperInfo.java
@@ -137,13 +137,13 @@
 
     @Override
     public void showPreview(Activity srcActivity, InlinePreviewIntentFactory factory,
-            int requestCode) {
+            int requestCode, boolean isAssetIdPresent) {
         try {
-            srcActivity.startActivityForResult(factory.newIntent(srcActivity, this), requestCode);
+            srcActivity.startActivityForResult(factory.newIntent(srcActivity, this,
+                    isAssetIdPresent), requestCode);
         } catch (ActivityNotFoundException |SecurityException e) {
             Log.e(TAG, "App isn't installed or ThemePicker doesn't have permission to launch", e);
         }
-
     }
 
     @Override
diff --git a/src/com/android/customization/module/CustomizationInjector.kt b/src/com/android/customization/module/CustomizationInjector.kt
index 5f8f9d3..8fd3768 100644
--- a/src/com/android/customization/module/CustomizationInjector.kt
+++ b/src/com/android/customization/module/CustomizationInjector.kt
@@ -16,6 +16,7 @@
 package com.android.customization.module
 
 import android.content.Context
+import android.content.res.Resources
 import androidx.activity.ComponentActivity
 import androidx.fragment.app.FragmentActivity
 import com.android.customization.model.theme.OverlayManagerCompat
@@ -78,5 +79,5 @@
         clockViewFactory: ClockViewFactory,
     ): ClockSettingsViewModel.Factory
 
-    fun getClockDescriptionUtils(): ClockDescriptionUtils
+    fun getClockDescriptionUtils(resources: Resources): ClockDescriptionUtils
 }
diff --git a/src/com/android/customization/module/DefaultCustomizationSections.java b/src/com/android/customization/module/DefaultCustomizationSections.java
index b408a89..bbe6bef 100644
--- a/src/com/android/customization/module/DefaultCustomizationSections.java
+++ b/src/com/android/customization/module/DefaultCustomizationSections.java
@@ -8,10 +8,8 @@
 import androidx.lifecycle.LifecycleOwner;
 import androidx.lifecycle.ViewModelProvider;
 
-import com.android.customization.model.color.ColorSectionController;
 import com.android.customization.model.grid.GridOptionsManager;
 import com.android.customization.model.grid.GridSectionController;
-import com.android.customization.model.mode.DarkModeSectionController;
 import com.android.customization.model.mode.DarkModeSnapshotRestorer;
 import com.android.customization.model.themedicon.ThemedIconSectionController;
 import com.android.customization.model.themedicon.ThemedIconSwitchProvider;
@@ -19,6 +17,7 @@
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconSnapshotRestorer;
 import com.android.customization.picker.clock.ui.view.ClockViewFactory;
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel;
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor;
 import com.android.customization.picker.color.ui.section.ColorSectionController2;
 import com.android.customization.picker.color.ui.viewmodel.ColorPickerViewModel;
 import com.android.customization.picker.notifications.ui.section.NotificationSectionController;
@@ -35,7 +34,6 @@
 import com.android.wallpaper.model.PermissionRequester;
 import com.android.wallpaper.model.WallpaperColorsViewModel;
 import com.android.wallpaper.model.WallpaperPreviewNavigator;
-import com.android.wallpaper.model.WallpaperSectionController;
 import com.android.wallpaper.module.CurrentWallpaperInfoFactory;
 import com.android.wallpaper.module.CustomizationSections;
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor;
@@ -61,6 +59,7 @@
     private final DarkModeSnapshotRestorer mDarkModeSnapshotRestorer;
     private final ThemedIconSnapshotRestorer mThemedIconSnapshotRestorer;
     private final ThemedIconInteractor mThemedIconInteractor;
+    private final ColorPickerInteractor mColorPickerInteractor;
 
     public DefaultCustomizationSections(
             ColorPickerViewModel.Factory colorPickerViewModelFactory,
@@ -73,7 +72,8 @@
             ClockViewFactory clockViewFactory,
             DarkModeSnapshotRestorer darkModeSnapshotRestorer,
             ThemedIconSnapshotRestorer themedIconSnapshotRestorer,
-            ThemedIconInteractor themedIconInteractor) {
+            ThemedIconInteractor themedIconInteractor,
+            ColorPickerInteractor colorPickerInteractor) {
         mColorPickerViewModelFactory = colorPickerViewModelFactory;
         mKeyguardQuickAffordancePickerInteractor = keyguardQuickAffordancePickerInteractor;
         mKeyguardQuickAffordancePickerViewModelFactory =
@@ -85,10 +85,11 @@
         mDarkModeSnapshotRestorer = darkModeSnapshotRestorer;
         mThemedIconSnapshotRestorer = themedIconSnapshotRestorer;
         mThemedIconInteractor = themedIconInteractor;
+        mColorPickerInteractor = colorPickerInteractor;
     }
 
     @Override
-    public List<CustomizationSectionController<?>> getRevampedUISectionControllersForScreen(
+    public List<CustomizationSectionController<?>> getSectionControllersForScreen(
             Screen screen,
             FragmentActivity activity,
             LifecycleOwner lifecycleOwner,
@@ -121,8 +122,10 @@
                         sectionNavigationController,
                         wallpaperInteractor,
                         mThemedIconInteractor,
+                        mColorPickerInteractor,
                         wallpaperManager,
-                        isTwoPaneAndSmallWidth)
+                        isTwoPaneAndSmallWidth,
+                        customizationPickerViewModel)
                         : new PreviewWithThemeSectionController(
                                 activity,
                                 lifecycleOwner,
@@ -133,8 +136,10 @@
                                 wallpaperPreviewNavigator,
                                 wallpaperInteractor,
                                 mThemedIconInteractor,
+                                mColorPickerInteractor,
                                 wallpaperManager,
-                                isTwoPaneAndSmallWidth));
+                                isTwoPaneAndSmallWidth,
+                                customizationPickerViewModel));
 
         sectionControllers.add(
                 new ConnectedSectionController(
@@ -202,57 +207,4 @@
 
         return sectionControllers;
     }
-
-    @Override
-    public List<CustomizationSectionController<?>> getAllSectionControllers(
-            FragmentActivity activity,
-            LifecycleOwner lifecycleOwner,
-            WallpaperColorsViewModel wallpaperColorsViewModel,
-            PermissionRequester permissionRequester,
-            WallpaperPreviewNavigator wallpaperPreviewNavigator,
-            CustomizationSectionNavigationController sectionNavigationController,
-            @Nullable Bundle savedInstanceState,
-            DisplayUtils displayUtils) {
-        List<CustomizationSectionController<?>> sectionControllers = new ArrayList<>();
-
-        // Wallpaper section.
-        sectionControllers.add(
-                new WallpaperSectionController(
-                activity,
-                lifecycleOwner,
-                permissionRequester,
-                wallpaperColorsViewModel,
-                mThemedIconInteractor.isActivatedAsLiveData(),
-                sectionNavigationController,
-                wallpaperPreviewNavigator,
-                savedInstanceState,
-                displayUtils));
-
-        // Theme color section.
-        sectionControllers.add(new ColorSectionController(
-                activity, wallpaperColorsViewModel, lifecycleOwner, savedInstanceState));
-
-        // Dark/Light theme section.
-        sectionControllers.add(new DarkModeSectionController(
-                activity,
-                lifecycleOwner.getLifecycle(),
-                mDarkModeSnapshotRestorer));
-
-        // Themed app icon section.
-        sectionControllers.add(new ThemedIconSectionController(
-                ThemedIconSwitchProvider.getInstance(activity),
-                mThemedIconInteractor,
-                savedInstanceState,
-                mThemedIconSnapshotRestorer));
-
-        // App grid section.
-        sectionControllers.add(
-                new GridSectionController(
-                        GridOptionsManager.getInstance(activity),
-                        sectionNavigationController,
-                        lifecycleOwner,
-                        /* isRevampedUiEnabled= */ false));
-
-        return sectionControllers;
-    }
 }
diff --git a/src/com/android/customization/module/StatsLogUserEventLogger.java b/src/com/android/customization/module/StatsLogUserEventLogger.java
index 32e0599..647cdc9 100644
--- a/src/com/android/customization/module/StatsLogUserEventLogger.java
+++ b/src/com/android/customization/module/StatsLogUserEventLogger.java
@@ -37,7 +37,6 @@
 import static com.android.wallpaper.util.LaunchSourceUtils.WALLPAPER_LAUNCH_SOURCE;
 
 import android.app.WallpaperManager;
-import android.content.Context;
 import android.content.Intent;
 import android.stats.style.StyleEnums;
 import android.text.TextUtils;
@@ -47,8 +46,6 @@
 import com.android.customization.model.color.ColorOption;
 import com.android.customization.model.grid.GridOption;
 import com.android.customization.model.theme.ThemeBundle;
-import com.android.wallpaper.module.Injector;
-import com.android.wallpaper.module.InjectorProvider;
 import com.android.wallpaper.module.NoOpUserEventLogger;
 import com.android.wallpaper.module.WallpaperPreferences;
 import com.android.wallpaper.module.WallpaperStatusChecker;
@@ -61,86 +58,59 @@
  */
 public class StatsLogUserEventLogger extends NoOpUserEventLogger implements ThemesUserEventLogger {
 
-    private static final String TAG = "StatsLogUserEventLogger";
-    private final Context mContext;
     private final WallpaperPreferences mPreferences;
     private final WallpaperStatusChecker mWallpaperStatusChecker;
 
-    public StatsLogUserEventLogger(Context appContext) {
-        mContext = appContext;
-        Injector injector = InjectorProvider.getInjector();
-        mPreferences = injector.getPreferences(appContext);
-        mWallpaperStatusChecker = injector.getWallpaperStatusChecker();
+    public StatsLogUserEventLogger(
+            WallpaperPreferences preferences,
+            WallpaperStatusChecker wallpaperStatusChecker) {
+        mPreferences = preferences;
+        mWallpaperStatusChecker = wallpaperStatusChecker;
     }
 
     @Override
     public void logAppLaunched(Intent launchSource) {
-        new SysUiStatsLogger()
-                .setAction(STYLE_UICHANGED__ACTION__APP_LAUNCHED)
+        new SysUiStatsLogger(STYLE_UICHANGED__ACTION__APP_LAUNCHED)
                 .setLaunchedPreference(getAppLaunchSource(launchSource))
                 .log();
     }
 
     @Override
     public void logResumed(boolean provisioned, boolean wallpaper) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.ONRESUME)
+        new SysUiStatsLogger(StyleEnums.ONRESUME)
                 .log();
     }
 
     @Override
     public void logStopped() {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.ONSTOP)
+        new SysUiStatsLogger(StyleEnums.ONSTOP)
                 .log();
     }
 
     @Override
     public void logActionClicked(String collectionId, int actionLabelResId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_EXPLORE)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EXPLORE)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .log();
     }
 
     @Override
     public void logIndividualWallpaperSelected(String collectionId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_SELECT)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_SELECT)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .log();
     }
 
     @Override
     public void logCategorySelected(String collectionId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_OPEN_CATEGORY)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_OPEN_CATEGORY)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .log();
     }
 
     @Override
-    public void logLiveWallpaperInfoSelected(String collectionId, @Nullable String wallpaperId) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.LIVE_WALLPAPER_INFO_SELECT)
-                .setWallpaperCategoryHash(getIdHashCode(collectionId))
-                .setWallpaperIdHash(getIdHashCode(wallpaperId))
-                .log();
-    }
-
-    @Override
-    public void logLiveWallpaperCustomizeSelected(String collectionId,
-            @Nullable String wallpaperId) {
-        new SysUiStatsLogger().setAction(StyleEnums.LIVE_WALLPAPER_CUSTOMIZE_SELECT)
-                .setWallpaperCategoryHash(getIdHashCode(collectionId))
-                .setWallpaperIdHash(getIdHashCode(wallpaperId))
-                .log();
-
-    }
-
-    @Override
     public void logSnapshot() {
-        final boolean isLockWallpaperSet = mWallpaperStatusChecker.isLockWallpaperSet(mContext);
+        final boolean isLockWallpaperSet = mWallpaperStatusChecker.isLockWallpaperSet();
         final String homeCollectionId = mPreferences.getHomeWallpaperCollectionId();
         final String homeRemoteId = mPreferences.getHomeWallpaperRemoteId();
         final String effects = mPreferences.getHomeWallpaperEffects();
@@ -151,7 +121,7 @@
         String lockWallpaperId = isLockWallpaperSet ? mPreferences.getLockWallpaperRemoteId()
                 : homeWallpaperId;
 
-        new SysUiStatsLogger().setAction(StyleEnums.SNAPSHOT)
+        new SysUiStatsLogger(StyleEnums.SNAPSHOT)
                 .setWallpaperCategoryHash(getIdHashCode(homeCollectionId))
                 .setWallpaperIdHash(getIdHashCode(homeWallpaperId))
                 .setLockWallpaperCategoryHash(getIdHashCode(lockCollectionId))
@@ -167,8 +137,7 @@
     @Override
     public void logWallpaperSet(String collectionId, @Nullable String wallpaperId,
             @Nullable String effects) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_APPLIED)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_APPLIED)
                 .setWallpaperCategoryHash(getIdHashCode(collectionId))
                 .setWallpaperIdHash(getIdHashCode(wallpaperId))
                 .setEffectIdHash(getIdHashCode(effects))
@@ -178,8 +147,7 @@
     @Override
     public void logEffectApply(String effect, @EffectStatus int status, long timeElapsedMillis,
             int resultCode) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_EFFECT_APPLIED)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_APPLIED)
                 .setEffectPreference(status)
                 .setEffectIdHash(getIdHashCode(effect))
                 .setTimeElapsed(timeElapsedMillis)
@@ -189,13 +157,22 @@
 
     @Override
     public void logEffectProbe(String effect, @EffectStatus int status) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.WALLPAPER_EFFECT_PROBE)
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_PROBE)
                 .setEffectPreference(status)
                 .setEffectIdHash(getIdHashCode(effect))
                 .log();
     }
 
+    @Override
+    public void logEffectForegroundDownload(String effect, @EffectStatus int status,
+            long timeElapsedMillis) {
+        new SysUiStatsLogger(StyleEnums.WALLPAPER_EFFECT_FG_DOWNLOAD)
+                .setEffectPreference(status)
+                .setEffectIdHash(getIdHashCode(effect))
+                .setTimeElapsed(timeElapsedMillis)
+                .log();
+    }
+
     @Nullable
     private String getThemePackage(ThemeBundle theme, String category) {
         Map<String, String> packages = theme.getPackagesByCategory();
@@ -204,8 +181,7 @@
 
     @Override
     public void logThemeSelected(ThemeBundle theme, boolean isCustomTheme) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_SELECT)
+        new SysUiStatsLogger(StyleEnums.PICKER_SELECT)
                 .setColorPackageHash(
                         Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_COLOR)))
                 .setFontPackageHash(Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_FONT)))
@@ -216,8 +192,7 @@
 
     @Override
     public void logThemeApplied(ThemeBundle theme, boolean isCustomTheme) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_APPLIED)
+        new SysUiStatsLogger(StyleEnums.PICKER_APPLIED)
                 .setColorPackageHash(
                         Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_COLOR)))
                 .setFontPackageHash(Objects.hashCode(getThemePackage(theme, OVERLAY_CATEGORY_FONT)))
@@ -228,8 +203,7 @@
 
     @Override
     public void logColorApplied(int action, ColorOption colorOption) {
-        new SysUiStatsLogger()
-                .setAction(action)
+        new SysUiStatsLogger(action)
                 .setColorPreference(colorOption.getIndex())
                 .setColorVariant(colorOption.getStyle().ordinal() + 1)
                 .log();
@@ -237,16 +211,14 @@
 
     @Override
     public void logGridSelected(GridOption grid) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_SELECT)
+        new SysUiStatsLogger(StyleEnums.PICKER_SELECT)
                 .setLauncherGrid(grid.cols)
                 .log();
     }
 
     @Override
     public void logGridApplied(GridOption grid) {
-        new SysUiStatsLogger()
-                .setAction(StyleEnums.PICKER_APPLIED)
+        new SysUiStatsLogger(StyleEnums.PICKER_APPLIED)
                 .setLauncherGrid(grid.cols)
                 .log();
     }
diff --git a/src/com/android/customization/module/SysUiStatsLogger.kt b/src/com/android/customization/module/SysUiStatsLogger.kt
index 318bf1f..8e97b0b 100644
--- a/src/com/android/customization/module/SysUiStatsLogger.kt
+++ b/src/com/android/customization/module/SysUiStatsLogger.kt
@@ -20,10 +20,8 @@
 import com.android.systemui.shared.system.SysUiStatsLog.STYLE_UI_CHANGED
 
 /** The builder for [SysUiStatsLog]. */
-class SysUiStatsLogger {
+class SysUiStatsLogger(val action: Int) {
 
-    private var atom = STYLE_UI_CHANGED
-    private var action = StyleEnums.DEFAULT_ACTION
     private var colorPackageHash = 0
     private var fontPackageHash = 0
     private var shapePackageHash = 0
@@ -46,85 +44,83 @@
     private var timeElapsedMillis = 0L
     private var effectResultCode = -1
 
-    fun setAction(action: Int) = apply { this.action = action }
-
-    fun setColorPackageHash(color_package_hash: Int) = apply {
-        this.colorPackageHash = color_package_hash
+    fun setColorPackageHash(colorPackageHash: Int) = apply {
+        this.colorPackageHash = colorPackageHash
     }
 
-    fun setFontPackageHash(font_package_hash: Int) = apply {
-        this.fontPackageHash = font_package_hash
+    fun setFontPackageHash(fontPackageHash: Int) = apply {
+        this.fontPackageHash = fontPackageHash
     }
 
-    fun setShapePackageHash(shape_package_hash: Int) = apply {
-        this.shapePackageHash = shape_package_hash
+    fun setShapePackageHash(shapePackageHash: Int) = apply {
+        this.shapePackageHash = shapePackageHash
     }
 
-    fun setClockPackageHash(clock_package_hash: Int) = apply {
-        this.clockPackageHash = clock_package_hash
+    fun setClockPackageHash(clockPackageHash: Int) = apply {
+        this.clockPackageHash = clockPackageHash
     }
 
-    fun setLauncherGrid(launcher_grid: Int) = apply { this.launcherGrid = launcher_grid }
+    fun setLauncherGrid(launcherGrid: Int) = apply { this.launcherGrid = launcherGrid }
 
-    fun setWallpaperCategoryHash(wallpaper_category_hash: Int) = apply {
-        this.wallpaperCategoryHash = wallpaper_category_hash
+    fun setWallpaperCategoryHash(wallpaperCategoryHash: Int) = apply {
+        this.wallpaperCategoryHash = wallpaperCategoryHash
     }
 
-    fun setWallpaperIdHash(wallpaper_id_hash: Int) = apply {
-        this.wallpaperIdHash = wallpaper_id_hash
+    fun setWallpaperIdHash(wallpaperIdHash: Int) = apply {
+        this.wallpaperIdHash = wallpaperIdHash
     }
 
-    fun setColorPreference(color_preference: Int) = apply {
-        this.colorPreference = color_preference
+    fun setColorPreference(colorPreference: Int) = apply {
+        this.colorPreference = colorPreference
     }
 
-    fun setLocationPreference(location_preference: Int) = apply {
-        this.locationPreference = location_preference
+    fun setLocationPreference(locationPreference: Int) = apply {
+        this.locationPreference = locationPreference
     }
 
-    fun setDatePreference(date_preference: Int) = apply { this.datePreference = date_preference }
+    fun setDatePreference(datePreference: Int) = apply { this.datePreference = datePreference }
 
-    fun setLaunchedPreference(launched_preference: Int) = apply {
-        this.launchedPreference = launched_preference
+    fun setLaunchedPreference(launchedPreference: Int) = apply {
+        this.launchedPreference = launchedPreference
     }
 
-    fun setEffectPreference(effect_preference: Int) = apply {
-        this.effectPreference = effect_preference
+    fun setEffectPreference(effectPreference: Int) = apply {
+        this.effectPreference = effectPreference
     }
 
-    fun setEffectIdHash(effect_id_hash: Int) = apply { this.effectIdHash = effect_id_hash }
+    fun setEffectIdHash(effectIdHash: Int) = apply { this.effectIdHash = effectIdHash }
 
-    fun setLockWallpaperCategoryHash(lock_wallpaper_category_hash: Int) = apply {
-        this.lockWallpaperCategoryHash = lock_wallpaper_category_hash
+    fun setLockWallpaperCategoryHash(lockWallpaperCategoryHash: Int) = apply {
+        this.lockWallpaperCategoryHash = lockWallpaperCategoryHash
     }
 
-    fun setLockWallpaperIdHash(lock_wallpaper_id_hash: Int) = apply {
-        this.lockWallpaperIdHash = lock_wallpaper_id_hash
+    fun setLockWallpaperIdHash(lockWallpaperIdHash: Int) = apply {
+        this.lockWallpaperIdHash = lockWallpaperIdHash
     }
 
-    fun setFirstLaunchDateSinceSetup(first_launch_date_since_setup: Int) = apply {
-        this.firstLaunchDateSinceSetup = first_launch_date_since_setup
+    fun setFirstLaunchDateSinceSetup(firstLaunchDateSinceSetup: Int) = apply {
+        this.firstLaunchDateSinceSetup = firstLaunchDateSinceSetup
     }
 
-    fun setFirstWallpaperApplyDateSinceSetup(first_wallpaper_apply_date_since_setup: Int) = apply {
-        this.firstWallpaperApplyDateSinceSetup = first_wallpaper_apply_date_since_setup
+    fun setFirstWallpaperApplyDateSinceSetup(firstWallpaperApplyDateSinceSetup: Int) = apply {
+        this.firstWallpaperApplyDateSinceSetup = firstWallpaperApplyDateSinceSetup
     }
 
-    fun setAppLaunchCount(app_launch_count: Int) = apply { this.appLaunchCount = app_launch_count }
+    fun setAppLaunchCount(appLaunchCount: Int) = apply { this.appLaunchCount = appLaunchCount }
 
-    fun setColorVariant(color_variant: Int) = apply { this.colorVariant = color_variant }
+    fun setColorVariant(colorVariant: Int) = apply { this.colorVariant = colorVariant }
 
-    fun setTimeElapsed(time_elapsed_millis: Long) = apply {
-      this.timeElapsedMillis = time_elapsed_millis
+    fun setTimeElapsed(timeElapsedMillis: Long) = apply {
+      this.timeElapsedMillis = timeElapsedMillis
     }
 
-    fun setEffectResultCode(effect_result_code: Int) = apply {
-        this.effectResultCode = effect_result_code
+    fun setEffectResultCode(effectResultCode: Int) = apply {
+        this.effectResultCode = effectResultCode
     }
 
     fun log() {
         SysUiStatsLog.write(
-            atom,
+            STYLE_UI_CHANGED,
             action,
             colorPackageHash,
             fontPackageHash,
diff --git a/src/com/android/customization/module/ThemePickerInjector.kt b/src/com/android/customization/module/ThemePickerInjector.kt
index cceb896..1486884 100644
--- a/src/com/android/customization/module/ThemePickerInjector.kt
+++ b/src/com/android/customization/module/ThemePickerInjector.kt
@@ -19,6 +19,7 @@
 import android.app.WallpaperManager
 import android.content.Context
 import android.content.Intent
+import android.content.res.Resources
 import android.net.Uri
 import android.os.Bundle
 import android.text.TextUtils
@@ -30,6 +31,7 @@
 import androidx.lifecycle.ViewModelProvider
 import com.android.customization.model.color.ColorCustomizationManager
 import com.android.customization.model.color.ColorOptionsProvider
+import com.android.customization.model.color.ColorOptionsProvider.COLOR_SOURCE_PRESET
 import com.android.customization.model.grid.GridOptionsManager
 import com.android.customization.model.grid.data.repository.GridRepositoryImpl
 import com.android.customization.model.grid.domain.interactor.GridInteractor
@@ -68,6 +70,9 @@
 import com.android.systemui.shared.clocks.ClockRegistry
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClient
 import com.android.systemui.shared.customization.data.content.CustomizationProviderClientImpl
+import com.android.wallpaper.config.BaseFlags
+import com.android.wallpaper.dispatchers.BackgroundDispatcher
+import com.android.wallpaper.dispatchers.MainDispatcher
 import com.android.wallpaper.model.LiveWallpaperInfo
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import com.android.wallpaper.model.WallpaperInfo
@@ -75,22 +80,30 @@
 import com.android.wallpaper.module.FragmentFactory
 import com.android.wallpaper.module.UserEventLogger
 import com.android.wallpaper.module.WallpaperPicker2Injector
-import com.android.wallpaper.module.WallpaperPreferences
 import com.android.wallpaper.picker.CustomizationPickerActivity
-import com.android.wallpaper.picker.ImagePreviewFragment
-import com.android.wallpaper.picker.LivePreviewFragment
-import com.android.wallpaper.picker.PreviewFragment
+import com.android.wallpaper.picker.ImagePreviewFragment2
+import com.android.wallpaper.picker.LivePreviewFragment2
+import com.android.wallpaper.picker.PreviewFragment2
 import com.android.wallpaper.picker.customization.data.content.WallpaperClientImpl
 import com.android.wallpaper.picker.customization.data.repository.WallpaperRepository
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.undo.domain.interactor.SnapshotRestorer
 import com.android.wallpaper.util.ScreenSizeCalculator
-import kotlinx.coroutines.Dispatchers
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.CoroutineDispatcher
+import kotlinx.coroutines.CoroutineScope
 
-open class ThemePickerInjector : WallpaperPicker2Injector(), CustomizationInjector {
+@Singleton
+open class ThemePickerInjector
+@Inject
+internal constructor(
+    @MainDispatcher private val mainScope: CoroutineScope,
+    @MainDispatcher private val mainDispatcher: CoroutineDispatcher,
+    @BackgroundDispatcher private val bgDispatcher: CoroutineDispatcher,
+) : WallpaperPicker2Injector(mainScope, bgDispatcher), CustomizationInjector {
     private var customizationSections: CustomizationSections? = null
     private var userEventLogger: UserEventLogger? = null
-    private var prefs: WallpaperPreferences? = null
     private var wallpaperInteractor: WallpaperInteractor? = null
     private var keyguardQuickAffordancePickerInteractor: KeyguardQuickAffordancePickerInteractor? =
         null
@@ -143,6 +156,7 @@
                     getDarkModeSnapshotRestorer(activity),
                     getThemedIconSnapshotRestorer(activity),
                     getThemedIconInteractor(),
+                    getColorPickerInteractor(activity, getWallpaperColorsViewModel()),
                 )
                 .also { customizationSections = it }
     }
@@ -165,32 +179,31 @@
         mode: Int,
         viewAsHome: Boolean,
         viewFullScreen: Boolean,
-        testingModeEnabled: Boolean
+        testingModeEnabled: Boolean,
+        isAssetIdPresent: Boolean
     ): Fragment {
-        return if (wallpaperInfo is LiveWallpaperInfo) LivePreviewFragment()
+        return if (wallpaperInfo is LiveWallpaperInfo) LivePreviewFragment2()
         else
-            ImagePreviewFragment().apply {
+            ImagePreviewFragment2().apply {
                 arguments =
                     Bundle().apply {
-                        putParcelable(PreviewFragment.ARG_WALLPAPER, wallpaperInfo)
-                        putInt(PreviewFragment.ARG_PREVIEW_MODE, mode)
-                        putBoolean(PreviewFragment.ARG_VIEW_AS_HOME, viewAsHome)
-                        putBoolean(PreviewFragment.ARG_FULL_SCREEN, viewFullScreen)
-                        putBoolean(PreviewFragment.ARG_TESTING_MODE_ENABLED, testingModeEnabled)
+                        putParcelable(PreviewFragment2.ARG_WALLPAPER, wallpaperInfo)
+                        putInt(PreviewFragment2.ARG_PREVIEW_MODE, mode)
+                        putBoolean(PreviewFragment2.ARG_VIEW_AS_HOME, viewAsHome)
+                        putBoolean(PreviewFragment2.ARG_FULL_SCREEN, viewFullScreen)
+                        putBoolean(PreviewFragment2.ARG_TESTING_MODE_ENABLED, testingModeEnabled)
                     }
             }
     }
 
     @Synchronized
     override fun getUserEventLogger(context: Context): ThemesUserEventLogger {
-        return if (userEventLogger != null) userEventLogger as ThemesUserEventLogger
-        else StatsLogUserEventLogger(context.applicationContext).also { userEventLogger = it }
-    }
-
-    @Synchronized
-    override fun getPreferences(context: Context): WallpaperPreferences {
-        return prefs
-            ?: DefaultCustomizationPreferences(context.applicationContext).also { prefs = it }
+        return userEventLogger as? ThemesUserEventLogger
+            ?: StatsLogUserEventLogger(
+                    getPreferences(context.applicationContext),
+                    getWallpaperStatusChecker(context.applicationContext),
+                )
+                .also { userEventLogger = it }
     }
 
     override fun getFragmentFactory(): FragmentFactory? {
@@ -244,7 +257,7 @@
                                     wallpaperManager = WallpaperManager.getInstance(appContext)
                                 ),
                             wallpaperPreferences = getPreferences(context = appContext),
-                            backgroundDispatcher = Dispatchers.IO,
+                            backgroundDispatcher = bgDispatcher,
                         ),
                     shouldHandleReload = {
                         TextUtils.equals(
@@ -284,7 +297,7 @@
         val client = getKeyguardQuickAffordancePickerProviderClient(context)
         val appContext = context.applicationContext
         return KeyguardQuickAffordancePickerInteractor(
-            KeyguardQuickAffordancePickerRepository(client, Dispatchers.IO),
+            KeyguardQuickAffordancePickerRepository(client, bgDispatcher),
             client
         ) {
             getKeyguardQuickAffordanceSnapshotRestorer(appContext)
@@ -295,7 +308,7 @@
         context: Context
     ): CustomizationProviderClient {
         return customizationProviderClient
-            ?: CustomizationProviderClientImpl(context.applicationContext, Dispatchers.IO).also {
+            ?: CustomizationProviderClientImpl(context.applicationContext, bgDispatcher).also {
                 customizationProviderClient = it
             }
     }
@@ -330,7 +343,7 @@
                     repository =
                         NotificationsRepository(
                             scope = getApplicationCoroutineScope(),
-                            backgroundDispatcher = Dispatchers.IO,
+                            backgroundDispatcher = bgDispatcher,
                             secureSettingsRepository = getSecureSettingsRepository(context),
                         ),
                     snapshotRestorer = { getNotificationsSnapshotRestorer(appContext) },
@@ -354,8 +367,8 @@
                 ?: ClockRegistryProvider(
                         context = context.applicationContext,
                         coroutineScope = getApplicationCoroutineScope(),
-                        mainDispatcher = Dispatchers.Main,
-                        backgroundDispatcher = Dispatchers.IO,
+                        mainDispatcher = mainDispatcher,
+                        backgroundDispatcher = bgDispatcher,
                     )
                     .also { clockRegistryProvider = it })
             .get()
@@ -372,7 +385,7 @@
                             secureSettingsRepository = getSecureSettingsRepository(appContext),
                             registry = getClockRegistry(appContext),
                             scope = getApplicationCoroutineScope(),
-                            mainDispatcher = Dispatchers.Main,
+                            mainDispatcher = mainDispatcher,
                         ),
                     snapshotRestorer = { getClockPickerSnapshotRestorer(appContext) },
                 )
@@ -394,7 +407,7 @@
         interactor: ClockPickerInteractor,
     ): ClockCarouselViewModel.Factory {
         return clockCarouselViewModelFactory
-            ?: ClockCarouselViewModel.Factory(interactor, Dispatchers.IO).also {
+            ?: ClockCarouselViewModel.Factory(interactor, bgDispatcher).also {
                 clockCarouselViewModelFactory = it
             }
     }
@@ -489,7 +502,7 @@
             ?: DarkModeSnapshotRestorer(
                     context = appContext,
                     manager = appContext.getSystemService(Context.UI_MODE_SERVICE) as UiModeManager,
-                    backgroundDispatcher = Dispatchers.IO,
+                    backgroundDispatcher = bgDispatcher,
                 )
                 .also { darkModeSnapshotRestorer = it }
     }
@@ -537,7 +550,7 @@
                 .also { clockSettingsViewModelFactory = it }
     }
 
-    override fun getClockDescriptionUtils(): ClockDescriptionUtils {
+    override fun getClockDescriptionUtils(resources: Resources): ClockDescriptionUtils {
         return clockDescriptionUtils
             ?: ThemePickerClockDescriptionUtils().also { clockDescriptionUtils = it }
     }
@@ -553,9 +566,7 @@
                 .also { gridScreenViewModelFactory = it }
     }
 
-    private fun getGridInteractor(
-        context: Context,
-    ): GridInteractor {
+    fun getGridInteractor(context: Context): GridInteractor {
         val appContext = context.applicationContext
         return gridInteractor
             ?: GridInteractor(
@@ -564,7 +575,9 @@
                         GridRepositoryImpl(
                             applicationScope = getApplicationCoroutineScope(),
                             manager = GridOptionsManager.getInstance(context),
-                            backgroundDispatcher = Dispatchers.IO,
+                            backgroundDispatcher = bgDispatcher,
+                            isGridApplyButtonEnabled =
+                                BaseFlags.get().isGridApplyButtonEnabled(appContext),
                         ),
                     snapshotRestorer = { getGridSnapshotRestorer(appContext) },
                 )
@@ -581,6 +594,12 @@
                 .also { gridSnapshotRestorer = it }
     }
 
+    override fun isCurrentSelectedColorPreset(context: Context): Boolean {
+        val colorManager =
+            ColorCustomizationManager.getInstance(context, OverlayManagerCompat(context))
+        return COLOR_SOURCE_PRESET == colorManager.currentColorSource
+    }
+
     companion object {
         @JvmStatic
         private val KEY_QUICK_AFFORDANCE_SNAPSHOT_RESTORER =
diff --git a/src/com/android/customization/picker/HorizontalTouchMovementAwareNestedScrollView.kt b/src/com/android/customization/picker/HorizontalTouchMovementAwareNestedScrollView.kt
deleted file mode 100644
index 06cf753..0000000
--- a/src/com/android/customization/picker/HorizontalTouchMovementAwareNestedScrollView.kt
+++ /dev/null
@@ -1,64 +0,0 @@
-/*
- * Copyright (C) 2023 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.customization.picker
-
-import android.content.Context
-import android.util.AttributeSet
-import android.view.MotionEvent
-import android.view.ViewConfiguration
-import androidx.core.widget.NestedScrollView
-import kotlin.math.abs
-
-/**
- * This nested scroll view will detect horizontal touch movements and stop vertical scrolls when a
- * horizontal touch movement is detected.
- */
-class HorizontalTouchMovementAwareNestedScrollView(context: Context, attrs: AttributeSet?) :
-    NestedScrollView(context, attrs) {
-
-    private var startXPosition = 0f
-    private var startYPosition = 0f
-    private var isHorizontalTouchMovement = false
-
-    override fun onInterceptTouchEvent(event: MotionEvent): Boolean {
-        when (event.action) {
-            MotionEvent.ACTION_DOWN -> {
-                startXPosition = event.x
-                startYPosition = event.y
-                isHorizontalTouchMovement = false
-            }
-            MotionEvent.ACTION_MOVE -> {
-                val xMoveDistance = abs(event.x - startXPosition)
-                val yMoveDistance = abs(event.y - startYPosition)
-                if (
-                    !isHorizontalTouchMovement &&
-                        xMoveDistance > yMoveDistance &&
-                        xMoveDistance > ViewConfiguration.get(context).scaledTouchSlop
-                ) {
-                    isHorizontalTouchMovement = true
-                }
-            }
-            else -> {}
-        }
-        return if (isHorizontalTouchMovement) {
-            // We only want to intercept the touch event when the touch moves more vertically than
-            // horizontally. So we return false.
-            false
-        } else {
-            super.onInterceptTouchEvent(event)
-        }
-    }
-}
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
index be6c6cb..004103f 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockPickerRepositoryImpl.kt
@@ -33,6 +33,7 @@
 import kotlinx.coroutines.flow.SharedFlow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.callbackFlow
+import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flowOn
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -52,11 +53,12 @@
     override val allClocks: Flow<List<ClockMetadataModel>> =
         callbackFlow {
                 fun send() {
+                    val activeClockId = registry.activeClockId
                     val allClocks =
-                        registry
-                            .getClocks()
-                            .filter { "NOT_IN_USE" !in it.clockId }
-                            .map { it.toModel() }
+                        registry.getClocks().map {
+                            it.toModel(isSelected = it.clockId == activeClockId)
+                        }
+
                     trySend(allClocks)
                 }
 
@@ -83,13 +85,14 @@
     override val selectedClock: Flow<ClockMetadataModel> =
         callbackFlow {
                 fun send() {
-                    val currentClockId = registry.currentClockId
+                    val activeClockId = registry.activeClockId
                     val metadata = registry.settings?.metadata
                     val model =
                         registry
                             .getClocks()
-                            .find { clockMetadata -> clockMetadata.clockId == currentClockId }
+                            .find { clockMetadata -> clockMetadata.clockId == activeClockId }
                             ?.toModel(
+                                isSelected = true,
                                 selectedColorId = metadata?.getSelectedColorId(),
                                 colorTone = metadata?.getColorTone()
                                         ?: ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
@@ -146,9 +149,10 @@
             )
             .map { setting -> setting == 1 }
             .map { isDynamic -> if (isDynamic) ClockSize.DYNAMIC else ClockSize.SMALL }
+            .distinctUntilChanged()
             .shareIn(
                 scope = scope,
-                started = SharingStarted.WhileSubscribed(),
+                started = SharingStarted.Eagerly,
                 replay = 1,
             )
 
@@ -176,6 +180,7 @@
 
     /** By default, [ClockMetadataModel] has no color information unless specified. */
     private fun ClockMetadata.toModel(
+        isSelected: Boolean,
         selectedColorId: String? = null,
         @IntRange(from = 0, to = 100) colorTone: Int = 0,
         @ColorInt seedColor: Int? = null,
@@ -183,6 +188,7 @@
         return ClockMetadataModel(
             clockId = clockId,
             name = name,
+            isSelected = isSelected,
             selectedColorId = selectedColorId,
             colorToneProgress = colorTone,
             seedColor = seedColor,
diff --git a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
index 652ffdd..b197edf 100644
--- a/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
+++ b/src/com/android/customization/picker/clock/data/repository/ClockRegistryProvider.kt
@@ -29,6 +29,7 @@
 import com.android.systemui.shared.plugins.PluginManagerImpl
 import com.android.systemui.shared.plugins.PluginPrefs
 import com.android.systemui.shared.system.UncaughtExceptionPreHandlerManager_Factory
+import com.android.wallpaper.module.InjectorProvider
 import java.util.concurrent.Executors
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.CoroutineScope
@@ -55,6 +56,8 @@
             DefaultClockProvider(context, LayoutInflater.from(context), context.resources),
             keepAllLoaded = true,
             subTag = "Picker",
+            isTransitClockEnabled =
+                InjectorProvider.getInjector().getFlags().isTransitClockEnabled(context)
         )
     }
 
diff --git a/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt b/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt
index bd87ba6..2522507 100644
--- a/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt
+++ b/src/com/android/customization/picker/clock/shared/model/ClockMetadataModel.kt
@@ -24,6 +24,7 @@
 data class ClockMetadataModel(
     val clockId: String,
     val name: String,
+    val isSelected: Boolean,
     val selectedColorId: String?,
     @IntRange(from = 0, to = 100) val colorToneProgress: Int,
     @ColorInt val seedColor: Int?,
diff --git a/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt b/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
index 89fac89..b55464d 100644
--- a/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
+++ b/src/com/android/customization/picker/clock/ui/binder/ClockCarouselViewBinder.kt
@@ -16,8 +16,8 @@
 package com.android.customization.picker.clock.ui.binder
 
 import android.content.Context
-import android.view.ViewGroup
-import android.widget.FrameLayout
+import android.content.res.Configuration
+import androidx.core.content.ContextCompat
 import androidx.core.view.isVisible
 import androidx.lifecycle.Lifecycle
 import androidx.lifecycle.LifecycleEventObserver
@@ -27,7 +27,6 @@
 import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
-import com.android.wallpaper.R
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewClickView
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.launch
@@ -38,7 +37,6 @@
     fun bind(
         context: Context,
         carouselView: ClockCarouselView,
-        singleClockView: ViewGroup,
         screenPreviewClickView: ScreenPreviewClickView,
         viewModel: ClockCarouselViewModel,
         clockViewFactory: ClockViewFactory,
@@ -46,6 +44,7 @@
         isTwoPaneAndSmallWidth: Boolean,
     ) {
         carouselView.setClockViewFactory(clockViewFactory)
+        carouselView.isVisible = true
         clockViewFactory.updateRegionDarkness()
         val carouselAccessibilityDelegate =
             CarouselAccessibilityDelegate(
@@ -61,12 +60,8 @@
             )
         screenPreviewClickView.accessibilityDelegate = carouselAccessibilityDelegate
 
-        val singleClockHostView =
-            singleClockView.requireViewById<FrameLayout>(R.id.single_clock_host_view)
         lifecycleOwner.lifecycleScope.launch {
             lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                launch { viewModel.isCarouselVisible.collect { carouselView.isVisible = it } }
-
                 launch {
                     combine(viewModel.selectedClockSize, viewModel.allClocks, ::Pair).collect {
                         (size, allClocks) ->
@@ -100,17 +95,11 @@
                 }
 
                 launch {
-                    viewModel.isSingleClockViewVisible.collect { singleClockView.isVisible = it }
-                }
-
-                launch {
-                    viewModel.clockId.collect { clockId ->
-                        singleClockHostView.removeAllViews()
-                        val clockView = clockViewFactory.getLargeView(clockId)
-                        // The clock view might still be attached to an existing parent. Detach
-                        // before adding to another parent.
-                        (clockView.parent as? ViewGroup)?.removeView(clockView)
-                        singleClockHostView.addView(clockView)
+                    val night =
+                        (context.resources.configuration.uiMode and
+                            Configuration.UI_MODE_NIGHT_MASK == Configuration.UI_MODE_NIGHT_YES)
+                    viewModel.getClockCardColorResId(night).collect {
+                        carouselView.setCarouselCardColor(ContextCompat.getColor(context, it))
                     }
                 }
             }
diff --git a/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt b/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
index 4f4bd1b..6e745d5 100644
--- a/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
+++ b/src/com/android/customization/picker/clock/ui/binder/ClockSettingsBinder.kt
@@ -33,6 +33,7 @@
 import androidx.recyclerview.widget.RecyclerView
 import com.android.customization.picker.clock.shared.ClockSize
 import com.android.customization.picker.clock.ui.adapter.ClockSettingsTabAdapter
+import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockHostView
 import com.android.customization.picker.clock.ui.view.ClockSizeRadioButtonGroup
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
@@ -199,7 +200,7 @@
                                     sizeOptions.radioButtonDynamic.isChecked = false
                                     sizeOptions.radioButtonSmall.isChecked = true
                                     clockHostView.doOnPreDraw {
-                                        it.pivotX = 0F
+                                        it.pivotX = ClockCarouselView.getCenteredHostViewPivotX(it)
                                         it.pivotY = 0F
                                     }
                                 }
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
deleted file mode 100644
index f138d6a..0000000
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragment.kt
+++ /dev/null
@@ -1,93 +0,0 @@
-package com.android.customization.picker.clock.ui.fragment
-
-import android.content.Context
-import android.os.Bundle
-import android.util.TypedValue
-import android.view.ContextThemeWrapper
-import android.view.LayoutInflater
-import android.view.View
-import android.view.ViewGroup
-import android.view.ViewGroup.LayoutParams.MATCH_PARENT
-import android.view.ViewGroup.LayoutParams.WRAP_CONTENT
-import android.widget.FrameLayout
-import android.widget.TextView
-import android.widget.Toast
-import androidx.core.view.setPadding
-import androidx.recyclerview.widget.LinearLayoutManager
-import androidx.recyclerview.widget.RecyclerView
-import com.android.customization.module.ThemePickerInjector
-import com.android.internal.annotations.VisibleForTesting
-import com.android.systemui.plugins.ClockMetadata
-import com.android.systemui.shared.clocks.ClockRegistry
-import com.android.wallpaper.R
-import com.android.wallpaper.module.InjectorProvider
-import com.android.wallpaper.picker.AppbarFragment
-
-class ClockCustomDemoFragment : AppbarFragment() {
-    @VisibleForTesting lateinit var recyclerView: RecyclerView
-    @VisibleForTesting lateinit var clockRegistry: ClockRegistry
-
-    override fun onCreateView(
-        inflater: LayoutInflater,
-        container: ViewGroup?,
-        savedInstanceState: Bundle?
-    ): View {
-        val view = inflater.inflate(R.layout.fragment_clock_custom_picker_demo, container, false)
-        setUpToolbar(view)
-        clockRegistry =
-            (InjectorProvider.getInjector() as ThemePickerInjector).getClockRegistry(
-                requireContext(),
-            )
-        val listInUse = clockRegistry.getClocks().filter { "NOT_IN_USE" !in it.clockId }
-
-        recyclerView = view.requireViewById(R.id.clock_preview_card_list_demo)
-        recyclerView.layoutManager = LinearLayoutManager(context, RecyclerView.VERTICAL, false)
-        recyclerView.adapter =
-            ClockRecyclerAdapter(listInUse, requireContext()) {
-                clockRegistry.currentClockId = it.clockId
-                Toast.makeText(context, "${it.name} selected", Toast.LENGTH_SHORT).show()
-            }
-        return view
-    }
-
-    override fun getDefaultTitle(): CharSequence {
-        return getString(R.string.clock_title)
-    }
-
-    internal class ClockRecyclerAdapter(
-        val list: List<ClockMetadata>,
-        val context: Context,
-        val onClockSelected: (ClockMetadata) -> Unit
-    ) : RecyclerView.Adapter<ClockRecyclerAdapter.ViewHolder>() {
-        class ViewHolder(val view: View, val textView: TextView, val onItemClicked: (Int) -> Unit) :
-            RecyclerView.ViewHolder(view) {
-            init {
-                itemView.setOnClickListener { onItemClicked(absoluteAdapterPosition) }
-            }
-        }
-
-        override fun onCreateViewHolder(viewGroup: ViewGroup, viewType: Int): ViewHolder {
-            val rootView = FrameLayout(viewGroup.context)
-            val textView =
-                TextView(ContextThemeWrapper(viewGroup.context, R.style.SectionTitleTextStyle))
-            textView.setPadding(ITEM_PADDING)
-            rootView.addView(textView)
-            val outValue = TypedValue()
-            context.theme.resolveAttribute(android.R.attr.selectableItemBackground, outValue, true)
-            rootView.setBackgroundResource(outValue.resourceId)
-            val lp = RecyclerView.LayoutParams(MATCH_PARENT, WRAP_CONTENT)
-            rootView.layoutParams = lp
-            return ViewHolder(rootView, textView) { onClockSelected(list[it]) }
-        }
-
-        override fun onBindViewHolder(viewHolder: ViewHolder, position: Int) {
-            viewHolder.textView.text = list[position].name
-        }
-
-        override fun getItemCount() = list.size
-
-        companion object {
-            val ITEM_PADDING = 40
-        }
-    }
-}
diff --git a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
index f4684d8..4805f37 100644
--- a/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
+++ b/src/com/android/customization/picker/clock/ui/fragment/ClockSettingsFragment.kt
@@ -20,8 +20,12 @@
 import android.view.View
 import android.view.ViewGroup
 import androidx.cardview.widget.CardView
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
+import androidx.transition.Transition
+import androidx.transition.doOnStart
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.clock.ui.binder.ClockSettingsBinder
 import com.android.systemui.shared.clocks.shared.model.ClockPreviewConstants
@@ -130,6 +134,8 @@
             viewLifecycleOwner,
         )
 
+        (returnTransition as? Transition)?.doOnStart { lockScreenView.isVisible = false }
+
         return view
     }
 
@@ -140,4 +146,8 @@
     override fun getToolbarColorId(): Int {
         return android.R.color.transparent
     }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
 }
diff --git a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
index 8764e54..0b62ffa 100644
--- a/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
+++ b/src/com/android/customization/picker/clock/ui/view/ClockCarouselView.kt
@@ -16,6 +16,7 @@
 package com.android.customization.picker.clock.ui.view
 
 import android.content.Context
+import android.content.res.ColorStateList
 import android.content.res.Resources
 import android.util.AttributeSet
 import android.view.LayoutInflater
@@ -71,17 +72,31 @@
         clockViewFactory = factory
     }
 
+    // This function is for the custom accessibility action to trigger a transition to the next
+    // carousel item. If the current item is the last item in the carousel, the next item
+    // will be the first item.
     fun transitionToNext() {
-        val index = (carousel.currentIndex + 1) % carousel.count
-        if (index < carousel.count && index > 0) {
-            carousel.transitionToIndex(index, 0)
+        if (carousel.count != 0) {
+            val index = (carousel.currentIndex + 1) % carousel.count
+            carousel.jumpToIndex(index)
+            // Explicitly called this since using transitionToIndex(index) leads to
+            // race-condition between announcement of content description of the correct clock-face
+            // and the selection of clock face itself
+            adapter.onNewItem(index)
         }
     }
 
+    // This function is for the custom accessibility action to trigger a transition to
+    // the previous carousel item. If the current item is the first item in the carousel,
+    // the previous item will be the last item.
     fun transitionToPrevious() {
-        val index = (carousel.currentIndex - 1) % carousel.count
-        if (index < carousel.count && index > 0) {
-            carousel.transitionToIndex(index, 0)
+        if (carousel.count != 0) {
+            val index = (carousel.currentIndex + carousel.count - 1) % carousel.count
+            carousel.jumpToIndex(index)
+            // Explicitly called this since using transitionToIndex(index) leads to
+            // race-condition between announcement of content description of the correct clock-face
+            // and the selection of clock face itself
+            adapter.onNewItem(index)
         }
     }
 
@@ -100,8 +115,15 @@
         }
 
         adapter = ClockCarouselAdapter(clockSize, clocks, clockViewFactory, onClockSelected)
+        carousel.isInfinite = clocks.size >= MIN_CLOCKS_TO_ENABLE_INFINITE_CAROUSEL
         carousel.setAdapter(adapter)
-        carousel.refresh()
+        val indexOfSelectedClock =
+            clocks
+                .indexOfFirst { it.isSelected }
+                // If not found, default to the first clock as selected:
+                .takeIf { it != -1 }
+                ?: 0
+        carousel.jumpToIndex(indexOfSelectedClock)
         motionLayout.setTransitionListener(
             object : MotionLayout.TransitionListener {
 
@@ -211,11 +233,13 @@
                         }
                             ?: return
                     offCenterClockHostView.doOnPreDraw {
-                        it.pivotX = progress * it.width / 2
+                        it.pivotX =
+                            progress * it.width / 2 + (1 - progress) * getCenteredHostViewPivotX(it)
                         it.pivotY = progress * it.height / 2
                     }
                     toCenterClockHostView.doOnPreDraw {
-                        it.pivotX = (1 - progress) * it.width / 2
+                        it.pivotX =
+                            (1 - progress) * it.width / 2 + progress * getCenteredHostViewPivotX(it)
                         it.pivotY = (1 - progress) * it.height / 2
                     }
                     offCenterClockFrame.translationX =
@@ -265,13 +289,25 @@
     fun setSelectedClockIndex(
         index: Int,
     ) {
-        // jumpToIndex to the same position can cause the views unnecessarily populate again.
-        // Only call jumpToIndex when the jump-to index is different from the current carousel.
-        if (index != carousel.currentIndex) {
+        // 1. setUpClockCarouselView() can possibly not be called before setSelectedClockIndex().
+        //    We need to check if index out of bound.
+        // 2. jumpToIndex() to the same position can cause the views unnecessarily populate again.
+        //    We only call jumpToIndex when the index is different from the current carousel.
+        if (index < carousel.count && index != carousel.currentIndex) {
             carousel.jumpToIndex(index)
         }
     }
 
+    fun setCarouselCardColor(color: Int) {
+        itemViewIds.forEach { id ->
+            val cardViewId = getClockCardViewId(id)
+            cardViewId?.let {
+                val cardView = motionLayout.requireViewById<View>(it)
+                cardView.backgroundTintList = ColorStateList.valueOf(color)
+            }
+        }
+    }
+
     private fun overrideScreenPreviewWidth() {
         val overrideWidth =
             context.resources.getDimensionPixelSize(
@@ -417,7 +453,7 @@
         ) {
             clockHostView.doOnPreDraw {
                 if (isMiddleView) {
-                    it.pivotX = 0F
+                    it.pivotX = getCenteredHostViewPivotX(it)
                     it.pivotY = 0F
                     clockView.translationX = 0F
                     clockView.translationY = 0F
@@ -446,6 +482,8 @@
     }
 
     companion object {
+        // The carousel needs to have at least 5 different clock faces to be infinite
+        const val MIN_CLOCKS_TO_ENABLE_INFINITE_CAROUSEL = 5
         const val CLOCK_CAROUSEL_VIEW_SCALE = 0.5f
 
         val itemViewIds =
@@ -507,6 +545,10 @@
             return rootViewId == R.id.item_view_2
         }
 
+        fun getCenteredHostViewPivotX(hostView: View): Float {
+            return if (hostView.isLayoutRtl) hostView.width.toFloat() else 0F
+        }
+
         private fun getTranslationDistance(
             hostLength: Int,
             frameLength: Int,
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt
index 708fa2f..9811426 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselItemViewModel.kt
@@ -20,15 +20,14 @@
 import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
 
-class ClockCarouselItemViewModel(val clockId: String) {
+class ClockCarouselItemViewModel(val clockId: String, val isSelected: Boolean) {
 
     /** Description for accessibility purposes when a clock is selected. */
     fun getContentDescription(resources: Resources): String {
         val clockContent =
             (InjectorProvider.getInjector() as? CustomizationInjector)
-                ?.getClockDescriptionUtils()
-                ?.getDescriptionResId(clockId)
-                ?.let { resources.getString(it) }
+                ?.getClockDescriptionUtils(resources)
+                ?.getDescription(clockId)
                 ?: ""
         return resources.getString(R.string.select_clock_action_description, clockContent)
     }
diff --git a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
index a4f9cc4..27c25a2 100644
--- a/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
+++ b/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModel.kt
@@ -15,11 +15,13 @@
  */
 package com.android.customization.picker.clock.ui.viewmodel
 
+import android.graphics.Color
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.viewModelScope
 import com.android.customization.picker.clock.domain.interactor.ClockPickerInteractor
 import com.android.customization.picker.clock.shared.ClockSize
+import com.android.wallpaper.R
 import kotlinx.coroutines.CoroutineDispatcher
 import kotlinx.coroutines.ExperimentalCoroutinesApi
 import kotlinx.coroutines.Job
@@ -27,7 +29,6 @@
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.SharingStarted
 import kotlinx.coroutines.flow.StateFlow
-import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.flow.mapLatest
@@ -39,8 +40,7 @@
  * Clock carousel view model that provides data for the carousel of clock previews. When there is
  * only one item, we should show a single clock preview instead of a carousel.
  */
-class ClockCarouselViewModel
-constructor(
+class ClockCarouselViewModel(
     private val interactor: ClockPickerInteractor,
     private val backgroundDispatcher: CoroutineDispatcher,
 ) : ViewModel() {
@@ -50,7 +50,7 @@
             .mapLatest { allClocks ->
                 // Delay to avoid the case that the full list of clocks is not initiated.
                 delay(CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-                allClocks.map { ClockCarouselItemViewModel(it.clockId) }
+                allClocks.map { ClockCarouselItemViewModel(it.clockId, it.isSelected) }
             }
             .stateIn(viewModelScope, SharingStarted.Eagerly, emptyList())
 
@@ -58,7 +58,37 @@
 
     val seedColor: Flow<Int?> = interactor.seedColor
 
-    val isCarouselVisible: Flow<Boolean> = allClocks.map { it.size > 1 }.distinctUntilChanged()
+    fun getClockCardColorResId(isDarkThemeEnabled: Boolean): Flow<Int> {
+        return interactor.seedColor.map {
+            it.let { seedColor ->
+                // if seedColor is null, default clock color is selected
+                if (seedColor == null) {
+                    if (isDarkThemeEnabled) {
+                        // In dark mode, use darkest surface container color
+                        R.color.system_surface_container_high
+                    } else {
+                        // In light mode, use lightest surface container color
+                        R.color.system_surface_bright
+                    }
+                } else {
+                    val luminance = Color.luminance(seedColor)
+                    if (isDarkThemeEnabled) {
+                        if (luminance <= CARD_COLOR_CHANGE_LUMINANCE_THRESHOLD_DARK_THEME) {
+                            R.color.system_surface_bright
+                        } else {
+                            R.color.system_surface_container_high
+                        }
+                    } else {
+                        if (luminance <= CARD_COLOR_CHANGE_LUMINANCE_THRESHOLD_LIGHT_THEME) {
+                            R.color.system_surface_bright
+                        } else {
+                            R.color.system_surface_container_highest
+                        }
+                    }
+                }
+            }
+        }
+    }
 
     @OptIn(ExperimentalCoroutinesApi::class)
     val selectedIndex: Flow<Int> =
@@ -77,15 +107,6 @@
             }
             .mapNotNull { it }
 
-    // Handle the case when there is only one clock in the carousel
-    val isSingleClockViewVisible: Flow<Boolean> =
-        allClocks.map { it.size == 1 }.distinctUntilChanged()
-
-    val clockId: Flow<String> =
-        allClocks
-            .map { allClockIds -> if (allClockIds.size == 1) allClockIds[0].clockId else null }
-            .mapNotNull { it }
-
     private var setSelectedClockJob: Job? = null
     fun setSelectedClock(clockId: String) {
         setSelectedClockJob?.cancel()
@@ -109,5 +130,7 @@
 
     companion object {
         const val CLOCKS_EVENT_UPDATE_DELAY_MILLIS: Long = 100
+        const val CARD_COLOR_CHANGE_LUMINANCE_THRESHOLD_LIGHT_THEME: Float = 0.85f
+        const val CARD_COLOR_CHANGE_LUMINANCE_THRESHOLD_DARK_THEME: Float = 0.03f
     }
 }
diff --git a/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt b/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt
index 9a0b66f..28ea4a3 100644
--- a/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt
+++ b/src/com/android/customization/picker/clock/utils/ClockDescriptionUtils.kt
@@ -15,14 +15,12 @@
  */
 package com.android.customization.picker.clock.utils
 
-import androidx.annotation.StringRes
-
 /** Provides clock description for accessibility purposes. */
 interface ClockDescriptionUtils {
 
     /**
-     * TODO (b/287507746) : Migrate description res ID to system UI or a shared library, instead of
-     * preserving the clock description at the Wallpaper Picker side.
+     * TODO (b/287507746) : Migrate the clock description to system UI or a shared library, instead
+     * of preserving at the Wallpaper Picker side.
      */
-    @StringRes fun getDescriptionResId(clockId: String): Int
+    fun getDescription(clockId: String): String
 }
diff --git a/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt b/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt
index 43f19b3..a04ebff 100644
--- a/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt
+++ b/src/com/android/customization/picker/clock/utils/ThemePickerClockDescriptionUtils.kt
@@ -15,12 +15,8 @@
  */
 package com.android.customization.picker.clock.utils
 
-import androidx.annotation.StringRes
-import com.android.wallpaper.R
-
 class ThemePickerClockDescriptionUtils : ClockDescriptionUtils {
-    @StringRes
-    override fun getDescriptionResId(clockId: String): Int {
-        return R.string.clock_title
+    override fun getDescription(clockId: String): String {
+        return ""
     }
 }
diff --git a/src/com/android/customization/picker/color/ColorSectionView.java b/src/com/android/customization/picker/color/ColorSectionView.java
deleted file mode 100644
index b8ba2e4..0000000
--- a/src/com/android/customization/picker/color/ColorSectionView.java
+++ /dev/null
@@ -1,33 +0,0 @@
-/*
- * Copyright (C) 2022 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.customization.picker.color;
-
-import android.content.Context;
-import android.util.AttributeSet;
-
-import androidx.annotation.Nullable;
-
-import com.android.wallpaper.picker.SectionView;
-
-/**
- * The class inherits from {@link SectionView} as the view representing the color section of the
- * customization picker.
- */
-public final class ColorSectionView extends SectionView {
-    public ColorSectionView(Context context, @Nullable AttributeSet attrs) {
-        super(context, attrs);
-    }
-}
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
index 7cf9fd0..fccaa65 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepository.kt
@@ -19,12 +19,15 @@
 import com.android.customization.picker.color.shared.model.ColorOptionModel
 import com.android.customization.picker.color.shared.model.ColorType
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.StateFlow
 
 /**
  * Abstracts access to application state related to functionality for selecting, picking, or setting
  * system color.
  */
 interface ColorPickerRepository {
+    /** Whether the system color is in the process of being updated */
+    val isApplyingSystemColor: StateFlow<Boolean>
 
     /** List of wallpaper and preset color options on the device, categorized by Color Type */
     val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>>
diff --git a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
index 41ef3a5..ce5c0c4 100644
--- a/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
+++ b/src/com/android/customization/picker/color/data/repository/ColorPickerRepositoryImpl.kt
@@ -27,7 +27,9 @@
 import com.android.wallpaper.model.WallpaperColorsModel
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.MutableStateFlow
 import kotlinx.coroutines.flow.StateFlow
+import kotlinx.coroutines.flow.asStateFlow
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.map
 import kotlinx.coroutines.suspendCancellableCoroutine
@@ -43,9 +45,17 @@
         wallpaperColorsViewModel.homeWallpaperColors
     private val lockWallpaperColors: StateFlow<WallpaperColorsModel?> =
         wallpaperColorsViewModel.lockWallpaperColors
+    private var selectedColorOption: MutableStateFlow<ColorOptionModel> =
+        MutableStateFlow(getCurrentColorOption())
+
+    private val _isApplyingSystemColor = MutableStateFlow(false)
+    override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
 
     override val colorOptions: Flow<Map<ColorType, List<ColorOptionModel>>> =
-        combine(homeWallpaperColors, lockWallpaperColors) { homeColors, lockColors ->
+        combine(homeWallpaperColors, lockWallpaperColors, selectedColorOption) {
+                homeColors,
+                lockColors,
+                _ ->
                 homeColors to lockColors
             }
             .map { (homeColors, lockColors) ->
@@ -109,17 +119,21 @@
                 }
             }
 
-    override suspend fun select(colorOptionModel: ColorOptionModel) =
+    override suspend fun select(colorOptionModel: ColorOptionModel) {
+        _isApplyingSystemColor.value = true
         suspendCancellableCoroutine { continuation ->
             colorManager.apply(
                 colorOptionModel.colorOption,
                 object : CustomizationManager.Callback {
                     override fun onSuccess() {
+                        _isApplyingSystemColor.value = false
+                        selectedColorOption.value = colorOptionModel
                         continuation.resumeWith(Result.success(Unit))
                     }
 
                     override fun onError(throwable: Throwable?) {
                         Log.w(TAG, "Apply theme with error", throwable)
+                        _isApplyingSystemColor.value = false
                         continuation.resumeWith(
                             Result.failure(throwable ?: Throwable("Error loading theme bundles"))
                         )
@@ -127,6 +141,7 @@
                 }
             )
         }
+    }
 
     override fun getCurrentColorOption(): ColorOptionModel {
         val overlays = colorManager.currentOverlays
diff --git a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
index 714129d..bb2ef9d 100644
--- a/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
+++ b/src/com/android/customization/picker/color/data/repository/FakeColorPickerRepository.kt
@@ -29,6 +29,9 @@
 
 class FakeColorPickerRepository(private val context: Context) : ColorPickerRepository {
 
+    private val _isApplyingSystemColor = MutableStateFlow(false)
+    override val isApplyingSystemColor = _isApplyingSystemColor.asStateFlow()
+
     private lateinit var selectedColorOption: ColorOptionModel
 
     private val _colorOptions =
diff --git a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
index 8c7a4b7..d3b2eba 100644
--- a/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
+++ b/src/com/android/customization/picker/color/domain/interactor/ColorPickerInteractor.kt
@@ -26,6 +26,8 @@
     private val repository: ColorPickerRepository,
     private val snapshotRestorer: Provider<ColorPickerSnapshotRestorer>,
 ) {
+    val isApplyingSystemColor = repository.isApplyingSystemColor
+
     /**
      * The newly selected color option for overwriting the current active option during an
      * optimistic update, the value is set to null when update fails
diff --git a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
index 2daefe4..ad81614 100644
--- a/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
+++ b/src/com/android/customization/picker/color/ui/binder/ColorSectionViewBinder.kt
@@ -112,22 +112,20 @@
             val optionSelectedView = itemView.requireViewById<ImageView>(R.id.option_selected)
 
             lifecycleOwner.lifecycleScope.launch {
-                lifecycleOwner.repeatOnLifecycle(Lifecycle.State.STARTED) {
-                    launch {
-                        item.isSelected.collect { isSelected ->
-                            optionSelectedView.isVisible = isSelected
-                        }
+                launch {
+                    item.isSelected.collect { isSelected ->
+                        optionSelectedView.isVisible = isSelected
                     }
-                    launch {
-                        item.onClicked.collect { onClicked ->
-                            itemView.setOnClickListener(
-                                if (onClicked != null) {
-                                    View.OnClickListener { onClicked.invoke() }
-                                } else {
-                                    null
-                                }
-                            )
-                        }
+                }
+                launch {
+                    item.onClicked.collect { onClicked ->
+                        itemView.setOnClickListener(
+                            if (onClicked != null) {
+                                View.OnClickListener { onClicked.invoke() }
+                            } else {
+                                null
+                            }
+                        )
                     }
                 }
             }
diff --git a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
index 78bfa43..4ef29d6 100644
--- a/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
+++ b/src/com/android/customization/picker/color/ui/fragment/ColorPickerFragment.kt
@@ -22,9 +22,13 @@
 import android.view.ViewGroup
 import android.widget.FrameLayout
 import androidx.cardview.widget.CardView
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
 import androidx.lifecycle.lifecycleScope
+import androidx.transition.Transition
+import androidx.transition.doOnStart
 import com.android.customization.model.mode.DarkModeSectionController
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.color.ui.binder.ColorPickerBinder
@@ -92,53 +96,58 @@
 
         savedInstanceState?.let { binding?.restoreInstanceState(it) }
 
-        ScreenPreviewBinder.bind(
-            activity = requireActivity(),
-            previewView = lockScreenView,
-            viewModel =
-                ScreenPreviewViewModel(
-                    previewUtils =
-                        PreviewUtils(
-                            context = requireContext(),
-                            authority =
-                                requireContext()
-                                    .getString(
-                                        R.string.lock_screen_preview_provider_authority,
-                                    ),
-                        ),
-                    wallpaperInfoProvider = { forceReload ->
-                        suspendCancellableCoroutine { continuation ->
-                            wallpaperInfoFactory.createCurrentWallpaperInfos(
-                                { homeWallpaper, lockWallpaper, _ ->
-                                    lifecycleScope.launch {
-                                        if (
-                                            wcViewModel.lockWallpaperColors.value
-                                                is WallpaperColorsModel.Loading
-                                        ) {
-                                            loadInitialColors(
-                                                wallpaperManager,
-                                                wcViewModel,
-                                                CustomizationSections.Screen.LOCK_SCREEN
-                                            )
+        val lockScreenPreviewBinder =
+            ScreenPreviewBinder.bind(
+                activity = requireActivity(),
+                previewView = lockScreenView,
+                viewModel =
+                    ScreenPreviewViewModel(
+                        previewUtils =
+                            PreviewUtils(
+                                context = requireContext(),
+                                authority =
+                                    requireContext()
+                                        .getString(
+                                            R.string.lock_screen_preview_provider_authority,
+                                        ),
+                            ),
+                        wallpaperInfoProvider = { forceReload ->
+                            suspendCancellableCoroutine { continuation ->
+                                wallpaperInfoFactory.createCurrentWallpaperInfos(
+                                    { homeWallpaper, lockWallpaper, _ ->
+                                        lifecycleScope.launch {
+                                            if (
+                                                wcViewModel.lockWallpaperColors.value
+                                                    is WallpaperColorsModel.Loading
+                                            ) {
+                                                loadInitialColors(
+                                                    wallpaperManager,
+                                                    wcViewModel,
+                                                    CustomizationSections.Screen.LOCK_SCREEN
+                                                )
+                                            }
                                         }
-                                    }
-                                    continuation.resume(lockWallpaper ?: homeWallpaper, null)
-                                },
-                                forceReload,
-                            )
-                        }
-                    },
-                    onWallpaperColorChanged = { colors ->
-                        wcViewModel.setLockWallpaperColors(colors)
-                    },
-                    wallpaperInteractor = injector.getWallpaperInteractor(requireContext()),
-                    screen = CustomizationSections.Screen.LOCK_SCREEN,
-                ),
-            lifecycleOwner = this,
-            offsetToStart =
-                displayUtils.isSingleDisplayOrUnfoldedHorizontalHinge(requireActivity()),
-            onWallpaperPreviewDirty = { activity?.recreate() },
-        )
+                                        continuation.resume(lockWallpaper ?: homeWallpaper, null)
+                                    },
+                                    forceReload,
+                                )
+                            }
+                        },
+                        onWallpaperColorChanged = { colors ->
+                            wcViewModel.setLockWallpaperColors(colors)
+                        },
+                        wallpaperInteractor = injector.getWallpaperInteractor(requireContext()),
+                        screen = CustomizationSections.Screen.LOCK_SCREEN,
+                    ),
+                lifecycleOwner = this,
+                offsetToStart =
+                    displayUtils.isSingleDisplayOrUnfoldedHorizontalHinge(requireActivity()),
+                onWallpaperPreviewDirty = { activity?.recreate() },
+            )
+        val shouldMirrorHomePreview =
+            wallpaperManager.getWallpaperInfo(WallpaperManager.FLAG_SYSTEM) != null &&
+                wallpaperManager.getWallpaperId(WallpaperManager.FLAG_LOCK) < 0
+        val mirrorSurface = if (shouldMirrorHomePreview) lockScreenPreviewBinder.surface() else null
         ScreenPreviewBinder.bind(
             activity = requireActivity(),
             previewView = homeScreenView,
@@ -185,6 +194,7 @@
             offsetToStart =
                 displayUtils.isSingleDisplayOrUnfoldedHorizontalHinge(requireActivity()),
             onWallpaperPreviewDirty = { activity?.recreate() },
+            mirrorSurface = mirrorSurface,
         )
         val darkModeToggleContainerView: FrameLayout =
             view.requireViewById(R.id.dark_mode_toggle_container)
@@ -197,6 +207,12 @@
                 .createView(requireContext())
         darkModeSectionView.background = null
         darkModeToggleContainerView.addView(darkModeSectionView)
+
+        (returnTransition as? Transition)?.doOnStart {
+            lockScreenView.isVisible = false
+            homeScreenView.isVisible = false
+        }
+
         return view
     }
 
@@ -236,4 +252,8 @@
     override fun getToolbarColorId(): Int {
         return android.R.color.transparent
     }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
 }
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
index a828f83..71dfe1d 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithClockCarouselSectionController.kt
@@ -23,7 +23,6 @@
 import android.view.TouchDelegate
 import android.view.View
 import android.view.View.OnAttachStateChangeListener
-import android.view.ViewGroup
 import android.view.ViewStub
 import androidx.activity.ComponentActivity
 import androidx.constraintlayout.helper.widget.Carousel
@@ -39,7 +38,9 @@
 import com.android.customization.picker.clock.ui.view.ClockCarouselView
 import com.android.customization.picker.clock.ui.view.ClockViewFactory
 import com.android.customization.picker.clock.ui.viewmodel.ClockCarouselViewModel
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.wallpaper.R
+import com.android.wallpaper.model.CustomizationSectionController
 import com.android.wallpaper.model.CustomizationSectionController.CustomizationSectionNavigationController
 import com.android.wallpaper.model.WallpaperColorsViewModel
 import com.android.wallpaper.model.WallpaperPreviewNavigator
@@ -49,6 +50,7 @@
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewClickView
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewSectionController
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewView
+import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationPickerViewModel
 import com.android.wallpaper.util.DisplayUtils
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.launch
@@ -70,8 +72,10 @@
     private val navigationController: CustomizationSectionNavigationController,
     wallpaperInteractor: WallpaperInteractor,
     themedIconInteractor: ThemedIconInteractor,
+    colorPickerInteractor: ColorPickerInteractor,
     wallpaperManager: WallpaperManager,
     private val isTwoPaneAndSmallWidth: Boolean,
+    customizationPickerViewModel: CustomizationPickerViewModel,
 ) :
     PreviewWithThemeSectionController(
         activity,
@@ -83,8 +87,10 @@
         wallpaperPreviewNavigator,
         wallpaperInteractor,
         themedIconInteractor,
+        colorPickerInteractor,
         wallpaperManager,
         isTwoPaneAndSmallWidth,
+        customizationPickerViewModel,
     ) {
 
     private val viewModel =
@@ -98,8 +104,11 @@
 
     override val hideLockScreenClockPreview = true
 
-    override fun createView(context: Context): ScreenPreviewView {
-        val view = super.createView(context)
+    override fun createView(
+        context: Context,
+        params: CustomizationSectionController.ViewCreationParams,
+    ): ScreenPreviewView {
+        val view = super.createView(context, params)
         if (screen == CustomizationSections.Screen.LOCK_SCREEN) {
             val screenPreviewClickView: ScreenPreviewClickView =
                 view.requireViewById(R.id.screen_preview_click_view)
@@ -146,12 +155,6 @@
                 guidelineEnd.layoutParams = layoutParams
             }
 
-            // TODO (b/270716937) We should handle the single clock case in the clock carousel
-            // itself
-            val singleClockViewStub: ViewStub = view.requireViewById(R.id.single_clock_view_stub)
-            singleClockViewStub.layoutResource = R.layout.single_clock_view
-            val singleClockView = singleClockViewStub.inflate() as ViewGroup
-
             /**
              * Only bind after [Carousel.onAttachedToWindow]. This is to avoid the race condition
              * that the flow emits before attached to window where [Carousel.mMotionLayout] is still
@@ -167,7 +170,6 @@
                                 ClockCarouselViewBinder.bind(
                                     context = context,
                                     carouselView = carouselView,
-                                    singleClockView = singleClockView,
                                     screenPreviewClickView = screenPreviewClickView,
                                     viewModel = viewModel,
                                     clockViewFactory = clockViewFactory,
diff --git a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
index 56c6c30..c4d6be4 100644
--- a/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
+++ b/src/com/android/customization/picker/preview/ui/section/PreviewWithThemeSectionController.kt
@@ -22,6 +22,7 @@
 import android.content.Context
 import androidx.lifecycle.LifecycleOwner
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.customization.picker.preview.ui.viewmodel.PreviewWithThemeViewModel
 import com.android.wallpaper.R
 import com.android.wallpaper.model.WallpaperColorsViewModel
@@ -30,6 +31,7 @@
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.section.ScreenPreviewSectionController
+import com.android.wallpaper.picker.customization.ui.viewmodel.CustomizationPickerViewModel
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.DisplayUtils
 import com.android.wallpaper.util.PreviewUtils
@@ -49,8 +51,10 @@
     wallpaperPreviewNavigator: WallpaperPreviewNavigator,
     private val wallpaperInteractor: WallpaperInteractor,
     private val themedIconInteractor: ThemedIconInteractor,
+    private val colorPickerInteractor: ColorPickerInteractor,
     wallpaperManager: WallpaperManager,
     isTwoPaneAndSmallWidth: Boolean,
+    customizationPickerViewModel: CustomizationPickerViewModel,
 ) :
     ScreenPreviewSectionController(
         activity,
@@ -62,7 +66,8 @@
         wallpaperPreviewNavigator,
         wallpaperInteractor,
         wallpaperManager,
-        isTwoPaneAndSmallWidth
+        isTwoPaneAndSmallWidth,
+        customizationPickerViewModel,
     ) {
     override fun createScreenPreviewViewModel(context: Context): ScreenPreviewViewModel {
         return PreviewWithThemeViewModel(
@@ -114,6 +119,7 @@
             initialExtrasProvider = { getInitialExtras(isOnLockScreen) },
             wallpaperInteractor = wallpaperInteractor,
             themedIconInteractor = themedIconInteractor,
+            colorPickerInteractor = colorPickerInteractor,
             screen = screen,
         )
     }
diff --git a/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt b/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
index 435878d..83f986d 100644
--- a/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
+++ b/src/com/android/customization/picker/preview/ui/viewmodel/PreviewWithThemeViewModel.kt
@@ -20,12 +20,14 @@
 import android.app.WallpaperColors
 import android.os.Bundle
 import com.android.customization.model.themedicon.domain.interactor.ThemedIconInteractor
+import com.android.customization.picker.color.domain.interactor.ColorPickerInteractor
 import com.android.wallpaper.model.WallpaperInfo
 import com.android.wallpaper.module.CustomizationSections
 import com.android.wallpaper.picker.customization.domain.interactor.WallpaperInteractor
 import com.android.wallpaper.picker.customization.ui.viewmodel.ScreenPreviewViewModel
 import com.android.wallpaper.util.PreviewUtils
 import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.combine
 
 /** A ThemePicker version of the [ScreenPreviewViewModel] */
 class PreviewWithThemeViewModel(
@@ -35,6 +37,7 @@
     onWallpaperColorChanged: (WallpaperColors?) -> Unit = {},
     wallpaperInteractor: WallpaperInteractor,
     private val themedIconInteractor: ThemedIconInteractor? = null,
+    colorPickerInteractor: ColorPickerInteractor? = null,
     screen: CustomizationSections.Screen,
 ) :
     ScreenPreviewViewModel(
@@ -46,4 +49,16 @@
         screen,
     ) {
     override fun workspaceUpdateEvents(): Flow<Boolean>? = themedIconInteractor?.isActivated
+
+    private val wallpaperIsLoading = super.isLoading
+
+    override val isLoading: Flow<Boolean> =
+        colorPickerInteractor?.let {
+            combine(wallpaperIsLoading, colorPickerInteractor.isApplyingSystemColor) {
+                wallpaperIsLoading,
+                colorIsLoading ->
+                wallpaperIsLoading || colorIsLoading
+            }
+        }
+            ?: wallpaperIsLoading
 }
diff --git a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
index 68367c8..091f484 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordancePickerBinder.kt
@@ -37,6 +37,7 @@
 import com.android.wallpaper.picker.common.icon.ui.viewmodel.Icon
 import com.android.wallpaper.picker.option.ui.adapter.OptionItemAdapter
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.flow.collectIndexed
 import kotlinx.coroutines.flow.combine
 import kotlinx.coroutines.flow.distinctUntilChanged
 import kotlinx.coroutines.flow.flatMapLatest
@@ -99,13 +100,18 @@
                                 selectedFlags.indexOfFirst { it }
                             }
                         }
-                        .collect { selectedPosition ->
+                        .collectIndexed { index, selectedPosition ->
                             // Scroll the view to show the first selected affordance.
                             if (selectedPosition != -1) {
                                 // We use "post" because we need to give the adapter item a pass to
                                 // update the view.
                                 affordancesView.post {
-                                    affordancesView.smoothScrollToPosition(selectedPosition)
+                                    if (index == 0) {
+                                        // don't animate on initial collection
+                                        affordancesView.scrollToPosition(selectedPosition)
+                                    } else {
+                                        affordancesView.smoothScrollToPosition(selectedPosition)
+                                    }
                                 }
                             }
                         }
diff --git a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt
index 28ad51a..7e1f4d3 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/binder/KeyguardQuickAffordanceSectionViewBinder.kt
@@ -48,7 +48,7 @@
 
         lifecycleOwner.lifecycleScope.launch {
             viewModel.summary
-                .flowWithLifecycle(lifecycleOwner.lifecycle, Lifecycle.State.RESUMED)
+                .flowWithLifecycle(lifecycleOwner.lifecycle, Lifecycle.State.STARTED)
                 .collectLatest { summary ->
                     TextViewBinder.bind(
                         view = descriptionView,
diff --git a/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt b/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
index d5f0d33..467e5a0 100644
--- a/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
+++ b/src/com/android/customization/picker/quickaffordance/ui/fragment/KeyguardQuickAffordancePickerFragment.kt
@@ -21,8 +21,12 @@
 import android.view.LayoutInflater
 import android.view.View
 import android.view.ViewGroup
+import androidx.core.content.ContextCompat
+import androidx.core.view.isVisible
 import androidx.lifecycle.ViewModelProvider
 import androidx.lifecycle.get
+import androidx.transition.Transition
+import androidx.transition.doOnStart
 import com.android.customization.module.ThemePickerInjector
 import com.android.customization.picker.quickaffordance.ui.binder.KeyguardQuickAffordancePickerBinder
 import com.android.customization.picker.quickaffordance.ui.binder.KeyguardQuickAffordancePreviewBinder
@@ -30,9 +34,7 @@
 import com.android.wallpaper.R
 import com.android.wallpaper.module.InjectorProvider
 import com.android.wallpaper.picker.AppbarFragment
-import kotlinx.coroutines.ExperimentalCoroutinesApi
 
-@OptIn(ExperimentalCoroutinesApi::class)
 class KeyguardQuickAffordancePickerFragment : AppbarFragment() {
     companion object {
         const val DESTINATION_ID = "quick_affordances"
@@ -77,6 +79,12 @@
             viewModel = viewModel,
             lifecycleOwner = this,
         )
+        postponeEnterTransition()
+        view.post { startPostponedEnterTransition() }
+        (returnTransition as? Transition)?.doOnStart {
+            // Hide preview during exit transition animation
+            view?.findViewById<View>(R.id.preview)?.isVisible = false
+        }
         return view
     }
 
@@ -87,4 +95,8 @@
     override fun getToolbarColorId(): Int {
         return android.R.color.transparent
     }
+
+    override fun getToolbarTextColor(): Int {
+        return ContextCompat.getColor(requireContext(), R.color.system_on_surface)
+    }
 }
diff --git a/src/com/android/customization/widget/GridTileDrawable.java b/src/com/android/customization/widget/GridTileDrawable.java
index 83cd0b5..b9d2036 100644
--- a/src/com/android/customization/widget/GridTileDrawable.java
+++ b/src/com/android/customization/widget/GridTileDrawable.java
@@ -55,8 +55,8 @@
         for (int r = 0; r < mRows; r++) {
             for (int c = 0; c < mCols; c++) {
                 int saveCount = canvas.save();
-                float x = (float) ((r * size / mRows) + SPACE_BETWEEN_ICONS);
-                float y = (float) ((c * size / mCols) + SPACE_BETWEEN_ICONS);
+                float y = (float) ((r * size / mRows) + SPACE_BETWEEN_ICONS);
+                float x = (float) ((c * size / mCols) + SPACE_BETWEEN_ICONS);
                 canvas.translate(x, y);
                 canvas.drawPath(mTransformedPath, mPaint);
                 canvas.restoreToCount(saveCount);
diff --git a/src/com/android/customization/picker/CustomizationPickerApplication.java b/src_override/com/android/customization/picker/CustomizationPickerApplication.java
similarity index 67%
rename from src/com/android/customization/picker/CustomizationPickerApplication.java
rename to src_override/com/android/customization/picker/CustomizationPickerApplication.java
index 178cfbf..f2000f6 100644
--- a/src/com/android/customization/picker/CustomizationPickerApplication.java
+++ b/src_override/com/android/customization/picker/CustomizationPickerApplication.java
@@ -17,15 +17,27 @@
 
 import android.app.Application;
 
-import com.android.customization.module.ThemePickerInjector;
+import com.android.customization.module.CustomizationInjector;
 import com.android.wallpaper.module.InjectorProvider;
 
-public class CustomizationPickerApplication extends Application {
+import dagger.hilt.android.HiltAndroidApp;
+
+import javax.inject.Inject;
+
+/**
+ * Application subclass that initializes the injector.
+ */
+@HiltAndroidApp(Application.class)
+public class CustomizationPickerApplication extends Hilt_CustomizationPickerApplication {
+
+    @Inject
+    CustomizationInjector mInjector;
+
     @Override
     public void onCreate() {
         super.onCreate();
 
         // Initialize the injector.
-        InjectorProvider.setInjector(new ThemePickerInjector());
+        InjectorProvider.setInjector(mInjector);
     }
 }
diff --git a/src_override/com/android/wallpaper/module/AppModule.kt b/src_override/com/android/wallpaper/module/AppModule.kt
new file mode 100644
index 0000000..95da101
--- /dev/null
+++ b/src_override/com/android/wallpaper/module/AppModule.kt
@@ -0,0 +1,44 @@
+/*
+ * Copyright (C) 2023 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.wallpaper.module
+
+import android.content.Context
+import com.android.customization.module.CustomizationInjector
+import com.android.customization.module.DefaultCustomizationPreferences
+import com.android.customization.module.ThemePickerInjector
+import dagger.Binds
+import dagger.Module
+import dagger.Provides
+import dagger.hilt.InstallIn
+import dagger.hilt.android.qualifiers.ApplicationContext
+import dagger.hilt.components.SingletonComponent
+import javax.inject.Singleton
+
+@Module
+@InstallIn(SingletonComponent::class)
+abstract class AppModule {
+    @Binds @Singleton abstract fun bindInjector(impl: ThemePickerInjector): CustomizationInjector
+
+    companion object {
+        @Provides
+        @Singleton
+        fun provideWallpaperPreferences(
+            @ApplicationContext context: Context
+        ): WallpaperPreferences {
+            return DefaultCustomizationPreferences(context)
+        }
+    }
+}
diff --git a/tests/common/Android.bp b/tests/common/Android.bp
new file mode 100644
index 0000000..2904fad
--- /dev/null
+++ b/tests/common/Android.bp
@@ -0,0 +1,39 @@
+// Copyright (C) 2023 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.
+//
+
+//
+// Build rule for WallpaperPicker2 tests
+//
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+android_library {
+    name: "ThemePickerTestLib",
+
+    defaults: ["ThemePicker_defaults"],
+    srcs: [
+        "src/com/android/customization/testing/**/*.java",
+        "src/com/android/customization/testing/**/*.kt",
+    ],
+    static_libs: [
+        "WallpaperPicker2TestLib",
+        "androidx.annotation_annotation",
+        "kotlinx_coroutines_test",
+        "truth-prebuilt",
+    ],
+
+    platform_apis: true,
+}
diff --git a/tests/common/AndroidManifest.xml b/tests/common/AndroidManifest.xml
new file mode 100644
index 0000000..81a2f03
--- /dev/null
+++ b/tests/common/AndroidManifest.xml
@@ -0,0 +1,21 @@
+<?xml version="1.0" encoding="utf-8"?>
+<!-- Copyright (C) 2023 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.
+-->
+<manifest xmlns:android="http://schemas.android.com/apk/res/android"
+    package="com.android.wallpaper">
+
+  <application>
+  </application>
+</manifest>
diff --git a/tests/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java b/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
similarity index 61%
rename from tests/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
rename to tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
index bcf5a5f..81890f0 100644
--- a/tests/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
+++ b/tests/common/src/com/android/customization/testing/TestDefaultCustomizationPreferences.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright (C) 2019 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.customization.testing;
 
 import android.content.Context;
diff --git a/tests/common/src/com/android/customization/testing/TestDrawableLayerResolver.java b/tests/common/src/com/android/customization/testing/TestDrawableLayerResolver.java
new file mode 100644
index 0000000..8b16299
--- /dev/null
+++ b/tests/common/src/com/android/customization/testing/TestDrawableLayerResolver.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright (C) 2019 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.customization.testing;
+
+import android.graphics.drawable.Drawable;
+import android.graphics.drawable.LayerDrawable;
+
+import com.android.wallpaper.module.DrawableLayerResolver;
+
+/**
+ * Test implementation of {@link DrawableLayerResolver}.
+ */
+public class TestDrawableLayerResolver implements DrawableLayerResolver {
+    @Override
+    public Drawable resolveLayer(LayerDrawable layerDrawable) {
+        return layerDrawable.getDrawable(0);
+    }
+}
diff --git a/tests/common/src/com/android/customization/testing/TestPackageStatusNotifier.java b/tests/common/src/com/android/customization/testing/TestPackageStatusNotifier.java
new file mode 100644
index 0000000..2aadae8
--- /dev/null
+++ b/tests/common/src/com/android/customization/testing/TestPackageStatusNotifier.java
@@ -0,0 +1,33 @@
+/*
+ * Copyright (C) 2019 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.customization.testing;
+
+import com.android.wallpaper.module.PackageStatusNotifier;
+
+/**
+ * Test implementation of {@link PackageStatusNotifier}.
+ */
+public class TestPackageStatusNotifier implements PackageStatusNotifier {
+    @Override
+    public void addListener(Listener listener, String action) {
+        // Do nothing
+    }
+
+    @Override
+    public void removeListener(Listener listener) {
+        // Do nothing
+    }
+}
diff --git a/tests/src/com/android/customization/testing/TestPluginManager.kt b/tests/common/src/com/android/customization/testing/TestPluginManager.kt
similarity index 100%
rename from tests/src/com/android/customization/testing/TestPluginManager.kt
rename to tests/common/src/com/android/customization/testing/TestPluginManager.kt
diff --git a/tests/src/com/android/customization/testing/TestThemeManager.java b/tests/common/src/com/android/customization/testing/TestThemeManager.java
similarity index 60%
rename from tests/src/com/android/customization/testing/TestThemeManager.java
rename to tests/common/src/com/android/customization/testing/TestThemeManager.java
index c4d25fb..5175b24 100644
--- a/tests/src/com/android/customization/testing/TestThemeManager.java
+++ b/tests/common/src/com/android/customization/testing/TestThemeManager.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright (C) 2019 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.customization.testing;
 
 import androidx.fragment.app.FragmentActivity;
diff --git a/tests/src/com/android/customization/testing/TestThemesUserEventLogger.java b/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.java
similarity index 63%
rename from tests/src/com/android/customization/testing/TestThemesUserEventLogger.java
rename to tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.java
index 2bb2082..22a5b94 100644
--- a/tests/src/com/android/customization/testing/TestThemesUserEventLogger.java
+++ b/tests/common/src/com/android/customization/testing/TestThemesUserEventLogger.java
@@ -1,3 +1,18 @@
+/*
+ * Copyright (C) 2019 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.customization.testing;
 
 import com.android.customization.model.color.ColorOption;
diff --git a/tests/robotests/Android.bp b/tests/robotests/Android.bp
index 4416e1c..ed684f4 100644
--- a/tests/robotests/Android.bp
+++ b/tests/robotests/Android.bp
@@ -2,13 +2,23 @@
 package {
     default_applicable_licenses: ["Android-Apache-2.0"],
 }
+
 android_robolectric_test {
     name: "ThemePickerRoboTests",
     srcs: [
         "src/**/*.java",
         "src/**/*.kt",
     ],
+    // TODO(b/291104503) Enable this test
+    exclude_srcs: ["src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt"],
     java_resource_dirs: ["config"],
+    static_libs: [
+        "WallpaperPicker2TestLib",
+        "androidx.test.rules",
+        "junit",
+        "kotlinx_coroutines_test",
+        "truth-prebuilt",
+    ],
     libs: [
         "androidx.test.core",
         "androidx.test.runner",
diff --git a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java b/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
deleted file mode 100644
index 820e641..0000000
--- a/tests/robotests/src/com/android/customization/model/color/ColorSectionControllerTest.java
+++ /dev/null
@@ -1,62 +0,0 @@
-/*
- * Copyright (C) 2022 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.customization.model.color;
-
-import static com.google.common.truth.Truth.assertThat;
-
-import androidx.appcompat.app.AppCompatActivity;
-
-import com.android.wallpaper.model.WallpaperColorsViewModel;
-
-import org.junit.Before;
-import org.junit.Test;
-import org.junit.runner.RunWith;
-import org.mockito.MockitoAnnotations;
-import org.robolectric.Robolectric;
-import org.robolectric.RobolectricTestRunner;
-import org.robolectric.annotation.Config;
-
-/**
- * Tests of {@link ColorSectionController}.
- */
-@RunWith(RobolectricTestRunner.class)
-@Config(manifest = Config.NONE)
-public final class ColorSectionControllerTest {
-
-    private AppCompatActivity mActivity;
-    private ColorSectionController mColorSectionController;
-
-    /**
-     * Set up the test case.
-     */
-    @Before
-    public void setUp() {
-        MockitoAnnotations.initMocks(this);
-        mActivity = Robolectric.buildActivity(AppCompatActivity.class).create().get();
-        mColorSectionController = new ColorSectionController(mActivity,
-                new WallpaperColorsViewModel(), mActivity, null);
-    }
-
-    /**
-     * isAvailable()'s test.
-     */
-    @Test
-    @Config(manifest = Config.NONE)
-    public void isAvailable_nullContext_shouldReturnFalse() {
-        assertThat(mColorSectionController.isAvailable(/* context= */ null)).isFalse();
-    }
-}
-
diff --git a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt b/tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
similarity index 90%
rename from tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
rename to tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
index 5953937..317ad3a 100644
--- a/tests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
+++ b/tests/robotests/src/com/android/customization/model/grid/data/repository/FakeGridRepository.kt
@@ -17,6 +17,8 @@
 
 package com.android.customization.model.grid.data.repository
 
+import com.android.customization.model.CustomizationManager
+import com.android.customization.model.grid.GridOption
 import com.android.customization.model.grid.shared.model.GridOptionItemModel
 import com.android.customization.model.grid.shared.model.GridOptionItemsModel
 import kotlinx.coroutines.CoroutineScope
@@ -51,6 +53,14 @@
         return options
     }
 
+    override fun getSelectedOption(): GridOption? = null
+
+    override fun applySelectedOption(callback: CustomizationManager.Callback) {}
+
+    override fun clearSelectedOption() {}
+
+    override fun isSelectedOptionApplied() = false
+
     fun setOptions(
         count: Int,
         selectedIndex: Int = 0,
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt b/tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridInteractorTest.kt
diff --git a/tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/grid/domain/interactor/GridSnapshotRestorerTest.kt
diff --git a/tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt b/tests/robotests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
rename to tests/robotests/src/com/android/customization/model/grid/ui/viewmodel/GridScreenViewModelTest.kt
diff --git a/tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/mode/DarkModeSnapshotRestorerTest.kt
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
index cbf1365..d4f24ee 100644
--- a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerInteractorTest.kt
@@ -32,11 +32,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ColorPickerInteractorTest {
     private lateinit var underTest: ColorPickerInteractor
     private lateinit var repository: FakeColorPickerRepository
diff --git a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
index 71a8f23..5f3e39e 100644
--- a/tests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/color/domain/interactor/ColorPickerSnapshotRestorerTest.kt
@@ -34,11 +34,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ColorPickerSnapshotRestorerTest {
 
     private lateinit var underTest: ColorPickerSnapshotRestorer
diff --git a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
index a29d559..9968c5f 100644
--- a/tests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/color/ui/viewmodel/ColorPickerViewModelTest.kt
@@ -43,11 +43,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ColorPickerViewModelTest {
     private lateinit var underTest: ColorPickerViewModel
     private lateinit var repository: FakeColorPickerRepository
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
similarity index 97%
rename from tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
index 3f22ced..35dbadd 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/data/repository/KeyguardQuickAffordancePickerRepositoryTest.kt
@@ -35,11 +35,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class KeyguardQuickAffordancePickerRepositoryTest {
 
     private lateinit var underTest: KeyguardQuickAffordancePickerRepository
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
index fea94dc..efe9f64 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/domain/interactor/KeyguardQuickAffordancePickerInteractorTest.kt
@@ -39,11 +39,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class KeyguardQuickAffordancePickerInteractorTest {
 
     private lateinit var underTest: KeyguardQuickAffordancePickerInteractor
diff --git a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
similarity index 98%
rename from tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
rename to tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
index 3f10674..f71bfc7 100644
--- a/tests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/model/picker/quickaffordance/ui/viewmodel/KeyguardQuickAffordancePickerViewModelTest.kt
@@ -19,8 +19,8 @@
 
 import android.content.Context
 import android.content.Intent
+import androidx.test.core.app.ApplicationProvider
 import androidx.test.filters.SmallTest
-import androidx.test.platform.app.InstrumentationRegistry
 import com.android.customization.picker.quickaffordance.data.repository.KeyguardQuickAffordancePickerRepository
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordancePickerInteractor
 import com.android.customization.picker.quickaffordance.domain.interactor.KeyguardQuickAffordanceSnapshotRestorer
@@ -57,11 +57,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class KeyguardQuickAffordancePickerViewModelTest {
 
     private lateinit var underTest: KeyguardQuickAffordancePickerViewModel
@@ -75,7 +75,7 @@
     @Before
     fun setUp() {
         InjectorProvider.setInjector(TestInjector())
-        context = InstrumentationRegistry.getInstrumentation().targetContext
+        context = ApplicationProvider.getApplicationContext()
         val testDispatcher = StandardTestDispatcher()
         testScope = TestScope(testDispatcher)
         Dispatchers.setMain(testDispatcher)
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt b/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
rename to tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconInteractorTest.kt
diff --git a/tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt b/tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
similarity index 100%
rename from tests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
rename to tests/robotests/src/com/android/customization/model/themedicon/domain/interactor/ThemedIconSnapshotRestorerTest.kt
diff --git a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt b/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
similarity index 84%
rename from tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
rename to tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
index bf2766d..95d7e35 100644
--- a/tests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/data/repository/FakeClockPickerRepository.kt
@@ -45,11 +45,12 @@
             val selectedClock = fakeClocks.find { clock -> clock.clockId == selectedClockId }
             checkNotNull(selectedClock)
             ClockMetadataModel(
-                selectedClock.clockId,
-                selectedClock.name,
-                selectedColor,
-                colorTone,
-                seedColor,
+                clockId = selectedClock.clockId,
+                name = selectedClock.name,
+                isSelected = true,
+                selectedColorId = selectedColor,
+                colorToneProgress = colorTone,
+                seedColor = seedColor,
             )
         }
 
@@ -81,10 +82,10 @@
         const val CLOCK_ID_3 = "clock3"
         val fakeClocks =
             listOf(
-                ClockMetadataModel(CLOCK_ID_0, "clock0", null, 50, null),
-                ClockMetadataModel(CLOCK_ID_1, "clock1", null, 50, null),
-                ClockMetadataModel(CLOCK_ID_2, "clock2", null, 50, null),
-                ClockMetadataModel(CLOCK_ID_3, "clock3", null, 50, null),
+                ClockMetadataModel(CLOCK_ID_0, "clock0", true, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_1, "clock1", false, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_2, "clock2", false, null, 50, null),
+                ClockMetadataModel(CLOCK_ID_3, "clock3", false, null, 50, null),
             )
         const val CLOCK_COLOR_ID = "RED"
         const val CLOCK_COLOR_TONE_PROGRESS = 87
diff --git a/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt b/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
similarity index 97%
rename from tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
index 1a7ebb5..c8e39be 100644
--- a/tests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/domain/interactor/ClockPickerInteractorTest.kt
@@ -17,11 +17,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockPickerInteractorTest {
 
     private lateinit var underTest: ClockPickerInteractor
diff --git a/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
deleted file mode 100644
index 0a54312..0000000
--- a/tests/robotests/src/com/android/customization/picker/clock/ui/fragment/ClockCustomDemoFragmentTest.kt
+++ /dev/null
@@ -1,75 +0,0 @@
-package com.android.customization.picker.clock.ui.fragment
-
-import android.view.View
-import androidx.appcompat.app.AppCompatActivity
-import androidx.recyclerview.widget.LinearLayoutManager
-import androidx.recyclerview.widget.RecyclerView
-import com.android.systemui.plugins.ClockMetadata
-import com.android.systemui.plugins.ClockSettings
-import com.android.systemui.plugins.PluginManager
-import com.android.systemui.shared.clocks.ClockRegistry
-import org.junit.Assert
-import org.junit.Before
-import org.junit.Ignore
-import org.junit.Test
-import org.junit.runner.RunWith
-import org.mockito.Mock
-import org.mockito.Mockito.verify
-import org.mockito.Mockito.`when` as whenever
-import org.mockito.MockitoAnnotations
-import org.robolectric.Robolectric
-import org.robolectric.RobolectricTestRunner
-import org.robolectric.annotation.Config
-
-/** Tests of [ClockCustomDemoFragment]. */
-@RunWith(RobolectricTestRunner::class)
-@Config(manifest = Config.NONE)
-@Ignore("b/270606895")
-class ClockCustomDemoFragmentTest {
-    private lateinit var mActivity: AppCompatActivity
-    private var mClockCustomDemoFragment: ClockCustomDemoFragment? = null
-    @Mock private lateinit var registry: ClockRegistry
-    @Mock private lateinit var mockPluginManager: PluginManager
-
-    private var settingValue: ClockSettings? = null
-
-    @Before
-    fun setUp() {
-        MockitoAnnotations.initMocks(this)
-        mActivity = Robolectric.buildActivity(AppCompatActivity::class.java).get()
-        mClockCustomDemoFragment = ClockCustomDemoFragment()
-        whenever(registry.getClocks())
-            .thenReturn(
-                listOf(
-                    ClockMetadata("CLOCK_1", "Clock 1"),
-                    ClockMetadata("CLOCK_2", "Clock 2"),
-                    ClockMetadata("CLOCK_NOT_IN_USE", "Clock not in use")
-                )
-            )
-
-        mClockCustomDemoFragment!!.clockRegistry = registry
-        mClockCustomDemoFragment!!.recyclerView = RecyclerView(mActivity)
-        mClockCustomDemoFragment!!.recyclerView.layoutManager =
-            LinearLayoutManager(mActivity, RecyclerView.VERTICAL, false)
-    }
-
-    @Test
-    fun testItemCount_getCorrectClockCount() {
-        Assert.assertEquals(3, mClockCustomDemoFragment!!.recyclerView.adapter!!.itemCount)
-    }
-
-    @Test
-    fun testClick_setCorrectClockId() {
-        mClockCustomDemoFragment!!
-            .recyclerView
-            .measure(View.MeasureSpec.UNSPECIFIED, View.MeasureSpec.UNSPECIFIED)
-        mClockCustomDemoFragment!!.recyclerView.layout(0, 0, 100, 10000)
-        val testPosition = 1
-        mClockCustomDemoFragment!!
-            .recyclerView
-            .findViewHolderForAdapterPosition(testPosition)
-            ?.itemView
-            ?.performClick()
-        verify(registry).currentClockId = "CLOCK_1"
-    }
-}
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
similarity index 70%
rename from tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
index c5eb796..ca6f8c7 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockCarouselViewModelTest.kt
@@ -37,22 +37,23 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockCarouselViewModelTest {
     private val repositoryWithMultipleClocks by lazy { FakeClockPickerRepository() }
     private val repositoryWithSingleClock by lazy {
         FakeClockPickerRepository(
             listOf(
                 ClockMetadataModel(
-                    "clock0",
-                    "clock0",
-                    null,
-                    ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
-                    null,
+                    clockId = "clock0",
+                    name = "clock0",
+                    isSelected = true,
+                    selectedColorId = null,
+                    colorToneProgress = ClockMetadataModel.DEFAULT_COLOR_TONE_PROGRESS,
+                    seedColor = null,
                 ),
             )
         )
@@ -87,38 +88,6 @@
         assertThat(observedSelectedIndex()).isEqualTo(2)
     }
 
-    @Test
-    fun multipleClockCase() = runTest {
-        underTest =
-            ClockCarouselViewModel(
-                getClockPickerInteractor(repositoryWithMultipleClocks),
-                testDispatcher
-            )
-        val observedIsCarouselVisible = collectLastValue(underTest.isCarouselVisible)
-        val observedIsSingleClockViewVisible = collectLastValue(underTest.isSingleClockViewVisible)
-
-        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-
-        assertThat(observedIsCarouselVisible()).isTrue()
-        assertThat(observedIsSingleClockViewVisible()).isFalse()
-    }
-
-    @Test
-    fun singleClockCase() = runTest {
-        underTest =
-            ClockCarouselViewModel(
-                getClockPickerInteractor(repositoryWithSingleClock),
-                testDispatcher
-            )
-        val observedIsCarouselVisible = collectLastValue(underTest.isCarouselVisible)
-        val observedIsSingleClockViewVisible = collectLastValue(underTest.isSingleClockViewVisible)
-
-        advanceTimeBy(ClockCarouselViewModel.CLOCKS_EVENT_UPDATE_DELAY_MILLIS)
-
-        assertThat(observedIsCarouselVisible()).isFalse()
-        assertThat(observedIsSingleClockViewVisible()).isTrue()
-    }
-
     private fun getClockPickerInteractor(repository: ClockPickerRepository): ClockPickerInteractor {
         return ClockPickerInteractor(
                 repository = repository,
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
similarity index 97%
rename from tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
index 293e393..19a704c 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSectionViewModelTest.kt
@@ -36,11 +36,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockSectionViewModelTest {
 
     private lateinit var clockColorMap: Map<String, ClockColorViewModel>
diff --git a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
similarity index 98%
rename from tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
index f58baf8..f09e977 100644
--- a/tests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/clock/ui/viewmodel/ClockSettingsViewModelTest.kt
@@ -28,11 +28,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class ClockSettingsViewModelTest {
 
     private lateinit var context: Context
diff --git a/tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
similarity index 100%
rename from tests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
rename to tests/robotests/src/com/android/customization/picker/notifications/data/repository/NotificationsRepositoryTest.kt
diff --git a/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
similarity index 97%
rename from tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
rename to tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
index 63334f2..5c3544a 100644
--- a/tests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
+++ b/tests/robotests/src/com/android/customization/picker/notifications/ui/viewmodel/NotificationSectionViewModelTest.kt
@@ -37,11 +37,11 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
-import org.junit.runners.JUnit4
+import org.robolectric.RobolectricTestRunner
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
-@RunWith(JUnit4::class)
+@RunWith(RobolectricTestRunner::class)
 class NotificationSectionViewModelTest {
 
     private lateinit var underTest: NotificationSectionViewModel
diff --git a/tests/src/com/android/customization/testing/TestDrawableLayerResolver.java b/tests/src/com/android/customization/testing/TestDrawableLayerResolver.java
deleted file mode 100644
index e507221..0000000
--- a/tests/src/com/android/customization/testing/TestDrawableLayerResolver.java
+++ /dev/null
@@ -1,16 +0,0 @@
-package com.android.customization.testing;
-
-import android.graphics.drawable.Drawable;
-import android.graphics.drawable.LayerDrawable;
-
-import com.android.wallpaper.module.DrawableLayerResolver;
-
-/**
- * Test implementation of {@link DrawableLayerResolver}.
- */
-public class TestDrawableLayerResolver implements DrawableLayerResolver {
-    @Override
-    public Drawable resolveLayer(LayerDrawable layerDrawable) {
-        return layerDrawable.getDrawable(0);
-    }
-}
diff --git a/tests/src/com/android/customization/testing/TestPackageStatusNotifier.java b/tests/src/com/android/customization/testing/TestPackageStatusNotifier.java
deleted file mode 100644
index 1e6a1a2..0000000
--- a/tests/src/com/android/customization/testing/TestPackageStatusNotifier.java
+++ /dev/null
@@ -1,18 +0,0 @@
-package com.android.customization.testing;
-
-import com.android.wallpaper.module.PackageStatusNotifier;
-
-/**
- * Test implementation of {@link PackageStatusNotifier}.
- */
-public class TestPackageStatusNotifier implements PackageStatusNotifier {
-    @Override
-    public void addListener(Listener listener, String action) {
-        // Do nothing
-    }
-
-    @Override
-    public void removeListener(Listener listener) {
-        // Do nothing
-    }
-}