diff --git a/assets/js/theme.js b/assets/js/theme.js
index f26bb92..551d159 100644
--- a/assets/js/theme.js
+++ b/assets/js/theme.js
@@ -1,49 +1,22 @@
 // Light / Dark theme toggle
 (function () {
   const defaultTheme = '{{ site.Params.theme.default | default `system`}}'
+  const themes = ["light", "dark"];
 
   const themeToggleButtons = document.querySelectorAll(".hextra-theme-toggle");
   const themeToggleOptions = document.querySelectorAll(".hextra-theme-toggle-options p");
 
-  function setSystemTheme() {
-    const prefersColorScheme = window.matchMedia("(prefers-color-scheme: dark)").matches ? "dark" : "light";
-
-    document.documentElement.classList.remove("dark", "light");
-    document.documentElement.classList.add(prefersColorScheme);
-
-    document.documentElement.style.colorScheme = prefersColorScheme;
-  }
-
   function applyTheme(theme) {
+    theme = themes.includes(theme) ? theme : "system";
+
     themeToggleButtons.forEach((btn) => btn.parentElement.dataset.theme = theme );
 
     localStorage.setItem("color-theme", theme);
   }
 
   function switchTheme(theme) {
-     switch (theme) {
-      case "light":
-        document.documentElement.classList.remove("dark");
-        document.documentElement.classList.add(theme);
-        document.documentElement.style.colorScheme = theme;
-
-        applyTheme(theme);
-
-        break;
-      case "dark":
-        document.documentElement.classList.remove("light");
-        document.documentElement.classList.add(theme);
-        document.documentElement.style.colorScheme = theme;
-
-        applyTheme(theme);
-
-        break;
-      default:
-        setSystemTheme();
-        applyTheme("system");
-
-        break;
-    }
+    setTheme(theme);
+    applyTheme(theme);
   }
 
   const colorTheme = "color-theme" in localStorage ? localStorage.getItem("color-theme") : defaultTheme;
@@ -112,7 +85,7 @@
   // Listen for system theme changes
   window.matchMedia("(prefers-color-scheme: dark)").addEventListener("change", () => {
     if (localStorage.getItem("color-theme") === "system") {
-      setSystemTheme();
+      setTheme("system");
     }
   });
 })();
diff --git a/layouts/_partials/head.html b/layouts/_partials/head.html
index 20941fb..ba6f115 100644
--- a/layouts/_partials/head.html
+++ b/layouts/_partials/head.html
@@ -59,6 +59,24 @@
     {{ partial "google-analytics.html" . -}}
   {{- end }}
 
+  <script>
+    // The section must not be in the theme.js file because it can create a quick flash (switch between light and dark).
+
+    function setTheme(theme) {
+      document.documentElement.classList.remove("light", "dark");
+
+      if (theme !== "light" && theme !== "dark") {
+       theme = window.matchMedia("(prefers-color-scheme: dark)").matches ? "dark" : "light";
+      }
+
+      document.documentElement.classList.add(theme);
+      document.documentElement.style.colorScheme = theme;
+    }
+
+    setTheme("color-theme" in localStorage ? localStorage.getItem("color-theme") : '{{ site.Params.theme.default | default `system`}}')
+
+  </script>
+
   <!-- Math engine -->
   {{ $noop := .WordCount -}}
   {{- $engine := site.Params.math.engine | default "katex" -}}