From adecf9386535141e6d5b8a91937158ba0b69f75c Mon Sep 17 00:00:00 2001
From: Bryn Pickering <17178478+brynpickering@users.noreply.github.com>
Date: Wed, 27 Dec 2023 14:37:37 +0000
Subject: [PATCH] MathJax -> KaTeX

---
 docs/overrides/main.html                    | 29 +++++++--------------
 mkdocs.yml                                  |  3 ---
 src/calliope/backend/latex_backend_model.py | 26 +++++++++++++++++-
 3 files changed, 34 insertions(+), 24 deletions(-)

diff --git a/docs/overrides/main.html b/docs/overrides/main.html
index ba9d5087..f8b73196 100644
--- a/docs/overrides/main.html
+++ b/docs/overrides/main.html
@@ -20,24 +20,13 @@
 
     {{ super() }}
 
-    <script type="text/javascript">
-        if (!('mathjaxConfig' in window)) {
-            window.MathJax = {
-                tex: {
-                    inlineMath: [["\\(", "\\)"]],
-                    displayMath: [["\\[", "\\]"]],
-                    processEscapes: true,
-                    processEnvironments: true
-                },
-                options: {
-                    ignoreHtmlClass: ".*|",
-                    processHtmlClass: "arithmatex"
-                }
-            };
-
-            document$.subscribe(() => {
-                MathJax.typesetPromise()
-            })
-        }
-    </script>
+    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.css" integrity="sha384-n8MVd4RsNIU0tAv4ct0nTaAbDJwPJzDEaqSD1odI+WdtXRGWt2kTvGFasHpSy3SV" crossorigin="anonymous">
+
+    <!-- The loading of KaTeX is deferred to speed up page rendering -->
+    <script defer src="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/katex.min.js" integrity="sha384-XjKyOOlGwcjNTAIQHIpgOno0Hl1YQqzUOEleOLALmuqehneUG+vnGctmUb0ZY0l8" crossorigin="anonymous"></script>
+
+    <!-- To automatically render math in text elements, include the auto-render extension: -->
+    <script defer src="https://cdn.jsdelivr.net/npm/katex@0.16.9/dist/contrib/auto-render.min.js" integrity="sha384-+VBxd3r6XgURycqtZ117nYw44OOcIax56Z4dCRWbxyPt0Koah1uHoK0o4+/RRE05" crossorigin="anonymous"
+        onload="renderMathInElement(document.body);"></script>
+
 {% endblock %}
diff --git a/mkdocs.yml b/mkdocs.yml
index d1a34632..7c17a601 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -56,9 +56,6 @@ markdown_extensions:
       preserve_tabs: true
   - pymdownx.snippets
 
-extra_javascript:
-  - https://polyfill.io/v3/polyfill.min.js?features=es6
-  - https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js
 
 nav:
   - Home: index.md
diff --git a/src/calliope/backend/latex_backend_model.py b/src/calliope/backend/latex_backend_model.py
index 5c1c8c41..dc2e3525 100644
--- a/src/calliope/backend/latex_backend_model.py
+++ b/src/calliope/backend/latex_backend_model.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import re
 import textwrap
 import typing
 from pathlib import Path
@@ -254,7 +255,7 @@ class LatexBackendModel(backend_model.BackendModelGenerator):
     {% if equation.expression != "" %}
 
     $$
-    {{ equation.expression | trim }}
+    {{ equation.expression | trim | escape_underscores | escape_text_in_text }}
     $$
     {% endif %}
     {% endfor %}
@@ -484,8 +485,31 @@ def _generate_math_string(
 
     @staticmethod
     def _render(template: str, **kwargs) -> str:
+        text_starter = r"\\text(?:bf|it)?"  # match one of `\text`, `\textit`, `\textbf`
+
+        def __escape_underscore(instring):
+            "KaTeX requires underscores in `\text{...}` blocks to be escaped."
+            return re.sub(
+                rf"{text_starter}{{.*?}}",
+                lambda x: x.group(0).replace("_", r"\_"),
+                instring,
+            )
+
+        def __escape_text_in_text(instring):
+            """KaTeX requires `\text{...}` blocks within `\text{...}` blocks to be placed within math blocks.
+
+            We use `\\(` as the math block descriptor.
+            """
+            return re.sub(
+                rf"{text_starter}{{(?:[^{{}}]*({text_starter}{{.*?}})[^{{}}]*?)?}}",
+                lambda x: x.group(0).replace(x.group(1), rf"\({x.group(1)}\)"),
+                instring,
+            )
+
         jinja_env = jinja2.Environment(trim_blocks=True, autoescape=False)
         jinja_env.filters["removesuffix"] = lambda val, remove: val.removesuffix(remove)
+        jinja_env.filters["escape_underscores"] = __escape_underscore
+        jinja_env.filters["escape_text_in_text"] = __escape_text_in_text
         return jinja_env.from_string(template).render(**kwargs)
 
     def _get_capacity_bounds(