diff --git a/parts/AndroidManifest.xml b/parts/AndroidManifest.xml
index 49811ff..cd917ce 100644
--- a/parts/AndroidManifest.xml
+++ b/parts/AndroidManifest.xml
@@ -48,7 +48,7 @@
         <activity
             android:name=".doze.DozeSettingsActivity"
             android:label="@string/ambient_display_title"
-            android:theme="@style/Theme.Main">
+            android:theme="@style/Theme.SubSettingsBase">
             <intent-filter>
                 <action android:name="org.lineageos.settings.device.DOZE_SETTINGS" />
                 <category android:name="android.intent.category.DEFAULT" />
diff --git a/parts/res/layout/doze.xml b/parts/res/layout/doze.xml
deleted file mode 100644
index 941cdf6..0000000
--- a/parts/res/layout/doze.xml
+++ /dev/null
@@ -1,28 +0,0 @@
-<?xml version="1.0" encoding="utf-8"?>
-<!--
-/*
-** Copyright 2014, 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.
-*/
--->
-
-<LinearLayout xmlns:android="http://schemas.android.com/apk/res/android"
-              android:orientation="vertical"
-              android:layout_height="match_parent"
-              android:layout_width="match_parent">
-
-    <include layout="@layout/switch_bar" />
-
-</LinearLayout>
-
diff --git a/parts/res/values/strings.xml b/parts/res/values/strings.xml
index 476cdb7..0818968 100644
--- a/parts/res/values/strings.xml
+++ b/parts/res/values/strings.xml
@@ -15,6 +15,9 @@
      limitations under the License.
 -->
 <resources>
+    <!-- Ambient settings -->
+    <string name="ambient_enable">Use ambient display</string>
+
     <string name="doze_brightness_title">Always on brightness</string>
     <string name="doze_brightness_summary">%s mode is enabled</string>
 </resources>
diff --git a/parts/res/xml/doze_settings.xml b/parts/res/xml/doze_settings.xml
index 01ef6c0..75710c3 100644
--- a/parts/res/xml/doze_settings.xml
+++ b/parts/res/xml/doze_settings.xml
@@ -15,7 +15,14 @@
      See the License for the specific language governing permissions and
      limitations under the License.
 -->
-<PreferenceScreen xmlns:android="http://schemas.android.com/apk/res/android">
+<PreferenceScreen
+    xmlns:android="http://schemas.android.com/apk/res/android"
+    android:title="@string/ambient_display_title">
+
+    <com.android.settingslib.widget.MainSwitchPreference
+        android:defaultValue="false"
+        android:key="doze_enable"
+        android:title="@string/ambient_enable" />
 
     <SwitchPreference
         android:key="always_on_display"
diff --git a/parts/src/org/lineageos/settings/doze/DozeSettingsActivity.java b/parts/src/org/lineageos/settings/doze/DozeSettingsActivity.java
index 18eed11..c1122f7 100644
--- a/parts/src/org/lineageos/settings/doze/DozeSettingsActivity.java
+++ b/parts/src/org/lineageos/settings/doze/DozeSettingsActivity.java
@@ -18,9 +18,11 @@
 package org.lineageos.settings.doze;
 
 import android.os.Bundle;
-import android.preference.PreferenceActivity;
 
-public class DozeSettingsActivity extends PreferenceActivity {
+import com.android.settingslib.collapsingtoolbar.CollapsingToolbarBaseActivity;
+import com.android.settingslib.collapsingtoolbar.R;
+
+public class DozeSettingsActivity extends CollapsingToolbarBaseActivity {
     private static final String TAG_DOZE = "doze";
 
     @Override
@@ -29,7 +31,7 @@ public class DozeSettingsActivity extends PreferenceActivity {
 
         getFragmentManager()
                 .beginTransaction()
-                .replace(android.R.id.content, new DozeSettingsFragment(), TAG_DOZE)
+                .replace(R.id.content_frame, new DozeSettingsFragment(), TAG_DOZE)
                 .commit();
     }
 }
diff --git a/parts/src/org/lineageos/settings/doze/DozeSettingsFragment.java b/parts/src/org/lineageos/settings/doze/DozeSettingsFragment.java
index 5588581..1e92e53 100644
--- a/parts/src/org/lineageos/settings/doze/DozeSettingsFragment.java
+++ b/parts/src/org/lineageos/settings/doze/DozeSettingsFragment.java
@@ -17,7 +17,6 @@
 
 package org.lineageos.settings.doze;
 
-import android.app.ActionBar;
 import android.app.Activity;
 import android.app.AlertDialog;
 import android.app.Dialog;
@@ -27,13 +26,7 @@ import android.content.DialogInterface;
 import android.content.SharedPreferences;
 import android.os.Bundle;
 import android.os.Handler;
-import android.view.LayoutInflater;
-import android.view.MenuItem;
-import android.view.View;
-import android.view.ViewGroup;
-import android.widget.CompoundButton;
 import android.widget.Switch;
-import android.widget.TextView;
 import androidx.preference.ListPreference;
 import androidx.preference.Preference;
 import androidx.preference.Preference.OnPreferenceChangeListener;
@@ -41,13 +34,15 @@ import androidx.preference.PreferenceCategory;
 import androidx.preference.PreferenceFragment;
 import androidx.preference.SwitchPreference;
 
+import com.android.settingslib.widget.MainSwitchPreference;
+import com.android.settingslib.widget.OnMainSwitchChangeListener;
+
 import org.lineageos.settings.R;
 import org.lineageos.settings.utils.FileUtils;
 
 public class DozeSettingsFragment extends PreferenceFragment
-        implements OnPreferenceChangeListener, CompoundButton.OnCheckedChangeListener {
-    private TextView mTextView;
-    private View mSwitchBar;
+        implements OnPreferenceChangeListener, OnMainSwitchChangeListener {
+    private MainSwitchPreference mSwitchBar;
 
     private SwitchPreference mAlwaysOnDisplayPreference;
     private ListPreference mDozeBrightnessPreference;
@@ -61,8 +56,6 @@ public class DozeSettingsFragment extends PreferenceFragment
     @Override
     public void onCreatePreferences(Bundle savedInstanceState, String rootKey) {
         addPreferencesFromResource(R.xml.doze_settings);
-        final ActionBar actionBar = getActivity().getActionBar();
-        actionBar.setDisplayHomeAsUpEnabled(true);
 
         SharedPreferences prefs =
                 getActivity().getSharedPreferences("doze_settings", Activity.MODE_PRIVATE);
@@ -72,6 +65,10 @@ public class DozeSettingsFragment extends PreferenceFragment
 
         boolean dozeEnabled = DozeUtils.isDozeEnabled(getActivity());
 
+        mSwitchBar = (MainSwitchPreference) findPreference(DozeUtils.DOZE_ENABLE);
+        mSwitchBar.addOnSwitchChangeListener(this);
+        mSwitchBar.setChecked(dozeEnabled);
+
         mAlwaysOnDisplayPreference = (SwitchPreference) findPreference(DozeUtils.ALWAYS_ON_DISPLAY);
         mAlwaysOnDisplayPreference.setEnabled(dozeEnabled);
         mAlwaysOnDisplayPreference.setChecked(DozeUtils.isAlwaysOnEnabled(getActivity()));
@@ -126,36 +123,6 @@ public class DozeSettingsFragment extends PreferenceFragment
         }
     }
 
-    @Override
-    public View onCreateView(
-            LayoutInflater inflater, ViewGroup container, Bundle savedInstanceState) {
-        final View view =
-                LayoutInflater.from(getContext()).inflate(R.layout.doze, container, false);
-        ((ViewGroup) view).addView(super.onCreateView(inflater, container, savedInstanceState));
-        return view;
-    }
-
-    @Override
-    public void onViewCreated(View view, Bundle savedInstanceState) {
-        super.onViewCreated(view, savedInstanceState);
-
-        boolean dozeEnabled = DozeUtils.isDozeEnabled(getActivity());
-
-        mTextView = view.findViewById(R.id.switch_text);
-        mTextView.setText(
-                getString(dozeEnabled ? R.string.switch_bar_on : R.string.switch_bar_off));
-
-        mSwitchBar = view.findViewById(R.id.switch_bar);
-        Switch switchWidget = mSwitchBar.findViewById(android.R.id.switch_widget);
-        switchWidget.setChecked(dozeEnabled);
-        switchWidget.setOnCheckedChangeListener(this);
-        mSwitchBar.setActivated(dozeEnabled);
-        mSwitchBar.setOnClickListener(v -> {
-            switchWidget.setChecked(!switchWidget.isChecked());
-            mSwitchBar.setActivated(switchWidget.isChecked());
-        });
-    }
-
     @Override
     public boolean onPreferenceChange(Preference preference, Object newValue) {
         if (DozeUtils.ALWAYS_ON_DISPLAY.equals(preference.getKey())) {
@@ -184,12 +151,11 @@ public class DozeSettingsFragment extends PreferenceFragment
     }
 
     @Override
-    public void onCheckedChanged(CompoundButton compoundButton, boolean isChecked) {
+    public void onSwitchChanged(Switch switchView, boolean isChecked) {
         DozeUtils.enableDoze(getActivity(), isChecked);
         DozeUtils.checkDozeService(getActivity());
 
-        mTextView.setText(getString(isChecked ? R.string.switch_bar_on : R.string.switch_bar_off));
-        mSwitchBar.setActivated(isChecked);
+        mSwitchBar.setChecked(isChecked);
 
         if (!isChecked) {
             DozeUtils.enableAlwaysOn(getActivity(), false);
@@ -209,15 +175,6 @@ public class DozeSettingsFragment extends PreferenceFragment
         mPocketPreference.setEnabled(isChecked);
     }
 
-    @Override
-    public boolean onOptionsItemSelected(MenuItem item) {
-        if (item.getItemId() == android.R.id.home) {
-            getActivity().onBackPressed();
-            return true;
-        }
-        return false;
-    }
-
     private static class HelpDialogFragment extends DialogFragment {
         @Override
         public Dialog onCreateDialog(Bundle savedInstanceState) {
diff --git a/parts/src/org/lineageos/settings/doze/DozeUtils.java b/parts/src/org/lineageos/settings/doze/DozeUtils.java
index eec61db..8d13cca 100644
--- a/parts/src/org/lineageos/settings/doze/DozeUtils.java
+++ b/parts/src/org/lineageos/settings/doze/DozeUtils.java
@@ -43,6 +43,7 @@ public final class DozeUtils {
 
     private static final String DOZE_INTENT = "com.android.systemui.doze.pulse";
 
+    protected static final String DOZE_ENABLE = "doze_enable";
     protected static final String ALWAYS_ON_DISPLAY = "always_on_display";
     protected static final String DOZE_BRIGHTNESS_KEY = "doze_brightness";
     protected static final String WAKE_ON_GESTURE_KEY = "wake_on_gesture";